You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "1fanwang (via GitHub)" <gi...@apache.org> on 2026/08/29 07:19:04 UTC
[PR] fix: prevent memory leak after failed Arrow vector import [datafusion-comet]
1fanwang opened a new pull request, #5539:
URL: https://github.com/apache/datafusion-comet/pull/5539
## Which issue does this PR close?
Closes https://github.com/apache/datafusion-comet/issues/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?
`importVector` closes vectors imported before the failure and releases Arrow structs from the first unconsumed column onward. Cleanup errors are attached to the original import exception instead of replacing it.
## How are these changes tested?
The regression exports the first of two columns, forces the second import to fail, checks the original exception, and verifies that allocator memory returns to zero.
<details>
<summary>Raw test output</summary>
```console
$ export JAVA_HOME=$(/usr/libexec/java_home -v 17)
$ cd native && cargo build --release && cd ..
$ git restore --source=98cd8c967995906500b12a709c25724b4cfba634 -- spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala
$ ./mvnw -Prelease -Pspark-4.1 test -Dtest=none -Dsuites=org.apache.comet.vector.NativeUtilSuite
java.lang.IllegalStateException: Memory was leaked by query. Memory leaked: (176)
Total number of tests run: 5
Tests: succeeded 4, failed 1, canceled 0, ignored 0, pending 0
BUILD FAILURE
$ git restore --source=HEAD -- spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala
$ ./mvnw -Prelease -Pspark-4.1 test -Dtest=none -Dsuites=org.apache.comet.vector.NativeUtilSuite
Total number of tests run: 5
Tests: succeeded 5, failed 0, canceled 0, ignored 0, pending 0
All tests passed.
Found 0 errors
Spotless.Scala is keeping 445 files clean
BUILD SUCCESS
```
</details>
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: prevent memory leak after failed Arrow vector import [datafusion-comet]
Posted by "1fanwang (via GitHub)" <gi...@apache.org>.
1fanwang commented on code in PR #5539:
URL: https://github.com/apache/datafusion-comet/pull/5539#discussion_r3886857057
##########
spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala:
##########
@@ -254,15 +254,29 @@ class NativeUtil {
def importVector(arrays: Array[ArrowArray], schemas: Array[ArrowSchema]): Seq[CometVector] = {
val arrayVectors = mutable.ArrayBuffer.empty[CometVector]
- (0 until arrays.length).foreach { i =>
- val arrowSchema = schemas(i)
- val arrowArray = arrays(i)
+ try {
+ (0 until arrays.length).foreach { i =>
+ val arrowSchema = schemas(i)
+ val arrowArray = arrays(i)
- arrayVectors += CometVector.getVector(
- importer.importVector(arrowArray, arrowSchema, dictionaryProvider),
- dictionaryProvider)
+ arrayVectors += CometVector.getVector(
+ importer.importVector(arrowArray, arrowSchema, dictionaryProvider),
+ dictionaryProvider)
+ }
+ arrayVectors.toSeq
+ } catch {
+ case failure: Throwable =>
+ val firstUnconsumed = arrayVectors.length
+ arrayVectors.foreach { vector =>
Review Comment:
Done in https://github.com/apache/datafusion-comet/commit/1f098bb49c90c63f21a94adafdaab482f25d5bc6. Cleanup now uses Arrow's AutoCloseables . NativeUtil implements AutoCloseable , and the test vectors use Using.resource .
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: prevent memory leak after failed Arrow vector import [datafusion-comet]
Posted by "peterxcli (via GitHub)" <gi...@apache.org>.
peterxcli commented on code in PR #5539:
URL: https://github.com/apache/datafusion-comet/pull/5539#discussion_r3886125731
##########
spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala:
##########
@@ -254,15 +254,29 @@ class NativeUtil {
def importVector(arrays: Array[ArrowArray], schemas: Array[ArrowSchema]): Seq[CometVector] = {
val arrayVectors = mutable.ArrayBuffer.empty[CometVector]
- (0 until arrays.length).foreach { i =>
- val arrowSchema = schemas(i)
- val arrowArray = arrays(i)
+ try {
+ (0 until arrays.length).foreach { i =>
+ val arrowSchema = schemas(i)
+ val arrowArray = arrays(i)
- arrayVectors += CometVector.getVector(
- importer.importVector(arrowArray, arrowSchema, dictionaryProvider),
- dictionaryProvider)
+ arrayVectors += CometVector.getVector(
Review Comment:
[P1] Keep the current `FieldVector` rollback-owned until publication
`arrayVectors` only owns earlier columns here. If `importer.importVector` fails after `field.createVector` / `importArray` has attached buffers, or if it succeeds and `CometVector.getVector` throws, the current `FieldVector` is unreachable. The C structs are already consumed, so `releaseArrowStructs` cannot recover it. I reproduced the latter with a valid `UInt4Vector`: allocator close reports 48 leaked bytes.
Please close the raw `FieldVector` on `ArrowImporter.importVector` failure and retain rollback ownership across wrapping and `+=`; commit only after it is appended. Add a failure-after-import regression—the current test fails during schema import before a `FieldVector` exists.
##########
spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala:
##########
@@ -254,15 +254,29 @@ class NativeUtil {
def importVector(arrays: Array[ArrowArray], schemas: Array[ArrowSchema]): Seq[CometVector] = {
val arrayVectors = mutable.ArrayBuffer.empty[CometVector]
- (0 until arrays.length).foreach { i =>
- val arrowSchema = schemas(i)
- val arrowArray = arrays(i)
+ try {
+ (0 until arrays.length).foreach { i =>
+ val arrowSchema = schemas(i)
+ val arrowArray = arrays(i)
- arrayVectors += CometVector.getVector(
- importer.importVector(arrowArray, arrowSchema, dictionaryProvider),
- dictionaryProvider)
+ arrayVectors += CometVector.getVector(
+ importer.importVector(arrowArray, arrowSchema, dictionaryProvider),
+ dictionaryProvider)
+ }
+ arrayVectors.toSeq
+ } catch {
+ case failure: Throwable =>
+ val firstUnconsumed = arrayVectors.length
+ arrayVectors.foreach { vector =>
Review Comment:
Non-blocking: Arrow 18.3 already provides `org.apache.arrow.util.AutoCloseables`; since `CometVector` is `AutoCloseable`, this loop can be `AutoCloseables.close(failure, arrayVectors.toSeq: _*)`. It closes every resource and attaches cleanup failures to the original. Caveat: it catches `Exception` rather than every `Throwable`, so keep the manual loop if continuing after a close-time `Error` is intentional. No new Comet helper is needed.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: prevent memory leak after failed Arrow vector import [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on PR #5539:
URL: https://github.com/apache/datafusion-comet/pull/5539#issuecomment-5485678397
Merged, thanks @1fanwang for the contribution and @andygrove @peterxcli for the review!
One nit in the above comments
> These tests catch the leaks indirectly. `withIsolatedStructAllocator` only isolates the C struct allocations, and the imported vector's buffers come from the global `CometArrowAllocator`.
Might worth adding a short code comment to explain `withIsolatedStructAllocator`.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: prevent memory leak after failed Arrow vector import [datafusion-comet]
Posted by "1fanwang (via GitHub)" <gi...@apache.org>.
1fanwang commented on code in PR #5539:
URL: https://github.com/apache/datafusion-comet/pull/5539#discussion_r3886856326
##########
spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala:
##########
@@ -254,15 +254,29 @@ class NativeUtil {
def importVector(arrays: Array[ArrowArray], schemas: Array[ArrowSchema]): Seq[CometVector] = {
val arrayVectors = mutable.ArrayBuffer.empty[CometVector]
- (0 until arrays.length).foreach { i =>
- val arrowSchema = schemas(i)
- val arrowArray = arrays(i)
+ try {
+ (0 until arrays.length).foreach { i =>
+ val arrowSchema = schemas(i)
+ val arrowArray = arrays(i)
- arrayVectors += CometVector.getVector(
- importer.importVector(arrowArray, arrowSchema, dictionaryProvider),
- dictionaryProvider)
+ arrayVectors += CometVector.getVector(
Review Comment:
Thanks!
Done in https://github.com/apache/datafusion-comet/commit/926be0dabc153a0c4fafc2c23e6f8c56dcfbf5e9. The current FieldVector stays rollback-owned through import, wrapping, and publication, with regressions for both failure paths.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: prevent memory leak after failed Arrow vector import [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #5539:
URL: https://github.com/apache/datafusion-comet/pull/5539#discussion_r3886912606
##########
spark/src/main/java/org/apache/arrow/c/ArrowImporter.java:
##########
@@ -50,10 +51,21 @@ Field importField(ArrowSchema schema, CDataDictionaryProvider provider) {
public FieldVector importVector(
ArrowArray array, ArrowSchema schema, CDataDictionaryProvider provider) {
- Field field = importField(schema, provider);
- FieldVector vector = field.createVector(allocator);
- ArrayImporter importer = new ArrayImporter(allocator, vector, provider);
- importer.importArray(array);
- return vector;
+ FieldVector vector = null;
+ try {
+ Field field = importField(schema, provider);
+ vector = field.createVector(allocator);
+ ArrayImporter importer = new ArrayImporter(allocator, vector, provider);
+ importer.importArray(array);
+ return vector;
+ } catch (RuntimeException | Error failure) {
+ if (vector != null) {
+ AutoCloseables.close(failure, vector);
+ }
+ if (!array.isClosed()) {
Review 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.
##########
spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala:
##########
@@ -253,16 +254,33 @@ class NativeUtil {
*/
def importVector(arrays: Array[ArrowArray], schemas: Array[ArrowSchema]): Seq[CometVector] = {
val arrayVectors = mutable.ArrayBuffer.empty[CometVector]
+ var firstUnconsumed = 0
- (0 until arrays.length).foreach { i =>
- val arrowSchema = schemas(i)
- val arrowArray = arrays(i)
-
- arrayVectors += CometVector.getVector(
- importer.importVector(arrowArray, arrowSchema, dictionaryProvider),
- dictionaryProvider)
+ try {
+ (0 until arrays.length).foreach { i =>
+ val arrowSchema = schemas(i)
+ val arrowArray = arrays(i)
+
+ firstUnconsumed = i + 1
Review 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.
##########
spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala:
##########
@@ -99,6 +103,99 @@ class NativeUtilSuite extends CometTestBase {
}
}
+ test("getNextBatch releases imported vectors and Arrow structs when vector import fails") {
+ withIsolatedStructAllocator { (nativeUtil, allocator, _) =>
+ Using.resource(new IntVector("value", allocator)) { vector =>
+ vector.allocateNew(4)
+ vector.setSafe(0, 42)
+ vector.setValueCount(1)
+ val failure = intercept[IllegalStateException] {
+ nativeUtil.getNextBatch(
+ 2,
+ (arrays, schemas) => {
+ Data.exportVector(
+ allocator,
+ vector,
+ null,
+ ArrowArray.wrap(arrays(0)),
+ ArrowSchema.wrap(schemas(0)))
+ vector.close()
+ 1L
+ })
+ }
+ assert(failure.getMessage == "Cannot import released ArrowSchema")
+ assert(failure.getSuppressed.isEmpty)
+ assert(allocator.getAllocatedMemory == 0)
+ }
+ }
+ }
+
+ test("getNextBatch releases an imported vector when its Comet wrapper rejects the type") {
+ withIsolatedStructAllocator { (nativeUtil, allocator, _) =>
+ Using.resource(new UInt4Vector("value", allocator)) { vector =>
+ vector.allocateNew(1)
+ vector.setSafe(0, 42)
+ vector.setValueCount(1)
+ val failure = intercept[UnsupportedOperationException] {
+ nativeUtil.getNextBatch(
+ 2,
+ (arrays, schemas) => {
+ Data.exportVector(
+ allocator,
+ vector,
+ null,
+ ArrowArray.wrap(arrays(0)),
+ ArrowSchema.wrap(schemas(0)))
+ vector.close()
+ 1L
+ })
+ }
+ assert(failure.getSuppressed.isEmpty)
+ assert(allocator.getAllocatedMemory == 0)
+ }
+ }
+ }
+
+ test("importVector releases partially imported vectors when Arrow array import fails") {
+ withIsolatedStructAllocator { (nativeUtil, allocator, _) =>
+ val intType = FieldType.nullable(new ArrowType.Int(32, true))
+ val stringType = FieldType.nullable(new ArrowType.Utf8())
+ val intStruct = StructVector.empty("int_struct", allocator)
+ val firstInt = intStruct.addOrGet("first", intType, classOf[IntVector])
+ val secondInt = intStruct.addOrGet("second", intType, classOf[IntVector])
+ intStruct.allocateNew()
+ intStruct.setIndexDefined(0)
+ firstInt.setSafe(0, 1)
+ secondInt.setSafe(0, 2)
+ intStruct.setValueCount(1)
+
+ val stringStruct = StructVector.empty("string_struct", allocator)
+ stringStruct.addOrGet("first", intType, classOf[IntVector])
+ stringStruct.addOrGet("second", stringType, classOf[VarCharVector])
+ stringStruct.allocateNew()
+ stringStruct.setValueCount(1)
+
+ val (arrays, schemas) = nativeUtil.allocateArrowStructs(2)
+ Data.exportVector(allocator, intStruct, null, arrays(0), schemas(0))
+ Data.exportVector(allocator, stringStruct, null, arrays(1), schemas(1))
+ intStruct.close()
+ stringStruct.close()
Review 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?
##########
spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala:
##########
@@ -253,16 +254,33 @@ class NativeUtil {
*/
def importVector(arrays: Array[ArrowArray], schemas: Array[ArrowSchema]): Seq[CometVector] = {
val arrayVectors = mutable.ArrayBuffer.empty[CometVector]
+ var firstUnconsumed = 0
- (0 until arrays.length).foreach { i =>
- val arrowSchema = schemas(i)
- val arrowArray = arrays(i)
-
- arrayVectors += CometVector.getVector(
- importer.importVector(arrowArray, arrowSchema, dictionaryProvider),
- dictionaryProvider)
+ try {
+ (0 until arrays.length).foreach { i =>
+ val arrowSchema = schemas(i)
+ val arrowArray = arrays(i)
+
+ firstUnconsumed = i + 1
+ val arrowVector = importer.importVector(arrowArray, arrowSchema, dictionaryProvider)
+ var cometVector: CometVector = null
+ try {
+ cometVector = CometVector.getVector(arrowVector, dictionaryProvider)
+ arrayVectors += cometVector
+ } catch {
+ case failure: Throwable =>
+ val rollback = if (cometVector == null) arrowVector else cometVector
Review 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:
```scala
val cometVector =
try CometVector.getVector(arrowVector, dictionaryProvider)
catch {
case failure: Throwable =>
AutoCloseables.close(failure, arrowVector)
throw failure
}
arrayVectors += cometVector
```
In code whose whole job is memory ownership, the fewer branches a reader has to hold in their head the better.
##########
spark/src/main/java/org/apache/arrow/c/ArrowImporter.java:
##########
@@ -50,10 +51,21 @@ Field importField(ArrowSchema schema, CDataDictionaryProvider provider) {
public FieldVector importVector(
ArrowArray array, ArrowSchema schema, CDataDictionaryProvider provider) {
- Field field = importField(schema, provider);
- FieldVector vector = field.createVector(allocator);
- ArrayImporter importer = new ArrayImporter(allocator, vector, provider);
- importer.importArray(array);
- return vector;
+ FieldVector vector = null;
+ try {
+ Field field = importField(schema, provider);
+ vector = field.createVector(allocator);
+ ArrayImporter importer = new ArrayImporter(allocator, vector, provider);
+ importer.importArray(array);
+ return vector;
+ } catch (RuntimeException | Error failure) {
+ if (vector != null) {
+ AutoCloseables.close(failure, vector);
Review 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.
##########
spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala:
##########
@@ -253,16 +254,33 @@ class NativeUtil {
*/
def importVector(arrays: Array[ArrowArray], schemas: Array[ArrowSchema]): Seq[CometVector] = {
Review 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.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: prevent memory leak after failed Arrow vector import [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao merged PR #5539:
URL: https://github.com/apache/datafusion-comet/pull/5539
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org