query/dist_plan/
planner.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//! [ExtensionPlanner] implementation for distributed planner
16
17use std::sync::Arc;
18
19use ahash::HashMap;
20use arrow_schema::SortOptions;
21use async_trait::async_trait;
22use catalog::CatalogManagerRef;
23use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
24use common_telemetry::debug;
25use datafusion::catalog::Session;
26use datafusion::common::Result;
27use datafusion::datasource::DefaultTableSource;
28use datafusion::execution::context::SessionState;
29use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
30use datafusion::physical_plan::ExecutionPlan;
31use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner};
32use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
33use datafusion_common::{DataFusionError, TableReference};
34use datafusion_expr::{LogicalPlan, UserDefinedLogicalNode};
35use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr};
36use datatypes::prelude::ConcreteDataType;
37use partition::expr::PartitionExpr;
38use partition::manager::{PartitionRuleManagerRef, create_partitions_from_region_routes};
39use session::context::QueryContext;
40use snafu::{OptionExt, ResultExt};
41use store_api::storage::RegionId;
42use table::TableRef;
43use table::metadata::TableInfo;
44pub use table::metadata::TableType;
45use table::table::adapter::DfTableProviderAdapter;
46use table::table_name::TableName;
47
48use crate::dist_plan::PredicateExtractor;
49use crate::dist_plan::merge_scan::{MergeScanExec, MergeScanLogicalPlan};
50use crate::dist_plan::merge_sort::{MergeSortExec, MergeSortLogicalPlan};
51use crate::dist_plan::region_pruner::ConstraintPruner;
52use crate::error::{CatalogSnafu, PartitionRuleManagerSnafu, TableNotFoundSnafu};
53use crate::optimizer::expr_depth;
54use crate::region_query::RegionQueryHandlerRef;
55
56/// Planner for converting merge sort logical plan to physical plan.
57///
58/// `MergeSortExec` always represents the distributed merge stage. It declares
59/// the required input ordering to DataFusion, so `EnforceSorting` inserts a
60/// `SortExec` below it when the input `MergeScanExec` cannot preserve per-region
61/// ordering, for example when one output partition may merge multiple region
62/// streams.
63pub struct MergeSortExtensionPlanner {}
64
65impl MergeSortExtensionPlanner {
66    fn ordering(
67        planner: &dyn PhysicalPlanner,
68        session: &dyn Session,
69        planning_ctx: &PhysicalPlanningContext,
70        merge_sort: &MergeSortLogicalPlan,
71    ) -> Result<LexOrdering> {
72        let ordering = merge_sort
73            .expr
74            .iter()
75            .map(|sort_expr| {
76                let physical_expr = planner.create_physical_expr(
77                    &sort_expr.expr,
78                    merge_sort.input.schema(),
79                    session,
80                    planning_ctx,
81                )?;
82                Ok(PhysicalSortExpr::new(
83                    physical_expr,
84                    SortOptions {
85                        descending: !sort_expr.asc,
86                        nulls_first: sort_expr.nulls_first,
87                    },
88                ))
89            })
90            .collect::<Result<Vec<_>>>()?;
91
92        LexOrdering::new(ordering).ok_or_else(|| {
93            DataFusionError::Internal(
94                "Expect MergeSort to have non-empty sort expressions".to_string(),
95            )
96        })
97    }
98}
99
100#[async_trait]
101impl ExtensionPlanner for MergeSortExtensionPlanner {
102    async fn plan_extension(
103        &self,
104        planner: &dyn PhysicalPlanner,
105        node: &dyn UserDefinedLogicalNode,
106        _logical_inputs: &[&LogicalPlan],
107        physical_inputs: &[Arc<dyn ExecutionPlan>],
108        session: &dyn Session,
109        planning_ctx: &PhysicalPlanningContext,
110    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
111        if let Some(merge_sort) = node.as_any().downcast_ref::<MergeSortLogicalPlan>() {
112            if let LogicalPlan::Extension(ext) = &merge_sort.input.as_ref()
113                && ext
114                    .node
115                    .as_any()
116                    .downcast_ref::<MergeScanLogicalPlan>()
117                    .is_some()
118            {
119                let input = physical_inputs.first().cloned().ok_or_else(|| {
120                    DataFusionError::Internal(
121                        "Expect MergeSort to have one physical input".to_string(),
122                    )
123                })?;
124                if input.downcast_ref::<MergeScanExec>().is_none() {
125                    return Err(DataFusionError::Internal(format!(
126                        "Expect MergeSort's input is a MergeScanExec, found {:?}",
127                        physical_inputs
128                    )));
129                }
130
131                let ordering = Self::ordering(planner, session, planning_ctx, merge_sort)?;
132                Ok(Some(Arc::new(MergeSortExec::new(
133                    ordering,
134                    input,
135                    merge_sort.fetch,
136                ))))
137            } else {
138                Ok(None)
139            }
140        } else {
141            Ok(None)
142        }
143    }
144}
145
146pub struct DistExtensionPlanner {
147    catalog_manager: CatalogManagerRef,
148    partition_rule_manager: PartitionRuleManagerRef,
149    region_query_handler: RegionQueryHandlerRef,
150    enable_per_region_metrics: bool,
151}
152
153impl DistExtensionPlanner {
154    pub fn new(
155        catalog_manager: CatalogManagerRef,
156        partition_rule_manager: PartitionRuleManagerRef,
157        region_query_handler: RegionQueryHandlerRef,
158        enable_per_region_metrics: bool,
159    ) -> Self {
160        Self {
161            catalog_manager,
162            partition_rule_manager,
163            region_query_handler,
164            enable_per_region_metrics,
165        }
166    }
167}
168
169#[async_trait]
170impl ExtensionPlanner for DistExtensionPlanner {
171    async fn plan_extension(
172        &self,
173        planner: &dyn PhysicalPlanner,
174        node: &dyn UserDefinedLogicalNode,
175        _logical_inputs: &[&LogicalPlan],
176        _physical_inputs: &[Arc<dyn ExecutionPlan>],
177        session: &dyn Session,
178        _planning_ctx: &PhysicalPlanningContext,
179    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
180        let Some(merge_scan) = node.as_any().downcast_ref::<MergeScanLogicalPlan>() else {
181            return Ok(None);
182        };
183
184        let input_plan = merge_scan.input();
185        let fallback = |logical_plan| async move {
186            let optimized_plan = self.optimize_input_logical_plan(session, logical_plan)?;
187            planner
188                .create_physical_plan(&optimized_plan, session)
189                .await
190                .map(Some)
191        };
192
193        if merge_scan.is_placeholder() {
194            // ignore placeholder
195            return fallback(input_plan).await;
196        }
197
198        let optimized_plan = input_plan;
199        let Some(table_name) = Self::extract_full_table_name(input_plan)? else {
200            // no relation found in input plan, going to execute them locally
201            return fallback(optimized_plan).await;
202        };
203
204        let Ok(regions) = self.get_regions(&table_name, input_plan).await else {
205            // no peers found, going to execute them locally
206            return fallback(optimized_plan).await;
207        };
208
209        // TODO(ruihang): generate different execution plans for different variant merge operation
210        let schema = merge_scan.schema().as_arrow();
211        let session_state = session
212            .as_any()
213            .downcast_ref::<SessionState>()
214            .ok_or_else(|| {
215                DataFusionError::Internal(
216                    "MergeScan requires a SessionState for physical planning".to_string(),
217                )
218            })?;
219        let query_ctx = session_state
220            .config()
221            .get_extension()
222            .unwrap_or_else(QueryContext::arc);
223        let merge_scan_plan = MergeScanExec::new(
224            session_state,
225            table_name,
226            regions,
227            input_plan.clone(),
228            schema,
229            self.region_query_handler.clone(),
230            query_ctx,
231            session.config().target_partitions(),
232            merge_scan.remote_dyn_filter_producer_id(),
233            self.enable_per_region_metrics,
234        )?;
235        Ok(Some(Arc::new(merge_scan_plan) as _))
236    }
237}
238
239impl DistExtensionPlanner {
240    /// Extract fully resolved table name from logical plan
241    fn extract_full_table_name(plan: &LogicalPlan) -> Result<Option<TableName>> {
242        let mut extractor = TableScanExtractor::default();
243        let _ = plan.visit(&mut extractor)?;
244        Ok(extractor.table_name)
245    }
246
247    async fn get_regions(
248        &self,
249        table_name: &TableName,
250        logical_plan: &LogicalPlan,
251    ) -> Result<Vec<RegionId>> {
252        let mut extractor = TableScanExtractor::default();
253        let _ = logical_plan.visit(&mut extractor)?;
254        // Resolving by name again could bind an authorized scan to a replacement table.
255        let table = match extractor.captured_table {
256            Some(table) => table,
257            None => self
258                .catalog_manager
259                .table(
260                    &table_name.catalog_name,
261                    &table_name.schema_name,
262                    &table_name.table_name,
263                    None,
264                )
265                .await
266                .context(CatalogSnafu)?
267                .with_context(|| TableNotFoundSnafu {
268                    table: table_name.to_string(),
269                })?,
270        };
271
272        let table_info = table.table_info();
273        let (physical_table_id, physical_table_route) = self
274            .partition_rule_manager
275            .find_physical_table_route_with_id(table_info.table_id())
276            .await
277            .context(PartitionRuleManagerSnafu)?;
278        let all_regions = physical_table_route
279            .region_routes
280            .iter()
281            .map(|r| RegionId::new(table_info.table_id(), r.region.id.region_number()))
282            .collect::<Vec<_>>();
283        let logical_partition_columns = partition_column_types(&table_info);
284        let partition_columns = logical_partition_columns
285            .iter()
286            .map(|(name, _)| name.clone())
287            .collect::<Vec<_>>();
288        debug!(
289            "DistExtensionPlanner: loaded table partition metadata, table: {}, table_id: {}, partition_key_indices: {:?}, partition_columns: {:?}, all_regions: {:?}",
290            table_name,
291            table_info.table_id(),
292            table_info.meta.partition_key_indices,
293            partition_columns,
294            all_regions,
295        );
296        if partition_columns.is_empty() {
297            return Ok(all_regions);
298        }
299        // Extract predicates from logical plan
300        let partition_expressions = match PredicateExtractor::extract_partition_expressions(
301            logical_plan,
302            &partition_columns,
303        ) {
304            Ok(expressions) => expressions,
305            Err(err) => {
306                common_telemetry::debug!(
307                    "Failed to extract partition expressions for table {} (id: {}), using all regions: {:?}",
308                    table_name,
309                    table.table_info().table_id(),
310                    err
311                );
312                return Ok(all_regions);
313            }
314        };
315
316        if partition_expressions.is_empty() {
317            return Ok(all_regions);
318        }
319
320        let Some(partition_column_types) = self
321            .partition_column_types_for_pruning(
322                table_name,
323                table_info.as_ref(),
324                physical_table_id,
325                &partition_expressions,
326                &all_regions,
327            )
328            .await
329        else {
330            return Ok(all_regions);
331        };
332
333        // Get partition information for the table if partition rule manager is available
334        let partitions = match create_partitions_from_region_routes(
335            table_info.table_id(),
336            &physical_table_route.region_routes,
337        ) {
338            Ok(partitions) => partitions,
339            Err(err) => {
340                common_telemetry::debug!(
341                    "Failed to get partition information for table {}, using all regions: {:?}",
342                    table_name,
343                    err
344                );
345                return Ok(all_regions);
346            }
347        };
348        if partitions.is_empty() {
349            return Ok(all_regions);
350        }
351        // Apply region pruning based on partition rules
352        let pruned_regions = match ConstraintPruner::prune_regions(
353            &partition_expressions,
354            &partitions,
355            partition_column_types,
356        ) {
357            Ok(regions) => regions,
358            Err(err) => {
359                common_telemetry::debug!(
360                    "Failed to prune regions for table {}, using all regions: {:?}",
361                    table_name,
362                    err
363                );
364                return Ok(all_regions);
365            }
366        };
367
368        common_telemetry::debug!(
369            "Region pruning for table {}: {} partition expressions applied, pruned from {} to {} regions",
370            table_name,
371            partition_expressions.len(),
372            all_regions.len(),
373            pruned_regions.len()
374        );
375
376        Ok(pruned_regions)
377    }
378
379    /// Resolves the partition-column types that are safe to use for region pruning.
380    ///
381    /// A logical metric table may not contain every physical partition column, either for backward
382    /// compatibility or because its physical table was repartitioned after the logical table was
383    /// created. Predicate extraction must remain bounded by the logical schema, while pruning
384    /// needs the physical datatypes to evaluate route expressions. Any lookup failure or
385    /// logical/physical datatype mismatch returns `None`, causing the caller to scan all regions.
386    async fn partition_column_types_for_pruning(
387        &self,
388        table_name: &TableName,
389        logical_table_info: &TableInfo,
390        physical_table_id: u32,
391        partition_expressions: &[PartitionExpr],
392        all_regions: &[RegionId],
393    ) -> Option<HashMap<String, ConcreteDataType>> {
394        let physical_partition_columns = if physical_table_id == logical_table_info.table_id() {
395            partition_column_types(logical_table_info)
396        } else {
397            match self
398                .catalog_manager
399                .table_info_by_id(physical_table_id)
400                .await
401            {
402                Ok(Some(physical_table_info)) => {
403                    partition_column_types(physical_table_info.as_ref())
404                }
405                Ok(None) => {
406                    debug!(
407                        "DistExtensionPlanner: physical table info not found for table {} (id: {}), using all regions: {:?}",
408                        table_name, physical_table_id, all_regions
409                    );
410                    return None;
411                }
412                Err(err) => {
413                    debug!(
414                        "DistExtensionPlanner: failed to load physical table info for table {} (id: {}): {}, using all regions: {:?}",
415                        table_name, physical_table_id, err, all_regions
416                    );
417                    return None;
418                }
419            }
420        };
421        let physical_column_types = physical_partition_columns
422            .into_iter()
423            .collect::<HashMap<_, _>>();
424        let logical_column_types = partition_column_types(logical_table_info)
425            .into_iter()
426            .collect::<HashMap<_, _>>();
427        let mut predicate_column_names = std::collections::HashSet::new();
428        for expression in partition_expressions {
429            expression.collect_column_names(&mut predicate_column_names);
430        }
431        if predicate_column_names
432            .iter()
433            .any(|name| logical_column_types.get(name) != physical_column_types.get(name))
434        {
435            debug!(
436                "DistExtensionPlanner: logical and physical partition metadata mismatch for table {} (physical id: {}), using all regions: {:?}",
437                table_name, physical_table_id, all_regions
438            );
439            return None;
440        }
441
442        Some(physical_column_types)
443    }
444
445    /// Input logical plan is analyzed. Thus only call logical optimizer to optimize it.
446    fn optimize_input_logical_plan(
447        &self,
448        session: &dyn Session,
449        plan: &LogicalPlan,
450    ) -> Result<LogicalPlan> {
451        expr_depth::validate_plan(plan)?;
452        let session_state = session
453            .as_any()
454            .downcast_ref::<SessionState>()
455            .ok_or_else(|| {
456                DataFusionError::Internal(
457                    "MergeScan requires a SessionState for logical optimization".to_string(),
458                )
459            })?;
460
461        session_state
462            .optimizer()
463            .optimize(plan.clone(), session_state, |_, _| {})
464    }
465}
466
467fn partition_column_types(table_info: &TableInfo) -> Vec<(String, ConcreteDataType)> {
468    table_info
469        .meta
470        .partition_columns()
471        .map(|column| (column.name.clone(), column.data_type.clone()))
472        .collect()
473}
474
475/// Extract the scan name and captured table identity from a logical plan.
476#[derive(Default)]
477struct TableScanExtractor {
478    pub table_name: Option<TableName>,
479    captured_table: Option<TableRef>,
480}
481
482impl TreeNodeVisitor<'_> for TableScanExtractor {
483    type Node = LogicalPlan;
484
485    fn f_down(&mut self, node: &Self::Node) -> Result<TreeNodeRecursion> {
486        match node {
487            LogicalPlan::TableScan(scan) => {
488                if let Some(source) = scan.source.downcast_ref::<DefaultTableSource>()
489                    && let Some(provider) = source
490                        .table_provider
491                        .downcast_ref::<DfTableProviderAdapter>()
492                {
493                    if provider.table().table_type() == TableType::Base {
494                        self.captured_table = Some(provider.table());
495                        let info = provider.table().table_info();
496                        self.table_name = Some(TableName::new(
497                            info.catalog_name.clone(),
498                            info.schema_name.clone(),
499                            info.name.clone(),
500                        ));
501                    }
502                    return Ok(TreeNodeRecursion::Stop);
503                }
504                match &scan.table_name {
505                    TableReference::Full {
506                        catalog,
507                        schema,
508                        table,
509                    } => {
510                        self.table_name = Some(TableName::new(
511                            catalog.to_string(),
512                            schema.to_string(),
513                            table.to_string(),
514                        ));
515                        Ok(TreeNodeRecursion::Stop)
516                    }
517                    // TODO(ruihang): Maybe the following two cases should not be valid
518                    TableReference::Partial { schema, table } => {
519                        self.table_name = Some(TableName::new(
520                            DEFAULT_CATALOG_NAME.to_string(),
521                            schema.to_string(),
522                            table.to_string(),
523                        ));
524                        Ok(TreeNodeRecursion::Stop)
525                    }
526                    TableReference::Bare { table } => {
527                        self.table_name = Some(TableName::new(
528                            DEFAULT_CATALOG_NAME.to_string(),
529                            DEFAULT_SCHEMA_NAME.to_string(),
530                            table.to_string(),
531                        ));
532                        Ok(TreeNodeRecursion::Stop)
533                    }
534                }
535            }
536            _ => Ok(TreeNodeRecursion::Continue),
537        }
538    }
539}
540
541#[cfg(test)]
542mod tests {
543    use std::collections::HashMap;
544    use std::sync::Arc;
545
546    use api::v1::region::{RemoteDynFilterUnregister, RemoteDynFilterUpdate};
547    use async_trait::async_trait;
548    use catalog::memory::MemoryCatalogManager;
549    use catalog::{CatalogManagerRef, RegisterTableRequest};
550    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
551    use common_meta::cache::new_table_route_cache;
552    use common_meta::key::TableMetadataManager;
553    use common_meta::key::table_route::TableRouteValue;
554    use common_meta::kv_backend::memory::MemoryKvBackend;
555    use common_meta::rpc::router::{Region, RegionRoute};
556    use common_query::request::QueryRequest;
557    use common_recordbatch::SendableRecordBatchStream;
558    use datafusion::datasource::DefaultTableSource;
559    use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, col as df_col, lit};
560    use datatypes::prelude::ConcreteDataType;
561    use datatypes::schema::{ColumnSchema, Schema};
562    use datatypes::value::Value;
563    use moka::future::CacheBuilder;
564    use partition::cache::new_partition_info_cache;
565    use partition::expr::{PartitionExpr, col as partition_col};
566    use partition::manager::PartitionRuleManager;
567    use session::ReadPreference;
568    use store_api::storage::RegionId;
569    use table::metadata::{TableInfo, TableInfoBuilder, TableMeta, TableType};
570    use table::table::adapter::DfTableProviderAdapter;
571    use table::table_name::TableName;
572    use table::test_util::EmptyTable;
573
574    use super::DistExtensionPlanner;
575    use crate::region_query::{RegionQueryHandler, RegionQueryTarget};
576
577    const LOGICAL_TABLE_ID: u32 = 1024;
578    const PHYSICAL_TABLE_ID: u32 = 2048;
579
580    struct UnusedRegionQueryHandler;
581
582    #[async_trait]
583    impl RegionQueryHandler for UnusedRegionQueryHandler {
584        async fn select_target(
585            &self,
586            _read_preference: ReadPreference,
587            _region_id: RegionId,
588        ) -> crate::error::Result<RegionQueryTarget> {
589            unreachable!("get_regions does not select region query targets")
590        }
591
592        async fn do_get(
593            &self,
594            _target: &RegionQueryTarget,
595            _request: QueryRequest,
596        ) -> crate::error::Result<SendableRecordBatchStream> {
597            unreachable!("get_regions does not query regions")
598        }
599
600        async fn handle_remote_dyn_filter_update(
601            &self,
602            _target: &RegionQueryTarget,
603            _query_id: String,
604            _update: RemoteDynFilterUpdate,
605        ) -> crate::error::Result<()> {
606            unreachable!("get_regions does not update dynamic filters")
607        }
608
609        async fn handle_remote_dyn_filter_unregister(
610            &self,
611            _target: &RegionQueryTarget,
612            _query_id: String,
613            _unregister: RemoteDynFilterUnregister,
614        ) -> crate::error::Result<()> {
615            unreachable!("get_regions does not unregister dynamic filters")
616        }
617    }
618
619    fn table_info(
620        table_id: u32,
621        name: &str,
622        columns: &[&str],
623        partition_keys: Vec<usize>,
624    ) -> TableInfo {
625        let schema = Arc::new(Schema::new(
626            columns
627                .iter()
628                .map(|name| ColumnSchema::new(*name, ConcreteDataType::string_datatype(), true))
629                .collect(),
630        ));
631        let meta = TableMeta {
632            schema,
633            primary_key_indices: vec![],
634            value_indices: vec![],
635            engine: "metric".to_string(),
636            next_column_id: columns.len() as u32,
637            options: Default::default(),
638            created_on: Default::default(),
639            updated_on: Default::default(),
640            partition_key_indices: partition_keys,
641            column_ids: (0..columns.len() as u32).collect(),
642        };
643        TableInfoBuilder::default()
644            .table_id(table_id)
645            .table_version(0)
646            .name(name.to_string())
647            .catalog_name(DEFAULT_CATALOG_NAME.to_string())
648            .schema_name(DEFAULT_SCHEMA_NAME.to_string())
649            .desc(None)
650            .table_type(TableType::Base)
651            .meta(meta)
652            .build()
653            .unwrap()
654    }
655
656    fn region_route(region_number: u32, expression: Option<PartitionExpr>) -> RegionRoute {
657        RegionRoute {
658            region: Region {
659                id: RegionId::new(PHYSICAL_TABLE_ID, region_number),
660                partition_expr: expression
661                    .map(|expression| expression.as_json_str().unwrap())
662                    .unwrap_or_default(),
663                ..Default::default()
664            },
665            ..Default::default()
666        }
667    }
668
669    async fn planner_and_plan(
670        physical_partition_keys: Vec<usize>,
671        expressions: Vec<Option<PartitionExpr>>,
672    ) -> (DistExtensionPlanner, LogicalPlan, TableName) {
673        let logical_info = table_info(LOGICAL_TABLE_ID, "logical", &["host"], vec![0]);
674        let physical_info = table_info(
675            PHYSICAL_TABLE_ID,
676            "physical",
677            &["host", "rack"],
678            physical_partition_keys,
679        );
680        let logical_table = EmptyTable::from_table_info(&logical_info);
681        let physical_table = EmptyTable::from_table_info(&physical_info);
682        let catalog_manager = MemoryCatalogManager::with_default_setup();
683        for table in [&logical_table, &physical_table] {
684            let info = table.table_info();
685            catalog_manager
686                .register_table_sync(RegisterTableRequest {
687                    catalog: info.catalog_name.clone(),
688                    schema: info.schema_name.clone(),
689                    table_name: info.name.clone(),
690                    table_id: info.table_id(),
691                    table: table.clone(),
692                })
693                .unwrap();
694        }
695
696        let backend = Arc::new(MemoryKvBackend::default());
697        let metadata_manager = TableMetadataManager::new(backend.clone());
698        let routes = expressions
699            .into_iter()
700            .enumerate()
701            .map(|(index, expression)| region_route(index as u32 + 1, expression))
702            .collect();
703        metadata_manager
704            .create_table_metadata(
705                physical_info,
706                TableRouteValue::physical(routes),
707                HashMap::new(),
708            )
709            .await
710            .unwrap();
711        metadata_manager
712            .create_table_metadata(
713                logical_info,
714                TableRouteValue::logical(PHYSICAL_TABLE_ID),
715                HashMap::new(),
716            )
717            .await
718            .unwrap();
719
720        let table_route_cache = Arc::new(new_table_route_cache(
721            "planner-test-routes".to_string(),
722            CacheBuilder::new(16).build(),
723            backend.clone(),
724        ));
725        let partition_info_cache = Arc::new(new_partition_info_cache(
726            "planner-test-partitions".to_string(),
727            CacheBuilder::new(16).build(),
728            table_route_cache.clone(),
729        ));
730        let partition_rule_manager = Arc::new(PartitionRuleManager::new(
731            backend,
732            table_route_cache,
733            partition_info_cache,
734        ));
735        let (resolved_physical_id, physical_route) = partition_rule_manager
736            .find_physical_table_route_with_id(PHYSICAL_TABLE_ID)
737            .await
738            .unwrap();
739        assert_eq!(PHYSICAL_TABLE_ID, resolved_physical_id);
740        let (resolved_logical_id, logical_route) = partition_rule_manager
741            .find_physical_table_route_with_id(LOGICAL_TABLE_ID)
742            .await
743            .unwrap();
744        assert_eq!(PHYSICAL_TABLE_ID, resolved_logical_id);
745        assert_eq!(physical_route.region_routes, logical_route.region_routes);
746        let catalog_manager: CatalogManagerRef = catalog_manager;
747        let planner = DistExtensionPlanner::new(
748            catalog_manager,
749            partition_rule_manager,
750            Arc::new(UnusedRegionQueryHandler),
751            false,
752        );
753        let table_source = Arc::new(DefaultTableSource::new(Arc::new(
754            DfTableProviderAdapter::new(logical_table),
755        )));
756        let plan = LogicalPlanBuilder::scan_with_filters("logical", table_source, None, vec![])
757            .unwrap()
758            .filter(df_col("host").eq(lit("a")))
759            .unwrap()
760            .build()
761            .unwrap();
762        (
763            planner,
764            plan,
765            TableName::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "logical"),
766        )
767    }
768
769    fn physical_partition_expressions() -> Vec<Option<PartitionExpr>> {
770        vec![
771            Some(partition_col("host").lt(Value::String("m".into()))),
772            Some(
773                partition_col("host")
774                    .gt_eq(Value::String("m".into()))
775                    .and(partition_col("rack").lt(Value::String("n".into()))),
776            ),
777            Some(
778                partition_col("host")
779                    .gt_eq(Value::String("m".into()))
780                    .and(partition_col("rack").gt_eq(Value::String("n".into()))),
781            ),
782        ]
783    }
784
785    #[tokio::test]
786    async fn region_routing_uses_captured_table() {
787        let (mut planner, plan, table_name) = planner_and_plan(vec![0], vec![None]).await;
788        planner.catalog_manager = MemoryCatalogManager::with_default_setup();
789        assert_eq!(
790            vec![RegionId::new(LOGICAL_TABLE_ID, 1)],
791            planner.get_regions(&table_name, &plan).await.unwrap()
792        );
793    }
794
795    #[tokio::test]
796    async fn logical_table_pruning_uses_physical_partition_datatypes() {
797        let (planner, plan, table_name) =
798            planner_and_plan(vec![0, 1], physical_partition_expressions()).await;
799
800        assert_eq!(
801            vec![RegionId::new(LOGICAL_TABLE_ID, 1)],
802            planner.get_regions(&table_name, &plan).await.unwrap()
803        );
804    }
805
806    #[tokio::test]
807    async fn missing_physical_partition_datatype_falls_back_to_all_logical_regions() {
808        let (planner, plan, table_name) =
809            planner_and_plan(vec![0], physical_partition_expressions()).await;
810
811        assert_all_logical_regions(planner.get_regions(&table_name, &plan).await.unwrap());
812    }
813
814    #[tokio::test]
815    async fn missing_route_partition_expression_falls_back_to_all_logical_regions() {
816        let mut expressions = physical_partition_expressions();
817        expressions[1] = None;
818        let (planner, plan, table_name) = planner_and_plan(vec![0, 1], expressions).await;
819
820        assert_all_logical_regions(planner.get_regions(&table_name, &plan).await.unwrap());
821    }
822
823    fn assert_all_logical_regions(mut regions: Vec<RegionId>) {
824        regions.sort_unstable();
825        assert_eq!(
826            vec![
827                RegionId::new(LOGICAL_TABLE_ID, 1),
828                RegionId::new(LOGICAL_TABLE_ID, 2),
829                RegionId::new(LOGICAL_TABLE_ID, 3),
830            ],
831            regions
832        );
833    }
834}
query/dist_plan/planner.rs

query/dist_plan/
planner.rs
