1#[cfg(test)]
16use std::cell::Cell;
17use std::sync::{Arc, Mutex};
18use std::time::Duration;
19
20use ahash::{HashMap, HashSet};
21use arrow_schema::{
22 ArrowError, DataType as ArrowDataType, Field, Schema as ArrowSchema,
23 SchemaRef as ArrowSchemaRef, SortOptions,
24};
25use async_stream::stream;
26use common_catalog::parse_catalog_and_schema_from_db_string;
27use common_error::ext::BoxedError;
28use common_plugins::GREPTIME_EXEC_READ_COST;
29use common_query::request::QueryRequest;
30use common_recordbatch::adapter::{RecordBatchMetrics, region_scan_output_bytes};
31use common_telemetry::tracing_context::TracingContext;
32use datafusion::execution::{SessionState, TaskContext};
33use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
34use datafusion::physical_plan::filter_pushdown::{
35 ChildPushdownResult, FilterPushdownPhase, FilterPushdownPropagation, PushedDown,
36};
37use datafusion::physical_plan::metrics::{
38 Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, MetricsSet, Time,
39};
40use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
41use datafusion::physical_plan::{
42 DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
43 SendableRecordBatchStream, StatisticsArgs, apply_expression_roots,
44};
45use datafusion_common::stats::Precision;
46use datafusion_common::tree_node::TreeNodeRecursion;
47use datafusion_common::{Column as ColumnExpr, DFSchemaRef, DataFusionError, Result, Statistics};
48use datafusion_expr::{Expr, Extension, FetchType, LogicalPlan, UserDefinedLogicalNodeCore};
49use datafusion_physical_expr::expressions::Column;
50use datafusion_physical_expr::{Distribution, EquivalenceProperties, PhysicalSortExpr};
51use datatypes::extension::json::is_any_json_extension_type;
52use futures_util::StreamExt;
53use greptime_proto::v1::region::RegionRequestHeader;
54use meter_core::data::ReadItem;
55use meter_macros::read_meter;
56use session::context::{
57 FLIGHT_METRICS_HEARTBEAT_INTERVAL, QueryContextRef,
58 SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY,
59};
60use store_api::metrics::{REGION_QUERY_CPU_TIME, REGION_QUERY_SCANNED_BYTES};
61use store_api::storage::RegionId;
62use table::table_name::TableName;
63use tokio::time;
64use tokio::time::Instant;
65use tracing::{Instrument, Span};
66
67use crate::dist_plan::analyzer::AliasMapping;
68use crate::dist_plan::analyzer::utils::patch_batch_timezone;
69use crate::dist_plan::dyn_filter_bridge::{
70 CapturedDynFilter, capture_remote_dyn_filters_for_pushdown,
71 query_context_with_refreshed_initial_dyn_filter_regs,
72 register_dyn_filter_subscribers_for_region, register_remote_dyn_filters,
73};
74use crate::dist_plan::{
75 FilterId, RemoteDynFilterProducerId, RemoteDynFilterRegistryLease, Subscriber,
76};
77use crate::metrics::{MERGE_SCAN_ERRORS_TOTAL, MERGE_SCAN_POLL_ELAPSED, MERGE_SCAN_REGIONS};
78use crate::options::{FlowQueryExtensions, remote_dyn_filter_pushdown_enabled_from_extensions};
79use crate::query_engine::QueryEngineState;
80use crate::region_query::RegionQueryHandlerRef;
81
82fn query_engine_state_from_task_context(context: &TaskContext) -> Option<Arc<QueryEngineState>> {
83 context.session_config().get_extension()
84}
85
86fn remote_plan_row_bound(plan: &LogicalPlan) -> Option<usize> {
92 match plan {
93 LogicalPlan::Limit(limit) => {
94 let input_bound = remote_plan_row_bound(&limit.input);
95 match limit.get_fetch_type() {
96 Ok(FetchType::Literal(Some(fetch))) => {
97 Some(input_bound.map_or(fetch, |bound| bound.min(fetch)))
98 }
99 _ => input_bound,
100 }
101 }
102 LogicalPlan::Sort(sort) => {
103 let input_bound = remote_plan_row_bound(&sort.input);
104 sort.fetch
105 .map(|fetch| input_bound.map_or(fetch, |bound| bound.min(fetch)))
106 .or(input_bound)
107 }
108 LogicalPlan::Projection(projection) => remote_plan_row_bound(&projection.input),
109 LogicalPlan::Filter(filter) => remote_plan_row_bound(&filter.input),
110 LogicalPlan::SubqueryAlias(alias) => remote_plan_row_bound(&alias.input),
111 LogicalPlan::Window(window) => remote_plan_row_bound(&window.input),
112 LogicalPlan::Repartition(repartition) => remote_plan_row_bound(&repartition.input),
113 LogicalPlan::Distinct(distinct) => remote_plan_row_bound(distinct.input()),
114 LogicalPlan::Aggregate(aggregate) => {
115 if aggregate
116 .group_expr
117 .iter()
118 .any(|expr| matches!(expr, Expr::GroupingSet(_)))
119 {
120 None
121 } else if aggregate.group_expr.is_empty() {
122 Some(1)
123 } else {
124 remote_plan_row_bound(&aggregate.input)
125 }
126 }
127 _ => None,
128 }
129}
130
131fn remote_dyn_filter_enabled(query_ctx: &QueryContextRef) -> Result<bool> {
132 remote_dyn_filter_pushdown_enabled_from_extensions(&query_ctx.extensions())
133 .map_err(|err| DataFusionError::External(Box::new(err)))
134}
135
136fn remote_schema_mismatch(message: impl Into<String>) -> DataFusionError {
137 DataFusionError::ArrowError(Box::new(ArrowError::SchemaError(message.into())), None)
138}
139
140fn record_merge_scan_schema_error() {
141 MERGE_SCAN_ERRORS_TOTAL.inc();
142
143 #[cfg(test)]
144 TEST_MERGE_SCAN_SCHEMA_ERRORS.with(|count| count.set(count.get() + 1));
145}
146
147#[cfg(test)]
148thread_local! {
149 static TEST_MERGE_SCAN_SCHEMA_ERRORS: Cell<u64> = const { Cell::new(0) };
153}
154
155#[cfg(test)]
156fn merge_scan_schema_error_count_for_test() -> u64 {
157 TEST_MERGE_SCAN_SCHEMA_ERRORS.with(Cell::get)
158}
159
160fn is_json_field(field: &Field) -> bool {
164 is_any_json_extension_type(field)
165 || field
166 .metadata()
167 .get(datatypes::schema::TYPE_KEY)
168 .map(String::as_ref)
169 == Some("Json")
170}
171
172fn fields_semantically_equal(expected_field: &Field, actual_field: &Field) -> bool {
182 expected_field.name() == actual_field.name()
183 && expected_field.data_type() == actual_field.data_type()
184 && expected_field.is_nullable() == actual_field.is_nullable()
185}
186
187fn json_fields_compatible(expected_field: &Field, actual_field: &Field) -> bool {
198 is_json_field(expected_field)
199 && is_json_field(actual_field)
200 && expected_field.name() == actual_field.name()
201 && expected_field.is_nullable() == actual_field.is_nullable()
202 && expected_field.metadata().get(datatypes::schema::TYPE_KEY)
204 == actual_field.metadata().get(datatypes::schema::TYPE_KEY)
205 && expected_field
207 .metadata()
208 .get(arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY)
209 == actual_field
210 .metadata()
211 .get(arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY)
212}
213
214fn validate_remote_schema(
226 expected: &ArrowSchema,
227 actual: &ArrowSchema,
228 source: &str,
229) -> Result<()> {
230 if expected.fields().len() != actual.fields().len() {
231 return Err(remote_schema_mismatch(format!(
232 "MergeScan {source} schema field count mismatch: expected {}, actual {}",
233 expected.fields().len(),
234 actual.fields().len()
235 )));
236 }
237
238 for (index, (expected_field, actual_field)) in expected
239 .fields()
240 .iter()
241 .zip(actual.fields().iter())
242 .enumerate()
243 {
244 if fields_semantically_equal(expected_field, actual_field) {
246 continue;
247 }
248
249 if json_fields_compatible(expected_field, actual_field) {
253 continue;
254 }
255
256 let timezone_only_difference = matches!(
258 (expected_field.data_type(), actual_field.data_type()),
259 (
260 ArrowDataType::Timestamp(expected_unit, expected_timezone),
261 ArrowDataType::Timestamp(actual_unit, actual_timezone),
262 ) if expected_unit == actual_unit
263 && expected_timezone != actual_timezone
264 && expected_field.name() == actual_field.name()
265 && expected_field.is_nullable() == actual_field.is_nullable()
266 );
267 if !timezone_only_difference {
268 return Err(remote_schema_mismatch(format!(
269 "MergeScan {source} schema field mismatch at position {index}: expected {:?}, actual {:?}",
270 expected_field, actual_field
271 )));
272 }
273 }
274
275 Ok(())
276}
277
278fn acquire_remote_dyn_filter_registry_lease(
279 context: &TaskContext,
280 query_ctx: &QueryContextRef,
281 captured_dyn_filters: &[CapturedDynFilter],
282) -> Option<RemoteDynFilterRegistryLease> {
283 if captured_dyn_filters.is_empty() {
284 return None;
285 }
286
287 let query_id = query_ctx.remote_query_id_value()?;
288 let query_engine_state = query_engine_state_from_task_context(context)?;
289 Some(
290 query_engine_state
291 .dyn_filter_registry_manager()
292 .acquire_lease(query_id),
293 )
294}
295
296fn register_remote_dyn_filters_for_region(
297 remote_dyn_filter_registry_lease: Option<&RemoteDynFilterRegistryLease>,
298 captured_dyn_filters: &[CapturedDynFilter],
299) {
300 if let Some(remote_dyn_filter_registry_lease) = remote_dyn_filter_registry_lease {
301 register_remote_dyn_filters(
302 remote_dyn_filter_registry_lease.registry(),
303 captured_dyn_filters,
304 );
305 }
306}
307
308struct SubscriberRollbackGuard<'a> {
309 registry: &'a crate::dist_plan::QueryDynFilterRegistry,
310 added: Vec<(FilterId, Subscriber)>,
311}
312
313impl<'a> SubscriberRollbackGuard<'a> {
314 fn new(
315 registry: &'a crate::dist_plan::QueryDynFilterRegistry,
316 added: Vec<(FilterId, Subscriber)>,
317 ) -> Self {
318 Self { registry, added }
319 }
320
321 fn disarm(&mut self) {
322 self.added.clear();
323 }
324}
325
326impl Drop for SubscriberRollbackGuard<'_> {
327 fn drop(&mut self) {
328 for (filter_id, subscriber) in &self.added {
329 self.registry.remove_subscriber(filter_id, subscriber);
330 }
331 }
332}
333
334#[derive(Debug, Hash, PartialEq, Eq, Clone)]
335pub struct MergeScanLogicalPlan {
336 input: LogicalPlan,
338 output_schema: DFSchemaRef,
340 is_placeholder: bool,
342 partition_cols: AliasMapping,
343 remote_dyn_filter_producer_id: Option<RemoteDynFilterProducerId>,
345}
346
347impl PartialOrd for MergeScanLogicalPlan {
348 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
349 let Self {
350 input,
351 output_schema,
352 is_placeholder,
353 partition_cols,
354 remote_dyn_filter_producer_id,
355 } = self;
356 let Self {
357 input: other_input,
358 output_schema: other_output_schema,
359 is_placeholder: other_is_placeholder,
360 partition_cols: other_partition_cols,
361 remote_dyn_filter_producer_id: other_remote_dyn_filter_producer_id,
362 } = other;
363
364 let ordering = (
365 input,
366 is_placeholder,
367 partition_cols,
368 remote_dyn_filter_producer_id,
369 )
370 .partial_cmp(&(
371 other_input,
372 other_is_placeholder,
373 other_partition_cols,
374 other_remote_dyn_filter_producer_id,
375 ));
376 match ordering {
377 Some(std::cmp::Ordering::Equal) if output_schema != other_output_schema => None,
378 ordering => ordering,
379 }
380 }
381}
382
383impl UserDefinedLogicalNodeCore for MergeScanLogicalPlan {
384 fn name(&self) -> &str {
385 Self::name()
386 }
387
388 fn inputs(&self) -> Vec<&LogicalPlan> {
391 vec![]
392 }
393
394 fn schema(&self) -> &datafusion_common::DFSchemaRef {
395 &self.output_schema
396 }
397
398 fn expressions(&self) -> Vec<datafusion_expr::Expr> {
400 vec![]
401 }
402
403 fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
404 write!(
405 f,
406 "MergeScan [is_placeholder={}, remote_input=[\n{}\n]]",
407 self.is_placeholder, self.input
408 )
409 }
410
411 fn with_exprs_and_inputs(
412 &self,
413 _exprs: Vec<datafusion::prelude::Expr>,
414 _inputs: Vec<LogicalPlan>,
415 ) -> Result<Self> {
416 Ok(self.clone())
417 }
418}
419
420impl MergeScanLogicalPlan {
421 pub fn new(input: LogicalPlan, is_placeholder: bool, partition_cols: AliasMapping) -> Self {
423 Self {
424 output_schema: input.schema().clone(),
425 input,
426 is_placeholder,
427 partition_cols,
428 remote_dyn_filter_producer_id: None,
429 }
430 }
431
432 pub(crate) fn with_output_schema(mut self, output_schema: DFSchemaRef) -> Self {
434 self.output_schema = output_schema;
435 self
436 }
437
438 pub(crate) fn with_remote_dyn_filter_producer_id(
439 mut self,
440 remote_dyn_filter_producer_id: RemoteDynFilterProducerId,
441 ) -> Self {
442 self.remote_dyn_filter_producer_id = Some(remote_dyn_filter_producer_id);
443 self
444 }
445
446 pub fn name() -> &'static str {
447 "MergeScan"
448 }
449
450 pub fn into_logical_plan(self) -> LogicalPlan {
452 LogicalPlan::Extension(Extension {
453 node: Arc::new(self),
454 })
455 }
456
457 pub fn is_placeholder(&self) -> bool {
458 self.is_placeholder
459 }
460
461 pub fn input(&self) -> &LogicalPlan {
462 &self.input
463 }
464
465 pub fn partition_cols(&self) -> &AliasMapping {
466 &self.partition_cols
467 }
468
469 pub fn remote_dyn_filter_producer_id(&self) -> Option<RemoteDynFilterProducerId> {
470 self.remote_dyn_filter_producer_id
471 }
472}
473
474#[derive(Clone)]
475pub struct MergeScanExec {
476 table: TableName,
477 regions: Vec<RegionId>,
478 plan: LogicalPlan,
479 arrow_schema: ArrowSchemaRef,
480 region_query_handler: RegionQueryHandlerRef,
481 metric: ExecutionPlanMetricsSet,
482 properties: Arc<PlanProperties>,
483 sub_stage_metrics: Arc<Mutex<HashMap<RegionId, RecordBatchMetrics>>>,
485 partition_metrics: Arc<Mutex<HashMap<usize, PartitionMetrics>>>,
487 query_ctx: QueryContextRef,
488 remote_dyn_filter_producer_id: Option<RemoteDynFilterProducerId>,
490 captured_remote_dyn_filters: Arc<Mutex<Vec<CapturedDynFilter>>>,
491 target_partition: usize,
492 partition_cols: AliasMapping,
493 enable_per_region_metrics: bool,
494}
495
496impl std::fmt::Debug for MergeScanExec {
497 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
498 f.debug_struct("MergeScanExec")
499 .field("table", &self.table)
500 .field("regions", &self.regions)
501 .field("plan", &self.plan)
502 .finish()
503 }
504}
505
506impl MergeScanExec {
507 #[allow(clippy::too_many_arguments)]
508 pub fn new(
509 session_state: &SessionState,
510 table: TableName,
511 regions: Vec<RegionId>,
512 plan: LogicalPlan,
513 arrow_schema: &ArrowSchema,
514 region_query_handler: RegionQueryHandlerRef,
515 query_ctx: QueryContextRef,
516 target_partition: usize,
517 partition_cols: AliasMapping,
518 remote_dyn_filter_producer_id: Option<RemoteDynFilterProducerId>,
519 enable_per_region_metrics: bool,
520 ) -> Result<Self> {
521 let arrow_schema = Arc::new(arrow_schema.clone());
522 let output_partition_count = Self::output_partition_count(regions.len(), target_partition);
523
524 let eq_properties = if let LogicalPlan::Sort(sort) = &plan
532 && output_partition_count >= regions.len()
533 {
534 let lex_ordering = sort
535 .expr
536 .iter()
537 .map(|sort_expr| {
538 let physical_expr = session_state
539 .create_physical_expr(sort_expr.expr.clone(), plan.schema())?;
540 Ok(PhysicalSortExpr::new(
541 physical_expr,
542 SortOptions {
543 descending: !sort_expr.asc,
544 nulls_first: sort_expr.nulls_first,
545 },
546 ))
547 })
548 .collect::<Result<Vec<_>>>()?;
549 EquivalenceProperties::new_with_orderings(arrow_schema.clone(), vec![lex_ordering])
550 } else {
551 EquivalenceProperties::new(arrow_schema.clone())
552 };
553
554 let partition_exprs = partition_cols
555 .iter()
556 .filter_map(|col| {
557 if let Some(first_alias) = col.1.first() {
558 session_state
559 .create_physical_expr(
560 Expr::Column(ColumnExpr::new_unqualified(
561 first_alias.name().to_string(),
562 )),
563 plan.schema(),
564 )
565 .ok()
566 } else {
567 None
568 }
569 })
570 .collect();
571 let partitioning = Partitioning::Hash(partition_exprs, output_partition_count);
572
573 let properties = Arc::new(PlanProperties::new(
574 eq_properties,
575 partitioning,
576 EmissionType::Incremental,
577 Boundedness::Bounded,
578 ));
579 Ok(Self {
580 table,
581 regions,
582 plan,
583 arrow_schema,
584 region_query_handler,
585 metric: ExecutionPlanMetricsSet::new(),
586 sub_stage_metrics: Arc::default(),
587 partition_metrics: Arc::default(),
588 properties,
589 query_ctx,
590 remote_dyn_filter_producer_id,
591 captured_remote_dyn_filters: Arc::default(),
592 target_partition,
593 partition_cols,
594 enable_per_region_metrics,
595 })
596 }
597
598 fn estimated_num_rows(&self) -> Precision<usize> {
603 if self.regions.is_empty() {
604 return Precision::Inexact(0);
605 }
606
607 let Some(rows_per_region) = remote_plan_row_bound(&self.plan) else {
608 return Precision::Absent;
609 };
610 rows_per_region
611 .checked_mul(self.regions.len())
612 .map_or(Precision::Absent, Precision::Inexact)
613 }
614
615 fn output_partition_count(num_regions: usize, target_partition: usize) -> usize {
617 num_regions.max(1).min(target_partition.max(1))
618 }
619
620 pub fn to_stream(
621 &self,
622 context: Arc<TaskContext>,
623 partition: usize,
624 ) -> Result<SendableRecordBatchStream> {
625 let regions = self.regions.clone();
627 let region_query_handler = self.region_query_handler.clone();
628 let metric = MergeScanMetric::new(&self.metric);
629 let arrow_schema = self.arrow_schema.clone();
630 let query_ctx = self.query_ctx.clone();
631 let sub_stage_metrics_moved = self.sub_stage_metrics.clone();
632 let partition_metrics_moved = self.partition_metrics.clone();
633 let plan = self.plan.clone();
634 let target_partition =
635 Self::output_partition_count(self.regions.len(), self.target_partition);
636 let remote_dyn_filter_enabled = remote_dyn_filter_enabled(&self.query_ctx)?;
637 let captured_remote_dyn_filters = if remote_dyn_filter_enabled {
638 self.captured_remote_dyn_filters()
639 } else {
640 Vec::new()
641 };
642 let dbname = context.task_id().unwrap_or_default();
643 let tracing_context = TracingContext::from_json(context.session_id().as_str());
644 let current_channel = self.query_ctx.channel();
645 let read_preference = self.query_ctx.read_preference();
646 let explain_verbose = self.query_ctx.explain_verbose();
647 let live_analyze_metrics = explain_verbose && self.query_ctx.live_analyze_metrics_enabled();
648 let remote_dyn_filter_registry_lease = acquire_remote_dyn_filter_registry_lease(
649 context.as_ref(),
650 &query_ctx,
651 &captured_remote_dyn_filters,
652 );
653
654 let stream = Box::pin(stream!({
655 let remote_dyn_filter_registry_lease = remote_dyn_filter_registry_lease;
656 if partition == 0 {
658 MERGE_SCAN_REGIONS.observe(regions.len() as f64);
659 }
660
661 let _finish_timer = metric.finish_time().timer();
662 let mut ready_timer = metric.ready_time().timer();
663 let mut first_consume_timer = Some(metric.first_consume_time().timer());
664
665 let partition_start = Instant::now();
667 let mut partition_ready_time: Option<Duration> = None;
668 let mut partition_first_consume_time: Option<Duration> = None;
669
670 for region_id in regions
671 .iter()
672 .skip(partition)
673 .step_by(target_partition)
674 .copied()
675 {
676 let region_span = tracing_context.attach(tracing::info_span!(
677 parent: &Span::current(),
678 "merge_scan_region",
679 region_id = %region_id,
680 partition = partition
681 ));
682 let region_start = Instant::now();
683 register_remote_dyn_filters_for_region(
684 remote_dyn_filter_registry_lease.as_ref(),
685 &captured_remote_dyn_filters,
686 );
687 let select_target_start = Instant::now();
688 let target = region_query_handler
689 .select_target(read_preference, region_id)
690 .instrument(region_span.clone())
691 .await
692 .map_err(|e| {
693 MERGE_SCAN_ERRORS_TOTAL.inc();
694 DataFusionError::External(Box::new(e))
695 })?;
696 let select_target_cost = select_target_start.elapsed();
697 let mut subscriber_rollback =
698 remote_dyn_filter_registry_lease.as_ref().map(|lease| {
699 SubscriberRollbackGuard::new(
700 lease.registry(),
701 register_dyn_filter_subscribers_for_region(
702 lease.registry(),
703 region_id,
704 target.clone(),
705 &captured_remote_dyn_filters,
706 ),
707 )
708 });
709 let mut region_query_ctx = query_context_with_refreshed_initial_dyn_filter_regs(
710 &query_ctx,
711 region_id,
712 &captured_remote_dyn_filters,
713 );
714 if live_analyze_metrics {
715 let remote_query_id = region_query_ctx.remote_query_id().map(str::to_string);
716 if let Some(remote_query_id) = remote_query_id {
717 region_query_ctx.set_extension(
718 SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY,
719 remote_query_id,
720 );
721 }
722 }
723 let request = QueryRequest {
724 header: Some(RegionRequestHeader {
725 tracing_context: tracing_context.to_w3c(),
726 dbname: dbname.clone(),
727 query_context: Some((®ion_query_ctx).into()),
728 }),
729 region_id,
730 plan: plan.clone(),
731 };
732 if explain_verbose {
733 common_telemetry::info!(
734 "Merge scan one region, partition: {}, region_id: {}",
735 partition,
736 region_id
737 );
738 }
739
740 let do_get_start = Instant::now();
741 let do_get_result = region_query_handler
742 .do_get(&target, request)
743 .instrument(region_span.clone())
744 .await;
745 if do_get_result.is_err() {
746 drop(subscriber_rollback.take());
747 }
748 let mut stream = do_get_result.map_err(|e| {
749 MERGE_SCAN_ERRORS_TOTAL.inc();
750 DataFusionError::External(Box::new(BoxedError::new(e)))
751 })?;
752
753 if let Some(subscriber_rollback) = subscriber_rollback.as_mut() {
754 subscriber_rollback.disarm();
755 }
756 let mut advertised_schema = stream.schema().arrow_schema().clone();
757 validate_remote_schema(
758 arrow_schema.as_ref(),
759 advertised_schema.as_ref(),
760 "advertised remote stream",
761 )
762 .inspect_err(|_| record_merge_scan_schema_error())?;
763 let do_get_cost = select_target_cost + do_get_start.elapsed();
764
765 if let Some(remote_dyn_filter_registry_lease) =
766 remote_dyn_filter_registry_lease.as_ref()
767 {
768 remote_dyn_filter_registry_lease
769 .ensure_fanout_task(region_query_handler.clone());
770 }
771
772 ready_timer.stop();
773 if partition_ready_time.is_none() {
774 partition_ready_time = Some(partition_start.elapsed());
775 }
776
777 let mut poll_duration = Duration::ZERO;
778 let mut poll_timer = Instant::now();
779 loop {
780 let batch = if live_analyze_metrics {
781 match time::timeout(
782 FLIGHT_METRICS_HEARTBEAT_INTERVAL,
783 stream.next().instrument(region_span.clone()),
784 )
785 .await
786 {
787 Ok(batch) => batch,
788 Err(_) => {
789 if let Some(metrics) = stream.metrics() {
790 let mut sub_stage_metrics =
791 sub_stage_metrics_moved.lock().unwrap();
792 sub_stage_metrics.insert(region_id, metrics);
793 }
794 continue;
795 }
796 }
797 } else {
798 stream.next().instrument(region_span.clone()).await
799 };
800 let Some(batch) = batch else {
801 if let Some(metrics) = stream.metrics() {
805 let load = region_scan_load(&metrics);
806 let (c, s) = parse_catalog_and_schema_from_db_string(&dbname);
807 let value = read_meter!(c, s, load, current_channel as u8);
808 metric.record_greptime_exec_cost(value as usize);
809 sub_stage_metrics_moved
810 .lock()
811 .unwrap()
812 .insert(region_id, metrics);
813 }
814 break;
815 };
816 let poll_elapsed = poll_timer.elapsed();
817 poll_duration += poll_elapsed;
818
819 let batch = batch
820 .map_err(|e| DataFusionError::External(Box::new(BoxedError::new(e))))?;
821 let df_batch = batch.into_df_record_batch();
822 if !Arc::ptr_eq(&advertised_schema, df_batch.schema_ref()) {
823 validate_remote_schema(
824 arrow_schema.as_ref(),
825 df_batch.schema_ref().as_ref(),
826 "remote record batch",
827 )
828 .inspect_err(|_| record_merge_scan_schema_error())?;
829 advertised_schema = df_batch.schema_ref().clone();
830 }
831 let batch =
832 patch_batch_timezone(arrow_schema.clone(), df_batch.columns().to_vec())?;
833 metric.record_output_batch_rows(batch.num_rows());
834 if let Some(mut first_consume_timer) = first_consume_timer.take() {
835 first_consume_timer.stop();
836 partition_first_consume_time = Some(partition_start.elapsed());
837 }
838
839 if let Some(metrics) = stream.metrics() {
840 let mut sub_stage_metrics = sub_stage_metrics_moved.lock().unwrap();
841 sub_stage_metrics.insert(region_id, metrics);
842 }
843
844 yield Ok(batch);
845 poll_timer = Instant::now();
847 }
848 if let Some(mut first_consume_timer) = first_consume_timer.take() {
851 first_consume_timer.stop();
852 partition_first_consume_time = Some(partition_start.elapsed());
853 }
854 let total_cost = region_start.elapsed();
855
856 let region_metrics = RegionMetrics {
858 region_id,
859 poll_duration,
860 do_get_cost,
861 total_cost,
862 };
863
864 {
866 let mut partition_metrics_guard = partition_metrics_moved.lock().unwrap();
867 let partition_metrics = partition_metrics_guard
868 .entry(partition)
869 .or_insert_with(|| PartitionMetrics::new(partition, explain_verbose));
870 partition_metrics.add_region_metrics(region_metrics);
871 }
872
873 if explain_verbose {
874 common_telemetry::info!(
875 "Merge scan finish one region, partition: {}, region_id: {}, poll_duration: {:?}, first_consume: {}, do_get_cost: {:?}",
876 partition,
877 region_id,
878 poll_duration,
879 metric.first_consume_time(),
880 do_get_cost
881 );
882 }
883
884 MERGE_SCAN_POLL_ELAPSED.observe(poll_duration.as_secs_f64());
885 }
886
887 ready_timer.stop();
890 if let Some(mut first_consume_timer) = first_consume_timer.take() {
891 first_consume_timer.stop();
892 }
893
894 let partition_finish_time = partition_start.elapsed();
896 {
897 let mut partition_metrics_guard = partition_metrics_moved.lock().unwrap();
898 if let Some(partition_metrics) = partition_metrics_guard.get_mut(&partition) {
899 partition_metrics.set_timings(
900 partition_ready_time.unwrap_or_default(),
901 partition_first_consume_time.unwrap_or_default(),
902 partition_finish_time,
903 );
904 partition_metrics.finish();
905 }
906 }
907 }));
908
909 Ok(Box::pin(RecordBatchStreamAdapter::new(
910 self.arrow_schema.clone(),
911 stream,
912 )))
913 }
914
915 pub fn try_with_new_distribution(&self, distribution: Distribution) -> Option<Self> {
916 let Distribution::KeyPartitioned(hash_exprs) = distribution else {
917 return None;
919 };
920
921 if let Partitioning::Hash(curr_dist, _) = &self.properties.partitioning
922 && curr_dist == &hash_exprs
923 {
924 return None;
926 }
927
928 let hash_expr_col_names: HashSet<_> = hash_exprs
929 .iter()
930 .filter_map(|expr| {
931 expr.downcast_ref::<Column>()
932 .map(|col_expr| col_expr.name())
933 })
934 .collect();
935
936 let covers_all_partition_cols = self.partition_cols.values().all(|aliases| {
937 aliases
938 .iter()
939 .any(|col| hash_expr_col_names.contains(col.name()))
940 });
941 if !covers_all_partition_cols {
942 return None;
943 }
944
945 let all_partition_col_aliases: HashSet<_> = self
946 .partition_cols
947 .values()
948 .flat_map(|aliases| aliases.iter().map(|c| c.name()))
949 .collect();
950 let overlaps: Vec<_> = hash_exprs
951 .iter()
952 .filter(|expr| {
953 expr.downcast_ref::<Column>()
954 .is_some_and(|col_expr| all_partition_col_aliases.contains(col_expr.name()))
955 })
956 .cloned()
957 .collect();
958
959 if overlaps.is_empty() {
960 return None;
961 }
962
963 Some(Self {
964 table: self.table.clone(),
965 regions: self.regions.clone(),
966 plan: self.plan.clone(),
967 arrow_schema: self.arrow_schema.clone(),
968 region_query_handler: self.region_query_handler.clone(),
969 metric: self.metric.clone(),
970 properties: Arc::new(PlanProperties::new(
971 self.properties.eq_properties.clone(),
972 Partitioning::Hash(overlaps, self.partition_count()),
973 self.properties.emission_type,
974 self.properties.boundedness,
975 )),
976 sub_stage_metrics: self.sub_stage_metrics.clone(),
977 partition_metrics: self.partition_metrics.clone(),
978 query_ctx: self.query_ctx.clone(),
979 remote_dyn_filter_producer_id: self.remote_dyn_filter_producer_id,
980 captured_remote_dyn_filters: self.captured_remote_dyn_filters.clone(),
981 target_partition: self.target_partition,
982 partition_cols: self.partition_cols.clone(),
983 enable_per_region_metrics: self.enable_per_region_metrics,
984 })
985 }
986
987 fn captured_remote_dyn_filters(&self) -> Vec<CapturedDynFilter> {
988 self.captured_remote_dyn_filters.lock().unwrap().clone()
989 }
990
991 pub fn sub_stage_metrics(&self) -> Vec<RecordBatchMetrics> {
992 let sub_stage_metrics = self.sub_stage_metrics.lock().unwrap();
993 let mut metrics: Vec<_> = sub_stage_metrics.iter().collect();
994 metrics.sort_unstable_by_key(|(region_id, _)| **region_id);
995 metrics
996 .into_iter()
997 .map(|(_, metrics)| metrics.clone())
998 .collect()
999 }
1000
1001 pub fn regions(&self) -> &[RegionId] {
1002 &self.regions
1003 }
1004
1005 pub fn is_flow_sink_scan(&self) -> bool {
1006 let Some(sink_table_id) =
1007 FlowQueryExtensions::parse_flow_extensions(&self.query_ctx.extensions())
1008 .ok()
1009 .flatten()
1010 .and_then(|extensions| extensions.sink_table_id)
1011 else {
1012 return false;
1013 };
1014
1015 !self.regions.is_empty()
1016 && self
1017 .regions
1018 .iter()
1019 .all(|region_id| region_id.table_id() == sink_table_id)
1020 }
1021
1022 pub fn partition_count(&self) -> usize {
1023 Self::output_partition_count(self.regions.len(), self.target_partition)
1024 }
1025
1026 pub fn region_count(&self) -> usize {
1027 self.regions.len()
1028 }
1029
1030 fn partition_metrics(&self) -> Vec<PartitionMetrics> {
1031 self.partition_metrics
1032 .lock()
1033 .unwrap()
1034 .values()
1035 .cloned()
1036 .collect()
1037 }
1038}
1039
1040#[cfg(test)]
1041impl MergeScanExec {
1042 fn remote_dyn_filter_producer_id(&self) -> Option<RemoteDynFilterProducerId> {
1043 self.remote_dyn_filter_producer_id
1044 }
1045}
1046
1047#[derive(Debug, Clone)]
1049struct RegionMetrics {
1050 region_id: RegionId,
1051 poll_duration: Duration,
1052 do_get_cost: Duration,
1053 total_cost: Duration,
1055}
1056
1057#[derive(Debug, Clone)]
1059struct PartitionMetrics {
1060 partition: usize,
1061 region_metrics: Vec<RegionMetrics>,
1062 total_poll_duration: Duration,
1063 total_do_get_cost: Duration,
1064 total_regions: usize,
1065 ready_time: Duration,
1067 first_consume_time: Duration,
1069 finish_time: Duration,
1071 explain_verbose: bool,
1072 finished: bool,
1073}
1074
1075impl PartitionMetrics {
1076 fn new(partition: usize, explain_verbose: bool) -> Self {
1077 Self {
1078 partition,
1079 region_metrics: Vec::new(),
1080 total_poll_duration: Duration::ZERO,
1081 total_do_get_cost: Duration::ZERO,
1082 total_regions: 0,
1083 ready_time: Duration::ZERO,
1084 first_consume_time: Duration::ZERO,
1085 finish_time: Duration::ZERO,
1086 explain_verbose,
1087 finished: false,
1088 }
1089 }
1090
1091 fn add_region_metrics(&mut self, region_metrics: RegionMetrics) {
1092 self.total_poll_duration += region_metrics.poll_duration;
1093 self.total_do_get_cost += region_metrics.do_get_cost;
1094 self.total_regions += 1;
1095 self.region_metrics.push(region_metrics);
1096 }
1097
1098 fn set_timings(
1100 &mut self,
1101 ready_time: Duration,
1102 first_consume_time: Duration,
1103 finish_time: Duration,
1104 ) {
1105 self.ready_time = ready_time;
1106 self.first_consume_time = first_consume_time;
1107 self.finish_time = finish_time;
1108 }
1109
1110 fn finish(&mut self) {
1112 if self.finished {
1113 return;
1114 }
1115 self.finished = true;
1116 self.log_metrics();
1117 }
1118
1119 fn log_metrics(&self) {
1121 if self.explain_verbose {
1122 common_telemetry::info!(
1123 "MergeScan partition {} finished: {} regions, total_poll_duration: {:?}, total_do_get_cost: {:?}, ready_time: {:?}, first_consume_time: {:?}, finish_time: {:?}",
1124 self.partition,
1125 self.total_regions,
1126 self.total_poll_duration,
1127 self.total_do_get_cost,
1128 self.ready_time,
1129 self.first_consume_time,
1130 self.finish_time
1131 );
1132 } else {
1133 common_telemetry::debug!(
1134 "MergeScan partition {} finished: {} regions, total_poll_duration: {:?}, total_do_get_cost: {:?}, ready_time: {:?}, first_consume_time: {:?}, finish_time: {:?}",
1135 self.partition,
1136 self.total_regions,
1137 self.total_poll_duration,
1138 self.total_do_get_cost,
1139 self.ready_time,
1140 self.first_consume_time,
1141 self.finish_time
1142 );
1143 }
1144 }
1145}
1146
1147impl Drop for PartitionMetrics {
1148 fn drop(&mut self) {
1149 if !self.finished {
1150 self.log_metrics();
1151 }
1152 }
1153}
1154
1155impl ExecutionPlan for MergeScanExec {
1156 fn schema(&self) -> ArrowSchemaRef {
1157 self.arrow_schema.clone()
1158 }
1159
1160 fn properties(&self) -> &Arc<PlanProperties> {
1161 &self.properties
1162 }
1163
1164 fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
1165 vec![]
1166 }
1167
1168 fn apply_expressions(
1169 &self,
1170 f: &mut dyn FnMut(
1171 &Arc<dyn datafusion_physical_expr::PhysicalExpr>,
1172 ) -> Result<TreeNodeRecursion>,
1173 ) -> Result<TreeNodeRecursion> {
1174 let captured_remote_dyn_filters = self.captured_remote_dyn_filters();
1175 apply_expression_roots(
1176 captured_remote_dyn_filters
1177 .into_iter()
1178 .map(|captured_dyn_filter| {
1179 captured_dyn_filter.alive_dyn_filter
1180 as Arc<dyn datafusion_physical_expr::PhysicalExpr>
1181 }),
1182 f,
1183 )
1184 }
1185
1186 fn with_new_children(
1189 self: Arc<Self>,
1190 _children: Vec<Arc<dyn ExecutionPlan>>,
1191 ) -> Result<Arc<dyn ExecutionPlan>> {
1192 Ok(self.clone())
1193 }
1194
1195 fn handle_child_pushdown_result(
1196 &self,
1197 _phase: FilterPushdownPhase,
1198 child_pushdown_result: ChildPushdownResult,
1199 _config: &datafusion::config::ConfigOptions,
1200 ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> {
1201 let parent_filters = child_pushdown_result
1202 .parent_filters
1203 .into_iter()
1204 .map(|filter| filter.filter)
1205 .collect::<Vec<_>>();
1206
1207 if !remote_dyn_filter_enabled(&self.query_ctx)? {
1208 self.captured_remote_dyn_filters.lock().unwrap().clear();
1212 let new_self = Arc::new(self.clone());
1213
1214 return Ok(FilterPushdownPropagation {
1215 filters: parent_filters.into_iter().map(|_| PushedDown::No).collect(),
1216 updated_node: Some(new_self),
1217 });
1218 }
1219
1220 let Some(remote_dyn_filter_producer_id) = self.remote_dyn_filter_producer_id else {
1221 common_telemetry::warn!(
1223 "MergeScan remote dynamic filter producer id is not assigned; skipping remote dynamic filter pushdown"
1224 );
1225 self.captured_remote_dyn_filters.lock().unwrap().clear();
1226 let new_self = Arc::new(self.clone());
1227
1228 return Ok(FilterPushdownPropagation {
1229 filters: parent_filters.into_iter().map(|_| PushedDown::No).collect(),
1230 updated_node: Some(new_self),
1231 });
1232 };
1233 let remote_dyn_filter_pushdown =
1234 capture_remote_dyn_filters_for_pushdown(remote_dyn_filter_producer_id, parent_filters);
1235 *self.captured_remote_dyn_filters.lock().unwrap() =
1236 remote_dyn_filter_pushdown.captured_dyn_filters;
1237 let new_self = Arc::new(self.clone());
1238
1239 Ok(FilterPushdownPropagation {
1240 filters: remote_dyn_filter_pushdown
1241 .pushed_down
1242 .into_iter()
1243 .map(|pushdown_ready| {
1244 if pushdown_ready {
1245 PushedDown::Yes
1246 } else {
1247 PushedDown::No
1248 }
1249 })
1250 .collect(),
1251 updated_node: Some(new_self),
1252 })
1253 }
1254
1255 fn execute(
1256 &self,
1257 partition: usize,
1258 context: Arc<TaskContext>,
1259 ) -> Result<SendableRecordBatchStream> {
1260 self.to_stream(context, partition)
1261 }
1262
1263 fn metrics(&self) -> Option<MetricsSet> {
1264 Some(self.metric.clone_inner())
1265 }
1266
1267 fn statistics_from_inputs(
1268 &self,
1269 _input_stats: &[Arc<Statistics>],
1270 args: &StatisticsArgs,
1271 ) -> Result<Arc<Statistics>> {
1272 if args.partition().is_some() {
1273 return Ok(Arc::new(Statistics::new_unknown(&self.arrow_schema)));
1274 }
1275
1276 let mut statistics = Statistics::new_unknown(&self.arrow_schema);
1277 statistics.num_rows = self.estimated_num_rows();
1278 Ok(Arc::new(statistics))
1279 }
1280
1281 fn name(&self) -> &str {
1282 "MergeScanExec"
1283 }
1284}
1285
1286impl DisplayAs for MergeScanExec {
1287 fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
1288 write!(f, "MergeScanExec: peers=[")?;
1289 for region_id in self.regions.iter() {
1290 write!(f, "{}, ", region_id)?;
1291 }
1292 write!(f, "]")?;
1293
1294 if matches!(t, DisplayFormatType::Verbose) {
1295 let partition_metrics = self.partition_metrics();
1296 if !partition_metrics.is_empty() {
1297 write!(f, ", metrics={{")?;
1298 for (i, pm) in partition_metrics.iter().enumerate() {
1299 if i > 0 {
1300 write!(f, ", ")?;
1301 }
1302 write!(
1303 f,
1304 "\"partition_{}\":{{\"regions\":{},\"total_poll_duration\":\"{:?}\",\"total_do_get_cost\":\"{:?}\",\"ready_time\":\"{:?}\",\"first_consume_time\":\"{:?}\",\"finish_time\":\"{:?}\",\"region_metrics\":[",
1305 pm.partition,
1306 pm.total_regions,
1307 pm.total_poll_duration,
1308 pm.total_do_get_cost,
1309 pm.ready_time,
1310 pm.first_consume_time,
1311 pm.finish_time
1312 )?;
1313 for (j, rm) in pm.region_metrics.iter().enumerate() {
1314 if j > 0 {
1315 write!(f, ",")?;
1316 }
1317 write!(
1318 f,
1319 "{{\"region_id\":\"{}\",\"poll_duration\":\"{:?}\",\"do_get_cost\":\"{:?}\",\"total_cost\":\"{:?}\"}}",
1320 rm.region_id, rm.poll_duration, rm.do_get_cost, rm.total_cost
1321 )?;
1322 }
1323 write!(f, "]}}")?;
1324 }
1325 write!(f, "}}")?;
1326 }
1327 }
1328
1329 Ok(())
1330 }
1331}
1332
1333fn region_scan_load(metrics: &RecordBatchMetrics) -> ReadItem {
1334 ReadItem {
1335 cpu_time: metrics.elapsed_compute as u64,
1336 table_scan: region_scan_output_bytes(metrics) as u64,
1337 }
1338}
1339
1340fn report_region_query_load(region_id: RegionId, load: &ReadItem) {
1341 let region_id = region_id.to_string();
1342 REGION_QUERY_CPU_TIME
1343 .with_label_values(&[®ion_id])
1344 .inc_by(load.cpu_time);
1345 REGION_QUERY_SCANNED_BYTES
1346 .with_label_values(&[®ion_id])
1347 .inc_by(load.table_scan);
1348}
1349
1350fn query_load_region_id(default_region_id: RegionId, metrics: &RecordBatchMetrics) -> RegionId {
1351 metrics
1352 .query_load_region_id
1353 .map(RegionId::from_u64)
1354 .unwrap_or(default_region_id)
1355}
1356
1357impl Drop for MergeScanExec {
1358 fn drop(&mut self) {
1359 if !self.enable_per_region_metrics {
1363 return;
1364 }
1365
1366 let metrics = self.sub_stage_metrics.lock().unwrap();
1367 for (region_id, metrics) in metrics.iter() {
1368 let load = region_scan_load(metrics);
1369 report_region_query_load(query_load_region_id(*region_id, metrics), &load);
1370 }
1371 }
1372}
1373
1374#[derive(Debug, Clone)]
1375struct MergeScanMetric {
1376 ready_time: Time,
1378 first_consume_time: Time,
1380 finish_time: Time,
1382 output_rows: Count,
1384
1385 greptime_exec_cost: Gauge,
1387}
1388
1389impl MergeScanMetric {
1390 pub fn new(metric: &ExecutionPlanMetricsSet) -> Self {
1391 Self {
1392 ready_time: MetricBuilder::new(metric).subset_time("ready_time", 1),
1393 first_consume_time: MetricBuilder::new(metric).subset_time("first_consume_time", 1),
1394 finish_time: MetricBuilder::new(metric).subset_time("finish_time", 1),
1395 output_rows: MetricBuilder::new(metric).output_rows(1),
1396 greptime_exec_cost: MetricBuilder::new(metric).gauge(GREPTIME_EXEC_READ_COST, 1),
1397 }
1398 }
1399
1400 pub fn ready_time(&self) -> &Time {
1401 &self.ready_time
1402 }
1403
1404 pub fn first_consume_time(&self) -> &Time {
1405 &self.first_consume_time
1406 }
1407
1408 pub fn finish_time(&self) -> &Time {
1409 &self.finish_time
1410 }
1411
1412 pub fn record_output_batch_rows(&self, num_rows: usize) {
1413 self.output_rows.add(num_rows);
1414 }
1415
1416 pub fn record_greptime_exec_cost(&self, metrics: usize) {
1417 self.greptime_exec_cost.add(metrics);
1418 }
1419}
1420
1421#[cfg(test)]
1422mod tests {
1423 use std::collections::{BTreeSet, HashMap as StdHashMap};
1424 use std::pin::Pin;
1425 use std::sync::atomic::{AtomicUsize, Ordering};
1426 use std::task::{Context, Poll};
1427
1428 use arrow::array::{Int64Array, TimestampMillisecondArray};
1429 use arrow_schema::{DataType as TestArrowDataType, Field, TimeUnit};
1430 use async_trait::async_trait;
1431 use common_base::Plugins;
1432 use common_error::ext::{ErrorExt, PlainError};
1433 use common_error::status_code::StatusCode;
1434 use common_meta::peer::Peer;
1435 use common_query::request::{
1436 INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, InitialDynFilterRegs,
1437 };
1438 use common_recordbatch::adapter::{PlanMetrics, RecordBatchMetrics};
1439 use common_recordbatch::{
1440 DfRecordBatch, EmptyRecordBatchStream, RecordBatch, RecordBatchStream,
1441 };
1442 use datafusion::config::ConfigOptions;
1443 use datafusion::execution::SessionStateBuilder;
1444 use datafusion::physical_plan::filter_pushdown::ChildFilterPushdownResult;
1445 use datafusion::physical_plan::repartition::RepartitionExec;
1446 use datafusion::physical_plan::{StatisticsArgs, StatisticsContext};
1447 use datafusion_common::TableReference;
1448 use datafusion_expr::{LogicalPlanBuilder, col, lit};
1449 use datafusion_physical_expr::expressions::{
1450 Column, DynamicFilterPhysicalExpr, lit as physical_lit,
1451 };
1452 use datafusion_physical_expr::{Distribution, PhysicalExpr};
1453 use datatypes::prelude::{ConcreteDataType, VectorRef};
1454 use datatypes::schema::{ColumnSchema, Schema};
1455 use datatypes::vectors::{Int64Vector, StringVector, TimestampMillisecondVector};
1456 use futures_util::{Stream, TryStreamExt};
1457 use session::ReadPreference;
1458 use session::context::QueryContext;
1459 use session::query_id::QueryId;
1460 use snafu::IntoError;
1461 use table::table::scan::REGION_SCAN_EXEC_NAME;
1462 use table::table_name::TableName;
1463 use tokio::sync::{Notify, oneshot};
1464 use uuid::Uuid;
1465
1466 use super::*;
1467 use crate::dist_plan::DynFilterRegistryManager;
1468 use crate::options::QueryOptions;
1469 use crate::query_engine::{QueryEngineContext, QueryEngineState};
1470 use crate::region_query::RegionQueryHandler;
1471
1472 fn test_target(id: u64) -> crate::region_query::RegionQueryTarget {
1473 crate::region_query::RegionQueryTarget::new(Peer {
1474 id,
1475 addr: format!("127.0.0.1:{id}"),
1476 })
1477 }
1478
1479 fn test_query_id(value: u128) -> QueryId {
1480 QueryId::from(Uuid::from_u128(value))
1481 }
1482
1483 fn merge_scan_exec_with_sorted_input(
1484 region_count: u64,
1485 target_partition: usize,
1486 ) -> MergeScanExec {
1487 let plan = LogicalPlanBuilder::empty(true)
1488 .project(vec![lit(1i64).alias("ts")])
1489 .unwrap()
1490 .sort(vec![col("ts").sort(false, true)])
1491 .unwrap()
1492 .build()
1493 .unwrap();
1494 let regions = (0..region_count)
1495 .map(|region_number| RegionId::new(1024, region_number as u32))
1496 .collect();
1497
1498 merge_scan_exec_with_plan(regions, plan, target_partition)
1499 }
1500
1501 fn merge_scan_exec_with_plan(
1502 regions: Vec<RegionId>,
1503 plan: LogicalPlan,
1504 target_partition: usize,
1505 ) -> MergeScanExec {
1506 let session_state = SessionStateBuilder::new().build();
1507 let schema = plan.schema().as_arrow().clone();
1508
1509 MergeScanExec::new(
1510 &session_state,
1511 TableName::new("catalog", "schema", "table"),
1514 regions,
1515 plan,
1516 &schema,
1517 Arc::new(TestRegionQueryHandler::default()),
1518 QueryContext::arc(),
1519 target_partition,
1520 AliasMapping::new(),
1521 None,
1522 false,
1523 )
1524 .unwrap()
1525 }
1526
1527 fn merge_scan_statistics(exec: &MergeScanExec) -> Arc<Statistics> {
1528 StatisticsContext::new()
1529 .compute(exec, &StatisticsArgs::new())
1530 .unwrap()
1531 }
1532
1533 fn task_context_with_engine_state(
1534 state: Arc<QueryEngineState>,
1535 query_ctx: QueryContextRef,
1536 ) -> Arc<TaskContext> {
1537 let mut session_state = state.session_state();
1538 session_state.config_mut().set_extension(state);
1539 QueryEngineContext::new(session_state, query_ctx).build_task_ctx()
1540 }
1541
1542 fn remote_dyn_filter_test_exec(
1543 handler: crate::region_query::RegionQueryHandlerRef,
1544 query_ctx: QueryContextRef,
1545 ) -> MergeScanExec {
1546 let session_state = SessionStateBuilder::new().build();
1547 let plan = LogicalPlanBuilder::empty(true)
1548 .project(vec![lit(1i32).alias("col1")])
1549 .unwrap()
1550 .build()
1551 .unwrap();
1552 let schema = plan.schema().as_arrow().clone();
1553
1554 MergeScanExec::new(
1555 &session_state,
1556 TableName::new("catalog", "schema", "table"),
1557 vec![RegionId::new(1024, 1)],
1558 plan,
1559 &schema,
1560 handler,
1561 query_ctx,
1562 1,
1563 AliasMapping::new(),
1564 Some(RemoteDynFilterProducerId::new(42)),
1565 false,
1566 )
1567 .unwrap()
1568 }
1569
1570 fn install_remote_dyn_filter(exec: &MergeScanExec) -> Arc<DynamicFilterPhysicalExpr> {
1571 let dyn_filter = Arc::new(DynamicFilterPhysicalExpr::new(
1572 vec![Arc::new(Column::new("host", 0)) as Arc<_>],
1573 physical_lit(true) as _,
1574 ));
1575 exec.handle_child_pushdown_result(
1576 FilterPushdownPhase::Post,
1577 ChildPushdownResult {
1578 parent_filters: vec![ChildFilterPushdownResult {
1579 filter: dyn_filter.clone() as Arc<dyn datafusion_physical_expr::PhysicalExpr>,
1580 child_results: vec![PushedDown::Yes],
1581 }],
1582 self_filters: Vec::new(),
1583 },
1584 &ConfigOptions::new(),
1585 )
1586 .unwrap();
1587 dyn_filter
1588 }
1589
1590 fn query_engine_state(
1591 handler: crate::region_query::RegionQueryHandlerRef,
1592 ) -> Arc<QueryEngineState> {
1593 Arc::new(QueryEngineState::new(
1594 catalog::memory::new_memory_catalog_manager().unwrap(),
1595 None,
1596 Some(handler),
1597 None,
1598 None,
1599 None,
1600 false,
1601 Plugins::default(),
1602 QueryOptions::default(),
1603 ))
1604 }
1605
1606 fn empty_record_batch_stream(
1607 request: &common_query::request::QueryRequest,
1608 ) -> common_recordbatch::SendableRecordBatchStream {
1609 let arrow_schema = request.plan.schema().as_arrow().clone();
1610 Box::pin(EmptyRecordBatchStream::new(Arc::new(
1611 datatypes::schema::Schema::try_from(Arc::new(arrow_schema)).unwrap(),
1612 )))
1613 }
1614
1615 fn pending_record_batch_stream(
1616 request: &common_query::request::QueryRequest,
1617 ) -> common_recordbatch::SendableRecordBatchStream {
1618 let stream = futures_util::stream::pending::<
1619 datafusion_common::Result<datafusion::arrow::record_batch::RecordBatch>,
1620 >();
1621 let arrow_schema = request.plan.schema().as_arrow().clone();
1622 let stream = RecordBatchStreamAdapter::new(Arc::new(arrow_schema), stream);
1623 Box::pin(
1624 common_recordbatch::adapter::RecordBatchStreamAdapter::try_new(Box::pin(stream))
1625 .unwrap(),
1626 )
1627 }
1628
1629 #[test]
1630 fn merge_scan_does_not_advertise_ordering_when_partition_may_merge_regions() {
1631 let exec = merge_scan_exec_with_sorted_input(3, 2);
1632
1633 assert!(
1634 exec.properties().output_ordering().is_none(),
1635 "target_partition < region_count means one output partition may concatenate multiple sorted region streams"
1636 );
1637 }
1638
1639 #[test]
1640 fn merge_scan_advertises_ordering_when_each_partition_reads_at_most_one_region() {
1641 let exec = merge_scan_exec_with_sorted_input(3, 3);
1642
1643 assert!(exec.properties().output_ordering().is_some());
1644 }
1645
1646 #[test]
1647 fn merge_scan_advertises_ordering_when_partitions_exceed_regions() {
1648 let exec = merge_scan_exec_with_sorted_input(3, 4);
1649
1650 assert!(exec.properties().output_ordering().is_some());
1651 }
1652
1653 #[test]
1654 fn merge_scan_reports_populated_partition_count() {
1655 let cases = [(0, 10, 1), (1, 10, 1), (3, 2, 2), (5, 10, 5), (3, 0, 1)];
1656
1657 for (region_count, target, expected) in cases {
1658 let exec = merge_scan_exec_with_sorted_input(region_count, target);
1659 assert_eq!(exec.partition_count(), expected);
1660 assert_eq!(
1661 exec.properties().output_partitioning().partition_count(),
1662 expected
1663 );
1664 }
1665 }
1666
1667 #[test]
1668 fn merge_scan_reports_only_deterministic_plan_bounds() {
1669 use datafusion::functions_aggregate::expr_fn::count;
1670 use datafusion_expr::GroupingSet;
1671
1672 let regions = vec![RegionId::new(1024, 1), RegionId::new(1024, 2)];
1673 let limited = LogicalPlanBuilder::empty(true)
1674 .project(vec![lit(1i64).alias("col")])
1675 .unwrap()
1676 .limit(0, Some(50))
1677 .unwrap()
1678 .build()
1679 .unwrap();
1680 assert_eq!(
1681 merge_scan_statistics(&merge_scan_exec_with_plan(regions.clone(), limited, 10))
1682 .num_rows,
1683 Precision::Inexact(100)
1684 );
1685
1686 let large_bound = i32::MAX as usize + 1;
1687 let large_limit = LogicalPlanBuilder::empty(true)
1688 .project(vec![lit(1i64).alias("col")])
1689 .unwrap()
1690 .limit(0, Some(large_bound))
1691 .unwrap()
1692 .build()
1693 .unwrap();
1694 assert_eq!(
1695 merge_scan_statistics(&merge_scan_exec_with_plan(
1696 vec![RegionId::new(1024, 1)],
1697 large_limit,
1698 10,
1699 ))
1700 .num_rows,
1701 Precision::Inexact(large_bound)
1702 );
1703
1704 let uncapped = LogicalPlanBuilder::empty(true)
1705 .project(vec![lit(1i64).alias("col")])
1706 .unwrap()
1707 .build()
1708 .unwrap();
1709 assert_eq!(
1710 merge_scan_statistics(&merge_scan_exec_with_plan(
1711 regions.clone(),
1712 uncapped.clone(),
1713 10,
1714 ))
1715 .num_rows,
1716 Precision::Absent
1717 );
1718 assert_eq!(
1719 merge_scan_statistics(&merge_scan_exec_with_plan(Vec::new(), uncapped, 10)).num_rows,
1720 Precision::Inexact(0)
1721 );
1722
1723 let global_aggregate = LogicalPlanBuilder::empty(true)
1724 .project(vec![lit(1i64).alias("col")])
1725 .unwrap()
1726 .limit(0, Some(0))
1727 .unwrap()
1728 .aggregate(Vec::<Expr>::new(), vec![count(lit(1))])
1729 .unwrap()
1730 .build()
1731 .unwrap();
1732 assert_eq!(
1733 merge_scan_statistics(&merge_scan_exec_with_plan(
1734 regions.clone(),
1735 global_aggregate,
1736 10,
1737 ))
1738 .num_rows,
1739 Precision::Inexact(2)
1740 );
1741
1742 let grouping_sets = LogicalPlanBuilder::empty(true)
1743 .project(vec![lit(1i64).alias("col")])
1744 .unwrap()
1745 .limit(0, Some(50))
1746 .unwrap()
1747 .aggregate(
1748 vec![Expr::GroupingSet(GroupingSet::GroupingSets(vec![
1749 vec![],
1750 vec![col("col")],
1751 ]))],
1752 Vec::<Expr>::new(),
1753 )
1754 .unwrap()
1755 .build()
1756 .unwrap();
1757 assert_eq!(
1758 merge_scan_statistics(&merge_scan_exec_with_plan(regions, grouping_sets, 10)).num_rows,
1759 Precision::Absent
1760 );
1761 }
1762
1763 #[test]
1764 fn sub_stage_metrics_are_sorted_by_region_id() {
1765 let exec = merge_scan_exec_with_sorted_input(0, 1);
1766 let higher_region = RegionId::new(1024, 2);
1767 let lower_region = RegionId::new(1024, 1);
1768 let higher_metrics = RecordBatchMetrics {
1769 plan_metrics: vec![PlanMetrics {
1770 plan: "higher region".to_string(),
1771 plan_name: "higher region".to_string(),
1772 level: 0,
1773 metrics: Vec::new(),
1774 }],
1775 ..Default::default()
1776 };
1777 let lower_metrics = RecordBatchMetrics {
1778 plan_metrics: vec![PlanMetrics {
1779 plan: "lower region".to_string(),
1780 plan_name: "lower region".to_string(),
1781 level: 0,
1782 metrics: Vec::new(),
1783 }],
1784 ..Default::default()
1785 };
1786
1787 let mut sub_stage_metrics = exec.sub_stage_metrics.lock().unwrap();
1788 sub_stage_metrics.insert(higher_region, higher_metrics);
1789 sub_stage_metrics.insert(lower_region, lower_metrics);
1790 drop(sub_stage_metrics);
1791
1792 let metrics = exec.sub_stage_metrics();
1793 let plans: Vec<_> = metrics
1794 .iter()
1795 .map(|metrics| metrics.plan_metrics[0].plan.as_str())
1796 .collect();
1797
1798 assert_eq!(plans, ["lower region", "higher region"]);
1799 }
1800
1801 #[test]
1802 fn remote_dyn_filter_producer_registration_defers_subscriber_registration() {
1803 let registry_manager = Arc::new(DynFilterRegistryManager::default());
1804 let query_ctx = QueryContext::arc();
1805 let query_id = query_ctx
1806 .remote_query_id_value()
1807 .expect("query context must have remote query id");
1808 let lease = registry_manager.acquire_lease(query_id);
1809 let region_id = RegionId::new(1024, 7);
1810 let dyn_filter = Arc::new(DynamicFilterPhysicalExpr::new(
1811 vec![Arc::new(Column::new("host", 0)) as Arc<_>],
1812 physical_lit(true) as _,
1813 )) as Arc<dyn datafusion_physical_expr::PhysicalExpr>;
1814 let captured = capture_remote_dyn_filters_for_pushdown(
1815 RemoteDynFilterProducerId::new(42),
1816 vec![dyn_filter],
1817 );
1818 assert_eq!(captured.captured_dyn_filters.len(), 1);
1819
1820 register_remote_dyn_filters_for_region(Some(&lease), &captured.captured_dyn_filters);
1821 let region_query_ctx = query_context_with_refreshed_initial_dyn_filter_regs(
1822 &query_ctx,
1823 region_id,
1824 &captured.captured_dyn_filters,
1825 );
1826
1827 let entries = lease.registry().entries();
1828 assert_eq!(entries.len(), 1);
1829 assert!(entries[0].subscribers().is_empty());
1830 assert!(
1831 !entries[0].fanout_started_for_test(),
1832 "fanout must start only after do_get succeeds"
1833 );
1834 assert!(
1835 region_query_ctx
1836 .extension(INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY)
1837 .is_some(),
1838 "initial RDF registrations must be present in the do_get query context"
1839 );
1840 }
1841
1842 #[tokio::test]
1843 async fn failed_do_get_preserves_status_code_and_rolls_back_subscriber() {
1844 let handler = Arc::new(FailingRegionQueryHandler::default());
1845 let query_ctx = QueryContext::arc();
1846 let state = Arc::new(QueryEngineState::new(
1847 catalog::memory::new_memory_catalog_manager().unwrap(),
1848 None,
1849 Some(handler.clone()),
1850 None,
1851 None,
1852 None,
1853 false,
1854 Plugins::default(),
1855 QueryOptions::default(),
1856 ));
1857 let lease = state
1858 .acquire_remote_dyn_filter_registry_lease(&query_ctx)
1859 .unwrap();
1860 handler.set_registry(lease.registry_arc_for_test());
1861 let task_ctx = task_context_with_engine_state(state, query_ctx.clone());
1862 let plan = LogicalPlanBuilder::empty(true)
1863 .project(vec![lit(1i32).alias("col1")])
1864 .unwrap()
1865 .build()
1866 .unwrap();
1867 let schema = plan.schema().as_arrow().clone();
1868 let exec = MergeScanExec::new(
1869 &SessionStateBuilder::new().build(),
1870 TableName::new("catalog", "schema", "table"),
1871 vec![RegionId::new(1024, 1)],
1872 plan,
1873 &schema,
1874 handler.clone(),
1875 query_ctx,
1876 1,
1877 AliasMapping::new(),
1878 Some(RemoteDynFilterProducerId::new(42)),
1879 false,
1880 )
1881 .unwrap();
1882 let dyn_filter = Arc::new(DynamicFilterPhysicalExpr::new(
1883 vec![Arc::new(Column::new("host", 0)) as Arc<_>],
1884 physical_lit(true) as _,
1885 )) as Arc<dyn datafusion_physical_expr::PhysicalExpr>;
1886 exec.handle_child_pushdown_result(
1887 FilterPushdownPhase::Post,
1888 ChildPushdownResult {
1889 parent_filters: vec![ChildFilterPushdownResult {
1890 filter: dyn_filter,
1891 child_results: vec![PushedDown::Yes],
1892 }],
1893 self_filters: Vec::new(),
1894 },
1895 &ConfigOptions::new(),
1896 )
1897 .unwrap();
1898
1899 let mut stream = common_recordbatch::adapter::RecordBatchStreamAdapter::try_new(
1900 exec.to_stream(task_ctx, 0).unwrap(),
1901 )
1902 .unwrap();
1903 let error = stream.next().await.unwrap().unwrap_err();
1904 assert_eq!(error.status_code(), StatusCode::RequestOutdated);
1905 assert_eq!(handler.do_get_calls.load(Ordering::SeqCst), 1);
1906 assert!(handler.saw_subscriber.load(Ordering::SeqCst));
1907
1908 let entries = lease.registry().entries();
1909 assert_eq!(entries.len(), 1);
1910 assert!(entries[0].subscribers().is_empty());
1911 assert!(!entries[0].fanout_started_for_test());
1912 }
1913
1914 #[tokio::test]
1915 async fn repartitioned_merge_scan_later_stream_error_preserves_status_code() {
1916 let region_id = RegionId::new(1024, 1);
1917 let handler = Arc::new(TestRegionQueryHandler::with_responses(vec![(
1918 region_id,
1919 int64_schema(&["a", "b"]),
1920 vec![Err(common_recordbatch::error::ExternalSnafu.into_error(
1921 BoxedError::new(PlainError::new(
1922 "neutral stream error".to_string(),
1923 StatusCode::RequestOutdated,
1924 )),
1925 ))],
1926 )]));
1927 let merge_scan = Arc::new(merge_scan_exec_with_handler(
1928 vec![region_id],
1929 expected_int64_schema(),
1930 handler,
1931 1,
1932 ));
1933 let repartition =
1934 RepartitionExec::try_new(merge_scan, Partitioning::RoundRobinBatch(2)).unwrap();
1935 assert_eq!(
1936 repartition
1937 .properties()
1938 .output_partitioning()
1939 .partition_count(),
1940 2
1941 );
1942
1943 let mut stream = common_recordbatch::adapter::RecordBatchStreamAdapter::try_new(
1944 repartition
1945 .execute(0, Arc::new(TaskContext::default()))
1946 .unwrap(),
1947 )
1948 .unwrap();
1949
1950 let error = stream.next().await.unwrap().unwrap_err();
1951 assert_eq!(error.status_code(), StatusCode::RequestOutdated);
1952 }
1953
1954 #[tokio::test]
1955 async fn aborting_pending_do_get_poll_rolls_back_subscriber_without_starting_fanout() {
1956 let handler = Arc::new(PendingDoGetHandler::default());
1957 let query_ctx = QueryContext::arc();
1958 let state = query_engine_state(handler.clone());
1959 let lease = state
1960 .acquire_remote_dyn_filter_registry_lease(&query_ctx)
1961 .unwrap();
1962 let exec = remote_dyn_filter_test_exec(handler.clone(), query_ctx.clone());
1963 install_remote_dyn_filter(&exec);
1964 let stream = exec
1965 .to_stream(task_context_with_engine_state(state, query_ctx), 0)
1966 .unwrap();
1967
1968 let poll = tokio::spawn(async move {
1969 let mut stream = stream;
1970 stream.next().await
1971 });
1972 handler.do_get_entered.notified().await;
1973 let entries = lease.registry().entries();
1974 assert_eq!(entries.len(), 1);
1975 assert_eq!(entries[0].subscribers().len(), 1);
1976 assert!(!entries[0].fanout_started_for_test());
1977
1978 poll.abort();
1979 assert!(poll.await.unwrap_err().is_cancelled());
1980
1981 assert!(entries[0].subscribers().is_empty());
1982 assert!(!entries[0].fanout_started_for_test());
1983 }
1984
1985 #[tokio::test]
1986 async fn failed_do_get_preserves_preexisting_duplicate_subscriber() {
1987 let handler = Arc::new(FailingRegionQueryHandler::default());
1988 let query_ctx = QueryContext::arc();
1989 let state = query_engine_state(handler.clone());
1990 let lease = state
1991 .acquire_remote_dyn_filter_registry_lease(&query_ctx)
1992 .unwrap();
1993 let exec = remote_dyn_filter_test_exec(handler.clone(), query_ctx.clone());
1994 install_remote_dyn_filter(&exec);
1995
1996 register_remote_dyn_filters_for_region(Some(&lease), &exec.captured_remote_dyn_filters());
1997 let entry = lease.registry().entries().pop().unwrap();
1998 let subscriber = Subscriber::new(RegionId::new(1024, 1), test_target(1));
1999 assert!(matches!(
2000 lease
2001 .registry()
2002 .register_subscriber(entry.filter_id(), subscriber.clone()),
2003 crate::dist_plan::SubscriberRegistration::Added
2004 ));
2005
2006 let mut stream = exec
2007 .to_stream(task_context_with_engine_state(state, query_ctx), 0)
2008 .unwrap();
2009 assert!(stream.next().await.unwrap().is_err());
2010
2011 assert_eq!(entry.subscribers(), vec![subscriber]);
2012 assert!(!entry.fanout_started_for_test());
2013 }
2014
2015 #[tokio::test]
2016 async fn failed_target_selection_registers_no_subscriber_or_fanout() {
2017 let handler = Arc::new(SelectTargetErrorHandler::default());
2018 let query_ctx = QueryContext::arc();
2019 let state = query_engine_state(handler.clone());
2020 let lease = state
2021 .acquire_remote_dyn_filter_registry_lease(&query_ctx)
2022 .unwrap();
2023 let exec = remote_dyn_filter_test_exec(handler.clone(), query_ctx.clone());
2024 install_remote_dyn_filter(&exec);
2025
2026 let mut stream = exec
2027 .to_stream(task_context_with_engine_state(state, query_ctx), 0)
2028 .unwrap();
2029 assert!(stream.next().await.unwrap().is_err());
2030
2031 assert_eq!(handler.select_target_calls.load(Ordering::SeqCst), 1);
2032 assert_eq!(handler.do_get_calls.load(Ordering::SeqCst), 0);
2033 let entries = lease.registry().entries();
2034 assert_eq!(entries.len(), 1);
2035 assert!(entries[0].subscribers().is_empty());
2036 assert!(!entries[0].fanout_started_for_test());
2037 }
2038
2039 #[tokio::test]
2040 async fn frozen_target_is_reused_for_do_get_update_and_unregister() {
2041 let (first_update_entered_tx, first_update_entered_rx) = oneshot::channel();
2042 let (release_first_update_tx, release_first_update_rx) = oneshot::channel();
2043 let (second_update_entered_tx, second_update_entered_rx) = oneshot::channel();
2044 let (release_second_update_tx, release_second_update_rx) = oneshot::channel();
2045 let (unregister_tx, unregister_rx) = oneshot::channel();
2046 let handler = Arc::new(RoutingRegionQueryHandler::new(
2047 first_update_entered_tx,
2048 release_first_update_rx,
2049 second_update_entered_tx,
2050 release_second_update_rx,
2051 unregister_tx,
2052 ));
2053 let query_ctx = QueryContext::arc();
2054 let state = query_engine_state(handler.clone());
2055 let lease = state
2056 .acquire_remote_dyn_filter_registry_lease(&query_ctx)
2057 .unwrap();
2058 let exec = remote_dyn_filter_test_exec(handler.clone(), query_ctx.clone());
2059 let dyn_filter = install_remote_dyn_filter(&exec);
2060
2061 let stream = exec
2062 .to_stream(task_context_with_engine_state(state, query_ctx), 0)
2063 .unwrap();
2064 let poll = tokio::spawn(async move {
2065 let mut stream = stream;
2066 stream.next().await
2067 });
2068 handler.do_get_entered.notified().await;
2069 assert_eq!(handler.do_get_targets(), vec![1]);
2070
2071 first_update_entered_rx.await.unwrap();
2072 dyn_filter.update(physical_lit(false) as _).unwrap();
2073 release_first_update_tx.send(()).unwrap();
2074 second_update_entered_rx.await.unwrap();
2075 assert_eq!(handler.update_targets(), vec![1, 1]);
2076
2077 poll.abort();
2078 assert!(poll.await.unwrap_err().is_cancelled());
2079 drop(exec);
2080 drop(dyn_filter);
2081 release_second_update_tx.send(()).unwrap();
2082 unregister_rx.await.unwrap();
2083 assert_eq!(handler.unregister_targets(), vec![1]);
2084 drop(lease);
2085 }
2086
2087 #[tokio::test]
2088 async fn immediate_eof_do_get_receives_refreshed_remote_dyn_filter_snapshot() {
2089 let handler = Arc::new(ImmediateEofRegionQueryHandler::default());
2090 let query_ctx = QueryContext::arc();
2091 let state = query_engine_state(handler.clone());
2092 let exec = remote_dyn_filter_test_exec(handler.clone(), query_ctx.clone());
2093 let dyn_filter = install_remote_dyn_filter(&exec);
2094 dyn_filter.update(physical_lit(false) as _).unwrap();
2095
2096 let mut stream = exec
2097 .to_stream(task_context_with_engine_state(state, query_ctx), 0)
2098 .unwrap();
2099 assert!(stream.next().await.is_none());
2100
2101 let registrations = handler.registrations();
2102 assert_eq!(registrations.regs.len(), 1);
2103 let snapshot = registrations.regs[0].initial_snapshot.as_ref().unwrap();
2104 assert!(snapshot.generation > 0);
2105 assert!(!snapshot.is_complete);
2106 assert_eq!(
2107 snapshot
2108 .payload
2109 .decode_datafusion_expr(
2110 &TaskContext::default(),
2111 &ArrowSchema::new(vec![arrow_schema::Field::new(
2112 "host",
2113 arrow_schema::DataType::Boolean,
2114 false,
2115 )]),
2116 common_query::request::REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES,
2117 )
2118 .unwrap()
2119 .to_string(),
2120 "false"
2121 );
2122 }
2123
2124 #[test]
2125 fn remote_dyn_filter_registry_cleanup_waits_for_last_query_scoped_stream_drop() {
2126 let registry_manager = Arc::new(DynFilterRegistryManager::default());
2127 let query_id = test_query_id(1);
2128
2129 let first = registry_manager.acquire_lease(query_id);
2130 let second = registry_manager.acquire_lease(query_id);
2131
2132 drop(first);
2133 assert_eq!(registry_manager.registry_count(), 1);
2134
2135 drop(second);
2136 assert_eq!(registry_manager.registry_count(), 0);
2137 }
2138
2139 #[test]
2140 fn remote_dyn_filter_registry_cleanup_shares_query_scope_across_independent_leases() {
2141 let registry_manager = Arc::new(DynFilterRegistryManager::default());
2142 let query_id = test_query_id(1);
2143
2144 let first_exec_like_lease = registry_manager.acquire_lease(query_id);
2145 let second_exec_like_lease = registry_manager.acquire_lease(query_id);
2146
2147 drop(first_exec_like_lease);
2148 assert_eq!(registry_manager.registry_count(), 1);
2149
2150 drop(second_exec_like_lease);
2151 assert_eq!(registry_manager.registry_count(), 0);
2152 }
2153
2154 struct TestRegionResponse {
2155 advertised_schema: Arc<Schema>,
2156 batches: Vec<common_recordbatch::error::Result<RecordBatch>>,
2157 }
2158
2159 #[derive(Default)]
2160 struct TestRegionQueryHandler {
2161 responses: HashMap<RegionId, TestRegionResponse>,
2162 }
2163
2164 impl TestRegionQueryHandler {
2165 fn new(responses: impl IntoIterator<Item = (RegionId, RecordBatch)>) -> Self {
2166 let responses = responses
2167 .into_iter()
2168 .map(|(region_id, batch)| {
2169 (
2170 region_id,
2171 TestRegionResponse {
2172 advertised_schema: batch.schema.clone(),
2173 batches: vec![Ok(batch)],
2174 },
2175 )
2176 })
2177 .collect();
2178 Self { responses }
2179 }
2180
2181 fn with_responses(
2182 responses: impl IntoIterator<
2183 Item = (
2184 RegionId,
2185 Arc<Schema>,
2186 Vec<common_recordbatch::error::Result<RecordBatch>>,
2187 ),
2188 >,
2189 ) -> Self {
2190 let responses = responses
2191 .into_iter()
2192 .map(|(region_id, advertised_schema, batches)| {
2193 (
2194 region_id,
2195 TestRegionResponse {
2196 advertised_schema,
2197 batches,
2198 },
2199 )
2200 })
2201 .collect();
2202 Self { responses }
2203 }
2204 }
2205
2206 struct TestRecordBatchStream {
2207 schema: Arc<Schema>,
2208 batches: Vec<common_recordbatch::error::Result<RecordBatch>>,
2209 }
2210
2211 impl Stream for TestRecordBatchStream {
2212 type Item = common_recordbatch::error::Result<RecordBatch>;
2213
2214 fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
2215 if self.batches.is_empty() {
2216 Poll::Ready(None)
2217 } else {
2218 Poll::Ready(Some(self.batches.remove(0)))
2219 }
2220 }
2221 }
2222
2223 impl RecordBatchStream for TestRecordBatchStream {
2224 fn schema(&self) -> Arc<Schema> {
2225 self.schema.clone()
2226 }
2227
2228 fn output_ordering(&self) -> Option<&[common_recordbatch::OrderOption]> {
2229 None
2230 }
2231
2232 fn metrics(&self) -> Option<RecordBatchMetrics> {
2233 None
2234 }
2235 }
2236
2237 #[derive(Default)]
2238 struct FailingRegionQueryHandler {
2239 do_get_calls: AtomicUsize,
2240 saw_subscriber: std::sync::atomic::AtomicBool,
2241 registry: std::sync::Mutex<Option<Arc<crate::dist_plan::QueryDynFilterRegistry>>>,
2242 }
2243
2244 impl FailingRegionQueryHandler {
2245 fn set_registry(&self, registry: Arc<crate::dist_plan::QueryDynFilterRegistry>) {
2246 *self.registry.lock().unwrap() = Some(registry);
2247 }
2248 }
2249
2250 #[async_trait]
2251 impl RegionQueryHandler for FailingRegionQueryHandler {
2252 async fn select_target(
2253 &self,
2254 _read_preference: ReadPreference,
2255 _region_id: RegionId,
2256 ) -> crate::error::Result<crate::region_query::RegionQueryTarget> {
2257 Ok(test_target(1))
2258 }
2259
2260 async fn do_get(
2261 &self,
2262 _target: &crate::region_query::RegionQueryTarget,
2263 _request: common_query::request::QueryRequest,
2264 ) -> crate::error::Result<common_recordbatch::SendableRecordBatchStream> {
2265 self.do_get_calls.fetch_add(1, Ordering::SeqCst);
2266 self.saw_subscriber.store(
2267 self.registry
2268 .lock()
2269 .unwrap()
2270 .as_ref()
2271 .is_some_and(|registry| {
2272 registry
2273 .entries()
2274 .iter()
2275 .any(|entry| !entry.subscribers().is_empty())
2276 }),
2277 Ordering::SeqCst,
2278 );
2279 Err(crate::error::Error::QueryExecution {
2280 source: BoxedError::new(PlainError::new(
2281 "neutral do_get error".to_string(),
2282 StatusCode::RequestOutdated,
2283 )),
2284 location: snafu::Location::default(),
2285 })
2286 }
2287
2288 async fn handle_remote_dyn_filter_update(
2289 &self,
2290 _target: &crate::region_query::RegionQueryTarget,
2291 _query_id: String,
2292 _update: api::v1::region::RemoteDynFilterUpdate,
2293 ) -> crate::error::Result<()> {
2294 unimplemented!("test only")
2295 }
2296
2297 async fn handle_remote_dyn_filter_unregister(
2298 &self,
2299 _target: &crate::region_query::RegionQueryTarget,
2300 _query_id: String,
2301 _unregister: api::v1::region::RemoteDynFilterUnregister,
2302 ) -> crate::error::Result<()> {
2303 unimplemented!("test only")
2304 }
2305 }
2306
2307 #[derive(Default)]
2308 struct PendingDoGetHandler {
2309 do_get_entered: Notify,
2310 never_complete: Notify,
2311 }
2312
2313 #[async_trait]
2314 impl RegionQueryHandler for PendingDoGetHandler {
2315 async fn select_target(
2316 &self,
2317 _read_preference: ReadPreference,
2318 _region_id: RegionId,
2319 ) -> crate::error::Result<crate::region_query::RegionQueryTarget> {
2320 Ok(test_target(1))
2321 }
2322
2323 async fn do_get(
2324 &self,
2325 _target: &crate::region_query::RegionQueryTarget,
2326 _request: common_query::request::QueryRequest,
2327 ) -> crate::error::Result<common_recordbatch::SendableRecordBatchStream> {
2328 self.do_get_entered.notify_one();
2329 self.never_complete.notified().await;
2330 unreachable!("the test aborts the pending do_get future")
2331 }
2332
2333 async fn handle_remote_dyn_filter_update(
2334 &self,
2335 _target: &crate::region_query::RegionQueryTarget,
2336 _query_id: String,
2337 _update: api::v1::region::RemoteDynFilterUpdate,
2338 ) -> crate::error::Result<()> {
2339 unreachable!("fanout must not start while do_get is pending")
2340 }
2341
2342 async fn handle_remote_dyn_filter_unregister(
2343 &self,
2344 _target: &crate::region_query::RegionQueryTarget,
2345 _query_id: String,
2346 _unregister: api::v1::region::RemoteDynFilterUnregister,
2347 ) -> crate::error::Result<()> {
2348 unreachable!("fanout must not start while do_get is pending")
2349 }
2350 }
2351
2352 #[derive(Default)]
2353 struct SelectTargetErrorHandler {
2354 select_target_calls: AtomicUsize,
2355 do_get_calls: AtomicUsize,
2356 }
2357
2358 #[async_trait]
2359 impl RegionQueryHandler for SelectTargetErrorHandler {
2360 async fn select_target(
2361 &self,
2362 _read_preference: ReadPreference,
2363 _region_id: RegionId,
2364 ) -> crate::error::Result<crate::region_query::RegionQueryTarget> {
2365 self.select_target_calls.fetch_add(1, Ordering::SeqCst);
2366 crate::error::UnimplementedSnafu {
2367 operation: "test target selection failure",
2368 }
2369 .fail()
2370 }
2371
2372 async fn do_get(
2373 &self,
2374 _target: &crate::region_query::RegionQueryTarget,
2375 _request: common_query::request::QueryRequest,
2376 ) -> crate::error::Result<common_recordbatch::SendableRecordBatchStream> {
2377 self.do_get_calls.fetch_add(1, Ordering::SeqCst);
2378 unreachable!("do_get must not run after select_target fails")
2379 }
2380
2381 async fn handle_remote_dyn_filter_update(
2382 &self,
2383 _target: &crate::region_query::RegionQueryTarget,
2384 _query_id: String,
2385 _update: api::v1::region::RemoteDynFilterUpdate,
2386 ) -> crate::error::Result<()> {
2387 unreachable!("fanout must not start after select_target fails")
2388 }
2389
2390 async fn handle_remote_dyn_filter_unregister(
2391 &self,
2392 _target: &crate::region_query::RegionQueryTarget,
2393 _query_id: String,
2394 _unregister: api::v1::region::RemoteDynFilterUnregister,
2395 ) -> crate::error::Result<()> {
2396 unreachable!("fanout must not start after select_target fails")
2397 }
2398 }
2399
2400 struct RoutingRegionQueryHandler {
2401 route: Mutex<crate::region_query::RegionQueryTarget>,
2402 do_get_entered: Notify,
2403 do_get_targets: Mutex<Vec<u64>>,
2404 update_targets: Mutex<Vec<u64>>,
2405 unregister_targets: Mutex<Vec<u64>>,
2406 update_calls: AtomicUsize,
2407 first_update_entered_tx: Mutex<Option<oneshot::Sender<()>>>,
2408 release_first_update_rx: Mutex<Option<oneshot::Receiver<()>>>,
2409 second_update_entered_tx: Mutex<Option<oneshot::Sender<()>>>,
2410 release_second_update_rx: Mutex<Option<oneshot::Receiver<()>>>,
2411 unregister_tx: Mutex<Option<oneshot::Sender<()>>>,
2412 }
2413
2414 impl RoutingRegionQueryHandler {
2415 fn new(
2416 first_update_entered_tx: oneshot::Sender<()>,
2417 release_first_update_rx: oneshot::Receiver<()>,
2418 second_update_entered_tx: oneshot::Sender<()>,
2419 release_second_update_rx: oneshot::Receiver<()>,
2420 unregister_tx: oneshot::Sender<()>,
2421 ) -> Self {
2422 Self {
2423 route: Mutex::new(test_target(1)),
2424 do_get_entered: Notify::new(),
2425 do_get_targets: Mutex::new(Vec::new()),
2426 update_targets: Mutex::new(Vec::new()),
2427 unregister_targets: Mutex::new(Vec::new()),
2428 update_calls: AtomicUsize::new(0),
2429 first_update_entered_tx: Mutex::new(Some(first_update_entered_tx)),
2430 release_first_update_rx: Mutex::new(Some(release_first_update_rx)),
2431 second_update_entered_tx: Mutex::new(Some(second_update_entered_tx)),
2432 release_second_update_rx: Mutex::new(Some(release_second_update_rx)),
2433 unregister_tx: Mutex::new(Some(unregister_tx)),
2434 }
2435 }
2436
2437 fn do_get_targets(&self) -> Vec<u64> {
2438 self.do_get_targets.lock().unwrap().clone()
2439 }
2440
2441 fn update_targets(&self) -> Vec<u64> {
2442 self.update_targets.lock().unwrap().clone()
2443 }
2444
2445 fn unregister_targets(&self) -> Vec<u64> {
2446 self.unregister_targets.lock().unwrap().clone()
2447 }
2448 }
2449
2450 #[async_trait]
2451 impl RegionQueryHandler for RoutingRegionQueryHandler {
2452 async fn select_target(
2453 &self,
2454 _read_preference: ReadPreference,
2455 _region_id: RegionId,
2456 ) -> crate::error::Result<crate::region_query::RegionQueryTarget> {
2457 let mut route = self.route.lock().unwrap();
2458 let target = route.clone();
2459 *route = test_target(2);
2460 Ok(target)
2461 }
2462
2463 async fn do_get(
2464 &self,
2465 target: &crate::region_query::RegionQueryTarget,
2466 request: common_query::request::QueryRequest,
2467 ) -> crate::error::Result<common_recordbatch::SendableRecordBatchStream> {
2468 self.do_get_targets.lock().unwrap().push(target.peer().id);
2469 self.do_get_entered.notify_one();
2470 Ok(pending_record_batch_stream(&request))
2471 }
2472
2473 async fn handle_remote_dyn_filter_update(
2474 &self,
2475 target: &crate::region_query::RegionQueryTarget,
2476 _query_id: String,
2477 _update: api::v1::region::RemoteDynFilterUpdate,
2478 ) -> crate::error::Result<()> {
2479 self.update_targets.lock().unwrap().push(target.peer().id);
2480 match self.update_calls.fetch_add(1, Ordering::SeqCst) {
2481 0 => {
2482 if let Some(tx) = self.first_update_entered_tx.lock().unwrap().take() {
2483 let _ = tx.send(());
2484 }
2485 let release = { self.release_first_update_rx.lock().unwrap().take() };
2486 if let Some(release) = release {
2487 let _ = release.await;
2488 }
2489 }
2490 1 => {
2491 if let Some(tx) = self.second_update_entered_tx.lock().unwrap().take() {
2492 let _ = tx.send(());
2493 }
2494 let release = { self.release_second_update_rx.lock().unwrap().take() };
2495 if let Some(release) = release {
2496 let _ = release.await;
2497 }
2498 }
2499 _ => {}
2500 }
2501 Ok(())
2502 }
2503
2504 async fn handle_remote_dyn_filter_unregister(
2505 &self,
2506 target: &crate::region_query::RegionQueryTarget,
2507 _query_id: String,
2508 _unregister: api::v1::region::RemoteDynFilterUnregister,
2509 ) -> crate::error::Result<()> {
2510 self.unregister_targets
2511 .lock()
2512 .unwrap()
2513 .push(target.peer().id);
2514 if let Some(tx) = self.unregister_tx.lock().unwrap().take() {
2515 let _ = tx.send(());
2516 }
2517 Ok(())
2518 }
2519 }
2520
2521 #[derive(Default)]
2522 struct ImmediateEofRegionQueryHandler {
2523 registrations: Mutex<Option<InitialDynFilterRegs>>,
2524 }
2525
2526 impl ImmediateEofRegionQueryHandler {
2527 fn registrations(&self) -> InitialDynFilterRegs {
2528 self.registrations.lock().unwrap().clone().unwrap()
2529 }
2530 }
2531
2532 #[async_trait]
2533 impl RegionQueryHandler for ImmediateEofRegionQueryHandler {
2534 async fn select_target(
2535 &self,
2536 _read_preference: ReadPreference,
2537 _region_id: RegionId,
2538 ) -> crate::error::Result<crate::region_query::RegionQueryTarget> {
2539 Ok(test_target(1))
2540 }
2541
2542 async fn do_get(
2543 &self,
2544 _target: &crate::region_query::RegionQueryTarget,
2545 request: common_query::request::QueryRequest,
2546 ) -> crate::error::Result<common_recordbatch::SendableRecordBatchStream> {
2547 let registrations = request
2548 .header
2549 .clone()
2550 .and_then(|header| header.query_context)
2551 .and_then(|query_context| {
2552 query_context
2553 .extensions
2554 .get(INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY)
2555 .cloned()
2556 })
2557 .map(|serialized| InitialDynFilterRegs::from_extension_value(&serialized).unwrap())
2558 .unwrap();
2559 *self.registrations.lock().unwrap() = Some(registrations);
2560 Ok(empty_record_batch_stream(&request))
2561 }
2562
2563 async fn handle_remote_dyn_filter_update(
2564 &self,
2565 _target: &crate::region_query::RegionQueryTarget,
2566 _query_id: String,
2567 _update: api::v1::region::RemoteDynFilterUpdate,
2568 ) -> crate::error::Result<()> {
2569 Ok(())
2570 }
2571
2572 async fn handle_remote_dyn_filter_unregister(
2573 &self,
2574 _target: &crate::region_query::RegionQueryTarget,
2575 _query_id: String,
2576 _unregister: api::v1::region::RemoteDynFilterUnregister,
2577 ) -> crate::error::Result<()> {
2578 Ok(())
2579 }
2580 }
2581
2582 #[async_trait]
2583 impl RegionQueryHandler for TestRegionQueryHandler {
2584 async fn select_target(
2585 &self,
2586 _read_preference: ReadPreference,
2587 _region_id: RegionId,
2588 ) -> crate::error::Result<crate::region_query::RegionQueryTarget> {
2589 Ok(test_target(1))
2590 }
2591
2592 async fn do_get(
2593 &self,
2594 _target: &crate::region_query::RegionQueryTarget,
2595 request: common_query::request::QueryRequest,
2596 ) -> crate::error::Result<common_recordbatch::SendableRecordBatchStream> {
2597 let response = self
2598 .responses
2599 .get(&request.region_id)
2600 .expect("test handler needs a response for every requested region");
2601 Ok(Box::pin(TestRecordBatchStream {
2602 schema: response.advertised_schema.clone(),
2603 batches: response
2604 .batches
2605 .iter()
2606 .map(|batch| match batch {
2607 Ok(batch) => Ok(batch.clone()),
2608 Err(error) => Err(common_recordbatch::error::ExternalSnafu.into_error(
2609 BoxedError::new(PlainError::new(
2610 error.to_string(),
2611 error.status_code(),
2612 )),
2613 )),
2614 })
2615 .collect(),
2616 }))
2617 }
2618
2619 async fn handle_remote_dyn_filter_update(
2620 &self,
2621 _target: &crate::region_query::RegionQueryTarget,
2622 _query_id: String,
2623 _update: api::v1::region::RemoteDynFilterUpdate,
2624 ) -> crate::error::Result<()> {
2625 unimplemented!("test only")
2626 }
2627
2628 async fn handle_remote_dyn_filter_unregister(
2629 &self,
2630 _target: &crate::region_query::RegionQueryTarget,
2631 _query_id: String,
2632 _unregister: api::v1::region::RemoteDynFilterUnregister,
2633 ) -> crate::error::Result<()> {
2634 unimplemented!("test only")
2635 }
2636 }
2637
2638 fn int64_schema(columns: &[&str]) -> Arc<Schema> {
2639 Arc::new(Schema::new(
2640 columns
2641 .iter()
2642 .map(|name| ColumnSchema::new(*name, ConcreteDataType::int64_datatype(), false))
2643 .collect(),
2644 ))
2645 }
2646
2647 fn record_batch(schema: Arc<Schema>, columns: Vec<VectorRef>) -> RecordBatch {
2648 RecordBatch::new(schema, columns).expect("test record batch must match its schema")
2649 }
2650
2651 fn expected_int64_schema() -> ArrowSchema {
2652 int64_schema(&["a", "b"]).arrow_schema().as_ref().clone()
2653 }
2654
2655 fn merge_scan_exec(
2656 responses: Vec<(RegionId, RecordBatch)>,
2657 expected_schema: ArrowSchema,
2658 target_partition: usize,
2659 ) -> MergeScanExec {
2660 let regions = responses.iter().map(|(region_id, _)| *region_id).collect();
2661 merge_scan_exec_with_handler(
2662 regions,
2663 expected_schema,
2664 Arc::new(TestRegionQueryHandler::new(responses)),
2665 target_partition,
2666 )
2667 }
2668
2669 fn merge_scan_exec_with_handler(
2670 regions: Vec<RegionId>,
2671 expected_schema: ArrowSchema,
2672 handler: Arc<TestRegionQueryHandler>,
2673 target_partition: usize,
2674 ) -> MergeScanExec {
2675 let plan = LogicalPlanBuilder::empty(true).build().unwrap();
2676 MergeScanExec::new(
2677 &SessionStateBuilder::new().build(),
2678 TableName::new("catalog", "schema", "table"),
2679 regions,
2680 plan,
2681 &expected_schema,
2682 handler,
2683 QueryContext::arc(),
2684 target_partition,
2685 AliasMapping::new(),
2686 None,
2687 false,
2688 )
2689 .unwrap()
2690 }
2691
2692 async fn collect_merge_scan(
2693 exec: MergeScanExec,
2694 ) -> datafusion_common::Result<Vec<DfRecordBatch>> {
2695 exec.execute(0, Arc::new(TaskContext::default()))?
2696 .try_collect()
2697 .await
2698 }
2699
2700 fn assert_int64_batch(batch: &DfRecordBatch, values: (i64, i64)) {
2701 assert_eq!(batch.schema().as_ref(), &expected_int64_schema());
2702 let a = batch
2703 .column(0)
2704 .as_any()
2705 .downcast_ref::<Int64Array>()
2706 .unwrap();
2707 let b = batch
2708 .column(1)
2709 .as_any()
2710 .downcast_ref::<Int64Array>()
2711 .unwrap();
2712 assert_eq!((a.value(0), b.value(0)), values);
2713 }
2714
2715 #[tokio::test]
2716 async fn qbs_merge_scan_remote_schema_identity_canonical_single_region() {
2717 let batch = record_batch(
2718 int64_schema(&["a", "b"]),
2719 vec![
2720 Arc::new(Int64Vector::from_slice([11])) as _,
2721 Arc::new(Int64Vector::from_slice([12])) as _,
2722 ],
2723 );
2724 let batches = collect_merge_scan(merge_scan_exec(
2725 vec![(RegionId::new(1024, 1), batch)],
2726 expected_int64_schema(),
2727 1,
2728 ))
2729 .await
2730 .unwrap();
2731 assert_eq!(batches.len(), 1);
2732 assert_int64_batch(&batches[0], (11, 12));
2733 }
2734
2735 #[tokio::test]
2736 async fn qbs_merge_scan_remote_schema_identity_canonical_two_regions() {
2737 let batch = || {
2738 record_batch(
2739 int64_schema(&["a", "b"]),
2740 vec![
2741 Arc::new(Int64Vector::from_slice([11])) as _,
2742 Arc::new(Int64Vector::from_slice([12])) as _,
2743 ],
2744 )
2745 };
2746 let batches = collect_merge_scan(merge_scan_exec(
2747 vec![
2748 (RegionId::new(1024, 1), batch()),
2749 (RegionId::new(1024, 2), batch()),
2750 ],
2751 expected_int64_schema(),
2752 1,
2753 ))
2754 .await
2755 .unwrap();
2756 assert_eq!(batches.len(), 2);
2757 for batch in &batches {
2758 assert_int64_batch(batch, (11, 12));
2759 }
2760 }
2761
2762 #[tokio::test]
2763 async fn qbs_merge_scan_remote_schema_identity_swapped_columns_never_relabels_positionally() {
2764 let batch = record_batch(
2765 int64_schema(&["b", "a"]),
2766 vec![
2767 Arc::new(Int64Vector::from_slice([2002])) as _,
2768 Arc::new(Int64Vector::from_slice([1002])) as _,
2769 ],
2770 );
2771 assert!(
2772 collect_merge_scan(merge_scan_exec(
2773 vec![(RegionId::new(1024, 1), batch)],
2774 expected_int64_schema(),
2775 1,
2776 ))
2777 .await
2778 .is_err()
2779 );
2780 }
2781
2782 #[tokio::test]
2783 async fn qbs_merge_scan_remote_schema_identity_allows_timestamp_timezone_only_patch() {
2784 let remote_arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
2785 "ts",
2786 TestArrowDataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
2787 false,
2788 )]));
2789 let remote_schema = Arc::new(Schema::try_from(remote_arrow_schema).unwrap());
2790 let timestamp_array: Arc<dyn arrow::array::Array> =
2791 Arc::new(TimestampMillisecondArray::from(vec![1002]).with_timezone("UTC"));
2792 let timestamp = TimestampMillisecondVector::try_from_arrow_array(timestamp_array).unwrap();
2793 let batch = record_batch(remote_schema, vec![Arc::new(timestamp) as _]);
2794 let expected_schema = ArrowSchema::new(vec![Field::new(
2795 "ts",
2796 TestArrowDataType::Timestamp(TimeUnit::Millisecond, Some("Asia/Shanghai".into())),
2797 false,
2798 )]);
2799 let batches = collect_merge_scan(merge_scan_exec(
2800 vec![(RegionId::new(1024, 1), batch)],
2801 expected_schema.clone(),
2802 1,
2803 ))
2804 .await
2805 .unwrap();
2806 assert_eq!(batches[0].schema().as_ref(), &expected_schema);
2807 }
2808
2809 #[tokio::test]
2810 async fn qbs_merge_scan_remote_schema_identity_rejects_incompatible_type() {
2811 let batch = record_batch(
2812 Arc::new(Schema::new(vec![
2813 ColumnSchema::new("a", ConcreteDataType::string_datatype(), false),
2814 ColumnSchema::new("b", ConcreteDataType::int64_datatype(), false),
2815 ])),
2816 vec![
2817 Arc::new(StringVector::from_slice(&["not-an-int"])) as _,
2818 Arc::new(Int64Vector::from_slice([12])) as _,
2819 ],
2820 );
2821 assert!(
2822 collect_merge_scan(merge_scan_exec(
2823 vec![(RegionId::new(1024, 1), batch)],
2824 expected_int64_schema(),
2825 1,
2826 ))
2827 .await
2828 .is_err()
2829 );
2830 }
2831
2832 #[tokio::test]
2833 async fn qbs_merge_scan_remote_schema_identity_rejects_too_few_columns() {
2834 let batch = record_batch(
2835 int64_schema(&["a"]),
2836 vec![Arc::new(Int64Vector::from_slice([11])) as _],
2837 );
2838 assert!(
2839 collect_merge_scan(merge_scan_exec(
2840 vec![(RegionId::new(1024, 1), batch)],
2841 expected_int64_schema(),
2842 1,
2843 ))
2844 .await
2845 .is_err()
2846 );
2847 }
2848
2849 #[tokio::test]
2850 async fn qbs_merge_scan_remote_schema_identity_rejects_too_many_columns() {
2851 let batch = record_batch(
2852 int64_schema(&["a", "b", "extra"]),
2853 vec![
2854 Arc::new(Int64Vector::from_slice([11])) as _,
2855 Arc::new(Int64Vector::from_slice([12])) as _,
2856 Arc::new(Int64Vector::from_slice([13])) as _,
2857 ],
2858 );
2859 assert!(
2860 collect_merge_scan(merge_scan_exec(
2861 vec![(RegionId::new(1024, 1), batch)],
2862 expected_int64_schema(),
2863 1,
2864 ))
2865 .await
2866 .is_err()
2867 );
2868 }
2869
2870 #[test]
2871 fn merge_scan_remote_schema_identity_allows_top_level_metadata_mismatch() {
2872 let fields = expected_int64_schema().fields().clone();
2873 let expected = ArrowSchema::new_with_metadata(
2874 fields.clone(),
2875 StdHashMap::from([("greptime:version".to_string(), "1".to_string())]),
2876 );
2877 let actual = ArrowSchema::new_with_metadata(
2878 fields,
2879 StdHashMap::from([("greptime:version".to_string(), "0".to_string())]),
2880 );
2881 assert!(validate_remote_schema(&expected, &actual, "test").is_ok());
2882 }
2883
2884 fn json_field_metadata(wire_form: bool, json_settings: &str) -> StdHashMap<String, String> {
2888 let mut metadata = StdHashMap::from([
2889 (datatypes::schema::TYPE_KEY.to_string(), "Json".to_string()),
2890 (
2891 arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY.to_string(),
2892 json_settings.to_string(),
2893 ),
2894 ]);
2895 if wire_form {
2896 metadata.insert(
2897 arrow_schema::extension::EXTENSION_TYPE_NAME_KEY.to_string(),
2898 "greptime.json".to_string(),
2899 );
2900 }
2901 metadata
2902 }
2903
2904 #[test]
2905 fn merge_scan_remote_schema_identity_accepts_json_wire_binary_vs_decoded_struct() {
2906 let json_settings = r#"{"json_settings":{"type_hints":[]}}"#;
2913 let expected = ArrowSchema::new(vec![
2914 Field::new("j", TestArrowDataType::Binary, true)
2915 .with_metadata(json_field_metadata(true, json_settings)),
2916 ]);
2917 let actual =
2918 ArrowSchema::new(vec![
2919 Field::new(
2920 "j",
2921 TestArrowDataType::Struct(arrow_schema::Fields::from(vec![Arc::new(
2922 Field::new("a", TestArrowDataType::Utf8View, true),
2923 )])),
2924 true,
2925 )
2926 .with_metadata(json_field_metadata(false, json_settings)),
2927 ]);
2928 assert!(validate_remote_schema(&expected, &actual, "test").is_ok());
2929 }
2930
2931 #[test]
2932 fn merge_scan_remote_schema_identity_accepts_json_decoded_struct_vs_wire_binary() {
2933 let json_settings = r#"{"json_settings":{"type_hints":[]}}"#;
2936 let expected =
2937 ArrowSchema::new(vec![
2938 Field::new(
2939 "j",
2940 TestArrowDataType::Struct(arrow_schema::Fields::from(vec![Arc::new(
2941 Field::new("a", TestArrowDataType::Utf8View, true),
2942 )])),
2943 true,
2944 )
2945 .with_metadata(json_field_metadata(false, json_settings)),
2946 ]);
2947 let actual = ArrowSchema::new(vec![
2948 Field::new("j", TestArrowDataType::Binary, true)
2949 .with_metadata(json_field_metadata(true, json_settings)),
2950 ]);
2951 assert!(validate_remote_schema(&expected, &actual, "test").is_ok());
2952 }
2953
2954 #[test]
2955 fn merge_scan_remote_schema_identity_rejects_json_fields_with_different_json_settings() {
2956 let expected = ArrowSchema::new(vec![
2960 Field::new("j", TestArrowDataType::Binary, true).with_metadata(json_field_metadata(
2961 true,
2962 r#"{"json_settings":{"type_hints":[["a",{"JsonType":"Int64"}]]}}"#,
2963 )),
2964 ]);
2965 let actual =
2966 ArrowSchema::new(vec![
2967 Field::new(
2968 "j",
2969 TestArrowDataType::Struct(arrow_schema::Fields::from(vec![Arc::new(
2970 Field::new("a", TestArrowDataType::Utf8View, true),
2971 )])),
2972 true,
2973 )
2974 .with_metadata(json_field_metadata(
2975 false,
2976 r#"{"json_settings":{"type_hints":[["a",{"JsonType":"String"}]]}}"#,
2977 )),
2978 ]);
2979 assert!(validate_remote_schema(&expected, &actual, "test").is_err());
2980 }
2981
2982 #[test]
2983 fn merge_scan_remote_schema_identity_ignores_json_extension_metadata_when_physical_type_matches()
2984 {
2985 let expected = ArrowSchema::new(vec![
2990 Field::new("j", TestArrowDataType::Binary, true).with_metadata(json_field_metadata(
2991 true,
2992 r#"{"json_settings":{"type_hints":[]}}"#,
2993 )),
2994 ]);
2995 let actual = ArrowSchema::new(vec![Field::new("j", TestArrowDataType::Binary, true)]);
2996 assert!(validate_remote_schema(&expected, &actual, "test").is_ok());
2997
2998 let plain_binary = ArrowSchema::new(vec![Field::new("v", TestArrowDataType::Binary, true)]);
3002 let plain_struct = ArrowSchema::new(vec![Field::new(
3003 "v",
3004 TestArrowDataType::Struct(arrow_schema::Fields::from(vec![Arc::new(Field::new(
3005 "a",
3006 TestArrowDataType::Utf8View,
3007 true,
3008 ))])),
3009 true,
3010 )]);
3011 assert!(validate_remote_schema(&plain_binary, &plain_struct, "test").is_err());
3012 }
3013
3014 #[test]
3015 fn merge_scan_remote_schema_identity_ignores_field_metadata_but_rejects_nullability_mismatch() {
3016 let expected = expected_int64_schema();
3017 let metadata_mismatch = ArrowSchema::new_with_metadata(
3018 vec![
3019 expected
3020 .field(0)
3021 .as_ref()
3022 .clone()
3023 .with_metadata(StdHashMap::from([(
3024 "remote".to_string(),
3025 "different".to_string(),
3026 )])),
3027 expected.field(1).as_ref().clone(),
3028 ],
3029 expected.metadata().clone(),
3030 );
3031 let nullability_mismatch = ArrowSchema::new_with_metadata(
3032 vec![
3033 expected.field(0).as_ref().clone().with_nullable(true),
3034 expected.field(1).as_ref().clone(),
3035 ],
3036 expected.metadata().clone(),
3037 );
3038 assert!(validate_remote_schema(&expected, &metadata_mismatch, "test").is_ok());
3041 assert!(validate_remote_schema(&expected, &nullability_mismatch, "test").is_err());
3042 }
3043
3044 #[test]
3045 fn merge_scan_remote_schema_identity_ignores_skipping_index_field_metadata() {
3046 let skipping_index_metadata = StdHashMap::from([(
3051 "greptime:skipping_index".to_string(),
3052 r#"{"granularity":1,"false-positive-rate-in-10000":100,"index-type":"BloomFilter"}"#
3053 .to_string(),
3054 )]);
3055 let expected =
3056 ArrowSchema::new(vec![Field::new("value", TestArrowDataType::Float64, true)]);
3057 let actual_with_metadata = ArrowSchema::new(vec![
3058 Field::new("value", TestArrowDataType::Float64, true)
3059 .with_metadata(skipping_index_metadata),
3060 ]);
3061 assert!(validate_remote_schema(&expected, &actual_with_metadata, "test").is_ok());
3062
3063 let wrong_type =
3065 ArrowSchema::new(vec![Field::new("value", TestArrowDataType::Int64, true)]);
3066 assert!(validate_remote_schema(&expected, &wrong_type, "test").is_err());
3067
3068 let wrong_name =
3070 ArrowSchema::new(vec![Field::new("other", TestArrowDataType::Float64, true)]);
3071 assert!(validate_remote_schema(&expected, &wrong_name, "test").is_err());
3072 }
3073
3074 #[test]
3075 fn merge_scan_remote_schema_identity_rejects_timestamp_timezone_plus_field_mismatches() {
3076 let expected = ArrowSchema::new(vec![Field::new(
3077 "ts",
3078 TestArrowDataType::Timestamp(TimeUnit::Millisecond, Some("Asia/Shanghai".into())),
3079 false,
3080 )]);
3081 let different_unit = ArrowSchema::new(vec![Field::new(
3082 "ts",
3083 TestArrowDataType::Timestamp(TimeUnit::Second, Some("UTC".into())),
3084 false,
3085 )]);
3086 let different_name = ArrowSchema::new(vec![Field::new(
3087 "other",
3088 TestArrowDataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
3089 false,
3090 )]);
3091 let nullability = ArrowSchema::new(vec![Field::new(
3092 "ts",
3093 TestArrowDataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
3094 true,
3095 )]);
3096 let metadata = ArrowSchema::new(vec![
3097 Field::new(
3098 "ts",
3099 TestArrowDataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
3100 false,
3101 )
3102 .with_metadata(StdHashMap::from([(
3103 "remote".to_string(),
3104 "different".to_string(),
3105 )])),
3106 ]);
3107 for actual in [&different_unit, &different_name, &nullability] {
3108 assert!(validate_remote_schema(&expected, actual, "test").is_err());
3109 }
3110 assert!(validate_remote_schema(&expected, &metadata, "test").is_ok());
3114 }
3115
3116 #[test]
3117 fn merge_scan_remote_schema_identity_returns_arrow_schema_error() {
3118 let expected = expected_int64_schema();
3119 let actual = ArrowSchema::new(vec![Field::new("other", TestArrowDataType::Int64, false)]);
3120 match validate_remote_schema(&expected, &actual, "test").unwrap_err() {
3121 DataFusionError::ArrowError(error, None) => match error.as_ref() {
3122 ArrowError::SchemaError(message) => {
3123 assert!(message.contains("field count mismatch"))
3124 }
3125 error => panic!("expected ArrowError::SchemaError, got {error:?}"),
3126 },
3127 error => panic!("expected DataFusionError::ArrowError(_, None), got {error:?}"),
3128 }
3129 }
3130
3131 #[tokio::test]
3132 async fn merge_scan_remote_schema_identity_rejects_incompatible_empty_advertised_schema() {
3133 let region_id = RegionId::new(1024, 1);
3134 let exec = merge_scan_exec_with_handler(
3135 vec![region_id],
3136 expected_int64_schema(),
3137 Arc::new(TestRegionQueryHandler::with_responses(vec![(
3138 region_id,
3139 int64_schema(&["a"]),
3140 vec![],
3141 )])),
3142 1,
3143 );
3144 let errors_before = merge_scan_schema_error_count_for_test();
3145 assert!(collect_merge_scan(exec).await.is_err());
3146 assert_eq!(merge_scan_schema_error_count_for_test(), errors_before + 1);
3147 }
3148
3149 #[tokio::test]
3150 async fn merge_scan_remote_schema_identity_allows_top_level_metadata_version_mismatch() {
3151 let fields = expected_int64_schema().fields().clone();
3152 let expected = ArrowSchema::new_with_metadata(
3153 fields.clone(),
3154 StdHashMap::from([("greptime:version".to_string(), "1".to_string())]),
3155 );
3156 let remote_schema = Arc::new(
3157 Schema::try_from(Arc::new(ArrowSchema::new_with_metadata(
3158 fields,
3159 StdHashMap::from([("greptime:version".to_string(), "0".to_string())]),
3160 )))
3161 .unwrap(),
3162 );
3163 let batch = record_batch(
3164 remote_schema.clone(),
3165 vec![
3166 Arc::new(Int64Vector::from_slice([11])) as _,
3167 Arc::new(Int64Vector::from_slice([12])) as _,
3168 ],
3169 );
3170 let batches = collect_merge_scan(merge_scan_exec_with_handler(
3171 vec![RegionId::new(1024, 1)],
3172 expected.clone(),
3173 Arc::new(TestRegionQueryHandler::with_responses(vec![(
3174 RegionId::new(1024, 1),
3175 remote_schema,
3176 vec![Ok(batch)],
3177 )])),
3178 1,
3179 ))
3180 .await
3181 .unwrap();
3182 assert_eq!(batches[0].schema().as_ref(), &expected);
3183 }
3184
3185 #[tokio::test]
3186 async fn merge_scan_remote_schema_identity_rejects_advertised_schema_inner_batch_mismatch() {
3187 let region_id = RegionId::new(1024, 1);
3188 let advertised_schema = int64_schema(&["a", "b"]);
3189 let inner_batch = record_batch(
3190 int64_schema(&["b", "a"]),
3191 vec![
3192 Arc::new(Int64Vector::from_slice([12])) as _,
3193 Arc::new(Int64Vector::from_slice([11])) as _,
3194 ],
3195 )
3196 .into_df_record_batch();
3197 let inner_schema = inner_batch.schema_ref().clone();
3198 let batch = RecordBatch::from_df_record_batch(advertised_schema.clone(), inner_batch);
3199 assert!(Arc::ptr_eq(
3200 advertised_schema.arrow_schema(),
3201 batch.schema.arrow_schema()
3202 ));
3203 assert!(!Arc::ptr_eq(
3204 advertised_schema.arrow_schema(),
3205 &inner_schema
3206 ));
3207 let exec = merge_scan_exec_with_handler(
3208 vec![region_id],
3209 expected_int64_schema(),
3210 Arc::new(TestRegionQueryHandler::with_responses(vec![(
3211 region_id,
3212 advertised_schema,
3213 vec![Ok(batch)],
3214 )])),
3215 1,
3216 );
3217 let errors_before = merge_scan_schema_error_count_for_test();
3218 assert!(collect_merge_scan(exec).await.is_err());
3219 assert_eq!(merge_scan_schema_error_count_for_test(), errors_before + 1);
3220 }
3221
3222 #[tokio::test]
3223 async fn merge_scan_remote_schema_identity_rejects_unchecked_inner_extra_column() {
3224 let region_id = RegionId::new(1024, 1);
3225 let advertised_schema = int64_schema(&["a", "b"]);
3226 let inner_batch = record_batch(
3227 int64_schema(&["a", "b", "extra"]),
3228 vec![
3229 Arc::new(Int64Vector::from_slice([11])) as _,
3230 Arc::new(Int64Vector::from_slice([12])) as _,
3231 Arc::new(Int64Vector::from_slice([13])) as _,
3232 ],
3233 )
3234 .into_df_record_batch();
3235 assert!(!Arc::ptr_eq(
3236 advertised_schema.arrow_schema(),
3237 inner_batch.schema_ref()
3238 ));
3239 let batch = RecordBatch::from_df_record_batch(advertised_schema.clone(), inner_batch);
3240 let exec = merge_scan_exec_with_handler(
3241 vec![region_id],
3242 expected_int64_schema(),
3243 Arc::new(TestRegionQueryHandler::with_responses(vec![(
3244 region_id,
3245 advertised_schema,
3246 vec![Ok(batch)],
3247 )])),
3248 1,
3249 );
3250 assert!(collect_merge_scan(exec).await.is_err());
3251 }
3252
3253 #[tokio::test]
3254 async fn merge_scan_remote_schema_identity_validates_structurally_equal_distinct_batch_schema()
3255 {
3256 let region_id = RegionId::new(1024, 1);
3257 let advertised_schema = int64_schema(&["a", "b"]);
3258 let inner_schema = Arc::new(
3259 Schema::try_from(Arc::new(advertised_schema.arrow_schema().as_ref().clone())).unwrap(),
3260 );
3261 let inner_batch = record_batch(
3262 inner_schema,
3263 vec![
3264 Arc::new(Int64Vector::from_slice([11])) as _,
3265 Arc::new(Int64Vector::from_slice([12])) as _,
3266 ],
3267 )
3268 .into_df_record_batch();
3269 assert_eq!(advertised_schema.arrow_schema(), inner_batch.schema_ref());
3270 assert!(!Arc::ptr_eq(
3271 advertised_schema.arrow_schema(),
3272 inner_batch.schema_ref()
3273 ));
3274 let batch = RecordBatch::from_df_record_batch(advertised_schema.clone(), inner_batch);
3275 let batches = collect_merge_scan(merge_scan_exec_with_handler(
3276 vec![region_id],
3277 expected_int64_schema(),
3278 Arc::new(TestRegionQueryHandler::with_responses(vec![(
3279 region_id,
3280 advertised_schema,
3281 vec![Ok(batch)],
3282 )])),
3283 1,
3284 ))
3285 .await
3286 .unwrap();
3287 assert_eq!(batches.len(), 1);
3288 assert_int64_batch(&batches[0], (11, 12));
3289 }
3290
3291 #[test]
3292 fn try_with_new_distribution_preserves_remote_dyn_filter_producer_id() {
3293 let remote_dyn_filter_producer_id = RemoteDynFilterProducerId::new(42);
3294
3295 let plan = LogicalPlanBuilder::empty(true)
3297 .project(vec![lit(1i32).alias("col1")])
3298 .unwrap()
3299 .build()
3300 .unwrap();
3301
3302 let schema = plan.schema().as_arrow().clone();
3303 let table = TableName::new("catalog", "schema", "table");
3304 let regions = vec![RegionId::new(1024, 1)];
3305 let query_ctx = QueryContext::arc();
3306
3307 let mut partition_cols = AliasMapping::new();
3309 partition_cols.insert(
3310 "col1".to_string(),
3311 BTreeSet::from([ColumnExpr::new(Some(TableReference::bare("table")), "col1")]),
3312 );
3313
3314 let session_state = SessionStateBuilder::new().build();
3315
3316 let handler = Arc::new(TestRegionQueryHandler::default());
3317 let target_partition = 2;
3318
3319 let exec = MergeScanExec::new(
3320 &session_state,
3321 table,
3322 regions,
3323 plan,
3324 &schema,
3325 handler,
3326 query_ctx,
3327 target_partition,
3328 partition_cols,
3329 Some(remote_dyn_filter_producer_id),
3330 false,
3331 )
3332 .unwrap();
3333
3334 assert_eq!(
3335 exec.remote_dyn_filter_producer_id(),
3336 Some(remote_dyn_filter_producer_id)
3337 );
3338
3339 let new_dist = Distribution::KeyPartitioned(vec![
3343 Arc::new(Column::new("col1", 0)),
3344 Arc::new(Column::new("col2", 1)),
3345 ]);
3346
3347 let cloned = exec
3348 .try_with_new_distribution(new_dist)
3349 .expect("expected a cloned exec with overlapping partition col");
3350
3351 assert_eq!(
3352 cloned.remote_dyn_filter_producer_id(),
3353 Some(remote_dyn_filter_producer_id),
3354 "try_with_new_distribution must preserve remote dynamic filter producer id"
3355 );
3356 }
3357
3358 #[test]
3359 fn merge_scan_apply_expressions_exposes_remote_dyn_filter_id() {
3360 let query_ctx = QueryContext::arc();
3361 let exec =
3362 remote_dyn_filter_test_exec(Arc::new(TestRegionQueryHandler::default()), query_ctx);
3363 let dyn_filter = install_remote_dyn_filter(&exec);
3364 let expected_expression_id = dyn_filter.expression_id();
3365 let mut expression_ids = Vec::new();
3366
3367 exec.apply_expressions(&mut |expr| {
3368 expression_ids.push(expr.expression_id());
3369 Ok(TreeNodeRecursion::Continue)
3370 })
3371 .unwrap();
3372
3373 assert_eq!(expression_ids, vec![expected_expression_id]);
3374 }
3375
3376 #[test]
3377 fn remote_dyn_filter_preflight_removes_parent_filter_after_dn_runtime_is_ready() {
3378 let remote_dyn_filter_producer_id = RemoteDynFilterProducerId::new(42);
3379 let plan = LogicalPlanBuilder::empty(true)
3380 .project(vec![lit(1i32).alias("col1")])
3381 .unwrap()
3382 .build()
3383 .unwrap();
3384
3385 let schema = plan.schema().as_arrow().clone();
3386 let table = TableName::new("catalog", "schema", "table");
3387 let regions = vec![RegionId::new(1024, 1)];
3388 let query_ctx = QueryContext::arc();
3389 let session_state = SessionStateBuilder::new().build();
3390 let handler = Arc::new(TestRegionQueryHandler::default());
3391 let exec = MergeScanExec::new(
3392 &session_state,
3393 table,
3394 regions,
3395 plan,
3396 &schema,
3397 handler,
3398 query_ctx,
3399 1,
3400 AliasMapping::new(),
3401 Some(remote_dyn_filter_producer_id),
3402 false,
3403 )
3404 .unwrap();
3405 let dyn_filter = Arc::new(DynamicFilterPhysicalExpr::new(
3406 vec![Arc::new(Column::new("host", 0)) as Arc<_>],
3407 physical_lit(true) as _,
3408 )) as Arc<dyn datafusion_physical_expr::PhysicalExpr>;
3409
3410 let propagation = exec
3411 .handle_child_pushdown_result(
3412 FilterPushdownPhase::Post,
3413 ChildPushdownResult {
3414 parent_filters: vec![ChildFilterPushdownResult {
3415 filter: dyn_filter,
3416 child_results: vec![PushedDown::Yes],
3417 }],
3418 self_filters: Vec::new(),
3419 },
3420 &ConfigOptions::new(),
3421 )
3422 .unwrap();
3423
3424 assert_eq!(exec.captured_remote_dyn_filters().len(), 1);
3425 assert!(matches!(propagation.filters.as_slice(), [PushedDown::Yes]));
3426 }
3427
3428 #[test]
3429 fn scan_output_bytes_uses_plan_name() {
3430 let metrics = RecordBatchMetrics {
3431 plan_metrics: vec![PlanMetrics {
3432 plan: "SeqScan: region=1".to_string(),
3433 plan_name: REGION_SCAN_EXEC_NAME.to_string(),
3434 level: 0,
3435 metrics: vec![("output_bytes".to_string(), 42)],
3436 }],
3437 ..Default::default()
3438 };
3439
3440 assert_eq!(region_scan_output_bytes(&metrics), 42);
3441 }
3442
3443 #[test]
3444 fn scan_output_bytes_defaults_to_zero_without_region_scan() {
3445 let metrics = RecordBatchMetrics {
3446 plan_metrics: vec![PlanMetrics {
3447 plan: "ProjectionExec".to_string(),
3448 plan_name: "ProjectionExec".to_string(),
3449 level: 0,
3450 metrics: vec![("output_bytes".to_string(), 42)],
3451 }],
3452 ..Default::default()
3453 };
3454
3455 assert_eq!(region_scan_output_bytes(&metrics), 0);
3456 }
3457
3458 #[test]
3459 fn scan_output_bytes_sums_multiple_region_scans() {
3460 let metrics = RecordBatchMetrics {
3461 plan_metrics: vec![
3462 PlanMetrics {
3463 plan: "RegionScanExec: region=1".to_string(),
3464 plan_name: REGION_SCAN_EXEC_NAME.to_string(),
3465 level: 0,
3466 metrics: vec![("output_bytes".to_string(), 42)],
3467 },
3468 PlanMetrics {
3469 plan: "RegionScanExec: region=2".to_string(),
3470 plan_name: REGION_SCAN_EXEC_NAME.to_string(),
3471 level: 0,
3472 metrics: vec![("output_bytes".to_string(), 18)],
3473 },
3474 ],
3475 ..Default::default()
3476 };
3477
3478 assert_eq!(region_scan_output_bytes(&metrics), 60);
3479 }
3480
3481 #[test]
3482 fn merge_scan_reports_region_query_load_on_drop() {
3483 use store_api::metrics::{REGION_QUERY_CPU_TIME, REGION_QUERY_SCANNED_BYTES};
3484
3485 let region_id = RegionId::new(1024, 10002);
3486 let region_id_label = region_id.to_string();
3487 let labels = [®ion_id_label];
3488 let _ = REGION_QUERY_CPU_TIME.remove_label_values(&labels);
3489 let _ = REGION_QUERY_SCANNED_BYTES.remove_label_values(&labels);
3490
3491 let plan = LogicalPlanBuilder::empty(true)
3492 .project(vec![lit(1i32).alias("col1")])
3493 .unwrap()
3494 .build()
3495 .unwrap();
3496 let schema = plan.schema().as_arrow().clone();
3497 let exec = MergeScanExec::new(
3498 &SessionStateBuilder::new().build(),
3499 TableName::new("catalog", "schema", "table"),
3500 vec![region_id],
3501 plan,
3502 &schema,
3503 Arc::new(TestRegionQueryHandler::default()),
3504 QueryContext::arc(),
3505 1,
3506 AliasMapping::new(),
3507 None,
3508 true,
3509 )
3510 .unwrap();
3511
3512 let metrics = RecordBatchMetrics {
3513 elapsed_compute: 42,
3514 plan_metrics: vec![PlanMetrics {
3515 plan: "RegionScanExec: region=1".to_string(),
3516 plan_name: REGION_SCAN_EXEC_NAME.to_string(),
3517 level: 0,
3518 metrics: vec![("output_bytes".to_string(), 24)],
3519 }],
3520 ..Default::default()
3521 };
3522 exec.sub_stage_metrics
3523 .lock()
3524 .unwrap()
3525 .insert(region_id, metrics);
3526
3527 assert_eq!(REGION_QUERY_CPU_TIME.with_label_values(&labels).get(), 0);
3528 assert_eq!(
3529 REGION_QUERY_SCANNED_BYTES.with_label_values(&labels).get(),
3530 0
3531 );
3532
3533 drop(exec);
3534
3535 assert_eq!(REGION_QUERY_CPU_TIME.with_label_values(&labels).get(), 42);
3536 assert_eq!(
3537 REGION_QUERY_SCANNED_BYTES.with_label_values(&labels).get(),
3538 24
3539 );
3540
3541 let _ = REGION_QUERY_CPU_TIME.remove_label_values(&labels);
3542 let _ = REGION_QUERY_SCANNED_BYTES.remove_label_values(&labels);
3543 }
3544
3545 #[test]
3546 fn merge_scan_reports_query_load_with_metrics_region_id() {
3547 use store_api::metrics::{REGION_QUERY_CPU_TIME, REGION_QUERY_SCANNED_BYTES};
3548
3549 let logical_region_id = RegionId::new(1024, 10002);
3550 let physical_region_id = RegionId::new(1024, 1);
3551 let logical_region_id_label = logical_region_id.to_string();
3552 let physical_region_id_label = physical_region_id.to_string();
3553 let logical_labels = [&logical_region_id_label];
3554 let physical_labels = [&physical_region_id_label];
3555 let _ = REGION_QUERY_CPU_TIME.remove_label_values(&logical_labels);
3556 let _ = REGION_QUERY_SCANNED_BYTES.remove_label_values(&logical_labels);
3557 let _ = REGION_QUERY_CPU_TIME.remove_label_values(&physical_labels);
3558 let _ = REGION_QUERY_SCANNED_BYTES.remove_label_values(&physical_labels);
3559
3560 let plan = LogicalPlanBuilder::empty(true)
3561 .project(vec![lit(1i32).alias("col1")])
3562 .unwrap()
3563 .build()
3564 .unwrap();
3565 let schema = plan.schema().as_arrow().clone();
3566 let exec = MergeScanExec::new(
3567 &SessionStateBuilder::new().build(),
3568 TableName::new("catalog", "schema", "table"),
3569 vec![logical_region_id],
3570 plan,
3571 &schema,
3572 Arc::new(TestRegionQueryHandler::default()),
3573 QueryContext::arc(),
3574 1,
3575 AliasMapping::new(),
3576 None,
3577 true,
3578 )
3579 .unwrap();
3580
3581 let metrics = RecordBatchMetrics {
3582 elapsed_compute: 42,
3583 query_load_region_id: Some(physical_region_id.as_u64()),
3584 plan_metrics: vec![PlanMetrics {
3585 plan: "RegionScanExec: region=1".to_string(),
3586 plan_name: REGION_SCAN_EXEC_NAME.to_string(),
3587 level: 0,
3588 metrics: vec![("output_bytes".to_string(), 24)],
3589 }],
3590 ..Default::default()
3591 };
3592 exec.sub_stage_metrics
3593 .lock()
3594 .unwrap()
3595 .insert(logical_region_id, metrics);
3596
3597 drop(exec);
3598
3599 assert_eq!(
3600 REGION_QUERY_CPU_TIME
3601 .with_label_values(&logical_labels)
3602 .get(),
3603 0
3604 );
3605 assert_eq!(
3606 REGION_QUERY_SCANNED_BYTES
3607 .with_label_values(&logical_labels)
3608 .get(),
3609 0
3610 );
3611 assert_eq!(
3612 REGION_QUERY_CPU_TIME
3613 .with_label_values(&physical_labels)
3614 .get(),
3615 42
3616 );
3617 assert_eq!(
3618 REGION_QUERY_SCANNED_BYTES
3619 .with_label_values(&physical_labels)
3620 .get(),
3621 24
3622 );
3623
3624 let _ = REGION_QUERY_CPU_TIME.remove_label_values(&logical_labels);
3625 let _ = REGION_QUERY_SCANNED_BYTES.remove_label_values(&logical_labels);
3626 let _ = REGION_QUERY_CPU_TIME.remove_label_values(&physical_labels);
3627 let _ = REGION_QUERY_SCANNED_BYTES.remove_label_values(&physical_labels);
3628 }
3629}