andygrove commented on PR #5027:
URL:
https://github.com/apache/datafusion-comet/pull/5027#issuecomment-5441628694
> **Note on this review:** this was generated by an LLM (Claude Code) at my
request while I worked through a review backlog. I have not verified the
individual findings myself. Please treat everything below as suggestions to
evaluate rather than as authoritative review feedback, and push back on
anything that is wrong or already handled.
Getting JVM UDF output vectors onto Spark's task accounting is worth doing,
and the distinction between imported native buffers (already accounted, keep on
the root allocator) and JVM-created output buffers (charge to the task) is the
right one. The retention logic that keeps the allocator alive until the last
FFI release, while dropping Spark accounting at task completion, is carefully
thought through.
Four things.
**The threading contract for `CometUDF` changed**
The class doc goes from:
> At any instant at most one thread is inside `evaluate()` for a given
`taskAttemptId`.
to:
> Calls for a task may arrive concurrently from different Tokio workers.
Implementations with mutable state are responsible for synchronizing
`evaluate()`.
That is a breaking change to the contract that user-written `CometUDF`
implementations were told they could rely on. Anyone who wrote a UDF holding
mutable per-task state, which the old wording explicitly invited, now has a
data race.
Was the old guarantee wrong all along, or does this PR change when
concurrent calls can happen? Either way this deserves its own line in the
description and probably a migration-guide note, because it is a much bigger
deal for users than the memory accounting is. Right now it reads like an
incidental doc edit.
**Lock ordering deserves to be written down**
`onPreAllocation` takes the `TaskMemoryManager` monitor, then the
`TaskState` monitor, then calls back into
`TaskMemoryManager.acquireExecutionMemory`. `taskCompleted` and `onRelease`
take the `TaskState` monitor and then call into `MemoryConsumer.freeMemory`,
which reaches the `MemoryManager` monitor. Spark's own `acquireExecutionMemory`
can call `spill()` on other consumers while holding the `TaskMemoryManager`
monitor, and an Arrow buffer release from such a path would re-enter
`onRelease`.
I worked through it and did not find an actual inversion, but it took a
while and I am not certain. Could you add a short comment naming the three
locks and the order they must always be taken in? Anything that acquires a
Spark memory lock from inside an Arrow allocation listener callback deserves
that much.
**On-heap Tungsten mode silently does nothing**
```java
this.consumer = taskMemoryManager.getTungstenMemoryMode() ==
MemoryMode.OFF_HEAP
? new TaskMemoryConsumer(taskMemoryManager) : null;
```
With on-heap Tungsten memory the whole accounting path is disabled. That is
probably correct, since Arrow's buffers are off-heap and charging them to
Spark's on-heap pool would be wrong. But it is not stated anywhere, and a user
running on-heap who reads the release notes will believe their UDF memory is
accounted. Could the class doc say so explicitly, and ideally log once at debug
level when accounting is skipped?
**The registration-order dependency**
`registerTask` has to be called before `CometExecIterator` registers its own
completion listener, because Spark runs listeners in reverse order. That is
documented in the javadoc, which is good, but it is enforced by nothing. If
someone reorders those two calls the failure is a use-after-free of an
allocator with buffers still exported, which will not reproduce reliably.
Is there a way to make that structural rather than conventional, for example
having `CometExecIterator` obtain the `TaskState` and register both listeners
itself? If not, is one of the tests in `CometUdfBridgeTest` specifically
asserting the ordering?
--
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]
