Avoid changing expression names during constant folding

datafusion/src/optimizer/constant_folding.rs

@@ -92,6 +92,10 @@ impl OptimizerRule for ConstantFolding {
                     .expressions()
                     .into_iter()
                     .map(|e| {
+                        // We need to keep original expression name, if any.
+                        // Constant folding should not change expression name.
+                        let name = &e.name(plan.schema());
+
                         // TODO iterate until no changes are made
                         // during rewrite (evaluating constants can
                         // enable new simplifications and
@@ -101,7 +105,30 @@ impl OptimizerRule for ConstantFolding {
                             // fold constants and then simplify
                             .rewrite(&mut const_evaluator)?
                             .rewrite(&mut simplifier)?;
-                        Ok(new_e)
+
+                        let new_name = &new_e.name(plan.schema());
+
+                        // Some plans will be candidates in projection pushdown rule to
+                        // trim expressions based on expression names. We need to keep
+                        // expression name for them.
+                        let is_plan_for_projection_pushdown = matches!(
+                            plan,
+                            LogicalPlan::Window { .. }
+                                | LogicalPlan::Aggregate { .. }
+                                | LogicalPlan::Union { .. }
+                        );
+
+                        if let (Ok(expr_name), Ok(new_expr_name)) = (name, new_name) {
+                            if expr_name != new_expr_name
+                                && is_plan_for_projection_pushdown
+                            {
+                                Ok(new_e.alias(expr_name))
+                            } else {
+                                Ok(new_e)
+                            }
+                        } else {
+                            Ok(new_e)
+                        }
                     })
                     .collect::<Result<Vec<_>>>()?;
 
@@ -733,7 +760,7 @@ mod tests {
             .build()?;
 
         let expected = "\
-        Aggregate: groupBy=[[#test.a, #test.c]], aggr=[[MAX(#test.b), MIN(#test.b)]]\
+        Aggregate: groupBy=[[#test.a, #test.c]], aggr=[[MAX(#test.b) AS MAX(test.b = Boolean(true)), MIN(#test.b)]]\
         \n  Projection: #test.a, #test.c, #test.b\
         \n    TableScan: test projection=None";
 

datafusion/src/physical_plan/planner.rs

@@ -1342,6 +1342,17 @@ impl DefaultPhysicalPlanner {
         }
     }
 
+    fn remove_aliases<'a>(&self, expr: &'a Expr) -> Result<(String, &'a Expr)> {
+        let (name, e) = match expr {
+            Expr::Alias(sub_expr, alias) => match &**sub_expr {
+                Expr::Alias(nested_expr, _) => self.remove_aliases(&*nested_expr)?,
+                _ => (alias.clone(), sub_expr.as_ref()),
+            },
+            _ => (physical_name(expr)?, expr),
+        };
+        Ok((name, e))
+    }
+
     /// Create an aggregate expression from a logical expression or an alias
     pub fn create_aggregate_expr(
         &self,
@@ -1350,11 +1361,8 @@ impl DefaultPhysicalPlanner {
         physical_input_schema: &Schema,
         ctx_state: &ExecutionContextState,
     ) -> Result<Arc<dyn AggregateExpr>> {
-        // unpack aliased logical expressions, e.g. "sum(col) as total"
-        let (name, e) = match e {
-            Expr::Alias(sub_expr, alias) => (alias.clone(), sub_expr.as_ref()),
-            _ => (physical_name(e)?, e),
-        };
+        // unpack (nested) aliased logical expressions, e.g. "sum(col) as total"
+        let (name, e) = self.remove_aliases(e)?;
 
         self.create_aggregate_expr_with_name(
             e,

datafusion/tests/sql.rs

@@ -1282,6 +1282,22 @@ async fn csv_query_approx_count() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn query_count_without_from() -> Result<()> {
+    let mut ctx = ExecutionContext::new();
+    let sql = "SELECT count(1 + 1)";
+    let actual = execute_to_batches(&mut ctx, sql).await;
+    let expected = vec![
+        "+----------------------------+",
+        "| COUNT(Int64(1) + Int64(1)) |",
+        "+----------------------------+",
+        "| 1                          |",
+        "+----------------------------+",
+    ];
+    assert_batches_eq!(expected, &actual);
+    Ok(())
+}
+
 #[tokio::test]
 async fn csv_query_array_agg() -> Result<()> {
     let mut ctx = ExecutionContext::new();