Repository navigation
CometShuffleManager treats Spark's own shuffles differently from SortShuffleManager #7488
Description
Describe the bug
CometShuffleManager handles shuffles that Comet does not replace differently from Spark's SortShuffleManager. These shuffles are RDD shuffles and Spark exchanges Comet leaves in place. They are not new, but they reach more applications once the plugin sets CometShuffleManager by default on Spark 4.0 and later (#7336). Found in review of that PR.
- Extra row copy in Spark exchanges.
ShuffleExchangeExec.needToCopyObjectsBeforeShufflereturns true for any manager that is not aSortShuffleManager, so Spark's own exchanges copy every row before the shuffle. sortShuffleManager.getWriterentries are never released. For a Spark shuffle dependency,getWriterdelegates to the innersortShuffleManager, which records each map task in its owntaskIdMapsForShuffle.unregisterShuffleclears onlyCometShuffleManager's map and never callssortShuffleManager.unregisterShuffle, so those entries stay on the executor heap until the application ends.- Blocks migrated in during decommissioning are never deleted. The
IndexShuffleBlockResolveris built with a freshConcurrentHashMapinstead oftaskIdMapsForShuffle. The resolver records blocks migrated in during decommissioning in that map, sounregisterShufflenever deletes them.
All three are in spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala.
Steps to reproduce
Read the code paths above. No reproduction has been run for any of them.
Expected behavior
A shuffle that Comet does not replace behaves as it does under SortShuffleManager: no extra row copy, and no state left behind after unregisterShuffle.
Activity
rich7420 commentedon Nov 10, 2026
Item 3 can be worse than a leak on Spark 4.x: the resolver can end up with no usable map at all.
shuffleBlockResolvertakes the first 2- or 3-parameter constructor thatgetDeclaredConstructorsreturns, and the JDK does not define that order. Spark 4.x has two 2-parameter constructors, andc.newInstance(conf, null)gives:(SparkConf, Map)on 4.0,(SparkConf, ConcurrentHashMap)on 4.1+:taskIdMapsForShuffleisnull(SparkConf, BlockManager)on 4.0:Collections.emptyMap()(SparkConf, BlockManager)on 4.1+: a fresh map, as item 3 describesWhen a shuffle block migrates in,
onCompleteinputShuffleBlockAsStreamrenames the file into place and then callstaskIdMapsForShuffle.computeIfAbsent(IndexShuffleBlockResolver.scala#L304-L319). That throws aNullPointerExceptiononnulland anUnsupportedOperationExceptiononemptyMap, so the block is never reported. The sending executor then stops migrating to the receiving executor and queues the block for another one, giving up after 3 failures by default. The renamed file stays on the receiving executor, untracked.I loaded
IndexShuffleBlockResolverin a standalone JVM (JDK 17) and printed its constructor order across 30 runs per version:(SparkConf, Map)/(SparkConf, ConcurrentHashMap)first(SparkConf, BlockManager)firstI haven't checked which constructor a real executor picks, or run decommissioning end to end. Spark 3.5 has a single constructor, so it isn't affected.
Changing only
case 3to passtaskIdMapsForShufflewouldn't fix this, because.headcan still pick a 2-parameter constructor. Calling the 3-parameter constructor withtaskIdMapsForShuffleshould fix both. Is anyone working on this? If not, I can open a PR for item 3.