sunchao commented on code in PR #7151: URL: https://github.com/apache/datafusion-comet/pull/7151#discussion_r4366956694
########## docs/source/contributor-guide/jvm_shuffle.md: ########## @@ -65,7 +65,7 @@ returns an `RDD[InternalRow]`. The rows flow directly to the JVM shuffle writers ### Comet Columnar Input -When the child plan is a Comet native operator (e.g., `CometHashAggregate`) but JVM shuffle is selected +When the child plan is a Comet operator (e.g., `CometHashAggregate`) but JVM shuffle is selected Review Comment: [P2] Broadening this paragraph to any Comet operator makes the conversion diagram incorrect for `CometSparkToColumnarExec` over Spark rows. In JVM shuffle mode, `child.execute()` forwards the underlying rows directly. The diagram assigns an Arrow-to-row conversion cost to a path that produces no input Arrow batch. Scope the diagram to children using `CometExec.doExecute()` and describe the row-forwarding bridge separately. Evidence: A Spark 4.1.3 probe requested `SinglePartition` over `CometSparkToColumnarExec(RDDScanExec)` with `spark.comet.shuffle.mode=jvm`. The planner selected `CometColumnarExchange`. Consuming its `inputRDD` returned both rows while the bridge’s `numOutputBatches` remained zero. [Bridge implementation](https://github.com/apache/datafusion-comet/blob/56396e38e2dd066e01256201b6b3dc4772dec970/spark/src/main/scala/org/apache/spark/sql/comet/CometSparkToColumnarExec.scala#L43-L61) ########## docs/source/user-guide/latest/understanding-comet-plans.md: ########## @@ -94,6 +94,8 @@ They serve different purposes and produce output in different places. | `spark.comet.explain.codegen.enabled` | `ExtendedExplainInfo` / SQL UI | A single `[COMET-INFO: JVM codegen dispatcher: <name1>, <name2>, ...]` segment naming each expression on the operator that was routed through the JVM codegen dispatcher. | | `spark.comet.explain.format` | Spark SQL UI (Spark 4.0 and newer) | Annotated plan or fallback-reason list, depending on `verbose` (default) or `fallback` value. | | `spark.comet.explain.native.enabled` | Executor logs, per task | The DataFusion plan with metrics, useful for inspecting Rust execution. | +| `spark.comet.explain.planOnly.enabled` | Driver log | The Comet plan and coverage summary, prefixed `[Comet plan-only]`, while Spark runs the original plan. | +| `spark.comet.explain.rules` | Driver log | Each plan transformation that Comet's planner rules make. | Review Comment: [P2] The new row promises a driver-log entry for every planner transformation. `RevertNativeForTransitionHeavyStages` can remove Comet operators without emitting a rule diff, so the log omits the rewrite someone investigating that fallback needs. Describe the before/after output from `CometScanRule`, `CometExecRule`, and `EliminateRedundantTransitions`, and point readers to `spark.comet.explain.fallback.log.enabled` for the reversion reason. Evidence: On Spark 4.1.3 with INFO logging, `spark.comet.explain.rules=true`, transition reversion enabled, and `maxTransitions=0`, the rule rewrote `ColumnarToRow(CometLocalTableScan)` to `LocalTableScan` with zero `Applying Rule` entries. An `EliminateRedundantTransitions` control emitted its rule diff. Enabling fallback-reason logging produced `Stage reverted: 1 C2R transitions exceed threshold 0`. [Reversion rule](https://github.com/apache/datafusion-comet/blob/56396e38e2dd066e01256201b6b3dc4772dec970/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala#L122-L152) ########## docs/source/contributor-guide/native_shuffle.md: ########## @@ -25,33 +25,35 @@ see [JVM Shuffle](jvm_shuffle.md). ## Overview -Native shuffle takes columnar input directly from Comet native operators and performs partitioning, +Native shuffle takes Arrow batches directly from Comet operators and performs partitioning, encoding, and writing in native Rust code. This avoids the columnar-to-row-to-columnar conversion overhead that JVM shuffle incurs. ``` -Comet Native (columnar) → Native Shuffle → Arrow IPC → columnar +Comet operator (Arrow batches) → Native Shuffle → Arrow IPC → columnar ``` Compare this to JVM shuffle's data path: ``` -Comet Native (columnar) → ColumnarToRowExec → rows → JVM Shuffle → Arrow IPC → columnar +Comet operator (Arrow batches) → ColumnarToRowExec → rows → JVM Shuffle → Arrow IPC → columnar ``` When `RevertNativeForTransitionHeavyStages` restores the map stage to Spark execution, the -native exchange stays in place. Its input still needs Arrow-backed Comet vectors, even if the -restored Spark operator supports columnar output. The rule adds `CometSparkToColumnarExec` to -convert either Spark rows or Spark columnar batches to Arrow before the native shuffle consumes -them. Spark's `RowToColumnarExec` alone does not satisfy this input contract. +native shuffle exchange usually stays in place. Its input still needs Arrow-backed Comet vectors, +even if the restored Spark operator supports columnar output. The rule adds `CometSparkToColumnarExec` +to convert either Spark rows or Spark columnar batches to Arrow before the native shuffle consumes +them. Spark's `RowToColumnarExec` alone does not satisfy this input contract. The exception is a +native shuffle whose input was converted by `spark.comet.convert.shuffleInput.enabled`: the rule +switches it back to JVM columnar shuffle, which reads the restored stage's rows. ## When Native Shuffle is Used Native shuffle (`CometExchange`) is selected when all of the following conditions are met: -1. **Shuffle mode allows native**: `spark.comet.shuffle.mode` is `native` or `auto`. +1. **Shuffle mode allows native shuffle**: `spark.comet.shuffle.mode` is `native` or `auto`. -2. **Child plan is a Comet native operator**: The child must be a `CometPlan` that produces +2. **Child plan is a Comet operator**: The child must be a `CometPlan` that produces Review Comment: [P2] The broader child criterion predicts `CometExchange` when Celeborn sees a `SinglePartition` exchange over an already-unwrapped `CometSparkToColumnarExec`. Even in explicit `native` mode with a supported integer schema, Celeborn declines that child and retains Spark’s exchange. Distinguish the local manager’s `CometPlan` check from Celeborn’s requirement for a `CometNativeExec` at selection time. Evidence: A Spark 4.1.3 Celeborn planning probe returned `None` for the columnar bridge with reason `Celeborn native shuffle requires a CometNativeExec child`. With the same schema and partitioning, a `CometNativeExec` control returned `Some(CometNativeShuffle)`. The repository test also asserts that the unwrapped case retains `ShuffleExchangeExec`. [Planner gate](https://github.com/apache/datafusion-comet/blob/56396e38e2dd066e01256201b6b3dc4772dec970/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala#L512-L536) · [Planning regression](https://github.com/apache/datafusion-comet/blob/56396e38e2dd066e01256201b6b3dc4772dec970/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala#L251-L265) ########## docs/source/contributor-guide/plugin_overview.md: ########## @@ -127,10 +140,26 @@ In the native code there is a `PhysicalPlanner` struct (in `planner.rs`) which c Apache DataFusion `ExecutionPlan`. In some cases, Comet provides specialized physical operators and expressions to override the DataFusion versions to ensure compatibility with Apache Spark. -The leaf nodes in the physical plan are always `ScanExec`. Each JVM-sourced input is exported once as an +Not every expression in a Rust plan runs in Rust. When an expression has no Rust implementation, or its Rust +implementation does not match Spark for the given input, its serde can route it through the JVM codegen dispatcher +instead of falling back to Spark (see the +[Implementation legend](../user-guide/latest/expressions.md#implementation-legend)). The Rust plan then evaluates the +expression with `JvmScalarUdfExpr`, which passes each batch's input columns through JNI and the Arrow C Data Interface +to `CometUdfBridge`, which hands them to `CometScalaUDFCodegen` on the JVM. There, Spark's own generated code for the expression, compiled into a batch +kernel by `CometBatchKernelCodegen`, runs over the Arrow vectors and returns an Arrow column. Scala and Java UDFs take +the same path. The operator around the expression stays Rust-implemented, so this JVM code runs inside the Comet +pipeline. `spark.comet.exec.scalaUDF.codegen.enabled`, which is on by default, turns the dispatcher on. When it is off, +these expressions fall back to Spark. + +A leaf node in the physical plan that takes Arrow batches from the JVM is a `ScanExec`. Each such input is exported once as an [Arrow C Stream](https://arrow.apache.org/docs/format/CStreamInterface.html) (`org.apache.arrow.c.ArrowArrayStream`), and `ScanExec` pulls each input batch through a single C callback rather than making a JNI call per batch. The -input could be a Comet native Parquet scan, a Spark exchange, or another native plan. +input could be a broadcast, a shuffle read decoded through JNI and imported into the JVM, the output of +`CometSparkToColumnarExec`, or the output of an earlier Comet block. When direct shuffle read is enabled and the input +is a Comet shuffle whose output types Comet supports, the input is instead a `ShuffleScanExec`, which pulls compressed Review Comment: [P2] With `spark.sql.adaptive.enabled=false` and `spark.comet.shuffle.directRead.enabled=true`, a final aggregate over a Comet hash shuffle with `BIGINT` output still reads through `ScanExec`. The paragraph predicts `ShuffleScanExec`, leading someone profiling the aggregate to expect the Arrow FFI crossing to disappear. Describe direct read as applying when the serialized input is a `ShuffleScan` leaf, and include the bare non-AQE exchange as a `ScanExec` case. Evidence: A Spark 4.1.3 serializer probe returned `SCAN` for the bare exchange and `SHUFFLE_SCAN` for its AQE stage with the same schema and direct-read setting. The checked-in regression groups a 50-partition range by `id % 997` and asserts those leaf kinds in the final aggregate. [Grouped-aggregate regression](https://github.com/apache/datafusion-comet/blob/56396e38e2dd066e01256201b6b3dc4772dec970/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala#L1791-L1821) -- 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]
