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