query/dist_plan/analyzer/
utils.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
15use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
16use std::sync::Arc;
17
18use arrow::array::ArrayRef;
19use arrow_schema::{ArrowError, DataType};
20use chrono::{DateTime, Utc};
21use datafusion::common::alias::AliasGenerator;
22use datafusion::config::ConfigOptions;
23use datafusion::error::Result as DfResult;
24use datafusion_common::Column;
25use datafusion_common::tree_node::{Transformed, TreeNode as _, TreeNodeRewriter};
26use datafusion_expr::expr::Alias;
27use datafusion_expr::{Expr, Extension, LogicalPlan};
28use datafusion_optimizer::simplify_expressions::SimplifyExpressions;
29use datafusion_optimizer::{OptimizerConfig, OptimizerRule as _};
30
31use crate::dist_plan::merge_sort::MergeSortLogicalPlan;
32use crate::plan::ExtractExpr as _;
33
34/// The `ConstEvaluator` in `SimplifyExpressions` might evaluate some UDFs early in the
35/// planning stage, by executing them directly. For example, the `database()` function.
36/// So the `ConfigOptions` here (which is set from the session context) should be present
37/// in the UDF's `ScalarFunctionArgs`. However, the default implementation in DataFusion
38/// seems to lost track on it: the `ConfigOptions` is recreated with its default values again.
39/// So we create a custom `OptimizerConfig` with the desired `ConfigOptions`
40/// to walk around the issue.
41/// TODO(LFC): Maybe use DataFusion's `OptimizerContext` again
42///   once https://github.com/apache/datafusion/pull/17742 is merged.
43pub(crate) struct PatchOptimizerContext {
44    pub(crate) inner: datafusion_optimizer::OptimizerContext,
45    pub(crate) config: Arc<ConfigOptions>,
46    /// Override for `query_execution_start_time()` — used during scheduled Flow
47    /// evaluation so that `SimplifyExpressions` does not constant-fold `now()`
48    /// into wall-clock literals.  When `None`, falls back to the inner context.
49    pub(crate) scheduled_time: Option<DateTime<Utc>>,
50}
51
52impl OptimizerConfig for PatchOptimizerContext {
53    fn query_execution_start_time(&self) -> Option<DateTime<Utc>> {
54        self.scheduled_time
55            .or_else(|| self.inner.query_execution_start_time())
56    }
57
58    fn alias_generator(&self) -> &Arc<AliasGenerator> {
59        self.inner.alias_generator()
60    }
61
62    fn options(&self) -> Arc<ConfigOptions> {
63        self.config.clone()
64    }
65}
66
67/// Simplify all expressions recursively in the plan tree
68/// which keeping the output schema unchanged
69pub(crate) struct PlanTreeExpressionSimplifier {
70    optimizer_context: PatchOptimizerContext,
71}
72
73impl PlanTreeExpressionSimplifier {
74    pub fn new(optimizer_context: PatchOptimizerContext) -> Self {
75        Self { optimizer_context }
76    }
77}
78
79impl TreeNodeRewriter for PlanTreeExpressionSimplifier {
80    type Node = LogicalPlan;
81    fn f_down(&mut self, plan: Self::Node) -> DfResult<Transformed<Self::Node>> {
82        let simp = SimplifyExpressions::new()
83            .rewrite(plan, &self.optimizer_context)?
84            .data;
85        Ok(Transformed::yes(simp))
86    }
87}
88
89/// A patch for substrait simply throw timezone away, so when decoding, if columns have different timezone then expected schema, use expected schema's timezone
90pub fn patch_batch_timezone(
91    expected_schema: arrow_schema::SchemaRef,
92    columns: Vec<ArrayRef>,
93) -> Result<arrow::record_batch::RecordBatch, ArrowError> {
94    let patched_columns: Vec<ArrayRef> = expected_schema
95        .fields()
96        .iter()
97        .zip(columns)
98        .map(|(expected_field, column)| {
99            let expected_type = expected_field.data_type();
100            let actual_type = column.data_type();
101
102            // Check if both are timestamp types with different timezones
103            match (expected_type, actual_type) {
104                (
105                    DataType::Timestamp(expected_unit, expected_tz),
106                    DataType::Timestamp(actual_unit, actual_tz),
107                ) if expected_unit == actual_unit && expected_tz != actual_tz => {
108                    // Cast the column to the expected timezone
109                    arrow::compute::cast(&column, expected_type)
110                }
111                _ => Ok(column),
112            }
113        })
114        .collect::<Result<Vec<_>, _>>()?;
115
116    arrow::record_batch::RecordBatch::try_new(expected_schema.clone(), patched_columns)
117}
118
119fn rewrite_column(
120    mapping: &BTreeMap<Column, BTreeSet<Column>>,
121    original_node: &LogicalPlan,
122    alias_node: &LogicalPlan,
123) -> impl Fn(Expr) -> DfResult<Transformed<Expr>> {
124    move |e: Expr| {
125        if let Expr::Column(col) = e {
126            if let Some(aliased_cols) = mapping.get(&col) {
127                // if multiple alias is available, just use first one
128                if let Some(aliased_col) = aliased_cols.iter().next() {
129                    Ok(Transformed::yes(Expr::Column(aliased_col.clone())))
130                } else {
131                    Err(datafusion_common::DataFusionError::Internal(format!(
132                        "PlanRewriter: expand: column {col} from {original_node}\n has empty alias set in plan: {alias_node}\n but expect at least one alias",
133                    )))
134                }
135            } else {
136                Err(datafusion_common::DataFusionError::Internal(format!(
137                    "PlanRewriter: expand: column {col} from {original_node}\n has no alias in plan: {alias_node}",
138                )))
139            }
140        } else {
141            Ok(Transformed::no(e))
142        }
143    }
144}
145
146/// Rewrite the expressions of the given merge sort plan from original columns(at merge sort's input plan) to aliased columns at the given aliased node
147pub fn rewrite_merge_sort_exprs(
148    merge_sort: &MergeSortLogicalPlan,
149    aliased_node: &LogicalPlan,
150) -> DfResult<LogicalPlan> {
151    let merge_sort = LogicalPlan::Extension(Extension {
152        node: Arc::new(merge_sort.clone()),
153    });
154
155    // tracking alias for sort exprs,
156    let sort_input = merge_sort.inputs().first().cloned().ok_or_else(|| {
157        datafusion_common::DataFusionError::Internal(format!(
158            "PlanRewriter: expand: merge sort stage has no input: {merge_sort}"
159        ))
160    })?;
161    let sort_exprs = merge_sort.expressions_consider_join();
162    let column_refs = sort_exprs
163        .iter()
164        .flat_map(|e| e.column_refs().into_iter().cloned())
165        .collect::<BTreeSet<_>>();
166    let column_alias_mapping = aliased_columns_for(&column_refs, aliased_node, Some(sort_input))?;
167    let aliased_sort_exprs = sort_exprs
168        .into_iter()
169        .map(|e| {
170            e.transform(rewrite_column(
171                &column_alias_mapping,
172                &merge_sort,
173                aliased_node,
174            ))
175        })
176        .map(|e| e.map(|e| e.data))
177        .collect::<DfResult<Vec<_>>>()?;
178    let new_merge_sort = merge_sort.with_new_exprs(
179        aliased_sort_exprs,
180        merge_sort.inputs().into_iter().cloned().collect(),
181    )?;
182    Ok(new_merge_sort)
183}
184
185/// Return all the original columns(at original node) for the given aliased columns at the aliased node
186///
187/// if `original_node` is None, it means original columns are from leaf node
188///
189/// Return value use `BTreeMap` to have deterministic order for choose first alias when multiple alias exist
190#[allow(unused)]
191pub fn original_column_for(
192    aliased_columns: &BTreeSet<Column>,
193    aliased_node: LogicalPlan,
194    original_node: Option<Arc<LogicalPlan>>,
195) -> DfResult<BTreeMap<Column, Column>> {
196    let schema_cols: BTreeSet<Column> = aliased_node.schema().columns().iter().cloned().collect();
197    let cur_aliases: BTreeMap<Column, Column> = aliased_columns
198        .iter()
199        .filter(|c| schema_cols.contains(c))
200        .map(|c| (c.clone(), c.clone()))
201        .collect();
202
203    if cur_aliases.is_empty() {
204        return Ok(BTreeMap::new());
205    }
206
207    original_column_for_inner(cur_aliases, &aliased_node, &original_node)
208}
209
210fn original_column_for_inner(
211    mut cur_aliases: BTreeMap<Column, Column>,
212    node: &LogicalPlan,
213    original_node: &Option<Arc<LogicalPlan>>,
214) -> DfResult<BTreeMap<Column, Column>> {
215    let mut current_node = node;
216
217    loop {
218        // Base case: check if we've reached the target node
219        if let Some(original_node) = original_node
220            && *current_node == **original_node
221        {
222            return Ok(cur_aliases);
223        } else if current_node.inputs().is_empty() {
224            // leaf node reached
225            return Ok(cur_aliases);
226        }
227
228        // Validate node has exactly one child
229        if current_node.inputs().len() != 1 {
230            return Err(datafusion::error::DataFusionError::Internal(format!(
231                "only accept plan with at most one child, found: {}",
232                current_node
233            )));
234        }
235
236        // Get alias layer and update aliases
237        let layer = get_alias_layer_from_node(current_node)?;
238        let mut new_aliases = BTreeMap::new();
239        for (start_alias, cur_alias) in cur_aliases {
240            if let Some(old_column) = layer.get_old_from_new(cur_alias.clone()) {
241                new_aliases.insert(start_alias, old_column);
242            }
243        }
244
245        // Move to child node and continue iteration
246        cur_aliases = new_aliases;
247        current_node = current_node.inputs()[0];
248    }
249}
250
251/// Return all the aliased columns(at aliased node) for the given original columns(at original node)
252///
253/// if `original_node` is None, it means original columns are from leaf node
254///
255/// Return value use `BTreeMap` to have deterministic order for choose first alias when multiple alias exist
256pub fn aliased_columns_for(
257    original_columns: &BTreeSet<Column>,
258    aliased_node: &LogicalPlan,
259    original_node: Option<&LogicalPlan>,
260) -> DfResult<BTreeMap<Column, BTreeSet<Column>>> {
261    let initial_aliases: BTreeMap<Column, BTreeSet<Column>> = {
262        if let Some(original) = &original_node {
263            let schema_cols: BTreeSet<Column> = original.schema().columns().into_iter().collect();
264            original_columns
265                .iter()
266                .filter(|c| schema_cols.contains(c))
267                .map(|c| (c.clone(), [c.clone()].into()))
268                .collect()
269        } else {
270            original_columns
271                .iter()
272                .map(|c| (c.clone(), [c.clone()].into()))
273                .collect()
274        }
275    };
276
277    if initial_aliases.is_empty() {
278        return Ok(BTreeMap::new());
279    }
280
281    aliased_columns_for_inner(initial_aliases, aliased_node, original_node)
282}
283
284fn aliased_columns_for_inner(
285    cur_aliases: BTreeMap<Column, BTreeSet<Column>>,
286    node: &LogicalPlan,
287    original_node: Option<&LogicalPlan>,
288) -> DfResult<BTreeMap<Column, BTreeSet<Column>>> {
289    // First, collect the path from current node to the target node
290    let mut path = Vec::new();
291    let mut current_node = node;
292
293    // Descend to the target node, collecting nodes along the way
294    loop {
295        // Base case: check if we've reached the target node
296        if let Some(original_node) = original_node
297            && *current_node == *original_node
298        {
299            break;
300        } else if current_node.inputs().is_empty() {
301            // leaf node reached
302            break;
303        }
304
305        // Validate node has exactly one child
306        if current_node.inputs().len() != 1 {
307            return Err(datafusion::error::DataFusionError::Internal(format!(
308                "only accept plan with at most one child, found: {}",
309                current_node
310            )));
311        }
312
313        // Add current node to path and move to child
314        path.push(current_node);
315        current_node = current_node.inputs()[0];
316    }
317
318    // Now apply alias layers in reverse order (from original to aliased)
319    let mut result = cur_aliases;
320    for &node_in_path in path.iter().rev() {
321        let layer = get_alias_layer_from_node(node_in_path)?;
322        let mut new_aliases = BTreeMap::new();
323        for (original_column, cur_alias_set) in result {
324            let mut new_alias_set = BTreeSet::new();
325            for cur_alias in cur_alias_set {
326                new_alias_set.extend(layer.get_new_from_old(cur_alias.clone()));
327            }
328            if !new_alias_set.is_empty() {
329                new_aliases.insert(original_column, new_alias_set);
330            }
331        }
332        result = new_aliases;
333    }
334
335    Ok(result)
336}
337
338/// Return a mapping of original column to all the aliased columns in current node of the plan
339/// TODO(discord9): also support merge scan node
340fn get_alias_layer_from_node(node: &LogicalPlan) -> DfResult<AliasLayer> {
341    match node {
342        LogicalPlan::Projection(proj) => Ok(get_alias_layer_from_exprs(&proj.expr)),
343        LogicalPlan::Aggregate(aggr) => Ok(get_alias_layer_from_exprs(&aggr.group_expr)),
344        LogicalPlan::SubqueryAlias(subquery_alias) => {
345            let mut layer = AliasLayer::default();
346            let old_columns = subquery_alias.input.schema().columns();
347            for old_column in old_columns {
348                let new_column = Column::new(
349                    Some(subquery_alias.alias.clone()),
350                    old_column.name().to_string(),
351                );
352                // mapping from old_column to new_column
353                layer.insert_alias(old_column, [new_column].into());
354            }
355            Ok(layer)
356        }
357        LogicalPlan::TableScan(scan) => {
358            let columns = scan.projected_schema.columns();
359            let mut layer = AliasLayer::default();
360            for col in columns {
361                layer.insert_alias(col.clone(), [col.clone()].into());
362            }
363            Ok(layer)
364        }
365        _ => {
366            let input_schema = node
367                .inputs()
368                .first()
369                .ok_or_else(|| {
370                    datafusion::error::DataFusionError::Internal(format!(
371                        "only accept plan with at most one child, found: {}",
372                        node
373                    ))
374                })?
375                .schema();
376            let output_schema = node.schema();
377            // only accept at most one child plan, and if not one of the above nodes,
378            // also shouldn't modify the schema or else alias scope tracker can't support them
379            if node.inputs().len() > 1 {
380                Err(datafusion::error::DataFusionError::Internal(format!(
381                    "only accept plan with at most one child, found: {}",
382                    node
383                )))
384            } else if node.inputs().len() == 1 {
385                if input_schema != output_schema {
386                    let input_columns = input_schema.columns();
387                    let all_input_is_in_output = input_columns
388                        .iter()
389                        .all(|c| output_schema.is_column_from_schema(c));
390                    if all_input_is_in_output {
391                        // all input is in output, so it's just adding some columns, we can do identity mapping for input columns
392                        let mut layer = AliasLayer::default();
393                        for col in input_columns {
394                            layer.insert_alias(col.clone(), [col.clone()].into());
395                        }
396                        Ok(layer)
397                    } else {
398                        // otherwise use the intersection of input and output
399                        // TODO(discord9): maybe just make this case unsupported for now?
400                        common_telemetry::debug!(
401                            "Might be unsupported plan for alias tracking, track alias anyway: {}",
402                            node
403                        );
404                        let input_columns = input_schema.columns();
405                        let output_columns =
406                            output_schema.columns().into_iter().collect::<HashSet<_>>();
407                        let common_columns: HashSet<Column> = input_columns
408                            .iter()
409                            .filter(|c| output_columns.contains(c))
410                            .cloned()
411                            .collect();
412
413                        let mut layer = AliasLayer::default();
414                        for col in &common_columns {
415                            layer.insert_alias(col.clone(), [col.clone()].into());
416                        }
417                        Ok(layer)
418                    }
419                } else {
420                    // identity mapping
421                    let mut layer = AliasLayer::default();
422                    for col in output_schema.columns() {
423                        layer.insert_alias(col.clone(), [col.clone()].into());
424                    }
425                    Ok(layer)
426                }
427            } else {
428                // unknown plan with no input, error msg
429                Err(datafusion::error::DataFusionError::Internal(format!(
430                    "Unsupported plan with no input: {}",
431                    node
432                )))
433            }
434        }
435    }
436}
437
438fn get_alias_layer_from_exprs(exprs: &[Expr]) -> AliasLayer {
439    let mut alias_mapping: HashMap<Column, HashSet<Column>> = HashMap::new();
440    for expr in exprs {
441        if let Expr::Alias(alias) = expr {
442            if let Some(column) = get_alias_original_column(alias) {
443                alias_mapping
444                    .entry(column.clone())
445                    .or_default()
446                    .insert(Column::new(alias.relation.clone(), alias.name.clone()));
447            }
448        } else if let Expr::Column(column) = expr {
449            // identity mapping
450            alias_mapping
451                .entry(column.clone())
452                .or_default()
453                .insert(column.clone());
454        }
455    }
456    let mut layer = AliasLayer::default();
457    for (old_column, new_columns) in alias_mapping {
458        layer.insert_alias(old_column, new_columns);
459    }
460    layer
461}
462
463#[derive(Default, Debug, Clone)]
464struct AliasLayer {
465    /// for convenient of querying, key is field's name
466    old_to_new: BTreeMap<Column, HashSet<Column>>,
467}
468
469impl AliasLayer {
470    pub fn insert_alias(&mut self, old_column: Column, new_columns: HashSet<Column>) {
471        self.old_to_new
472            .entry(old_column)
473            .or_default()
474            .extend(new_columns);
475    }
476
477    pub fn get_new_from_old(&self, old_column: Column) -> HashSet<Column> {
478        let mut res_cols = HashSet::new();
479        for (old, new_cols) in self.old_to_new.iter() {
480            if old.name() == old_column.name() {
481                match (&old.relation, &old_column.relation) {
482                    (Some(o), Some(c)) => {
483                        if o.resolved_eq(c) {
484                            res_cols.extend(new_cols.clone());
485                        }
486                    }
487                    _ => {
488                        // if any of the two relation is None, meaning not fully qualified, just match name
489                        res_cols.extend(new_cols.clone());
490                    }
491                }
492            }
493        }
494        res_cols
495    }
496
497    pub fn get_old_from_new(&self, new_column: Column) -> Option<Column> {
498        for (old, new_set) in &self.old_to_new {
499            if new_set.iter().any(|n| {
500                if n.name() != new_column.name() {
501                    return false;
502                }
503                match (&n.relation, &new_column.relation) {
504                    (Some(r1), Some(r2)) => r1.resolved_eq(r2),
505                    _ => true,
506                }
507            }) {
508                return Some(old.clone());
509            }
510        }
511        None
512    }
513}
514
515fn get_alias_original_column(alias: &Alias) -> Option<Column> {
516    let mut cur_alias = alias;
517    while let Expr::Alias(inner_alias) = cur_alias.expr.as_ref() {
518        cur_alias = inner_alias;
519    }
520    if let Expr::Column(column) = cur_alias.expr.as_ref() {
521        return Some(column.clone());
522    }
523
524    None
525}
526
527/// Mapping of original column in table to all the alias at current node
528pub type AliasMapping = BTreeMap<String, BTreeSet<Column>>;
529
530#[cfg(test)]
531mod tests {
532    use std::sync::Arc;
533
534    use common_telemetry::init_default_ut_logging;
535    use datafusion::datasource::DefaultTableSource;
536    use datafusion::functions_aggregate::min_max::{max, min};
537    use datafusion_expr::{LogicalPlanBuilder, col};
538    use pretty_assertions::assert_eq;
539    use table::table::adapter::DfTableProviderAdapter;
540
541    use super::*;
542    use crate::dist_plan::analyzer::test::TestTable;
543
544    fn qcol(name: &str) -> Column {
545        Column::from_qualified_name(name)
546    }
547
548    #[test]
549    fn proj_multi_layered_alias_tracker() {
550        // use logging for better debugging
551        init_default_ut_logging();
552        let test_table = TestTable::table_with_name(0, "t".to_string());
553        let table_source = Arc::new(DefaultTableSource::new(Arc::new(
554            DfTableProviderAdapter::new(test_table),
555        )));
556        let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![])
557            .unwrap()
558            .project(vec![
559                col("number"),
560                col("pk3").alias("pk1"),
561                col("pk3").alias("pk2"),
562            ])
563            .unwrap()
564            .project(vec![
565                col("number"),
566                col("pk2").alias("pk4"),
567                col("pk1").alias("pk5"),
568            ])
569            .unwrap()
570            .build()
571            .unwrap();
572
573        let child = plan.inputs()[0].clone();
574
575        assert_eq!(
576            aliased_columns_for(&[qcol("pk1"), qcol("pk2")].into(), &plan, Some(&child)).unwrap(),
577            [
578                (qcol("pk1"), [qcol("pk5")].into()),
579                (qcol("pk2"), [qcol("pk4")].into())
580            ]
581            .into()
582        );
583
584        // columns not in the plan should return empty mapping
585        assert_eq!(
586            aliased_columns_for(&[qcol("pk1"), qcol("pk2")].into(), &plan, Some(&plan)).unwrap(),
587            [].into()
588        );
589
590        assert_eq!(
591            aliased_columns_for(&[qcol("t.pk3")].into(), &plan, Some(&child)).unwrap(),
592            [].into()
593        );
594
595        assert_eq!(
596            original_column_for(&[qcol("pk5"), qcol("pk4")].into(), plan.clone(), None).unwrap(),
597            [(qcol("pk5"), qcol("t.pk3")), (qcol("pk4"), qcol("t.pk3"))].into()
598        );
599
600        assert_eq!(
601            aliased_columns_for(&[qcol("pk3")].into(), &plan, None).unwrap(),
602            [(qcol("pk3"), [qcol("pk5"), qcol("pk4")].into())].into()
603        );
604        assert_eq!(
605            original_column_for(&[qcol("pk1"), qcol("pk2")].into(), child.clone(), None).unwrap(),
606            [(qcol("pk1"), qcol("t.pk3")), (qcol("pk2"), qcol("t.pk3"))].into()
607        );
608
609        assert_eq!(
610            aliased_columns_for(&[qcol("pk3")].into(), &child, None).unwrap(),
611            [(qcol("pk3"), [qcol("pk1"), qcol("pk2")].into())].into()
612        );
613
614        assert_eq!(
615            original_column_for(
616                &[qcol("pk4"), qcol("pk5")].into(),
617                plan.clone(),
618                Some(Arc::new(child.clone()))
619            )
620            .unwrap(),
621            [(qcol("pk4"), qcol("pk2")), (qcol("pk5"), qcol("pk1"))].into()
622        );
623    }
624
625    #[test]
626    fn sort_subquery_alias_layered_tracker() {
627        let test_table = TestTable::table_with_name(0, "t".to_string());
628        let table_source = Arc::new(DefaultTableSource::new(Arc::new(
629            DfTableProviderAdapter::new(test_table),
630        )));
631
632        let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![])
633            .unwrap()
634            .sort(vec![col("t.number").sort(true, false)])
635            .unwrap()
636            .alias("a")
637            .unwrap()
638            .build()
639            .unwrap();
640
641        let sort_plan = plan.inputs()[0].clone();
642        let scan_plan = sort_plan.inputs()[0].clone();
643
644        // Test aliased_columns_for from scan to final plan
645        assert_eq!(
646            aliased_columns_for(&[qcol("t.number")].into(), &plan, Some(&scan_plan)).unwrap(),
647            [(qcol("t.number"), [qcol("a.number")].into())].into()
648        );
649
650        // Test aliased_columns_for from sort to final plan
651        assert_eq!(
652            aliased_columns_for(&[qcol("t.number")].into(), &plan, Some(&sort_plan)).unwrap(),
653            [(qcol("t.number"), [qcol("a.number")].into())].into()
654        );
655
656        // Test aliased_columns_for from leaf to final plan
657        assert_eq!(
658            aliased_columns_for(&[qcol("t.number")].into(), &plan, None).unwrap(),
659            [(qcol("t.number"), [qcol("a.number")].into())].into()
660        );
661
662        // Test original_column_for from final plan to scan
663        assert_eq!(
664            original_column_for(
665                &[qcol("a.number")].into(),
666                plan.clone(),
667                Some(Arc::new(scan_plan.clone()))
668            )
669            .unwrap(),
670            [(qcol("a.number"), qcol("t.number"))].into()
671        );
672
673        // Test original_column_for from final plan to sort
674        assert_eq!(
675            original_column_for(
676                &[qcol("a.number")].into(),
677                plan.clone(),
678                Some(Arc::new(sort_plan.clone()))
679            )
680            .unwrap(),
681            [(qcol("a.number"), qcol("t.number"))].into()
682        );
683    }
684
685    #[test]
686    fn proj_alias_layered_tracker() {
687        // use logging for better debugging
688        init_default_ut_logging();
689        let test_table = TestTable::table_with_name(0, "t".to_string());
690        let table_source = Arc::new(DefaultTableSource::new(Arc::new(
691            DfTableProviderAdapter::new(test_table),
692        )));
693        let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![])
694            .unwrap()
695            .project(vec![
696                col("number"),
697                col("pk3").alias("pk1"),
698                col("pk2").alias("pk3"),
699            ])
700            .unwrap()
701            .project(vec![
702                col("number"),
703                col("pk1").alias("pk2"),
704                col("pk3").alias("pk1"),
705            ])
706            .unwrap()
707            .build()
708            .unwrap();
709
710        let first_proj = plan.inputs()[0].clone();
711        let scan_plan = first_proj.inputs()[0].clone();
712
713        // Test original_column_for from final plan to scan
714        assert_eq!(
715            original_column_for(
716                &[qcol("pk1")].into(),
717                plan.clone(),
718                Some(Arc::new(scan_plan.clone()))
719            )
720            .unwrap(),
721            [(qcol("pk1"), qcol("t.pk2"))].into()
722        );
723
724        // Test original_column_for from final plan to first projection
725        assert_eq!(
726            original_column_for(
727                &[qcol("pk1")].into(),
728                plan.clone(),
729                Some(Arc::new(first_proj.clone()))
730            )
731            .unwrap(),
732            [(qcol("pk1"), qcol("pk3"))].into()
733        );
734
735        // Test original_column_for from final plan to leaf
736        assert_eq!(
737            original_column_for(
738                &[qcol("pk1")].into(),
739                plan.clone(),
740                Some(Arc::new(plan.clone()))
741            )
742            .unwrap(),
743            [(qcol("pk1"), qcol("pk1"))].into()
744        );
745
746        // Test aliased_columns_for from scan to first projection
747        assert_eq!(
748            aliased_columns_for(&[qcol("t.pk2")].into(), &first_proj, Some(&scan_plan)).unwrap(),
749            [(qcol("t.pk2"), [qcol("pk3")].into())].into()
750        );
751
752        // Test aliased_columns_for from first projection to final plan
753        assert_eq!(
754            aliased_columns_for(&[qcol("pk3")].into(), &plan, Some(&first_proj)).unwrap(),
755            [(qcol("pk3"), [qcol("pk1")].into())].into()
756        );
757
758        // Test aliased_columns_for from scan to final plan
759        assert_eq!(
760            aliased_columns_for(&[qcol("t.pk2")].into(), &plan, Some(&scan_plan)).unwrap(),
761            [(qcol("t.pk2"), [qcol("pk1")].into())].into()
762        );
763
764        // Test aliased_columns_for from leaf to final plan
765        assert_eq!(
766            aliased_columns_for(&[qcol("pk2")].into(), &plan, None).unwrap(),
767            [(qcol("pk2"), [qcol("pk1")].into())].into()
768        );
769    }
770
771    #[test]
772    fn proj_alias_relation_layered_tracker() {
773        // use logging for better debugging
774        init_default_ut_logging();
775        let test_table = TestTable::table_with_name(0, "t".to_string());
776        let table_source = Arc::new(DefaultTableSource::new(Arc::new(
777            DfTableProviderAdapter::new(test_table),
778        )));
779        let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![])
780            .unwrap()
781            .project(vec![
782                col("number"),
783                col("pk3").alias_qualified(Some("b"), "pk1"),
784                col("pk2").alias_qualified(Some("a"), "pk1"),
785            ])
786            .unwrap()
787            .build()
788            .unwrap();
789
790        let scan_plan = plan.inputs()[0].clone();
791
792        // Test aliased_columns_for from scan to projection
793        assert_eq!(
794            aliased_columns_for(&[qcol("t.pk2")].into(), &plan, Some(&scan_plan)).unwrap(),
795            [(qcol("t.pk2"), [qcol("a.pk1")].into())].into()
796        );
797    }
798
799    #[test]
800    fn proj_alias_aliased_aggr() {
801        // use logging for better debugging
802        init_default_ut_logging();
803        let test_table = TestTable::table_with_name(0, "t".to_string());
804        let table_source = Arc::new(DefaultTableSource::new(Arc::new(
805            DfTableProviderAdapter::new(test_table),
806        )));
807        let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![])
808            .unwrap()
809            .project(vec![
810                col("number"),
811                col("pk1").alias("pk3"),
812                col("pk2").alias("pk4"),
813            ])
814            .unwrap()
815            .project(vec![
816                col("number"),
817                col("pk3").alias("pk42"),
818                col("pk4").alias("pk43"),
819            ])
820            .unwrap()
821            .aggregate(vec![col("pk42"), col("pk43")], vec![min(col("number"))])
822            .unwrap()
823            .build()
824            .unwrap();
825
826        let aggr_plan = plan.clone();
827        let second_proj = aggr_plan.inputs()[0].clone();
828        let first_proj = second_proj.inputs()[0].clone();
829        let scan_plan = first_proj.inputs()[0].clone();
830
831        // Test aliased_columns_for from scan to final plan
832        assert_eq!(
833            aliased_columns_for(&[qcol("t.pk1")].into(), &plan, Some(&scan_plan)).unwrap(),
834            [(qcol("t.pk1"), [qcol("pk42")].into())].into()
835        );
836
837        // Test aliased_columns_for from scan to first projection
838        assert_eq!(
839            aliased_columns_for(&[Column::from_name("pk1")].into(), &first_proj, None).unwrap(),
840            [(Column::from_name("pk1"), [qcol("pk3")].into())].into()
841        );
842    }
843
844    #[test]
845    fn aggr_aggr_alias() {
846        // use logging for better debugging
847        init_default_ut_logging();
848        let test_table = TestTable::table_with_name(0, "t".to_string());
849        let table_source = Arc::new(DefaultTableSource::new(Arc::new(
850            DfTableProviderAdapter::new(test_table),
851        )));
852        let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![])
853            .unwrap()
854            .aggregate(vec![col("pk1"), col("pk2")], vec![max(col("number"))])
855            .unwrap()
856            .aggregate(
857                vec![col("pk1"), col("pk2")],
858                vec![min(col("max(t.number)"))],
859            )
860            .unwrap()
861            .build()
862            .unwrap();
863
864        let second_aggr = plan.clone();
865        let first_aggr = second_aggr.inputs()[0].clone();
866        let scan_plan = first_aggr.inputs()[0].clone();
867
868        // Test aliased_columns_for from scan to final plan (identity mapping for aggregates)
869        assert_eq!(
870            aliased_columns_for(&[qcol("t.pk1")].into(), &plan, Some(&scan_plan)).unwrap(),
871            [(qcol("t.pk1"), [qcol("t.pk1")].into())].into()
872        );
873
874        // Test aliased_columns_for from scan to first aggregate
875        assert_eq!(
876            aliased_columns_for(&[qcol("t.pk1")].into(), &first_aggr, Some(&scan_plan)).unwrap(),
877            [(qcol("t.pk1"), [qcol("t.pk1")].into())].into()
878        );
879
880        // Test aliased_columns_for from first aggregate to final plan
881        assert_eq!(
882            aliased_columns_for(&[qcol("t.pk1")].into(), &plan, Some(&first_aggr)).unwrap(),
883            [(qcol("t.pk1"), [qcol("t.pk1")].into())].into()
884        );
885
886        // Test aliased_columns_for from leaf to final plan
887        assert_eq!(
888            aliased_columns_for(&[Column::from_name("pk1")].into(), &plan, None).unwrap(),
889            [(Column::from_name("pk1"), [qcol("t.pk1")].into())].into()
890        );
891    }
892
893    #[test]
894    fn aggr_aggr_alias_projection() {
895        // use logging for better debugging
896        init_default_ut_logging();
897        let test_table = TestTable::table_with_name(0, "t".to_string());
898        let table_source = Arc::new(DefaultTableSource::new(Arc::new(
899            DfTableProviderAdapter::new(test_table),
900        )));
901        let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![])
902            .unwrap()
903            .aggregate(vec![col("pk1"), col("pk2")], vec![max(col("number"))])
904            .unwrap()
905            .aggregate(
906                vec![col("pk1"), col("pk2")],
907                vec![min(col("max(t.number)"))],
908            )
909            .unwrap()
910            .project(vec![
911                col("pk1").alias("pk11"),
912                col("pk2").alias("pk22"),
913                col("min(max(t.number))").alias("min_max_number"),
914            ])
915            .unwrap()
916            .build()
917            .unwrap();
918
919        let proj_plan = plan.clone();
920        let second_aggr = proj_plan.inputs()[0].clone();
921
922        // Test original_column_for from projection to second aggregate for aggr gen column
923        assert_eq!(
924            original_column_for(
925                &[Column::from_name("min_max_number")].into(),
926                plan.clone(),
927                Some(Arc::new(second_aggr.clone()))
928            )
929            .unwrap(),
930            [(
931                Column::from_name("min_max_number"),
932                Column::from_name("min(max(t.number))")
933            )]
934            .into()
935        );
936
937        // Test aliased_columns_for from second aggregate to projection
938        assert_eq!(
939            aliased_columns_for(
940                &[Column::from_name("min(max(t.number))")].into(),
941                &plan,
942                Some(&second_aggr)
943            )
944            .unwrap(),
945            [(
946                Column::from_name("min(max(t.number))"),
947                [Column::from_name("min_max_number")].into()
948            )]
949            .into()
950        );
951    }
952}
query/dist_plan/analyzer/utils.rs

query/dist_plan/analyzer/
utils.rs
