andygrove opened a new issue, #5534:
URL: https://github.com/apache/datafusion-comet/issues/5534
Discovered during post-merge review of #5531.
**Component**:
`spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala:209`
#5531 added `releaseArrowStructs` to two of the three exits from
`getNextBatch`: the `func` failure
path (line 197) and the EOF path (line 206). The third exit,
`importVector(arrays, schemas)` on
line 209, is not guarded.
When `importer.importVector` throws part-way through the column loop, every
struct from the failing
column onward keeps its populated C data with the release callback never
invoked, and all
`2 * numCols` `ArrowArray` / `ArrowSchema` wrapper buffers leak from the
allocator.
This is on every native operator path, not just shuffle. The most reachable
trigger is an allocator
OOM inside `importVector` under memory pressure, which is exactly when
leaking is worst.
## Reproduction
Add to `NativeUtilSuite`, reusing the harness #5531 introduced:
```scala
test("getNextBatch leaks Arrow structs when importVector fails") {
withIsolatedStructAllocator { (nativeUtil, allocator, _) =>
val vector = new IntVector("value", allocator)
vector.allocateNew(4)
vector.setSafe(0, 42)
vector.setValueCount(1)
intercept[Throwable] {
nativeUtil.getNextBatch(2, (arrays, schemas) => {
// Column 0 exported; column 1 left as a freshly allocated struct,
which is what native
// leaves behind if it populates some columns and the importer then
rejects one.
Data.exportVector(allocator, vector, null,
ArrowArray.wrap(arrays(0)), ArrowSchema.wrap(schemas(0)))
vector.close()
1L
})
}
assert(allocator.getAllocatedMemory == 0) // fails: 176
}
}
```
Observed:
```
importVector failure = java.lang.IllegalStateException: Cannot
import released ArrowSchema
allocated after failure = 176
```
The success path is clean by comparison, so this is specific to the import
failure:
```
success allocated before close = 48
success allocated after close = 0
```
## Suggested fix
Wrap the `importVector` call in the same `try` /
`releaseArrowStructs(arrays, schemas, failure)`
pattern the other two exits use. The comment at
`NativeBatchDecoderIterator.scala:105` already
treats "Arrow import failures" as a first-class case, so this looks like an
oversight rather than a
decision.
Reproduced on Spark 4.1.3 / Scala 2.13 / JDK 17 at merge commit `98cd8c967`.
--
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]