dongjoon-hyun commented on code in PR #59146:
URL: https://github.com/apache/spark/pull/59146#discussion_r4148252126
##########
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:
The cached `InMemoryTableScanExec` shape that I suggested in the previous
round turns out not to exercise the columnar path. Since SPARK-37369,
`InMemoryTableScanExec.supportsRowBased` is always `true`, so no
`ColumnarToRowExec` is inserted, and the query reads the cache row-wise through
`InputAdapter.inputRDD` (`child.execute()`).
https://github.com/apache/spark/blob/a026a0be4c88e7ace59d1da7c18ab73b0d81edde/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala#L90
So the job never creates the scan's columnar RDD. This line creates it after
`checkAnswer`, and line 98 fails without the fix only because of that. I ran
the same query on Apache Spark 4.3.0, which doesn't have this fix and has the
same code paths as `master`:
- The executed plan has no `ColumnarToRow`. `*(1) HashAggregate` reads the
`InMemoryTableScan` through an `InputAdapter`.
- After the query runs, `executeColumnarRDD` is not materialized in any node.
- Before this line, the serialized `(stageRDD, taskFunc)` contains only
`MapPartitionsRDD` and `ShuffledRowRDD`, i.e., no upstream RDD. After this
line, it contains the columnar RDD and its whole lineage down to
`ParallelCollectionRDD`.
So the test doesn't reproduce a leak of this query, although the PR
description presents it as "a real cached-DataFrame aggregation across a
shuffle". `assert(scan.supportsColumnar)` doesn't prove that the columnar path
is used either. Could you assert that the job itself created `columnarRDD`,
like the "or the test exercises nothing" guards in
`DataFrameSetOperationsSuite`?
```scala
def lineage(rdd: RDD[_]): Seq[RDD[_]] = rdd +:
rdd.dependencies.flatMap(d => lineage(d.rdd))
assert(lineage(exchanges.head.inputRDD).exists(_ eq columnarRDD),
"the job must create the columnar RDD, or the test exercises
nothing")
```
This assertion fails with the cached DataFrame. A vectorized Parquet scan
passes it because `ColumnarToRowExec.inputRDDs()` materializes
`FileSourceScanExec.executeColumnarRDD` during the job. For Parquet, note that
`FileSourceScanExec.inputRDD` is still serialized after this fix, as discussed
in the previous round, so the identity check covers only the wrapper RDD.
Otherwise, please update the test comments and the PR description to say that
this line stands in for a columnar consumer, e.g., a plugin operator.
##########
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:
`serializedStageRDD` is always `true`. `stageRDD` is `_1` of the tuple
passed to `writeObject`, so `replaceObject` always sees it while the hook is
enabled, and this line only checks that the hook works.
Nothing checks that the downstream stage captures the upstream plan, which
is the precondition of line 98. Today the only path is
`ctx.addReferenceObj("plan", this)` of the final `HashAggregateExec`:
https://github.com/apache/spark/blob/a026a0be4c88e7ace59d1da7c18ab73b0d81edde/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala#L589
If it goes away, e.g., when whole-plan capture is reduced as discussed in
the previous round, line 98 passes even without this fix. I simulated it on
Apache Spark 4.3.0 by writing the final `HashAggregateExec` as `null` in the
same stream. Neither `scan` nor `columnarRDD` was serialized, so the assertions
of this test would pass without the fix. How about tracking the scan node
instead of `stageRDD`?
```scala
var serializedScan = false
...
serializedScan ||= obj eq scan
...
assert(serializedScan, "the downstream stage no longer captures the
upstream plan")
```
##########
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:
nit: This is not from this PR, but the Scaladoc of `executeColumnar()` right
below refers to `doColumnarExecute` twice (lines 227 and 230), which doesn't
exist. It should be this `doExecuteColumnar`.
##########
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()
Review Comment:
nit: `cached.count()` isn't needed because `cache()` already registers the
relation and the query builds the buffers on its first run. In addition,
`checkAnswer` runs the aggregation twice: `materializedRdd.count()` on a
separate `rddQueryExecution` plan, and then `collect()` on this plan. Only the
latter matters for the checks below, so `QueryTest.checkAnswer(query,
Seq(Row(0L, 4L), Row(1L, 4L)), checkToRDD = false)` would be enough.
##########
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:
nit: The serialized bytes are never read. `OutputStream.nullOutputStream()`
avoids buffering the whole task binary and makes the intent clearer, like the
`NullOutputStream` of `SerializationDebugger.ListObjectOutputStream`.
```scala
import java.io.{ObjectOutputStream, OutputStream}
...
val out = new ObjectOutputStream(OutputStream.nullOutputStream()) {
```
--
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]
