feat: Propagate orderings through struct-producing projections (#21218)

## Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->

- Closes #21217

## What changes are included in this PR?

<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->

- Adds `ScalarUDFImpl::struct_field_mapping`
- Adds logic in `ProjectionMapping` to decompose struct-producing
functions into their field-level mapping entries so that orderings
propagate through struct projections
- Adds unit tests/SLT

## Are these changes tested?

Yes.

## Are there any user-facing changes?

N/A

---------

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Author: rkrishn7

datafusion/expr/src/lib.rs

@@ -126,7 +126,9 @@ pub use udaf::{
     udaf_default_schema_name, udaf_default_window_function_display_name,
     udaf_default_window_function_schema_name,
 };
-pub use udf::{ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl};
+pub use udf::{
+    ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, StructFieldMapping,
+};
 pub use udwf::{LimitEffect, ReversedUDWF, WindowUDF, WindowUDFImpl};
 pub use window_frame::{WindowFrame, WindowFrameBound, WindowFrameUnits};
 

datafusion/expr/src/udf.rs

@@ -38,6 +38,25 @@ use std::fmt::Debug;
 use std::hash::{Hash, Hasher};
 use std::sync::Arc;
 
+/// Describes how a struct-producing UDF's output fields correspond to its
+/// input arguments. This enables the optimizer to propagate orderings
+/// through struct projections (e.g., so that sorting by a struct field
+/// can be recognized as equivalent to sorting by the source column).
+///
+/// See [`ScalarUDFImpl::struct_field_mapping`] for details.
+pub struct StructFieldMapping {
+    /// The UDF used to construct field access expressions on the output.
+    /// For example, the `get_field` UDF for accessing struct fields.
+    pub field_accessor: Arc<ScalarUDF>,
+    /// For each output field: the literal arguments to pass to the
+    /// `field_accessor` UDF (after the base expression), and the index
+    /// of the corresponding input argument that produces the field's value.
+    ///
+    /// For `named_struct('a', col1, 'b', col2)`, this would be:
+    /// `[(["a"], 1), (["b"], 3)]` — field `"a"` comes from arg index 1.
+    pub fields: Vec<(Vec<ScalarValue>, usize)>,
+}
+
 /// Logical representation of a Scalar User Defined Function.
 ///
 /// A scalar function produces a single row output for each row of input. This
@@ -305,6 +324,14 @@ impl ScalarUDF {
         self.inner.evaluate_bounds(inputs)
     }
 
+    /// See [`ScalarUDFImpl::struct_field_mapping`] for more details.
+    pub fn struct_field_mapping(
+        &self,
+        literal_args: &[Option<ScalarValue>],
+    ) -> Option<StructFieldMapping> {
+        self.inner.struct_field_mapping(literal_args)
+    }
+
     /// Updates bounds for child expressions, given a known interval for this
     /// function. This is used to propagate constraints down through an expression
     /// tree.
@@ -961,6 +988,25 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any {
         not_impl_err!("Function {} does not implement coerce_types", self.name())
     }
 
+    /// For struct-producing functions, return how output fields map to input
+    /// arguments. This enables the optimizer to propagate orderings through
+    /// struct projections.
+    ///
+    /// `literal_args[i]` is `Some(value)` if argument `i` is a known literal,
+    /// allowing extraction of field names from arguments like
+    /// `named_struct('field_name', value, ...)`.
+    ///
+    /// For example, `named_struct('a', col1, 'b', col2)` would return a
+    /// mapping indicating that output field `'a'` (accessed via
+    /// `get_field(output, 'a')`) corresponds to input argument `col1` at
+    /// index 1, and field `'b'` corresponds to `col2` at index 3.
+    fn struct_field_mapping(
+        &self,
+        _literal_args: &[Option<ScalarValue>],
+    ) -> Option<StructFieldMapping> {
+        None
+    }
+
     /// Returns the documentation for this Scalar UDF.
     ///
     /// Documentation can be accessed programmatically as well as generating
@@ -1109,6 +1155,13 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl {
         self.inner.propagate_constraints(interval, inputs)
     }
 
+    fn struct_field_mapping(
+        &self,
+        literal_args: &[Option<ScalarValue>],
+    ) -> Option<StructFieldMapping> {
+        self.inner.struct_field_mapping(literal_args)
+    }
+
     fn output_ordering(&self, inputs: &[ExprProperties]) -> Result<SortProperties> {
         self.inner.output_ordering(inputs)
     }

datafusion/functions/src/core/named_struct.rs

@@ -15,11 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use super::getfield::GetFieldFunc;
 use arrow::array::StructArray;
 use arrow::datatypes::{DataType, Field, FieldRef, Fields};
-use datafusion_common::{Result, exec_err, internal_err};
+use datafusion_common::{Result, ScalarValue, exec_err, internal_err};
 use datafusion_expr::{
-    ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs,
+    ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF,
+    StructFieldMapping,
 };
 use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
 use datafusion_macros::user_doc;
@@ -174,4 +176,31 @@ impl ScalarUDFImpl for NamedStructFunc {
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
+
+    fn struct_field_mapping(
+        &self,
+        literal_args: &[Option<ScalarValue>],
+    ) -> Option<StructFieldMapping> {
+        if literal_args.is_empty() || !literal_args.len().is_multiple_of(2) {
+            return None;
+        }
+
+        let mut fields = Vec::with_capacity(literal_args.len() / 2);
+        for (i, chunk) in literal_args.chunks(2).enumerate() {
+            match chunk {
+                [Some(ScalarValue::Utf8(Some(name))), _] => {
+                    fields.push((
+                        vec![ScalarValue::Utf8(Some(name.clone()))],
+                        i * 2 + 1, // index of the value argument
+                    ));
+                }
+                _ => return None,
+            }
+        }
+
+        Some(StructFieldMapping {
+            field_accessor: Arc::new(ScalarUDF::from(GetFieldFunc::new())),
+            fields,
+        })
+    }
 }

datafusion/physical-expr/src/equivalence/properties/dependency.rs

@@ -1564,4 +1564,102 @@ mod tests {
 
         Ok(())
     }
+
+    /// Test that orderings propagate through struct-producing projections.
+    ///
+    /// When a projection creates a struct via `named_struct('a', col_a, ...)`,
+    /// the output should preserve the ordering of `col_a` as an ordering on
+    /// `get_field(col("s"), "a")`. This enables sort elimination when the
+    /// framework sorts by a struct field that corresponds to an already-sorted
+    /// input column.
+    #[test]
+    fn test_ordering_propagation_through_named_struct() -> Result<()> {
+        use crate::expressions::Literal;
+        use datafusion_common::ScalarValue;
+        use datafusion_functions::core::{get_field, named_struct};
+
+        let input_schema = Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int32, true),
+            Field::new("b", DataType::Int32, true),
+        ]));
+
+        let col_a = col("a", &input_schema)?;
+        let col_b = col("b", &input_schema)?;
+        let config = Arc::new(ConfigOptions::new());
+
+        // Build: named_struct('a', col_a, 'b', col_b) AS s
+        let named_struct_udf = named_struct();
+        let named_struct_expr = Arc::new(ScalarFunctionExpr::new(
+            "named_struct",
+            named_struct_udf,
+            vec![
+                Arc::new(Literal::new(ScalarValue::Utf8(Some("a".to_string())))),
+                Arc::clone(&col_a),
+                Arc::new(Literal::new(ScalarValue::Utf8(Some("b".to_string())))),
+                Arc::clone(&col_b),
+            ],
+            Arc::new(Field::new(
+                "named_struct",
+                DataType::Struct(
+                    vec![
+                        Field::new("a", DataType::Int32, true),
+                        Field::new("b", DataType::Int32, true),
+                    ]
+                    .into(),
+                ),
+                true,
+            )),
+            Arc::clone(&config),
+        )) as Arc<dyn PhysicalExpr>;
+
+        // Projection: named_struct(...) AS s
+        let proj_exprs = vec![(named_struct_expr, "s".to_string())];
+        let projection_mapping = ProjectionMapping::try_new(proj_exprs, &input_schema)?;
+
+        // Input is ordered by [a ASC]
+        let mut input_properties = EquivalenceProperties::new(Arc::clone(&input_schema));
+        let sort_a = PhysicalSortExpr::new(
+            Arc::clone(&col_a),
+            SortOptions {
+                descending: false,
+                nulls_first: false,
+            },
+        );
+        input_properties.add_orderings([vec![sort_a]]);
+
+        // Project through the named_struct
+        let out_schema = output_schema(&projection_mapping, &input_schema)?;
+        let out_properties = input_properties.project(&projection_mapping, out_schema);
+
+        // Build the sort expression: get_field(col("s"), "a")
+        // This is what the framework would generate for ORDER BY s.a
+        let get_field_udf = get_field();
+        let col_s = Arc::new(Column::new("s", 0)) as Arc<dyn PhysicalExpr>;
+        let get_field_expr = Arc::new(ScalarFunctionExpr::new(
+            "get_field",
+            get_field_udf,
+            vec![
+                Arc::clone(&col_s),
+                Arc::new(Literal::new(ScalarValue::Utf8(Some("a".to_string())))),
+            ],
+            Arc::new(Field::new("a", DataType::Int32, true)),
+            Arc::clone(&config),
+        )) as Arc<dyn PhysicalExpr>;
+
+        let sort_get_field_a = PhysicalSortExpr::new(
+            get_field_expr,
+            SortOptions {
+                descending: false,
+                nulls_first: false,
+            },
+        );
+
+        // The output should satisfy ordering by get_field(s, "a")
+        assert!(
+            out_properties.ordering_satisfy(vec![sort_get_field_a])?,
+            "Output should be ordered by get_field(s, 'a') since input is ordered by col_a"
+        );
+
+        Ok(())
+    }
 }

datafusion/physical-expr/src/projection.rs

@@ -22,10 +22,11 @@ use std::sync::Arc;
 
 use crate::PhysicalExpr;
 use crate::expressions::{Column, Literal};
+use crate::scalar_function::ScalarFunctionExpr;
 use crate::utils::collect_columns;
 
 use arrow::array::{RecordBatch, RecordBatchOptions};
-use arrow::datatypes::{Field, Schema, SchemaRef};
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion_common::stats::{ColumnStatistics, Precision};
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_common::{
@@ -1063,9 +1064,72 @@ impl ProjectionMapping {
                 None => Ok(Transformed::no(e)),
             })
             .data()?;
-            map.entry(source_expr)
+            map.entry(Arc::clone(&source_expr))
                 .or_default()
-                .push((target_expr, expr_idx));
+                .push((Arc::clone(&target_expr), expr_idx));
+
+            // For struct-producing functions (e.g. named_struct), decompose
+            // into field-level mapping entries so that orderings propagate
+            // through struct projections. For example, if the projection has
+            // `named_struct('ticker', p.ticker, ...) AS details`, this adds:
+            //   p.ticker → get_field(col("details"), "ticker")
+            // enabling the optimizer to know that sorting by
+            // `details.ticker` is equivalent to sorting by `p.ticker`.
+            if let Some(func_expr) =
+                source_expr.as_any().downcast_ref::<ScalarFunctionExpr>()
+            {
+                let literal_args: Vec<Option<ScalarValue>> = func_expr
+                    .args()
+                    .iter()
+                    .map(|arg| {
+                        arg.as_any()
+                            .downcast_ref::<Literal>()
+                            .map(|l| l.value().clone())
+                    })
+                    .collect();
+
+                if let Some(field_mapping) =
+                    func_expr.fun().struct_field_mapping(&literal_args)
+                    && let DataType::Struct(struct_fields) = func_expr.return_type()
+                {
+                    for (accessor_args, source_arg_idx) in &field_mapping.fields {
+                        let value_expr = Arc::clone(&func_expr.args()[*source_arg_idx]);
+
+                        // Build accessor args: [target_col, ...field_name_literals]
+                        let mut accessor_fn_args: Vec<Arc<dyn PhysicalExpr>> =
+                            vec![Arc::clone(&target_expr)];
+                        accessor_fn_args.extend(accessor_args.iter().map(|sv| {
+                            Arc::new(Literal::new(sv.clone())) as Arc<dyn PhysicalExpr>
+                        }));
+
+                        // Look up the field's return type from the struct schema
+                        let return_field = accessor_args
+                            .first()
+                            .and_then(|sv| sv.try_as_str().flatten())
+                            .and_then(|field_name| {
+                                struct_fields
+                                    .iter()
+                                    .find(|f| f.name() == field_name)
+                                    .cloned()
+                            });
+
+                        if let Some(return_field) = return_field {
+                            let field_access_expr = Arc::new(ScalarFunctionExpr::new(
+                                field_mapping.field_accessor.name(),
+                                Arc::clone(&field_mapping.field_accessor),
+                                accessor_fn_args,
+                                return_field,
+                                Arc::new(func_expr.config_options().clone()),
+                            ))
+                                as Arc<dyn PhysicalExpr>;
+
+                            map.entry(value_expr)
+                                .or_default()
+                                .push((field_access_expr, expr_idx));
+                        }
+                    }
+                }
+            }
         }
         Ok(Self { map })
     }
@@ -1219,8 +1283,10 @@ pub(crate) mod tests {
             let data_type = source.data_type(input_schema)?;
             let nullable = source.nullable(input_schema)?;
             for (target, _) in targets.iter() {
+                // Skip non-Column targets (e.g. struct field decomposition
+                // entries which are ScalarFunctionExpr targets).
                 let Some(column) = target.as_any().downcast_ref::<Column>() else {
-                    return plan_err!("Expects to have column");
+                    continue;
                 };
                 fields.push(Field::new(column.name(), data_type.clone(), nullable));
             }