andygrove commented on code in PR #5004:
URL: https://github.com/apache/datafusion-comet/pull/5004#discussion_r3645542426
##########
native/shuffle/src/partitioners/single_partition.rs:
##########
@@ -19,66 +19,27 @@ use crate::metrics::ShufflePartitionerMetrics;
use crate::partitioners::ShufflePartitioner;
use crate::writers::PartitionWriter;
use arrow::array::RecordBatch;
-use datafusion::common::DataFusionError;
use std::iter;
use tokio::time::Instant;
-/// A partitioner that writes all shuffle data to a single file and a single
index file
+/// A partitioner that writes all shuffle data to a single file and a single
index file.
Review Comment:
Reworded the comment. It now says a `>= batch_size` batch on an empty buffer
is passed through verbatim as one (possibly oversized) block and that block
boundaries depend on how the input is chunked, so it no longer disagrees with
#5003. Dropped the same chunk-independence phrasing from the PR description.
##########
native/shuffle/src/shuffle_writer.rs:
##########
@@ -924,6 +923,86 @@ mod test {
.sum()
}
+ /// The single-partition partitioner streams batches straight to the
writer's
+ /// `BatchCoalescer` without any concat/buffer layer of its own. Feeding
batches
+ /// of varying sizes must still produce `batch_size`-row IPC blocks (plus
a tail)
+ /// with every row roundtripping in order.
+ #[tokio::test]
+ #[cfg_attr(miri, ignore)]
+ async fn single_partition_coalesces_to_batch_size_in_order() {
Review Comment:
Reworked the test to lead with a 250-row batch on an empty buffer, so it now
takes the passthrough branch. It asserts blocks `[250, 100, 30]`: the 250-row
block can only appear if the bypass fires (plain coalescing would split it into
100/100/50), which pins down the post-#5003 behavior instead of leaving it to
buffer state.
##########
native/shuffle/src/partitioners/single_partition.rs:
##########
@@ -19,66 +19,27 @@ use crate::metrics::ShufflePartitionerMetrics;
use crate::partitioners::ShufflePartitioner;
use crate::writers::PartitionWriter;
use arrow::array::RecordBatch;
-use datafusion::common::DataFusionError;
use std::iter;
use tokio::time::Instant;
-/// A partitioner that writes all shuffle data to a single file and a single
index file
+/// A partitioner that writes all shuffle data to a single file and a single
index file.
+///
+/// Batches are streamed straight to the long-lived `BufBatchWriter` inside the
+/// [`PartitionWriter`], whose internal `BatchCoalescer` coalesces small
batches into
+/// `batch_size`-row IPC blocks across calls. The partitioner therefore does
no buffering
+/// or concatenation of its own; doing so would copy every row an extra time
before the
+/// coalescer copies it again.
pub(crate) struct SinglePartitionShufflePartitioner<T: PartitionWriter> {
partition_writer: T,
- /// Batches that are smaller than the batch size and to be concatenated
- buffered_batches: Vec<RecordBatch>,
- /// Number of rows in the concatenating batches
- num_buffered_rows: usize,
/// Metrics for the repartitioner
metrics: ShufflePartitionerMetrics,
- /// The configured batch size
- batch_size: usize,
}
impl<T: PartitionWriter> SinglePartitionShufflePartitioner<T> {
- pub(crate) fn try_new(
Review Comment:
Confirmed. The `use datafusion::common::DataFusionError` import was removed
along with `concat_buffered_batches`, so nothing else in the file references
it, and clippy is clean.
--
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]
