andygrove opened a new pull request, #5998: URL: https://github.com/apache/datafusion-comet/pull/5998
## Which issue does this PR close? Part of #5997. This is the first of the four items listed there, and the only one that needs no native changes. ## Rationale for this change `CometArrowAllocator` is a single process-wide `RootAllocator(Long.MaxValue)`. Child allocators are cut from it for FFI stream export, broadcast coalescing, `CometSparkToColumnarExec`, the cached batch serializer, codegen output, and the Python runner. Every byte it hands out is real off-heap memory resident in the container, and none of it is visible to Spark's `TaskMemoryManager` or to Comet's native memory pool. The [memory management guide](https://github.com/apache/datafusion-comet/blob/main/docs/source/contributor-guide/memory_management.md) already lists this as an open problem: "unbounded and accounted by nobody". That makes these bytes a blind spot exactly when it matters, which is when an executor is killed for exceeding its container limit. Gluten solves the same problem by attaching an allocation listener to its Arrow root allocator so that allocations are reserved from Spark in fixed-size blocks. This PR closes the reporting half of the gap. It deliberately does not close the enforcing half: Arrow allocation on these paths cannot fail today, and making it fail is a behavioural change that deserves its own PR once we have numbers on what the real volumes are. ## What changes are included in this PR? - `CometArrowAllocationListener`, an Arrow `AllocationListener` that charges each allocation to a Spark `MemoryConsumer` belonging to the task that made it, so the bytes appear in `TaskMemoryManager.showMemoryUsage` and are arbitrated against Spark's other off-heap consumers. - The listener is attached to `CometArrowAllocator`. Arrow's `BaseAllocator.newChildAllocator` passes the parent's listener down, so one attachment point covers every call site including the three that cut child allocators. - Reservations are grown and shrunk in whole blocks (1 MiB by default, `spark.comet.arrowAllocator.accounting.blockSize`) because Arrow allocates per buffer and `acquireExecutionMemory` takes locks. Only block-crossing changes reach Spark. - `spark.comet.arrowAllocator.accounting.enabled` (default true) turns the reporting off. These are read from the `SparkConf` rather than declared as `CometConf` entries because the listener is constructed on first touch of a package-object `val`, which can happen before any `SparkSession` exists and on executors where `SQLConf` does not carry Comet's settings. - The new suite is registered in both `pr_build_linux.yml` and `pr_build_macos.yml`. Three cases are deliberately no-ops, each for a different reason: - **No active task.** Broadcast coalescing and the cached batch serializer can allocate from the driver or a non-task thread, where there is no task to charge. - **On-heap mode.** Comet's on-heap mode exists so the Spark SQL suite can run without off-heap memory configured, and charging an off-heap consumer there would be wrong. Note this differs from `CometUnifiedShuffleMemoryAllocator`, which throws in that situation; throwing here would break those tests. - **A buffer released after its allocating task has finished.** The allocator is process-wide precisely because buffers can outlive the task that created them, so the task's reservation is dropped at task end and later releases are ignored rather than double-counted. ## How are these changes tested? A new suite, `CometArrowAllocationListenerSuite`, covering: - a sub-block allocation reserves exactly one block, and a second sub-block allocation does not ask Spark again - a request larger than one block reserves enough to cover it - releasing returns whole blocks, leaving nothing reserved - no active task is a no-op rather than an error, which is the property that keeps the driver-side paths working - on-heap mode is not accounted -- 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]
