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

query/dist_plan/
planner.rs
