sunchao commented on code in PR #6018: URL: https://github.com/apache/datafusion-comet/pull/6018#discussion_r4050194494
########## docs/source/contributor-guide/native_shuffle.md: ########## @@ -172,6 +194,63 @@ Native shuffle (`CometExchange`) is selected when all of the following condition 4. Arrow FFI transfers the `RecordBatch` to JVM as a `ColumnarBatch`. +## Direct Read (ShuffleScan) + +Direct read lets a native operator consume shuffle output without the batch ever being decoded in +the JVM or crossing Arrow FFI. It is controlled by `spark.comet.shuffle.directRead.enabled`, which +defaults to `true` and requires `spark.comet.shuffle.enabled`. It applies to both native shuffle and +JVM columnar shuffle, because both write the same Arrow IPC block format. + +### How the path is selected + +`CometExchangeSink.shouldUseShuffleScan` (`spark/src/main/scala/org/apache/comet/serde/operator/CometSink.scala`) +decides during plan serialization. When the sink's input is a Comet shuffle exchange, it emits a +`ShuffleScan` operator in place of the usual `Scan`. If any output type fails +`supportedSinkDataType`, it records the fallback reason `Unsupported data type for shuffle direct read` +and the slot serializes as a regular `Scan` instead. Review Comment: ### Correctness [P2] Describe the actual unsupported-type fallback When `supportedSinkDataType` rejects a shuffle output, [`convertToShuffleScan`](https://github.com/apache/datafusion-comet/blob/e2f054991f28c49f85e53e8bb4a35985ca297be6/spark/src/main/scala/org/apache/comet/serde/operator/CometSink.scala#L104-L140) returns `None`. `convert` does not retry `super.convert`. The ordinary `Scan` conversion also rejects the same types. The [AQE caller](https://github.com/apache/datafusion-comet/blob/e2f054991f28c49f85e53e8bb4a35985ca297be6/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala#L455-L464) retains the original Spark stage on that failure. This paragraph therefore gives readers the wrong native-conversion boundary when diagnosing unsupported schemas. Please describe the failed native conversion here, and reserve the regular `Scan` description for inputs/configurations where `shouldUseShuffleScan` is false. ########## .ai/skills/review-comet-shuffle-pr/SKILL.md: ########## @@ -0,0 +1,189 @@ +--- +name: review-comet-shuffle-pr +description: Use when reviewing a DataFusion Comet pull request that touches native or JVM columnar shuffle, the shuffle writers and readers, partitioning, the Arrow IPC block format, shuffle compression, or the Celeborn integration. Load alongside review-comet-pr. +argument-hint: <pr-number> +--- + +<!-- +Licensed to the Apache Software Foundation (ASF) under one +or more contributor license agreements. See the NOTICE file +distributed with this work for additional information +regarding copyright ownership. The ASF licenses this file +to you under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance +with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +--> + +Shuffle-specific review for Comet PR #$ARGUMENTS. + +**REQUIRED BACKGROUND:** Use `review-comet-pr` for PR metadata, existing comments, CI, the review +bar, and the output format. This skill only covers shuffle. + +## Read the Contributor Guide First + +| Doc | What you need from it | +| ---------------------------------------------------- | ---------------------------------------------------------------------- | +| `docs/source/contributor-guide/native_shuffle.md` | Selection rules, architecture, partitioning, block format, spilling | +| `docs/source/contributor-guide/jvm_shuffle.md` | Writer variants, handle selection, the row-based path, spill mechanics | +| `docs/source/contributor-guide/memory_management.md` | Where shuffle memory comes from, which differs between the two paths | + +**Read both shuffle docs even if the PR only touches one path.** The two implementations share the +manager, the dependency, the reader, and the on-disk format, and a change to one side of a shared +piece is the most common way to break the other. + +## 1. Which Implementation + +| Implementation | Selected when | +| ------------------------------------- | -------------------------------------------------------------------------------------------------------------- | +| Native, `CometExchange` | `shuffle.mode` is `native` or `auto`, child is a `CometPlan`, supported partitioning, primitive partition keys | +| JVM columnar, `CometColumnarExchange` | `shuffle.mode` is `jvm`, or the child is row-based, or partition keys are complex types | + +Complex types are fully supported as **data** columns in both. The primitive-only restriction +applies to **partition keys** for `HashPartitioning` and `RangePartitioning` only. Review Comment: ### Correctness [P2] Include the opt-in nested hash-key path The primitive-only rule is not unconditional for `HashPartitioning`. With [`spark.comet.shuffle.native.partitioning.hash.nested.enabled=true`](https://github.com/apache/datafusion-comet/blob/e2f054991f28c49f85e53e8bb4a35985ca297be6/spark/src/main/scala/org/apache/comet/CometConf.scala#L422-L436), the [type gate](https://github.com/apache/datafusion-comet/blob/e2f054991f28c49f85e53e8bb4a35985ca297be6/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala#L411-L462) admits supported structs and arrays recursively, and maps on Spark 4.0+ subject to normalization/expression support. The [existing config test](https://github.com/apache/datafusion-comet/blob/e2f054991f28c49f85e53e8bb4a35985ca297be6/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala#L908-L919) expects a native exchange when enabled and none when disabled. The table and this checklist would classify a valid, covered native plan as an incorrect fallback decisi on. Please qualify the default-disabled nested-hash path and distinguish it from the complex-key restriction for range partitioning. ########## .ai/skills/review-comet-shuffle-pr/SKILL.md: ########## @@ -0,0 +1,189 @@ +--- +name: review-comet-shuffle-pr +description: Use when reviewing a DataFusion Comet pull request that touches native or JVM columnar shuffle, the shuffle writers and readers, partitioning, the Arrow IPC block format, shuffle compression, or the Celeborn integration. Load alongside review-comet-pr. +argument-hint: <pr-number> +--- + +<!-- +Licensed to the Apache Software Foundation (ASF) under one +or more contributor license agreements. See the NOTICE file +distributed with this work for additional information +regarding copyright ownership. The ASF licenses this file +to you under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance +with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +--> + +Shuffle-specific review for Comet PR #$ARGUMENTS. + +**REQUIRED BACKGROUND:** Use `review-comet-pr` for PR metadata, existing comments, CI, the review +bar, and the output format. This skill only covers shuffle. + +## Read the Contributor Guide First + +| Doc | What you need from it | +| ---------------------------------------------------- | ---------------------------------------------------------------------- | +| `docs/source/contributor-guide/native_shuffle.md` | Selection rules, architecture, partitioning, block format, spilling | +| `docs/source/contributor-guide/jvm_shuffle.md` | Writer variants, handle selection, the row-based path, spill mechanics | +| `docs/source/contributor-guide/memory_management.md` | Where shuffle memory comes from, which differs between the two paths | + +**Read both shuffle docs even if the PR only touches one path.** The two implementations share the +manager, the dependency, the reader, and the on-disk format, and a change to one side of a shared +piece is the most common way to break the other. + +## 1. Which Implementation + +| Implementation | Selected when | +| ------------------------------------- | -------------------------------------------------------------------------------------------------------------- | +| Native, `CometExchange` | `shuffle.mode` is `native` or `auto`, child is a `CometPlan`, supported partitioning, primitive partition keys | +| JVM columnar, `CometColumnarExchange` | `shuffle.mode` is `jvm`, or the child is row-based, or partition keys are complex types | + +Complex types are fully supported as **data** columns in both. The primitive-only restriction +applies to **partition keys** for `HashPartitioning` and `RangePartitioning` only. + +- [ ] A PR that widens what native shuffle supports updates the fallback conditions in + `CometShuffleExchangeExec` **and** both docs' "When X is Used" lists +- [ ] A PR that narrows support does not silently move workloads onto the slower path. The JVM path + costs a columnar to row to columnar round trip through `ColumnarToRowExec`. +- [ ] Fallback decisions stay consistent across a stage. `CometShuffleFallbackStickinessSuite` + exists because they did not once. + +## 2. Spark Compatibility of Partitioning + +Partitioning is where shuffle silently produces wrong answers rather than failing. + +- [ ] **Hash partitioning uses Murmur3 with seed 42** and `partition_id = hash % num_partitions`, + matching Spark. Any change to the hash, the seed, or the modulo changes which rows land in + which partition, which breaks a join between a Comet-shuffled side and a Spark-shuffled side. +- [ ] **Round robin is hash-based on purpose.** Comet assigns partitions from a Murmur3 hash rather + than cycling row by row, because determinism across task retries is required for correctness + under fault tolerance. A PR that implements "true" round robin to fix skew breaks that. The + known cost is that low-cardinality data distributes unevenly, and that is the accepted + trade-off. +- [ ] **Range partitioning bounds come from the driver.** Spark's `RangePartitioner` samples and + computes boundaries, they are serialized into the native plan, and native does a binary + search over comparable-row-format keys. A change to the comparison or the row encoding must + match Spark's ordering exactly, including nulls and signed zero. +- [ ] The JVM path uses Spark's own partitioner via `partitioner.getPartition(key)`, so it inherits + Spark's semantics for free. A PR that reimplements partitioning on that path is solving a + problem that does not exist. + +## 3. On-Disk and On-Wire Format + +Writer and reader must change together, and they are in different languages. + +The block layout is an 8-byte compressed length header, an 8-byte field count header, then the +compressed Arrow IPC stream. It is written by the native `ShuffleBlockWriter` and read by +`NativeBatchDecoderIterator` calling `Native.decodeShuffleBlock()`. + +- [ ] A format change updates the writer, the reader, and the Celeborn reader path +- [ ] A format change is not silently incompatible with shuffle files written by a previous version + in the same cluster during a rolling deployment. If it is, the PR needs to say so. +- [ ] Compression codec changes apply uniformly to all partitions, and each partition stays + independently decompressible so reads can parallelize +- [ ] The commit path still works. Native records the byte offset where each partition begins plus + the total length, `CometNativeShuffleWriter` fetches them with + `Native.getShufflePartitionOffsets`, converts them to partition lengths, and commits through + Spark's `IndexShuffleBlockResolver.writeMetadataFileAndCommit`. Offsets and lengths are easy + to confuse and the failure is a corrupt index file rather than an exception. +- [ ] Checksums via `CometShuffleChecksumSupport` still cover what Spark expects + +## 4. Memory and Spilling + +Shuffle is the largest memory consumer in most queries, and the two paths draw from different +budgets. + +**Native shuffle** uses the DataFusion memory pool. Partitions spill when the pool denies an +allocation, or when buffered bytes reach `spark.comet.shuffle.native.maxBufferBytes`, which +defaults to `0`, meaning the fixed limit is disabled and memory pressure is the only trigger. Each +partition has its own spill file and multiple spills for a partition are concatenated when the +final output is written. Review Comment: ### Correctness [P2] Document one shared spill file for local native shuffle The current local writer creates one [`PartitionedSpill`](https://github.com/apache/datafusion-comet/blob/e2f054991f28c49f85e53e8bb4a35985ca297be6/native/shuffle/src/writers/local/local_partition_writer.rs#L127-L146) for all output partitions. It owns a single spill file and records [per-partition byte ranges](https://github.com/apache/datafusion-comet/blob/e2f054991f28c49f85e53e8bb4a35985ca297be6/native/shuffle/src/writers/local/spill.rs#L54-L66). the existing [`spilling_every_partition_creates_one_file` test](https://github.com/apache/datafusion-comet/blob/e2f054991f28c49f85e53e8bb4a35985ca297be6/native/shuffle/src/writers/local/local_partition_writer.rs#L615-L641) explicitly asserts one file after repeatedly spilling 64 partitions. Teaching reviewers to expect a file per partition misstates the resource ownership and spill lifecycle they need to check. Please describe the shared file and ordered per-partition ranges, and make the same correction to the newly added `spill.rs` ta ble entry and spill paragraph in `native_shuffle.md`. ########## .ai/skills/review-comet-ffi-pr/SKILL.md: ########## @@ -0,0 +1,159 @@ +--- +name: review-comet-ffi-pr +description: Use when reviewing a DataFusion Comet pull request that crosses the JVM/native boundary, touching Arrow C Data or C Stream interface code, batch export and import, CometExecIterator, ScanExec, NativeUtil, CometVector subclasses, or jni_api. Load alongside review-comet-pr. +argument-hint: <pr-number> +--- + +<!-- +Licensed to the Apache Software Foundation (ASF) under one +or more contributor license agreements. See the NOTICE file +distributed with this work for additional information +regarding copyright ownership. The ASF licenses this file +to you under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance +with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +--> + +FFI-specific review for Comet PR #$ARGUMENTS. + +**REQUIRED BACKGROUND:** Use `review-comet-pr` for PR metadata, existing comments, CI, the review +bar, and the output format. This skill only covers the JVM/native boundary. + +Bugs in this area do not produce wrong answers, they produce segfaults, leaks, and use-after-free +under load, often only on one platform or only when an operator buffers batches. Review it with +that in mind. + +## Read the Contributor Guide First + +| Doc | What you need from it | +| ---------------------------------------------------- | --------------------------------------------------------------------------- | +| `docs/source/contributor-guide/ffi.md` | Both data-flow directions, ownership rules, lifecycle, alignment workaround | +| `docs/source/contributor-guide/memory_management.md` | The "Crossing the FFI boundary" section: who is charged for a batch's bytes | + +Read `ffi.md` in full before the diff. The two directions have different ownership semantics and +reviewing one with the other's mental model is the most common way to miss a bug. + +## The Two Directions + +| Direction | Mechanism | Who owns the data | +| ----------------------------------- | ------------------------------------------ | ------------------------------------------------------------------ | +| JVM to native (`ScanExec`) | Arrow C **Stream**, one per partition | Native takes ownership by reference count when it imports a batch | +| Native to JVM (`CometExecIterator`) | Arrow C **Data**, one array pair per batch | Native allocates, JVM holds pointers and must `close()` to release | + +## 1. Ownership and Lifetime + +- [ ] **JVM to native: no defensive deep copies.** The C Stream transfers ownership by reference + count, so native can buffer imported batches in `SortExec` or the shuffle writer without + copying. A new `.clone()` of the data, a `copy_array`, or a "to be safe" deep copy on this + path is a real throughput cost. Ask what it is protecting against. +- [ ] **JVM to native: dropping the reader is the release.** When `ScanExec` drops its + `AlignedArrowStreamReader`, the stream's release callback fires synchronously back into the + JVM and closes the `ArrowReader` and its `VectorSchemaRoot`. Anything that extends the + reader's lifetime, stores it somewhere longer-lived, or drops it early changes when JVM + off-heap buffers are freed. +- [ ] **JVM to native: buffering pins JVM memory.** An operator that holds many imported batches + keeps the corresponding JVM-side off-heap buffers alive. A change that makes an operator + buffer more is also a memory change. +- [ ] **Native to JVM: every export needs a matching close.** Native allocates, the JVM wraps the + pointers in `ArrowBuf`s, and the bytes are freed only when the JVM calls `close()`. Trace the + new export to the `close()` that releases it, including on the exception path. +- [ ] **Release callbacks run on the error path too.** A batch exported and then abandoned because + the query failed still has to be released. Check the failure and cancellation paths, not just + the happy path. +- [ ] **No unwinding across `extern "C"`.** A Rust panic crossing the FFI boundary is undefined + behavior. New `#[no_mangle] extern "system"` entry points must not let a panic escape, and + must not `unwrap()` on anything an input can make fail. +- [ ] **Null and error checks on every pointer received from the other side.** + +## 2. The Stream Design + +The JVM exports each per-partition iterator **once** as an `ArrowArrayStream`, and native pulls +every batch through the stream's `get_next` callback. There is no per-batch JNI call and no +per-column FFI export on this path. + +A change that reintroduces per-batch or per-column export on the JVM-to-native path is a +performance regression even if it is correct. Flag it and ask why the stream could not carry it. + +The reader implementations in `CometNativeArrowSource.scala` are `RowArrowReader` for +`Iterator[InternalRow]`, `SparkColumnarArrowReader` for a non-Arrow `ColumnarBatch`, and +`ColumnarBatchArrowReader` for an Arrow-backed `ColumnarBatch`, which transfers `VectorSchemaRoot` +ownership. A new input shape needs a reader, not a special case elsewhere. + +## 3. Vector Types and Export Dispatch + +`NativeUtil.exportBatch()` matches on the concrete vector type. The `CometVector` hierarchy is +`CometDecodedVector` with Plain, Dictionary, List, Map, and Struct subclasses, plus +`CometSelectionVector` and `CometDelegateVector`. + +- [ ] A new `CometVector` subclass has a case in `exportBatch()` +- [ ] The case ordering is right. `CometSelectionVector` must be matched **before** the general + `CometVector` case, or the selection is silently dropped and the exported batch has the wrong + rows. +- [ ] Selection vectors are applied where `scan.rs` expects them, in `ScanExec::get_next()` Review Comment: ### Correctness [P2] Replace obsolete selection-vector review checks Neither `CometSelectionVector` nor `CometDelegateVector` exists in the current Spark/common/native sources. [`NativeUtil.exportBatch`](https://github.com/apache/datafusion-comet/blob/e2f054991f28c49f85e53e8bb4a35985ca297be6/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala#L119-L165) dispatches directly to `CometVector` or `ConstantColumnVector`, and [`ScanExec`](https://github.com/apache/datafusion-comet/blob/e2f054991f28c49f85e53e8bb4a35985ca297be6/native/core/src/execution/operators/scan.rs#L115-L175) imports batches through `get_next_batch`/`pull_next`, with no `get_next()` selection-vector path. These mandatory checks would tell the reviewer to flag the current generic `CometVector` case as dropping rows and to inspect a nonexistent import hook. Please update this section and the preceding hierarchy to the current export/import model. -- 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]
