jayzhan211 commented on code in PR #24392:
URL: https://github.com/apache/datafusion/pull/24392#discussion_r3791213948
##########
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'm not sure about this casting — should it happen inside the accumulator,
or should the types already be correct before they enter it, like at the
logical layer, so we don't need to cast again here?
##########
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:
If the array isn't filtered much — barely anything changes, or only 1-2
items get removed — would copying the whole array every time still cause high
overhead?
--
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]
