andygrove commented on code in PR #7234:
URL: https://github.com/apache/datafusion-comet/pull/7234#discussion_r4404671549
##########
native/shuffle/src/writers/buf_batch_writer.rs:
##########
@@ -118,7 +154,15 @@ impl<S: Borrow<ShuffleBlockWriter>, W: Write>
BufBatchWriter<S, W> {
write_time: &Time,
) -> datafusion::common::Result<usize> {
self.check_scratch(scratch);
- let batch_size = self.batch_size;
+ let Some(batch_size) = self.coalesce_batch_size else {
+ return self.write_batch_to_buffer(
+ batch,
+ scratch,
+ codec_context,
+ encode_time,
+ write_time,
+ );
+ };
let coalescer = self.coalescer.get_or_insert_with(|| {
// Enable BatchCoalescer's zero-copy passthrough for batches that
are already big
Review Comment:
You're right, the old wording described `PartitionedBatchIterator` output,
which no longer reaches this path. I rewrote the comment around the only
remaining caller, the single-partition path: it streams its input batches
straight in and routes `>= batch_size` batches here directly, so those are
written as one oversized block rather than split at `batch_size` (6fe4add2d0).
##########
native/shuffle/src/writers/partition_writer.rs:
##########
@@ -54,6 +54,10 @@ pub(crate) trait PartitionWriter: Send {
///
/// Must be called exactly once per partition, in ascending id order, so
the
/// writer can lay partitions out contiguously and record their offsets.
+ ///
+ /// `iter` should yield maximal `batch_size` chunks plus at most one
smaller tail, as
+ /// `PartitionedBatchIterator` does. [`LocalPartitionWriter`] writes each
batch as its own
+ /// block in multi-partition mode, so many small batches would become many
small blocks.
Review Comment:
Agreed, `write` has the same property, since `PartitionedSpill::write` also
writes each batch as its own block and spilled blocks are copied to the output
as they are. I moved the note to the `PartitionWriter` trait docs so it covers
both `write` and `finish_partition`, and it now says why: many small batches
become many small blocks (6fe4add2d0).
##########
native/shuffle/src/writers/local/local_partition_writer.rs:
##########
@@ -502,6 +527,52 @@ mod tests {
}
}
+ /// Partition offsets are tracked arithmetically instead of read back from
the output
+ /// file, so the index must still describe the data file exactly: every
offset lands on a
+ /// block boundary and each partition's blocks decode to the rows written
for it, with a
+ /// spilled prefix and in-memory tail both accounted for.
+ #[test]
+ fn offsets_match_data_file_with_spilled_and_in_memory_batches() {
Review Comment:
You're right that `spilled_partitions_read_back_in_write_order` already
covers the spilled prefix plus in-memory tail with arithmetic offsets, so my
test added nothing. I replaced it with
`offsets_match_data_file_for_spill_only_and_empty_partitions`. It uses four
partitions: a spill-only partition, an empty one, an in-memory-only one, and an
empty last one. It asserts that empty partitions have zero-length ranges, that
the final offset equals the file length, and that every non-empty range decodes
to the expected row count (6fe4add2d0).
##########
native/shuffle/src/writers/local/local_partition_writer.rs:
##########
@@ -345,11 +358,23 @@ impl PartitionWriter for LocalPartitionWriter {
)?;
writer.writer_stream_position()?
}
- DataOutput::Multi { output_writer, .. } => {
+ DataOutput::Multi {
+ output_writer,
+ bytes_written,
+ ..
+ } => {
let mut write_timer = metrics.write_time.timer();
output_writer.flush()?;
let pos = output_writer.stream_position()?;
write_timer.stop();
+ // The offsets were derived from the running total; the file
must agree,
+ // or the index would point readers at the wrong bytes.
+ if pos != *bytes_written {
+ return Err(DataFusionError::Execution(format!(
+ "shuffle write error: data file holds {pos} bytes but
the partition \
+ offsets account for {bytes_written}"
+ )));
+ }
Review Comment:
Good call, nothing exercised it. I added
`finish_all_fails_when_data_file_disagrees_with_offsets`, following your
recipe. It uses a 64-byte write buffer and fails partition 0 after its first
block reached the output. It then finishes partition 1 and checks that
`finish_all` returns the "partition offsets account for" error and that no
offsets are published (6fe4add2d0).
--
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]
