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

query/dist_plan/
planner.rs
