fix: duplicate ordering state field names in partial aggregates

datafusion/common/src/dfschema.rs

@@ -1138,11 +1138,7 @@ impl TryFrom<SchemaRef> for DFSchema {
             field_qualifiers: vec![None; field_count],
             functional_dependencies: FunctionalDependencies::empty(),
         };
-        // Without checking names, because schema here may have duplicate field names.
-        // For example, Partial AggregateMode will generate duplicate field names from
-        // state_fields.
-        // See <https://github.com/apache/datafusion/issues/17715>
-        // dfschema.check_names()?;
+        dfschema.check_names()?;
         Ok(dfschema)
     }
 }

datafusion/core/tests/dataframe/mod.rs

@@ -7611,7 +7611,7 @@ async fn execute_logical_plan_rejects_duplicate_unqualified_names_in_replace_vie
 }
 
 #[tokio::test]
-async fn test_duplicate_state_fields_for_dfschema_construct() -> Result<()> {
+async fn test_partial_aggregate_state_fields_have_unique_names() -> Result<()> {
     let ctx = SessionContext::new();
 
     // Simple schema with just the fields we need
@@ -7686,28 +7686,22 @@ async fn test_duplicate_state_fields_for_dfschema_construct() -> Result<()> {
     )
     .expect("Failed to build partial agg");
 
-    // Assert that the schema field names match the expected names
     let expected_field_names = vec![
         "date",
         "ticker",
         "first_value(value)[first_value]",
-        "timestamp@0",
+        "first_value(value)[ordering_0_timestamp@0]",
         "first_value(value)[first_value_is_set]",
         "last_value(value)[last_value]",
-        "timestamp@0",
+        "last_value(value)[ordering_0_timestamp@0]",
         "last_value(value)[last_value_is_set]",
     ];
 
     let binding = partial_agg.schema();
     let actual_field_names: Vec<_> = binding.fields().iter().map(|f| f.name()).collect();
     assert_eq!(actual_field_names, expected_field_names);
 
-    // Ensure that DFSchema::try_from does not fail
-    let partial_agg_exec_schema = DFSchema::try_from(partial_agg.schema());
-    assert!(
-        partial_agg_exec_schema.is_ok(),
-        "Expected get AggregateExec schema to succeed with duplicate state fields"
-    );
+    DFSchema::try_from(partial_agg.schema())?;
 
     Ok(())
 }

datafusion/expr/src/udaf.rs

@@ -609,7 +609,12 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any {
         Ok(fields
             .into_iter()
             .map(Arc::new)
-            .chain(args.ordering_fields.to_vec())
+            .chain(args.ordering_fields.iter().enumerate().map(|(idx, field)| {
+                Arc::new(field.as_ref().clone().with_name(format_state_name(
+                    args.name,
+                    &format!("ordering_{idx}_{}", field.name()),
+                )))
+            }))
             .collect())
     }
 
@@ -1716,7 +1721,7 @@ pub enum SetMonotonicity {
 #[cfg(test)]
 mod test {
     use crate::{AggregateUDF, AggregateUDFImpl};
-    use arrow::datatypes::{DataType, FieldRef};
+    use arrow::datatypes::{DataType, Field, FieldRef};
     use datafusion_common::Result;
     use datafusion_expr_common::accumulator::Accumulator;
     use datafusion_expr_common::signature::{Signature, Volatility};
@@ -1725,6 +1730,7 @@ mod test {
     };
     use std::cmp::Ordering;
     use std::hash::{DefaultHasher, Hash, Hasher};
+    use std::sync::Arc;
 
     #[derive(Debug, Clone, PartialEq, Eq, Hash)]
     struct AMeanUdf {
@@ -1801,6 +1807,44 @@ mod test {
         }
     }
 
+    #[derive(Debug, Clone, PartialEq, Eq, Hash)]
+    struct DefaultStateFieldsUdf {
+        signature: Signature,
+    }
+
+    impl DefaultStateFieldsUdf {
+        fn new() -> Self {
+            Self {
+                signature: Signature::uniform(
+                    1,
+                    vec![DataType::Float64],
+                    Volatility::Immutable,
+                ),
+            }
+        }
+    }
+
+    impl AggregateUDFImpl for DefaultStateFieldsUdf {
+        fn name(&self) -> &str {
+            "default_state_fields"
+        }
+
+        fn signature(&self) -> &Signature {
+            &self.signature
+        }
+
+        fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+            unimplemented!()
+        }
+
+        fn accumulator(
+            &self,
+            _acc_args: AccumulatorArgs,
+        ) -> Result<Box<dyn Accumulator>> {
+            unimplemented!()
+        }
+    }
+
     #[test]
     fn test_partial_eq() {
         let a1 = AggregateUDF::from(AMeanUdf::new());
@@ -1829,4 +1873,37 @@ mod test {
         value.hash(hasher);
         hasher.finish()
     }
+
+    #[test]
+    fn test_default_state_fields_namespaces_ordering_fields() -> Result<()> {
+        let udf = DefaultStateFieldsUdf::new();
+
+        let input_fields = vec![Arc::new(Field::new("value", DataType::Float64, true))];
+
+        let ordering_fields = vec![
+            Arc::new(Field::new("timestamp@0", DataType::Int64, true)),
+            Arc::new(Field::new("timestamp@0", DataType::Int64, true)),
+        ];
+
+        let fields = udf.state_fields(StateFieldsArgs {
+            name: "my_agg(value)",
+            input_fields: &input_fields,
+            return_field: Arc::new(Field::new("result", DataType::Float64, true)),
+            ordering_fields: &ordering_fields,
+            is_distinct: false,
+        })?;
+
+        let names: Vec<_> = fields.iter().map(|f| f.name().as_str()).collect();
+
+        assert_eq!(
+            names,
+            vec![
+                "my_agg(value)[value]",
+                "my_agg(value)[ordering_0_timestamp@0]",
+                "my_agg(value)[ordering_1_timestamp@0]",
+            ]
+        );
+
+        Ok(())
+    }
 }

datafusion/ffi/src/udaf/mod.rs

@@ -814,6 +814,9 @@ mod tests {
         );
 
         let a_field = Arc::new(Field::new("a", DataType::Float64, true));
+
+        let expected_ordering_field =
+            Arc::new(a_field.as_ref().clone().with_name("a[ordering_0_a]"));
         let state_fields = foreign_udaf.state_fields(StateFieldsArgs {
             name: "a",
             input_fields: &[Field::new("f", DataType::Float64, true).into()],
@@ -823,7 +826,7 @@ mod tests {
         })?;
 
         assert_eq!(state_fields.len(), 3);
-        assert_eq!(state_fields[1], a_field);
+        assert_eq!(state_fields[1], expected_ordering_field);
         Ok(())
     }
 

datafusion/functions-aggregate/src/first_last.rs

@@ -332,7 +332,12 @@ impl AggregateUDFImpl for FirstValue {
             )
             .into(),
         ];
-        fields.extend(args.ordering_fields.iter().cloned());
+        fields.extend(args.ordering_fields.iter().enumerate().map(|(idx, field)| {
+            Arc::new(field.as_ref().clone().with_name(format_state_name(
+                args.name,
+                &format!("ordering_{idx}_{}", field.name()),
+            )))
+        }));
         fields.push(
             Field::new(
                 format_state_name(args.name, "first_value_is_set"),
@@ -1246,7 +1251,12 @@ impl AggregateUDFImpl for LastValue {
             )
             .into(),
         ];
-        fields.extend(args.ordering_fields.iter().cloned());
+        fields.extend(args.ordering_fields.iter().enumerate().map(|(idx, field)| {
+            Arc::new(field.as_ref().clone().with_name(format_state_name(
+                args.name,
+                &format!("ordering_{idx}_{}", field.name()),
+            )))
+        }));
         fields.push(
             Field::new(
                 format_state_name(args.name, "last_value_is_set"),