Repository navigation
Description
Describe the bug
When a Comet plan runs in a job that is not a SQL execution, the native plans of its map stages
resolve every spark.comet.* execution setting to its default. That includes settings passed with
--conf at submit time. The driver plans the query with the session's values, and the result
stage gets them too, so nothing looks wrong until the map stage behaves differently. There is no
warning.
Affected: df.rdd with AQE off, and df.queryExecution.toRdd with AQE on or off, followed by an
action. A df.rdd action with AQE on, and every DataFrame action or write, resolve the right values.
Steps to reproduce
Spark 4.1.1, Comet main at 7d8d9b1544, local[2].
log4j2.properties, to print the memory settings each native plan resolves (the thread name
carries the stage):
rootLogger.level = warn
rootLogger.appenderRef.console.ref = console
appender.console.type = Console
appender.console.name = console
appender.console.layout.type = PatternLayout
appender.console.layout.pattern = [%t] %c{1}: %m%n
logger.cei.name = org.apache.comet.CometExecIterator
logger.cei.level = debugspark-shell --master 'local[2]' --jars comet-spark-spark4.1_2.13-1.2.0-SNAPSHOT.jar \
--driver-java-options "-Dlog4j2.configurationFile=file:$PWD/log4j2.properties" \
--conf spark.plugins=org.apache.spark.CometPlugin \
--conf spark.shuffle.manager=org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager \
--conf spark.memory.offHeap.enabled=true --conf spark.memory.offHeap.size=2g \
--conf spark.comet.exec.memoryPool=greedy_unified \
--conf spark.comet.exec.memoryPool.fraction=0.5spark.range(0, 200000, 1, 4).selectExpr("id % 1000 as k", "id as v").write.mode("overwrite").parquet("/tmp/t")
spark.conf.set("spark.sql.adaptive.enabled", "false")
spark.read.parquet("/tmp/t").groupBy("k").agg(sum("v")).rdd.foreachPartition(it => it.foreach(_ => ()))The 2 map-stage plans (scan, partial aggregate, native shuffle write) log the defaults, the 200
result-stage plans log the submitted values:
[Executor task launch worker for task 1.0 in stage 2.0 (TID 6)] CometExecIterator: memoryPoolType=fair_unified, offHeapSize=2048 MB, memoryFraction=1.0, memoryLimit=2048 MB
[Executor task launch worker for task 117.0 in stage 3.0 (TID 7)] CometExecIterator: memoryPoolType=greedy_unified, offHeapSize=2048 MB, memoryFraction=0.5, memoryLimit=1024 MB
The same query run five ways (driver: greedy_unified, fraction 0.5):
| how the query runs | map stage | result stage |
|---|---|---|
df.collect(), AQE on |
greedy_unified, 0.5 |
greedy_unified, 0.5 |
df.queryExecution.toRdd + action, AQE off |
fair_unified, 1.0 |
greedy_unified, 0.5 |
df.rdd + action, AQE off |
fair_unified, 1.0 |
greedy_unified, 0.5 |
df.rdd + action, AQE on |
greedy_unified, 0.5 |
greedy_unified, 0.5 |
df.queryExecution.toRdd + action, AQE on |
fair_unified, 1.0 |
greedy_unified, 0.5 |
Expected behavior
Every native plan resolves the same spark.comet.* values the driver planned the query with,
however the job is launched.
Additional context
Why it happens:
- On an executor,
SQLConf.getis aReadOnlySQLConfthat reads only the task's local
properties, with no fallback to theSparkConf. Spark fills them for jobs that run inside
SQLExecution.withNewExecutionId. - For RDD-level jobs,
SQLExecutionRDD(SPARK-28939) re-applies the session conf, but only in the
stage that computes that RDD, the result stage.Dataset.rddbuilds the RDD inside a SQL
execution, which covers the map stages only when AQE runs them while the RDD is built. CometExecIteratorreads its settings on the executor when it creates a native plan:
serializeCometSQLConfs()takes everyspark.comet.*entry fromSQLConf.get.getAllConfs
(including thespark.comet.datafusion.*passthrough) andgetMemoryConfigreads the pool type
and fraction (CometExecIterator.scalalines 123, 138, 876-878 and 920-928 onmain
a828088500).
We hit this while reproducing a memory issue. With greedy_unified,
respectDataFusionConfigs=true, enable_migration_aggregate=false and
sort_spill_reservation_bytes=256MiB set, a toRdd job's map stage ran with fair_unified and the
migrated aggregate, and every task failed with an ExternalSorterMerge allocation error. The same
query written with df.write passed.
A possible fix, not tried: resolve these settings on the driver where Comet builds its RDDs and
ship them with the tasks, instead of reading SQLConf.get on the executor.
Willingness to contribute
I can contribute a fix for this bug independently
Activity
fix: apply the session's Comet configs in map stages of jobs outside …
fix: [branch-1.2] apply the session's Comet configs in map stages of …