mbutrovich commented on code in PR #7234:
URL: https://github.com/apache/datafusion-comet/pull/7234#discussion_r4391647561
##########
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:
Could we add a test that reaches this error? Nothing in the PR makes `pos`
and `bytes_written` disagree, so the guard is untested.
As I read `finish_partition`, a partition that fails after its first block
reached `output_writer` gets there. With `partition_writer_with(&batch, 2, 64,
...)`, the 64-byte write buffer hands the first block to `output_writer` as
soon as it is encoded. Then an injected `Err` from the iterator, as in
[`finish_partition_error_drains_recycled_buffer`](https://github.com/apache/datafusion-comet/blob/c5ce2b151da606bc3a8a34b93f8b307556b3f5ed/native/shuffle/src/writers/local/local_partition_writer.rs#L502-L528),
returns before line 339 adds that partition's bytes to the total. Finishing
partition 1 and calling `finish_all` should then return this error and publish
no offsets.
##########
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:
Should this go in the trait docs instead, so it covers `write` too? In
multi-partition mode `LocalPartitionWriter::write` hands its batches to
`PartitionedSpill::write`, which now also writes each batch as its own block
([spill.rs](https://github.com/apache/datafusion-comet/blob/c5ce2b151da606bc3a8a34b93f8b307556b3f5ed/native/shuffle/src/writers/local/spill.rs#L89-L93)).
A caller that stages many small batches through `write` would get many small
spilled blocks, and those are copied to the output as they are.
##########
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:
What does this test cover that
[`spilled_partitions_read_back_in_write_order`](https://github.com/apache/datafusion-comet/blob/c5ce2b151da606bc3a8a34b93f8b307556b3f5ed/native/shuffle/src/writers/local/local_partition_writer.rs#L618-L673)
doesn't? That test already spills every partition, finishes each with an
in-memory batch, and decodes each partition's byte range against the expected
rows, with both a 1 MiB and a 64-byte write buffer. With the 1 MiB buffer the
output is no longer flushed between partitions, so it already checks the
arithmetic offsets.
The cases neither test has are the boundaries of the running total: a
partition with only spilled data and no in-memory batch, and an empty partition
between two non-empty ones. Should this test cover those instead?
--
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]
