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