peterxcli opened a new pull request, #6552: URL: https://github.com/apache/datafusion-comet/pull/6552
## Which issue does this PR close? Closes #. ## Rationale for this change The JVM codegen dispatcher (`spark.comet.exec.scalaUDF.codegen.enabled`, on by default) leaks off-heap Arrow memory on every batch whose output is a top-level struct. Examples are a ScalaUDF that returns a case class or a tuple, and `from_json`. Arrow Java's `StructVector` builds its writer in a field initializer, `private final NullableStructWriter writer = new NullableStructWriter(this)`. The `NullableStructWriter` constructor walks `getField().getChildren()` and calls `addOrGet` and then `allocateNewSafe()` for each child. `allocateOutput` builds the struct as `new RenamedStructVector(field, allocator)` from the export field, which carries the children. So by the time the constructor returns, the struct already holds a default-capacity child for each field. `allocateOutput` then calls `initializeChildrenFromFields`, which goes through `AbstractStructVector.add` and `putVector`. Under the default `CONFLICT_REPLACE` policy, that removes the existing children from the struct without closing them. Closing the struct later closes only the replacement children. The originals are still referenced by the writer, and their buffers are never released. That is 65,536 bytes per batch for `struct<name:string,age:int>` and 81,920 bytes for `struct<_1:bigint,_2:string>`. On `main` the output comes from `CometArrowAllocator`, a root allocator with no limit that is never closed, so nothing reports the leak. It showed up with the per-task allocator from #5027. With `-Darrow.memory.debug.allocator=true`, every outstanding ledger was allocated under `NullableStructWriter.<init>` ← `StructVector.<init>` ← `RenamedStructVector.<init>` ← `CometBatchKernelCodegenOutput.allocateOutput` ← `CometScalaUDFCodegen.evaluate`. Arrow 19.0.0 builds the writer the same way, so an upgrade would not fix this. List and Map outputs are not affected. `ListVector` and `MapVector` create their writers on demand, and the struct nested under them is created from a field without children. ## What changes are included in this PR? - In the struct case of `allocateOutput`, close the children the writer created before `initializeChildrenFromFields` replaces them. The writer keeps its references to the closed vectors. Comet never uses it, because the generated kernel writes through `getChildByOrdinal`. - `allocateOutput` takes the allocator as a parameter, defaulting to `CometArrowAllocator`, so a test can check what a closed output leaves allocated. Existing callers are unchanged. ## How are these changes tested? A new test in `CometCodegenSuite` allocates outputs of several types from a fresh child allocator of `CometArrowAllocator`, closes each one, and asserts that the allocator is back to 0 bytes. Without the fix, it fails on the first type with `Memory was leaked by query. Memory leaked: (65536)`. Measured per type before the fix: | Output type | Bytes left after `close()` | | --- | --- | | `struct<name:string,age:int>` | 65,536 | | `struct<_1:bigint,_2:string>` | 81,920 | | `struct<inner:struct<name:string,age:int>,tags:array<string>,attrs:map<string,int>>` | 34,304 | | `array<struct<name:string,age:int>>` | 0 | | `map<string,struct<name:string,age:int>>` | 0 | | `string` | 0 | With the fix, every type is back to 0. `CometCodegenSuite` passes (105 tests), and so do the other codegen dispatcher suites (`CometCodegenSourceSuite`, `CometCodegenFuzzSuite`, `CometCodegenHOFSuite`, `CometJsonJvmSuite`, `CometJsonExpressionSuite` and `CometScalaUDFClassLoaderSuite`, 108 tests). ## Backporting `branch-1.0` and `branch-1.1` have the same `allocateOutput` code, use Arrow 18.3.0 and enable the dispatcher by default, so they leak the same way. This may need backporting to both (`backport-1.0`, `backport-1.1`); see the [backporting guide](https://github.com/apache/datafusion-comet/blob/main/docs/source/contributor-guide/backporting.md). The commit applies cleanly to both branches. ## Relation to #5603 #5603 moves this allocation into `NativeUtil.createVector` and builds the struct from a field without children, which also stops the writer from allocating. This PR is the small fix for `main` and the release branches. If #5603 lands first, this PR reduces to the regression test. -- 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]
