WIP: Make extract_projection_and_having handle multiple projections/havings.

Author: srh

rust/cubestore/cubestore/src/queryplanner/topk/plan.rs

@@ -8,6 +8,7 @@ use crate::queryplanner::udfs::{
 };
 use datafusion::arrow::compute::SortOptions;
 use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::common::tree_node::{Transformed, TreeNode};
 use datafusion::error::DataFusionError;
 use datafusion::execution::SessionState;
 use datafusion::logical_expr::expr::{AggregateFunction, Alias, ScalarFunction};
@@ -78,7 +79,7 @@ pub fn materialize_topk(p: LogicalPlan) -> Result<LogicalPlan, DataFusionError>
 
 /// Returns Ok(None) when materialization failed (without error) and the original plan should be returned.
 fn materialize_topk_under_limit_sort(fetch: usize, sort_expr: &Vec<SortExpr>, sort_input: &Arc<LogicalPlan>) -> Result<Option<LogicalPlan>, DataFusionError> {
-    let projection = extract_projection_and_having(&sort_input);
+    let projection = extract_projections_and_havings(&sort_input)?;
 
     let aggregate = projection.as_ref().map(|p| p.input).unwrap_or(sort_input);
     match aggregate.as_ref() {
@@ -238,79 +239,182 @@ fn extract_aggregate_fun(e: &Expr) -> Option<(TopKAggregateFunction, &Vec<Expr>)
 
 #[derive(Debug)]
 struct ColumnProjection<'a> {
+    // The (sole) column index within `input.schema()` that the post_projection expr uses.
     input_columns: Vec<usize>,
     input: &'a Arc<LogicalPlan>,
+    // Output schema (after applying `having_expr` and then `post_projection` and then aliases).  In
+    // other words, this saves the top level projection's aliases.
     schema: &'a DFSchemaRef,
+    // Defined on `input` schema.  Excludes Expr::Aliases necessary to produce the output schema, `schema`.
     post_projection: Vec<Expr>,
+    // Defined on `input` schema
     having_expr: Option<Expr>,
 }
 
-fn extract_having(p: &Arc<LogicalPlan>) -> (Option<Expr>, &Arc<LogicalPlan>) {
-    // Filter's "having" flag is not relevant to us.  It is used by DF to get the proper wildcard
-    // expansion behavior in the analysis pass (before LogicalPlan optimizations, and before we
-    // materialize the topk node here).
-    match p.as_ref() {
-        LogicalPlan::Filter(Filter { predicate, input, having: _, .. }) => (Some(predicate.clone()), input),
-        _ => (None, p),
-    }
-}
 
-fn extract_projection_and_having(p: &LogicalPlan) -> Option<ColumnProjection> {
-    match p {
+fn extract_projections_and_havings(p: &Arc<LogicalPlan>) -> Result<Option<ColumnProjection>, DataFusionError> {
+    // Goal:  Deal with arbitrary series of Projection and Filter, where the Projections are column
+    // projections (or cardinality(column)), on top of an underlying node.
+    //
+    // Real world example:  p = Projection > Filter > Projection > Aggregation
+    //
+    // Because the Sort node above p is defined in terms of the projection outputs, it needs those
+    // outputs remapped to projection inputs.
+
+    match p.as_ref() {
         LogicalPlan::Projection(Projection {
             expr,
             input,
             schema,
             ..
         }) => {
             let in_schema = input.schema();
-            let mut input_columns = Vec::with_capacity(expr.len());
-            let mut post_projection = Vec::with_capacity(expr.len());
+            let mut input_columns: Vec<usize> = Vec::with_capacity(expr.len());
+
+            // Check that this projection is a column (or cardinality(column)) projection first.
             for e in expr {
-                fn make_column(in_field_qualifier: Option<&TableReference>, in_field: &Field) -> datafusion::common::Column {
-                    datafusion::common::Column {
-                        relation: in_field_qualifier.map(|tr| tr.clone()),
-                        name: in_field.name().clone(),
-                    }
-                }
                 match e {
                     Expr::Alias(Alias { expr: box Expr::Column(c), relation: _, name: _ }) | Expr::Column(c) => {
-                        let fi = field_index(in_schema, c.relation.as_ref(), &c.name)?;
+                        // TODO upgrade DF: These field_index errors should return Err (right?) as
+                        // we do in the Filter case only because it's in a transform_up with map_data.
+                        let Some(fi) = field_index(in_schema, c.relation.as_ref(), &c.name) else {
+                            return Ok(None);
+                        };
                         input_columns.push(fi);
-                        let (in_field_qualifier, in_field) = in_schema.qualified_field(fi);
-                        post_projection.push(Expr::Column(make_column(in_field_qualifier, in_field)));
-                    }
-                    Expr::Alias(Alias { expr: box Expr::ScalarFunction(ScalarFunction { func, args }), relation: _, name: _ })
+                    },
+                    Expr::Alias(Alias { expr: box Expr::ScalarFunction(ScalarFunction { func, args }), relation: _, name: _})
                     | Expr::ScalarFunction(ScalarFunction { func, args }) => match func.name() {
-                        // TODO upgrade DF: use as_any() or something
+                        // TODO upgrade DF: Use as_any() or something
                         "cardinality" => match &args[0] {
                             Expr::Column(c) => {
-                                let fi = field_index(in_schema, c.relation.as_ref(), &c.name)?;
+                                let Some(fi) = field_index(in_schema, c.relation.as_ref(), &c.name) else {
+                                    return Ok(None);
+                                };
                                 input_columns.push(fi);
-                                let (in_field_qualifier, in_field) = in_schema.qualified_field(fi);
-                                post_projection.push(Expr::ScalarFunction(ScalarFunction {
-                                    func: scalar_udf_by_kind(CubeScalarUDFKind::HllCardinality),
-                                    args: vec![Expr::Column(make_column(in_field_qualifier, in_field))],
-                                }));
                             }
-                            _ => return None,
+                            _ => return Ok(None),
                         },
-                        _ => return None,
+                        _ => return Ok(None),
                     },
+                    _ => return Ok(None),
+                };
+            }
 
-                    _ => return None,
-                }
+            // Now recurse.
+            let inner_column_projection = extract_projections_and_havings(input)?;
+            let Some(inner_column_projection) = inner_column_projection else {
+                return Ok(None);
+            };
+
+            // Now apply our projection on top of the recursion
+
+            // input_columns[i] is the (sole) column number of `input.schema()` used by expr[i].
+            // inner_column_projection[j] is the (sole) column number of the presumed underlying `aggregate.schema()` used by inner expr j.
+            // So inner_column_projection[input_columns[i]] is the column number of the presumed underlying `aggregate.schema()` used by expr[i].
+
+            let mut deep_input_columns = Vec::with_capacity(expr.len());
+            for i in 0..expr.len() {
+                let j = input_columns[i];
+                deep_input_columns.push(inner_column_projection.input_columns[j]);
             }
-            let (having_expr, input) = extract_having(input);
-            Some(ColumnProjection {
-                input_columns,
-                input,
+
+            let mut new_post_projection = Vec::with_capacity(expr.len());
+
+            // And our projection's Column expressions need to be replaced with the inner post_projection expressions.
+            for (i, e) in expr.iter().enumerate() {
+                let new_e = e.clone().transform_up(|node| {
+                    node.unalias_nested().transform_data(|node|
+                        match node {
+                            Expr::Column(c) => {
+                                let replacement: Expr = inner_column_projection.post_projection[input_columns[i]].clone();
+                                // Transformed::yes/no doesn't matter here.
+                                // let unequal = &replacement != &node;
+                                Ok(Transformed::yes(replacement))
+                            },
+                            _ => Ok(Transformed::no(node))
+                        }
+                    )
+                })?;
+                new_post_projection.push(new_e.data);
+            }
+
+            let column_projection = ColumnProjection {
+                input_columns: deep_input_columns,
+                input: inner_column_projection.input,
                 schema,
+                post_projection: new_post_projection,
+                having_expr: inner_column_projection.having_expr,
+            };
+
+            return Ok(Some(column_projection));
+        },
+        LogicalPlan::Filter(Filter {
+            predicate,
+            input,
+            having: _,
+            ..
+        }) => {
+            // Filter's "having" flag is not relevant to us.  It is used by DF to get the proper wildcard
+            // expansion behavior in the analysis pass (before LogicalPlan optimizations, and before we
+            // materialize the topk node here).
+
+            // First, recurse.
+            let inner_column_projection = extract_projections_and_havings(input)?;
+            let Some(inner_column_projection) = inner_column_projection else {
+                return Ok(None);
+            };
+
+            let in_schema = input.schema();
+
+            // Our filter's columns, defined in terms of in_schema, need to be mapped to inner_column_projection.input.schema().
+            let transformed_predicate = predicate.clone().transform_up(|node| {
+                node.unalias_nested().transform_data(|node|
+                    match node {
+                        Expr::Column(c) => {
+                            let Some(fi) = field_index(in_schema, c.relation.as_ref(), &c.name) else {
+                                return Err(DataFusionError::Plan(format!("Field '{}' not found", &c.name)));  // TODO upgrade DF: spend brain cycles on error message
+                            };
+                            let replacement = inner_column_projection.post_projection[fi].clone();
+                            // Transformed::yes/no doesn't matter here.
+                            // let unequal = &replacement != &node;
+                            Ok(Transformed::yes(replacement))
+                        },
+                        _ => Ok(Transformed::no(node))
+                    }
+                )
+            })?.data;
+
+            let column_projection = ColumnProjection {
+                input_columns: inner_column_projection.input_columns,
+                input: inner_column_projection.input,
+                schema: inner_column_projection.schema,
+                post_projection: inner_column_projection.post_projection,
+                having_expr: Some(if let Some(previous_predicate) = inner_column_projection.having_expr {
+                    previous_predicate.and(transformed_predicate)
+                } else {
+                    transformed_predicate
+                }),
+            };
+
+            return Ok(Some(column_projection));
+        }
+        _ => {
+            let in_schema = p.schema();
+            let post_projection: Vec<Expr> = in_schema.iter().map(|(in_field_qualifier, in_field)|
+                Expr::Column(datafusion::common::Column {
+                    relation: in_field_qualifier.cloned(),
+                    name: in_field.name().clone(),
+                })
+            ).collect();
+            let column_projection = ColumnProjection {
+                input_columns: (0..post_projection.len()).collect(),
+                input: p,
+                schema: in_schema,
                 post_projection,
-                having_expr,
-            })
+                having_expr: None,
+            };
+            return Ok(Some(column_projection));
         }
-        _ => None,
     }
 }