mbutrovich commented on code in PR #5004:
URL: https://github.com/apache/datafusion-comet/pull/5004#discussion_r3640934789
##########
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:
The "independent of input chunking" in the PR description reasoning holds
only while `biggest_coalesce_batch_size` is unset. Given #5003 sets it, would
it be worth dropping the chunk-independence phrasing so this comment and #5003
do not disagree about what the coalescer guarantees?
##########
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:
Asserts `[100, 100, 100, 50]`. I traced it under #5003's bypass too and it
still passes, because the 250-row batch arrives on a 60-row buffer and takes
the "small buffer, coalesce normally" branch (`coalesce.rs:471-475`) rather
than the passthrough. So neither PR's test exercises the divergent case, a `>=
batch_size` batch on an empty buffer. A test with the large batch first would
pin down the intended 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:
Good cleanup, `new` is infallible now that concat is gone. Quick check that
`DataFusionError` is unused elsewhere in the file, though lint will catch it.
--
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]
