fix: prevent memory leak after failed Arrow vector import - #5539
fix: prevent memory leak after failed Arrow vector import#55391fanwang wants to merge 3 commits into
Conversation
Signed-off-by: 1fanwang <1fannnw@gmail.com>
peterxcli
left a comment
There was a problem hiding this comment.
I reproduced one remaining ownership leak on 79d97e1, so I am requesting changes.
A valid UInt4Vector imports successfully, then CometVector.getVector rejects its unsigned logical type. The current FieldVector has not been appended, while its Arrow structs have already been consumed. A focused Spark 4.1 NativeUtilSuite regression leaves 48 bytes allocated; the other five tests pass. The same missing rollback exists when ArrowImporter.importArray fails after field.createVector.
Please keep the current FieldVector rollback-owned until import, wrapping, and publication complete, and add a post-import failure regression. Arrow 18.3's AutoCloseables is already available; no new cleanup abstraction is needed.
Non-blocking: NativeUtil documents that it must be closed but does not extend AutoCloseable; extending it would make that lifecycle contract explicit. The test-local IntVector can use Using.resource.
CI has not run: the workflow runs stopped as action_required. Beyond this ownership gap, I found no other correctness or concurrency issue in this two-file diff.
Signed-off-by: 1fanwang <1fannnw@gmail.com>
Generated-by: GitHub Copilot CLI (GPT-5.6 Sol) Signed-off-by: 1fanwang <1fannnw@gmail.com>
andygrove
left a comment
There was a problem hiding this comment.
Thanks for taking this on, and thanks for the detailed raw test output in the description. The core fix looks right to me. I walked the Arrow 18.3.0 ArrayImporter source and the ownership reasoning holds on every exit I could find. importArray moves the caller's ArrowArray into an internal ownedArray and closes the source before doImport can fail, so the array.isClosed() guard correctly separates "the importer took it" from "it never got that far". Calling release() on a freshly allocated struct whose release callback is NULL is safe, which the EOF path merged in #5531 already depends on.
CI has not run here. The workflow runs are sitting in action_required and need a maintainer to approve them. Since this touches every native operator path I ran the relevant pieces locally against 1f098bb4:
| Check | Result |
|---|---|
test-compile, Spark 3.5 / Scala 2.12 |
pass |
NativeUtilSuite, Spark 3.5 / Scala 2.12 |
7/7 pass |
NativeUtilSuite, Spark 4.1 / Scala 2.13 |
7/7 pass |
spotless:check |
pass |
I also checked the regressions are not vacuous. Reverting just NativeUtil.scala and ArrowImporter.java to apache/main while keeping the new tests gives three failures, all Memory was leaked by query, at 176, 288 and 33592 bytes. Each passes with the fix in place. NativeUtilSuite is already registered in both pr_build_linux.yml and pr_build_macos.yml, so no workflow change is needed.
One thing worth recording for whoever touches the harness next. These tests catch the leaks indirectly. withIsolatedStructAllocator only isolates the C struct allocations, and the imported vector's buffers come from the global CometArrowAllocator. An unreleased import keeps the exported source buffers pinned on the isolated allocator, and that is what trips the assertion. I would not change it, since asserting on the shared global allocator would be flaky. It is just worth a note somewhere so nobody later simplifies the harness and quietly makes the assertions blind.
I have left a few comments inline. The dictionary-encoded case is the one I would most like your read on.
| return vector; | ||
| } catch (RuntimeException | Error failure) { | ||
| if (vector != null) { | ||
| AutoCloseables.close(failure, vector); |
There was a problem hiding this comment.
For a dictionary-encoded column, ArrayImporter.doImport loads the dictionary values into the provider's vector before it loads the main data, and those buffers reference the same ReferenceCountedArrowArray as the column itself. If the main data import then fails, closing vector here does not drop those references, so the C release callback only fires when NativeUtil.close() closes the provider. The same applies when CometVector.getVector throws on a dictionary column, since the rollback there is only the indices vector.
Does that match your reading? It is a deferred release rather than a permanent leak, and it is bounded by task lifetime. But the trigger this is aimed at is allocator pressure, which is exactly when holding a native batch until end of task hurts most. Would clearing the failing column's dictionary vector be enough? A test with a dictionary-encoded column would pin the behavior down either way.
| if (vector != null) { | ||
| AutoCloseables.close(failure, vector); | ||
| } | ||
| if (!array.isClosed()) { |
There was a problem hiding this comment.
isClosed() is package-private and annotated @VisibleForTesting in Arrow. It works here because this class shares Arrow's package, and I do not think there is a better signal available.
Could you add a short comment saying what it distinguishes, namely that ArrayImporter.importArray closes the source array before it can fail in doImport? Without that, a future Arrow upgrade that drops the method is likely to get fixed by deleting the guard, and that turns this into a double release rather than a leak.
| val arrowSchema = schemas(i) | ||
| val arrowArray = arrays(i) | ||
|
|
||
| firstUnconsumed = i + 1 |
There was a problem hiding this comment.
Setting this before the import is the right call, but it is worth a comment saying why. It relies on importer.importVector consuming both the schema and the array on every exit, success or failure. The schema goes through importField's finally and the array through the new catch in ArrowImporter. That is what makes it safe to exclude column i from releaseArrowStructs.
Right now that invariant lives entirely in ArrowImporter and there is nothing here pointing at it. Getting it wrong in either direction gives you a leak or a double release.
| @@ -253,16 +254,33 @@ class NativeUtil { | |||
| */ | |||
| def importVector(arrays: Array[ArrowArray], schemas: Array[ArrowSchema]): Seq[CometVector] = { | |||
There was a problem hiding this comment.
Could you extend the scaladoc to say that on failure this method releases everything it was given? getNextBatch does its own cleanup on the other two exits, so the next person reading it could reasonably wrap this call in cleanup too and cause a double release.
| arrayVectors += cometVector | ||
| } catch { | ||
| case failure: Throwable => | ||
| val rollback = if (cometVector == null) arrowVector else cometVector |
There was a problem hiding this comment.
The var plus null check plus rollback makes the reader reason about a case that cannot happen, since arrayVectors += cometVector will not throw. This is equivalent and drops both:
val cometVector =
try CometVector.getVector(arrowVector, dictionaryProvider)
catch {
case failure: Throwable =>
AutoCloseables.close(failure, arrowVector)
throw failure
}
arrayVectors += cometVectorIn code whose whole job is memory ownership, the fewer branches a reader has to hold in their head the better.
| Data.exportVector(allocator, intStruct, null, arrays(0), schemas(0)) | ||
| Data.exportVector(allocator, stringStruct, null, arrays(1), schemas(1)) | ||
| intStruct.close() | ||
| stringStruct.close() |
There was a problem hiding this comment.
intStruct and stringStruct are closed unconditionally before the try, so if either Data.exportVector throws they leak and the failure surfaces as an allocator complaint rather than the real error. You use Using.resource in the two tests above. Could this one do the same?
peterxcli
left a comment
There was a problem hiding this comment.
lgtm now, thanks for the update!
Which issue does this PR close?
Closes #5534.
Rationale for this change
Failed Arrow imports leak memory in Comet executors. When a later column fails after earlier columns succeeded, the allocator reports
Memory was leaked by query. Memory leaked: (176), which compounds failures caused by memory pressure.What changes are included in this PR?
ArrowImporterkeeps a newFieldVectorrollback-owned until array import succeeds.NativeUtil.importVectordoes the same until Comet wrapping and publication complete, then closes earlier vectors and untouched C Data structs if a later column fails. Cleanup errors are attached to the original exception instead of replacing it.How are these changes tested?
The regressions cover schema import failure, array import after buffers are attached, and wrapper rejection after a successful import. Each path preserves the original exception, adds no synthetic suppressed failures, and returns allocator memory to zero.
Raw test output