fix: string_agg not respecting ORDER BY

Author: nuno-faria

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,7 +31,7 @@ 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, SingleRowListArrayBuilder};
 use datafusion_common::{exec_err, internal_err, Result, ScalarValue};
 use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
 use datafusion_expr::utils::format_state_name;
@@ -78,12 +78,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 +146,16 @@ impl AggregateUDFImpl for ArrayAgg {
         Ok(fields)
     }
 
+    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 +208,7 @@ impl AggregateUDFImpl for ArrayAgg {
             &data_type,
             &ordering_dtypes,
             ordering,
+            self.is_input_pre_ordered,
             acc_args.is_reversed,
             ignore_nulls,
         )
@@ -518,6 +531,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 +546,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 +557,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..]);
 
@@ -629,6 +668,9 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator {
         let mut partition_ordering_values = vec![];
 
         // Existing values should be merged also.
+        if !self.is_input_pre_ordered {
+            self.sort();
+        }
         partition_values.push(self.values.clone().into());
         partition_ordering_values.push(self.ordering_values.clone().into());
 
@@ -679,13 +721,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(),

datafusion/sqllogictest/test_files/aggregate.slt

@@ -6028,6 +6028,84 @@ GROUP BY dummy
 ----
 text1
 
+
+# Test string_agg with ORDER BY clasuses (issue #17011)
+statement ok
+create table t (k varchar, v int);
+
+statement ok
+insert into t values ('a', 2), ('b', 3), ('c', 1), ('d', null);
+
+query T
+select string_agg(k, ',' order by k) from t;
+----
+a,b,c,d
+
+query T
+select string_agg(k, ',' order by k desc) from t;
+----
+d,c,b,a
+
+query T
+select string_agg(k, ',' order by v) from t;
+----
+c,a,b,d
+
+query T
+select string_agg(k, ',' order by v nulls first) from t;
+----
+d,c,a,b
+
+query T
+select string_agg(k, ',' order by v desc) from t;
+----
+d,b,a,c
+
+query T
+select string_agg(k, ',' order by v desc nulls last) from t;
+----
+b,a,c,d
+
+query T
+-- odd indexes should appear first, ties solved by v
+select string_agg(k, ',' order by v % 2 == 0, v) from t;
+----
+c,b,a,d
+
+query T
+-- odd indexes should appear first, ties solved by v desc
+select string_agg(k, ',' order by v % 2 == 0, v desc) from t;
+----
+b,c,a,d
+
+query T
+select string_agg(k, ',' order by
+  case
+    when k = 'a' then 3
+    when k = 'b' then 0
+    when k = 'c' then 2
+    when k = 'd' then 1 
+  end)
+from t;
+----
+b,d,c,a
+
+query T
+select string_agg(k, ',' order by
+  case
+    when k = 'a' then 3
+    when k = 'b' then 0
+    when k = 'c' then 2
+    when k = 'd' then 1 
+  end desc)
+from t;
+----
+a,c,d,b
+
+statement ok
+drop table t;
+
+
 # Tests for aggregating with NaN values
 statement ok
 CREATE TABLE float_table (