Embucket changes

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: JanKaul

Cargo.lock

@@ -188,7 +188,9 @@ regex = "1.12"
 rstest = "0.26.1"
 serde_json = "1"
 sha2 = "^0.10.9"
-sqlparser = { version = "0.61.0", default-features = false, features = ["std", "visitor"] }
+sqlparser = { git = "https://github.com/Embucket/datafusion-sqlparser-rs.git", branch = "embucket-sync-df53.0.0-parser0.61.0", features = [
+    "visitor",
+] }
 strum = "0.28.0"
 strum_macros = "0.28.0"
 tempfile = "3"

Cargo.toml

@@ -324,6 +324,10 @@ pub struct ParquetMetadataFunc {}
 impl TableFunctionImpl for ParquetMetadataFunc {
     fn call_with_args(&self, args: TableFunctionArgs) -> Result<Arc<dyn TableProvider>> {
         let exprs = args.exprs();
+        if exprs.is_empty() {
+            return plan_err!("parquet_metadata requires string argument as its input");
+        }
+
         let filename = match exprs.first() {
             Some(Expr::Literal(ScalarValue::Utf8(Some(s)), _)) => s, // single quote: parquet_metadata('x.parquet')
             Some(Expr::Column(Column { name, .. })) => name, // double quote: parquet_metadata("x.parquet")

FETCH_HEAD

@@ -319,7 +319,9 @@ impl Session for SessionState {
 }
 
 impl SessionState {
-    pub(crate) fn resolve_table_ref(
+    /// Resolve a [`TableReference`] into a [`ResolvedTableReference`] using
+    /// the session's configured default catalog and schema.
+    pub fn resolve_table_ref(
         &self,
         table_ref: impl Into<TableReference>,
     ) -> ResolvedTableReference {
@@ -541,8 +543,9 @@ impl SessionState {
         query.statement_to_plan(statement)
     }
 
+    /// Get the parser options
     #[cfg(feature = "sql")]
-    fn get_parser_options(&self) -> ParserOptions {
+    pub fn get_parser_options(&self) -> ParserOptions {
         let sql_parser_options = &self.config.options().sql_parser;
 
         ParserOptions {
@@ -1864,9 +1867,11 @@ impl From<SessionState> for SessionStateBuilder {
 /// This is used so the SQL planner can access the state of the session without
 /// having a direct dependency on the [`SessionState`] struct (and core crate)
 #[cfg(feature = "sql")]
-struct SessionContextProvider<'a> {
-    state: &'a SessionState,
-    tables: HashMap<ResolvedTableReference, Arc<dyn TableSource>>,
+pub struct SessionContextProvider<'a> {
+    /// The session state
+    pub state: &'a SessionState,
+    /// The tables available in the session
+    pub tables: HashMap<ResolvedTableReference, Arc<dyn TableSource>>,
 }
 
 #[cfg(feature = "sql")]
@@ -1901,14 +1906,15 @@ impl ContextProvider for SessionContextProvider<'_> {
     fn get_table_function_source(
         &self,
         name: &str,
-        args: Vec<Expr>,
+        args: Vec<(Expr, Option<String>)>,
     ) -> datafusion_common::Result<Arc<dyn TableSource>> {
         use datafusion_catalog::TableFunctionArgs;
 
+        let name = name.to_ascii_lowercase();
         let tbl_func = self
             .state
             .table_functions
-            .get(name)
+            .get(&name)
             .cloned()
             .ok_or_else(|| plan_datafusion_err!("table function '{name}' not found"))?;
         let simplify_context = SimplifyContext::builder()
@@ -1919,11 +1925,11 @@ impl ContextProvider for SessionContextProvider<'_> {
             .build();
         let simplifier = ExprSimplifier::new(simplify_context);
         let schema = DFSchema::empty();
-        let args = args
+        let args: Vec<Expr> = args
             .into_iter()
-            .map(|arg| {
+            .map(|(expr, _named_param)| {
                 simplifier
-                    .coerce(arg, &schema)
+                    .coerce(expr, &schema)
                     .and_then(|e| simplifier.simplify(e))
             })
             .collect::<datafusion_common::Result<Vec<_>>>()?;

datafusion-cli/src/functions.rs

@@ -86,9 +86,9 @@ use datafusion_expr::expr_rewriter::unnormalize_cols;
 use datafusion_expr::logical_plan::builder::wrap_projection_for_join_if_necessary;
 use datafusion_expr::utils::{expr_to_columns, split_conjunction};
 use datafusion_expr::{
-    Analyze, BinaryExpr, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension,
-    FetchType, Filter, JoinType, Operator, RecursiveQuery, SkipType, StringifiedPlan,
-    WindowFrame, WindowFrameBound, WriteOp,
+    Analyze, BinaryExpr, Cast, DescribeTable, DmlStatement, Explain, ExplainFormat,
+    Extension, FetchType, Filter, JoinType, LogicalPlanBuilder, Operator, RecursiveQuery,
+    SkipType, StringifiedPlan, WindowFrame, WindowFrameBound, WriteOp,
 };
 use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
 use datafusion_physical_expr::expressions::Literal;
@@ -110,6 +110,8 @@ use itertools::{Itertools, multiunzip};
 use log::debug;
 use tokio::sync::Mutex;
 
+use datafusion_physical_plan::collect;
+
 /// Physical query planner that converts a `LogicalPlan` to an
 /// `ExecutionPlan` suitable for execution.
 #[async_trait]
@@ -1253,7 +1255,62 @@ impl DefaultPhysicalPlanner {
                     options.clone(),
                 )?)
             }
+            LogicalPlan::Pivot(pivot) => {
+                return if !pivot.pivot_values.is_empty() {
+                    let input = Arc::new(pivot.input.as_ref().clone());
+                    let agg_plan = transform_pivot_to_aggregate(
+                        &input,
+                        &pivot.aggregate_expr,
+                        &pivot.pivot_column,
+                        &pivot.pivot_values,
+                        pivot.default_on_null_expr.as_ref(),
+                    )?;
+
+                    self.create_physical_plan(&agg_plan, session_state).await
+                } else if let Some(subquery) = &pivot.value_subquery {
+                    let optimized_subquery = session_state.optimize(subquery.as_ref())?;
+
+                    let subquery_physical_plan = self
+                        .create_physical_plan(&optimized_subquery, session_state)
+                        .await?;
 
+                    let subquery_results = collect(
+                        Arc::clone(&subquery_physical_plan),
+                        session_state.task_ctx(),
+                    )
+                    .await?;
+
+                    let mut pivot_values = Vec::new();
+                    for batch in subquery_results.iter() {
+                        if batch.num_columns() != 1 {
+                            return plan_err!(
+                                "Pivot subquery must return a single column"
+                            );
+                        }
+
+                        let column = batch.column(0);
+                        for row_idx in 0..batch.num_rows() {
+                            if !column.is_null(row_idx) {
+                                pivot_values
+                                    .push(ScalarValue::try_from_array(column, row_idx)?);
+                            }
+                        }
+                    }
+
+                    let input = Arc::new(pivot.input.as_ref().clone());
+                    let agg_plan = transform_pivot_to_aggregate(
+                        &input,
+                        &pivot.aggregate_expr,
+                        &pivot.pivot_column,
+                        &pivot_values,
+                        pivot.default_on_null_expr.as_ref(),
+                    )?;
+
+                    self.create_physical_plan(&agg_plan, session_state).await
+                } else {
+                    plan_err!("PIVOT operation requires at least one value to pivot on")
+                };
+            }
             // 2 Children
             LogicalPlan::Join(Join {
                 left: original_left,
@@ -2142,7 +2199,8 @@ fn extract_dml_filters(
             | LogicalPlan::Ddl(_)
             | LogicalPlan::Copy(_)
             | LogicalPlan::Unnest(_)
-            | LogicalPlan::RecursiveQuery(_) => {
+            | LogicalPlan::RecursiveQuery(_)
+            | LogicalPlan::Pivot(_) => {
                 // No filters to extract from leaf/meta plans
             }
             // Plans with inputs (may contain filters in children)
@@ -2498,6 +2556,148 @@ pub fn create_aggregate_expr_and_maybe_filter(
     )
 }
 
+/// Transform a PIVOT operation into a more standard Aggregate + Projection plan
+/// For known pivot values, we create a projection that includes "IS NOT DISTINCT FROM" conditions
+///
+/// For example, for SUM(amount) PIVOT(quarter FOR quarter in ('2023_Q1', '2023_Q2')), we create:
+/// - SUM(amount) FILTER (WHERE quarter IS NOT DISTINCT FROM '2023_Q1') AS "2023_Q1"
+/// - SUM(amount) FILTER (WHERE quarter IS NOT DISTINCT FROM '2023_Q2') AS "2023_Q2"
+///
+/// If DEFAULT ON NULL is specified, each aggregate expression is wrapped with an outer projection that
+/// applies COALESCE to the results.
+pub fn transform_pivot_to_aggregate(
+    input: &Arc<LogicalPlan>,
+    aggregate_expr: &Expr,
+    pivot_column: &Column,
+    pivot_values: &[ScalarValue],
+    default_on_null_expr: Option<&Expr>,
+) -> Result<LogicalPlan> {
+    let df_schema = input.schema();
+
+    let all_columns: Vec<Column> = df_schema.columns();
+
+    // Filter to include only columns we want for GROUP BY
+    // (exclude pivot column and aggregate expression columns)
+    let group_by_columns: Vec<Expr> = all_columns
+        .into_iter()
+        .filter(|col: &Column| {
+            col.name != pivot_column.name
+                && !aggregate_expr
+                    .column_refs()
+                    .iter()
+                    .any(|agg_col| agg_col.name == col.name)
+        })
+        .map(|col: Column| Expr::Column(col))
+        .collect();
+
+    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::Cast(Cast::new(
+                Box::new(Expr::Literal(value.clone(), None)),
+                pivot_col_type.clone(),
+            ))),
+        ));
+
+        let filtered_agg = match aggregate_expr {
+            Expr::AggregateFunction(agg) => {
+                let mut new_params = agg.params.clone();
+                new_params.filter = Some(Box::new(filter_condition));
+                Expr::AggregateFunction(AggregateFunction {
+                    func: Arc::clone(&agg.func),
+                    params: new_params,
+                })
+            }
+            _ => {
+                return plan_err!(
+                    "Unsupported aggregate expression should always be AggregateFunction"
+                );
+            }
+        };
+
+        // Use the pivot value as the column name
+        let field_name = value.to_string().trim_matches('\'').to_string();
+        let aliased_agg = Expr::Alias(Alias {
+            expr: Box::new(filtered_agg),
+            relation: None,
+            name: field_name,
+            metadata: None,
+        });
+
+        aggregate_exprs.push(aliased_agg);
+    }
+
+    // Create the plan with the aggregate
+    let aggregate_plan = builder
+        .aggregate(group_by_columns, aggregate_exprs)?
+        .build()?;
+
+    // If DEFAULT ON NULL is specified, add a projection to apply COALESCE
+    if let Some(default_expr) = default_on_null_expr {
+        let schema = aggregate_plan.schema();
+        let mut projection_exprs = Vec::new();
+
+        for field in schema.fields() {
+            if !pivot_values
+                .iter()
+                .any(|v| field.name() == v.to_string().trim_matches('\''))
+            {
+                projection_exprs.push(Expr::Column(Column::from_name(field.name())));
+            }
+        }
+
+        // Apply COALESCE to aggregate columns
+        for value in pivot_values {
+            let field_name = value.to_string().trim_matches('\'').to_string();
+            let aggregate_col = Expr::Column(Column::from_name(&field_name));
+
+            // Create COALESCE expression using CASE: CASE WHEN col IS NULL THEN default_value ELSE col END
+            let coalesce_expr = Expr::Case(datafusion_expr::expr::Case {
+                expr: None,
+                when_then_expr: vec![(
+                    Box::new(Expr::IsNull(Box::new(aggregate_col.clone()))),
+                    Box::new(default_expr.clone()),
+                )],
+                else_expr: Some(Box::new(aggregate_col)),
+            });
+
+            let aliased_coalesce = Expr::Alias(Alias {
+                expr: Box::new(coalesce_expr),
+                relation: None,
+                name: field_name,
+                metadata: None,
+            });
+
+            projection_exprs.push(aliased_coalesce);
+        }
+
+        // Apply the projection
+        LogicalPlanBuilder::from(aggregate_plan)
+            .project(projection_exprs)?
+            .build()
+    } else {
+        Ok(aggregate_plan)
+    }
+}
+
 impl DefaultPhysicalPlanner {
     /// Handles capturing the various plans for EXPLAIN queries
     ///
@@ -2887,6 +3087,32 @@ impl DefaultPhysicalPlanner {
             .collect::<Result<Vec<_>>>()?;
 
         let num_input_columns = input_exec.schema().fields().len();
+        // When we detect a PIVOT-derived plan with a value_subquery, ensure all generated columns are preserved
+        if let LogicalPlan::Pivot(pivot) = input.as_ref()
+            && pivot.value_subquery.is_some()
+            && let Some(agg_exec) = input_exec.downcast_ref::<AggregateExec>()
+        {
+            let schema = input_exec.schema();
+            let group_by_len = agg_exec.group_expr().expr().len();
+
+            if group_by_len < schema.fields().len() {
+                let mut all_exprs = physical_exprs.clone();
+
+                for (i, field) in schema.fields().iter().enumerate().skip(group_by_len) {
+                    if !physical_exprs.iter().any(|(_, name)| name == field.name()) {
+                        all_exprs.push((
+                            Arc::new(datafusion_physical_expr::expressions::Column::new(
+                                field.name(),
+                                i,
+                            )) as Arc<dyn PhysicalExpr>,
+                            field.name().to_string(),
+                        ));
+                    }
+                }
+
+                return Ok(Arc::new(ProjectionExec::try_new(all_exprs, input_exec)?));
+            }
+        }
 
         match self.try_plan_async_exprs(
             num_input_columns,