feat: carry source span on WindowFunction for diagnostics

Window aggregate / window UDF calls (e.g. `sum(first_name) OVER ()`)
previously passed `None` as the span into `verify_function_arguments`, so
the resulting diagnostic could not point at the SQL call site.

Add a `spans` field to `WindowFunction` (mirroring `ScalarFunction` /
`AggregateFunction`), populate it in `SqlToRel`, and thread it through
`ExprSchemable` so window function argument diagnostics carry the span.
Add a regression test for the window function path.

Addresses PR review feedback.

Author: choplin

datafusion-examples/examples/udf/advanced_udwf.rs

@@ -201,6 +201,7 @@ impl WindowUDFImpl for SimplifySmoothItUdf {
                     distinct: window_function.params.distinct,
                     filter: window_function.params.filter,
                 },
+                spans: Default::default(),
             }))
         };
 

datafusion/core/src/physical_planner.rs

@@ -2439,6 +2439,7 @@ pub fn create_window_expr_with_name(
                         distinct,
                         filter,
                     },
+                ..
             } = window_fun.as_ref();
             let physical_args =
                 create_physical_exprs(args, logical_schema, execution_props)?;

datafusion/expr/src/expr.rs

@@ -1245,6 +1245,8 @@ pub struct WindowFunction {
     /// Name of the function
     pub fun: WindowFunctionDefinition,
     pub params: WindowFunctionParams,
+    /// Original source code location, if known
+    pub spans: Spans,
 }
 
 #[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
@@ -1280,9 +1282,15 @@ impl WindowFunction {
                 null_treatment: None,
                 distinct: false,
             },
+            spans: Spans::new(),
         }
     }
 
+    /// Returns a mutable reference to the spans
+    pub fn spans_mut(&mut self) -> &mut Spans {
+        &mut self.spans
+    }
+
     /// Returns this window function's simplification hook, if any.
     ///
     /// See [`WindowFunctionSimplification`] for more information
@@ -2308,6 +2316,7 @@ impl Expr {
             Expr::Not(inner) | Expr::Negative(inner) => inner.spans(),
             Expr::ScalarFunction(func) => Some(&func.spans),
             Expr::AggregateFunction(func) => Some(&func.spans),
+            Expr::WindowFunction(func) => Some(&func.spans),
             _ => None,
         }
     }
@@ -2582,6 +2591,7 @@ impl NormalizeEq for Expr {
                             null_treatment: self_null_treatment,
                             distinct: self_distinct,
                         },
+                    ..
                 } = left.as_ref();
                 let WindowFunction {
                     fun: other_fun,
@@ -2595,6 +2605,7 @@ impl NormalizeEq for Expr {
                             null_treatment: other_null_treatment,
                             distinct: other_distinct,
                         },
+                    ..
                 } = other.as_ref();
 
                 self_fun.name() == other_fun.name()
@@ -2856,6 +2867,7 @@ impl HashNode for Expr {
                             null_treatment,
                             distinct,
                         },
+                    ..
                 } = window_fun.as_ref();
                 fun.hash(state);
                 window_frame.hash(state);
@@ -3175,7 +3187,7 @@ impl Display for SchemaDisplay<'_> {
                 Ok(())
             }
             Expr::WindowFunction(window_fun) => {
-                let WindowFunction { fun, params } = window_fun.as_ref();
+                let WindowFunction { fun, params, .. } = window_fun.as_ref();
                 match fun {
                     WindowFunctionDefinition::AggregateUDF(fun) => {
                         match fun.window_function_schema_name(params) {
@@ -3615,7 +3627,7 @@ impl Display for Expr {
                 fmt_function(f, fun.name(), false, &fun.args, true)
             }
             Expr::WindowFunction(window_fun) => {
-                let WindowFunction { fun, params } = window_fun.as_ref();
+                let WindowFunction { fun, params, .. } = window_fun.as_ref();
                 match fun {
                     WindowFunctionDefinition::AggregateUDF(fun) => {
                         match fun.window_function_display_name(params) {

datafusion/expr/src/expr_schema.rs

@@ -539,7 +539,7 @@ impl ExprSchemable for Expr {
                 let WindowFunction {
                     fun,
                     params: WindowFunctionParams { args, .. },
-                    ..
+                    spans,
                 } = window_function.as_ref();
 
                 let fields = args
@@ -548,14 +548,20 @@ impl ExprSchemable for Expr {
                     .collect::<Result<Vec<_>>>()?;
                 match fun {
                     WindowFunctionDefinition::AggregateUDF(udaf) => {
-                        let new_fields =
-                            verify_function_arguments(udaf.as_ref(), &fields, None)?;
+                        let new_fields = verify_function_arguments(
+                            udaf.as_ref(),
+                            &fields,
+                            spans.first(),
+                        )?;
                         let return_field = udaf.return_field(&new_fields)?;
                         Ok(return_field)
                     }
                     WindowFunctionDefinition::WindowUDF(udwf) => {
-                        let new_fields =
-                            verify_function_arguments(udwf.as_ref(), &fields, None)?;
+                        let new_fields = verify_function_arguments(
+                            udwf.as_ref(),
+                            &fields,
+                            spans.first(),
+                        )?;
                         let return_field = udwf
                             .field(WindowUDFFieldArgs::new(&new_fields, &schema_name))?;
                         Ok(return_field)

datafusion/expr/src/logical_plan/plan.rs

@@ -2809,6 +2809,7 @@ impl Window {
                 let WindowFunction {
                     fun: WindowFunctionDefinition::WindowUDF(udwf),
                     params: WindowFunctionParams { partition_by, .. },
+                    ..
                 } = window_fun.as_ref()
                 else {
                     return None;

datafusion/expr/src/tree_node.rs

@@ -275,6 +275,7 @@ impl TreeNode for Expr {
                             null_treatment,
                             distinct,
                         },
+                    spans,
                 } = *window_fun;
 
                 (args, partition_by, order_by, filter)
@@ -292,6 +293,7 @@ impl TreeNode for Expr {
                                     null_treatment,
                                     distinct,
                                 },
+                                spans,
                             }))
                         },
                     )?

datafusion/functions-window/src/planner.rs

@@ -17,7 +17,7 @@
 
 //! SQL planning extensions like [`WindowFunctionPlanner`]
 
-use datafusion_common::Result;
+use datafusion_common::{Result, Spans};
 use datafusion_expr::{
     Expr,
     expr::{WindowFunction, WindowFunctionParams},
@@ -56,6 +56,7 @@ impl ExprPlanner for WindowFunctionPlanner {
                 null_treatment,
                 distinct,
             },
+            spans: Spans::new(),
         });
 
         let saved_name = NamePreserver::new_for_projection().save(&origin_expr);
@@ -75,6 +76,7 @@ impl ExprPlanner for WindowFunctionPlanner {
                     distinct,
                     filter,
                 },
+            ..
         } = *window_fun;
         let raw_expr = RawWindowExpr {
             func_def: fun,
@@ -116,6 +118,7 @@ impl ExprPlanner for WindowFunctionPlanner {
                     null_treatment,
                     distinct,
                 },
+                spans: Spans::new(),
             });
             let new_expr = saved_name.restore(new_expr);
 

datafusion/optimizer/src/analyzer/type_coercion.rs

@@ -30,9 +30,9 @@ use arrow::temporal_conversions::SECONDS_IN_DAY;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter};
 use datafusion_common::{
-    Column, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, TableReference,
-    exec_err, internal_datafusion_err, internal_err, not_impl_err, plan_datafusion_err,
-    plan_err,
+    Column, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, Spans,
+    TableReference, exec_err, internal_datafusion_err, internal_err, not_impl_err,
+    plan_datafusion_err, plan_err,
 };
 use datafusion_expr::expr::{
     self, AggregateFunctionParams, Alias, Between, BinaryExpr, Case, Exists,
@@ -746,6 +746,7 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> {
                             null_treatment,
                             distinct,
                         },
+                    ..
                 } = *window_fun;
                 let window_frame =
                     coerce_window_frame(window_frame, self.schema, &order_by)?;
@@ -775,6 +776,7 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> {
                         null_treatment,
                         distinct,
                     },
+                    spans: Spans::new(),
                 });
                 Ok(Transformed::yes(new_expr))
             }

datafusion/proto/src/logical_plan/to_proto.rs

@@ -324,6 +324,7 @@ pub fn serialize_expr(
                         distinct,
                         filter,
                     },
+                ..
             } = window_fun.as_ref();
             let mut buf = Vec::new();
             let window_function = match fun {

datafusion/sql/src/expr/function.rs

@@ -19,7 +19,7 @@ use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
 
 use arrow::datatypes::{DataType, FieldRef};
 use datafusion_common::{
-    DFSchema, Dependency, Diagnostic, HashSet, Result, Span, datatype::FieldExt,
+    DFSchema, Dependency, Diagnostic, HashSet, Result, Span, Spans, datatype::FieldExt,
     internal_datafusion_err, internal_err, not_impl_err, plan_datafusion_err, plan_err,
 };
 use datafusion_expr::{
@@ -682,7 +682,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                     distinct,
                 } = window_expr;
 
-                let inner = WindowFunction {
+                let mut inner = WindowFunction {
                     fun: func_def,
                     params: expr::WindowFunctionParams {
                         args,
@@ -693,7 +693,13 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                         null_treatment,
                         distinct,
                     },
+                    spans: Spans::new(),
                 };
+                if self.options.collect_spans
+                    && let Some(span) = Span::try_from_sqlparser_span(sql_parser_span)
+                {
+                    inner.spans_mut().add_span(span);
+                }
 
                 if name.eq_ignore_ascii_case(inner.fun.name()) {
                     return Ok(Expr::WindowFunction(Box::new(inner)));