lyne7-sc commented on code in PR #24392:
URL: https://github.com/apache/datafusion/pull/24392#discussion_r3791990068
##########
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)?
+ } else {
+ return exec_err!(
+ "ordered array_agg payload has type {}, expected {}",
+ values.data_type(),
+ self.value_type
+ );
+ };
+ if let Some(column) = ordering_values.first() {
+ assert_eq_or_internal_err!(
+ column.len(),
+ values.len(),
+ "ordered array_agg payload and ordering columns must have
equal lengths"
+ );
}
- let val = &values[0];
- let ord = &values[1..];
- let nulls = if self.ignore_nulls {
- val.logical_nulls()
+ let nulls = ignore_nulls
+ .then(|| values.logical_nulls())
+ .flatten()
+ .filter(|nulls| nulls.null_count() > 0);
+ let (values, filtered_ordering_values) = if let Some(nulls) = nulls {
+ let mask: BooleanArray = nulls.iter().map(Some).collect();
+ let values = filter(values.as_ref(), &mask)?;
+ let ordering_values = ordering_values
+ .iter()
+ .map(|column| filter(column.as_ref(), &mask))
+ .collect::<std::result::Result<Vec<_>, _>>()?;
+ (values, Some(ordering_values))
} else {
- None
+ (values, None)
};
-
- let nulls = nulls.as_ref();
- if nulls.is_none_or(|nulls| nulls.null_count() < val.len()) {
- for i in 0..val.len() {
- if nulls.is_none_or(|nulls| nulls.is_valid(i)) {
- self.values
- .push(ScalarValue::try_from_array(val,
i)?.compacted());
- self.ordering_values.push(
- get_row_at_idx(ord, i)?
- .into_iter()
- .map(|v| v.compacted())
- .collect(),
- )
- }
- }
+ let ordering_values = filtered_ordering_values
+ .as_deref()
+ .unwrap_or(ordering_values);
+ // Detach the stored payload from potentially oversized backing
buffers.
+ let values = make_array(copy_array_data(&values.to_data()));
Review Comment:
Yes, I had the same concern. When I tried removing this copy locally and ran
[`does_not_over_account_memory_ordered`](https://github.com/apache/datafusion/blob/bc03e0c1f26a01db4d86ca84835440d270e4e73f/datafusion/functions-aggregate/src/array_agg.rs#L1912-L1934),
the accumulator's reported size increased from 2295 to 7243 bytes.
This makes me concerned that removing it could introduce a memory regression
in cases where retained ArrayRefs keep much larger backing buffers alive,
especially given the related memory explosion reported in #20788 with the
previous ordered ARRAY_AGG implementation.
On the performance side, I also compared the two variants using the
benchmark, and the difference was relatively small for this workload.
It may still be worth exploring a conditional policy, although there does
not seem to be a simple way to determine whether an arbitrary nested arrow
array is already compact.
--
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]
