Fix incompatible type in IS NOT DISTINCT filter for PIVOT (#25)

* Fix incompattible type in IS NOT DISTINCT filter

* Cargo fmt

* Cargo clippy

Author: Vedin

datafusion/core/src/physical_planner.rs

@@ -76,8 +76,8 @@ use datafusion_expr::expr::{
 use datafusion_expr::expr_rewriter::unnormalize_cols;
 use datafusion_expr::logical_plan::builder::wrap_projection_for_join_if_necessary;
 use datafusion_expr::{
-    Analyze, BinaryExpr, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension,
-    FetchType, Filter, JoinType, LogicalPlanBuilder, RecursiveQuery, SkipType,
+    Analyze, BinaryExpr, Cast, DescribeTable, DmlStatement, Explain, ExplainFormat,
+    Extension, FetchType, Filter, JoinType, LogicalPlanBuilder, RecursiveQuery, SkipType,
     StringifiedPlan, WindowFrame, WindowFrameBound, WriteOp,
 };
 use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
@@ -1773,16 +1773,31 @@ pub fn transform_pivot_to_aggregate(
         .map(|col: datafusion_common::Column| Expr::Column(col))
         .collect();
 
-    let builder = LogicalPlanBuilder::from(Arc::unwrap_or_clone(input));
+    let builder = LogicalPlanBuilder::from(Arc::unwrap_or_clone(Arc::clone(&input)));
 
     // Create the aggregate plan with filtered aggregates
     let mut aggregate_exprs = Vec::new();
 
+    let input_schema = input.schema();
+    let pivot_col_idx = match input_schema.index_of_column(pivot_column) {
+        Ok(idx) => idx,
+        Err(_) => {
+            return plan_err!(
+                "Pivot column '{}' does not exist in input schema",
+                pivot_column
+            )
+        }
+    };
+    let pivot_col_type = input_schema.field(pivot_col_idx).data_type();
+
     for value in &pivot_values {
         let filter_condition = Expr::BinaryExpr(BinaryExpr::new(
             Box::new(Expr::Column(pivot_column.clone())),
             Operator::IsNotDistinctFrom,
-            Box::new(Expr::Literal(value.clone())),
+            Box::new(Expr::Cast(Cast::new(
+                Box::new(Expr::Literal(value.clone())),
+                pivot_col_type.clone(),
+            ))),
         ));
 
         let filtered_agg = match aggregate_expr {

datafusion/sqllogictest/test_files/pivot.slt

@@ -68,6 +68,19 @@ ORDER BY empid;
 2 39500 90700
 3 1001 1001
 
+# PIVOT with cast to pivot column type
+query TIII
+SELECT *
+FROM quarterly_sales
+PIVOT(SUM(amount) FOR empid IN (1,2,3))
+ORDER BY quarter;
+----
+2023_Q1 10400 39500 NULL
+2023_Q2 8000 90700 NULL
+2023_Q3 11000 12000 2700
+2023_Q4 18000 5300 28900
+
+
 # PIVOT with automatic detection of all distinct column values using ANY
 query TIII
 SELECT *
@@ -119,7 +132,7 @@ ORDER BY empid;
 
 
 # Non-existent column in the FOR clause
-query error DataFusion error: Schema error: No field named non_existent_column\. Valid fields are quarterly_sales\.empid, quarterly_sales\.amount, quarterly_sales\.quarter\.
+query error DataFusion error: Error during planning: Pivot column 'non_existent_column' does not exist in input schema
 SELECT *
 FROM quarterly_sales
 PIVOT(SUM(amount) FOR non_existent_column IN ('2023_Q1', '2023_Q2'))