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

Reply via email to