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: [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]
