dongjoon-hyun commented on code in PR #59146:
URL: https://github.com/apache/spark/pull/59146#discussion_r4140449874
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +42,26 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("columnar execution cache is not captured by whole-stage codegen
references") {
+ val source = ColumnarCacheSerializationSource()
+ val join = SortMergeJoinExec(
+ Seq(Literal(1)), Seq(Literal(1)), Inner, None,
+ ColumnarToRowExec(source), LocalTableScanExec(Nil, Nil, None))
+ val (ctx, code) = WholeStageCodegenExec(join)(0).doCodeGen()
+ assert(ctx.references.exists(_.asInstanceOf[AnyRef] eq join))
+
+ val rdd = source.executeColumnar()
+ assert(source.executeColumnar() eq rdd)
+ val factory = new WholeStageCodegenEvaluatorFactory(
+ Right(code), join.longMetric("numOutputRows"), ctx.references.toArray)
+ val serializer = SparkEnv.get.closureSerializer.newInstance()
+ // The cached RDD holds non-serializable driver state, but the generated
evaluator does not
+ // need it. Input iterators are passed separately when the evaluator runs.
+
serializer.deserialize[WholeStageCodegenEvaluatorFactory](serializer.serialize(factory))
+ val readback =
serializer.deserialize[SortMergeJoinExec](serializer.serialize(join))
+ readback.cleanupResources()
Review Comment:
These two lines can never fail on their own:
- `join` is already in `factory.references` (line 51), so line 60 serializes
the same object graph. Without the fix, line 60 fails first.
- No node in this tree (`SortMergeJoinExec`, `ColumnarToRowExec`,
`LocalTableScanExec`, `ColumnarCacheSerializationSource`) overrides
`cleanupResources()`, so it only recurses into the children and never touches
`executeColumnarRDD`.
Could you remove them, along with the "deserialized-plan cleanup" claim in
the PR description?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala:
##########
@@ -218,6 +218,7 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with
Logging with Serializ
executeBroadcastBcast.get.asInstanceOf[broadcast.Broadcast[T]]
}
+ @transient
private val executeColumnarRDD = LazyTry {
doExecuteColumnar()
Review Comment:
`AQEShuffleReadExec` is a case where this change has no effect at all. Its
`doExecuteColumnar()` (and `doExecute()`) returns `shuffleRDD`, which is a
non-transient `private lazy val`:
https://github.com/apache/spark/blob/a026a0be4c88e7ace59d1da7c18ab73b0d81edde/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEShuffleReadExec.scala#L268-L284
So the very object cached in `executeColumnarRDD` is still serialized
through `shuffleRDD` when a downstream stage's `SortExec`, `HashAggregateExec`,
`SortMergeJoinExec`, or `ShuffledHashJoinExec` captures a subtree that reaches
an earlier stage's `AQEShuffleReadExec` through the non-transient
`ShuffleQueryStageExec.plan`. For the columnar path this only happens with
plugin `ShuffleExchangeLike` implementations, but the row path has bypassed the
`@transient executeRDD` in the same way since SPARK-48195. `@transient private
lazy val shuffleRDD` (like `partitionDataSizes` and `metrics` in the same
class) would close it.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +42,26 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("columnar execution cache is not captured by whole-stage codegen
references") {
+ val source = ColumnarCacheSerializationSource()
+ val join = SortMergeJoinExec(
+ Seq(Literal(1)), Seq(Literal(1)), Inner, None,
+ ColumnarToRowExec(source), LocalTableScanExec(Nil, Nil, None))
Review Comment:
This plan shape is a same-stage case where the fix does not change the job
outcome. If it were planned and executed, this child would become
`InputAdapter(WholeStageCodegenExec(ColumnarToRowExec(InputAdapter(source))))`,
and `ColumnarToRowExec.inputRDDs()` returns the cached RDD itself. So the
cached RDD is a parent of the stage RDD and is serialized in the same task
binary anyway:
- If the cached RDD is not serializable and has partitions, the job still
fails with `Task not serializable`, with or without `@transient`.
- If it is serializable, the plan reference costs no extra bytes because of
Java serialization back-references.
The test passes because it serializes the factory in isolation, without the
stage RDD lineage. The fix really matters for cross-stage capture, e.g., a
cached DataFrame (columnar `InMemoryTableScanExec`) behind a shuffle under a
`SortMergeJoinExec` or `HashAggregateExec`. Could you add a test with such a
query that asserts the upstream RDD is not serialized into the downstream stage
closure? (A Parquet variant would still fail because of
`FileSourceScanExec.inputRDD`.)
In addition, the comment at lines 58-59 ("Input iterators are passed
separately") is misleading because those iterators are computed from exactly
this RDD's lineage. The chain in the PR description also skips a hop: it is the
`executeColumnarRDD` of `ColumnarToRowExec.child` (the source), not of
`ColumnarToRowExec`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala:
##########
@@ -218,6 +218,7 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with
Logging with Serializ
executeBroadcastBcast.get.asInstanceOf[broadcast.Broadcast[T]]
}
+ @transient
private val executeColumnarRDD = LazyTry {
Review Comment:
This makes `executeColumnarRDD` transient, but for Spark's own columnar
scans the same RDD chain is still reachable from the plan through non-transient
subclass caches:
- `FileSourceScanExec.inputRDD` and `maybeCoalesceInputRDD`.
`doExecuteColumnar()` is `maybeCoalesceInputRDD.mapPartitionsInternal { ... }`.
- `BatchScanExec.inputRDD` and the streaming DSv2 scans
(`MicroBatchScanExec`, `ContinuousScanExec`, `RealTimeStreamScanExec`).
`doExecuteColumnar()` is `inputRDD.map { ... }`.
https://github.com/apache/spark/blob/a026a0be4c88e7ace59d1da7c18ab73b0d81edde/sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala#L752
For example, in `spark.read.parquet(path).groupBy("k").count()` with the
vectorized reader, the final `HashAggregateExec` still drags the `FileScanRDD`
into the reduce-stage task binary through `ShuffleQueryStageExec.plan` (or
`ShuffleExchangeExec.child` without AQE). This PR removes only the thin wrapper
RDD there. If `ColumnarCacheSerializationSource` cached its RDD in a `lazy val
inputRDD` like the real scans do, the new test would still fail after this fix.
These fields are only read on the driver, so marking them `@transient` like
`ShuffleExchangeExec.inputRDD` and `LocalTableScanExec.rdd` should be safe.
Could you include them in this PR, or narrow the PR title and description?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala:
##########
@@ -218,6 +218,7 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with
Logging with Serializ
executeBroadcastBcast.get.asInstanceOf[broadcast.Broadcast[T]]
}
+ @transient
Review Comment:
This is a regression from SPARK-48195
([d2e8c1c](https://github.com/apache/spark/commit/d2e8c1cb60e34a1c7e92374c07d682aa5ca79145),
#48037), which added `@transient` only to `executeRDD`. It was first released
in [v4.0.0](https://github.com/apache/spark/releases/tag/v4.0.0) (2025-05-19),
so `branch-4.0` through `branch-4.3` have the same issue. Could you mention
this in the PR description so that we can consider backporting? The description
also still says "Draft pending full source-build validation."
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala:
##########
@@ -218,6 +218,7 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with
Logging with Serializ
executeBroadcastBcast.get.asInstanceOf[broadcast.Broadcast[T]]
Review Comment:
SPARK-48195 introduced three `LazyTry` caches, and with this PR
`executeBroadcastBcast` becomes the only non-transient one. So every plan node
serialized through codegen `plan` references still carries an unused `LazyTry`
with its initializer lambda, and materialized broadcast nodes carry
`Success(TorrentBroadcast)`. This is harmless in practice because only the
broadcast handle is serialized. However, a deserialized plan can never reach it
either, since `executeQuery` dereferences the `@transient session` first.
Adding `@transient` to `executeBroadcastBcast` would make the three caches
consistent. WDYT?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -268,6 +290,22 @@ case class ColumnarOp(child: SparkPlan) extends
UnaryExecNode {
copy(child = newChild)
}
+case class ColumnarCacheSerializationSource() extends LeafExecNode {
+ override val output: Seq[Attribute] = Seq(AttributeReference("id",
IntegerType)())
+ override val supportsColumnar: Boolean = true
+
+ override protected def doExecuteColumnar(): RDD[ColumnarBatch] = {
+ new RDD[ColumnarBatch](sparkContext, Nil) {
+ val driverOnlyState: AnyRef = new Object
+ override def compute(split: Partition, context: TaskContext):
Iterator[ColumnarBatch] =
+ Iterator.empty
+ override protected def getPartitions: Array[Partition] = Array.empty
+ }
Review Comment:
nit: This reimplements `org.apache.spark.rdd.EmptyRDD`, which is
`private[spark]` and already used in `sql/core`. The following keeps the
non-serializable sentinel, and the `Partition` and `TaskContext` imports can be
removed.
```scala
new EmptyRDD[ColumnarBatch](sparkContext) {
val driverOnlyState: AnyRef = new Object
}
```
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -268,6 +290,22 @@ case class ColumnarOp(child: SparkPlan) extends
UnaryExecNode {
copy(child = newChild)
}
+case class ColumnarCacheSerializationSource() extends LeafExecNode {
Review Comment:
nit: With an empty parameter list, `TreeNode.makeCopy` (`allowEmptyArgs =
false`) cannot find a constructor. So canonicalizing this node or any of its
ancestors (`canonicalized`, `sameResult`) throws `INTERNAL_ERROR: Couldn't find
a valid constructor`. The current test doesn't canonicalize, but it would break
as soon as someone adds a SPARK-23731-style check. In addition, all instances
are equal to each other while their `output` expression IDs differ. `case class
ColumnarCacheSerializationSource(output: Seq[Attribute])` would avoid both.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +42,26 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("columnar execution cache is not captured by whole-stage codegen
references") {
Review Comment:
nit: Could you add the JIRA ID prefix like the other regression tests in
this suite?
```suggestion
test("SPARK-59884: columnar execution cache is not captured by whole-stage
codegen references") {
```
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +42,26 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("columnar execution cache is not captured by whole-stage codegen
references") {
+ val source = ColumnarCacheSerializationSource()
+ val join = SortMergeJoinExec(
+ Seq(Literal(1)), Seq(Literal(1)), Inner, None,
+ ColumnarToRowExec(source), LocalTableScanExec(Nil, Nil, None))
+ val (ctx, code) = WholeStageCodegenExec(join)(0).doCodeGen()
+ assert(ctx.references.exists(_.asInstanceOf[AnyRef] eq join))
Review Comment:
This assertion points at the root cause, which is one level above the
patched field. `SortMergeJoinExec`, `SortExec`, `HashAggregateExec`, and
`ShuffledHashJoinExec` put `this` into the codegen references
(`ctx.addReferenceObj("plan", this)`). As a result, the whole subtree is
serialized into each task, and it crosses stage boundaries through the
non-transient `QueryStageExec.plan`, `QueryStageExec._canonicalized`, and
`ShuffleExchangeExec.child`. Downstream tasks therefore also deserialize all
upstream stage plans and register their `SQLMetric`s. The non-codegen paths
(the `SortExec.doExecute` closure and `SortMergeJoinEvaluatorFactory(left,
right)`) capture the subtree too.
A field-level `@transient` fixes one symptom. It may be worth a follow-up
JIRA to (1) make all execution caches in `SparkPlan` and its subclasses
transient consistently, and (2) avoid serializing plans across stage
boundaries, e.g., by keeping only `output`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +42,26 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("columnar execution cache is not captured by whole-stage codegen
references") {
+ val source = ColumnarCacheSerializationSource()
+ val join = SortMergeJoinExec(
+ Seq(Literal(1)), Seq(Literal(1)), Inner, None,
+ ColumnarToRowExec(source), LocalTableScanExec(Nil, Nil, None))
+ val (ctx, code) = WholeStageCodegenExec(join)(0).doCodeGen()
+ assert(ctx.references.exists(_.asInstanceOf[AnyRef] eq join))
+
+ val rdd = source.executeColumnar()
+ assert(source.executeColumnar() eq rdd)
+ val factory = new WholeStageCodegenEvaluatorFactory(
+ Right(code), join.longMetric("numOutputRows"), ctx.references.toArray)
Review Comment:
nit: The second parameter of `WholeStageCodegenEvaluatorFactory` is
`durationMs`, and production passes `longMetric("pipelineTime")` of
`WholeStageCodegenExec`. This passes the join's `numOutputRows` instead. It's
harmless here because no evaluator is created, but an evaluation would add the
pipeline time to the row count. How about keeping `val wscg =
WholeStageCodegenExec(join)(0)` and passing `wscg.longMetric("pipelineTime")`?
This is unnecessary if the test is simplified as suggested above.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +42,26 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("columnar execution cache is not captured by whole-stage codegen
references") {
+ val source = ColumnarCacheSerializationSource()
+ val join = SortMergeJoinExec(
Review Comment:
nit: The SMJ, `doCodeGen()`, and hand-built factory scaffolding isn't needed
for the property under test, and line 51 couples the test to SMJ codegen
internals. If SMJ stops referencing `this` in the future, this test will break
even though the property improves. A direct round trip fails in the same way
without the fix (`LazyTry.tryT = Success(rdd)` -> `driverOnlyState`), like the
SPARK-23731 and SPARK-27418 tests in this suite:
```scala
val source = ColumnarCacheSerializationSource()
source.executeColumnar() // Materialize the cached columnar RDD.
val serializer = SparkEnv.get.closureSerializer.newInstance()
// Parent plans can be captured by task closures, e.g.,
SortMergeJoinExec adds itself to
// the whole-stage codegen references.
serializer.deserialize[SparkPlan](serializer.serialize(ColumnarToRowExec(source)))
```
--
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]
