query/dist_plan/analyzer/
fallback.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//! Fallback dist plan analyzer, which will only push down table scan node
16//! This is used when `PlanRewriter` produce errors when trying to rewrite the plan
17//! This is a temporary solution, and will be removed once we have a more robust plan rewriter
18//!
19
20use std::collections::BTreeSet;
21
22use common_telemetry::debug;
23use datafusion::datasource::DefaultTableSource;
24use datafusion_common::Result as DfResult;
25use datafusion_common::tree_node::{Transformed, TreeNodeRewriter};
26use datafusion_expr::LogicalPlan;
27use table::metadata::TableType;
28use table::table::adapter::DfTableProviderAdapter;
29
30use crate::dist_plan::MergeScanLogicalPlan;
31use crate::dist_plan::analyzer::{AliasMapping, OTHER_PHY_PART_COL_PLACEHOLDER};
32
33/// FallbackPlanRewriter is a plan rewriter that will only push down table scan node
34/// This is used when `PlanRewriter` produce errors when trying to rewrite the plan
35/// This is a temporary solution, and will be removed once we have a more robust plan rewriter
36/// It will traverse the logical plan and rewrite table scan node to merge scan node
37#[derive(Debug, Clone, Default)]
38pub struct FallbackPlanRewriter;
39
40impl TreeNodeRewriter for FallbackPlanRewriter {
41    type Node = LogicalPlan;
42
43    fn f_down(
44        &mut self,
45        plan: Self::Node,
46    ) -> DfResult<datafusion_common::tree_node::Transformed<Self::Node>> {
47        if let LogicalPlan::TableScan(table_scan) = &plan {
48            let partition_cols = if let Some(source) = table_scan
49                .source
50                .as_any()
51                .downcast_ref::<DefaultTableSource>()
52            {
53                if let Some(provider) = source
54                    .table_provider
55                    .as_any()
56                    .downcast_ref::<DfTableProviderAdapter>()
57                {
58                    if provider.table().table_type() == TableType::Base {
59                        let info = provider.table().table_info();
60                        let partition_key_indices = info.meta.partition_key_indices.clone();
61                        let schema = info.meta.schema.clone();
62                        let partition_cols = partition_key_indices
63                            .into_iter()
64                            .map(|index| schema.column_name_by_index(index).to_string())
65                            .collect::<Vec<String>>();
66                        debug!(
67                            "FallbackPlanRewriter: table {} has partition columns: {:?}",
68                            info.name, partition_cols
69                        );
70                        Some(partition_cols
71                                .into_iter()
72                                .map(|c| {
73                                    if c == OTHER_PHY_PART_COL_PLACEHOLDER {
74                                        // for placeholder, just return a empty alias
75                                        return Ok((c.clone(), BTreeSet::new()));
76                                    }
77                                    let index =
78                                        plan.schema().index_of_column_by_name(None, &c).ok_or_else(|| {
79                                            datafusion_common::DataFusionError::Internal(
80                                                format!(
81                                                    "PlanRewriter: maybe_set_partitions: column {c} not found in schema of plan: {plan}"
82                                                ),
83                                            )
84                                        })?;
85                                    let column = plan.schema().columns().get(index).cloned().ok_or_else(|| {
86                                        datafusion_common::DataFusionError::Internal(format!(
87                                            "PlanRewriter: maybe_set_partitions: column index {index} out of bounds in schema of plan: {plan}"
88                                        ))
89                                    })?;
90                                    Ok((c.clone(), BTreeSet::from([column])))
91                                })
92                                .collect::<DfResult<AliasMapping>>()?)
93                    } else {
94                        None
95                    }
96                } else {
97                    None
98                }
99            } else {
100                None
101            };
102            let node = MergeScanLogicalPlan::new(
103                plan,
104                false,
105                // at this stage, the partition cols should be set
106                // treat it as non-partitioned if None
107                partition_cols.clone().unwrap_or_default(),
108            )
109            .into_logical_plan();
110            Ok(Transformed::yes(node))
111        } else {
112            Ok(Transformed::no(plan))
113        }
114    }
115}
query/dist_plan/analyzer/fallback.rs

query/dist_plan/analyzer/
fallback.rs
