andygrove commented on PR #5531:
URL:
https://github.com/apache/datafusion-comet/pull/5531#issuecomment-5458127105
I ran a post-merge pass over this with a local build, so everything below is
retrospective rather than a request to change anything here. Three follow-ups
filed: #5534, #5535 and #5536.
Starting with what held up, because I went looking for problems and mostly
did not find them. I checked every Celeborn member this reflects against the
real `apache/celeborn` sources, and they all line up. The 15-argument
`readPartition` matches positionally in 0.6.0, 0.6.3 and 0.7.0, including
`shuffleId` before `appShuffleId`, and `ReduceFileGroups.{partitionGroups,
mapAttempts, pushFailedBatches}` are public fields in all three. On 0.6.3 and
0.7.0 `CelebornShuffleHandle.stageRerunEnabled` exists on a class that really
does extend `BaseShuffleHandle`, and `object CelebornShuffleReader` really does
register the broadcast `GetReducerFileGroupResponse` deserializer in its static
initialiser, so the comment justifying that `loadClass` is accurate. Passing
`null` for `exceptionMaker` is right as well, since Celeborn's own
`reportShuffleFetchFailure` sits inside an `exceptionMaker != null` guard,
which makes `reportFetchFailure` here the single reporter rather than a double
one. Local
ly I ran 92 shuffle and 192 core Rust tests, 58 in
`CometCelebornShuffleReaderSuite` and 113 across the local shuffle suites on
Spark 4.1 / Scala 2.13 / JDK 17, all green.
On #5534, `getNextBatch` now releases the Arrow structs when `func` throws
and when native returns EOF, but the `importVector` call is still outside the
guard. If `importer.importVector` throws part-way through the column loop, the
structs from the failing column onward keep their populated C data with the
release callback never invoked, and the wrapper buffers leak from the allocator
too. I reproduced it with the suite's own `withIsolatedStructAllocator`: export
column 0, leave column 1 as a freshly allocated struct, and after the resulting
`IllegalStateException: Cannot import released ArrowSchema` the allocator still
reports 176 bytes outstanding. The success path is clean by comparison, 48
bytes before `batch.close()` and 0 after, so this is specific to the import
failure. The comment at `NativeBatchDecoderIterator.scala:105` already treats
Arrow import failures as a first-class case, so I think this one just got
missed. It is worth fixing beyond shuffle too, since `getNextBat
ch` backs every native operator and an allocator OOM inside `importVector` is
the most likely way to hit it.
On #5535, could a benchmark for the remote read path go into the 8/n
description? Dropping `with_skip_validation(true)` for remote frames is the
right call, but it is not cheap and the cost is not written down anywhere.
Release build, one `ShuffleBlockWriter` frame of 20 columns by 8192 rows, 300
iterations:
| codec | local (us/blk) | validated (us/blk) | ratio |
| -------- | -------------- | ------------------ | ----- |
| None | 38.8 | 175.2 | 4.52x |
| Lz4Frame | 225.8 | 356.6 | 1.58x |
| Snappy | 558.5 | 696.5 | 1.25x |
| Zstd(1) | 636.9 | 773.1 | 1.21x |
The local path itself did not regress, and measured against the pre-PR
implementation it is actually a little faster, from -8.8% on `None` to +1.1% on
Snappy, so the table above is purely the price of the new reader. Since the
whole reason to read Celeborn frames natively is to beat Celeborn's own row
decoder, that seems worth stating so 8/n can be judged against it. Two things
on the same path also look hoistable if it matters.
`decodeShuffleBlockWithValidation` re-decodes the `ShuffleScan` protobuf and
rebuilds `expected_types` on every block, and `env.convert_byte_array` copies
the serialized schema into a fresh `Vec` each time, and both are constant per
reducer.
On #5536, `validate_remote_schema` unwraps exactly one dictionary shape,
top-level `Dictionary(Int32, Utf8 | Binary)`, and treats everything else as a
mismatch. The local reader is much more permissive, since `unpack_dictionary`
casts any top-level dictionary to its value type and `cast_and_stamp_schema`
absorbs nested ones. So the two readers disagree about what a valid frame is,
and only the remote one turns the disagreement into a `FetchFailedException`
and a map-stage rerun that reproduces the same frame. I checked whether this is
live and it is not. I instrumented the decoder to log any wire schema
containing a dictionary at any depth and ran 125 tests across
`CometFuzzTestSuite`, `CometShuffleSuite` and `CometNativeShuffleSuite` with
zero hits, because `SchemaAlignExec` and `ScanExec::unpack_dictionary_type`
between them keep dictionaries off the wire entirely today. What worries me is
that nothing holds that in place. `ShuffleBlockWriter` genuinely encodes nested
dictionari
es, and on a frame the writer produced I get `Shuffle block type mismatch at
column 0: got List(Dictionary(Int32, Utf8)) but expected List(Utf8)` from the
validator while `arrow::compute::cast` handles the same frame fine. Would it be
worth either recursing the dictionary unwrap through the nested arms so the
validator accepts whatever the encoder can emit, or adding a test that pins
"the writer never emits a dictionary the reader would reject" so a drift breaks
the build instead of production?
Two smaller things I did not file. There are now three places that parse the
same 16-byte frame header with different rules: `NativeFrameInputStream` checks
the length floor, the ceiling, the field-count range and the expected field
count, `CometShuffleBlockIterator` checks the floor and ceiling and discards
the field count, and `NativeBatchDecoderIterator.readNextBlock` checks only the
ceiling. The third has no `compressedLength < 12` check, so a length below 8
reaches `dataBuf.limit(negative)` and throws `IllegalArgumentException` rather
than a classified corruption error. On the Celeborn path
`NativeFrameInputStream` catches it first so nothing is broken, but the local
path has no equivalent guard and three validators that are meant to agree tend
to stop agreeing. Separately, `NativeUtil.releaseArrowStructs` and
`NativeBatchDecoderIterator.close` now contain the same nine-line "collect the
first failure and `addSuppressed` the rest" helper, character for character
apart from wh
at they release. Both are new here and the series looks likely to want a
third, so a small helper in `org.apache.comet.util` next to `ClassLoaders`
might be worth it.
Last one is a process note for 8/n rather than anything about this code.
Three pieces of this PR change behaviour for every Comet user who will never
run Celeborn: the new strictness in `read_ipc_compressed` around empty streams,
multiple batches, trailing bytes and truncated LZ4, the `NativeUtil`
Arrow-struct release paths, and the `NativeBatchDecoderIterator` switch from an
eager constructor fetch to a lazy one plus no longer swallowing `EOFException`
as clean EOF. I ran the local shuffle and fuzz suites against all of it and
found no regression, so I do not think anything is broken. But that is 335
lines that mention Celeborn nowhere, the `NativeUtil` part is a standalone leak
fix with a standalone test, and the `ipc.rs` part tightens a path every query
touches. Would you consider peeling shared-path changes like those out of the
Celeborn series going forward? They would each get their own bisect point and
could be reviewed on their own terms.
--
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]
