andygrove commented on PR #5027: URL: https://github.com/apache/datafusion-comet/pull/5027#issuecomment-5876350753
This is a light fully automated review since there are so many PRs open. 1. `chargedOutputSize` (`spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java:395`) assumes every byte a task-allocator ledger owns was charged in `onPreAllocation`. But Arrow's `transferOwnership` moves ownership through `forceAllocate`, which never calls the listener, and imported inputs are owned by a ledger on `CometArrowImportAllocator`. So if a UDF returns an input moved into `allocator` with `getTransferPair(allocator)` and `transfer()`, or an aligned `splitAndTransfer` slice of one (the comment at line 284 already expects input slices), the ref-count check still passes and `releaseExportedCharge` frees bytes this consumer never acquired. Spark releases from the task-wide total, so that comes out of the task's native reservations, and `accounted` goes negative. Closing a `transfer()`-ed input inside the UDF does the same through `onRelease`. Could releases be capped at what this consumer holds, with a test where another consumer holds memory while the UDF returns a transferred input? 2. `taskCompleted` (`CometUdfBridge.java:614`) drops UDF instances without closing anything, and `closeIfIdle` waits for the allocator to be empty. `CometUDF.scala:52` still invites scratch buffers in fields, and line 38 now asks for temporary buffers to come from `allocator`. A UDF holding one leaves the allocator non-empty for good. Arrow's root keeps every unclosed child in `childAllocators`, so each such task pins its `TaskState`, `TaskContext` and `TaskMemoryManager` for the executor's lifetime, where before only the off-heap bytes leaked. Would a `close()` hook on `CometUDF`, called once the task is complete and nothing is in flight, make sense? 3. `docs/source/contributor-guide/memory_management.md:133` still routes `CometBatchKernelCodegenOutput` to `CometArrowAllocator` and "accounted by nobody", the child allocator list at line 112 has no per-task UDF allocator, and line 167 calls `CometTaskMemoryManager` the one place Comet acts as a Spark `MemoryConsumer`. Could this PR update those? -- 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]
