Support multiple ordered array_agg aggregations (#16625)

* Validate states shape in merge_batch

Due to `..` in the pattern, the
`OrderSensitiveArrayAggAccumulator::merge_batch` did not validate it's
not receiving additional states columns it ignores. Update the code to
check number of inputs.

* Support multiple ordered array_agg

Before the change, `array_agg` with ordering would depend on input being
ordered. As a result, it was impossible to do two or more `array_agg(x
ORDER BY ...)` with incompatible ordering. This change moves ordering
responsibility into `OrderSensitiveArrayAggAccumulator`. When input is
pre-ordered (beneficial ordering), no additional work is done. However,
when it's not, `array_agg` accumulator will order the data on its own.

* Generate sorts based on aggregations soft requirements

The sorting consideration before aggregations did respect only ordered
aggregation functions with `AggregateOrderSensitivity::HardRequirement`.
This change includes sorting expectations from
`AggregateOrderSensitivity::Beneficial` functions. When beneficial
ordered function requirements are not satisfied, no error is raised,
they are considered in the second pass only.

* Fix reversing first_value, last_value

Upon reversing, a schema and field mismatch would happen.

* Revert "Fix reversing first_value, last_value"

This reverts commit 9b7e94d.

* sort array_agg input the old way whenever possible

* revert some now unnecessary change

* Improve doc for SoftRequiement

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

* Add comment for include_soft_requirement

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

* Document include_soft_requirement param

* fmt

* doc fix

---------

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

Author: findepi

datafusion/ffi/src/udaf/mod.rs

@@ -589,6 +589,7 @@ impl AggregateUDFImpl for ForeignAggregateUDF {
 pub enum FFI_AggregateOrderSensitivity {
     Insensitive,
     HardRequirement,
+    SoftRequirement,
     Beneficial,
 }
 
@@ -597,6 +598,7 @@ impl From<FFI_AggregateOrderSensitivity> for AggregateOrderSensitivity {
         match value {
             FFI_AggregateOrderSensitivity::Insensitive => Self::Insensitive,
             FFI_AggregateOrderSensitivity::HardRequirement => Self::HardRequirement,
+            FFI_AggregateOrderSensitivity::SoftRequirement => Self::SoftRequirement,
             FFI_AggregateOrderSensitivity::Beneficial => Self::Beneficial,
         }
     }
@@ -607,6 +609,7 @@ impl From<AggregateOrderSensitivity> for FFI_AggregateOrderSensitivity {
         match value {
             AggregateOrderSensitivity::Insensitive => Self::Insensitive,
             AggregateOrderSensitivity::HardRequirement => Self::HardRequirement,
+            AggregateOrderSensitivity::SoftRequirement => Self::SoftRequirement,
             AggregateOrderSensitivity::Beneficial => Self::Beneficial,
         }
     }
@@ -748,6 +751,7 @@ mod tests {
     fn test_round_trip_all_order_sensitivities() {
         test_round_trip_order_sensitivity(AggregateOrderSensitivity::Insensitive);
         test_round_trip_order_sensitivity(AggregateOrderSensitivity::HardRequirement);
+        test_round_trip_order_sensitivity(AggregateOrderSensitivity::SoftRequirement);
         test_round_trip_order_sensitivity(AggregateOrderSensitivity::Beneficial);
     }
 }

datafusion/functions-aggregate-common/src/order.rs

@@ -22,9 +22,20 @@ pub enum AggregateOrderSensitivity {
     /// Ordering at the input is not important for the result of the aggregator.
     Insensitive,
     /// Indicates that the aggregate expression has a hard requirement on ordering.
-    /// The aggregator can not produce a correct result unless its ordering
+    /// The aggregator cannot produce a correct result unless its ordering
     /// requirement is satisfied.
     HardRequirement,
+    /// Indicates that the aggregator is more efficient when the input is ordered
+    /// but can still produce its result correctly regardless of the input ordering.
+    /// This is similar to, but stronger than, [`Self::Beneficial`].
+    ///
+    /// Similarly to [`Self::HardRequirement`], when possible DataFusion will insert
+    /// a `SortExec`, to reorder the input to match the SoftRequirement. However,
+    /// when such a `SortExec` cannot be inserted, (for example, due to conflicting
+    /// [`Self::HardRequirement`] with other ordered aggregates in the query),
+    /// the aggregate function will still execute, without the preferred order, unlike
+    /// with [`Self::HardRequirement`]
+    SoftRequirement,
     /// Indicates that ordering is beneficial for the aggregate expression in terms
     /// of evaluation efficiency. The aggregator can produce its result efficiently
     /// when its required ordering is satisfied; however, it can still produce the
@@ -38,7 +49,7 @@ impl AggregateOrderSensitivity {
     }
 
     pub fn is_beneficial(&self) -> bool {
-        self.eq(&AggregateOrderSensitivity::Beneficial)
+        matches!(self, Self::SoftRequirement | Self::Beneficial)
     }
 
     pub fn hard_requires(&self) -> bool {

datafusion/functions-aggregate/src/array_agg.rs

@@ -19,7 +19,7 @@
 
 use std::cmp::Ordering;
 use std::collections::{HashSet, VecDeque};
-use std::mem::{size_of, size_of_val};
+use std::mem::{size_of, size_of_val, take};
 use std::sync::Arc;
 
 use arrow::array::{
@@ -31,14 +31,17 @@ use arrow::datatypes::{DataType, Field, FieldRef, Fields};
 
 use datafusion_common::cast::as_list_array;
 use datafusion_common::scalar::copy_array_data;
-use datafusion_common::utils::{get_row_at_idx, SingleRowListArrayBuilder};
+use datafusion_common::utils::{
+    compare_rows, get_row_at_idx, take_function_args, SingleRowListArrayBuilder,
+};
 use datafusion_common::{exec_err, internal_err, Result, ScalarValue};
 use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion_expr::utils::format_state_name;
 use datafusion_expr::{
     Accumulator, AggregateUDFImpl, Documentation, Signature, Volatility,
 };
 use datafusion_functions_aggregate_common::merge_arrays::merge_ordered_arrays;
+use datafusion_functions_aggregate_common::order::AggregateOrderSensitivity;
 use datafusion_functions_aggregate_common::utils::ordering_fields;
 use datafusion_macros::user_doc;
 use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
@@ -78,12 +81,14 @@ This aggregation function can only mix DISTINCT and ORDER BY if the ordering exp
 /// ARRAY_AGG aggregate expression
 pub struct ArrayAgg {
     signature: Signature,
+    is_input_pre_ordered: bool,
 }
 
 impl Default for ArrayAgg {
     fn default() -> Self {
         Self {
             signature: Signature::any(1, Volatility::Immutable),
+            is_input_pre_ordered: false,
         }
     }
 }
@@ -144,6 +149,20 @@ impl AggregateUDFImpl for ArrayAgg {
         Ok(fields)
     }
 
+    fn order_sensitivity(&self) -> AggregateOrderSensitivity {
+        AggregateOrderSensitivity::SoftRequirement
+    }
+
+    fn with_beneficial_ordering(
+        self: Arc<Self>,
+        beneficial_ordering: bool,
+    ) -> Result<Option<Arc<dyn AggregateUDFImpl>>> {
+        Ok(Some(Arc::new(Self {
+            signature: self.signature.clone(),
+            is_input_pre_ordered: beneficial_ordering,
+        })))
+    }
+
     fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
         let data_type = acc_args.exprs[0].data_type(acc_args.schema)?;
         let ignore_nulls =
@@ -196,6 +215,7 @@ impl AggregateUDFImpl for ArrayAgg {
             &data_type,
             &ordering_dtypes,
             ordering,
+            self.is_input_pre_ordered,
             acc_args.is_reversed,
             ignore_nulls,
         )
@@ -518,6 +538,8 @@ pub(crate) struct OrderSensitiveArrayAggAccumulator {
     datatypes: Vec<DataType>,
     /// Stores the ordering requirement of the `Accumulator`.
     ordering_req: LexOrdering,
+    /// Whether the input is known to be pre-ordered
+    is_input_pre_ordered: bool,
     /// Whether the aggregation is running in reverse.
     reverse: bool,
     /// Whether the aggregation should ignore null values.
@@ -531,6 +553,7 @@ impl OrderSensitiveArrayAggAccumulator {
         datatype: &DataType,
         ordering_dtypes: &[DataType],
         ordering_req: LexOrdering,
+        is_input_pre_ordered: bool,
         reverse: bool,
         ignore_nulls: bool,
     ) -> Result<Self> {
@@ -541,11 +564,34 @@ impl OrderSensitiveArrayAggAccumulator {
             ordering_values: vec![],
             datatypes,
             ordering_req,
+            is_input_pre_ordered,
             reverse,
             ignore_nulls,
         })
     }
 
+    fn sort(&mut self) {
+        let sort_options = self
+            .ordering_req
+            .iter()
+            .map(|sort_expr| sort_expr.options)
+            .collect::<Vec<_>>();
+        let mut values = take(&mut self.values)
+            .into_iter()
+            .zip(take(&mut self.ordering_values))
+            .collect::<Vec<_>>();
+        let mut delayed_cmp_err = Ok(());
+        values.sort_by(|(_, left_ordering), (_, right_ordering)| {
+            compare_rows(left_ordering, right_ordering, &sort_options).unwrap_or_else(
+                |err| {
+                    delayed_cmp_err = Err(err);
+                    Ordering::Equal
+                },
+            )
+        });
+        (self.values, self.ordering_values) = values.into_iter().unzip();
+    }
+
     fn evaluate_orderings(&self) -> Result<ScalarValue> {
         let fields = ordering_fields(&self.ordering_req, &self.datatypes[1..]);
 
@@ -616,9 +662,8 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator {
         // inside `ARRAY_AGG` list, we will receive an `Array` that stores values
         // received from its ordering requirement expression. (This information
         // is necessary for during merging).
-        let [array_agg_values, agg_orderings, ..] = &states else {
-            return exec_err!("State should have two elements");
-        };
+        let [array_agg_values, agg_orderings] =
+            take_function_args("OrderSensitiveArrayAggAccumulator::merge_batch", states)?;
         let Some(agg_orderings) = agg_orderings.as_list_opt::<i32>() else {
             return exec_err!("Expects to receive a list array");
         };
@@ -629,8 +674,11 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator {
         let mut partition_ordering_values = vec![];
 
         // Existing values should be merged also.
-        partition_values.push(self.values.clone().into());
-        partition_ordering_values.push(self.ordering_values.clone().into());
+        if !self.is_input_pre_ordered {
+            self.sort();
+        }
+        partition_values.push(take(&mut self.values).into());
+        partition_ordering_values.push(take(&mut self.ordering_values).into());
 
         // Convert array to Scalars to sort them easily. Convert back to array at evaluation.
         let array_agg_res = ScalarValue::convert_array_to_scalar_vec(array_agg_values)?;
@@ -679,13 +727,21 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator {
     }
 
     fn state(&mut self) -> Result<Vec<ScalarValue>> {
+        if !self.is_input_pre_ordered {
+            self.sort();
+        }
+
         let mut result = vec![self.evaluate()?];
         result.push(self.evaluate_orderings()?);
 
         Ok(result)
     }
 
     fn evaluate(&mut self) -> Result<ScalarValue> {
+        if !self.is_input_pre_ordered {
+            self.sort();
+        }
+
         if self.values.is_empty() {
             return Ok(ScalarValue::new_null_list(
                 self.datatypes[0].clone(),