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: [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]
