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;
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                                    let index =
74                                        plan.schema().index_of_column_by_name(None, &c).ok_or_else(|| {
75                                            datafusion_common::DataFusionError::Internal(
76                                                format!(
77                                                    "PlanRewriter: maybe_set_partitions: column {c} not found in schema of plan: {plan}"
78                                                ),
79                                            )
80                                        })?;
81                                    let column = plan.schema().columns().get(index).cloned().ok_or_else(|| {
82                                        datafusion_common::DataFusionError::Internal(format!(
83                                            "PlanRewriter: maybe_set_partitions: column index {index} out of bounds in schema of plan: {plan}"
84                                        ))
85                                    })?;
86                                    Ok((c.clone(), BTreeSet::from([column])))
87                                })
88                                .collect::<DfResult<AliasMapping>>()?)
89                    } else {
90                        None
91                    }
92                } else {
93                    None
94                }
95            } else {
96                None
97            };
98            let node = MergeScanLogicalPlan::new(
99                plan,
100                false,
101                // at this stage, the partition cols should be set
102                // treat it as non-partitioned if None
103                partition_cols.clone().unwrap_or_default(),
104            )
105            .into_logical_plan();
106            Ok(Transformed::yes(node))
107        } else {
108            Ok(Transformed::no(plan))
109        }
110    }
111}
query/dist_plan/analyzer/fallback.rs

query/dist_plan/analyzer/
fallback.rs
