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