sunchao commented on code in PR #6018: URL: https://github.com/apache/datafusion-comet/pull/6018#discussion_r4050194506
########## .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`. -- 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]
