You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "sunchao (via GitHub)" <gi...@apache.org> on 2026/09/30 00:18:35 UTC
[PR] [SQL][WIP] Avoid serializing cached columnar RDDs in whole-stage closures [spark]
sunchao opened a new pull request, #59146:
URL: https://github.com/apache/spark/pull/59146
### What changes were proposed in this pull request?
Mark `SparkPlan.executeColumnarRDD` transient, matching `executeRDD`, so serializing a plan retained by whole-stage code generation does not capture its cached columnar RDD.
Add a regression in `SparkPlanSuite` that uses real sort-merge join codegen references, primes the columnar cache with an RDD containing a nonserializable sentinel, and checks driver-side cache reuse, evaluator serialization/deserialization, and deserialized-plan cleanup.
Draft pending JIRA linkage and full source-build validation.
### Why are the changes needed?
`SortMergeJoinExec` adds its plan to codegen references for executor cleanup. If a descendant columnar plan has initialized its execution cache, serializing the evaluator follows that cached RDD even though the evaluator receives executable inputs separately. This captures unnecessary driver-side objects and can cause `NotSerializableException`.
The regression fails before the change through `WholeStageCodegenEvaluatorFactory.references -> SortMergeJoinExec -> ColumnarToRowExec -> executeColumnarRDD -> cached RDD -> sentinel`. Excluding the cache fixes that path while preserving driver-side reuse.
### Does this PR introduce _any_ user-facing change?
Yes. Whole-stage evaluator serialization no longer captures a referenced plan's cached columnar execution RDD. Query results and public APIs are unchanged. This fixes the demonstrated capture path; it does not claim to eliminate all sources of large task closures.
### How was this patch tested?
- Compiled the exact modified `SparkPlan.scala` and `SparkPlanSuite.scala` from master at `a026a0be4c88e7ace59d1da7c18ab73b0d81edde` against public Spark 4.2.0 dependencies with Scala 2.13.18 and ScalaTest 3.2.16. The new regression passed.
- A control removing only `@transient` failed with the expected `NotSerializableException` through `executeColumnarRDD`.
- The complete `SparkPlanSuite` in this class overlay passed 12 of 13 tests. The existing SPARK-58120 `BatchScanExec` canonicalization test failed identically with unmodified `SparkPlan`; the cached 4.2.0 `BatchScanExec` predates that master fix. This is not a full current-master build.
- Scoped Scalastyle with the repository configuration passed for both modified files: zero errors and warnings. `git diff --check` passed.
- Attempted `build/sbt 'sql/testOnly org.apache.spark.sql.execution.SparkPlanSuite'`. Project loading was blocked by missing BOM dependencies and a read-only local Maven cache. Full source-build validation remains outstanding.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-6).
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Avoid serializing cached columnar RDDs in whole-stage closures [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on PR #59146:
URL: https://github.com/apache/spark/pull/59146#issuecomment-5901810927
cc @viirya @cloud-fan @peter-toth @dongjoon-hyun
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #59146:
URL: https://github.com/apache/spark/pull/59146#discussion_r4146675270
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +44,63 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("SPARK-59884: columnar execution cache is not serialized with a plan") {
+ val source = ColumnarCacheSerializationSource(Seq(AttributeReference("id", IntegerType)()))
+ val rdd = source.executeColumnar()
+ assert(source.executeColumnar() eq rdd)
+ val serializer = SparkEnv.get.closureSerializer.newInstance()
+ // Plans can be captured by task closures, but their execution caches belong to the driver.
+ serializer.deserialize[ColumnarToRowExec](serializer.serialize(ColumnarToRowExec(source)))
+ }
+
+ test("SPARK-59884: columnar execution cache does not cross a shuffle in the task binary") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ val cached = spark.range(0, 8, 1, 2).selectExpr("id % 2 AS key").cache()
+ try {
+ cached.count()
+ val query = cached.groupBy("key").count()
+ checkAnswer(query, Seq(Row(0L, 4L), Row(1L, 4L)))
+
+ val plan = query.queryExecution.executedPlan
+ val exchanges = plan.collect { case exchange: ShuffleExchangeExec => exchange }
+ assert(exchanges.size == 1)
+ val scan = exchanges.head.child.collectFirst {
+ case scan: InMemoryTableScanExec => scan
+ }.get
+ assert(scan.supportsColumnar)
+ val columnarRDD = scan.executeColumnar()
+ val stageRDD = plan.execute()
+
+ var serializedStageRDD = false
+ var serializedColumnarRDD = false
+ val out = new ObjectOutputStream(new ByteArrayOutputStream()) {
+ enableReplaceObject(true)
+ override protected def replaceObject(obj: AnyRef): AnyRef = {
+ serializedStageRDD ||= obj eq stageRDD
+ serializedColumnarRDD ||= obj eq columnarRDD
+ obj
+ }
+ }
+ try {
+ // DAGScheduler serializes the complete (stage.rdd, stage.func) tuple for a result task.
+ // The shuffle cuts off the upstream RDD lineage, including the columnar scan's cache.
Review Comment:
Updated in [489b75d955ee](https://github.com/apache/spark/commit/489b75d955eea7a95848a6b98bfaf5a32edee2e6) using the suggested wording: the upstream scan RDD is outside the downstream stage's input lineage, and captured plans must not reintroduce it through their columnar execution caches. This is a comment-only change; executable code is unchanged and whitespace, ASCII, and line-length checks passed.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #59146:
URL: https://github.com/apache/spark/pull/59146#discussion_r4170495763
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +44,64 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("SPARK-59884: columnar execution cache is not serialized with a plan") {
+ val source = ColumnarCacheSerializationSource(Seq(AttributeReference("id", IntegerType)()))
+ val rdd = source.executeColumnar()
+ assert(source.executeColumnar() eq rdd)
+ val serializer = SparkEnv.get.closureSerializer.newInstance()
+ // Plans can be captured by task closures, but their execution caches belong to the driver.
+ serializer.deserialize[ColumnarToRowExec](serializer.serialize(ColumnarToRowExec(source)))
+ }
+
+ test("SPARK-59884: columnar execution cache does not cross a shuffle in the task binary") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ val cached = spark.range(0, 8, 1, 2).selectExpr("id % 2 AS key").cache()
+ try {
+ cached.count()
Review Comment:
The replacement Parquet regression has no cache warm-up action and now calls `QueryTest.checkAnswer(..., checkToRDD = false)`, so the aggregation runs only on the plan whose task binary is inspected. Updated in e3b6ec55d570bf2c7d2d79061c414796d4f2ae7a.
##########
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:
Fixed both Scaladoc references to `doExecuteColumnar` in e3b6ec55d570bf2c7d2d79061c414796d4f2ae7a.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun closed pull request #59146: [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient
URL: https://github.com/apache/spark/pull/59146
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on PR #59146:
URL: https://github.com/apache/spark/pull/59146#issuecomment-6026791338
Thanks @dongjoon-hyun @peter-toth @HyukjinKwon @viirya !!
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #59146:
URL: https://github.com/apache/spark/pull/59146#discussion_r4141045160
##########
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:
Added the SPARK-59884 prefix to both regression tests.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #59146:
URL: https://github.com/apache/spark/pull/59146#issuecomment-6023902380
Since this is filed as `Improvement`, this landed only at master/4.x, @sunchao .
<img width="1246" height="443" alt="Screenshot 2026-10-06 at 12 28 46" src="https://github.com/user-attachments/assets/7ddd3326-63ff-49c7-a4c4-4e4050265ae7" />
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #59146:
URL: https://github.com/apache/spark/pull/59146#issuecomment-6023874047
Please make backporting PRs for release branches, @sunchao .
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on PR #59146:
URL: https://github.com/apache/spark/pull/59146#issuecomment-5962984773
Updated in [e3b6ec55](https://github.com/apache/spark/commit/e3b6ec55d570bf2c7d2d79061c414796d4f2ae7a). The cross-shuffle regression now uses vectorized Parquet and checks both preconditions: the query uses the scan's columnar wrapper RDD, and the downstream task serialization visits the upstream scan plan. Applied the three cleanups and replied to all five inline comments. The PR description now states the wrapper-only scope and has a concise testing section, addressing the description feedback as well.
Validation: recompiled the two changed files and ran the complete SparkPlanSuite: **14 passed, 0 failed**. Removing only the `executeColumnarRDD` annotation made both regressions fail for the intended reasons: the direct round trip reached the nonserializable sentinel, and the shuffle test found the upstream wrapper RDD. Scoped Scalastyle passed with 0 errors and 0 warnings; `git diff --check` passed.
The local runs used Scala 2.13.18 / Java 17 and the exact exported test classpath from [CI run 37074734073](https://github.com/sunchao/spark/actions/runs/37074734073), whose compile artifact was built from CI merge 46131903a74c52e9c457400f445efe7fcb38f9a7 (Apache master 099c3ee2e2e71ef8c826ee06d35081cb6d82b29c plus this PR). The two source files were verified against that base before recompiling. The source build passed; the broader CI test jobs are still running.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Avoid serializing cached columnar RDDs in whole-stage closures [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #59146:
URL: https://github.com/apache/spark/pull/59146#issuecomment-6023862443
**Merge Summary:**
- merged into master https://github.com/apache/spark/commit/895c75f0032a93ac29dff18d6720783ba1f22a5c
- merged into branch-4.x https://github.com/apache/spark/commit/7f923c25bbb8c806c156b461f4645d9a15116779
*Posted by `merge_spark_pr.py`*
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #59146:
URL: https://github.com/apache/spark/pull/59146#discussion_r4170607801
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +44,64 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("SPARK-59884: columnar execution cache is not serialized with a plan") {
+ val source = ColumnarCacheSerializationSource(Seq(AttributeReference("id", IntegerType)()))
+ val rdd = source.executeColumnar()
+ assert(source.executeColumnar() eq rdd)
+ val serializer = SparkEnv.get.closureSerializer.newInstance()
+ // Plans can be captured by task closures, but their execution caches belong to the driver.
+ serializer.deserialize[ColumnarToRowExec](serializer.serialize(ColumnarToRowExec(source)))
+ }
+
+ test("SPARK-59884: columnar execution cache does not cross a shuffle in the task binary") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ val cached = spark.range(0, 8, 1, 2).selectExpr("id % 2 AS key").cache()
+ try {
+ cached.count()
+ val query = cached.groupBy("key").count()
+ checkAnswer(query, Seq(Row(0L, 4L), Row(1L, 4L)))
+
+ val plan = query.queryExecution.executedPlan
+ val exchanges = plan.collect { case exchange: ShuffleExchangeExec => exchange }
+ assert(exchanges.size == 1)
+ val scan = exchanges.head.child.collectFirst {
+ case scan: InMemoryTableScanExec => scan
+ }.get
+ assert(scan.supportsColumnar)
+ val columnarRDD = scan.executeColumnar()
Review Comment:
You're right: the previous cached-DataFrame query read rows, and my explicit call created the columnar RDD after execution. Replaced it with a vectorized Parquet aggregation in e3b6ec55d570bf2c7d2d79061c414796d4f2ae7a. The test now asserts that the same columnar RDD is present in the shuffle exchange's input lineage, proving the query used it. The comments and PR description explicitly limit this check to the columnar wrapper; `FileSourceScanExec.inputRDD` remains a separate serialization path.
Validation: the complete SparkPlanSuite passed all 14 tests. With only `executeColumnarRDD`'s `@transient` removed, both regressions fail; the Parquet test reaches the assertion that the upstream columnar RDD was serialized.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +44,64 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("SPARK-59884: columnar execution cache is not serialized with a plan") {
+ val source = ColumnarCacheSerializationSource(Seq(AttributeReference("id", IntegerType)()))
+ val rdd = source.executeColumnar()
+ assert(source.executeColumnar() eq rdd)
+ val serializer = SparkEnv.get.closureSerializer.newInstance()
+ // Plans can be captured by task closures, but their execution caches belong to the driver.
+ serializer.deserialize[ColumnarToRowExec](serializer.serialize(ColumnarToRowExec(source)))
+ }
+
+ test("SPARK-59884: columnar execution cache does not cross a shuffle in the task binary") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ val cached = spark.range(0, 8, 1, 2).selectExpr("id % 2 AS key").cache()
+ try {
+ cached.count()
+ val query = cached.groupBy("key").count()
+ checkAnswer(query, Seq(Row(0L, 4L), Row(1L, 4L)))
+
+ val plan = query.queryExecution.executedPlan
+ val exchanges = plan.collect { case exchange: ShuffleExchangeExec => exchange }
+ assert(exchanges.size == 1)
+ val scan = exchanges.head.child.collectFirst {
+ case scan: InMemoryTableScanExec => scan
+ }.get
+ assert(scan.supportsColumnar)
+ val columnarRDD = scan.executeColumnar()
+ val stageRDD = plan.execute()
+
+ var serializedStageRDD = false
+ var serializedColumnarRDD = false
+ val out = new ObjectOutputStream(new ByteArrayOutputStream()) {
+ enableReplaceObject(true)
+ override protected def replaceObject(obj: AnyRef): AnyRef = {
+ serializedStageRDD ||= obj eq stageRDD
+ serializedColumnarRDD ||= obj eq columnarRDD
+ obj
+ }
+ }
+ try {
+ // DAGScheduler serializes the complete (stage.rdd, stage.func) tuple for a result task.
+ // The upstream scan RDD is outside the downstream stage's input lineage.
+ // Captured plans must not reintroduce it through their columnar execution caches.
+ val taskFunc = (_: TaskContext, rows: Iterator[InternalRow]) => rows.size
+ out.writeObject((stageRDD, taskFunc))
+ } finally {
+ out.close()
+ }
+ assert(serializedStageRDD)
Review Comment:
Replaced `serializedStageRDD` with `serializedScan` in e3b6ec55d570bf2c7d2d79061c414796d4f2ae7a, and assert that the downstream task's serialized object graph visits the exact upstream scan plan. If whole-plan capture is removed later, this guard fails instead of allowing the cache-exclusion assertion to pass vacuously. The full suite passes, and removing only the transient annotation makes the wrapper-exclusion assertion fail after this guard passes.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #59146:
URL: https://github.com/apache/spark/pull/59146#discussion_r4170496866
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +44,64 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("SPARK-59884: columnar execution cache is not serialized with a plan") {
+ val source = ColumnarCacheSerializationSource(Seq(AttributeReference("id", IntegerType)()))
+ val rdd = source.executeColumnar()
+ assert(source.executeColumnar() eq rdd)
+ val serializer = SparkEnv.get.closureSerializer.newInstance()
+ // Plans can be captured by task closures, but their execution caches belong to the driver.
+ serializer.deserialize[ColumnarToRowExec](serializer.serialize(ColumnarToRowExec(source)))
+ }
+
+ test("SPARK-59884: columnar execution cache does not cross a shuffle in the task binary") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ val cached = spark.range(0, 8, 1, 2).selectExpr("id % 2 AS key").cache()
+ try {
+ cached.count()
+ val query = cached.groupBy("key").count()
+ checkAnswer(query, Seq(Row(0L, 4L), Row(1L, 4L)))
+
+ val plan = query.queryExecution.executedPlan
+ val exchanges = plan.collect { case exchange: ShuffleExchangeExec => exchange }
+ assert(exchanges.size == 1)
+ val scan = exchanges.head.child.collectFirst {
+ case scan: InMemoryTableScanExec => scan
+ }.get
+ assert(scan.supportsColumnar)
+ val columnarRDD = scan.executeColumnar()
+ val stageRDD = plan.execute()
+
+ var serializedStageRDD = false
+ var serializedColumnarRDD = false
+ val out = new ObjectOutputStream(new ByteArrayOutputStream()) {
Review Comment:
Switched to `ObjectOutputStream(OutputStream.nullOutputStream())` and updated the imports in e3b6ec55d570bf2c7d2d79061c414796d4f2ae7a. The test only observes the objects visited during serialization.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient [spark]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #59146:
URL: https://github.com/apache/spark/pull/59146#discussion_r4141043704
##########
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:
Updated in [d9fe008b13f](https://github.com/apache/spark/commit/d9fe008b13fdb8be4bc80b3386f045a159ac0481). Narrowed the title and description to the base SparkPlan columnar execution cache. The description now explicitly calls out the remaining scan-owned inputRDD/maybeCoalesceInputRDD paths. The new cross-shuffle regression uses a built-in cached DataFrame and makes no claim that this fixes Parquet scan capture.
##########
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:
Agreed. AQEShuffleReadExec.shuffleRDD remains an independent reference to the same RDD, so this annotation does not remove that retention path. I documented this limitation and kept that cache change for separate follow-up coverage.
##########
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:
Agreed; the factory-only test did not demonstrate the cross-stage benefit. Added a cached-DataFrame aggregation with a real shuffle and verified its results. The test then serializes the complete downstream result-stage RDD/function tuple in the same shape used by DAGScheduler, recording whether the upstream InMemoryTableScanExec columnar RDD is visited. It is absent with the fix and present when only this annotation is reverted. The focused unit test is now a direct plan round trip. Removed the misleading input-iterator comment and rewrote the description around the actual cross-stage path. The revised SparkPlanSuite passes all 14 tests against the exact Spark 5.0 CI-built artifacts; both new tests fail with only the annotation removed.
##########
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:
The consistency change makes sense as a follow-up. I kept this PR focused on the columnar RDD retention regression and its distinguishing tests; executeBroadcastBcast remains unchanged. The description now makes that scope explicit.
##########
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:
Agreed that reducing whole-plan capture deserves follow-up work. This PR now explicitly limits itself to the base columnar cache. A broader change will also need to preserve the operator methods and cleanup behavior used by the captured plans, so I kept it separate from this fix.
##########
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:
Removed the redundant join round trip and cleanup call, and removed the deserialized-plan cleanup claim from the description.
##########
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:
Simplified the focused regression to a direct ColumnarToRowExec serialization round trip after initializing the source cache. The separate cached-DataFrame regression covers serialization across a real shuffle boundary.
##########
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:
Added SPARK-48195/#48037 and its originating commit to the description, including Spark 4.0.0 as the first affected release and branches 4.0 through 4.3 as backport candidates. Removed the stale draft sentence and replaced the old validation section with the revised tests' results.
##########
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:
Removed the hand-built evaluator factory when simplifying the test, which also removes the incorrect metric argument.
##########
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:
Reused EmptyRDD for the sentinel fixture and removed the Partition import. TaskContext is still needed by the new cross-shuffle test's result-task function.
##########
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:
Changed the helper to case class ColumnarCacheSerializationSource(output: Seq[Attribute]), so the output attributes participate in copying and equality.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #59146:
URL: https://github.com/apache/spark/pull/59146#discussion_r4148252126
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +44,64 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("SPARK-59884: columnar execution cache is not serialized with a plan") {
+ val source = ColumnarCacheSerializationSource(Seq(AttributeReference("id", IntegerType)()))
+ val rdd = source.executeColumnar()
+ assert(source.executeColumnar() eq rdd)
+ val serializer = SparkEnv.get.closureSerializer.newInstance()
+ // Plans can be captured by task closures, but their execution caches belong to the driver.
+ serializer.deserialize[ColumnarToRowExec](serializer.serialize(ColumnarToRowExec(source)))
+ }
+
+ test("SPARK-59884: columnar execution cache does not cross a shuffle in the task binary") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ val cached = spark.range(0, 8, 1, 2).selectExpr("id % 2 AS key").cache()
+ try {
+ cached.count()
+ val query = cached.groupBy("key").count()
+ checkAnswer(query, Seq(Row(0L, 4L), Row(1L, 4L)))
+
+ val plan = query.queryExecution.executedPlan
+ val exchanges = plan.collect { case exchange: ShuffleExchangeExec => exchange }
+ assert(exchanges.size == 1)
+ val scan = exchanges.head.child.collectFirst {
+ case scan: InMemoryTableScanExec => scan
+ }.get
+ assert(scan.supportsColumnar)
+ val columnarRDD = scan.executeColumnar()
Review Comment:
The cached `InMemoryTableScanExec` shape that I suggested in the previous round turns out not to exercise the columnar path. Since SPARK-37369, `InMemoryTableScanExec.supportsRowBased` is always `true`, so no `ColumnarToRowExec` is inserted, and the query reads the cache row-wise through `InputAdapter.inputRDD` (`child.execute()`).
https://github.com/apache/spark/blob/a026a0be4c88e7ace59d1da7c18ab73b0d81edde/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala#L90
So the job never creates the scan's columnar RDD. This line creates it after `checkAnswer`, and line 98 fails without the fix only because of that. I ran the same query on Apache Spark 4.3.0, which doesn't have this fix and has the same code paths as `master`:
- The executed plan has no `ColumnarToRow`. `*(1) HashAggregate` reads the `InMemoryTableScan` through an `InputAdapter`.
- After the query runs, `executeColumnarRDD` is not materialized in any node.
- Before this line, the serialized `(stageRDD, taskFunc)` contains only `MapPartitionsRDD` and `ShuffledRowRDD`, i.e., no upstream RDD. After this line, it contains the columnar RDD and its whole lineage down to `ParallelCollectionRDD`.
So the test doesn't reproduce a leak of this query, although the PR description presents it as "a real cached-DataFrame aggregation across a shuffle". `assert(scan.supportsColumnar)` doesn't prove that the columnar path is used either. Could you assert that the job itself created `columnarRDD`, like the "or the test exercises nothing" guards in `DataFrameSetOperationsSuite`?
```scala
def lineage(rdd: RDD[_]): Seq[RDD[_]] = rdd +: rdd.dependencies.flatMap(d => lineage(d.rdd))
assert(lineage(exchanges.head.inputRDD).exists(_ eq columnarRDD),
"the job must create the columnar RDD, or the test exercises nothing")
```
This assertion fails with the cached DataFrame. A vectorized Parquet scan passes it because `ColumnarToRowExec.inputRDDs()` materializes `FileSourceScanExec.executeColumnarRDD` during the job. For Parquet, note that `FileSourceScanExec.inputRDD` is still serialized after this fix, as discussed in the previous round, so the identity check covers only the wrapper RDD. Otherwise, please update the test comments and the PR description to say that this line stands in for a columnar consumer, e.g., a plugin operator.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +44,64 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("SPARK-59884: columnar execution cache is not serialized with a plan") {
+ val source = ColumnarCacheSerializationSource(Seq(AttributeReference("id", IntegerType)()))
+ val rdd = source.executeColumnar()
+ assert(source.executeColumnar() eq rdd)
+ val serializer = SparkEnv.get.closureSerializer.newInstance()
+ // Plans can be captured by task closures, but their execution caches belong to the driver.
+ serializer.deserialize[ColumnarToRowExec](serializer.serialize(ColumnarToRowExec(source)))
+ }
+
+ test("SPARK-59884: columnar execution cache does not cross a shuffle in the task binary") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ val cached = spark.range(0, 8, 1, 2).selectExpr("id % 2 AS key").cache()
+ try {
+ cached.count()
+ val query = cached.groupBy("key").count()
+ checkAnswer(query, Seq(Row(0L, 4L), Row(1L, 4L)))
+
+ val plan = query.queryExecution.executedPlan
+ val exchanges = plan.collect { case exchange: ShuffleExchangeExec => exchange }
+ assert(exchanges.size == 1)
+ val scan = exchanges.head.child.collectFirst {
+ case scan: InMemoryTableScanExec => scan
+ }.get
+ assert(scan.supportsColumnar)
+ val columnarRDD = scan.executeColumnar()
+ val stageRDD = plan.execute()
+
+ var serializedStageRDD = false
+ var serializedColumnarRDD = false
+ val out = new ObjectOutputStream(new ByteArrayOutputStream()) {
+ enableReplaceObject(true)
+ override protected def replaceObject(obj: AnyRef): AnyRef = {
+ serializedStageRDD ||= obj eq stageRDD
+ serializedColumnarRDD ||= obj eq columnarRDD
+ obj
+ }
+ }
+ try {
+ // DAGScheduler serializes the complete (stage.rdd, stage.func) tuple for a result task.
+ // The upstream scan RDD is outside the downstream stage's input lineage.
+ // Captured plans must not reintroduce it through their columnar execution caches.
+ val taskFunc = (_: TaskContext, rows: Iterator[InternalRow]) => rows.size
+ out.writeObject((stageRDD, taskFunc))
+ } finally {
+ out.close()
+ }
+ assert(serializedStageRDD)
Review Comment:
`serializedStageRDD` is always `true`. `stageRDD` is `_1` of the tuple passed to `writeObject`, so `replaceObject` always sees it while the hook is enabled, and this line only checks that the hook works.
Nothing checks that the downstream stage captures the upstream plan, which is the precondition of line 98. Today the only path is `ctx.addReferenceObj("plan", this)` of the final `HashAggregateExec`:
https://github.com/apache/spark/blob/a026a0be4c88e7ace59d1da7c18ab73b0d81edde/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala#L589
If it goes away, e.g., when whole-plan capture is reduced as discussed in the previous round, line 98 passes even without this fix. I simulated it on Apache Spark 4.3.0 by writing the final `HashAggregateExec` as `null` in the same stream. Neither `scan` nor `columnarRDD` was serialized, so the assertions of this test would pass without the fix. How about tracking the scan node instead of `stageRDD`?
```scala
var serializedScan = false
...
serializedScan ||= obj eq scan
...
assert(serializedScan, "the downstream stage no longer captures the upstream plan")
```
##########
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:
nit: This is not from this PR, but the Scaladoc of `executeColumnar()` right below refers to `doColumnarExecute` twice (lines 227 and 230), which doesn't exist. It should be this `doExecuteColumnar`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +44,64 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("SPARK-59884: columnar execution cache is not serialized with a plan") {
+ val source = ColumnarCacheSerializationSource(Seq(AttributeReference("id", IntegerType)()))
+ val rdd = source.executeColumnar()
+ assert(source.executeColumnar() eq rdd)
+ val serializer = SparkEnv.get.closureSerializer.newInstance()
+ // Plans can be captured by task closures, but their execution caches belong to the driver.
+ serializer.deserialize[ColumnarToRowExec](serializer.serialize(ColumnarToRowExec(source)))
+ }
+
+ test("SPARK-59884: columnar execution cache does not cross a shuffle in the task binary") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ val cached = spark.range(0, 8, 1, 2).selectExpr("id % 2 AS key").cache()
+ try {
+ cached.count()
Review Comment:
nit: `cached.count()` isn't needed because `cache()` already registers the relation and the query builds the buffers on its first run. In addition, `checkAnswer` runs the aggregation twice: `materializedRdd.count()` on a separate `rddQueryExecution` plan, and then `collect()` on this plan. Only the latter matters for the checks below, so `QueryTest.checkAnswer(query, Seq(Row(0L, 4L), Row(1L, 4L)), checkToRDD = false)` would be enough.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +44,64 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("SPARK-59884: columnar execution cache is not serialized with a plan") {
+ val source = ColumnarCacheSerializationSource(Seq(AttributeReference("id", IntegerType)()))
+ val rdd = source.executeColumnar()
+ assert(source.executeColumnar() eq rdd)
+ val serializer = SparkEnv.get.closureSerializer.newInstance()
+ // Plans can be captured by task closures, but their execution caches belong to the driver.
+ serializer.deserialize[ColumnarToRowExec](serializer.serialize(ColumnarToRowExec(source)))
+ }
+
+ test("SPARK-59884: columnar execution cache does not cross a shuffle in the task binary") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ val cached = spark.range(0, 8, 1, 2).selectExpr("id % 2 AS key").cache()
+ try {
+ cached.count()
+ val query = cached.groupBy("key").count()
+ checkAnswer(query, Seq(Row(0L, 4L), Row(1L, 4L)))
+
+ val plan = query.queryExecution.executedPlan
+ val exchanges = plan.collect { case exchange: ShuffleExchangeExec => exchange }
+ assert(exchanges.size == 1)
+ val scan = exchanges.head.child.collectFirst {
+ case scan: InMemoryTableScanExec => scan
+ }.get
+ assert(scan.supportsColumnar)
+ val columnarRDD = scan.executeColumnar()
+ val stageRDD = plan.execute()
+
+ var serializedStageRDD = false
+ var serializedColumnarRDD = false
+ val out = new ObjectOutputStream(new ByteArrayOutputStream()) {
Review Comment:
nit: The serialized bytes are never read. `OutputStream.nullOutputStream()` avoids buffering the whole task binary and makes the intent clearer, like the `NullOutputStream` of `SerializationDebugger.ListObjectOutputStream`.
```scala
import java.io.{ObjectOutputStream, OutputStream}
...
val out = new ObjectOutputStream(OutputStream.nullOutputStream()) {
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59884][SQL] Mark SparkPlan's columnar execution cache transient [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #59146:
URL: https://github.com/apache/spark/pull/59146#discussion_r4142207769
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -40,6 +44,63 @@ import org.apache.spark.util.ThreadUtils
class SparkPlanSuite extends SharedSparkSession {
+ test("SPARK-59884: columnar execution cache is not serialized with a plan") {
+ val source = ColumnarCacheSerializationSource(Seq(AttributeReference("id", IntegerType)()))
+ val rdd = source.executeColumnar()
+ assert(source.executeColumnar() eq rdd)
+ val serializer = SparkEnv.get.closureSerializer.newInstance()
+ // Plans can be captured by task closures, but their execution caches belong to the driver.
+ serializer.deserialize[ColumnarToRowExec](serializer.serialize(ColumnarToRowExec(source)))
+ }
+
+ test("SPARK-59884: columnar execution cache does not cross a shuffle in the task binary") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ val cached = spark.range(0, 8, 1, 2).selectExpr("id % 2 AS key").cache()
+ try {
+ cached.count()
+ val query = cached.groupBy("key").count()
+ checkAnswer(query, Seq(Row(0L, 4L), Row(1L, 4L)))
+
+ val plan = query.queryExecution.executedPlan
+ val exchanges = plan.collect { case exchange: ShuffleExchangeExec => exchange }
+ assert(exchanges.size == 1)
+ val scan = exchanges.head.child.collectFirst {
+ case scan: InMemoryTableScanExec => scan
+ }.get
+ assert(scan.supportsColumnar)
+ val columnarRDD = scan.executeColumnar()
+ val stageRDD = plan.execute()
+
+ var serializedStageRDD = false
+ var serializedColumnarRDD = false
+ val out = new ObjectOutputStream(new ByteArrayOutputStream()) {
+ enableReplaceObject(true)
+ override protected def replaceObject(obj: AnyRef): AnyRef = {
+ serializedStageRDD ||= obj eq stageRDD
+ serializedColumnarRDD ||= obj eq columnarRDD
+ obj
+ }
+ }
+ try {
+ // DAGScheduler serializes the complete (stage.rdd, stage.func) tuple for a result task.
+ // The shuffle cuts off the upstream RDD lineage, including the columnar scan's cache.
Review Comment:
Nit: could we distinguish the RDD lineage boundary from the plan-cache reference here? The shuffle excludes the upstream scan RDD from the downstream input lineage, but the captured plan can still retain it through its execution cache—that is the path this patch removes. Perhaps:
```scala
// The upstream scan RDD is outside the downstream stage's input lineage.
// Captured plans must not reintroduce it through their columnar execution caches.
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org