revert change recompute_schema for LogicalPlan::Values

Author: dqkqd

datafusion/core/tests/execution/logical_plan.rs

@@ -22,35 +22,45 @@ use arrow::array::Int64Array;
 use arrow::datatypes::{DataType, Field, Schema};
 use datafusion::datasource::{provider_as_source, ViewTable};
 use datafusion::execution::session_state::SessionStateBuilder;
-use datafusion_common::{DFSchema, DFSchemaRef, Result, ScalarValue};
+use datafusion_common::{Column, DFSchema, DFSchemaRef, Result, ScalarValue, Spans};
 use datafusion_execution::TaskContext;
 use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams};
-use datafusion_expr::logical_plan::LogicalPlan;
+use datafusion_expr::logical_plan::{LogicalPlan, Values};
 use datafusion_expr::{
-    col, Aggregate, AggregateUDF, EmptyRelation, Expr, LogicalPlanBuilder, UNNAMED_TABLE,
+    Aggregate, AggregateUDF, EmptyRelation, Expr, LogicalPlanBuilder, UNNAMED_TABLE,
 };
 use datafusion_functions_aggregate::count::Count;
 use datafusion_physical_plan::collect;
 use insta::assert_snapshot;
+use std::collections::HashMap;
 use std::fmt::Debug;
 use std::ops::Deref;
 use std::sync::Arc;
 
 #[tokio::test]
 async fn count_only_nulls() -> Result<()> {
     // Input: VALUES (NULL), (NULL), (NULL) AS _(col)
-    let input = LogicalPlanBuilder::values(vec![
-        vec![Expr::Literal(ScalarValue::Null, None)],
-        vec![Expr::Literal(ScalarValue::Null, None)],
-        vec![Expr::Literal(ScalarValue::Null, None)],
-    ])?
-    .project(vec![col("column1").alias("col")])?
-    .build()?;
-    let input_col_ref = col("col");
+    let input_schema = Arc::new(DFSchema::from_unqualified_fields(
+        vec![Field::new("col", DataType::Null, true)].into(),
+        HashMap::new(),
+    )?);
+    let input = Arc::new(LogicalPlan::Values(Values {
+        schema: input_schema,
+        values: vec![
+            vec![Expr::Literal(ScalarValue::Null, None)],
+            vec![Expr::Literal(ScalarValue::Null, None)],
+            vec![Expr::Literal(ScalarValue::Null, None)],
+        ],
+    }));
+    let input_col_ref = Expr::Column(Column {
+        relation: None,
+        name: "col".to_string(),
+        spans: Spans::new(),
+    });
 
     // Aggregation: count(col) AS count
     let aggregate = LogicalPlan::Aggregate(Aggregate::try_new(
-        input.into(),
+        input,
         vec![],
         vec![Expr::AggregateFunction(AggregateFunction {
             func: Arc::new(AggregateUDF::new_from_impl(Count::new())),

datafusion/expr/src/logical_plan/plan.rs

@@ -632,8 +632,9 @@ impl LogicalPlan {
             }) => Projection::try_new(expr, input).map(LogicalPlan::Projection),
             LogicalPlan::Dml(_) => Ok(self),
             LogicalPlan::Copy(_) => Ok(self),
-            LogicalPlan::Values(Values { schema: _, values }) => {
-                LogicalPlanBuilder::values(values)?.build()
+            LogicalPlan::Values(Values { schema, values }) => {
+                // todo it isn't clear why the schema is not recomputed here
+                Ok(LogicalPlan::Values(Values { schema, values }))
             }
             LogicalPlan::Filter(Filter { predicate, input }) => {
                 Filter::try_new(predicate, input).map(LogicalPlan::Filter)
@@ -1473,7 +1474,13 @@ impl LogicalPlan {
             })?
             // always recompute the schema to ensure the changed in the schema's field should be
             // poplulated to the plan's parent
-            .map_data(|plan| plan.recompute_schema())
+            .map_data(|plan| plan.recompute_schema())?
+            .map_data(|plan| match plan {
+                LogicalPlan::Values(Values { values, schema: _ }) => {
+                    LogicalPlanBuilder::values(values)?.build()
+                }
+                _ => Ok(plan),
+            })
         })
         .map(|res| res.data)
     }