Skip to main content

query/dist_plan/
merge_scan.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#[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
86/// Returns a deterministic upper bound on rows emitted by the remote plan for
87/// one region.
88///
89/// Only explicit caps and row-non-increasing wrappers are followed. Nodes that
90/// may multiply rows fail open, so the result remains a sound upper bound.
91fn 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    // Prometheus counters are process-global and tests run concurrently. This
150    // companion counter is incremented at the exact production increment site
151    // and is scoped to the polling test thread, making delta assertions stable.
152    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
160/// Returns true when the field is a JSON column, identified either by its
161/// Arrow extension type (`greptime.json` / `greptime.json2`) or by its
162/// `greptime:type=Json` marker.
163fn 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
172/// Returns true when two fields describe the same column semantics: same
173/// name, same data type, and same nullability.
174///
175/// Arrow `Field` metadata is auxiliary information (indexing, encoding,
176/// compression hints, ...) and does not participate in the comparison. This
177/// applies to every field, JSON columns included: JSON identity keys
178/// (`greptime:type`, `ARROW:extension:metadata`) are never compared here.
179/// The wire/decoded physical-type difference of JSON columns is exempted
180/// separately by [`json_fields_compatible`].
181fn 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
187/// Returns true when two fields are semantically equivalent JSON columns.
188///
189/// A JSON column is carried on the wire in its binary-encoded form (e.g.
190/// `Binary` + `ARROW:extension:name=greptime.json` / `greptime:type=Json`) and,
191/// after decoding on the remote side, in a concretized structured form (e.g.
192/// `Struct(...)` / `List(...)` carrying the same extension metadata). Both forms
193/// describe the same column, so the raw arrow data type must not be compared
194/// directly. The JSON extension identity (`greptime:type`) and the JSON2
195/// settings (`ARROW:extension:metadata`, e.g. type hints) are the semantic
196/// parts of the field and must match.
197fn 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        // Both must carry the same JSON marker.
203        && expected_field.metadata().get(datatypes::schema::TYPE_KEY)
204            == actual_field.metadata().get(datatypes::schema::TYPE_KEY)
205        // JSON2 settings (type hints etc.) must match.
206        && 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
214/// Validates the remote schema before positional column handling.
215///
216/// Field metadata (indexing/encoding/compression hints) is non-semantic at
217/// this boundary and is ignored: two fields are equal when their name, data
218/// type, and nullability match (see [`fields_semantically_equal`]). A
219/// timestamp timezone difference is the only intentional exception, accepted
220/// when the fields share the same timestamp unit, distinct timezones, and
221/// equal name and nullability. Top-level Arrow schema metadata is
222/// non-semantic at this boundary as well. JSON columns are compared
223/// semantically (see [`json_fields_compatible`]) because their wire and
224/// decoded representations use different physical arrow types.
225fn 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        // Field metadata does not participate in the semantic comparison.
245        if fields_semantically_equal(expected_field, actual_field) {
246            continue;
247        }
248
249        // JSON columns are equivalent in their binary wire form and their
250        // decoded structured form; compare them semantically instead of
251        // comparing the raw arrow data type.
252        if json_fields_compatible(expected_field, actual_field) {
253            continue;
254        }
255
256        // Intentionally mirrors Arrow Field equality properties, except timezone.
257        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    /// In logical plan phase it only contains one input
337    input: LogicalPlan,
338    /// Schema exposed to the local stage.
339    output_schema: DFSchemaRef,
340    /// If this plan is a placeholder
341    is_placeholder: bool,
342    partition_cols: AliasMapping,
343    /// Assigned after dist-plan rewriting so rewriters only deal with plan shape.
344    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    // Prevent further optimization.
389    // The input can be retrieved by `self.input()`
390    fn inputs(&self) -> Vec<&LogicalPlan> {
391        vec![]
392    }
393
394    fn schema(&self) -> &datafusion_common::DFSchemaRef {
395        &self.output_schema
396    }
397
398    // Prevent further optimization
399    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    /// Creates a merge scan with the input plan's schema.
422    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    /// Replaces the schema exposed to the local stage.
433    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    /// Create a [LogicalPlan::Extension] node from this merge scan plan
451    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    /// Metrics from sub stages
484    sub_stage_metrics: Arc<Mutex<HashMap<RegionId, RecordBatchMetrics>>>,
485    /// Metrics for each partition
486    partition_metrics: Arc<Mutex<HashMap<usize, PartitionMetrics>>>,
487    query_ctx: QueryContextRef,
488    /// Optional because RDF must fail open: missing ids skip RDF but keep normal query execution.
489    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        // States the output ordering of the plan.
525        //
526        // When the input plan is a sort, we can use the sort ordering as the output ordering
527        // if the target partition is greater than the number of regions, which means we won't
528        // break the ordering on merging (of MergeScan).
529        //
530        // Otherwise, we need to use the default ordering.
531        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    /// Conservative row-count upper bound for all selected regions.
599    ///
600    /// This is not an expected cardinality: DataFusion receives it as an
601    /// inexact estimate only because `Statistics` has no upper-bound precision.
602    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    /// Number of partitions populated by the region striping in [`Self::to_stream`].
616    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        // prepare states to move
626        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            // only report metrics once for each MergeScan
657            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            // Per-partition timings, scoped to this partition's stream for `EXPLAIN VERBOSE`.
666            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((&region_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                        // The remote Flight stream publishes its terminal metrics
802                        // immediately after EOF. Capture them before leaving the
803                        // loop so the final verbose snapshot is not lost.
804                        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                    // reset poll timer
846                    poll_timer = Instant::now();
847                }
848                // Also stop on an exhausted stream that yielded no batch. The `take()`
849                // guard ensures it only records once, on the first such region.
850                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                // Record region metrics and push to global partition_metrics
857                let region_metrics = RegionMetrics {
858                    region_id,
859                    poll_duration,
860                    do_get_cost,
861                    total_cost,
862                };
863
864                // Push RegionMetrics to global partition_metrics immediately after scanning this region
865                {
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            // Stop the global timers for partitions with no region, otherwise they keep
888            // running until drop and inflate the shared metrics. No-op otherwise.
889            ready_timer.stop();
890            if let Some(mut first_consume_timer) = first_consume_timer.take() {
891                first_consume_timer.stop();
892            }
893
894            // Finish partition metrics and log results
895            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            // not applicable
918            return None;
919        };
920
921        if let Partitioning::Hash(curr_dist, _) = &self.properties.partitioning
922            && curr_dist == &hash_exprs
923        {
924            // No need to change the distribution
925            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/// Metrics for a region of a partition.
1048#[derive(Debug, Clone)]
1049struct RegionMetrics {
1050    region_id: RegionId,
1051    poll_duration: Duration,
1052    do_get_cost: Duration,
1053    /// Total cost to scan the region.
1054    total_cost: Duration,
1055}
1056
1057/// Metrics for a partition of a MergeScanExec.
1058#[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    /// Time until this partition's scan is ready to emit data.
1066    ready_time: Duration,
1067    /// Time until this partition's first stream poll resolves (a batch or exhausted).
1068    first_consume_time: Duration,
1069    /// Time until this partition's scan finishes execution.
1070    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    /// Set the per-partition timings captured during streaming.
1099    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    /// Finish the partition metrics and log the results.
1111    fn finish(&mut self) {
1112        if self.finished {
1113            return;
1114        }
1115        self.finished = true;
1116        self.log_metrics();
1117    }
1118
1119    /// Log partition metrics based on explain_verbose level.
1120    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    // DataFusion will swap children unconditionally.
1187    // But since this node is leaf node, it's safe to just return self.
1188    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            // Reject remote pushdown instead of pretending success: this keeps
1209            // DataFusion/local dynamic filter semantics intact while disabling
1210            // only FE -> DN remote dynamic filter propagation.
1211            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            // Missing RDF identity disables only RDF, not normal execution.
1222            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(&[&region_id])
1344        .inc_by(load.cpu_time);
1345    REGION_QUERY_SCANNED_BYTES
1346        .with_label_values(&[&region_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        // Per-region Prometheus metrics can have high cardinality, so they are
1360        // controlled by `enable_per_region_metrics`. Region-owned counters for
1361        // heartbeat reporting are updated on datanodes when query metrics resolve.
1362        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    /// Nanosecond elapsed till the scan operator is ready to emit data
1377    ready_time: Time,
1378    /// Nanosecond elapsed till the first record batch emitted from the scan operator gets consumed
1379    first_consume_time: Time,
1380    /// Nanosecond elapsed till the scan operator finished execution
1381    finish_time: Time,
1382    /// Count of rows fetched from remote
1383    output_rows: Count,
1384
1385    /// Gauge for greptime plan execution cost metrics for output
1386    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            // The table name is not relevant to these ordering metadata tests;
1512            // `MergeScanExec::new` requires one to model the production plan.
1513            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    /// Builds the metadata of a JSON column field. `wire_form` mirrors the
2885    /// binary-encoded representation (adds `ARROW:extension:name`), while the
2886    /// decoded structured form only carries the semantic keys.
2887    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        // The merge-scan (expected) side carries a JSON2 column in its binary
2907        // wire form (Binary + `ARROW:extension:name`); the remote decoded
2908        // stream carries the concretized structured form (Struct + the same
2909        // semantic extension metadata, without the arrow extension name). This
2910        // mirrors the failing `json2_limit` case: the two representations must
2911        // validate as equal.
2912        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        // The reverse direction: expected side is the decoded structured form
2934        // while the actual remote stream advertises the binary wire form.
2935        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        // JSON2 settings (type hints) are semantic: two JSON columns with
2957        // different settings describe different logical structures and must be
2958        // rejected.
2959        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        // Field metadata (including the JSON extension identity) is not part
2986        // of the semantic field identity: a JSON column in its binary wire
2987        // form and a plain Binary column share the same name, physical type,
2988        // and nullability, so they validate as equal.
2989        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        // The JSON compatibility exemption is scoped to JSON columns only: a
2999        // pair of non-JSON fields with different physical types is still
3000        // rejected even though one side is Binary.
3001        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        // Field metadata is auxiliary and does not participate in the
3039        // semantic field comparison.
3040        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        // Regression: a remote stream may advertise auxiliary field metadata
3047        // (e.g. `greptime:skipping_index`) that the expected schema does not
3048        // carry. Field metadata is not part of the semantic field identity
3049        // (name + data type + nullability) and must not fail validation.
3050        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        // A real data type mismatch is still rejected.
3064        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        // A name mismatch is still rejected.
3069        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        // Field metadata is non-semantic and is ignored even for timezone
3111        // pairs: a metadata difference on an otherwise timezone-only
3112        // difference is accepted.
3113        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        // Build a plan whose schema contains "col1"
3296        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        // Non-empty partition_cols so try_with_new_distribution can detect an overlap
3308        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        // A distribution that differs from the current partitioning but shares a
3340        // column name present in partition_cols, so try_with_new_distribution
3341        // produces a clone instead of returning None.
3342        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 = [&region_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}