lyne7-sc commented on code in PR #24392:
URL: https://github.com/apache/datafusion/pull/24392#discussion_r3791969164
##########
datafusion/functions-aggregate/src/array_agg.rs:
##########
@@ -1221,246 +1232,382 @@ impl OrderSensitiveArrayAggAccumulator {
pub fn try_new(
datatype: &DataType,
ordering_dtypes: &[DataType],
- ordering_req: LexOrdering,
+ ordering_req: &LexOrdering,
is_input_pre_ordered: bool,
reverse: bool,
ignore_nulls: bool,
) -> Result<Self> {
- let mut datatypes = vec![datatype.clone()];
- datatypes.extend(ordering_dtypes.iter().cloned());
+ assert_eq_or_internal_err!(
+ ordering_dtypes.len(),
+ ordering_req.len(),
+ "ordered array_agg requires one datatype per ordering expression"
+ );
+ let ordering_fields =
+ Fields::from(ordering_fields(ordering_req, ordering_dtypes));
+ let sort_fields = ordering_dtypes
+ .iter()
+ .zip(ordering_req.iter())
+ .map(|(data_type, sort_expr)| {
+ SortField::new_with_options(data_type.clone(),
sort_expr.options)
+ })
+ .collect();
+ let ordering_converter = RowConverter::new(sort_fields)?;
+ let ordering_rows = ordering_converter.empty_rows(0, 0);
Ok(Self {
- values: vec![],
- ordering_values: vec![],
- datatypes,
- ordering_req,
+ batches: vec![],
+ entries: vec![],
+ ordering_rows,
+ sorted_runs: vec![],
+ sorted_entry_indices: None,
+ ordering_converter,
+ value_type: datatype.clone(),
+ ordering_fields,
is_input_pre_ordered,
reverse,
ignore_nulls,
+ can_extend_preordered_run: false,
})
}
- fn sort(&mut self) {
- let sort_options = self
- .ordering_req
+ fn append_input_batch(
+ &mut self,
+ values: &ArrayRef,
+ ordering_values: &[ArrayRef],
+ ) -> Result<()> {
+ let Some(entry_range) =
+ self.store_batch(values, ordering_values, self.ignore_nulls)?
+ else {
+ return Ok(());
+ };
+ if self.is_input_pre_ordered {
+ if self.can_extend_preordered_run {
+ self.sorted_runs
+ .last_mut()
+ .expect("an extendable preordered run must exist")
+ .end = entry_range.end;
+ } else {
+ self.sorted_runs.push(entry_range);
+ }
+ }
+ self.can_extend_preordered_run = self.is_input_pre_ordered;
+ Ok(())
+ }
+
+ fn append_sorted_run(
+ &mut self,
+ values: &ArrayRef,
+ ordering_values: &[ArrayRef],
+ ) -> Result<()> {
+ if let Some(entry_range) = self.store_batch(values, ordering_values,
false)?
+ && entry_range
+ .clone()
+ .zip(entry_range.start + 1..entry_range.end)
+ .all(|(left, right)| self.ordering_row(left) <=
self.ordering_row(right))
+ {
+ self.sorted_runs.push(entry_range);
+ }
+ self.can_extend_preordered_run = false;
+ Ok(())
+ }
+
+ fn ordering_row(&self, entry_idx: usize) -> Row<'_> {
+ self.ordering_rows.row(entry_idx)
+ }
+
+ fn merge_sorted_runs(&self, unsorted_indices: Vec<usize>) -> Vec<usize> {
+ let unsorted_run = (!unsorted_indices.is_empty())
+ .then(|| Either::Right(unsorted_indices.into_iter()));
+ self.sorted_runs
.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
- },
- )
+ .cloned()
+ .map(Either::Left)
+ .chain(unsorted_run)
+ .kmerge_by(|left, right| {
+ self.ordering_row(*left)
+ .cmp(&self.ordering_row(*right))
+ .then_with(|| left.cmp(right))
+ .is_lt()
+ })
+ .collect()
+ }
+
+ fn ensure_sorted_indices(&mut self) {
+ if self.sorted_entry_indices.is_some() {
+ return;
+ }
+
+ let sorted_len = self.sorted_runs.iter().map(|run|
run.len()).sum::<usize>();
+ let mut unsorted_indices = Vec::with_capacity(self.entries.len() -
sorted_len);
+ let mut next_unsorted = 0;
+ for run in &self.sorted_runs {
+ debug_assert!(run.start >= next_unsorted);
+ debug_assert!(run.end <= self.entries.len());
+ unsorted_indices.extend(next_unsorted..run.start);
+ next_unsorted = run.end;
+ }
+ unsorted_indices.extend(next_unsorted..self.entries.len());
+ unsorted_indices.sort_by(|left, right| {
+ self.ordering_row(*left)
+ .cmp(&self.ordering_row(*right))
+ .then_with(|| left.cmp(right))
});
- (self.values, self.ordering_values) = values.into_iter().unzip();
+ self.sorted_entry_indices =
Some(self.merge_sorted_runs(unsorted_indices));
}
- fn evaluate_orderings(&self) -> Result<ScalarValue> {
- let fields = ordering_fields(&self.ordering_req, &self.datatypes[1..]);
+ fn select_values(&self, sorted_indices: &[usize], reverse: bool) ->
Result<ArrayRef> {
+ if sorted_indices.is_empty() {
+ return Ok(new_empty_array(&self.value_type));
+ }
+
+ // A common preordered case is a consecutive range in one input batch.
+ // Return a zero-copy slice instead of invoking interleave.
+ if !reverse {
+ let first = self.entries[sorted_indices[0]];
+ let is_contiguous =
sorted_indices.iter().enumerate().all(|(offset, idx)| {
+ let entry = self.entries[*idx];
+ entry.batch_idx == first.batch_idx
+ && entry.row_idx == first.row_idx + offset
+ });
+ if is_contiguous {
+ return Ok(self.batches[first.batch_idx]
+ .slice(first.row_idx, sorted_indices.len()));
+ }
+ }
- let column_wise_ordering_values = if self.ordering_values.is_empty() {
- fields
+ let sources = self
+ .batches
+ .iter()
+ .map(|batch| batch.as_ref())
+ .collect::<Vec<_>>();
+ let indices = if reverse {
+ sorted_indices
.iter()
- .map(|f| new_empty_array(f.data_type()))
+ .rev()
+ .map(|idx| {
+ let entry = self.entries[*idx];
+ (entry.batch_idx, entry.row_idx)
+ })
.collect::<Vec<_>>()
} else {
- (0..fields.len())
- .map(|i| {
- let column_values: Box<dyn Iterator<Item = ScalarValue>> =
if self
- .reverse
- {
- Box::new(self.ordering_values.iter().rev().map(|x|
x[i].clone()))
- } else {
- Box::new(self.ordering_values.iter().map(|x|
x[i].clone()))
- };
- ScalarValue::iter_to_array(column_values)
+ sorted_indices
+ .iter()
+ .map(|idx| {
+ let entry = self.entries[*idx];
+ (entry.batch_idx, entry.row_idx)
})
- .collect::<Result<_>>()?
+ .collect::<Vec<_>>()
};
+ Ok(arrow::compute::interleave(&sources, &indices)?)
+ }
- let ordering_array = StructArray::try_new(
- Fields::from(fields),
- column_wise_ordering_values,
- None,
- )?;
+ fn evaluate_orderings(
+ &self,
+ sorted_indices: &[usize],
+ reverse: bool,
+ ) -> Result<ScalarValue> {
+ let indices = if reverse {
+ Either::Left(sorted_indices.iter().rev())
+ } else {
+ Either::Right(sorted_indices.iter())
+ };
+ let mut columns = self
+ .ordering_converter
+ .convert_rows(indices.map(|idx| self.ordering_row(*idx)))?;
+
+ // RowConverter decodes dictionary values to their physical type. State
+ // fields, however, are required to retain their declared logical type.
+ for (column, field) in columns.iter_mut().zip(&self.ordering_fields) {
+ if column.data_type() != field.data_type() {
+ *column = cast(column.as_ref(), field.data_type())?;
+ }
+ }
+
+ let ordering_array =
+ StructArray::try_new(self.ordering_fields.clone(), columns, None)?;
Ok(SingleRowListArrayBuilder::new(Arc::new(ordering_array)).build_list_scalar())
}
-}
-impl Accumulator for OrderSensitiveArrayAggAccumulator {
- fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
- if values.is_empty() {
- return Ok(());
+ fn store_batch(
+ &mut self,
+ values: &ArrayRef,
+ ordering_values: &[ArrayRef],
+ ignore_nulls: bool,
+ ) -> Result<Option<Range<usize>>> {
+ let values = if values.data_type() == &self.value_type {
+ Arc::clone(values)
+ } else if self.value_type.contains(values.data_type()) {
+ cast(values.as_ref(), &self.value_type)?
Review Comment:
I think this could ideally be handled at a common physical-expression
boundary rather than by each accumulator.
However, `self.value_type` comes from the declared aggregate argument type,
while the concrete `ArrayRef` is produced later during physical expression
evaluation. It seems that exact type equality is not currently validated at
that boundary, as also noted in #24022. #24029 handles such differences by
normalizing the final array in `ScalarValue::new_list`.
Since this columnar implementation bypasses `ScalarValue::new_list` and
`interleave` expects consistent source types, normalizing compatible inputs
here seems like a reasonable way to preserve the existing behavior.
The accumulator-side nested-nullability case is exercised by
[`ordered_aggregate_nested_nullability_mismatch_issue_24022`](https://github.com/apache/datafusion/blob/bc03e0c1f26a01db4d86ca84835440d270e4e73f/datafusion/functions-aggregate/src/array_agg.rs#L1937-L2005).
The test expects the result to use the accumulator's declared nullable type
even when the concrete input field is non-nullable.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
