WIP: Use ColumnProjection::has_projection, avoid some spurious projection nodes, clean up logic

Author: srh

rust/cubestore/cubestore/src/queryplanner/planning.rs

@@ -2058,9 +2058,8 @@ pub mod tests {
 
         assert_eq!(
             pretty_printers::pp_plan(&plan),
-            "Projection, [s.orders.order_customer:order_customer, sum(s.orders.order_amount)]\
-            \n  ClusterAggregateTopK, limit: 10\
-            \n    Scan s.orders, source: CubeTable(index: by_customer:3:[]:sort_on[order_customer]), fields: [order_customer, order_amount]"
+            "ClusterAggregateTopK, limit: 10\
+            \n  Scan s.orders, source: CubeTable(index: by_customer:3:[]:sort_on[order_customer]), fields: [order_customer, order_amount]"
         );
 
         // Projections should be handled properly.

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

@@ -83,8 +83,11 @@ fn materialize_topk_under_limit_sort(
     sort_input: &Arc<LogicalPlan>,
 ) -> Result<Option<LogicalPlan>, DataFusionError> {
     let projection = extract_projections_and_havings(&sort_input)?;
+    let Some(projection) = projection else {
+        return Ok(None);
+    };
 
-    let aggregate = projection.as_ref().map(|p| p.input).unwrap_or(sort_input);
+    let aggregate: &Arc<LogicalPlan> = projection.input;
     match aggregate.as_ref() {
         LogicalPlan::Aggregate(Aggregate {
             input: cluster_send,
@@ -109,7 +112,7 @@ fn materialize_topk_under_limit_sort(
                 group_expr.len(),
                 &sort_expr,
                 sort_input.schema(),
-                projection.as_ref().map(|c| c.input_columns.as_slice()),
+                projection.input_columns.as_slice(),
             )? {
                 sort_columns = sc;
             } else {
@@ -130,17 +133,13 @@ fn materialize_topk_under_limit_sort(
                             group_expr: group_expr.clone(),
                             aggregate_expr: aggr_expr.clone(),
                             order_by: sort_columns,
-                            having_expr: projection
-                                .as_ref()
-                                .map_or(None, |p| p.having_expr.clone()),
+                            having_expr: projection.having_expr.clone(),
                             schema: aggregate_schema.clone(),
                             snapshots: cs.snapshots.clone(),
                         }),
                     });
-                    // TODO upgrade DF: Projection should always be Some, now.  We are also in some
-                    // cases (such as when there was no projection node) creating a spurious
-                    // projection.
-                    if let Some(p) = projection {
+                    if projection.has_projection {
+                        let p = projection;
                         let out_schema = p.schema;
                         let mut expr = Vec::with_capacity(p.input_columns.len());
                         for out_i in 0..p.input_columns.len() {
@@ -274,6 +273,9 @@ struct ColumnProjection<'a> {
     post_projection: Vec<Expr>,
     // Defined on `input` schema
     having_expr: Option<Expr>,
+    // True if there is some sort of projection seen.
+    has_projection: bool,
+
 }
 
 fn extract_projections_and_havings(
@@ -372,6 +374,7 @@ fn extract_projections_and_havings(
                 schema,
                 post_projection: new_post_projection,
                 having_expr: inner_column_projection.having_expr,
+                has_projection: true,
             };
 
             return Ok(Some(column_projection));
@@ -423,6 +426,7 @@ fn extract_projections_and_havings(
                         transformed_predicate
                     },
                 ),
+                has_projection: inner_column_projection.has_projection,
             };
 
             return Ok(Some(column_projection));
@@ -444,6 +448,7 @@ fn extract_projections_and_havings(
                 schema: in_schema,
                 post_projection,
                 having_expr: None,
+                has_projection: false,
             };
             return Ok(Some(column_projection));
         }
@@ -454,7 +459,7 @@ fn extract_sort_columns(
     group_key_len: usize,
     sort_expr: &[SortExpr],
     schema: &DFSchema,
-    projection: Option<&[usize]>,
+    projection: &[usize],
 ) -> Result<Option<Vec<SortColumn>>, DataFusionError> {
     let mut sort_columns = Vec::with_capacity(sort_expr.len());
     for e in sort_expr {
@@ -466,9 +471,7 @@ fn extract_sort_columns(
         match expr {
             Expr::Column(c) => {
                 let mut index = field_index(schema, c.relation.as_ref(), &c.name)?;
-                if let Some(p) = projection {
-                    index = p[index];
-                }
+                index = projection[index];
                 if index < group_key_len {
                     return Ok(None);
                 }