diff --git a/Cargo.lock b/Cargo.lock index e3d0dfb69e..5f0227d1ca 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11464,6 +11464,7 @@ dependencies = [ "datafusion-functions", "datafusion-optimizer", "datafusion-physical-expr", + "datafusion-proto", "datafusion-sql", "datatypes", "either", diff --git a/src/client/src/region.rs b/src/client/src/region.rs index 72321b2659..19c173bd05 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -293,11 +293,8 @@ impl RegionRequester { query_id: impl Into, update: RemoteDynFilterUpdate, ) -> Result { - 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, unregister: RemoteDynFilterUnregister, ) -> Result { - 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, + 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, + 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) -> 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(_)) + )); + } } diff --git a/src/common/query/Cargo.toml b/src/common/query/Cargo.toml index b28e56fe8e..9706098826 100644 --- a/src/common/query/Cargo.toml +++ b/src/common/query/Cargo.toml @@ -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" diff --git a/src/common/query/src/request.rs b/src/common/query/src/request.rs index 22437bde71..121292709f 100644 --- a/src/common/query/src/request.rs +++ b/src/common/query/src/request.rs @@ -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), -} - -mod base64_bytes { - use base64::Engine; - use base64::prelude::BASE64_STANDARD; - use serde::de::Error; - use serde::{Deserialize, Deserializer, Serializer}; - - pub fn serialize(bytes: &[u8], serializer: S) -> Result - where - S: Serializer, - { - serializer.serialize_str(&BASE64_STANDARD.encode(bytes)) - } - - pub fn deserialize<'de, D>(deserializer: D) -> Result, 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), } 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> { let Self::Datafusion(bytes) = self; @@ -124,6 +108,34 @@ impl DynFilterPayload { } } +fn encode_physical_expr_to_bytes(expr: &Arc) -> DataFusionResult> { + 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> { + 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) -> DataFusionRe /// schema inconsistency that should be surfaced loudly. fn validate_decoded_payload_expr( expr: &Arc, - input_schema: &Schema, + input_schema: &datafusion::arrow::datatypes::Schema, ) -> DataFusionResult<()> { expr.apply(|node| { if let Some(column) = node.as_any().downcast_ref::() { diff --git a/src/common/query/src/request/base64_serde.rs b/src/common/query/src/request/base64_serde.rs new file mode 100644 index 0000000000..9de0db2f1e --- /dev/null +++ b/src/common/query/src/request/base64_serde.rs @@ -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, base64::DecodeError> { + BASE64_STANDARD.decode(encoded) +} + +pub(crate) mod bytes { + use serde::de::Error; + use serde::{Deserialize, Deserializer, Serializer}; + + pub fn serialize(value: &[u8], serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_str(&super::encode(value)) + } + + pub fn deserialize<'de, D>(deserializer: D) -> Result, 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(values: &[Vec], serializer: S) -> Result + where + S: Serializer, + { + values + .iter() + .map(|bytes| super::encode(bytes)) + .collect::>() + .serialize(serializer) + } + + pub fn deserialize<'de, D>(deserializer: D) -> Result>, D::Error> + where + D: Deserializer<'de>, + { + Vec::::deserialize(deserializer)? + .into_iter() + .map(|encoded| { + super::decode(&encoded).map_err(|error| { + D::Error::custom(format!("invalid base64 bytes vector item: {error}")) + }) + }) + .collect() + } +} diff --git a/src/common/query/src/request/initial_remote_dyn_filter_reg.rs b/src/common/query/src/request/initial_remote_dyn_filter_reg.rs new file mode 100644 index 0000000000..4d73a5eef2 --- /dev/null +++ b/src/common/query/src/request/initial_remote_dyn_filter_reg.rs @@ -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, +} + +impl InitialDynFilterRegs { + pub fn new(regs: Vec) -> 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 { + serde_json::to_string(self) + } + + pub fn from_extension_value(value: &str) -> serde_json::Result { + let regs = serde_json::from_str::(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>, + /// 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, +} + +#[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, child_exprs_datafusion_proto: Vec>) -> 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, + children: &[Arc], + ) -> DataFusionResult { + let child_exprs_datafusion_proto = children + .iter() + .map(encode_physical_expr_to_bytes) + .collect::>>()?; + + 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>> { + self.child_exprs_datafusion_proto + .iter() + .map(|expr_bytes| { + decode_physical_expr_from_bytes( + expr_bytes, + task_ctx, + input_schema, + max_payload_bytes, + ) + }) + .collect::>>() + } +} + +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 = + 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::().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], + ) + .unwrap(); + + let err = reg + .decode_children(&TaskContext::default(), &schema, 1024) + .unwrap_err(); + + assert!(matches!(err, DataFusionError::Plan(_))); + } +} diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index ce831353d1..eb7d2dfa6c 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -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)) } diff --git a/src/frontend/src/instance/region_query.rs b/src/frontend/src/instance/region_query.rs index 56874f5777..1523e52341 100644 --- a/src/frontend/src/instance/region_query.rs +++ b/src/frontend/src/instance/region_query.rs @@ -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(()) + } } diff --git a/src/query/Cargo.toml b/src/query/Cargo.toml index b92921d29b..9d9f8b24b0 100644 --- a/src/query/Cargo.toml +++ b/src/query/Cargo.toml @@ -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 diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index a173f2aee8..a8d60d42f4 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -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) { diff --git a/src/query/src/dist_plan.rs b/src/query/src/dist_plan.rs index 2a6ba0dfbb..084ca7d1e5 100644 --- a/src/query/src/dist_plan.rs +++ b/src/query/src/dist_plan.rs @@ -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, +}; diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index 6eb8b5eac6..ba0d5a70c0 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -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 { 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> { @@ -224,6 +225,57 @@ impl DistPlannerAnalyzer { spans: Default::default(), }) } + + fn assign_merge_scan_remote_dyn_filter_producer_ids( + plan: LogicalPlan, + ) -> DfResult { + 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> { + let LogicalPlan::Extension(extension) = &node else { + return Ok(Transformed::no(node)); + }; + let Some(merge_scan) = extension + .node + .as_any() + .downcast_ref::() + 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 diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 9f72df55a5..6321337858 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -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, +) { + 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, +} + +impl TreeNodeRewriter for MergeScanRemoteDynFilterProducerIdCollector<'_> { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> DfResult> { + if let LogicalPlan::Extension(extension) = &node + && let Some(merge_scan) = extension + .node + .as_any() + .downcast_ref::() + { + 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, +) { + 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::>(); + + 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(); diff --git a/src/query/src/dist_plan/dyn_filter_bridge.rs b/src/query/src/dist_plan/dyn_filter_bridge.rs new file mode 100644 index 0000000000..546aa20778 --- /dev/null +++ b/src/query/src/dist_plan/dyn_filter_bridge.rs @@ -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, +} + +#[derive(Debug, Clone)] +pub(crate) struct RemoteDynFilterPushdown { + pub(crate) captured_dyn_filters: Vec, + /// Preflight result per parent filter. + pub(crate) pushed_down: Vec, +} + +pub(crate) fn capture_remote_dyn_filters_for_pushdown( + remote_dyn_filter_producer_id: RemoteDynFilterProducerId, + parent_filters: Vec>, +) -> 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, +) -> Option> { + (expr as Arc) + .downcast::() + .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, +) -> Result { + let children = alive_dyn_filter + .children() + .into_iter() + .cloned() + .collect::>(); + 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 { + 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( + ¤t, + 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(&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 { + Ok(arrow_schema::DataType::Boolean) + } + + fn nullable( + &self, + _input_schema: &arrow_schema::Schema, + ) -> datafusion_common::Result { + Ok(false) + } + + fn evaluate( + &self, + _batch: &common_recordbatch::DfRecordBatch, + ) -> datafusion_common::Result { + Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(true)))) + } + + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> datafusion_common::Result> { + 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 { + 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, + Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("host", 1)) as Arc<_>], + lit(true) as _, + )) as Arc, + Arc::new(Column::new("zone", 2)) as Arc, + Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("pod", 3)) as Arc<_>], + lit(true) as _, + )) as Arc, + ]; + + 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, + Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("host", 1)) as Arc<_>], + lit(true) as _, + )) as Arc, + Arc::new(Column::new("zone", 2)) as Arc, + ]; + + 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]; + + 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]; + + 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]; + + 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, + test_dyn_filter_with_snapshot_payload("pod", 1, 40 * 1024) + as Arc, + ]; + + 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 + }) + .collect::>(); + + 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 + }) + .collect::>(); + + 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(®istry, first_region_id, &captured_dyn_filters); + register_dyn_filters_for_region(®istry, 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( + ®istry, + region_id, + &[make_filter(test_remote_dyn_filter_producer_id(42))], + ); + register_dyn_filters_for_region( + ®istry, + 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::()); + } + + #[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() + ); + } +} diff --git a/src/query/src/dist_plan/filter_id.rs b/src/query/src/dist_plan/filter_id.rs new file mode 100644 index 0000000000..d3e7929336 --- /dev/null +++ b/src/query/src/dist_plan/filter_id.rs @@ -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 { + 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 { + 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 { + let mut parts = s.split(':'); + let remote_dyn_filter_producer_id = parts + .next() + .ok_or(ParseFilterIdError)? + .parse::() + .map_err(|_| ParseFilterIdError)?; + let producer_local_ordinal = parts + .next() + .ok_or(ParseFilterIdError)? + .parse::() + .map_err(|_| ParseFilterIdError)?; + let children_fingerprint = parts + .next() + .ok_or(ParseFilterIdError)? + .parse::() + .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], +) -> Result { + 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], +) -> Result { + 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> { + names + .iter() + .enumerate() + .map(|(index, name)| Arc::new(Column::new(name, index)) as Arc) + .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::().unwrap(), filter_id); + } + + #[test] + fn filter_id_rejects_malformed_strings() { + assert!("".parse::().is_err()); + assert!("0000000000000001:3".parse::().is_err()); + assert!("0000000000000001:3:zzzz".parse::().is_err()); + assert!( + "0000000000000001:3:0000000000000abc:extra" + .parse::() + .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")); + } +} diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index e1dd635de7..2201d0c52f 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -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> { + context.session_config().get_extension() +} + +fn acquire_remote_dyn_filter_registry_lease( + context: &TaskContext, + query_ctx: &QueryContextRef, + captured_dyn_filters: &[CapturedDynFilter], +) -> Option { + 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, } 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 { + self.remote_dyn_filter_producer_id + } } +#[derive(Clone)] pub struct MergeScanExec { table: TableName, regions: Vec, @@ -146,6 +193,9 @@ pub struct MergeScanExec { /// Metrics for each partition partition_metrics: Arc>>, query_ctx: QueryContextRef, + /// Optional because RDF must fail open: missing ids skip RDF but keep normal query execution. + remote_dyn_filter_producer_id: Option, + captured_remote_dyn_filters: Arc>>, 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, ) -> Result { // 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((®ion_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 { + self.captured_remote_dyn_filters.lock().unwrap().clone() + } + pub fn sub_stage_metrics(&self) -> Vec { self.sub_stage_metrics .lock() @@ -517,6 +598,13 @@ impl MergeScanExec { } } +#[cfg(test)] +impl MergeScanExec { + fn remote_dyn_filter_producer_id(&self) -> Option { + 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>> { + let parent_filters = child_pushdown_result + .parent_filters + .into_iter() + .map(|filter| filter.filter) + .collect::>(); + 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 { + 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; + + 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])); + } +} diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index 5864ea7bd0..e6776b7039 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -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 _)) } diff --git a/src/query/src/dist_plan/remote_dyn_filter_registry.rs b/src/query/src/dist_plan/remote_dyn_filter_registry.rs new file mode 100644 index 0000000000..4c33ddd960 --- /dev/null +++ b/src/query/src/dist_plan/remote_dyn_filter_registry.rs @@ -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), + /// The filter already existed; this contains the previously registered entry. + Existing(Arc), +} + +/// 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, + subscribers: RwLock>, +} + +#[derive(Debug)] +struct QueryDynFilterRegistryInner { + entries: HashMap>, +} + +impl DynFilterEntry { + pub fn new(filter_id: FilterId, alive_dyn_filter: Arc) -> 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> { + self.alive_dyn_filter.upgrade() + } + + pub fn subscribers(&self) -> Vec { + 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, +} + +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> { + self.inner + .read() + .unwrap() + .entries + .values() + .cloned() + .collect() + } + + pub fn remote_dyn_filter(&self, filter_id: &FilterId) -> Option> { + self.inner.read().unwrap().entries.get(filter_id).cloned() + } + + pub fn register_remote_dyn_filter( + &self, + filter_id: FilterId, + alive_dyn_filter: Arc, + ) -> 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, + /// Always `Some` while the lease is alive. + /// + /// `Option` lets `Drop` release the strong `Arc` before pruning the weak index. + registry: Option>, +} + +impl RemoteDynFilterRegistryLease { + fn new( + registry_manager: Arc, + registry: Arc, + ) -> 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(®istry); + + // 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, ®istry_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>>, +} + +impl DynFilterRegistryManager { + #[cfg(test)] + fn get(&self, query_id: &QueryId) -> Option> { + 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> { + self.registries.write().unwrap().remove(query_id) + } + + fn remove_if_dropped_registry( + &self, + query_id: &QueryId, + dropped_registry: &Weak, + ) -> Option> { + 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, + ) { + 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, 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 { + 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(®istry)); + 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 { + 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::>(); + + barrier.wait(); + let leases = handles + .into_iter() + .map(|handle| handle.join().unwrap()) + .collect::>(); + + 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); + } +} diff --git a/src/query/src/metrics.rs b/src/query/src/metrics.rs index 7696291480..b1483faea2 100644 --- a/src/query/src/metrics.rs +++ b/src/query/src/metrics.rs @@ -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 { 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)]) -> RecordBatchMetrics { @@ -439,6 +459,7 @@ mod tests { query_ctx, 1, BTreeMap::>::new(), + Some(RemoteDynFilterProducerId::new(0)), ) .unwrap(), ) diff --git a/src/query/src/optimizer/pass_distribution.rs b/src/query/src/optimizer/pass_distribution.rs index 1004252636..e9cd5f028d 100644 --- a/src/query/src/optimizer/pass_distribution.rs +++ b/src/query/src/optimizer/pass_distribution.rs @@ -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 { 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() } diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index 4262428091..ad2ef92bb9 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -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, function_state: Arc, scalar_functions: Arc>>, aggr_functions: Arc>>, @@ -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 { + self.dyn_filter_registry_manager.clone() + } + + pub fn acquire_remote_dyn_filter_registry_lease( + &self, + query_ctx: &QueryContextRef, + ) -> Option { + 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 { 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() + ); + } +} diff --git a/src/query/src/region_query.rs b/src/query/src/region_query.rs index 2519ccc25f..0e6b724337 100644 --- a/src/query/src/region_query.rs +++ b/src/query/src/region_query.rs @@ -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; + + 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; diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index 7ae881adea..4669ade118 100644 --- a/src/servers/src/grpc/greptime_handler.rs +++ b/src/servers/src/grpc/greptime_handler.rs @@ -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] diff --git a/src/servers/src/http/hints.rs b/src/servers/src/http/hints.rs index 406f226738..32598d7bf7 100644 --- a/src/servers/src/http/hints.rs +++ b/src/servers/src/http/hints.rs @@ -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 + ); + } } diff --git a/src/session/src/hints.rs b/src/session/src/hints.rs index db51bc6279..c971de5967 100644 --- a/src/session/src/hints.rs +++ b/src/session/src/hints.rs @@ -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)); } }