sunchao commented on code in PR #5615:
URL: https://github.com/apache/datafusion-comet/pull/5615#discussion_r3921441698
##########
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:
[P2] Scope prepared shuffle data to the shuffle-manager lifetime
For local or embedded callers that stop and recreate `SparkContext` in the
same JVM and Comet classloader, each context restarts shuffle IDs at zero. This
singleton survives, and neither Comet shuffle manager clears it in `stop()`.
When successive contexts perform native shuffles with new scan keys, those keys
accumulate in the same inner map if the reused IDs stay within the 16-entry
limit. An exact-source component probe retained 128 prepared commons under ID
`0` in one outer entry. The finalized-byte guard prevents stale reads but does
not remove old keys. Could this store be owned by the manager lifetime or
explicitly cleared on stop, with a recreated-context regression test? Otherwise
successive contexts can keep retaining more commons despite the outer LRU bound.
--
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]
