kosiew commented on code in PR #25358:
URL: https://github.com/apache/datafusion/pull/25358#discussion_r4257400598
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -957,6 +1033,144 @@ impl Debug for ExternalSorter {
}
}
+struct SortStage {
+ columns: Vec<usize>,
+ expressions: Vec<PhysicalSortExpr>,
+}
+
+fn staged_sort_batch_chunked(
+ batch: &RecordBatch,
+ expressions: &LexOrdering,
+ batch_size: usize,
+ stages: &[SortStage],
+) -> Result<Vec<RecordBatch>> {
+ let indices = staged_sort_indices(batch, stages)?;
+ IncrementalSortIterator::new(batch.clone(), expressions.clone(),
batch_size)
+ .with_sorted_indices(indices)
+ .collect()
+}
+
+/// Compute a complete permutation by evaluating successive key groups only
for ties.
+fn staged_sort_indices(batch: &RecordBatch, stages: &[SortStage]) ->
Result<UInt32Array> {
+ let mut indices = Vec::new();
+ // Ranges address positions in `indices`, whose values address the input
batch.
+ let mut unresolved: Vec<_> =
std::iter::once(0..batch.num_rows()).collect();
+ for (index, stage) in stages.iter().enumerate() {
+ if unresolved.is_empty() {
+ break;
+ }
+ unresolved = refine_tied_groups(
+ batch,
+ stage,
+ &mut indices,
+ &unresolved,
+ index == 0,
+ index + 1 < stages.len(),
+ )?;
+ }
+
+ Ok(UInt32Array::from(indices))
+}
+
+/// Evaluate keys once for all unresolved rows, then sort each tied group
independently.
+fn refine_tied_groups(
+ batch: &RecordBatch,
+ stage: &SortStage,
+ indices: &mut Vec<u32>,
+ unresolved: &[Range<usize>],
+ first_stage: bool,
+ find_ties: bool,
+) -> Result<Vec<Range<usize>>> {
+ let selected = if first_stage {
+ batch.clone()
+ } else {
+ let mut selection =
Vec::with_capacity(unresolved.iter().map(Range::len).sum());
+ for range in unresolved {
+ selection.extend_from_slice(&indices[range.clone()]);
+ }
+ arrow::compute::take_record_batch(
+ &batch.project(&stage.columns)?,
Review Comment:
When a later stage uses a column-free expression, `stage.columns` is empty
and `take_record_batch` fails because the empty-column batch loses its row
count. Could we preserve the selected row count using
`RecordBatchOptions::with_row_count(Some(selection.len()))` or retain the
full-batch gather for this case, and add tests with a zero-argument UDF and
tied keys for both single and concatenated batches?
--
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]
