rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-07-07 06:20:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add remote dynamic filter frontend registration (#8148)

Browse Source 
* feat: filter id

Signed-off-by: discord9 <discord9@163.com>

* feat: dyn filter registry

Signed-off-by: discord9 <discord9@163.com>

* feat: filter id&refactor to type

Signed-off-by: discord9 <discord9@163.com>

* feat: merge scan register dyn filter(not send yet)

Signed-off-by: discord9 <discord9@163.com>

* feat: init reg dyn filter

Signed-off-by: discord9 <discord9@163.com>

* wip: remote dyn filter task 03

Signed-off-by: discord9 <discord9@163.com>

* fix: resolve remote dyn filter rebase fallout

Signed-off-by: discord9 <discord9@163.com>

* chore: keep remote dyn filter docs local

Signed-off-by: discord9 <discord9@163.com>

* chore: remove stale filter id allow

Signed-off-by: discord9 <discord9@163.com>

* chore: clippy

Signed-off-by: discord9 <discord9@163.com>

* chore: fix remote dyn filter import style

Signed-off-by: discord9 <discord9@163.com>

* chore: fix query metrics test fallout

Signed-off-by: discord9 <discord9@163.com>

* fix: exclude region from remote dyn filter id

Signed-off-by: discord9 <discord9@163.com>

* chore: import

Signed-off-by: discord9 <discord9@163.com>

* refactor: rm some to latter

Signed-off-by: discord9 <discord9@163.com>

* feat: add initial dyn filter snapshot

Signed-off-by: discord9 <discord9@163.com>

* refactor: per review

Signed-off-by: discord9 <discord9@163.com>

* docs: better comment, rm some slop

Signed-off-by: discord9 <discord9@163.com>

* chore: per review

Signed-off-by: discord9 <discord9@163.com>

---------

Signed-off-by: discord9 <discord9@163.com>
This commit is contained in:
discord9
2026-06-04 20:31:42 +08:00
committed by GitHub
parent 42eeeaa514
commit 6b7772e457
25 changed files with 2964 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
Cargo.lock generated
View File

@@ -11464,6 +11464,7 @@ dependencies = [
"datafusion-functions",
"datafusion-optimizer",
"datafusion-physical-expr",
"datafusion-proto",
"datafusion-sql",
"datatypes",
"either",

67
src/client/src/region.rs
View File

@@ -293,11 +293,8 @@ impl RegionRequester {
query_id: impl Into<String>,
update: RemoteDynFilterUpdate,
) -> Result<RegionResponse> {
self.handle_inner(build_remote_dyn_filter_request(
query_id.into(),
remote_dyn_filter_request::Action::Update(update),
))
.await
self.handle_inner(build_remote_dyn_filter_update_request(query_id, update))
.await
}
pub async fn handle_remote_dyn_filter_unregister(
@@ -305,14 +302,33 @@ impl RegionRequester {
query_id: impl Into<String>,
unregister: RemoteDynFilterUnregister,
) -> Result<RegionResponse> {
self.handle_inner(build_remote_dyn_filter_request(
query_id.into(),
remote_dyn_filter_request::Action::Unregister(unregister),
self.handle_inner(build_remote_dyn_filter_unregister_request(
query_id, unregister,
))
.await
}
}
pub fn build_remote_dyn_filter_update_request(
query_id: impl Into<String>,
update: RemoteDynFilterUpdate,
) -> RegionRequest {
build_remote_dyn_filter_request(
query_id.into(),
remote_dyn_filter_request::Action::Update(update),
)
}
pub fn build_remote_dyn_filter_unregister_request(
query_id: impl Into<String>,
unregister: RemoteDynFilterUnregister,
) -> RegionRequest {
build_remote_dyn_filter_request(
query_id.into(),
remote_dyn_filter_request::Action::Unregister(unregister),
)
}
fn build_remote_dyn_filter_request(
query_id: String,
action: remote_dyn_filter_request::Action,
@@ -357,7 +373,9 @@ pub fn check_response_header(header: &Option<ResponseHeader>) -> Result<()> {
#[cfg(test)]
mod test {
use api::v1::Status as PbStatus;
use api::v1::region::{RemoteDynFilterUpdate, region_request, remote_dyn_filter_request};
use api::v1::region::{
RemoteDynFilterUnregister, RemoteDynFilterUpdate, region_request, remote_dyn_filter_request,
};
use super::*;
use crate::Error::{IllegalDatabaseResponse, Server};
@@ -410,14 +428,14 @@ mod test {
#[test]
fn test_build_remote_dyn_filter_request_sets_header_and_body() {
let request = build_remote_dyn_filter_request(
"query-1".to_string(),
remote_dyn_filter_request::Action::Update(RemoteDynFilterUpdate {
let request = build_remote_dyn_filter_update_request(
"query-1",
RemoteDynFilterUpdate {
filter_id: "filter-1".to_string(),
payload: vec![1, 2, 3],
generation: 7,
is_complete: false,
}),
},
);
request.header.expect("remote dyn filter header must exist");
@@ -433,4 +451,27 @@ mod test {
Some(remote_dyn_filter_request::Action::Update(_))
));
}
#[test]
fn test_build_remote_dyn_filter_unregister_request_sets_header_and_body() {
let request = build_remote_dyn_filter_unregister_request(
"query-1",
RemoteDynFilterUnregister {
filter_id: "filter-9".to_string(),
},
);
request.header.expect("remote dyn filter header must exist");
let body = request.body.expect("remote dyn filter body must exist");
let region_request::Body::RemoteDynFilter(remote_request) = body else {
panic!("expected remote dyn filter request body");
};
assert_eq!(remote_request.query_id, "query-1");
assert!(matches!(
remote_request.action,
Some(remote_dyn_filter_request::Action::Unregister(_))
));
}
}

1
src/common/query/Cargo.toml
View File

@@ -28,6 +28,7 @@ datatypes.workspace = true
once_cell.workspace = true
prost.workspace = true
serde.workspace = true
serde_json.workspace = true
snafu.workspace = true
sqlparser.workspace = true
sqlparser_derive = "0.1"

68
src/common/query/src/request.rs
View File

@@ -12,10 +12,12 @@
// See the License for the specific language governing permissions and
// limitations under the License.
mod base64_serde;
mod initial_remote_dyn_filter_reg;
use std::sync::Arc;
use api::v1::region::RegionRequestHeader;
use datafusion::arrow::datatypes::Schema;
use datafusion::execution::TaskContext;
use datafusion::physical_expr::expressions::Column;
use datafusion::physical_plan::PhysicalExpr;
@@ -32,6 +34,12 @@ use serde::{Deserialize, Serialize};
use store_api::storage::RegionId;
/// Current wire-format version for remote dynamic filter payload updates.
pub use self::initial_remote_dyn_filter_reg::{
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY,
INITIAL_REMOTE_DYN_FILTER_REGS_MAX_TOTAL_PROTO_BYTES, InitialDynFilterReg,
InitialDynFilterRegs, InitialDynFilterSnapshot,
};
pub const DYN_FILTER_PROTOCOL_VERSION: u32 = 1;
/// Serialized predicate payload for remote dynamic filter updates.
@@ -49,31 +57,7 @@ pub const DYN_FILTER_PROTOCOL_VERSION: u32 = 1;
pub enum DynFilterPayload {
/// A serialized DataFusion [`PhysicalExpr`] encoded as a protobuf
/// [`PhysicalExprNode`].
Datafusion(#[serde(with = "base64_bytes")] Vec<u8>),
}
mod base64_bytes {
use base64::Engine;
use base64::prelude::BASE64_STANDARD;
use serde::de::Error;
use serde::{Deserialize, Deserializer, Serializer};
pub fn serialize<S>(bytes: &[u8], serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(&BASE64_STANDARD.encode(bytes))
}
pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
where
D: Deserializer<'de>,
{
let encoded = String::deserialize(deserializer)?;
BASE64_STANDARD.decode(encoded).map_err(|err| {
D::Error::custom(format!("invalid base64 dynamic filter payload: {err}"))
})
}
Datafusion(#[serde(with = "base64_serde::bytes")] Vec<u8>),
}
impl DynFilterPayload {
@@ -107,7 +91,7 @@ impl DynFilterPayload {
pub fn decode_datafusion_expr(
&self,
task_ctx: &TaskContext,
input_schema: &Schema,
input_schema: &datafusion::arrow::datatypes::Schema,
max_payload_bytes: usize,
) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
let Self::Datafusion(bytes) = self;
@@ -124,6 +108,34 @@ impl DynFilterPayload {
}
}
fn encode_physical_expr_to_bytes(expr: &Arc<dyn PhysicalExpr>) -> DataFusionResult<Vec<u8>> {
let codec = DefaultPhysicalExtensionCodec {};
let proto = serialize_physical_expr(expr, &codec)?;
let mut bytes = Vec::new();
proto.encode(&mut bytes).map_err(|e| {
DataFusionError::Internal(format!("Failed to encode PhysicalExprNode: {e}"))
})?;
Ok(bytes)
}
pub(crate) fn decode_physical_expr_from_bytes(
bytes: &[u8],
task_ctx: &TaskContext,
input_schema: &datafusion::arrow::datatypes::Schema,
max_payload_bytes: usize,
) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
validate_payload_size(bytes.len(), max_payload_bytes)?;
let codec = DefaultPhysicalExtensionCodec {};
let proto = PhysicalExprNode::decode(bytes).map_err(|e| {
DataFusionError::Internal(format!("Failed to decode PhysicalExprNode: {e}"))
})?;
let expr = parse_physical_expr(&proto, task_ctx, input_schema, &codec)?;
validate_supported_payload_expr(&expr)?;
validate_decoded_payload_expr(&expr, input_schema)?;
Ok(expr)
}
fn validate_payload_size(
payload_size_bytes: usize,
max_payload_bytes: usize,
@@ -161,7 +173,7 @@ fn validate_supported_payload_expr(expr: &Arc<dyn PhysicalExpr>) -> DataFusionRe
/// schema inconsistency that should be surfaced loudly.
fn validate_decoded_payload_expr(
expr: &Arc<dyn PhysicalExpr>,
input_schema: &Schema,
input_schema: &datafusion::arrow::datatypes::Schema,
) -> DataFusionResult<()> {
expr.apply(|node| {
if let Some(column) = node.as_any().downcast_ref::<Column>() {

78
src/common/query/src/request/base64_serde.rs Normal file
View File

@@ -0,0 +1,78 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//! Serde adapters for byte fields encoded as base64 strings in JSON.
use base64::Engine;
use base64::prelude::BASE64_STANDARD;
fn encode(bytes: &[u8]) -> String {
BASE64_STANDARD.encode(bytes)
}
fn decode(encoded: &str) -> Result<Vec<u8>, base64::DecodeError> {
BASE64_STANDARD.decode(encoded)
}
pub(crate) mod bytes {
use serde::de::Error;
use serde::{Deserialize, Deserializer, Serializer};
pub fn serialize<S>(value: &[u8], serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(&super::encode(value))
}
pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
where
D: Deserializer<'de>,
{
let encoded = String::deserialize(deserializer)?;
super::decode(&encoded).map_err(|err| {
D::Error::custom(format!("invalid base64 dynamic filter payload: {err}"))
})
}
}
pub(crate) mod bytes_vec {
use serde::de::Error;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
pub fn serialize<S>(values: &[Vec<u8>], serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
values
.iter()
.map(|bytes| super::encode(bytes))
.collect::<Vec<_>>()
.serialize(serializer)
}
pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<Vec<u8>>, D::Error>
where
D: Deserializer<'de>,
{
Vec::<String>::deserialize(deserializer)?
.into_iter()
.map(|encoded| {
super::decode(&encoded).map_err(|error| {
D::Error::custom(format!("invalid base64 bytes vector item: {error}"))
})
})
.collect()
}
}

397
src/common/query/src/request/initial_remote_dyn_filter_reg.rs Normal file
View File

@@ -0,0 +1,397 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashSet;
use std::sync::Arc;
use datafusion::arrow::datatypes::Schema;
use datafusion::execution::TaskContext;
use datafusion::physical_plan::PhysicalExpr;
use datafusion_common::Result as DataFusionResult;
use serde::{Deserialize, Serialize};
use crate::request::{
DynFilterPayload, decode_physical_expr_from_bytes, encode_physical_expr_to_bytes,
};
pub const INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY: &str =
"initial_remote_dyn_filter_registrations";
pub const INITIAL_REMOTE_DYN_FILTER_REGS_MAX_COUNT: usize = 64;
/// Raw encoded registration byte budget for initial remote dynamic filter registrations.
///
/// Counts proto payload bytes before JSON/base64 expansion, not the final extension size.
pub const INITIAL_REMOTE_DYN_FILTER_REGS_MAX_TOTAL_PROTO_BYTES: usize = 64 * 1024;
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct InitialDynFilterRegs {
#[serde(rename = "registrations")]
pub regs: Vec<InitialDynFilterReg>,
}
impl InitialDynFilterRegs {
pub fn new(regs: Vec<InitialDynFilterReg>) -> Self {
Self { regs }
}
pub fn is_empty(&self) -> bool {
self.regs.is_empty()
}
pub fn total_encoded_registration_bytes(&self) -> usize {
self.regs
.iter()
.map(InitialDynFilterReg::encoded_registration_bytes)
.sum()
}
pub fn validate_default_bounds(&self) -> Result<(), String> {
self.validate_bounds(
INITIAL_REMOTE_DYN_FILTER_REGS_MAX_COUNT,
INITIAL_REMOTE_DYN_FILTER_REGS_MAX_TOTAL_PROTO_BYTES,
)
}
pub fn validate_bounds(
&self,
max_count: usize,
max_total_proto_bytes: usize,
) -> Result<(), String> {
if self.regs.len() > max_count {
return Err(format!(
"InitialDynFilterRegs contains {} registrations, which exceeds the configured limit of {}",
self.regs.len(),
max_count
));
}
let total_registration_bytes = self.total_encoded_registration_bytes();
if total_registration_bytes > max_total_proto_bytes {
return Err(format!(
"InitialDynFilterRegs contains {} total encoded registration bytes, which exceeds the configured limit of {}",
total_registration_bytes, max_total_proto_bytes
));
}
let mut seen_filter_ids = HashSet::with_capacity(self.regs.len());
for reg in &self.regs {
if !seen_filter_ids.insert(reg.filter_id.as_str()) {
return Err(format!(
"InitialDynFilterRegs contains duplicate filter_id '{}'",
reg.filter_id
));
}
}
Ok(())
}
pub fn to_extension_value(&self) -> serde_json::Result<String> {
serde_json::to_string(self)
}
pub fn from_extension_value(value: &str) -> serde_json::Result<Self> {
let regs = serde_json::from_str::<Self>(value)?;
regs.validate_default_bounds()
.map_err(serde::de::Error::custom)?;
Ok(regs)
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct InitialDynFilterReg {
pub filter_id: String,
#[serde(with = "super::base64_serde::bytes_vec")]
pub child_exprs_datafusion_proto: Vec<Vec<u8>>,
/// Optional producer-side predicate snapshot captured at initial registration time.
///
/// This is only an initial pending update for the remote runtime filter. It is not part of
/// registration identity; identity is carried by `filter_id` and child expressions.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub initial_snapshot: Option<InitialDynFilterSnapshot>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct InitialDynFilterSnapshot {
pub payload: DynFilterPayload,
/// Producer-side generation used to ignore stale snapshots.
pub generation: u64,
/// Whether this snapshot completes the dynamic filter stream.
pub is_complete: bool,
}
impl InitialDynFilterReg {
pub fn new(filter_id: impl Into<String>, child_exprs_datafusion_proto: Vec<Vec<u8>>) -> Self {
Self {
filter_id: filter_id.into(),
child_exprs_datafusion_proto,
initial_snapshot: None,
}
}
pub fn with_initial_snapshot(mut self, initial_snapshot: InitialDynFilterSnapshot) -> Self {
self.initial_snapshot = Some(initial_snapshot);
self
}
pub fn from_filter_id_and_children(
filter_id: impl Into<String>,
children: &[Arc<dyn PhysicalExpr>],
) -> DataFusionResult<Self> {
let child_exprs_datafusion_proto = children
.iter()
.map(encode_physical_expr_to_bytes)
.collect::<DataFusionResult<Vec<_>>>()?;
Ok(Self::new(filter_id, child_exprs_datafusion_proto))
}
pub fn encoded_child_expr_bytes(&self) -> usize {
self.child_exprs_datafusion_proto.iter().map(Vec::len).sum()
}
pub fn encoded_registration_bytes(&self) -> usize {
self.encoded_child_expr_bytes()
+ self
.initial_snapshot
.as_ref()
.map(InitialDynFilterSnapshot::encoded_payload_bytes)
.unwrap_or(0)
}
pub fn decode_children(
&self,
task_ctx: &TaskContext,
input_schema: &Schema,
max_payload_bytes: usize,
) -> DataFusionResult<Vec<Arc<dyn PhysicalExpr>>> {
self.child_exprs_datafusion_proto
.iter()
.map(|expr_bytes| {
decode_physical_expr_from_bytes(
expr_bytes,
task_ctx,
input_schema,
max_payload_bytes,
)
})
.collect::<DataFusionResult<Vec<_>>>()
}
}
impl InitialDynFilterSnapshot {
pub fn new(payload: DynFilterPayload, generation: u64, is_complete: bool) -> Self {
Self {
payload,
generation,
is_complete,
}
}
pub fn encoded_payload_bytes(&self) -> usize {
match &self.payload {
DynFilterPayload::Datafusion(bytes) => bytes.len(),
}
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::physical_expr::expressions::Column;
use datafusion::physical_plan::PhysicalExpr;
use datafusion_common::DataFusionError;
use super::*;
#[test]
fn initial_dyn_filter_regs_json_round_trip() {
let regs = InitialDynFilterRegs::new(vec![
InitialDynFilterReg::new("filter-a", vec![vec![1, 2, 3]]),
InitialDynFilterReg::new("filter-b", vec![vec![4, 5]]),
]);
let encoded = regs.to_extension_value().unwrap();
let json: serde_json::Value = serde_json::from_str(&encoded).unwrap();
let decoded = InitialDynFilterRegs::from_extension_value(&encoded).unwrap();
assert_eq!(
json["registrations"][0]["child_exprs_datafusion_proto"],
serde_json::json!(["AQID"])
);
assert_eq!(
json["registrations"][1]["child_exprs_datafusion_proto"],
serde_json::json!(["BAU="])
);
assert_eq!(decoded, regs);
}
#[test]
fn initial_dyn_filter_regs_json_round_trip_with_snapshot() {
let regs = InitialDynFilterRegs::new(vec![
InitialDynFilterReg::new("filter-a", vec![vec![1, 2, 3]]).with_initial_snapshot(
InitialDynFilterSnapshot::new(DynFilterPayload::Datafusion(vec![4, 5, 6]), 7, true),
),
]);
let encoded = regs.to_extension_value().unwrap();
let json: serde_json::Value = serde_json::from_str(&encoded).unwrap();
let decoded = InitialDynFilterRegs::from_extension_value(&encoded).unwrap();
assert_eq!(
json["registrations"][0]["child_exprs_datafusion_proto"],
serde_json::json!(["AQID"])
);
assert_eq!(
json["registrations"][0]["initial_snapshot"]["payload"],
serde_json::json!({"kind":"datafusion","payload":"BAUG"})
);
assert_eq!(decoded, regs);
assert_eq!(
decoded.regs[0]
.initial_snapshot
.as_ref()
.unwrap()
.generation,
7
);
assert!(
decoded.regs[0]
.initial_snapshot
.as_ref()
.unwrap()
.is_complete
);
}
#[test]
fn initial_dyn_filter_reg_json_defaults_missing_snapshot_to_none() {
let decoded = InitialDynFilterRegs::from_extension_value(
r#"{"registrations":[{"filter_id":"filter-a","child_exprs_datafusion_proto":["AQID"]}]}"#,
)
.unwrap();
assert_eq!(decoded.regs.len(), 1);
assert!(decoded.regs[0].initial_snapshot.is_none());
}
#[test]
fn initial_dyn_filter_reg_encoded_registration_bytes_include_snapshot_payload() {
let reg = InitialDynFilterReg::new("filter-a", vec![vec![1, 2, 3], vec![4]])
.with_initial_snapshot(InitialDynFilterSnapshot::new(
DynFilterPayload::Datafusion(vec![5, 6]),
2,
false,
));
assert_eq!(reg.encoded_child_expr_bytes(), 4);
assert_eq!(reg.encoded_registration_bytes(), 6);
}
#[test]
fn initial_dyn_filter_regs_validate_bounds_rejects_duplicate_filter_ids() {
let regs = InitialDynFilterRegs::new(vec![
InitialDynFilterReg::new("filter-a", vec![vec![1]]),
InitialDynFilterReg::new("filter-a", vec![vec![2]]),
]);
let err = regs.validate_bounds(8, 1024).unwrap_err();
assert!(err.contains("duplicate filter_id 'filter-a'"));
}
#[test]
fn initial_dyn_filter_regs_from_extension_value_validates_default_bounds() {
let value = r#"{"registrations":[{"filter_id":"filter-a","child_exprs_datafusion_proto":["AQ=="]},{"filter_id":"filter-a","child_exprs_datafusion_proto":["Ag=="]}]}"#;
let err = InitialDynFilterRegs::from_extension_value(value).unwrap_err();
assert!(err.to_string().contains("duplicate filter_id 'filter-a'"));
}
#[test]
fn initial_dyn_filter_regs_validate_bounds_rejects_too_many_regs() {
let regs = InitialDynFilterRegs::new(vec![
InitialDynFilterReg::new("filter-a", vec![vec![1]]),
InitialDynFilterReg::new("filter-b", vec![vec![2]]),
]);
let err = regs.validate_bounds(1, 1024).unwrap_err();
assert!(err.contains("exceeds the configured limit of 1"));
}
#[test]
fn initial_dyn_filter_regs_validate_bounds_rejects_total_proto_bytes_over_limit() {
let regs = InitialDynFilterRegs::new(vec![
InitialDynFilterReg::new("filter-a", vec![vec![1, 2, 3]]),
InitialDynFilterReg::new("filter-b", vec![vec![4, 5, 6]]),
]);
let err = regs.validate_bounds(8, 5).unwrap_err();
assert!(err.contains("6 total encoded registration bytes"));
}
#[test]
fn initial_dyn_filter_regs_validate_bounds_rejects_snapshot_bytes_over_limit() {
let regs = InitialDynFilterRegs::new(vec![
InitialDynFilterReg::new("filter-a", vec![vec![1]]).with_initial_snapshot(
InitialDynFilterSnapshot::new(
DynFilterPayload::Datafusion(vec![2, 3, 4]),
2,
false,
),
),
]);
let err = regs.validate_bounds(8, 3).unwrap_err();
assert!(err.contains("4 total encoded registration bytes"));
}
#[test]
fn initial_dyn_filter_reg_round_trips_child_exprs() {
let schema = Schema::new(vec![Field::new("host", DataType::Utf8, false)]);
let child: Arc<dyn PhysicalExpr> =
Arc::new(Column::new_with_schema("host", &schema).unwrap());
let reg = InitialDynFilterReg::from_filter_id_and_children("filter-1", &[child]).unwrap();
let decoded = reg
.decode_children(&TaskContext::default(), &schema, 1024)
.unwrap();
let decoded = decoded[0].as_any().downcast_ref::<Column>().unwrap();
assert_eq!(reg.filter_id, "filter-1");
assert_eq!(decoded.name(), "host");
assert_eq!(decoded.index(), 0);
}
#[test]
fn initial_dyn_filter_reg_decode_rejects_column_name_index_mismatch() {
let schema = Schema::new(vec![Field::new("host", DataType::Utf8, false)]);
let reg = InitialDynFilterReg::from_filter_id_and_children(
"filter-1",
&[Arc::new(Column::new("service", 0)) as Arc<dyn PhysicalExpr>],
)
.unwrap();
let err = reg
.decode_children(&TaskContext::default(), &schema, 1024)
.unwrap_err();
assert!(matches!(err, DataFusionError::Plan(_)));
}
}

1
src/datanode/src/region_server.rs
View File

@@ -298,7 +298,6 @@ impl RegionServer {
query_ctx.clone(),
)
.await?;
let stream = wrap_flow_region_watermark_stream(stream, region_id, &query_ctx);
Ok(maybe_guard_stream(stream, permit))
}

75
src/frontend/src/instance/region_query.rs
View File

@@ -14,7 +14,11 @@
use std::sync::Arc;
use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate};
use async_trait::async_trait;
use client::region::{
build_remote_dyn_filter_unregister_request, build_remote_dyn_filter_update_request,
};
use common_error::ext::BoxedError;
use common_meta::node_manager::NodeManagerRef;
use common_query::request::QueryRequest;
@@ -24,6 +28,7 @@ use query::error::{RegionQuerySnafu, Result as QueryResult};
use query::region_query::RegionQueryHandler;
use session::ReadPreference;
use snafu::ResultExt;
use store_api::storage::RegionId;
use crate::error::{FindRegionPeerSnafu, RequestQuerySnafu, Result};
@@ -56,6 +61,30 @@ impl RegionQueryHandler for FrontendRegionQueryHandler {
.map_err(BoxedError::new)
.context(RegionQuerySnafu)
}
async fn handle_remote_dyn_filter_update(
&self,
region_id: RegionId,
query_id: String,
update: RemoteDynFilterUpdate,
) -> QueryResult<()> {
self.handle_remote_dyn_filter_update_inner(region_id, query_id, update)
.await
.map_err(BoxedError::new)
.context(RegionQuerySnafu)
}
async fn handle_remote_dyn_filter_unregister(
&self,
region_id: RegionId,
query_id: String,
unregister: RemoteDynFilterUnregister,
) -> QueryResult<()> {
self.handle_remote_dyn_filter_unregister_inner(region_id, query_id, unregister)
.await
.map_err(BoxedError::new)
.context(RegionQuerySnafu)
}
}
impl FrontendRegionQueryHandler {
@@ -82,4 +111,50 @@ impl FrontendRegionQueryHandler {
.await
.context(RequestQuerySnafu)
}
async fn handle_remote_dyn_filter_update_inner(
&self,
region_id: RegionId,
query_id: String,
update: RemoteDynFilterUpdate,
) -> Result<()> {
let peer = &self
.partition_manager
.find_region_leader(region_id)
.await
.context(FindRegionPeerSnafu {
region_id,
read_preference: ReadPreference::Leader,
})?;
let client = self.node_manager.datanode(peer).await;
client
.handle(build_remote_dyn_filter_update_request(query_id, update))
.await
.context(RequestQuerySnafu)?;
Ok(())
}
async fn handle_remote_dyn_filter_unregister_inner(
&self,
region_id: RegionId,
query_id: String,
unregister: RemoteDynFilterUnregister,
) -> Result<()> {
let peer = &self
.partition_manager
.find_region_leader(region_id)
.await
.context(FindRegionPeerSnafu {
region_id,
read_preference: ReadPreference::Leader,
})?;
let client = self.node_manager.datanode(peer).await;
client
.handle(build_remote_dyn_filter_unregister_request(
query_id, unregister,
))
.await
.context(RequestQuerySnafu)?;
Ok(())
}
}

1
src/query/Cargo.toml
View File

@@ -44,6 +44,7 @@ datafusion-expr-common.workspace = true
datafusion-functions.workspace = true
datafusion-optimizer.workspace = true
datafusion-physical-expr.workspace = true
datafusion-proto.workspace = true
datafusion-sql.workspace = true
datatypes.workspace = true
either.workspace = true

1
src/query/src/datafusion.rs
View File

@@ -475,6 +475,7 @@ impl QueryEngine for DatafusionQueryEngine {
fn engine_context(&self, query_ctx: QueryContextRef) -> QueryEngineContext {
let mut state = self.state.session_state();
state.config_mut().set_extension(query_ctx.clone());
state.config_mut().set_extension(self.state.clone());
// note that hints in "x-greptime-hints" is automatically parsed
// and set to query context's extension, so we can get it from query context.
if let Some(parallelism) = query_ctx.extension(QUERY_PARALLELISM_HINT) {

8
src/query/src/dist_plan.rs
View File

@@ -14,14 +14,22 @@
mod analyzer;
mod commutativity;
mod dyn_filter_bridge;
mod filter_id;
mod merge_scan;
mod merge_sort;
mod planner;
mod predicate_extractor;
mod region_pruner;
mod remote_dyn_filter_registry;
pub use analyzer::{DistPlannerAnalyzer, DistPlannerOptions};
pub use filter_id::{FilterFingerprint, FilterId, ParseFilterIdError, RemoteDynFilterProducerId};
pub use merge_scan::{MergeScanExec, MergeScanLogicalPlan};
pub use planner::{DistExtensionPlanner, MergeSortExtensionPlanner};
pub use predicate_extractor::PredicateExtractor;
pub use region_pruner::ConstraintPruner;
pub use remote_dyn_filter_registry::{
DynFilterEntry, DynFilterRegistryManager, EntryRegistration, QueryDynFilterRegistry,
RemoteDynFilterRegistryLease, Subscriber, SubscriberRegistration,
};

56
src/query/src/dist_plan/analyzer.rs
View File

@@ -31,6 +31,7 @@ use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
use table::metadata::TableType;
use table::table::adapter::DfTableProviderAdapter;
use crate::dist_plan::RemoteDynFilterProducerId;
use crate::dist_plan::analyzer::utils::{
PatchOptimizerContext, PlanTreeExpressionSimplifier, aliased_columns_for,
rewrite_merge_sort_exprs,
@@ -146,14 +147,14 @@ impl DistPlannerAnalyzer {
let plan = plan.transform(&Self::inspect_plan_with_subquery)?;
let mut rewriter = PlanRewriter::default();
let result = plan.data.rewrite(&mut rewriter)?.data;
Ok(result)
Self::assign_merge_scan_remote_dyn_filter_producer_ids(result)
}
/// Use fallback plan rewriter to rewrite the plan and only push down table scan nodes
fn use_fallback(&self, plan: LogicalPlan) -> DfResult<LogicalPlan> {
let mut rewriter = fallback::FallbackPlanRewriter;
let result = plan.rewrite(&mut rewriter)?.data;
Ok(result)
Self::assign_merge_scan_remote_dyn_filter_producer_ids(result)
}
fn inspect_plan_with_subquery(plan: LogicalPlan) -> DfResult<Transformed<LogicalPlan>> {
@@ -224,6 +225,57 @@ impl DistPlannerAnalyzer {
spans: Default::default(),
})
}
fn assign_merge_scan_remote_dyn_filter_producer_ids(
plan: LogicalPlan,
) -> DfResult<LogicalPlan> {
let mut assigner = MergeScanRemoteDynFilterProducerIdAssigner::default();
Ok(plan.rewrite_with_subqueries(&mut assigner)?.data)
}
}
#[derive(Debug, Default)]
struct RemoteDynFilterProducerIdAllocator {
next_remote_dyn_filter_producer_id: u64,
}
impl RemoteDynFilterProducerIdAllocator {
fn allocate(&mut self) -> RemoteDynFilterProducerId {
self.next_remote_dyn_filter_producer_id += 1;
RemoteDynFilterProducerId::new(self.next_remote_dyn_filter_producer_id)
}
}
/// Assigns query-local RDF producer ids to visible `MergeScan` nodes after plan rewriting.
#[derive(Debug, Default)]
struct MergeScanRemoteDynFilterProducerIdAssigner {
remote_dyn_filter_producer_id_allocator: RemoteDynFilterProducerIdAllocator,
}
impl TreeNodeRewriter for MergeScanRemoteDynFilterProducerIdAssigner {
type Node = LogicalPlan;
fn f_up(&mut self, node: Self::Node) -> DfResult<Transformed<Self::Node>> {
let LogicalPlan::Extension(extension) = &node else {
return Ok(Transformed::no(node));
};
let Some(merge_scan) = extension
.node
.as_any()
.downcast_ref::<MergeScanLogicalPlan>()
else {
return Ok(Transformed::no(node));
};
Ok(Transformed::yes(
merge_scan
.clone()
.with_remote_dyn_filter_producer_id(
self.remote_dyn_filter_producer_id_allocator.allocate(),
)
.into_logical_plan(),
))
}
}
/// Status of the rewriter to mark if the current pass is expanded

137
src/query/src/dist_plan/analyzer/test.rs
View File

@@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::BTreeSet;
use std::pin::Pin;
use std::sync::Arc;
@@ -29,9 +30,10 @@ use datafusion::functions_aggregate::expr_fn::avg;
use datafusion::functions_aggregate::min_max::{max, min};
use datafusion::prelude::SessionContext;
use datafusion_common::{JoinType, ScalarValue};
use datafusion_expr::expr::ScalarFunction;
use datafusion_expr::expr::{Exists, ScalarFunction};
use datafusion_expr::{
AggregateUDF, Expr, ExprSchemable as _, LogicalPlanBuilder, Operator, binary_expr, col, lit,
AggregateUDF, Expr, ExprSchemable as _, LogicalPlanBuilder, Operator, Subquery, binary_expr,
col, lit,
};
use datafusion_functions::datetime::date_bin;
use datafusion_functions::datetime::expr_fn::now;
@@ -53,6 +55,50 @@ use table::{Table, TableRef};
use super::*;
fn collect_merge_scan_remote_dyn_filter_producer_ids(
plan: &LogicalPlan,
producer_ids: &mut BTreeSet<RemoteDynFilterProducerId>,
) {
let mut producer_id_list = Vec::new();
collect_merge_scan_remote_dyn_filter_producer_id_list(plan, &mut producer_id_list);
producer_ids.extend(producer_id_list);
}
struct MergeScanRemoteDynFilterProducerIdCollector<'a> {
producer_ids: &'a mut Vec<RemoteDynFilterProducerId>,
}
impl TreeNodeRewriter for MergeScanRemoteDynFilterProducerIdCollector<'_> {
type Node = LogicalPlan;
fn f_up(&mut self, node: Self::Node) -> DfResult<Transformed<Self::Node>> {
if let LogicalPlan::Extension(extension) = &node
&& let Some(merge_scan) = extension
.node
.as_any()
.downcast_ref::<MergeScanLogicalPlan>()
{
self.producer_ids.push(
merge_scan
.remote_dyn_filter_producer_id()
.expect("MergeScan remote dynamic filter producer id must be assigned"),
);
}
Ok(Transformed::no(node))
}
}
fn collect_merge_scan_remote_dyn_filter_producer_id_list(
plan: &LogicalPlan,
producer_ids: &mut Vec<RemoteDynFilterProducerId>,
) {
let _ = plan
.clone()
.rewrite_with_subqueries(&mut MergeScanRemoteDynFilterProducerIdCollector { producer_ids })
.unwrap();
}
pub(crate) struct TestTable;
impl TestTable {
@@ -1360,6 +1406,60 @@ fn test_simplify_select_now_expression() {
assert_eq!(expected, normalized);
}
#[test]
fn sibling_merge_scans_have_unique_remote_dyn_filter_producer_ids() {
init_default_ut_logging();
let left_table = TestTable::table_with_name(0, "left_table".to_string());
let right_table = TestTable::table_with_name(1, "right_table".to_string());
let left_source = Arc::new(DefaultTableSource::new(Arc::new(
DfTableProviderAdapter::new(left_table),
)));
let right_source = Arc::new(DefaultTableSource::new(Arc::new(
DfTableProviderAdapter::new(right_table),
)));
let left_sorted =
LogicalPlanBuilder::scan_with_filters("left_table", left_source, None, vec![])
.unwrap()
.sort(vec![col("pk1").sort(true, false)])
.unwrap()
.build()
.unwrap();
let right_sorted =
LogicalPlanBuilder::scan_with_filters("right_table", right_source, None, vec![])
.unwrap()
.sort(vec![col("pk1").sort(true, false)])
.unwrap()
.build()
.unwrap();
let plan = LogicalPlanBuilder::from(left_sorted)
.cross_join(right_sorted)
.unwrap()
.build()
.unwrap();
let config = ConfigOptions::default();
let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
let mut producer_ids = Vec::new();
collect_merge_scan_remote_dyn_filter_producer_id_list(&result, &mut producer_ids);
let unique_producer_ids = producer_ids.iter().copied().collect::<BTreeSet<_>>();
assert!(
producer_ids.len() >= 2,
"Expected at least 2 RemoteDynFilterProducerIds, got {}: {producer_ids:?}",
producer_ids.len()
);
assert_eq!(
producer_ids.len(),
unique_producer_ids.len(),
"Expected all sibling RemoteDynFilterProducerIds to be unique, got ids: {producer_ids:?}"
);
}
#[test]
fn test_simplify_now_expression() {
init_default_ut_logging();
@@ -1823,6 +1923,39 @@ fn transform_sort_subquery_alias() {
assert_eq!(expected, result.to_string());
}
#[test]
fn remote_dyn_filter_producer_ids_do_not_collide_between_subquery_and_outer_plan() {
let test_table = TestTable::table_with_name(0, "numbers".to_string());
let table_source = Arc::new(DefaultTableSource::new(Arc::new(
DfTableProviderAdapter::new(test_table),
)));
let subquery_plan =
LogicalPlanBuilder::scan_with_filters("inner", table_source.clone(), None, vec![])
.unwrap()
.build()
.unwrap();
let subquery = Subquery {
subquery: Arc::new(subquery_plan),
outer_ref_columns: Default::default(),
spans: Default::default(),
};
let outer_plan = LogicalPlanBuilder::scan_with_filters("outer", table_source, None, vec![])
.unwrap()
.filter(Expr::Exists(Exists {
subquery,
negated: false,
}))
.unwrap()
.build()
.unwrap();
let rewritten = DistPlannerAnalyzer {}.try_push_down(outer_plan).unwrap();
let mut producer_ids = BTreeSet::new();
collect_merge_scan_remote_dyn_filter_producer_ids(&rewritten, &mut producer_ids);
assert_eq!(producer_ids.len(), 2);
}
#[test]
fn date_bin_ts_group_by() {
init_default_ut_logging();

665
src/query/src/dist_plan/dyn_filter_bridge.rs Normal file
View File

@@ -0,0 +1,665 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::any::Any;
use std::sync::Arc;
use common_query::request::{
DynFilterPayload, INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY,
INITIAL_REMOTE_DYN_FILTER_REGS_MAX_TOTAL_PROTO_BYTES, InitialDynFilterReg,
InitialDynFilterRegs, InitialDynFilterSnapshot,
};
use datafusion_common::Result;
use datafusion_physical_expr::PhysicalExpr;
use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr;
use session::context::{QueryContext, QueryContextRef};
use store_api::storage::RegionId;
use crate::dist_plan::filter_id::build_remote_dyn_filter_id;
use crate::dist_plan::{FilterId, QueryDynFilterRegistry, RemoteDynFilterProducerId, Subscriber};
#[derive(Debug, Clone)]
pub(crate) struct CapturedDynFilter {
filter_id: FilterId,
initial_registration: InitialDynFilterReg,
pub(crate) alive_dyn_filter: Arc<DynamicFilterPhysicalExpr>,
}
#[derive(Debug, Clone)]
pub(crate) struct RemoteDynFilterPushdown {
pub(crate) captured_dyn_filters: Vec<CapturedDynFilter>,
/// Preflight result per parent filter.
pub(crate) pushed_down: Vec<bool>,
}
pub(crate) fn capture_remote_dyn_filters_for_pushdown(
remote_dyn_filter_producer_id: RemoteDynFilterProducerId,
parent_filters: Vec<Arc<dyn datafusion::physical_plan::PhysicalExpr>>,
) -> RemoteDynFilterPushdown {
let mut pushed_down = Vec::with_capacity(parent_filters.len());
let mut captured_dyn_filters = Vec::new();
for (producer_local_ordinal, filter) in parent_filters.into_iter().enumerate() {
let Some(alive_dyn_filter) = downcast_dynamic_filter(filter) else {
pushed_down.push(false);
continue;
};
match build_captured_dyn_filter(
remote_dyn_filter_producer_id,
producer_local_ordinal,
alive_dyn_filter,
) {
Ok(captured_dyn_filter) => {
pushed_down.push(true);
captured_dyn_filters.push(captured_dyn_filter);
}
Err(error) => {
common_telemetry::warn!(error; "Remote dyn filter is not pushed down because initial registration cannot be built");
pushed_down.push(false);
}
}
}
if let Err(error) = validate_initial_registrations_for_pushdown(&captured_dyn_filters) {
common_telemetry::warn!(error; "Remote dyn filters are not pushed down because initial registrations are invalid");
return RemoteDynFilterPushdown {
captured_dyn_filters: Vec::new(),
pushed_down: vec![false; pushed_down.len()],
};
}
RemoteDynFilterPushdown {
captured_dyn_filters,
pushed_down,
}
}
fn downcast_dynamic_filter(
expr: Arc<dyn datafusion::physical_plan::PhysicalExpr>,
) -> Option<Arc<DynamicFilterPhysicalExpr>> {
(expr as Arc<dyn Any + Send + Sync + 'static>)
.downcast::<DynamicFilterPhysicalExpr>()
.ok()
}
pub(crate) fn register_dyn_filters_for_region(
registry: &QueryDynFilterRegistry,
region_id: RegionId,
captured_dyn_filters: &[CapturedDynFilter],
) {
for captured_dyn_filter in captured_dyn_filters {
let _ = registry.register_remote_dyn_filter(
captured_dyn_filter.filter_id.clone(),
captured_dyn_filter.alive_dyn_filter.clone(),
);
let _ = registry
.register_subscriber(&captured_dyn_filter.filter_id, Subscriber::new(region_id));
}
}
fn build_captured_dyn_filter(
remote_dyn_filter_producer_id: RemoteDynFilterProducerId,
producer_local_ordinal: usize,
alive_dyn_filter: Arc<DynamicFilterPhysicalExpr>,
) -> Result<CapturedDynFilter> {
let children = alive_dyn_filter
.children()
.into_iter()
.cloned()
.collect::<Vec<_>>();
let filter_id = build_remote_dyn_filter_id(
remote_dyn_filter_producer_id,
producer_local_ordinal,
&children,
)?;
let initial_registration =
InitialDynFilterReg::from_filter_id_and_children(filter_id.to_string(), &children)?;
Ok(CapturedDynFilter {
filter_id,
initial_registration: attach_initial_snapshot(initial_registration, &alive_dyn_filter),
alive_dyn_filter,
})
}
fn validate_initial_registrations_for_pushdown(
captured_dyn_filters: &[CapturedDynFilter],
) -> std::result::Result<(), String> {
let regs = build_initial_dyn_filter_regs_for_region(captured_dyn_filters);
regs.validate_default_bounds()?;
regs.to_extension_value()
.map_err(|error| error.to_string())?;
Ok(())
}
fn attach_initial_snapshot(
initial_registration: InitialDynFilterReg,
alive_dyn_filter: &DynamicFilterPhysicalExpr,
) -> InitialDynFilterReg {
let Some(initial_snapshot) = initial_snapshot(alive_dyn_filter) else {
return initial_registration;
};
initial_registration.with_initial_snapshot(initial_snapshot)
}
fn initial_snapshot(
alive_dyn_filter: &DynamicFilterPhysicalExpr,
) -> Option<InitialDynFilterSnapshot> {
let generation = alive_dyn_filter.snapshot_generation();
let current = match alive_dyn_filter.current() {
Ok(current) => current,
Err(error) => {
common_telemetry::warn!(error; "Failed to read remote dyn filter initial snapshot");
return None;
}
};
let payload = match DynFilterPayload::from_datafusion_expr(
&current,
INITIAL_REMOTE_DYN_FILTER_REGS_MAX_TOTAL_PROTO_BYTES,
) {
Ok(payload) => payload,
Err(error) => {
common_telemetry::warn!(error; "Failed to encode remote dyn filter initial snapshot");
return None;
}
};
// Current DataFusion exposes `wait_complete()`, but no non-blocking completion getter.
let is_complete = false;
Some(InitialDynFilterSnapshot::new(
payload,
generation,
is_complete,
))
}
fn build_initial_dyn_filter_regs_for_region(
captured_dyn_filters: &[CapturedDynFilter],
) -> InitialDynFilterRegs {
InitialDynFilterRegs::new(
captured_dyn_filters
.iter()
.map(|captured| captured.initial_registration.clone())
.collect(),
)
}
pub(crate) fn query_context_with_initial_dyn_filter_regs(
query_ctx: &QueryContextRef,
region_id: RegionId,
captured_dyn_filters: &[CapturedDynFilter],
) -> QueryContext {
let mut region_query_ctx = query_ctx.as_ref().clone();
let regs = build_initial_dyn_filter_regs_for_region(captured_dyn_filters);
if regs.is_empty() {
return region_query_ctx;
}
if let Err(error) = regs.validate_default_bounds() {
common_telemetry::warn!(error; "Dropping initial remote dyn filter registrations for region {} that exceed configured bounds", region_id);
return region_query_ctx;
}
match regs.to_extension_value() {
Ok(serialized) => region_query_ctx.set_extension(
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY,
serialized,
),
Err(error) => {
common_telemetry::warn!(error; "Failed to serialize initial remote dyn filter registrations");
}
}
region_query_ctx
}
#[cfg(test)]
mod tests {
use std::fmt;
use std::hash::{Hash, Hasher};
use datafusion::execution::TaskContext;
use datafusion_common::ScalarValue;
use datafusion_expr::ColumnarValue;
use datafusion_physical_expr::expressions::{Column, lit};
use session::query_id::QueryId;
use uuid::Uuid;
use super::*;
#[derive(Debug)]
struct UnserializableExpr;
impl fmt::Display for UnserializableExpr {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "unserializable_expr")
}
}
impl Hash for UnserializableExpr {
fn hash<H: Hasher>(&self, state: &mut H) {
"unserializable_expr".hash(state);
}
}
impl PartialEq for UnserializableExpr {
fn eq(&self, _other: &Self) -> bool {
true
}
}
impl Eq for UnserializableExpr {}
impl datafusion_physical_expr::PhysicalExpr for UnserializableExpr {
fn as_any(&self) -> &dyn Any {
self
}
fn data_type(
&self,
_input_schema: &arrow_schema::Schema,
) -> datafusion_common::Result<arrow_schema::DataType> {
Ok(arrow_schema::DataType::Boolean)
}
fn nullable(
&self,
_input_schema: &arrow_schema::Schema,
) -> datafusion_common::Result<bool> {
Ok(false)
}
fn evaluate(
&self,
_batch: &common_recordbatch::DfRecordBatch,
) -> datafusion_common::Result<ColumnarValue> {
Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))))
}
fn children(&self) -> Vec<&Arc<dyn datafusion_physical_expr::PhysicalExpr>> {
Vec::new()
}
fn with_new_children(
self: Arc<Self>,
_children: Vec<Arc<dyn datafusion_physical_expr::PhysicalExpr>>,
) -> datafusion_common::Result<Arc<dyn datafusion_physical_expr::PhysicalExpr>> {
Ok(self)
}
fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{self}")
}
}
fn test_query_id(value: u128) -> QueryId {
QueryId::from(Uuid::from_u128(value))
}
fn test_remote_dyn_filter_producer_id(value: u64) -> RemoteDynFilterProducerId {
RemoteDynFilterProducerId::new(value)
}
fn test_captured_dyn_filter(
remote_dyn_filter_producer_id: RemoteDynFilterProducerId,
producer_local_ordinal: usize,
column_name: &str,
column_index: usize,
) -> CapturedDynFilter {
build_captured_dyn_filter(
remote_dyn_filter_producer_id,
producer_local_ordinal,
Arc::new(DynamicFilterPhysicalExpr::new(
vec![Arc::new(Column::new(column_name, column_index)) as Arc<_>],
lit(true) as _,
)),
)
.unwrap()
}
fn test_dyn_filter_with_snapshot_payload(
column_name: &str,
column_index: usize,
payload_bytes: usize,
) -> Arc<DynamicFilterPhysicalExpr> {
let dyn_filter = Arc::new(DynamicFilterPhysicalExpr::new(
vec![Arc::new(Column::new(column_name, column_index)) as Arc<_>],
lit(true) as _,
));
dyn_filter
.update(lit(ScalarValue::Utf8(Some("x".repeat(payload_bytes)))) as _)
.unwrap();
dyn_filter
}
#[test]
fn capture_remote_dyn_filters_for_pushdown_preserves_parent_filter_ordinals() {
let parent_filters = vec![
Arc::new(Column::new("service", 0)) as Arc<dyn datafusion::physical_plan::PhysicalExpr>,
Arc::new(DynamicFilterPhysicalExpr::new(
vec![Arc::new(Column::new("host", 1)) as Arc<_>],
lit(true) as _,
)) as Arc<dyn datafusion::physical_plan::PhysicalExpr>,
Arc::new(Column::new("zone", 2)) as Arc<dyn datafusion::physical_plan::PhysicalExpr>,
Arc::new(DynamicFilterPhysicalExpr::new(
vec![Arc::new(Column::new("pod", 3)) as Arc<_>],
lit(true) as _,
)) as Arc<dyn datafusion::physical_plan::PhysicalExpr>,
];
let remote_dyn_filter_producer_id = test_remote_dyn_filter_producer_id(42);
let captured =
capture_remote_dyn_filters_for_pushdown(remote_dyn_filter_producer_id, parent_filters)
.captured_dyn_filters;
assert_eq!(captured.len(), 2);
assert_eq!(
captured[0].filter_id.remote_dyn_filter_producer_id(),
remote_dyn_filter_producer_id
);
assert_eq!(
captured[1].filter_id.remote_dyn_filter_producer_id(),
remote_dyn_filter_producer_id
);
assert_eq!(captured[0].filter_id.producer_ordinal(), 1);
assert_eq!(captured[1].filter_id.producer_ordinal(), 3);
}
#[test]
fn capture_remote_dyn_filters_for_pushdown_marks_only_valid_initial_regs() {
let parent_filters = vec![
Arc::new(Column::new("service", 0)) as Arc<dyn datafusion::physical_plan::PhysicalExpr>,
Arc::new(DynamicFilterPhysicalExpr::new(
vec![Arc::new(Column::new("host", 1)) as Arc<_>],
lit(true) as _,
)) as Arc<dyn datafusion::physical_plan::PhysicalExpr>,
Arc::new(Column::new("zone", 2)) as Arc<dyn datafusion::physical_plan::PhysicalExpr>,
];
let remote_dyn_filter_producer_id = test_remote_dyn_filter_producer_id(42);
let pushdown =
capture_remote_dyn_filters_for_pushdown(remote_dyn_filter_producer_id, parent_filters);
assert_eq!(pushdown.pushed_down, vec![false, true, false]);
assert_eq!(pushdown.captured_dyn_filters.len(), 1);
assert_eq!(
pushdown.captured_dyn_filters[0]
.filter_id
.remote_dyn_filter_producer_id(),
remote_dyn_filter_producer_id
);
assert_eq!(
pushdown.captured_dyn_filters[0]
.filter_id
.producer_ordinal(),
1
);
assert!(
pushdown.captured_dyn_filters[0]
.initial_registration
.initial_snapshot
.is_some()
);
}
#[test]
fn capture_remote_dyn_filters_for_pushdown_rejects_unencodable_registration() {
let parent_filters = vec![Arc::new(DynamicFilterPhysicalExpr::new(
vec![Arc::new(UnserializableExpr) as Arc<_>],
lit(true) as _,
))
as Arc<dyn datafusion::physical_plan::PhysicalExpr>];
let pushdown = capture_remote_dyn_filters_for_pushdown(
test_remote_dyn_filter_producer_id(42),
parent_filters,
);
assert_eq!(pushdown.pushed_down, vec![false]);
assert!(pushdown.captured_dyn_filters.is_empty());
}
#[test]
fn capture_remote_dyn_filters_for_pushdown_attaches_initial_snapshot() {
let parent_filters = vec![Arc::new(DynamicFilterPhysicalExpr::new(
vec![Arc::new(Column::new("host", 1)) as Arc<_>],
lit(true) as _,
))
as Arc<dyn datafusion::physical_plan::PhysicalExpr>];
let pushdown = capture_remote_dyn_filters_for_pushdown(
test_remote_dyn_filter_producer_id(42),
parent_filters,
);
assert_eq!(pushdown.pushed_down, vec![true]);
assert!(
pushdown.captured_dyn_filters[0]
.initial_registration
.initial_snapshot
.is_some()
);
}
#[test]
fn capture_remote_dyn_filters_for_pushdown_attaches_initial_snapshot_after_update() {
let dyn_filter = Arc::new(DynamicFilterPhysicalExpr::new(
vec![Arc::new(Column::new("host", 1)) as Arc<_>],
lit(true) as _,
));
dyn_filter.update(lit(false) as _).unwrap();
let parent_filters = vec![dyn_filter as Arc<dyn datafusion::physical_plan::PhysicalExpr>];
let pushdown = capture_remote_dyn_filters_for_pushdown(
test_remote_dyn_filter_producer_id(42),
parent_filters,
);
assert_eq!(pushdown.pushed_down, vec![true]);
let snapshot = pushdown.captured_dyn_filters[0]
.initial_registration
.initial_snapshot
.as_ref()
.unwrap();
assert_eq!(snapshot.generation, 2);
assert!(!snapshot.is_complete);
assert!(matches!(
snapshot.payload,
DynFilterPayload::Datafusion(ref bytes) if !bytes.is_empty()
));
}
#[test]
fn capture_remote_dyn_filters_for_pushdown_rejects_oversized_snapshots() {
let parent_filters = vec![
test_dyn_filter_with_snapshot_payload("host", 0, 40 * 1024)
as Arc<dyn datafusion::physical_plan::PhysicalExpr>,
test_dyn_filter_with_snapshot_payload("pod", 1, 40 * 1024)
as Arc<dyn datafusion::physical_plan::PhysicalExpr>,
];
let pushdown = capture_remote_dyn_filters_for_pushdown(
test_remote_dyn_filter_producer_id(42),
parent_filters,
);
assert_eq!(pushdown.pushed_down, vec![false, false]);
assert!(pushdown.captured_dyn_filters.is_empty());
}
#[test]
fn capture_remote_dyn_filters_for_pushdown_rejects_too_many_regs_with_snapshots() {
const TOO_MANY_INITIAL_REGS: usize = 65;
let parent_filters = (0..TOO_MANY_INITIAL_REGS)
.map(|ordinal| {
test_dyn_filter_with_snapshot_payload(&format!("host_{ordinal}"), ordinal, 1)
as Arc<dyn datafusion::physical_plan::PhysicalExpr>
})
.collect::<Vec<_>>();
let pushdown = capture_remote_dyn_filters_for_pushdown(
test_remote_dyn_filter_producer_id(42),
parent_filters,
);
assert!(pushdown.captured_dyn_filters.is_empty());
assert_eq!(pushdown.pushed_down, vec![false; TOO_MANY_INITIAL_REGS]);
}
#[test]
fn capture_remote_dyn_filters_for_pushdown_rejects_regs_exceeding_bounds() {
const TOO_MANY_INITIAL_REGS: usize = 65;
let parent_filters = (0..TOO_MANY_INITIAL_REGS)
.map(|_| {
Arc::new(DynamicFilterPhysicalExpr::new(
vec![Arc::new(Column::new("host", 0)) as Arc<_>],
lit(true) as _,
)) as Arc<dyn datafusion::physical_plan::PhysicalExpr>
})
.collect::<Vec<_>>();
let pushdown = capture_remote_dyn_filters_for_pushdown(
test_remote_dyn_filter_producer_id(42),
parent_filters,
);
assert!(pushdown.captured_dyn_filters.is_empty());
assert_eq!(pushdown.pushed_down, vec![false; TOO_MANY_INITIAL_REGS]);
}
#[test]
fn register_dyn_filters_for_region_reuses_existing_entry() {
let registry = QueryDynFilterRegistry::new(test_query_id(1));
let captured_dyn_filters = vec![test_captured_dyn_filter(
test_remote_dyn_filter_producer_id(42),
2,
"host",
0,
)];
let first_region_id = RegionId::new(1024, 7);
let second_region_id = RegionId::new(1024, 8);
register_dyn_filters_for_region(&registry, first_region_id, &captured_dyn_filters);
register_dyn_filters_for_region(&registry, second_region_id, &captured_dyn_filters);
assert_eq!(registry.entry_count(), 1);
let entry = registry.entries().pop().unwrap();
assert_eq!(
entry.filter_id().remote_dyn_filter_producer_id(),
test_remote_dyn_filter_producer_id(42)
);
assert_eq!(entry.filter_id().producer_ordinal(), 2);
let subscribers = entry.subscribers();
assert_eq!(subscribers.len(), 2);
assert!(
subscribers
.iter()
.any(|subscriber| subscriber.region_id() == first_region_id)
);
assert!(
subscribers
.iter()
.any(|subscriber| subscriber.region_id() == second_region_id)
);
}
#[test]
fn register_dyn_filters_for_region_keeps_independent_producer_ids_distinct() {
let registry = QueryDynFilterRegistry::new(test_query_id(1));
let region_id = RegionId::new(1024, 7);
let make_filter = |remote_dyn_filter_producer_id| {
test_captured_dyn_filter(remote_dyn_filter_producer_id, 2, "host", 0)
};
register_dyn_filters_for_region(
&registry,
region_id,
&[make_filter(test_remote_dyn_filter_producer_id(42))],
);
register_dyn_filters_for_region(
&registry,
region_id,
&[make_filter(test_remote_dyn_filter_producer_id(43))],
);
assert_eq!(registry.entry_count(), 2);
}
#[test]
fn query_context_includes_region_initial_dyn_filter_regs() {
let captured_dyn_filters = vec![test_captured_dyn_filter(
test_remote_dyn_filter_producer_id(42),
2,
"host",
0,
)];
let region_id = RegionId::new(1024, 7);
let query_ctx = QueryContext::arc();
let region_query_ctx = query_context_with_initial_dyn_filter_regs(
&query_ctx,
region_id,
&captured_dyn_filters,
);
let extension = region_query_ctx
.extension(INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY)
.unwrap();
let regs = InitialDynFilterRegs::from_extension_value(extension).unwrap();
let decoded_children = regs.regs[0]
.decode_children(
&TaskContext::default(),
&arrow_schema::Schema::new(vec![arrow_schema::Field::new(
"host",
arrow_schema::DataType::Utf8,
false,
)]),
1024,
)
.unwrap();
assert_eq!(regs.regs.len(), 1);
assert_eq!(
regs.regs[0].filter_id,
captured_dyn_filters[0].filter_id.to_string()
);
assert_eq!(decoded_children.len(), 1);
assert!(decoded_children[0].as_any().is::<Column>());
}
#[test]
fn query_context_drops_initial_regs_when_duplicate_filter_ids_exceed_bounds() {
let captured_dyn_filters = vec![
test_captured_dyn_filter(test_remote_dyn_filter_producer_id(42), 2, "host", 0),
test_captured_dyn_filter(test_remote_dyn_filter_producer_id(42), 2, "host", 0),
];
let region_id = RegionId::new(1024, 7);
let query_ctx = QueryContext::arc();
let region_query_ctx = query_context_with_initial_dyn_filter_regs(
&query_ctx,
region_id,
&captured_dyn_filters,
);
assert!(
region_query_ctx
.extension(INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY)
.is_none()
);
}
}

322
src/query/src/dist_plan/filter_id.rs Normal file
View File

@@ -0,0 +1,322 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::fmt::{Display, Formatter};
use std::hash::{DefaultHasher, Hasher};
use std::num::ParseIntError;
use std::str::FromStr;
use std::sync::Arc;
use datafusion_common::{DataFusionError, Result};
use datafusion_physical_expr::PhysicalExpr;
use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec;
use datafusion_proto::physical_plan::to_proto::serialize_physical_expr;
use prost::Message;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct FilterFingerprint(u64);
impl FilterFingerprint {
pub fn new(value: u64) -> Self {
Self(value)
}
pub fn get(self) -> u64 {
self.0
}
}
impl Display for FilterFingerprint {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "{:016x}", self.0)
}
}
impl FromStr for FilterFingerprint {
type Err = ParseIntError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
Ok(Self(u64::from_str_radix(s, 16)?))
}
}
/// Query-local identity for one remote dynamic filter producer.
///
/// Distinguishes independent producers(MergeScan node) that may share ordinals and child fingerprints.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct RemoteDynFilterProducerId(u64);
impl RemoteDynFilterProducerId {
pub fn new(value: u64) -> Self {
Self(value)
}
pub fn get(self) -> u64 {
self.0
}
}
impl Display for RemoteDynFilterProducerId {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "{:016x}", self.0)
}
}
impl FromStr for RemoteDynFilterProducerId {
type Err = ParseIntError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
Ok(Self(u64::from_str_radix(s, 16)?))
}
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct FilterId {
remote_dyn_filter_producer_id: RemoteDynFilterProducerId,
producer_ordinal: u32,
children_fingerprint: FilterFingerprint,
}
// NOTE(remote-dyn-filter): FilterId is source-generated and propagated; consumers should not
// recompute it from local scan state.
impl FilterId {
pub fn new(
remote_dyn_filter_producer_id: RemoteDynFilterProducerId,
producer_ordinal: u32,
children_fingerprint: FilterFingerprint,
) -> Self {
Self {
remote_dyn_filter_producer_id,
producer_ordinal,
children_fingerprint,
}
}
pub fn producer_ordinal(&self) -> u32 {
self.producer_ordinal
}
pub fn remote_dyn_filter_producer_id(&self) -> RemoteDynFilterProducerId {
self.remote_dyn_filter_producer_id
}
pub fn children_fingerprint(&self) -> FilterFingerprint {
self.children_fingerprint
}
}
impl Display for FilterId {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{}:{}:{}",
self.remote_dyn_filter_producer_id, self.producer_ordinal, self.children_fingerprint
)
}
}
impl FromStr for FilterId {
type Err = ParseFilterIdError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
let mut parts = s.split(':');
let remote_dyn_filter_producer_id = parts
.next()
.ok_or(ParseFilterIdError)?
.parse::<RemoteDynFilterProducerId>()
.map_err(|_| ParseFilterIdError)?;
let producer_local_ordinal = parts
.next()
.ok_or(ParseFilterIdError)?
.parse::<u32>()
.map_err(|_| ParseFilterIdError)?;
let children_fingerprint = parts
.next()
.ok_or(ParseFilterIdError)?
.parse::<FilterFingerprint>()
.map_err(|_| ParseFilterIdError)?;
if parts.next().is_some() {
return Err(ParseFilterIdError);
}
Ok(Self::new(
remote_dyn_filter_producer_id,
producer_local_ordinal,
children_fingerprint,
))
}
}
#[derive(Debug)]
pub struct ParseFilterIdError;
impl Display for ParseFilterIdError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "invalid filter id")
}
}
/// Builds the query-local remote dynamic filter identity.
///
/// Identity is producer id + local ordinal + child fingerprint; routing stays outside.
pub(crate) fn build_remote_dyn_filter_id(
remote_dyn_filter_producer_id: RemoteDynFilterProducerId,
producer_local_ordinal: usize,
children: &[Arc<dyn PhysicalExpr>],
) -> Result<FilterId> {
let children_fingerprint = canonicalize_dyn_filter_children(children)?;
let producer_local_ordinal = u32::try_from(producer_local_ordinal).map_err(|err| {
let _ = err;
DataFusionError::Execution("producer ordinal out of range for filter id".to_string())
})?;
Ok(FilterId::new(
remote_dyn_filter_producer_id,
producer_local_ordinal,
children_fingerprint,
))
}
fn canonicalize_dyn_filter_children(
children: &[Arc<dyn PhysicalExpr>],
) -> Result<FilterFingerprint> {
let codec = DefaultPhysicalExtensionCodec {};
let mut hasher = DefaultHasher::new();
hasher.write_usize(children.len());
for child in children {
let proto = serialize_physical_expr(child, &codec)?;
let mut bytes = Vec::new();
proto
.encode(&mut bytes)
.map_err(|e| DataFusionError::External(Box::new(e)))?;
hasher.write_usize(bytes.len());
hasher.write(&bytes);
}
Ok(FilterFingerprint::new(hasher.finish()))
}
#[cfg(test)]
mod tests {
use datafusion_physical_expr::expressions::Column;
use super::*;
fn test_children(names: &[&str]) -> Vec<Arc<dyn PhysicalExpr>> {
names
.iter()
.enumerate()
.map(|(index, name)| Arc::new(Column::new(name, index)) as Arc<dyn PhysicalExpr>)
.collect()
}
#[test]
fn filter_id_round_trips_through_string() {
let filter_id = FilterId::new(
RemoteDynFilterProducerId::new(42),
3,
FilterFingerprint::new(0xabc),
);
let encoded = filter_id.to_string();
assert_eq!(encoded, "000000000000002a:3:0000000000000abc");
assert_eq!(encoded.parse::<FilterId>().unwrap(), filter_id);
}
#[test]
fn filter_id_rejects_malformed_strings() {
assert!("".parse::<FilterId>().is_err());
assert!("0000000000000001:3".parse::<FilterId>().is_err());
assert!("0000000000000001:3:zzzz".parse::<FilterId>().is_err());
assert!(
"0000000000000001:3:0000000000000abc:extra"
.parse::<FilterId>()
.is_err()
);
}
#[test]
fn remote_dyn_filter_id_is_stable_for_equivalent_children() {
let remote_dyn_filter_producer_id = RemoteDynFilterProducerId::new(42);
let first = build_remote_dyn_filter_id(
remote_dyn_filter_producer_id,
3,
&test_children(&["host", "pod"]),
)
.unwrap();
let second = build_remote_dyn_filter_id(
remote_dyn_filter_producer_id,
3,
&test_children(&["host", "pod"]),
)
.unwrap();
assert_eq!(first, second);
}
#[test]
fn remote_dyn_filter_id_changes_when_remote_dyn_filter_producer_changes() {
let children = test_children(&["host", "pod"]);
let baseline =
build_remote_dyn_filter_id(RemoteDynFilterProducerId::new(42), 3, &children).unwrap();
let different_producer_id =
build_remote_dyn_filter_id(RemoteDynFilterProducerId::new(43), 3, &children).unwrap();
assert_ne!(baseline, different_producer_id);
}
#[test]
fn remote_dyn_filter_id_changes_when_identity_inputs_change() {
let children = test_children(&["host", "pod"]);
let remote_dyn_filter_producer_id = RemoteDynFilterProducerId::new(42);
let baseline =
build_remote_dyn_filter_id(remote_dyn_filter_producer_id, 3, &children).unwrap();
let different_ordinal =
build_remote_dyn_filter_id(remote_dyn_filter_producer_id, 4, &children).unwrap();
let different_children = build_remote_dyn_filter_id(
remote_dyn_filter_producer_id,
3,
&test_children(&["pod", "host"]),
)
.unwrap();
assert_ne!(baseline, different_ordinal);
assert_ne!(baseline, different_children);
}
#[test]
fn remote_dyn_filter_id_supports_empty_children() {
let remote_dyn_filter_producer_id = RemoteDynFilterProducerId::new(42);
let first = build_remote_dyn_filter_id(remote_dyn_filter_producer_id, 1, &[]).unwrap();
let second = build_remote_dyn_filter_id(remote_dyn_filter_producer_id, 1, &[]).unwrap();
assert_eq!(first, second);
assert_eq!(
first.remote_dyn_filter_producer_id(),
remote_dyn_filter_producer_id
);
assert_eq!(first.producer_ordinal(), 1);
assert_eq!(first.children_fingerprint(), second.children_fingerprint());
}
#[test]
fn remote_dyn_filter_id_rejects_out_of_range_producer_ordinal() {
let error = build_remote_dyn_filter_id(RemoteDynFilterProducerId::new(42), usize::MAX, &[])
.unwrap_err()
.to_string();
assert!(error.contains("producer ordinal out of range for filter id"));
}
}

341
src/query/src/dist_plan/merge_scan.rs
View File

@@ -26,6 +26,9 @@ use common_recordbatch::adapter::RecordBatchMetrics;
use common_telemetry::tracing_context::TracingContext;
use datafusion::execution::{SessionState, TaskContext};
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::filter_pushdown::{
ChildPushdownResult, FilterPushdownPhase, FilterPushdownPropagation, PushedDown,
};
use datafusion::physical_plan::metrics::{
Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, MetricsSet, Time,
};
@@ -50,10 +53,38 @@ use tracing::{Instrument, Span};
use crate::dist_plan::analyzer::AliasMapping;
use crate::dist_plan::analyzer::utils::patch_batch_timezone;
use crate::dist_plan::dyn_filter_bridge::{
CapturedDynFilter, capture_remote_dyn_filters_for_pushdown,
query_context_with_initial_dyn_filter_regs, register_dyn_filters_for_region,
};
use crate::dist_plan::{RemoteDynFilterProducerId, RemoteDynFilterRegistryLease};
use crate::metrics::{MERGE_SCAN_ERRORS_TOTAL, MERGE_SCAN_POLL_ELAPSED, MERGE_SCAN_REGIONS};
use crate::options::FlowQueryExtensions;
use crate::query_engine::QueryEngineState;
use crate::region_query::RegionQueryHandlerRef;
fn query_engine_state_from_task_context(context: &TaskContext) -> Option<Arc<QueryEngineState>> {
context.session_config().get_extension()
}
fn acquire_remote_dyn_filter_registry_lease(
context: &TaskContext,
query_ctx: &QueryContextRef,
captured_dyn_filters: &[CapturedDynFilter],
) -> Option<RemoteDynFilterRegistryLease> {
if captured_dyn_filters.is_empty() {
return None;
}
let query_id = query_ctx.remote_query_id_value()?;
let query_engine_state = query_engine_state_from_task_context(context)?;
Some(
query_engine_state
.dyn_filter_registry_manager()
.acquire_lease(query_id),
)
}
#[derive(Debug, Hash, PartialOrd, PartialEq, Eq, Clone)]
pub struct MergeScanLogicalPlan {
/// In logical plan phase it only contains one input
@@ -61,6 +92,8 @@ pub struct MergeScanLogicalPlan {
/// If this plan is a placeholder
is_placeholder: bool,
partition_cols: AliasMapping,
/// Assigned after dist-plan rewriting so rewriters only deal with plan shape.
remote_dyn_filter_producer_id: Option<RemoteDynFilterProducerId>,
}
impl UserDefinedLogicalNodeCore for MergeScanLogicalPlan {
@@ -106,9 +139,18 @@ impl MergeScanLogicalPlan {
input,
is_placeholder,
partition_cols,
remote_dyn_filter_producer_id: None,
}
}
pub(crate) fn with_remote_dyn_filter_producer_id(
mut self,
remote_dyn_filter_producer_id: RemoteDynFilterProducerId,
) -> Self {
self.remote_dyn_filter_producer_id = Some(remote_dyn_filter_producer_id);
self
}
pub fn name() -> &'static str {
"MergeScan"
}
@@ -131,8 +173,13 @@ impl MergeScanLogicalPlan {
pub fn partition_cols(&self) -> &AliasMapping {
&self.partition_cols
}
pub fn remote_dyn_filter_producer_id(&self) -> Option<RemoteDynFilterProducerId> {
self.remote_dyn_filter_producer_id
}
}
#[derive(Clone)]
pub struct MergeScanExec {
table: TableName,
regions: Vec<RegionId>,
@@ -146,6 +193,9 @@ pub struct MergeScanExec {
/// Metrics for each partition
partition_metrics: Arc<Mutex<HashMap<usize, PartitionMetrics>>>,
query_ctx: QueryContextRef,
/// Optional because RDF must fail open: missing ids skip RDF but keep normal query execution.
remote_dyn_filter_producer_id: Option<RemoteDynFilterProducerId>,
captured_remote_dyn_filters: Arc<Mutex<Vec<CapturedDynFilter>>>,
target_partition: usize,
partition_cols: AliasMapping,
}
@@ -172,6 +222,7 @@ impl MergeScanExec {
query_ctx: QueryContextRef,
target_partition: usize,
partition_cols: AliasMapping,
remote_dyn_filter_producer_id: Option<RemoteDynFilterProducerId>,
) -> Result<Self> {
// TODO(CookiePieWw): Initially we removed the metadata from the schema in #2000, but we have to
// keep it for #4619 to identify json type in src/datatypes/src/schema/column_schema.rs.
@@ -244,6 +295,8 @@ impl MergeScanExec {
partition_metrics: Arc::default(),
properties,
query_ctx,
remote_dyn_filter_producer_id,
captured_remote_dyn_filters: Arc::default(),
target_partition,
partition_cols,
})
@@ -264,13 +317,20 @@ impl MergeScanExec {
let partition_metrics_moved = self.partition_metrics.clone();
let plan = self.plan.clone();
let target_partition = self.target_partition;
let captured_remote_dyn_filters = self.captured_remote_dyn_filters();
let dbname = context.task_id().unwrap_or_default();
let tracing_context = TracingContext::from_json(context.session_id().as_str());
let current_channel = self.query_ctx.channel();
let read_preference = self.query_ctx.read_preference();
let explain_verbose = self.query_ctx.explain_verbose();
let remote_dyn_filter_registry_lease = acquire_remote_dyn_filter_registry_lease(
context.as_ref(),
&query_ctx,
&captured_remote_dyn_filters,
);
let stream = Box::pin(stream!({
let remote_dyn_filter_registry_lease = remote_dyn_filter_registry_lease;
// only report metrics once for each MergeScan
if partition == 0 {
MERGE_SCAN_REGIONS.observe(regions.len() as f64);
@@ -286,17 +346,32 @@ impl MergeScanExec {
.step_by(target_partition)
.copied()
{
if let Some(remote_dyn_filter_registry_lease) =
remote_dyn_filter_registry_lease.as_ref()
{
register_dyn_filters_for_region(
remote_dyn_filter_registry_lease.registry(),
region_id,
&captured_remote_dyn_filters,
);
}
let region_span = tracing_context.attach(tracing::info_span!(
parent: &Span::current(),
"merge_scan_region",
region_id = %region_id,
partition = partition
));
let region_query_ctx = query_context_with_initial_dyn_filter_regs(
&query_ctx,
region_id,
&captured_remote_dyn_filters,
);
let request = QueryRequest {
header: Some(RegionRequestHeader {
tracing_context: tracing_context.to_w3c(),
dbname: dbname.clone(),
query_context: Some(query_ctx.as_ref().into()),
query_context: Some((&region_query_ctx).into()),
}),
region_id,
plan: plan.clone(),
@@ -464,11 +539,17 @@ impl MergeScanExec {
sub_stage_metrics: self.sub_stage_metrics.clone(),
partition_metrics: self.partition_metrics.clone(),
query_ctx: self.query_ctx.clone(),
remote_dyn_filter_producer_id: self.remote_dyn_filter_producer_id,
captured_remote_dyn_filters: self.captured_remote_dyn_filters.clone(),
target_partition: self.target_partition,
partition_cols: self.partition_cols.clone(),
})
}
fn captured_remote_dyn_filters(&self) -> Vec<CapturedDynFilter> {
self.captured_remote_dyn_filters.lock().unwrap().clone()
}
pub fn sub_stage_metrics(&self) -> Vec<RecordBatchMetrics> {
self.sub_stage_metrics
.lock()
@@ -517,6 +598,13 @@ impl MergeScanExec {
}
}
#[cfg(test)]
impl MergeScanExec {
fn remote_dyn_filter_producer_id(&self) -> Option<RemoteDynFilterProducerId> {
self.remote_dyn_filter_producer_id
}
}
/// Metrics for a region of a partition.
#[derive(Debug, Clone)]
struct RegionMetrics {
@@ -624,6 +712,50 @@ impl ExecutionPlan for MergeScanExec {
Ok(self.clone())
}
fn handle_child_pushdown_result(
&self,
_phase: FilterPushdownPhase,
child_pushdown_result: ChildPushdownResult,
_config: &datafusion::config::ConfigOptions,
) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> {
let parent_filters = child_pushdown_result
.parent_filters
.into_iter()
.map(|filter| filter.filter)
.collect::<Vec<_>>();
let Some(remote_dyn_filter_producer_id) = self.remote_dyn_filter_producer_id else {
// Missing RDF identity disables only RDF, not normal execution.
common_telemetry::warn!(
"MergeScan remote dynamic filter producer id is not assigned; skipping remote dynamic filter pushdown"
);
self.captured_remote_dyn_filters.lock().unwrap().clear();
let new_self = Arc::new(self.clone());
return Ok(FilterPushdownPropagation {
filters: parent_filters.into_iter().map(|_| PushedDown::No).collect(),
updated_node: Some(new_self),
});
};
let remote_dyn_filter_pushdown =
capture_remote_dyn_filters_for_pushdown(remote_dyn_filter_producer_id, parent_filters);
*self.captured_remote_dyn_filters.lock().unwrap() =
remote_dyn_filter_pushdown.captured_dyn_filters;
let new_self = Arc::new(self.clone());
Ok(FilterPushdownPropagation {
filters: remote_dyn_filter_pushdown
.pushed_down
.into_iter()
.map(|_pushdown_ready| {
// TODO(discord9): Return `PushedDown::Yes` after datanodes consume RDF
// registrations and pending updates. Until then, keep the parent-side filter.
PushedDown::No
})
.collect(),
updated_node: Some(new_self),
})
}
fn execute(
&self,
partition: usize,
@@ -731,3 +863,210 @@ impl MergeScanMetric {
self.greptime_exec_cost.add(metrics);
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeSet;
use async_trait::async_trait;
use datafusion::config::ConfigOptions;
use datafusion::execution::SessionStateBuilder;
use datafusion::physical_plan::filter_pushdown::ChildFilterPushdownResult;
use datafusion_common::TableReference;
use datafusion_expr::{LogicalPlanBuilder, lit};
use datafusion_physical_expr::Distribution;
use datafusion_physical_expr::expressions::{
Column, DynamicFilterPhysicalExpr, lit as physical_lit,
};
use session::ReadPreference;
use session::context::QueryContext;
use session::query_id::QueryId;
use table::table_name::TableName;
use uuid::Uuid;
use super::*;
use crate::dist_plan::DynFilterRegistryManager;
use crate::region_query::RegionQueryHandler;
fn test_query_id(value: u128) -> QueryId {
QueryId::from(Uuid::from_u128(value))
}
#[test]
fn remote_dyn_filter_registry_cleanup_waits_for_last_query_scoped_stream_drop() {
let registry_manager = Arc::new(DynFilterRegistryManager::default());
let query_id = test_query_id(1);
let first = registry_manager.acquire_lease(query_id);
let second = registry_manager.acquire_lease(query_id);
drop(first);
assert_eq!(registry_manager.registry_count(), 1);
drop(second);
assert_eq!(registry_manager.registry_count(), 0);
}
#[test]
fn remote_dyn_filter_registry_cleanup_shares_query_scope_across_independent_leases() {
let registry_manager = Arc::new(DynFilterRegistryManager::default());
let query_id = test_query_id(1);
let first_exec_like_lease = registry_manager.acquire_lease(query_id);
let second_exec_like_lease = registry_manager.acquire_lease(query_id);
drop(first_exec_like_lease);
assert_eq!(registry_manager.registry_count(), 1);
drop(second_exec_like_lease);
assert_eq!(registry_manager.registry_count(), 0);
}
struct TestRegionQueryHandler;
#[async_trait]
impl RegionQueryHandler for TestRegionQueryHandler {
async fn do_get(
&self,
_read_preference: ReadPreference,
_request: common_query::request::QueryRequest,
) -> crate::error::Result<common_recordbatch::SendableRecordBatchStream> {
unimplemented!("test only")
}
async fn handle_remote_dyn_filter_update(
&self,
_region_id: RegionId,
_query_id: String,
_update: api::v1::region::RemoteDynFilterUpdate,
) -> crate::error::Result<()> {
unimplemented!("test only")
}
async fn handle_remote_dyn_filter_unregister(
&self,
_region_id: RegionId,
_query_id: String,
_unregister: api::v1::region::RemoteDynFilterUnregister,
) -> crate::error::Result<()> {
unimplemented!("test only")
}
}
#[test]
fn try_with_new_distribution_preserves_remote_dyn_filter_producer_id() {
let remote_dyn_filter_producer_id = RemoteDynFilterProducerId::new(42);
// Build a plan whose schema contains "col1"
let plan = LogicalPlanBuilder::empty(true)
.project(vec![lit(1i32).alias("col1")])
.unwrap()
.build()
.unwrap();
let schema = plan.schema().as_arrow().clone();
let table = TableName::new("catalog", "schema", "table");
let regions = vec![RegionId::new(1024, 1)];
let query_ctx = QueryContext::arc();
// Non-empty partition_cols so try_with_new_distribution can detect an overlap
let mut partition_cols = AliasMapping::new();
partition_cols.insert(
"col1".to_string(),
BTreeSet::from([ColumnExpr::new(Some(TableReference::bare("table")), "col1")]),
);
let session_state = SessionStateBuilder::new().build();
let handler = Arc::new(TestRegionQueryHandler);
let target_partition = 2;
let exec = MergeScanExec::new(
&session_state,
table,
regions,
plan,
&schema,
handler,
query_ctx,
target_partition,
partition_cols,
Some(remote_dyn_filter_producer_id),
)
.unwrap();
assert_eq!(
exec.remote_dyn_filter_producer_id(),
Some(remote_dyn_filter_producer_id)
);
// A distribution that differs from the current partitioning but shares a
// column name present in partition_cols, so try_with_new_distribution
// produces a clone instead of returning None.
let new_dist = Distribution::HashPartitioned(vec![
Arc::new(Column::new("col1", 0)),
Arc::new(Column::new("col2", 1)),
]);
let cloned = exec
.try_with_new_distribution(new_dist)
.expect("expected a cloned exec with overlapping partition col");
assert_eq!(
cloned.remote_dyn_filter_producer_id(),
Some(remote_dyn_filter_producer_id),
"try_with_new_distribution must preserve remote dynamic filter producer id"
);
}
#[test]
fn remote_dyn_filter_preflight_keeps_parent_filter_until_dn_runtime_is_ready() {
let remote_dyn_filter_producer_id = RemoteDynFilterProducerId::new(42);
let plan = LogicalPlanBuilder::empty(true)
.project(vec![lit(1i32).alias("col1")])
.unwrap()
.build()
.unwrap();
let schema = plan.schema().as_arrow().clone();
let table = TableName::new("catalog", "schema", "table");
let regions = vec![RegionId::new(1024, 1)];
let query_ctx = QueryContext::arc();
let session_state = SessionStateBuilder::new().build();
let handler = Arc::new(TestRegionQueryHandler);
let exec = MergeScanExec::new(
&session_state,
table,
regions,
plan,
&schema,
handler,
query_ctx,
1,
AliasMapping::new(),
Some(remote_dyn_filter_producer_id),
)
.unwrap();
let dyn_filter = Arc::new(DynamicFilterPhysicalExpr::new(
vec![Arc::new(Column::new("host", 0)) as Arc<_>],
physical_lit(true) as _,
)) as Arc<dyn datafusion_physical_expr::PhysicalExpr>;
let propagation = exec
.handle_child_pushdown_result(
FilterPushdownPhase::Post,
ChildPushdownResult {
parent_filters: vec![ChildFilterPushdownResult {
filter: dyn_filter,
child_results: vec![PushedDown::Yes],
}],
self_filters: Vec::new(),
},
&ConfigOptions::new(),
)
.unwrap();
assert_eq!(exec.captured_remote_dyn_filters().len(), 1);
assert!(matches!(propagation.filters.as_slice(), [PushedDown::No]));
}
}

1
src/query/src/dist_plan/planner.rs
View File

@@ -178,6 +178,7 @@ impl ExtensionPlanner for DistExtensionPlanner {
query_ctx,
session_state.config().target_partitions(),
merge_scan.partition_cols().clone(),
merge_scan.remote_dyn_filter_producer_id(),
)?;
Ok(Some(Arc::new(merge_scan_plan) as _))
}

578
src/query/src/dist_plan/remote_dyn_filter_registry.rs Normal file
View File

@@ -0,0 +1,578 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, RwLock, Weak};
use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr;
use session::query_id::QueryId;
use store_api::storage::RegionId;
use crate::dist_plan::FilterId;
/// Routing metadata for a remote dynamic filter subscriber.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Subscriber {
region_id: RegionId,
}
impl Subscriber {
pub fn new(region_id: RegionId) -> Self {
Self { region_id }
}
pub fn region_id(&self) -> RegionId {
self.region_id
}
}
/// Result of registering a remote dynamic filter entry.
#[derive(Debug, Clone)]
pub enum EntryRegistration {
Inserted(Arc<DynFilterEntry>),
/// The filter already existed; this contains the previously registered entry.
Existing(Arc<DynFilterEntry>),
}
/// Result of registering a subscriber under an existing filter entry.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SubscriberRegistration {
Added,
Duplicate,
MissingFilter,
}
/// A registered query-local remote dynamic filter entry.
///
/// The frontend query owns the strong DataFusion filter handle until the query finishes; the
/// registry only keeps a weak reference for later updates.
#[derive(Debug)]
pub struct DynFilterEntry {
filter_id: FilterId,
alive_dyn_filter: Weak<DynamicFilterPhysicalExpr>,
subscribers: RwLock<HashSet<Subscriber>>,
}
#[derive(Debug)]
struct QueryDynFilterRegistryInner {
entries: HashMap<FilterId, Arc<DynFilterEntry>>,
}
impl DynFilterEntry {
pub fn new(filter_id: FilterId, alive_dyn_filter: Arc<DynamicFilterPhysicalExpr>) -> Self {
Self {
filter_id,
alive_dyn_filter: Arc::downgrade(&alive_dyn_filter),
subscribers: RwLock::new(HashSet::new()),
}
}
pub fn filter_id(&self) -> &FilterId {
&self.filter_id
}
pub fn upgrade_alive_dyn_filter(&self) -> Option<Arc<DynamicFilterPhysicalExpr>> {
self.alive_dyn_filter.upgrade()
}
pub fn subscribers(&self) -> Vec<Subscriber> {
self.subscribers.read().unwrap().iter().cloned().collect()
}
pub fn register_subscriber(&self, subscriber: Subscriber) -> bool {
let mut subscribers = self.subscribers.write().unwrap();
subscribers.insert(subscriber)
}
}
/// Query-scoped registry that owns all remote dynamic filters for one query.
#[derive(Debug)]
pub struct QueryDynFilterRegistry {
query_id: QueryId,
inner: RwLock<QueryDynFilterRegistryInner>,
}
impl QueryDynFilterRegistry {
pub fn new(query_id: QueryId) -> Self {
Self {
query_id,
inner: RwLock::new(QueryDynFilterRegistryInner {
entries: HashMap::new(),
}),
}
}
pub fn query_id(&self) -> QueryId {
self.query_id
}
pub fn entry_count(&self) -> usize {
self.inner.read().unwrap().entries.len()
}
pub fn entries(&self) -> Vec<Arc<DynFilterEntry>> {
self.inner
.read()
.unwrap()
.entries
.values()
.cloned()
.collect()
}
pub fn remote_dyn_filter(&self, filter_id: &FilterId) -> Option<Arc<DynFilterEntry>> {
self.inner.read().unwrap().entries.get(filter_id).cloned()
}
pub fn register_remote_dyn_filter(
&self,
filter_id: FilterId,
alive_dyn_filter: Arc<DynamicFilterPhysicalExpr>,
) -> EntryRegistration {
let mut inner = self.inner.write().unwrap();
if let Some(existing) = inner.entries.get(&filter_id) {
return EntryRegistration::Existing(existing.clone());
}
let entry = Arc::new(DynFilterEntry::new(filter_id.clone(), alive_dyn_filter));
inner.entries.insert(filter_id, entry.clone());
EntryRegistration::Inserted(entry)
}
pub fn register_subscriber(
&self,
filter_id: &FilterId,
subscriber: Subscriber,
) -> SubscriberRegistration {
let Some(entry) = self.inner.read().unwrap().entries.get(filter_id).cloned() else {
return SubscriberRegistration::MissingFilter;
};
if entry.register_subscriber(subscriber) {
SubscriberRegistration::Added
} else {
SubscriberRegistration::Duplicate
}
}
}
/// Stream-scoped lease that keeps a query registry alive.
///
/// Production code owns registries through this lease; the manager only keeps a weak index.
#[derive(Debug)]
pub struct RemoteDynFilterRegistryLease {
registry_manager: Arc<DynFilterRegistryManager>,
/// Always `Some` while the lease is alive.
///
/// `Option` lets `Drop` release the strong `Arc` before pruning the weak index.
registry: Option<Arc<QueryDynFilterRegistry>>,
}
impl RemoteDynFilterRegistryLease {
fn new(
registry_manager: Arc<DynFilterRegistryManager>,
registry: Arc<QueryDynFilterRegistry>,
) -> Self {
Self {
registry_manager,
registry: Some(registry),
}
}
pub fn registry(&self) -> &QueryDynFilterRegistry {
self.registry
.as_deref()
.expect("remote dyn filter registry lease must hold a registry")
}
#[cfg(test)]
pub(crate) fn ptr_eq(&self, other: &Self) -> bool {
Arc::ptr_eq(
self.registry.as_ref().unwrap(),
other.registry.as_ref().unwrap(),
)
}
}
impl Drop for RemoteDynFilterRegistryLease {
fn drop(&mut self) {
let Some(registry) = self.registry.take() else {
return;
};
let query_id = registry.query_id();
let registry_weak = Arc::downgrade(&registry);
// Release this lease before pruning; concurrent drops must not observe each other's strong refs.
drop(registry);
let _ = self
.registry_manager
.remove_if_dropped_registry(&query_id, &registry_weak);
}
}
/// Query-engine manager for query-scoped remote dynamic filter registries.
///
/// Weak index only; active streams own registries through [`RemoteDynFilterRegistryLease`].
#[derive(Debug, Default)]
pub struct DynFilterRegistryManager {
registries: RwLock<HashMap<QueryId, Weak<QueryDynFilterRegistry>>>,
}
impl DynFilterRegistryManager {
#[cfg(test)]
fn get(&self, query_id: &QueryId) -> Option<Arc<QueryDynFilterRegistry>> {
let (registry, stale_entry) = {
let registries = self.registries.read().unwrap();
let registry = registries.get(query_id)?;
(registry.upgrade(), registry.clone())
};
if registry.is_none() {
self.remove_stale_entry(query_id, &stale_entry);
}
registry
}
#[cfg(test)]
fn remove(&self, query_id: &QueryId) -> Option<Weak<QueryDynFilterRegistry>> {
self.registries.write().unwrap().remove(query_id)
}
fn remove_if_dropped_registry(
&self,
query_id: &QueryId,
dropped_registry: &Weak<QueryDynFilterRegistry>,
) -> Option<Weak<QueryDynFilterRegistry>> {
let mut registries = self
.registries
.write()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let current = registries.get(query_id)?;
// `ptr_eq` protects a newer registry for the same query id; `upgrade` ensures it is dead.
if current.ptr_eq(dropped_registry) && current.upgrade().is_none() {
registries.remove(query_id)
} else {
None
}
}
#[cfg(test)]
fn remove_stale_entry(
&self,
query_id: &QueryId,
stale_registry: &Weak<QueryDynFilterRegistry>,
) {
let mut registries = self.registries.write().unwrap();
let Some(current) = registries.get(query_id) else {
return;
};
if current.ptr_eq(stale_registry) && current.upgrade().is_none() {
registries.remove(query_id);
}
}
/// Acquires the stream-owned registry lease for `query_id`.
///
/// Returns a lease holding a strong registry reference.
pub fn acquire_lease(self: &Arc<Self>, query_id: QueryId) -> RemoteDynFilterRegistryLease {
let registry = self.get_or_init(query_id);
RemoteDynFilterRegistryLease::new(self.clone(), registry)
}
fn get_or_init(&self, query_id: QueryId) -> Arc<QueryDynFilterRegistry> {
let mut registries = self.registries.write().unwrap();
if let Some(registry) = registries.get(&query_id).and_then(Weak::upgrade) {
return registry;
}
let registry = Arc::new(QueryDynFilterRegistry::new(query_id));
registries.insert(query_id, Arc::downgrade(&registry));
registry
}
#[cfg(test)]
pub fn registry_count(&self) -> usize {
// Test snapshot helper; lifecycle decisions use lease-owned Arcs and weak pruning.
self.registries
.read()
.unwrap()
.values()
.filter(|registry| registry.strong_count() > 0)
.count()
}
#[cfg(test)]
fn weak_entry_count(&self) -> usize {
self.registries.read().unwrap().len()
}
}
#[cfg(test)]
mod tests {
use std::sync::Barrier;
use std::thread;
use datafusion_physical_expr::expressions::{Column, lit};
use uuid::Uuid;
use super::*;
use crate::dist_plan::{FilterFingerprint, RemoteDynFilterProducerId};
fn test_query_id(value: u128) -> QueryId {
QueryId::from(Uuid::from_u128(value))
}
fn test_filter_id(producer_ordinal: u32) -> FilterId {
FilterId::new(
RemoteDynFilterProducerId::new(42),
producer_ordinal,
FilterFingerprint::new(0xabc),
)
}
fn test_dyn_filter(names: &[&str]) -> Arc<DynamicFilterPhysicalExpr> {
let children = names
.iter()
.enumerate()
.map(|(index, name)| Arc::new(Column::new(name, index)) as _)
.collect();
Arc::new(DynamicFilterPhysicalExpr::new(children, lit(true) as _))
}
#[test]
fn registry_manager_returns_same_registry_for_same_query() {
let manager = Arc::new(DynFilterRegistryManager::default());
let query_id = test_query_id(1);
let first = manager.acquire_lease(query_id);
let second = manager.acquire_lease(query_id);
assert!(first.ptr_eq(&second));
assert_eq!(manager.registry_count(), 1);
assert_eq!(manager.weak_entry_count(), 1);
}
#[test]
fn registry_manager_removes_registry_for_query() {
let manager = Arc::new(DynFilterRegistryManager::default());
let query_id = test_query_id(1);
let lease = manager.acquire_lease(query_id);
assert!(
manager
.remove(&query_id)
.unwrap()
.ptr_eq(&Arc::downgrade(lease.registry.as_ref().unwrap()))
);
assert!(manager.get(&query_id).is_none());
assert_eq!(manager.registry_count(), 0);
assert_eq!(manager.weak_entry_count(), 0);
}
#[test]
fn registry_manager_lease_waits_for_last_query_scoped_stream() {
let manager = Arc::new(DynFilterRegistryManager::default());
let query_id = test_query_id(1);
let first = manager.acquire_lease(query_id);
let second = manager.acquire_lease(query_id);
assert!(first.ptr_eq(&second));
assert_eq!(manager.registry_count(), 1);
assert_eq!(manager.weak_entry_count(), 1);
drop(first);
assert_eq!(manager.registry_count(), 1);
assert_eq!(manager.weak_entry_count(), 1);
drop(second);
assert_eq!(manager.registry_count(), 0);
assert_eq!(manager.weak_entry_count(), 0);
}
#[test]
fn registry_manager_lease_does_not_remove_reacquired_registry() {
let manager = Arc::new(DynFilterRegistryManager::default());
let query_id = test_query_id(1);
let first = manager.acquire_lease(query_id);
drop(first);
assert_eq!(manager.registry_count(), 0);
assert_eq!(manager.weak_entry_count(), 0);
let second = manager.acquire_lease(query_id);
assert_eq!(manager.registry_count(), 1);
assert_eq!(manager.weak_entry_count(), 1);
drop(second);
assert_eq!(manager.registry_count(), 0);
assert_eq!(manager.weak_entry_count(), 0);
}
#[test]
fn registry_manager_concurrent_final_lease_drop_cleans_weak_entry() {
let manager = Arc::new(DynFilterRegistryManager::default());
let query_id = test_query_id(1);
let first = manager.acquire_lease(query_id);
let second = manager.acquire_lease(query_id);
let barrier = Arc::new(Barrier::new(3));
let first_barrier = barrier.clone();
let first_drop = thread::spawn(move || {
first_barrier.wait();
drop(first);
});
let second_barrier = barrier.clone();
let second_drop = thread::spawn(move || {
second_barrier.wait();
drop(second);
});
barrier.wait();
first_drop.join().unwrap();
second_drop.join().unwrap();
assert_eq!(manager.registry_count(), 0);
assert_eq!(manager.weak_entry_count(), 0);
}
#[test]
fn registry_manager_concurrent_first_acquire_shares_registry() {
let manager = Arc::new(DynFilterRegistryManager::default());
let query_id = test_query_id(1);
let worker_count = 8;
let barrier = Arc::new(Barrier::new(worker_count + 1));
let handles = (0..worker_count)
.map(|_| {
let manager = manager.clone();
let barrier = barrier.clone();
thread::spawn(move || {
barrier.wait();
manager.acquire_lease(query_id)
})
})
.collect::<Vec<_>>();
barrier.wait();
let leases = handles
.into_iter()
.map(|handle| handle.join().unwrap())
.collect::<Vec<_>>();
let first = leases.first().unwrap();
assert!(leases.iter().all(|lease| first.ptr_eq(lease)));
assert_eq!(manager.registry_count(), 1);
assert_eq!(manager.weak_entry_count(), 1);
drop(leases);
assert_eq!(manager.registry_count(), 0);
assert_eq!(manager.weak_entry_count(), 0);
}
#[test]
fn registry_manager_drop_racing_acquire_does_not_leave_stale_entry() {
let manager = Arc::new(DynFilterRegistryManager::default());
let query_id = test_query_id(1);
for _ in 0..64 {
let old_lease = manager.acquire_lease(query_id);
let barrier = Arc::new(Barrier::new(3));
let drop_barrier = barrier.clone();
let drop_thread = thread::spawn(move || {
drop_barrier.wait();
drop(old_lease);
});
let acquire_manager = manager.clone();
let acquire_barrier = barrier.clone();
let acquire_thread = thread::spawn(move || {
acquire_barrier.wait();
acquire_manager.acquire_lease(query_id)
});
barrier.wait();
drop_thread.join().unwrap();
let new_lease = acquire_thread.join().unwrap();
assert_eq!(manager.registry_count(), 1);
assert_eq!(manager.weak_entry_count(), 1);
drop(new_lease);
assert_eq!(manager.registry_count(), 0);
assert_eq!(manager.weak_entry_count(), 0);
}
}
#[test]
fn registry_manager_old_drop_cannot_remove_replacement_registry() {
let manager = Arc::new(DynFilterRegistryManager::default());
let query_id = test_query_id(1);
let old_lease = manager.acquire_lease(query_id);
let old_registry = Arc::downgrade(old_lease.registry.as_ref().unwrap());
drop(old_lease);
assert_eq!(manager.registry_count(), 0);
assert_eq!(manager.weak_entry_count(), 0);
let replacement_lease = manager.acquire_lease(query_id);
assert_eq!(manager.registry_count(), 1);
assert_eq!(manager.weak_entry_count(), 1);
assert!(
manager
.remove_if_dropped_registry(&query_id, &old_registry)
.is_none(),
"old registry cleanup must not remove the replacement weak entry"
);
assert_eq!(manager.registry_count(), 1);
assert_eq!(manager.weak_entry_count(), 1);
drop(replacement_lease);
assert_eq!(manager.registry_count(), 0);
assert_eq!(manager.weak_entry_count(), 0);
}
#[test]
fn registry_stores_filter_and_deduplicates_subscribers() {
let registry = QueryDynFilterRegistry::new(test_query_id(1));
let filter = test_dyn_filter(&["host"]);
let filter_id = test_filter_id(1);
let entry = match registry.register_remote_dyn_filter(filter_id.clone(), filter.clone()) {
EntryRegistration::Inserted(entry) => entry,
other => panic!("unexpected registration result: {other:?}"),
};
assert_eq!(entry.filter_id(), &filter_id);
assert_eq!(registry.entry_count(), 1);
let subscriber = Subscriber::new(RegionId::new(1024, 1));
assert_eq!(
registry.register_subscriber(&filter_id, subscriber.clone()),
SubscriberRegistration::Added
);
assert_eq!(
registry.register_subscriber(&filter_id, subscriber),
SubscriberRegistration::Duplicate
);
assert_eq!(entry.subscribers().len(), 1);
}
}

21
src/query/src/metrics.rs
View File

@@ -383,6 +383,7 @@ mod tests {
use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;
use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate};
use async_trait::async_trait;
use datafusion::arrow::datatypes::Schema as ArrowSchema;
use datafusion::execution::SessionStateBuilder;
@@ -394,6 +395,7 @@ mod tests {
use table::table_name::TableName;
use super::*;
use crate::dist_plan::RemoteDynFilterProducerId;
use crate::options::{FLOW_RETURN_REGION_SEQ, FLOW_SINK_TABLE_ID};
use crate::region_query::RegionQueryHandler;
@@ -408,6 +410,24 @@ mod tests {
) -> Result<SendableRecordBatchStream> {
unreachable!("metrics tests should not execute remote queries")
}
async fn handle_remote_dyn_filter_update(
&self,
_region_id: RegionId,
_query_id: String,
_update: RemoteDynFilterUpdate,
) -> Result<()> {
unreachable!("metrics tests should not send remote dyn filter updates")
}
async fn handle_remote_dyn_filter_unregister(
&self,
_region_id: RegionId,
_query_id: String,
_unregister: RemoteDynFilterUnregister,
) -> Result<()> {
unreachable!("metrics tests should not send remote dyn filter unregisters")
}
}
fn metrics_with_region_watermarks(entries: &[(u64, Option<u64>)]) -> RecordBatchMetrics {
@@ -439,6 +459,7 @@ mod tests {
query_ctx,
1,
BTreeMap::<String, BTreeSet<datafusion_common::Column>>::new(),
Some(RemoteDynFilterProducerId::new(0)),
)
.unwrap(),
)

21
src/query/src/optimizer/pass_distribution.rs
View File

@@ -133,6 +133,7 @@ impl PassDistribution {
mod tests {
use std::collections::{BTreeMap, BTreeSet};
use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate};
use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit};
use async_trait::async_trait;
use common_query::request::QueryRequest;
@@ -153,6 +154,7 @@ mod tests {
use table::table_name::TableName;
use super::*;
use crate::dist_plan::RemoteDynFilterProducerId;
use crate::error::Result as QueryResult;
use crate::region_query::RegionQueryHandler;
@@ -167,6 +169,24 @@ mod tests {
) -> QueryResult<SendableRecordBatchStream> {
unreachable!("pass distribution tests should not execute remote queries")
}
async fn handle_remote_dyn_filter_update(
&self,
_region_id: RegionId,
_query_id: String,
_update: RemoteDynFilterUpdate,
) -> QueryResult<()> {
unreachable!("pass distribution tests should not send remote dyn filter updates")
}
async fn handle_remote_dyn_filter_unregister(
&self,
_region_id: RegionId,
_query_id: String,
_unregister: RemoteDynFilterUnregister,
) -> QueryResult<()> {
unreachable!("pass distribution tests should not send remote dyn filter unregisters")
}
}
#[test]
@@ -272,6 +292,7 @@ mod tests {
QueryContext::arc(),
32,
partition_cols,
Some(RemoteDynFilterProducerId::new(1)),
)
.unwrap()
}

110
src/query/src/query_engine/state.rs
View File

@@ -50,12 +50,14 @@ use datafusion_optimizer::analyzer::function_rewrite::ApplyFunctionRewrites;
use datafusion_optimizer::optimizer::Optimizer;
use partition::manager::PartitionRuleManagerRef;
use promql::extension_plan::PromExtensionPlanner;
use session::context::QueryContextRef;
use table::TableRef;
use table::table::adapter::DfTableProviderAdapter;
use crate::QueryEngineContext;
use crate::dist_plan::{
DistExtensionPlanner, DistPlannerAnalyzer, DistPlannerOptions, MergeSortExtensionPlanner,
DistExtensionPlanner, DistPlannerAnalyzer, DistPlannerOptions, DynFilterRegistryManager,
MergeSortExtensionPlanner, RemoteDynFilterRegistryLease,
};
use crate::metrics::{QUERY_MEMORY_POOL_REJECTED_TOTAL, QUERY_MEMORY_POOL_USAGE_BYTES};
use crate::optimizer::ExtensionAnalyzerRule;
@@ -84,6 +86,7 @@ use crate::region_query::RegionQueryHandlerRef;
pub struct QueryEngineState {
df_context: SessionContext,
catalog_manager: CatalogManagerRef,
dyn_filter_registry_manager: Arc<DynFilterRegistryManager>,
function_state: Arc<FunctionState>,
scalar_functions: Arc<RwLock<HashMap<String, ScalarFunctionFactory>>>,
aggr_functions: Arc<RwLock<HashMap<String, AggregateUDF>>>,
@@ -228,7 +231,7 @@ impl QueryEngineState {
.with_query_planner(Arc::new(DfQueryPlanner::new(
catalog_list.clone(),
partition_rule_manager,
region_query_handler,
region_query_handler.clone(),
)))
.with_optimizer_rules(optimizer.rules)
.with_physical_optimizer_rules(physical_optimizer.rules)
@@ -240,6 +243,7 @@ impl QueryEngineState {
Self {
df_context,
catalog_manager: catalog_list,
dyn_filter_registry_manager: Arc::new(DynFilterRegistryManager::default()),
function_state: Arc::new(FunctionState {
table_mutation_handler,
procedure_service_handler,
@@ -396,6 +400,22 @@ impl QueryEngineState {
&self.catalog_manager
}
pub fn dyn_filter_registry_manager(&self) -> Arc<DynFilterRegistryManager> {
self.dyn_filter_registry_manager.clone()
}
pub fn acquire_remote_dyn_filter_registry_lease(
&self,
query_ctx: &QueryContextRef,
) -> Option<RemoteDynFilterRegistryLease> {
let query_id = query_ctx.remote_query_id_value()?;
Some(
self.dyn_filter_registry_manager
.clone()
.acquire_lease(query_id),
)
}
pub fn function_state(&self) -> Arc<FunctionState> {
self.function_state.clone()
}
@@ -577,3 +597,89 @@ impl MemoryPool for MetricsMemoryPool {
self.inner.memory_limit()
}
}
#[cfg(test)]
mod tests {
use common_base::Plugins;
use session::context::QueryContext;
use super::*;
use crate::options::QueryOptions;
fn new_query_engine_state() -> QueryEngineState {
QueryEngineState::new(
catalog::memory::new_memory_catalog_manager().unwrap(),
None,
None,
None,
None,
None,
false,
Plugins::default(),
QueryOptions::default(),
)
}
#[test]
fn query_engine_state_reuses_query_scoped_dyn_filter_registry_lease() {
let state = new_query_engine_state();
let query_ctx = QueryContext::arc();
let first = state
.acquire_remote_dyn_filter_registry_lease(&query_ctx)
.unwrap();
let second = state
.acquire_remote_dyn_filter_registry_lease(&query_ctx)
.unwrap();
assert!(first.ptr_eq(&second));
assert_eq!(state.dyn_filter_registry_manager().registry_count(), 1);
assert_eq!(
first.registry().query_id(),
query_ctx.remote_query_id_value().unwrap()
);
}
#[test]
fn query_engine_state_relies_on_query_context_remote_query_id_contract() {
let state = new_query_engine_state();
let query_ctx = QueryContext::arc();
assert!(query_ctx.remote_query_id_value().is_some());
let lease = state
.acquire_remote_dyn_filter_registry_lease(&query_ctx)
.unwrap();
assert_eq!(
lease.registry().query_id(),
query_ctx.remote_query_id_value().unwrap()
);
assert_eq!(state.dyn_filter_registry_manager().registry_count(), 1);
}
#[test]
fn query_engine_state_separates_registries_for_different_query_contexts() {
let state = new_query_engine_state();
let first_query_ctx = QueryContext::arc();
let second_query_ctx = QueryContext::arc();
let first = state
.acquire_remote_dyn_filter_registry_lease(&first_query_ctx)
.unwrap();
let second = state
.acquire_remote_dyn_filter_registry_lease(&second_query_ctx)
.unwrap();
assert!(!first.ptr_eq(&second));
assert_eq!(state.dyn_filter_registry_manager().registry_count(), 2);
assert_eq!(
first.registry().query_id(),
first_query_ctx.remote_query_id_value().unwrap()
);
assert_eq!(
second.registry().query_id(),
second_query_ctx.remote_query_id_value().unwrap()
);
}
}

16
src/query/src/region_query.rs
View File

@@ -14,12 +14,14 @@
use std::sync::Arc;
use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate};
use async_trait::async_trait;
use common_meta::node_manager::NodeManagerRef;
use common_query::request::QueryRequest;
use common_recordbatch::SendableRecordBatchStream;
use partition::manager::PartitionRuleManagerRef;
use session::ReadPreference;
use store_api::storage::RegionId;
use crate::error::Result;
@@ -42,6 +44,20 @@ pub trait RegionQueryHandler: Send + Sync {
read_preference: ReadPreference,
request: QueryRequest,
) -> Result<SendableRecordBatchStream>;
async fn handle_remote_dyn_filter_update(
&self,
region_id: RegionId,
query_id: String,
update: RemoteDynFilterUpdate,
) -> Result<()>;
async fn handle_remote_dyn_filter_unregister(
&self,
region_id: RegionId,
query_id: String,
unregister: RemoteDynFilterUnregister,
) -> Result<()>;
}
pub type RegionQueryHandlerRef = Arc<dyn RegionQueryHandler>;

18
src/servers/src/grpc/greptime_handler.rs
View File

@@ -286,7 +286,9 @@ impl Drop for RequestTimer {
mod tests {
use chrono::FixedOffset;
use common_time::Timezone;
use session::hints::REMOTE_QUERY_ID_EXTENSION_KEY;
use session::hints::{
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY,
};
use super::*;
@@ -303,6 +305,14 @@ mod tests {
vec![
("auto_create_table".to_string(), "true".to_string()),
("read_preference".to_string(), "leader".to_string()),
(
REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
"spoofed".to_string(),
),
(
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY.to_string(),
"spoofed-regs".to_string(),
),
],
)
.unwrap();
@@ -327,6 +337,12 @@ mod tests {
query_context.remote_query_id(),
Some(extensions[1].1.as_str())
);
assert_ne!(query_context.remote_query_id(), Some("spoofed"));
assert!(
query_context
.extension(INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY)
.is_none()
);
}
#[test]

22
src/servers/src/http/hints.rs
View File

@@ -45,8 +45,11 @@ fn apply_hints(query_ctx: &mut QueryContext, hints: Vec<(String, String)>) {
#[cfg(test)]
mod tests {
use common_query::request::INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY as COMMON_INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY;
use session::context::{QueryContextBuilder, generate_remote_query_id};
use session::hints::REMOTE_QUERY_ID_EXTENSION_KEY;
use session::hints::{
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY,
};
use super::apply_hints;
@@ -67,6 +70,10 @@ mod tests {
REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
"spoofed".to_string(),
),
(
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY.to_string(),
"spoofed-regs".to_string(),
),
("ttl".to_string(), "7d".to_string()),
],
);
@@ -75,6 +82,19 @@ mod tests {
query_ctx.remote_query_id(),
Some(original_query_id.as_str())
);
assert!(
query_ctx
.extension(INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY)
.is_none()
);
assert_eq!(query_ctx.extension("ttl"), Some("7d"));
}
#[test]
fn test_initial_dyn_filter_registration_key_matches_common_query_constant() {
assert_eq!(
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY,
COMMON_INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY
);
}
}

10
src/session/src/hints.rs
View File

@@ -17,9 +17,14 @@ pub const HINTS_KEY: &str = "x-greptime-hints";
/// Deprecated, use `HINTS_KEY` instead. Notes if "x-greptime-hints" is set, keys with this prefix will be ignored.
pub const HINTS_KEY_PREFIX: &str = "x-greptime-hint-";
pub const REMOTE_QUERY_ID_EXTENSION_KEY: &str = "remote_query_id";
pub const INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY: &str =
"initial_remote_dyn_filter_registrations";
pub const READ_PREFERENCE_HINT: &str = "read_preference";
pub const RESERVED_EXTENSION_KEYS: [&str; 1] = [REMOTE_QUERY_ID_EXTENSION_KEY];
pub const RESERVED_EXTENSION_KEYS: [&str; 2] = [
REMOTE_QUERY_ID_EXTENSION_KEY,
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY,
];
/// Deprecated, use `HINTS_KEY` instead.
pub const HINT_KEYS: [&str; 7] = [
@@ -43,6 +48,9 @@ mod tests {
#[test]
fn test_is_reserved_extension_key() {
assert!(is_reserved_extension_key(REMOTE_QUERY_ID_EXTENSION_KEY));
assert!(is_reserved_extension_key(
INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY
));
assert!(!is_reserved_extension_key(READ_PREFERENCE_HINT));
}
}