sunchao commented on code in PR #5027:
URL: https://github.com/apache/datafusion-comet/pull/5027#discussion_r4433048773
##########
spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java:
##########
@@ -317,10 +375,30 @@ public static void releasePlan(long taskAttemptId, long
planId) {
/** The instance of {@code udfClassName} that task {@code taskAttemptId}
holds, or null. */
@VisibleForTesting
public static CometUDF instanceFor(long taskAttemptId, String udfClassName) {
- ConcurrentHashMap<String, CometUDF> perTask = INSTANCES.get(taskAttemptId);
+ ConcurrentHashMap<String, CometUDF> perTask = instancesFor(taskAttemptId);
return perTask == null ? null : perTask.get(udfClassName);
}
+ /** Sentinel key for calls that carry no {@link TaskContext} (unit tests,
direct driver). */
+ private static final long NO_TASK_ID = -1L;
+
+ /**
+ * The UDF instances of task {@code taskAttemptId}, or null when the task
holds none. The
+ * instances live in that task's {@link TaskState}, which owns their close
and their Arrow
+ * accounting, so this resolves the id against {@link #TASKS} rather than a
second cache.
+ */
+ private static ConcurrentHashMap<String, CometUDF> instancesFor(long
taskAttemptId) {
+ if (taskAttemptId == NO_TASK_ID) {
+ return NO_TASK_INSTANCES;
+ }
+ for (TaskState state : TASKS.values()) {
Review Comment:
Releasing a plan now scans every live task state for a matching numeric
attempt ID, even though registration keys the map by `TaskContext`. Releasing
one plan for each of T active tasks can take O(T²) lookups. It also makes
distinct contexts with the same attempt ID indistinguishable while both states
are live. Could teardown pass the captured context and look up the same key
used at registration?
--
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]
