dwsmith1983 commented on code in PR #5615:
URL: https://github.com/apache/datafusion-comet/pull/5615#discussion_r3924448154
##########
spark/src/main/scala/org/apache/spark/sql/comet/operators.scala:
##########
@@ -136,9 +230,52 @@ private[comet] object PlanDataInjector extends Logging {
* reference rather than rebuilt; only the root-to-scan paths are rebuilt.
*/
def injectPlanData(
+ op: Operator,
+ commonByKey: Map[String, Array[Byte]],
+ partitionByKey: Map[String, Array[Byte]]): Operator =
+ injectPlanData(op, commonByKey, partitionByKey, null)
+
+ /**
+ * Injects planning data into a cached base plan, memoizing each scan's
prepared common inside
+ * the plan's own cache entry so it is prepared once per stage rather than
once per task.
+ */
+ def injectPlanData(
+ cachedPlan: CachedPlanData,
+ commonByKey: Map[String, Array[Byte]],
+ partitionByKey: Map[String, Array[Byte]]): Operator =
+ injectPlanData(cachedPlan.plan, commonByKey, partitionByKey,
cachedPlan.preparedCommons)
+
+ // The native shuffle writer's unified plan differs per task (output paths),
so those plans
+ // never pass through parseBasePlan. Prepared commons for that path are
scoped to the
+ // shuffleId instead: one shuffle stage's scans still share a single
eviction unit.
+ private val shufflePreparedCommons = java.util.Collections.synchronizedMap(
+ new LinkedHashMap[Integer, ConcurrentHashMap[String, PreparedCommon]](4,
0.75f, true) {
Review Comment:
Fixed in f24e4fcc6, at both boundaries you named. unregisterShuffle now
releases that shuffle's prepared commons in both managers, which is the precise
lifetime (verified the ContextCleaner path in the 3.5.9 bytecode:
doCleanupShuffle to BlockManagerStorageEndpoint RemoveShuffle to
ShuffleManager.unregisterShuffle, and neither Comet manager delegates that call
away), and stop() clears the whole store as the safety net for recreated
contexts. Your reused-id scenario is the regression: two manager instances with
a stop between them, shuffle 0 holds exactly the second context's keys instead
of four. There is also a real SparkContext suite that runs a native shuffle,
proves the cleaner path releases the entry, stops the context, recreates a
session, and asserts the store holds only the new context's keys. The base plan
cache is left alone deliberately since it is keyed by plan bytes rather than a
per-context counter, so a recreated context either hits or evicts through the
existing bo
und. One residual note: a straggling map task calling injection after
unregister would re-insert an empty inner map for that id, bounded by the outer
LRU and cleared on stop, so I did not add a tombstone.
--
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]
