Repository navigation
NativeUtil.getNextBatch leaks Arrow structs when importVector fails #5534
Description
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:
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.
Activity
pingzh commentedon Aug 31, 2026
Follow-up to #5352 (Support Apache Celeborn remote shuffle service), discovered during post-merge review of #5531.