kosiew commented on code in PR #24979:
URL: https://github.com/apache/datafusion/pull/24979#discussion_r3987827160
##########
datafusion/physical-plan/src/sorts/partial_sort.rs:
##########
@@ -583,50 +798,142 @@ impl PartialSortStream {
}
}
- /// Returns a sorted RecordBatch from in_mem_batches and clears
in_mem_batches
- ///
- /// If fetch is specified for PartialSortStream `sort_in_mem_batch` will
limit
- /// the last RecordBatch returned and will mark the stream as closed
- fn sort_in_mem_batch(self: &mut Pin<&mut Self>) -> Result<RecordBatch> {
- let input_batch = self.in_mem_batch.clone();
- self.in_mem_batch = RecordBatch::new_empty(self.schema());
- let result = sort_batch(&input_batch, &self.expr, self.fetch)?;
+ fn sort_completed_prefixes(
+ self: &mut Pin<&mut Self>,
+ completed: CompletedPrefixes,
+ ) -> Result<RecordBatch> {
+ let result = if completed.num_prefixes() == 1 {
+ self.sort_single_prefix(completed.into_batches())?
+ } else {
+ self.sort_multiple_prefixes(&completed)?
+ };
+
if let Some(remaining_fetch) = self.fetch {
- // remaining_fetch - result.num_rows() is always be >= 0
- // because result length of sort_batch with limit cannot be
- // more than the requested limit
self.fetch = Some(remaining_fetch - result.num_rows());
- if remaining_fetch == result.num_rows() {
- self.is_closed = true;
- }
}
Ok(result)
}
- /// Return the end index of the second last partition if the batch
- /// can be partitioned based on its already sorted columns
- ///
- /// Return None if the batch cannot be partitioned, which means the
- /// batch does not have the information for a safe sort
- fn get_slice_point(
+ fn sort_single_prefix(
&self,
- common_prefix_len: usize,
- batch: &RecordBatch,
- ) -> Result<Option<usize>> {
- let common_prefix_sort_keys = (0..common_prefix_len)
- .map(|idx| self.expr[idx].evaluate_to_sort_column(batch))
+ completed_batches: Vec<RecordBatch>,
+ ) -> Result<RecordBatch> {
+ // A single prefix may span batches, so concatenate it at most once.
+ let batch = if completed_batches.len() == 1 {
+ completed_batches.into_iter().next().unwrap()
+ } else {
+ concat_batches(&self.schema(), &completed_batches)?
+ };
+
+ if let Some(suffix_ordering) = &self.suffix_ordering {
+ sort_batch(&batch, suffix_ordering, self.fetch)
+ } else {
+ let row_count = self
+ .fetch
+ .unwrap_or_else(|| batch.num_rows())
+ .min(batch.num_rows());
+ Ok(batch.slice(0, row_count))
+ }
+ }
+
+ fn sort_multiple_prefixes(
+ &self,
+ completed: &CompletedPrefixes,
+ ) -> Result<RecordBatch> {
+ // Evaluate suffix expressions once per source batch, then sort each
+ // prefix independently and materialize the result with one interleave.
+ let sort_columns_by_batch = completed
+ .batches
+ .iter()
+ .map(|batch| {
+ self.suffix_ordering
+ .iter()
+ .flat_map(|exprs| exprs.iter())
+ .map(|expr| expr.evaluate_to_sort_column(batch))
+ .collect::<Result<Vec<_>>>()
+ })
.collect::<Result<Vec<_>>>()?;
- let partition_points =
- evaluate_partition_ranges(batch.num_rows(),
&common_prefix_sort_keys)?;
- // If partition points are [0..100], [100..200], [200..300]
- // we should return 200, which is the safest and furthest partition
boundary
- // Please note that we shouldn't return 300 (which is number of rows
in the batch),
- // because this boundary may change with new data.
- if partition_points.len() >= 2 {
- Ok(Some(partition_points[partition_points.len() - 2].end))
+ let mut remaining_fetch = self.fetch.unwrap_or(usize::MAX);
+ let mut interleave_indices = vec![];
+
+ for prefix_range in completed.prefix_ranges() {
+ if remaining_fetch == 0 {
+ break;
+ }
+
+ let row_count = prefix_range.len();
+ let prefix_fetch = remaining_fetch.min(row_count);
+ let sorted_indices = if let Some(suffix_ordering) =
&self.suffix_ordering {
+ let sort_columns = completed.gather_sort_columns(
+ prefix_range.clone(),
+ &sort_columns_by_batch,
+ suffix_ordering,
+ )?;
+ let fetch = (prefix_fetch < row_count).then_some(prefix_fetch);
+ lexsort_to_indices(&sort_columns, fetch)?
+ .values()
+ .iter()
+ .map(|idx| *idx as usize)
+ .collect::<Vec<_>>()
+ } else {
+ (0..prefix_fetch).collect()
+ };
+
+ interleave_indices.extend(
+ sorted_indices
+ .into_iter()
+ .map(|index| completed.source_row(prefix_range.start +
index)),
+ );
+ remaining_fetch -= prefix_fetch;
+ }
+
+ if completed.batches[0].num_columns() == 0 {
+ let options =
+
RecordBatchOptions::new().with_row_count(Some(interleave_indices.len()));
+ Ok(RecordBatch::try_new_with_options(
+ self.schema(),
+ vec![],
+ &options,
+ )?)
} else {
- Ok(None)
+ let completed_batches =
completed.batches.iter().collect::<Vec<_>>();
+ Ok(interleave_record_batch(
+ &completed_batches,
+ &interleave_indices,
+ )?)
+ }
+ }
+
+ fn get_prefix_ranges(&self, batch: &RecordBatch) ->
Result<Vec<Range<usize>>> {
+ let common_prefix_sort_keys = (0..self.common_prefix_length)
+ .map(|idx| self.expr[idx].evaluate_to_sort_column(batch))
+ .collect::<Result<Vec<_>>>()?;
+ evaluate_partition_ranges(batch.num_rows(), &common_prefix_sort_keys)
+ }
+
+ fn prefix_changed_at_batch_boundary(
Review Comment:
Could we add a cross-batch regression test where the prefix key contains
repeated NULL and NaN values, including `nulls_first` and descending sort
options? It would be useful to compare the result against `SortExec` as well.
`prefix_changed_at_batch_boundary` now determines boundary equality
independently using `make_comparator`, while the in-batch groups come from
`evaluate_partition_ranges`, so a test like this would help make sure both
paths stay consistent for these special values.
--
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]
