diff --git a/src/datanode/src/region_server/registrations.rs b/src/datanode/src/region_server/registrations.rs index 2eb86a9ff704..427e739884c7 100644 --- a/src/datanode/src/region_server/registrations.rs +++ b/src/datanode/src/region_server/registrations.rs @@ -161,7 +161,7 @@ fn record_unregister_outcome(outcome: RemoteDynFilterUpdateOutcome) { #[derive(Debug, Clone)] struct PendingDynFilterUpdate { - payload: Vec, + payload: DynFilterPayload, generation: u64, is_complete: bool, } @@ -169,14 +169,12 @@ struct PendingDynFilterUpdate { impl PendingDynFilterUpdate { fn from_initial_reg(reg: &InitialDynFilterReg) -> Option { let snapshot = reg.initial_snapshot.as_ref()?; - match &snapshot.payload { - DynFilterPayload::Datafusion(payload) => Some(Self { - payload: payload.clone(), - generation: snapshot.generation, - is_complete: snapshot.is_complete, - }), - _ => None, - } + // Accept any payload variant for pending; FE may send non-Datafusion in the future. + Some(Self { + payload: snapshot.payload.clone(), + generation: snapshot.generation, + is_complete: snapshot.is_complete, + }) } } @@ -190,14 +188,20 @@ struct RemoteDynFilterEpochState { struct RemoteDynFilterState { filter: Arc, input_schema: SchemaRef, + registered_children: Vec>, epoch: Mutex, } impl RemoteDynFilterState { - fn new(filter: Arc, input_schema: SchemaRef) -> Self { + fn new( + filter: Arc, + input_schema: SchemaRef, + registered_children: Vec>, + ) -> Self { Self { filter, input_schema, + registered_children, epoch: Mutex::new(RemoteDynFilterEpochState { generation: None, is_complete: false, @@ -211,7 +215,7 @@ impl RemoteDynFilterState { fn apply_update( &self, - payload: &[u8], + payload: &DynFilterPayload, generation: u64, is_complete: bool, ) -> RemoteDynFilterUpdateOutcome { @@ -239,7 +243,12 @@ impl RemoteDynFilterState { return RemoteDynFilterUpdateOutcome::AlreadyComplete; } - let expr = match decode_update_payload(payload, self.input_schema.as_ref()) { + let expr = match payload.decode_expr_with_registered_children( + remote_dyn_filter_task_context(), + self.input_schema.as_ref(), + &self.registered_children, + REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES, + ) { Ok(expr) => expr, Err(error) => { warn!(error; "Failed to decode remote dynamic filter update payload"); @@ -324,7 +333,7 @@ impl RegisteredDynFilter { fn buffer_update( &mut self, - payload: &[u8], + payload: &DynFilterPayload, generation: u64, is_complete: bool, ) -> RemoteDynFilterUpdateOutcome { @@ -348,7 +357,7 @@ impl RegisteredDynFilter { } self.pending_update = Some(PendingDynFilterUpdate { - payload: payload.to_vec(), + payload: payload.clone(), generation, is_complete, }); @@ -357,7 +366,7 @@ impl RegisteredDynFilter { fn apply_or_buffer_update( &mut self, - payload: &[u8], + payload: &DynFilterPayload, generation: u64, is_complete: bool, ) -> RemoteDynFilterUpdateOutcome { @@ -399,6 +408,7 @@ impl RegisteredDynFilter { let runtime = Arc::new(RemoteDynFilterState::new( filter, Arc::new(input_schema.clone()), + children.clone(), )); if let Some(pending) = self.pending_update.take() { let outcome = runtime.apply_update( @@ -537,7 +547,7 @@ pub(super) fn apply_remote_dyn_filter_update( regs_by_query: &RemoteDynFilterRegistry, query_id: &QueryId, filter_id: &RemoteDynFilterId, - payload: &[u8], + payload: &DynFilterPayload, generation: u64, is_complete: bool, ) -> RemoteDynFilterUpdateOutcome { @@ -546,7 +556,7 @@ pub(super) fn apply_remote_dyn_filter_update( "Ignored oversized remote dynamic filter update, query_id: {}, filter_id: {}, payload_size: {}, max_payload_size: {}", query_id, filter_id, - payload.len(), + payload.encoded_payload_bytes(), REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES ); let outcome = RemoteDynFilterUpdateOutcome::PayloadTooLarge; @@ -554,7 +564,7 @@ pub(super) fn apply_remote_dyn_filter_update( return outcome; } - REMOTE_DYN_FILTER_PAYLOAD_BYTES.observe(payload.len() as f64); + REMOTE_DYN_FILTER_PAYLOAD_BYTES.observe(payload.encoded_payload_bytes() as f64); let Some(query_regs) = regs_by_query.get_query(query_id) else { warn!( @@ -665,15 +675,8 @@ pub(super) fn remove_initial_dyn_filter_regs( } } -fn decode_update_payload( - payload: &[u8], - input_schema: &Schema, -) -> DataFusionResult> { - DynFilterPayload::Datafusion(payload.to_vec()).decode_datafusion_expr( - remote_dyn_filter_task_context(), - input_schema, - REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES, - ) +fn validate_update_payload_size(payload: &DynFilterPayload) -> bool { + payload.encoded_payload_bytes() <= REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES } fn remote_dyn_filter_task_context() -> &'static TaskContext { @@ -689,10 +692,6 @@ fn remote_dyn_filter_task_context() -> &'static TaskContext { }) } -fn validate_update_payload_size(payload: &[u8]) -> bool { - payload.len() <= REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/datanode/src/region_server/remote_dyn_filter.rs b/src/datanode/src/region_server/remote_dyn_filter.rs index 2034917fd254..529e470b412d 100644 --- a/src/datanode/src/region_server/remote_dyn_filter.rs +++ b/src/datanode/src/region_server/remote_dyn_filter.rs @@ -20,6 +20,7 @@ use api::region::RegionResponse; use api::v1::region::RemoteDynFilterRequest; use api::v1::region::remote_dyn_filter_request::Action; use common_base::Plugins; +use common_query::request::DynFilterPayload; use common_recordbatch::adapter::RecordBatchMetrics; use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream}; use common_telemetry::{debug, warn}; @@ -180,16 +181,36 @@ impl RegionServer { return error::MissingRequiredFieldSnafu { name: "filter_id" }.fail(); } - if request.payload.is_empty() { + if request.payload.is_empty() && request.typed_payload.is_none() { return error::MissingRequiredFieldSnafu { name: "payload" }.fail(); } let filter_id = RemoteDynFilterId::new(request.filter_id.clone()); + + // Prefer typed payload if available; fall back to legacy payload. + let dyn_filter_payload = if let Some(ref typed_proto) = request.typed_payload { + match DynFilterPayload::from_region_proto_payload(typed_proto.clone()) { + Ok(typed) => typed, + Err(error) if !request.payload.is_empty() => { + warn!(error; "Failed to decode typed remote dyn filter payload, falling back to legacy Datafusion payload"); + DynFilterPayload::Datafusion(request.payload.clone()) + } + Err(_error) => { + warn!(_error; "Failed to decode typed remote dyn filter payload with no legacy fallback"); + let outcome = RemoteDynFilterUpdateOutcome::DecodeFailed; + self.log_remote_dyn_filter_update_outcome(query_id, &filter_id, outcome); + return Ok(RegionResponse::new(0)); + } + } + } else { + DynFilterPayload::Datafusion(request.payload.clone()) + }; + let outcome = apply_remote_dyn_filter_update( &self.inner.initial_remote_dyn_filter_registrations, query_id, &filter_id, - &request.payload, + &dyn_filter_payload, request.generation, request.is_complete, ); @@ -357,9 +378,15 @@ mod tests { REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES, }; use common_recordbatch::RecordBatches; - use datafusion::arrow::datatypes::Schema as ArrowSchema; + use datafusion::arrow::datatypes::{ + DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema, + }; use datafusion::physical_plan::PhysicalExpr; - use datafusion::physical_plan::expressions::{DynamicFilterPhysicalExpr, lit as physical_lit}; + use datafusion::physical_plan::expressions::{ + Column, DynamicFilterPhysicalExpr, lit as physical_lit, + }; + use datafusion::physical_plan::joins::join_hash_map::{JoinHashMapType, JoinHashMapU32}; + use datafusion::physical_plan::joins::{HashTableLookupExpr, Map, SeededRandomState}; use datafusion_common::DFSchema; use datafusion_expr::EmptyRelation; use datatypes::prelude::{ConcreteDataType, VectorRef}; @@ -776,8 +803,14 @@ mod tests { let query_id = test_remote_query_id(); let filter_id = RemoteDynFilterId::new("filter-1"); - let outcome = - apply_remote_dyn_filter_update(®s_by_query, &query_id, &filter_id, &[1], 1, false); + let outcome = apply_remote_dyn_filter_update( + ®s_by_query, + &query_id, + &filter_id, + &test_datafusion_payload(vec![1]), + 1, + false, + ); assert_eq!(outcome, RemoteDynFilterUpdateOutcome::MissingRegistration); let outcome = unregister_remote_dyn_filter(®s_by_query, &query_id, &filter_id); @@ -802,36 +835,67 @@ mod tests { ); // First update with generation 1: should be Buffered (no runtime installed yet) - let outcome = - apply_remote_dyn_filter_update(®s_by_query, &query_id, &filter_id, &[1], 1, false); + let outcome = apply_remote_dyn_filter_update( + ®s_by_query, + &query_id, + &filter_id, + &test_datafusion_payload(vec![1]), + 1, + false, + ); assert_eq!(outcome, RemoteDynFilterUpdateOutcome::Buffered); // Update with generation 0: should be Stale (older than pending generation 1) - let outcome = - apply_remote_dyn_filter_update(®s_by_query, &query_id, &filter_id, &[2], 0, false); + let outcome = apply_remote_dyn_filter_update( + ®s_by_query, + &query_id, + &filter_id, + &test_datafusion_payload(vec![2]), + 0, + false, + ); assert_eq!(outcome, RemoteDynFilterUpdateOutcome::Stale); // Update with generation 1 again: should be Idempotent (same generation) - let outcome = - apply_remote_dyn_filter_update(®s_by_query, &query_id, &filter_id, &[3], 1, false); + let outcome = apply_remote_dyn_filter_update( + ®s_by_query, + &query_id, + &filter_id, + &test_datafusion_payload(vec![3]), + 1, + false, + ); assert_eq!(outcome, RemoteDynFilterUpdateOutcome::Idempotent); } - fn datafusion_payload_bytes(expr: Arc) -> Vec { - match DynFilterPayload::from_datafusion_expr(&expr, REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES) - .unwrap() - { - DynFilterPayload::Datafusion(bytes) => bytes, - _ => unreachable!( - "DynFilterPayload::from_datafusion_expr only returns datafusion payloads" - ), - } + fn datafusion_payload_bytes(expr: Arc) -> DynFilterPayload { + DynFilterPayload::from_datafusion_expr(&expr, REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES).unwrap() + } + + fn test_datafusion_payload(bytes: Vec) -> DynFilterPayload { + DynFilterPayload::Datafusion(bytes) + } + + fn test_lookup_expr(child: Arc, hashes: Vec) -> Arc { + let mut hash_map = JoinHashMapU32::with_capacity(hashes.len().max(1)); + hash_map.update_from_iter(Box::new(hashes.iter().enumerate()), 0); + let map = Arc::new(Map::HashMap(Box::new(hash_map))); + Arc::new(HashTableLookupExpr::new( + vec![child], + SeededRandomState::with_seeds(1, 2, 3, 4), + map, + "lookup".to_string(), + )) } fn empty_arrow_schema() -> ArrowSchema { ArrowSchema::empty() } + fn one_column_arrow_schema() -> ArrowSchema { + ArrowSchema::new(vec![ArrowField::new("id", ArrowDataType::Int32, false)]) + } + fn register_empty_remote_dyn_filter( regs_by_query: &RemoteDynFilterRegistry, query_id: &QueryId, @@ -860,7 +924,7 @@ mod tests { fn initial_remote_dyn_filter_snapshot_initializes_runtime_filter() { let regs_by_query = RemoteDynFilterRegistry::new(); let query_id = test_remote_query_id(); - let payload = DynFilterPayload::Datafusion(datafusion_payload_bytes(physical_lit(false))); + let payload = datafusion_payload_bytes(physical_lit(false)); let regs = InitialDynFilterRegs::new(vec![ InitialDynFilterReg::new("filter-1", vec![]) .with_initial_snapshot(InitialDynFilterSnapshot::new(payload, 7, false)), @@ -883,6 +947,45 @@ mod tests { assert_eq!(format!("{}", exprs[0]), "DynamicFilter [ false ]"); } + #[test] + fn bloom_initial_remote_dyn_filter_snapshot_initializes_runtime_filter() { + let regs_by_query = RemoteDynFilterRegistry::new(); + let query_id = test_remote_query_id(); + let schema = one_column_arrow_schema(); + let id_col = Arc::new(Column::new("id", 0)) as Arc; + let children = vec![Arc::clone(&id_col)]; + let lookup = test_lookup_expr(Arc::clone(&id_col), vec![10, 20, 30]); + let payload = DynFilterPayload::from_datafusion_expr_with_registered_children( + &lookup, + &children, + REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES, + &schema, + ) + .unwrap(); + assert!(matches!(payload, DynFilterPayload::JoinHashBloom(_))); + + let reg = InitialDynFilterReg::from_filter_id_and_children("filter-1", &children) + .unwrap() + .with_initial_snapshot(InitialDynFilterSnapshot::new(payload, 7, false)); + let regs = InitialDynFilterRegs::new(vec![reg]); + + register_initial_dyn_filter_regs( + ®s_by_query, + &query_id, + test_remote_dyn_filter_region_id(), + ®s, + ); + let exprs = + remote_dyn_filter_exprs_for_initial_regs(®s_by_query, &query_id, ®s, &schema); + + assert_eq!(exprs.len(), 1); + assert!( + format!("{}", exprs[0]).contains("bloom_probe"), + "expected Bloom initial snapshot to initialize runtime filter, got {}", + exprs[0] + ); + } + fn only_remote_dyn_filter( regs_by_query: &RemoteDynFilterRegistry, query_id: &QueryId, @@ -908,7 +1011,8 @@ mod tests { let filter_id = RemoteDynFilterId::new("filter-1"); register_empty_remote_dyn_filter(®s_by_query, &query_id); - let oversized = vec![0; REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES + 1]; + let oversized = + DynFilterPayload::Datafusion(vec![0; REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES + 1]); let outcome = apply_remote_dyn_filter_update( ®s_by_query, &query_id, @@ -920,8 +1024,14 @@ mod tests { assert_eq!(outcome, RemoteDynFilterUpdateOutcome::PayloadTooLarge); // The rejected generation must not become the pending generation. - let outcome = - apply_remote_dyn_filter_update(®s_by_query, &query_id, &filter_id, &[1], 0, false); + let outcome = apply_remote_dyn_filter_update( + ®s_by_query, + &query_id, + &filter_id, + &test_datafusion_payload(vec![1]), + 0, + false, + ); assert_eq!(outcome, RemoteDynFilterUpdateOutcome::Buffered); } diff --git a/src/query/src/dist_plan/dyn_filter_bridge.rs b/src/query/src/dist_plan/dyn_filter_bridge.rs index 2937924d5e7e..b0f20c5e3551 100644 --- a/src/query/src/dist_plan/dyn_filter_bridge.rs +++ b/src/query/src/dist_plan/dyn_filter_bridge.rs @@ -15,6 +15,7 @@ use std::any::Any; use std::sync::Arc; +use arrow_schema::Schema as ArrowSchema; use common_query::request::{ DynFilterPayload, INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, InitialDynFilterReg, InitialDynFilterRegs, InitialDynFilterSnapshot, REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES, @@ -33,6 +34,7 @@ pub(crate) struct CapturedDynFilter { filter_id: FilterId, initial_registration: InitialDynFilterReg, pub(crate) alive_dyn_filter: Arc, + pub(crate) input_schema: arrow_schema::SchemaRef, } #[derive(Debug, Clone)] @@ -45,6 +47,7 @@ pub(crate) struct RemoteDynFilterPushdown { pub(crate) fn capture_remote_dyn_filters_for_pushdown( remote_dyn_filter_producer_id: RemoteDynFilterProducerId, parent_filters: Vec>, + input_schema: &ArrowSchema, ) -> RemoteDynFilterPushdown { let mut pushed_down = Vec::with_capacity(parent_filters.len()); let mut captured_dyn_filters = Vec::new(); @@ -59,6 +62,7 @@ pub(crate) fn capture_remote_dyn_filters_for_pushdown( remote_dyn_filter_producer_id, producer_local_ordinal, alive_dyn_filter, + input_schema, ) { Ok(captured_dyn_filter) => { pushed_down.push(true); @@ -102,6 +106,7 @@ pub(crate) fn register_dyn_filters_for_region( let _ = registry.register_remote_dyn_filter( captured_dyn_filter.filter_id.clone(), captured_dyn_filter.alive_dyn_filter.clone(), + captured_dyn_filter.input_schema.clone(), ); let _ = registry .register_subscriber(&captured_dyn_filter.filter_id, Subscriber::new(region_id)); @@ -112,6 +117,7 @@ fn build_captured_dyn_filter( remote_dyn_filter_producer_id: RemoteDynFilterProducerId, producer_local_ordinal: usize, alive_dyn_filter: Arc, + input_schema: &ArrowSchema, ) -> Result { let children = alive_dyn_filter .children() @@ -128,8 +134,14 @@ fn build_captured_dyn_filter( Ok(CapturedDynFilter { filter_id, - initial_registration: attach_initial_snapshot(initial_registration, &alive_dyn_filter), + initial_registration: attach_initial_snapshot( + initial_registration, + &alive_dyn_filter, + &children, + input_schema, + ), alive_dyn_filter, + input_schema: Arc::new(input_schema.clone()), }) } @@ -146,8 +158,12 @@ fn validate_initial_registrations_for_pushdown( fn attach_initial_snapshot( initial_registration: InitialDynFilterReg, alive_dyn_filter: &DynamicFilterPhysicalExpr, + registered_children: &[Arc], + input_schema: &ArrowSchema, ) -> InitialDynFilterReg { - let Some(initial_snapshot) = initial_snapshot(alive_dyn_filter) else { + let Some(initial_snapshot) = + initial_snapshot(alive_dyn_filter, registered_children, input_schema) + else { return initial_registration; }; @@ -156,6 +172,8 @@ fn attach_initial_snapshot( fn initial_snapshot( alive_dyn_filter: &DynamicFilterPhysicalExpr, + registered_children: &[Arc], + input_schema: &ArrowSchema, ) -> Option { let generation = alive_dyn_filter.snapshot_generation(); let current = match alive_dyn_filter.current() { @@ -166,9 +184,11 @@ fn initial_snapshot( } }; - let payload = match DynFilterPayload::from_datafusion_expr( + let payload = match DynFilterPayload::from_datafusion_expr_with_registered_children( ¤t, + registered_children, REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES, + input_schema, ) { Ok(payload) => payload, Err(error) => { @@ -177,8 +197,7 @@ fn initial_snapshot( } }; - // Current DataFusion exposes `wait_complete()`, but no non-blocking completion getter. - let is_complete = false; + let is_complete = alive_dyn_filter.is_complete(); Some(InitialDynFilterSnapshot::new( payload, generation, @@ -232,6 +251,8 @@ mod tests { use std::hash::{Hash, Hasher}; use datafusion::execution::TaskContext; + use datafusion::physical_plan::joins::join_hash_map::{JoinHashMapType, JoinHashMapU32}; + use datafusion::physical_plan::joins::{HashTableLookupExpr, Map, SeededRandomState}; use datafusion_common::ScalarValue; use datafusion_expr::ColumnarValue; use datafusion_physical_expr::expressions::{Column, lit}; @@ -313,6 +334,15 @@ mod tests { RemoteDynFilterProducerId::new(value) } + fn test_arrow_schema() -> ArrowSchema { + ArrowSchema::new(vec![ + arrow_schema::Field::new("service", arrow_schema::DataType::Utf8, true), + arrow_schema::Field::new("host", arrow_schema::DataType::Utf8, true), + arrow_schema::Field::new("zone", arrow_schema::DataType::Utf8, true), + arrow_schema::Field::new("pod", arrow_schema::DataType::Utf8, true), + ]) + } + fn test_captured_dyn_filter( remote_dyn_filter_producer_id: RemoteDynFilterProducerId, producer_local_ordinal: usize, @@ -326,6 +356,7 @@ mod tests { vec![Arc::new(Column::new(column_name, column_index)) as Arc<_>], lit(true) as _, )), + &test_arrow_schema(), ) .unwrap() } @@ -345,6 +376,18 @@ mod tests { dyn_filter } + fn test_lookup_expr(child: Arc, hashes: Vec) -> Arc { + let mut hash_map = JoinHashMapU32::with_capacity(hashes.len().max(1)); + hash_map.update_from_iter(Box::new(hashes.iter().enumerate()), 0); + let map = Arc::new(Map::HashMap(Box::new(hash_map))); + Arc::new(HashTableLookupExpr::new( + vec![child], + SeededRandomState::with_seeds(1, 2, 3, 4), + map, + "lookup".to_string(), + )) + } + #[test] fn capture_remote_dyn_filters_for_pushdown_preserves_parent_filter_ordinals() { let parent_filters = vec![ @@ -361,9 +404,12 @@ mod tests { ]; 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; + let captured = capture_remote_dyn_filters_for_pushdown( + remote_dyn_filter_producer_id, + parent_filters, + &test_arrow_schema(), + ) + .captured_dyn_filters; assert_eq!(captured.len(), 2); assert_eq!( @@ -390,8 +436,11 @@ mod tests { ]; 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); + let pushdown = capture_remote_dyn_filters_for_pushdown( + remote_dyn_filter_producer_id, + parent_filters, + &test_arrow_schema(), + ); assert_eq!(pushdown.pushed_down, vec![false, true, false]); assert_eq!(pushdown.captured_dyn_filters.len(), 1); @@ -426,6 +475,7 @@ mod tests { let pushdown = capture_remote_dyn_filters_for_pushdown( test_remote_dyn_filter_producer_id(42), parent_filters, + &test_arrow_schema(), ); assert_eq!(pushdown.pushed_down, vec![false]); @@ -443,6 +493,7 @@ mod tests { let pushdown = capture_remote_dyn_filters_for_pushdown( test_remote_dyn_filter_producer_id(42), parent_filters, + &test_arrow_schema(), ); assert_eq!(pushdown.pushed_down, vec![true]); @@ -466,6 +517,7 @@ mod tests { let pushdown = capture_remote_dyn_filters_for_pushdown( test_remote_dyn_filter_producer_id(42), parent_filters, + &test_arrow_schema(), ); assert_eq!(pushdown.pushed_down, vec![true]); @@ -482,6 +534,64 @@ mod tests { )); } + #[test] + fn capture_remote_dyn_filters_for_pushdown_preserves_initial_snapshot_completion() { + 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(); + dyn_filter.mark_complete(); + 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, + &test_arrow_schema(), + ); + + let snapshot = pushdown.captured_dyn_filters[0] + .initial_registration + .initial_snapshot + .as_ref() + .unwrap(); + assert_eq!(snapshot.generation, 2); + assert!(snapshot.is_complete); + } + + #[test] + fn capture_remote_dyn_filters_for_pushdown_can_attach_bloom_initial_snapshot() { + let host = Arc::new(Column::new("host", 1)) as Arc; + let dyn_filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&host)], + lit(true) as _, + )); + dyn_filter + .update(test_lookup_expr(Arc::clone(&host), vec![10, 20, 30])) + .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, + &test_arrow_schema(), + ); + + assert_eq!(pushdown.pushed_down, vec![true]); + let snapshot = pushdown.captured_dyn_filters[0] + .initial_registration + .initial_snapshot + .as_ref() + .unwrap(); + match &snapshot.payload { + DynFilterPayload::JoinHashBloom(bloom) => { + assert_eq!(bloom.join_key_child_indices, vec![0]); + assert_ne!(bloom.hash_compat_fingerprint, 0); + } + other => panic!("expected Bloom initial snapshot, got {other:?}"), + } + } + #[test] fn capture_remote_dyn_filters_for_pushdown_rejects_oversized_snapshots() { let oversized_total_snapshot_bytes = REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES * 3 / 5; @@ -495,6 +605,7 @@ mod tests { let pushdown = capture_remote_dyn_filters_for_pushdown( test_remote_dyn_filter_producer_id(42), parent_filters, + &test_arrow_schema(), ); assert_eq!(pushdown.pushed_down, vec![false, false]); @@ -515,6 +626,7 @@ mod tests { let pushdown = capture_remote_dyn_filters_for_pushdown( test_remote_dyn_filter_producer_id(42), parent_filters, + &test_arrow_schema(), ); assert!(pushdown.captured_dyn_filters.is_empty()); @@ -537,6 +649,7 @@ mod tests { let pushdown = capture_remote_dyn_filters_for_pushdown( test_remote_dyn_filter_producer_id(42), parent_filters, + &test_arrow_schema(), ); assert!(pushdown.captured_dyn_filters.is_empty()); diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 262e391b556a..1fab5ac625a3 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -829,8 +829,11 @@ impl ExecutionPlan for MergeScanExec { updated_node: Some(new_self), }); }; - let remote_dyn_filter_pushdown = - capture_remote_dyn_filters_for_pushdown(remote_dyn_filter_producer_id, parent_filters); + let remote_dyn_filter_pushdown = capture_remote_dyn_filters_for_pushdown( + remote_dyn_filter_producer_id, + parent_filters, + self.arrow_schema.as_ref(), + ); *self.captured_remote_dyn_filters.lock().unwrap() = remote_dyn_filter_pushdown.captured_dyn_filters; let new_self = Arc::new(self.clone()); @@ -1057,6 +1060,11 @@ mod tests { let captured = capture_remote_dyn_filters_for_pushdown( RemoteDynFilterProducerId::new(42), vec![dyn_filter], + &ArrowSchema::new(vec![arrow_schema::Field::new( + "host", + arrow_schema::DataType::Utf8, + true, + )]), ); assert_eq!(captured.captured_dyn_filters.len(), 1); diff --git a/src/query/src/dist_plan/remote_dyn_filter_registry.rs b/src/query/src/dist_plan/remote_dyn_filter_registry.rs index 5abac13a57a8..03f2d7cb91b3 100644 --- a/src/query/src/dist_plan/remote_dyn_filter_registry.rs +++ b/src/query/src/dist_plan/remote_dyn_filter_registry.rs @@ -18,20 +18,18 @@ use std::sync::{Arc, Mutex, RwLock, Weak}; use std::time::Duration; use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate}; +use arrow_schema::SchemaRef as ArrowSchemaRef; use common_query::request::{DynFilterPayload, REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES}; use common_runtime::spawn_global; use common_telemetry::{debug, warn}; use datafusion_physical_expr::PhysicalExpr; -use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; +use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit as physical_lit}; use session::query_id::QueryId; use store_api::storage::RegionId; use tokio::sync::{Notify, watch}; use crate::dist_plan::FilterId; -use crate::metrics::{ - REMOTE_DYN_FILTER_ENCODE_TOTAL, REMOTE_DYN_FILTER_PAYLOAD_BYTES, - REMOTE_DYN_FILTER_UPDATE_RPC_TOTAL, -}; +use crate::metrics::REMOTE_DYN_FILTER_UPDATE_RPC_TOTAL; use crate::region_query::RegionQueryHandlerRef; const REMOTE_DYN_FILTER_RECONCILE_INTERVAL: Duration = Duration::from_secs(1); @@ -78,6 +76,7 @@ pub struct DynFilterEntry { subscribers: RwLock>, state: Mutex, subscriber_changed: Notify, + input_schema: ArrowSchemaRef, } #[derive(Debug, Default)] @@ -93,13 +92,18 @@ struct QueryDynFilterRegistryInner { } impl DynFilterEntry { - pub fn new(filter_id: FilterId, producer_filter: Arc) -> Self { + pub fn new( + filter_id: FilterId, + producer_filter: Arc, + input_schema: ArrowSchemaRef, + ) -> Self { Self { filter_id, producer_filter: Arc::downgrade(&producer_filter), subscribers: RwLock::new(HashSet::new()), state: Mutex::new(DynFilterEntryState::default()), subscriber_changed: Notify::new(), + input_schema, } } @@ -130,6 +134,14 @@ impl DynFilterEntry { true } + fn should_send_generation(&self, generation: u64, is_complete: bool) -> bool { + if is_complete { + return true; + } + + generation > self.state.lock().unwrap().last_sent_generation + } + fn try_mark_unregistered(&self) -> bool { let mut state = self.state.lock().unwrap(); if state.unregistered { @@ -213,13 +225,18 @@ impl QueryDynFilterRegistry { &self, filter_id: FilterId, producer_filter: Arc, + input_schema: ArrowSchemaRef, ) -> 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(), producer_filter)); + let entry = Arc::new(DynFilterEntry::new( + filter_id.clone(), + producer_filter, + input_schema, + )); inner.entries.insert(filter_id, entry.clone()); EntryRegistration::Inserted(entry) } @@ -366,6 +383,74 @@ async fn run_entry_fanout( unregister_entry_once_for_query(®ion_query_handler, query_id, &entry).await; } +/// Encoded payload pair for a fanout update. +#[derive(Debug, Clone)] +struct EncodedRemoteDynFilterUpdate { + typed_payload: DynFilterPayload, + legacy_payload: Vec, +} + +fn encode_legacy_datafusion_payload(expr: &Arc) -> Option> { + match DynFilterPayload::from_datafusion_expr(expr, REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES) { + Ok(DynFilterPayload::Datafusion(bytes)) => Some(bytes), + Ok(_) | Err(_) => None, + } +} + +fn encode_legacy_datafusion_or_true(current: &Arc) -> Option> { + encode_legacy_datafusion_payload(current).or_else(|| { + encode_legacy_datafusion_payload(&(physical_lit(true) as Arc)) + }) +} + +fn encode_remote_dyn_filter_update( + current: &Arc, + filter: &DynamicFilterPhysicalExpr, + entry: &DynFilterEntry, +) -> Option { + let children = filter.children().into_iter().cloned().collect::>(); + + let typed = match DynFilterPayload::from_datafusion_expr_with_registered_children( + current, + &children, + REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES, + entry.input_schema.as_ref(), + ) { + Ok(typed) => typed, + Err(error) => { + warn!(error; "Failed to encode typed remote dynamic filter update"); + return None; + } + }; + + let legacy_payload = match &typed { + DynFilterPayload::Datafusion(bytes) => bytes.clone(), + DynFilterPayload::JoinHashBloom(_) => { + // Build a safe legacy fallback for old DNs that cannot parse JoinHashBloom. + // Try Datafusion encoding first; if that fails, use lit(true). + let Some(bytes) = encode_legacy_datafusion_or_true(current) else { + warn!("Failed to encode legacy fallback payload for old DNs"); + return None; + }; + bytes + } + _ => { + warn!("Unsupported DynFilterPayload variant for fanout"); + return None; + } + }; + + if legacy_payload.is_empty() { + warn!("Failed to encode non-empty legacy remote dynamic filter payload for old DNs"); + return None; + } + + Some(EncodedRemoteDynFilterUpdate { + typed_payload: typed, + legacy_payload, + }) +} + async fn fanout_snapshot_for_query( query_id: QueryId, region_query_handler: &RegionQueryHandlerRef, @@ -379,9 +464,17 @@ async fn fanout_snapshot_for_query( return true; }; - // The entry-global watermark advances before best-effort fanout. A timed-out - // subscriber may miss this generation; later/complete snapshots supersede it, - // and RDF only prunes. + if !entry.should_send_generation(generation, is_complete) { + return true; + } + + let Some(encoded) = encode_remote_dyn_filter_update(¤t, filter, entry) else { + // Failed to encode; do not advance watermark. + return true; + }; + + // Advance watermark after encoding succeeds, so a failed encode doesn't + // skip a generation. if !is_complete && !entry.mark_generation_sent(generation) { return true; } @@ -390,39 +483,13 @@ async fn fanout_snapshot_for_query( let _ = entry.mark_generation_sent(generation); } - let payload = - match DynFilterPayload::from_datafusion_expr(¤t, REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES) - { - Ok(DynFilterPayload::Datafusion(payload)) => { - REMOTE_DYN_FILTER_ENCODE_TOTAL - .with_label_values(&["success"]) - .inc(); - REMOTE_DYN_FILTER_PAYLOAD_BYTES.observe(payload.len() as f64); - payload - } - Ok(_) => { - REMOTE_DYN_FILTER_ENCODE_TOTAL - .with_label_values(&["unsupported"]) - .inc(); - warn!("Ignored unsupported remote dynamic filter producer payload"); - return true; - } - Err(error) => { - REMOTE_DYN_FILTER_ENCODE_TOTAL - .with_label_values(&["error"]) - .inc(); - warn!(error; "Failed to encode remote dynamic filter producer snapshot"); - return true; - } - }; - fanout_update_for_query( query_id, region_query_handler, entry, generation, is_complete, - payload, + encoded, lifecycle_rx, control_rpc_timeout, ) @@ -437,7 +504,7 @@ async fn fanout_update_for_query( entry: &DynFilterEntry, generation: u64, is_complete: bool, - payload: Vec, + encoded: EncodedRemoteDynFilterUpdate, lifecycle_rx: &mut watch::Receiver<()>, control_rpc_timeout: Duration, ) -> bool { @@ -447,10 +514,10 @@ async fn fanout_update_for_query( for subscriber in entry.subscribers() { let update = RemoteDynFilterUpdate { filter_id: filter_id.clone(), - payload: payload.clone(), + payload: encoded.legacy_payload.clone(), generation, is_complete, - typed_payload: None, + typed_payload: Some(encoded.typed_payload.to_region_proto_payload()), }; match await_control_rpc_or_lifecycle_close( @@ -808,6 +875,8 @@ mod tests { use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate}; use async_trait::async_trait; use common_query::request::QueryRequest; + use datafusion::physical_plan::joins::join_hash_map::{JoinHashMapType, JoinHashMapU32}; + use datafusion::physical_plan::joins::{HashTableLookupExpr, Map, SeededRandomState}; use datafusion_physical_expr::expressions::{Column, lit}; use session::ReadPreference; use uuid::Uuid; @@ -825,6 +894,7 @@ mod tests { generation: u64, is_complete: bool, payload: Vec, + typed_payload: Option, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -920,6 +990,7 @@ mod tests { generation: update.generation, is_complete: update.is_complete, payload: update.payload, + typed_payload: update.typed_payload, }); if should_block { self.update_blocked.notify_one(); @@ -965,6 +1036,31 @@ mod tests { Arc::new(DynamicFilterPhysicalExpr::new(children, lit(true) as _)) } + fn test_lookup_expr(child: Arc, hashes: Vec) -> Arc { + let mut hash_map = JoinHashMapU32::with_capacity(hashes.len().max(1)); + hash_map.update_from_iter(Box::new(hashes.iter().enumerate()), 0); + let map = Arc::new(Map::HashMap(Box::new(hash_map))); + Arc::new(HashTableLookupExpr::new( + vec![child], + SeededRandomState::with_seeds(1, 2, 3, 4), + map, + "lookup".to_string(), + )) + } + + fn test_arrow_schema(names: &[&str]) -> ArrowSchemaRef { + Arc::new(arrow_schema::Schema::new( + names + .iter() + .map(|name| arrow_schema::Field::new(*name, arrow_schema::DataType::Utf8, true)) + .collect::>(), + )) + } + + fn test_entry_schema() -> ArrowSchemaRef { + test_arrow_schema(&["host"]) + } + #[test] fn registry_manager_returns_same_registry_for_same_query() { let manager = Arc::new(DynFilterRegistryManager::default()); @@ -1164,7 +1260,11 @@ mod tests { 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()) { + let entry = match registry.register_remote_dyn_filter( + filter_id.clone(), + filter.clone(), + test_entry_schema(), + ) { EntryRegistration::Inserted(entry) => entry, other => panic!("unexpected registration result: {other:?}"), }; @@ -1190,7 +1290,11 @@ mod tests { let registry = Arc::new(QueryDynFilterRegistry::new(query_id)); 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()) { + let entry = match registry.register_remote_dyn_filter( + filter_id.clone(), + filter.clone(), + test_entry_schema(), + ) { EntryRegistration::Inserted(entry) => entry, other => panic!("unexpected registration result: {other:?}"), }; @@ -1214,6 +1318,7 @@ mod tests { assert_eq!(updates[0].generation, filter.snapshot_generation()); assert!(!updates[0].is_complete); assert!(!updates[0].payload.is_empty()); + assert!(updates[0].typed_payload.is_some()); registry .fanout_snapshot(&handler_ref, &entry, filter.as_ref(), false) @@ -1251,6 +1356,59 @@ mod tests { assert_eq!(entry.subscribers().len(), 2); } + #[tokio::test] + async fn fanout_sends_join_hash_bloom_typed_payload_with_legacy_fallback() { + let query_id = test_query_id(10); + let registry = Arc::new(QueryDynFilterRegistry::new(query_id)); + 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(), + test_entry_schema(), + ) { + EntryRegistration::Inserted(entry) => entry, + other => panic!("unexpected registration result: {other:?}"), + }; + let subscriber = Subscriber::new(RegionId::new(1024, 7)); + assert_eq!( + registry.register_subscriber(&filter_id, subscriber.clone()), + SubscriberRegistration::Added + ); + + let host = Arc::new(Column::new("host", 0)) as Arc; + filter + .update(test_lookup_expr(Arc::clone(&host), vec![10, 20, 30])) + .unwrap(); + + let handler = Arc::new(RecordingRegionQueryHandler::default()); + let handler_ref = handler.clone() as RegionQueryHandlerRef; + registry + .fanout_snapshot(&handler_ref, &entry, filter.as_ref(), false) + .await; + + let updates = handler.updates(); + assert_eq!(updates.len(), 1); + assert_eq!(updates[0].region_id, subscriber.region_id()); + assert!( + !updates[0].payload.is_empty(), + "legacy fallback payload must stay non-empty for old DNs" + ); + + let typed = updates[0] + .typed_payload + .clone() + .expect("typed payload must be populated"); + let typed = DynFilterPayload::from_region_proto_payload(typed).unwrap(); + match typed { + DynFilterPayload::JoinHashBloom(bloom) => { + assert_eq!(bloom.join_key_child_indices, vec![0]); + assert_ne!(bloom.hash_compat_fingerprint, 0); + } + other => panic!("expected JoinHashBloom typed payload, got {other:?}"), + } + } + #[tokio::test] async fn fanout_task_waits_for_dynamic_filter_notifications() { let query_id = test_query_id(3); @@ -1259,9 +1417,11 @@ mod tests { let registry_weak = Arc::downgrade(lease.registry.as_ref().unwrap()); let filter = test_dyn_filter(&["host"]); let filter_id = test_filter_id(1); - let _ = lease - .registry() - .register_remote_dyn_filter(filter_id.clone(), filter.clone()); + let _ = lease.registry().register_remote_dyn_filter( + filter_id.clone(), + filter.clone(), + test_entry_schema(), + ); let subscriber = Subscriber::new(RegionId::new(1024, 7)); assert_eq!( lease @@ -1305,10 +1465,11 @@ mod tests { let registry_weak = Arc::downgrade(lease.registry.as_ref().unwrap()); let filter = test_dyn_filter(&["host"]); let filter_id = test_filter_id(1); - let entry = match lease - .registry() - .register_remote_dyn_filter(filter_id.clone(), filter.clone()) - { + let entry = match lease.registry().register_remote_dyn_filter( + filter_id.clone(), + filter.clone(), + test_entry_schema(), + ) { EntryRegistration::Inserted(entry) => entry, other => panic!("unexpected registration result: {other:?}"), }; @@ -1347,9 +1508,11 @@ mod tests { let registry_weak = Arc::downgrade(lease.registry.as_ref().unwrap()); let filter = test_dyn_filter(&["host"]); let filter_id = test_filter_id(1); - let _ = lease - .registry() - .register_remote_dyn_filter(filter_id.clone(), filter.clone()); + let _ = lease.registry().register_remote_dyn_filter( + filter_id.clone(), + filter.clone(), + test_entry_schema(), + ); let first_subscriber = Subscriber::new(RegionId::new(1024, 7)); assert_eq!( lease @@ -1401,9 +1564,11 @@ mod tests { let registry_weak = Arc::downgrade(lease.registry.as_ref().unwrap()); let filter = test_dyn_filter(&["host"]); let filter_id = test_filter_id(1); - let _ = lease - .registry() - .register_remote_dyn_filter(filter_id.clone(), filter.clone()); + let _ = lease.registry().register_remote_dyn_filter( + filter_id.clone(), + filter.clone(), + test_entry_schema(), + ); let subscriber = Subscriber::new(RegionId::new(1024, 7)); assert_eq!( lease @@ -1434,9 +1599,11 @@ mod tests { let registry_weak = Arc::downgrade(lease.registry.as_ref().unwrap()); let filter = test_dyn_filter(&["host"]); let filter_id = test_filter_id(1); - let _ = lease - .registry() - .register_remote_dyn_filter(filter_id.clone(), filter.clone()); + let _ = lease.registry().register_remote_dyn_filter( + filter_id.clone(), + filter.clone(), + test_entry_schema(), + ); let subscriber = Subscriber::new(RegionId::new(1024, 7)); assert_eq!( lease @@ -1475,9 +1642,11 @@ mod tests { let registry_weak = Arc::downgrade(lease.registry.as_ref().unwrap()); let filter = test_dyn_filter(&["host"]); let filter_id = test_filter_id(1); - let _ = lease - .registry() - .register_remote_dyn_filter(filter_id.clone(), filter.clone()); + let _ = lease.registry().register_remote_dyn_filter( + filter_id.clone(), + filter.clone(), + test_entry_schema(), + ); let subscriber = Subscriber::new(RegionId::new(1024, 7)); assert_eq!( lease @@ -1506,7 +1675,11 @@ mod tests { let registry = QueryDynFilterRegistry::new(query_id); 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()) { + let entry = match registry.register_remote_dyn_filter( + filter_id.clone(), + filter.clone(), + test_entry_schema(), + ) { EntryRegistration::Inserted(entry) => entry, other => panic!("unexpected registration result: {other:?}"), }; @@ -1596,7 +1769,7 @@ mod tests { let registry = QueryDynFilterRegistry::new(query_id); let filter = test_dyn_filter(&["host"]); let filter_id = test_filter_id(1); - let _ = registry.register_remote_dyn_filter(filter_id.clone(), filter); + let _ = registry.register_remote_dyn_filter(filter_id.clone(), filter, test_entry_schema()); let subscriber = Subscriber::new(RegionId::new(1024, 7)); assert_eq!( registry.register_subscriber(&filter_id, subscriber.clone()),