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

Author: rkrishn7

datafusion/expr/src/lib.rs

@@ -125,7 +125,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

@@ -36,6 +36,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
@@ -344,6 +363,14 @@ impl ScalarUDF {
         self.inner.documentation()
     }
 
+    /// 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)
+    }
+
     /// Return true if this function is an async function
     pub fn as_async(&self) -> Option<&AsyncScalarUDF> {
         self.inner().as_any().downcast_ref::<AsyncScalarUDF>()
@@ -846,6 +873,25 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync {
     fn documentation(&self) -> Option<&Documentation> {
         None
     }
+
+    /// 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
+    }
 }
 
 /// ScalarUDF that adds an alias to the underlying function. It is better to
@@ -964,6 +1010,13 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl {
     fn documentation(&self) -> Option<&Documentation> {
         self.inner.documentation()
     }
+
+    fn struct_field_mapping(
+        &self,
+        literal_args: &[Option<ScalarValue>],
+    ) -> Option<StructFieldMapping> {
+        self.inner.struct_field_mapping(literal_args)
+    }
 }
 
 #[cfg(test)]

datafusion/functions/src/core/named_struct.rs

@@ -17,11 +17,14 @@
 
 use arrow::array::StructArray;
 use arrow::datatypes::{DataType, Field, FieldRef, Fields};
-use datafusion_common::{Result, exec_err, internal_err};
+use datafusion_common::{exec_err, internal_err, Result, ScalarValue};
 use datafusion_expr::{
-    ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs,
+    ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF,
+    StructFieldMapping,
 };
 use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
+
+use super::getfield::GetFieldFunc;
 use datafusion_macros::user_doc;
 use std::any::Any;
 use std::sync::Arc;
@@ -177,4 +180,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

@@ -1528,4 +1528,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

@@ -20,12 +20,13 @@
 use std::ops::Deref;
 use std::sync::Arc;
 
+use crate::scalar_function::ScalarFunctionExpr;
 use crate::PhysicalExpr;
 use crate::expressions::{Column, Literal};
 use crate::utils::collect_columns;
 
 use arrow::array::{RecordBatch, RecordBatchOptions};
-use arrow::datatypes::{Field, Schema, SchemaRef};
+use arrow::datatypes::{Field, Schema, SchemaRef, DataType};
 use datafusion_common::stats::{ColumnStatistics, Precision};
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_common::{
@@ -952,9 +953,73 @@ 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 })
     }
@@ -1110,8 +1175,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));
             }