[comphead]

Which issue does this PR close?

Closes #7301.

Rationale for this change

A job launched from a query's RDD (Dataset.rdd with AQE off, QueryExecution.toRdd with AQE on or off) runs its map stages outside a SQL execution. On an executor, SQLConf.get then reads task local properties that Spark has not filled, and SQLExecutionRDD (SPARK-28939) re-applies the session's configs only in the final stage. Comet reads its settings on the executor when it builds a native plan and when it reads a shuffle, so those map stages ran with every spark.comet.* setting at its default. That includes the memory pool, the batch size and the configs passed to native. There was no warning.

What changes are included in this PR?

• CapturedSQLConf captures the session's spark.comet.* configs on the driver. withConf makes them what SQLConf.get returns on the task thread, using SQLConf.withExistingConf as SQLExecutionRDD does. It does nothing when a conf is already set on the thread, or when Spark passed the task the session's configs as local properties (inside a SQL execution, or under SQLExecution.withSQLConfPropagated).
• It applies where Comet builds native plans or reads shuffles on executors: CometExecRDD.compute, the native shuffle writer and the Celeborn writer's frame size (the configs ride on NativeShuffleSpec), the serialized-plan CometExec.getCometIterator used by the native limit and top-K, and CometShuffledBatchRDD.compute / computeAsShuffleBlockIterator.
• The withSQLConfPropagated workarounds in CometJoinSuite and CometTaskMetricsSuite are no longer needed and are removed.

Not covered here: the JVM columnar shuffle writers, tracked in #7303. CometColumnarShuffleSuite keeps its workaround until then. Native writes run inside a SQL execution, so they are not affected.

How are these changes tested?

New CometNativeShuffleSuite tests run a query from its RDD with AQE off, one per site:

• the native shuffle writer, a native plan in a map stage feeding the JVM shuffle, the native limit, and the native top-K each check, from the line CometExecIterator logs when it creates a plan, that every native plan resolved the session's memory pool
• a Comet shuffle read in an intermediate stage under Spark operators runs with spark.comet.shuffle.readBufferSize=0, which fails only if the read sees the session's value

The same query shapes in spark-shell (Spark 4.1.1): before this change, map-stage plans resolve fair_unified and the intermediate shuffle read succeeds. After it, every plan resolves the session's greedy_unified and the read fails as collect() does. The probe from #7301 resolves the submitted pool and fraction in the map stage in all five scenarios.

[andygrove]

Thanks @comphead. Capturing the session's configs on the driver and applying them the way SQLExecutionRDD does is a clean fix, and leaving tasks inside a SQL execution on the existing path keeps the change small. The repro in #7301 is very clear.

I compared against branch-1.1. The same bug is there, so this is a fix and not a behavior regression. I did not find a correctness or performance concern in the diff itself. I have two things I would like addressed before merge.

The new test only exercises the writer. input_batches == 100 shows that the iterator built in CometNativeShuffleWriter.write saw batchSize=10. Nothing covers CometNativeShuffleInputRDD.compute, CometExecRDD.compute, the limit closure in CometExecUtils, or the top-K and final closures in CometTakeOrderedAndProjectExec. If any of those wrappers were reverted, I think the suite would stay green. Could you add a case for each, following the same pattern? A toRdd over a LIMIT and over an ORDER BY ... LIMIT would cover the closures. A map stage whose child is a CometExecRDD would cover that one.

Some executor-side reads still happen outside these wrappers. The description mentions the JVM columnar writers. CometShuffledRowRDD.compute, computeAsShuffleBlockIterator and CometBlockStoreShuffleReader also read CometConf values such as COMET_SHUFFLE_READ_BUFFER_SIZE and COMET_TRACING_ENABLED. When one of them is the RDD of an intermediate reduce stage in a toRdd job, would it still see defaults? Inside CometExecRDD it is covered because the inputs resolve within withConf. Could you confirm which of these can run unwrapped? Then either wrap them here, or open an issue and link it from the description, since a "not covered here" note without an issue tends not to get done.

One smaller question. The sqlConf: CapturedSQLConf = CapturedSQLConf.capture() defaults on NativeShuffleSpec and CometNativeShuffleInputRDD mean the Arrow-stream path picks up the capture silently. That is correct today because both are built on the driver. Would you make the parameter required, so each site states that it captures on the driver? A default that reads SQLConf.get would capture the executor's defaults if someone later built one of these on an executor, and the bug would come back unnoticed.

[comphead]

Thanks @andygrove. All three are addressed in 04cdf52.

Tests per site. The single writer test is replaced by one test per site, each running the query from toRdd with AQE off:

• the native shuffle writer
• a native plan (CometExecRDD) in a map stage feeding the JVM shuffle
• LIMIT (CometCollectLimitExec, which goes through getNativeLimitRDD)
• ORDER BY ... LIMIT (CometTakeOrderedAndProjectExec, top-K and final)

Each one checks the line CometExecIterator logs when it creates a plan, and asserts that every native plan resolved the session's greedy_unified rather than the default fair_unified. A batch count only covered the plan fused into the writer. This check covers every native plan in the job. I ran the same query shapes in spark-shell: before the PR, each shape had map-stage plans on fair_unified (for LIMIT and top-K, 6 of 7 plans). With the PR, all of them resolve greedy_unified.

Shuffle reads outside the wrappers. Confirmed, they could run unwrapped. When a Comet shuffle read has no native plan above it in its stage, it read the defaults. An example is an intermediate stage where Spark's aggregate or project consumes a CometColumnarExchange. I checked it with spark.comet.shuffle.readBufferSize=0, which fails only where an executor reads the session's value. That stage failed under collect() and passed under toRdd. CometShuffledBatchRDD now captures the configs and applies them in compute and computeAsShuffleBlockIterator, and a test covers that case. Since shuffle reads now apply the configs themselves, the wrap in CometNativeShuffleInputRDD.compute became redundant (its only Comet config reads were those shuffle reads), so I removed it. That RDD now only carries the captured configs to the writer. The JVM columnar writers are tracked in #7303, with a repro, and linked from the description.

Defaults. The NativeShuffleSpec default was already gone. sqlConf on CometNativeShuffleInputRDD is now required, and both call sites pass CapturedSQLConf.capture(). CometExecRDD and CometShuffledBatchRDD capture in their constructors. An RDD can only be constructed on the driver, since it needs the SparkContext, so that stays driver-side.

I haven't run the new tests locally. They will run in this PR's CI.

[rich7420]

@comphead, thanks for the patch.

[rich7420]

Set the appender threshold to DEBUG; INFO drops these messages.

[comphead]

Good catch, thanks. LogAppender keeps INFO and above by default, so the check saw no lines. That is very likely why the [shuffle] job failed. Fixed in bcf4862 with appender.setThreshold(Level.DEBUG). Locally the five tests pass, and all five fail with CapturedSQLConf.withConf turned into a no-op.

[rich7420]

Native Celeborn reads maxFrameBytes before this wrapper. Apply captured configs during destination construction too.

[comphead]

Right, getWriter read spark.comet.shuffle.rss.maxFrameBytes from the session before the writer applied anything. Fixed in bcf4862: the captured configs now live on NativeShuffleSpec, which getWriter reaches through the dependency, and the frame-size read runs inside withConf. The writer takes them from the same spec. The other two limits in CelebornShufflePusherFactory.create (frame and in-flight bytes) come from the executor's SparkConf, not the session, so they are unchanged. I did not add a Celeborn test for this read.

[andygrove]

Thanks @comphead. The three asks from the earlier round are addressed. There is a test per site, the shuffle reads apply the captured configs, and sqlConf is required on the spec and the input RDD. The toRdd repro and the per-site checks of the logged memory pool make the fix easy to trust. I compared against branch-1.2 and the only change is that map stages of toRdd and .rdd jobs now see the session's spark.comet.* values, which is the intended fix.

I added the run-spark-4.1-tests label, since Spark's own suites use .rdd and toRdd a lot and this changes which SQLConf the executor-side code sees.

Nit, optional: CapturedSQLConf.withConf decides it is inside a SQL execution by checking that some captured key is a task local property. Spark only propagates the keys that are currently set. If a capture holds spark.comet.x and the session later unsets it, would a rerun map stage inside a SQL execution see no Comet key in its local properties and apply the stale capture? shuffleDependency is a lazy val on the cached executedPlan, so the capture can outlive the setting. It needs a map stage recompute, so it is narrow. Would also skipping when SQLExecution.EXECUTION_ID_KEY is set make that case exact? A small unit test for the skip conditions would pin it down.