From e3f12ccbff622bcba901dcf9241ecf51fdfaf6e7 Mon Sep 17 00:00:00 2001 From: jackylee-ch Date: Mon, 7 Sep 2026 11:02:34 +0800 Subject: [PATCH 1/2] fix: release blob file handles on resolve error paths resolveBatch closed each BlobFile with try-with-resources inside the loop, so the count-mismatch throw abandoned the whole returned list, the null check abandoned everything after the current index, and a read() failure abandoned the tail. BlobFile holds a native handle with no cleaner, so an abandoned one lives until the JVM exits, and Spark retries the task in the same JVM by default. Close the handles once each in a finally instead. resolve() had the same shape for a multi-element return, so it gets the same treatment. --- .../spark/utils/BlobReferenceResolver.java | 62 +++++++++++-------- 1 file changed, 36 insertions(+), 26 deletions(-) diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/BlobReferenceResolver.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/BlobReferenceResolver.java index 7f90e4fac..fd9503313 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/BlobReferenceResolver.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/BlobReferenceResolver.java @@ -80,11 +80,15 @@ public byte[] resolve(BlobReference ref) throws IOException { List rowAddresses = new ArrayList<>(1); rowAddresses.add(ref.getRowAddress()); List blobs = dataset.takeBlobs(rowAddresses, ref.getColumnName()); - if (blobs.isEmpty()) { - return new byte[0]; - } - try (BlobFile blob = blobs.get(0)) { - return blob.read(); + try { + if (blobs.isEmpty()) { + return new byte[0]; + } + return blobs.get(0).read(); + } finally { + for (BlobFile blob : blobs) { + CloseableUtil.closeQuietly(blob); + } } } @@ -137,31 +141,37 @@ public Map resolveBatch(List indices, List addresses = group.distinctAddresses; // requested order List blobs = dataset.takeBlobs(addresses, group.columnName); - // takeBlobs must return exactly one BlobFile per requested address, in order. A mismatch - // means the selection hit deleted/null-descriptor rows, in which case positional mapping - // would skew and silently write the wrong bytes into the target table — fail loudly instead. - if (blobs.size() != addresses.size()) { - throw new IOException( - String.format( - "takeBlobs returned %d blobs for %d requested addresses (column=%s, dataset=%s); " - + "cannot map results to rows", - blobs.size(), addresses.size(), group.columnName, group.datasetUri)); - } - - for (int i = 0; i < addresses.size(); i++) { - BlobFile blob = blobs.get(i); - if (blob == null) { + // Every handle takeBlobs returned is released in the finally below, including on the two + // throws: BlobFile wraps a native handle with no cleaner, so an abandoned one is only + // reclaimed when the JVM exits, and Spark retries the task in the same JVM. + try { + // takeBlobs must return exactly one BlobFile per requested address, in order. A mismatch + // means the selection hit deleted/null-descriptor rows, in which case positional mapping + // would skew and silently write the wrong bytes into the target table — fail loudly. + if (blobs.size() != addresses.size()) { throw new IOException( String.format( - "takeBlobs returned a null blob for address %d (column=%s, dataset=%s)", - addresses.get(i), group.columnName, group.datasetUri)); + "takeBlobs returned %d blobs for %d requested addresses (column=%s, dataset=%s); " + + "cannot map results to rows", + blobs.size(), addresses.size(), group.columnName, group.datasetUri)); } - byte[] data; - try (BlobFile b = blob) { - data = b.read(); + + for (int i = 0; i < addresses.size(); i++) { + BlobFile blob = blobs.get(i); + if (blob == null) { + throw new IOException( + String.format( + "takeBlobs returned a null blob for address %d (column=%s, dataset=%s)", + addresses.get(i), group.columnName, group.datasetUri)); + } + byte[] data = blob.read(); + for (int vectorIndex : group.indicesByAddress.get(addresses.get(i))) { + resolved.put(vectorIndex, data); + } } - for (int vectorIndex : group.indicesByAddress.get(addresses.get(i))) { - resolved.put(vectorIndex, data); + } finally { + for (BlobFile blob : blobs) { + CloseableUtil.closeQuietly(blob); } } } From cb514bac1df73e3a8d7a87c03702a0d5aa6a3b7b Mon Sep 17 00:00:00 2001 From: jackylee-ch Date: Tue, 15 Sep 2026 15:31:38 +0800 Subject: [PATCH 2/2] test: cover blob resolver cleanup on failures Exercise successful reads, count mismatches, null blobs and read failures with real Lance blob v2 datasets. Capture the returned handles through a shared package-private takeBlobs helper and verify their native owners are released before the resolver closes. The old per-item cleanup fails all three batch error regressions. Blob/Join suites pass on Spark 3.5/Scala 2.12 and Spark 4.1/Scala 2.13 (149 tests each, no skips); Checkstyle and Spotless pass. --- .../spark/utils/BlobReferenceResolver.java | 10 +- .../utils/BlobReferenceResolverTest.scala | 191 ++++++++++++++++++ 2 files changed, 197 insertions(+), 4 deletions(-) create mode 100644 lance-spark-base_2.12/src/test/scala/org/lance/spark/utils/BlobReferenceResolverTest.scala diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/BlobReferenceResolver.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/BlobReferenceResolver.java index fd9503313..71f744a56 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/BlobReferenceResolver.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/BlobReferenceResolver.java @@ -76,10 +76,9 @@ public BlobReferenceResolver(Map sourceContexts) { * @throws IOException if reading the blob fails */ public byte[] resolve(BlobReference ref) throws IOException { - Dataset dataset = getOrOpenDataset(ref.getDatasetUri()); List rowAddresses = new ArrayList<>(1); rowAddresses.add(ref.getRowAddress()); - List blobs = dataset.takeBlobs(rowAddresses, ref.getColumnName()); + List blobs = takeBlobs(ref.getDatasetUri(), rowAddresses, ref.getColumnName()); try { if (blobs.isEmpty()) { return new byte[0]; @@ -137,9 +136,8 @@ public Map resolveBatch(List indices, List addresses = group.distinctAddresses; // requested order - List blobs = dataset.takeBlobs(addresses, group.columnName); + List blobs = takeBlobs(group.datasetUri, addresses, group.columnName); // Every handle takeBlobs returned is released in the finally below, including on the two // throws: BlobFile wraps a native handle with no cleaner, so an abandoned one is only @@ -178,6 +176,10 @@ public Map resolveBatch(List indices, List takeBlobs(String datasetUri, List addresses, String columnName) { + return getOrOpenDataset(datasetUri).takeBlobs(addresses, columnName); + } + private Dataset getOrOpenDataset(String datasetUri) { return datasetCache.computeIfAbsent(datasetUri, this::openDataset); } diff --git a/lance-spark-base_2.12/src/test/scala/org/lance/spark/utils/BlobReferenceResolverTest.scala b/lance-spark-base_2.12/src/test/scala/org/lance/spark/utils/BlobReferenceResolverTest.scala new file mode 100644 index 000000000..b5f212f22 --- /dev/null +++ b/lance-spark-base_2.12/src/test/scala/org/lance/spark/utils/BlobReferenceResolverTest.scala @@ -0,0 +1,191 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.utils + +import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.{LargeVarBinaryVector, VectorSchemaRoot} +import org.apache.arrow.vector.complex.StructVector +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema} +import org.junit.jupiter.api.Assertions._ +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.io.TempDir +import org.lance.{BlobFile, Dataset} +import org.lance.spark.write.SingleBatchArrowReader + +import java.io.IOException +import java.nio.file.{Files, Path, Paths} +import java.util.{Arrays, Collections, List => JList} + +import scala.collection.JavaConverters._ + +class BlobReferenceResolverTest { + @TempDir var tempDir: Path = _ + + private val data = Array[Byte](1, 2, 3) + + // JNI take_rust_field clears this field when close releases the native owner. + private def handle(blob: BlobFile): Long = { + val field = classOf[BlobFile].getDeclaredField("nativeBlobHandle") + field.setAccessible(true) + field.getLong(blob) + } + + private class RecordingResolver(afterTake: JList[BlobFile] => Unit) + extends BlobReferenceResolver { + var acquired: Seq[BlobFile] = Seq.empty + + override def takeBlobs( + uri: String, + addresses: JList[java.lang.Long], + column: String): JList[BlobFile] = { + val blobs = super.takeBlobs(uri, addresses, column) + acquired = blobs.asScala.filter(_ != null).toVector + assertTrue(acquired.nonEmpty) + acquired.foreach(blob => assertNotEquals(0L, handle(blob))) + afterTake(blobs) + blobs + } + + def assertReleased(): Unit = + acquired.foreach(blob => assertEquals(0L, handle(blob), "native blob handle leaked")) + + override def close(): Unit = { + acquired.filter(blob => handle(blob) != 0L).foreach(_.close()) + super.close() + } + } + + private def withSource(nullable: Boolean = false)( + body: (String, JList[BlobReference]) => Unit): Unit = { + val uri = tempDir.resolve("source.lance").toString + val field = new Field( + "data", + new FieldType( + true, + ArrowType.Struct.INSTANCE, + null, + Collections.singletonMap( + BlobUtils.ARROW_EXTENSION_NAME_KEY, + BlobUtils.ARROW_EXTENSION_BLOB_V2)), + Arrays.asList( + Field.nullable("data", ArrowType.LargeBinary.INSTANCE), + Field.nullable("uri", ArrowType.Utf8.INSTANCE))) + val allocator = new RootAllocator() + val root = VectorSchemaRoot.create(new Schema(Collections.singletonList(field)), allocator) + val reader = new SingleBatchArrowReader(allocator, root) + try { + root.allocateNew() + val vector = root.getVector(0).asInstanceOf[StructVector] + val bytes = vector.getChild("data").asInstanceOf[LargeVarBinaryVector] + (0 until 3).foreach { i => + if (nullable && i == 1) { + vector.setNull(i) + bytes.setNull(i) + } else { + vector.setIndexDefined(i) + bytes.setSafe(i, data) + } + } + root.setRowCount(3) + val dataset = Dataset.write() + .allocator(allocator) + .reader(reader) + .uri(uri) + .dataStorageVersion("2.2") + .execute() + try { + assertEquals(1, dataset.getFragments.size()) + val fragmentId = dataset.getFragments.get(0).getId.toLong + val refs = (0 until 3) + .map(i => new BlobReference(uri, "data", (fragmentId << 32) | i.toLong)) + .asJava + body(uri, refs) + } finally dataset.close() + } finally { + reader.close() + root.close() + allocator.close() + } + } + + private def indices(refs: JList[BlobReference]): JList[Integer] = + (0 until refs.size()).map(Integer.valueOf).asJava + + private def deleteDataFiles(uri: String): Unit = { + val paths = Files.walk(Paths.get(uri).resolve("data")) + try { + val files = paths.iterator().asScala.filter(Files.isRegularFile(_)).toVector + assertFalse(files.isEmpty) + files.foreach(Files.delete) + } finally paths.close() + } + + @Test + def batchSuccessReleasesAllHandles(): Unit = withSource() { (_, refs) => + val resolver = new RecordingResolver(_ => ()) + try { + val result = resolver.resolveBatch(indices(refs), refs) + assertEquals(3, result.size()) + result.values().asScala.foreach(bytes => assertArrayEquals(data, bytes)) + resolver.assertReleased() + } finally resolver.close() + } + + @Test + def nullBlobReleasesVisitedAndUnvisitedHandles(): Unit = withSource(nullable = true) { + (_, refs) => + val resolver = new RecordingResolver(blobs => assertNull(blobs.get(1))) + try { + val error = assertThrows( + classOf[IOException], + () => resolver.resolveBatch(indices(refs), refs)) + assertTrue(error.getMessage.contains("takeBlobs returned a null blob")) + resolver.assertReleased() + } finally resolver.close() + } + + @Test + def countMismatchReleasesReturnedHandles(): Unit = withSource() { (_, refs) => + val resolver = new RecordingResolver(blobs => blobs.remove(blobs.size() - 1).close()) + try { + val error = assertThrows( + classOf[IOException], + () => resolver.resolveBatch(indices(refs), refs)) + assertTrue(error.getMessage.contains("takeBlobs returned 2 blobs for 3 requested addresses")) + resolver.assertReleased() + } finally resolver.close() + } + + @Test + def batchReadFailureReleasesAllHandles(): Unit = withSource() { (uri, refs) => + val resolver = new RecordingResolver(_ => deleteDataFiles(uri)) + try { + val error = assertThrows( + classOf[IOException], + () => resolver.resolveBatch(indices(refs), refs)) + assertTrue(error.getMessage.contains("Not found"), error.getMessage) + resolver.assertReleased() + } finally resolver.close() + } + + @Test + def singleReadFailureReleasesHandle(): Unit = withSource() { (uri, refs) => + val resolver = new RecordingResolver(_ => deleteDataFiles(uri)) + try { + val error = assertThrows(classOf[IOException], () => resolver.resolve(refs.get(0))) + assertTrue(error.getMessage.contains("Not found"), error.getMessage) + resolver.assertReleased() + } finally resolver.close() + } +}