diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala index 266772fdb..5a229d1d9 100755 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala @@ -16,7 +16,6 @@ package org.apache.spark.sql.execution.datasources.v2 import com.fasterxml.jackson.databind.ObjectMapper import com.fasterxml.jackson.databind.node.ObjectNode import org.apache.arrow.c.{ArrowArrayStream, Data} -import org.apache.arrow.vector.VectorSchemaRoot import org.apache.arrow.vector.ipc.ArrowReader import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.InternalRow @@ -24,7 +23,6 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, GenericInternalRow} import org.apache.spark.sql.catalyst.plans.logical.{AddIndexOutputType, LanceNamedArgument} import org.apache.spark.sql.connector.catalog.{Identifier, TableCatalog} import org.apache.spark.sql.types.StructType -import org.apache.spark.sql.util.LanceArrowUtils import org.apache.spark.sql.util.LanceSerializeUtil.{decode, encode} import org.apache.spark.unsafe.types.UTF8String import org.lance.{CommitBuilder, Dataset, Transaction} @@ -32,9 +30,8 @@ import org.lance.index.{Index, IndexOptions, IndexParams, IndexType} import org.lance.index.scalar.{BTreeIndexParams, ScalarIndexParams} import org.lance.operation.{CreateIndex => AddIndexOperation} import org.lance.spark.{BaseLanceNamespaceSparkCatalog, LanceDataset, LanceRuntime, LanceSparkReadOptions} -import org.lance.spark.arrow.LanceArrowWriter import org.lance.spark.utils.{CloseableUtil, FieldPathUtils, Utils} -import org.lance.spark.write.SingleBatchArrowReader +import org.lance.spark.write.{ArrowBatchChunker, ChunkedArrowReader} import java.time.Instant import java.util.{Collections, Optional, UUID} @@ -761,31 +758,31 @@ case class RangeBTreeIndexBuilder( val fragmentIdOrdinal = 2 val allocator = LanceRuntime.allocator() - val data = - VectorSchemaRoot.create(LanceArrowUtils.toArrowSchema(streamSchema, "UTC", false), allocator) - val writer = LanceArrowWriter.create(data, streamSchema) - + // Accumulate (value, row id) into narrow Arrow batches, each kept under Arrow's 2 GiB + // 32-bit-offset cap. A fragment group whose value column exceeds 2 GiB would overflow a + // single VarChar/VarBinary vector; chunking into multiple batches avoids that, and lance-core + // re-zones the multi-batch stream into an index identical to a single-batch handoff. + val chunker = new ArrowBatchChunker(streamSchema, allocator) val fragmentIds = scala.collection.mutable.LinkedHashSet[java.lang.Integer]() - // Write the indexed value and row id of each row to the Arrow stream. - try { - while (rowsIter.hasNext) { - val row = rowsIter.next() - writer.field(0).write(row, 0) - writer.field(1).write(row, 1) - fragmentIds += java.lang.Integer.valueOf(row.getInt(fragmentIdOrdinal)) + // Write the indexed value and row id of each row; collect fragment ids for the segment's + // coverage in the same pass so they are known before the stream is handed to lance-core. + val batches = + try { + while (rowsIter.hasNext) { + val row = rowsIter.next() + chunker.write(row) + fragmentIds += java.lang.Integer.valueOf(row.getInt(fragmentIdOrdinal)) + } + chunker.finish() + } catch { + case e: Throwable => + chunker.close() + throw e } - writer.finish() - } catch { - case e: Throwable => - CloseableUtil.closeQuietly(data) - throw e - } - // No rows are written - if (data.getRowCount == 0) { - data.close() + if (batches.isEmpty) { return Iterator(encode(None: Option[Index])) } @@ -795,7 +792,7 @@ case class RangeBTreeIndexBuilder( try { stream = ArrowArrayStream.allocateNew(allocator) - reader = new SingleBatchArrowReader(allocator, data) + reader = new ChunkedArrowReader(allocator, chunker.arrowSchema, batches) dataset = Utils.openDatasetBuilder( decode[LanceSparkReadOptions](encodedReadOptions)) @@ -832,7 +829,7 @@ case class RangeBTreeIndexBuilder( } finally { CloseableUtil.closeQuietly(stream) CloseableUtil.closeQuietly(reader) - CloseableUtil.closeQuietly(data) + batches.foreach(b => CloseableUtil.closeQuietly(b)) CloseableUtil.closeQuietly(dataset) } } diff --git a/lance-spark-base_2.12/src/main/scala/org/lance/spark/write/ArrowBatchChunker.scala b/lance-spark-base_2.12/src/main/scala/org/lance/spark/write/ArrowBatchChunker.scala new file mode 100644 index 000000000..13c0733c9 --- /dev/null +++ b/lance-spark-base_2.12/src/main/scala/org/lance/spark/write/ArrowBatchChunker.scala @@ -0,0 +1,230 @@ +/* + * 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.write + +import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.arrow.vector.types.FloatingPointPrecision +import org.apache.arrow.vector.types.pojo.{ArrowType, Schema} +import org.apache.arrow.vector.util.OversizedAllocationException +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.LanceArrowUtils +import org.lance.spark.arrow.LanceArrowWriter + +import scala.collection.JavaConverters._ +import scala.collection.mutable.ArrayBuffer + +/** + * Accumulates [[InternalRow]]s into a sequence of narrow Arrow [[VectorSchemaRoot]] batches, + * rolling to a new batch before any variable-width (Utf8/Binary) column would push its data + * buffer past Arrow's 32-bit-offset 2 GiB cap. The whole partition would otherwise be buffered + * into a single `VectorSchemaRoot` and overflow that cap the moment a value column exceeds 2 GiB; + * splitting the same rows across capped batches keeps peak memory unchanged while staying within + * the cap. Index-type agnostic — any build that pre-processes rows into an Arrow stream can use + * it. + * + * Drive it with one [[write]] per row, then [[finish]] to get the batches; on the error path call + * [[close]]. The caller owns and must close the returned batches. + * + * @param sparkSchema schema of the rows written (Arrow batches use the narrow encoding of it) + * @param allocator allocator for the batch vectors + * @param maxBatchBytes soft per-column data-buffer budget; default just under the 2 GiB hard cap + */ +class ArrowBatchChunker( + sparkSchema: StructType, + allocator: BufferAllocator, + maxBatchBytes: Long = ArrowBatchChunker.DefaultMaxBatchBytes) { + + require(maxBatchBytes > 0, s"maxBatchBytes must be positive, got $maxBatchBytes") + + /** Narrow Arrow schema every emitted batch uses. */ + val arrowSchema: Schema = + LanceArrowUtils.toArrowSchema(sparkSchema, "UTC", errorOnDuplicatedFieldNames = false) + + // Ordinals of narrow var-width columns (Utf8/Binary) subject to the 32-bit data-buffer cap, + // each tagged Utf8 (string) vs Binary so the per-row byte size is read from the right accessor. + // 64-bit Large* columns (blob / large_var_char) and fixed-width columns are excluded: Large* + // offsets don't hit the 32-bit cap, and fixed-width buffers are bounded by maxRows below. + private val (cappedOrdinals, cappedIsString): (Array[Int], Array[Boolean]) = { + val ords = ArrayBuffer[Int]() + val isStr = ArrayBuffer[Boolean]() + arrowSchema.getFields.asScala.zipWithIndex.foreach { + case (field, i) => + field.getType match { + case _: ArrowType.Utf8 => ords += i; isStr += true + case _: ArrowType.Binary => ords += i; isStr += false + case _ => + } + } + (ords.toArray, isStr.toArray) + } + + // Cap each batch's row count so fixed-width buffers and var-width offset/validity buffers also + // stay under the 32-bit cap, not just the var-width data buffers. perRowFixedCost is the widest + // per-row fixed contribution across columns (e.g. an 8-byte row id, a 4-byte narrow offset). + private val maxRows: Int = { + val perRowFixedCost = math.max( + 1L, + arrowSchema.getFields.asScala + .map(f => ArrowBatchChunker.fixedWidthBytes(f.getType).toLong) + .max) + math.max(1, (maxBatchBytes / perRowFixedCost).min(Int.MaxValue.toLong).toInt) + } + + private val batches = ArrayBuffer[VectorSchemaRoot]() + private var current: VectorSchemaRoot = _ + private var writer: LanceArrowWriter = _ + private var rowCount: Int = 0 + private val running: Array[Long] = new Array[Long](cappedOrdinals.length) + private var finished: Boolean = false + + private def startBatch(): Unit = { + current = VectorSchemaRoot.create(arrowSchema, allocator) + writer = LanceArrowWriter.create(current, sparkSchema) + rowCount = 0 + java.util.Arrays.fill(running, 0L) + } + + private def sealBatch(): Unit = { + writer.finish() + batches += current + current = null + writer = null + } + + // Var-width byte size this row contributes to each capped column (0 for nulls). + private def rowVarBytes(row: InternalRow): Array[Long] = { + val out = new Array[Long](cappedOrdinals.length) + var k = 0 + while (k < cappedOrdinals.length) { + val ord = cappedOrdinals(k) + out(k) = + if (row.isNullAt(ord)) 0L + else if (cappedIsString(k)) row.getUTF8String(ord).numBytes().toLong + else row.getBinary(ord).length.toLong + k += 1 + } + out + } + + /** + * Append one row, first rolling to a new batch if it would push a capped column's data buffer + * past `maxBatchBytes` or the batch past `maxRows`. A single value larger than the 32-bit cap + * can never fit a narrow vector, so it fails fast rather than overflowing to a negative offset. + */ + def write(row: InternalRow): Unit = { + require(!finished, "write called after finish") + if (current == null) { + startBatch() + } + + val rowBytes = rowVarBytes(row) + var k = 0 + while (k < rowBytes.length) { + if (rowBytes(k) > ArrowBatchChunker.HardCapBytes) { + throw new OversizedAllocationException( + s"row holds ${rowBytes(k)} bytes in column " + + s"'${arrowSchema.getFields.get(cappedOrdinals(k)).getName}', over the " + + s"${ArrowBatchChunker.HardCapBytes}-byte limit of a 32-bit-offset Arrow vector; " + + "a single index value cannot exceed 2 GiB") + } + k += 1 + } + + var wouldExceedBytes = false + var i = 0 + while (i < running.length && !wouldExceedBytes) { + if (running(i) + rowBytes(i) > maxBatchBytes) { + wouldExceedBytes = true + } + i += 1 + } + + if (rowCount > 0 && (wouldExceedBytes || rowCount >= maxRows)) { + sealBatch() + startBatch() + } + + writer.write(row) + rowCount += 1 + var j = 0 + while (j < running.length) { + running(j) += rowBytes(j) + j += 1 + } + } + + /** + * Seal the final batch and return all batches, each within budget. Empty if no rows were + * written. The caller owns the returned roots and must close them. + */ + def finish(): Array[VectorSchemaRoot] = { + require(!finished, "finish called twice") + finished = true + if (current != null) { + if (rowCount > 0) { + sealBatch() + } else { + current.close() + current = null + writer = null + } + } + batches.toArray + } + + /** Close the in-progress batch and every sealed batch. Use on the error path before rethrow. */ + def close(): Unit = { + if (current != null) { + current.close() + current = null + writer = null + } + batches.foreach(_.close()) + batches.clear() + } +} + +object ArrowBatchChunker { + + /** A 32-bit-offset Arrow vector addresses at most `Integer.MAX_VALUE` data bytes. */ + val HardCapBytes: Long = Integer.MAX_VALUE.toLong + + /** + * Default soft budget per batch: just under the 2 GiB hard cap, leaving headroom so one more + * typical row won't cross it before the chunker rolls. + */ + val DefaultMaxBatchBytes: Long = HardCapBytes - (64L << 20) + + // Per-row byte cost of a column's fixed-width buffers (data for fixed types, offset slot for + // var types) — used only to bound rows per batch so non-data buffers stay under the 32-bit cap. + private def fixedWidthBytes(t: ArrowType): Int = t match { + case _: ArrowType.Utf8 | _: ArrowType.Binary => 4 + case _: ArrowType.LargeUtf8 | _: ArrowType.LargeBinary => 8 + case i: ArrowType.Int => math.max(1, i.getBitWidth / 8) + case fp: ArrowType.FloatingPoint => + fp.getPrecision match { + case FloatingPointPrecision.HALF => 2 + case FloatingPointPrecision.SINGLE => 4 + case FloatingPointPrecision.DOUBLE => 8 + } + case _: ArrowType.Bool => 1 + case d: ArrowType.Decimal => math.max(1, d.getBitWidth / 8) + case _: ArrowType.Date => 8 + case _: ArrowType.Timestamp => 8 + case fsb: ArrowType.FixedSizeBinary => math.max(1, fsb.getByteWidth) + case _ => 8 + } +} diff --git a/lance-spark-base_2.12/src/main/scala/org/lance/spark/write/ChunkedArrowReader.scala b/lance-spark-base_2.12/src/main/scala/org/lance/spark/write/ChunkedArrowReader.scala new file mode 100644 index 000000000..39729f91e --- /dev/null +++ b/lance-spark-base_2.12/src/main/scala/org/lance/spark/write/ChunkedArrowReader.scala @@ -0,0 +1,62 @@ +/* + * 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.write + +import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.vector.{VectorLoader, VectorSchemaRoot, VectorUnloader} +import org.apache.arrow.vector.ipc.ArrowReader +import org.apache.arrow.vector.types.pojo.Schema + +/** + * An [[ArrowReader]] that emits a fixed sequence of pre-built [[VectorSchemaRoot]] batches, one + * per [[loadNextBatch]] call. Generalizes [[SingleBatchArrowReader]] from one batch to N so an + * index build can hand a partition to lance-core split across several sub-2-GiB Arrow batches + * rather than one, sidestepping Arrow's 32-bit-offset (2 GiB) cap on variable-width columns. + * lance-core consumes the whole stream per index, so the built index matches a single-batch + * handoff. + * + * Index-type agnostic: any build that pre-processes rows into Arrow batches can use it. The + * batches and `schema` are owned by the caller (mirroring [[SingleBatchArrowReader]]); this + * reader closes neither. + */ +class ChunkedArrowReader( + allocator: BufferAllocator, + schema: Schema, + batches: Array[VectorSchemaRoot]) + extends ArrowReader(allocator) { + + private var nextBatch: Int = 0 + + override protected def readSchema(): Schema = schema + + override def loadNextBatch(): Boolean = { + if (nextBatch >= batches.length) { + return false + } + prepareLoadNextBatch() + val source = batches(nextBatch) + val recordBatch = new VectorUnloader(source).getRecordBatch() + try { + new VectorLoader(getVectorSchemaRoot()).load(recordBatch) + } finally { + recordBatch.close() + } + nextBatch += 1 + true + } + + override def bytesRead(): Long = 0L + + override protected def closeReadSource(): Unit = {} +} diff --git a/lance-spark-base_2.12/src/test/scala/org/lance/spark/write/ArrowBatchChunkerSuite.scala b/lance-spark-base_2.12/src/test/scala/org/lance/spark/write/ArrowBatchChunkerSuite.scala new file mode 100644 index 000000000..88f2307db --- /dev/null +++ b/lance-spark-base_2.12/src/test/scala/org/lance/spark/write/ArrowBatchChunkerSuite.scala @@ -0,0 +1,168 @@ +/* + * 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.write + +import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.{BigIntVector, VarBinaryVector, VectorSchemaRoot} +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow +import org.apache.spark.sql.types.{BinaryType, LongType, StructType} +import org.junit.jupiter.api.Assertions._ +import org.junit.jupiter.api.Test + +import scala.collection.mutable.ArrayBuffer + +/** + * Unit tests for [[ArrowBatchChunker]]: the byte-budget and row-count boundaries that split a + * row stream into narrow Arrow batches, each under Arrow's 32-bit-offset (2 GiB) cap. Pure JVM + * (no SparkSession, no Lance native library). + */ +class ArrowBatchChunkerSuite { + + private val binarySchema = new StructType() + .add("value", BinaryType, nullable = false) + .add("id", LongType, nullable = false) + + private val longSchema = new StructType() + .add("value", LongType, nullable = false) + .add("id", LongType, nullable = false) + + private def binaryRow(id: Long, size: Int): GenericInternalRow = + new GenericInternalRow(Array[Any](Array.fill[Byte](size)(id.toByte), id)) + + private def longRow(id: Long): GenericInternalRow = + new GenericInternalRow(Array[Any](id, id)) + + // Per-batch row counts and the flattened ids across all batches, in emission order. + private def drain(batches: Array[VectorSchemaRoot]): (List[Int], List[Long]) = { + val counts = batches.map(_.getRowCount).toList + val ids = ArrayBuffer[Long]() + batches.foreach { root => + val idVec = root.getVector("id").asInstanceOf[BigIntVector] + var r = 0 + while (r < root.getRowCount) { + ids += idVec.get(r) + r += 1 + } + } + (counts, ids.toList) + } + + @Test + def splitsBinaryAtByteBudget(): Unit = { + val allocator = new RootAllocator(Long.MaxValue) + val chunker = new ArrowBatchChunker(binarySchema, allocator, maxBatchBytes = 250L) + try { + (0 until 6).foreach(i => chunker.write(binaryRow(i.toLong, 100))) + val batches = chunker.finish() + try { + val (counts, ids) = drain(batches) + assertEquals(List(2, 2, 2), counts, "250B budget holds two 100B rows per batch") + assertEquals((0L until 6L).toList, ids, "rows must survive in order") + assertTrue( + batches.forall(_.getVector("value").isInstanceOf[VarBinaryVector]), + "value column stays narrow Binary (no widening, no narrowing)") + } finally { + batches.foreach(_.close()) + } + } finally { + allocator.close() + } + } + + @Test + def singleBatchWhenUnderBudget(): Unit = { + val allocator = new RootAllocator(Long.MaxValue) + val chunker = new ArrowBatchChunker(binarySchema, allocator, maxBatchBytes = 10000L) + try { + (0 until 6).foreach(i => chunker.write(binaryRow(i.toLong, 100))) + val batches = chunker.finish() + try { + val (counts, ids) = drain(batches) + assertEquals(List(6), counts) + assertEquals((0L until 6L).toList, ids) + } finally { + batches.foreach(_.close()) + } + } finally { + allocator.close() + } + } + + @Test + def isolatesOverBudgetRow(): Unit = { + val allocator = new RootAllocator(Long.MaxValue) + val chunker = new ArrowBatchChunker(binarySchema, allocator, maxBatchBytes = 250L) + try { + chunker.write(binaryRow(0L, 300)) // over budget -> its own batch + chunker.write(binaryRow(1L, 100)) + chunker.write(binaryRow(2L, 100)) + val batches = chunker.finish() + try { + val (counts, ids) = drain(batches) + assertEquals(List(1, 2), counts) + assertEquals(List(0L, 1L, 2L), ids) + } finally { + batches.foreach(_.close()) + } + } finally { + allocator.close() + } + } + + @Test + def chunksFixedWidthByRowCap(): Unit = { + val allocator = new RootAllocator(Long.MaxValue) + // No var-width column, so only the row-count ceiling can roll. perRowFixedCost = 8 (BigInt), + // so maxBatchBytes=16 yields maxRows=2. + val chunker = new ArrowBatchChunker(longSchema, allocator, maxBatchBytes = 16L) + try { + (0 until 5).foreach(i => chunker.write(longRow(i.toLong))) + val batches = chunker.finish() + try { + val (counts, ids) = drain(batches) + assertEquals(List(2, 2, 1), counts) + assertEquals((0L until 5L).toList, ids) + } finally { + batches.foreach(_.close()) + } + } finally { + allocator.close() + } + } + + @Test + def emptyWhenNoRows(): Unit = { + val allocator = new RootAllocator(Long.MaxValue) + val chunker = new ArrowBatchChunker(binarySchema, allocator) + try { + assertEquals(0, chunker.finish().length) + } finally { + allocator.close() + } + } + + @Test + def writeAfterFinishThrows(): Unit = { + val allocator = new RootAllocator(Long.MaxValue) + val chunker = new ArrowBatchChunker(binarySchema, allocator) + try { + chunker.finish() + assertThrows( + classOf[IllegalArgumentException], + () => chunker.write(binaryRow(0L, 1))) + } finally { + allocator.close() + } + } +} diff --git a/lance-spark-base_2.12/src/test/scala/org/lance/spark/write/ChunkedArrowReaderSuite.scala b/lance-spark-base_2.12/src/test/scala/org/lance/spark/write/ChunkedArrowReaderSuite.scala new file mode 100644 index 000000000..fad983716 --- /dev/null +++ b/lance-spark-base_2.12/src/test/scala/org/lance/spark/write/ChunkedArrowReaderSuite.scala @@ -0,0 +1,96 @@ +/* + * 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.write + +import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.{BigIntVector, VarBinaryVector, VectorSchemaRoot} +import org.apache.spark.sql.types.{BinaryType, LongType, StructType} +import org.apache.spark.sql.util.LanceArrowUtils +import org.junit.jupiter.api.Assertions._ +import org.junit.jupiter.api.Test + +import scala.collection.mutable.ArrayBuffer + +/** + * Unit tests for [[ChunkedArrowReader]]: emits each pre-built batch once, in order, then reports + * exhaustion. Pure JVM (no SparkSession, no Lance native library, no Arrow C interface). + */ +class ChunkedArrowReaderSuite { + + private val sparkSchema = new StructType() + .add("value", BinaryType, nullable = false) + .add("id", LongType, nullable = false) + + private def newBatch(allocator: RootAllocator, ids: Seq[Long]): VectorSchemaRoot = { + val root = VectorSchemaRoot.create( + LanceArrowUtils.toArrowSchema(sparkSchema, "UTC", false), + allocator) + val value = root.getVector("value").asInstanceOf[VarBinaryVector] + val id = root.getVector("id").asInstanceOf[BigIntVector] + value.allocateNew() + id.allocateNew() + ids.zipWithIndex.foreach { + case (v, i) => + value.setSafe(i, Array.fill[Byte](1)(v.toByte)) + id.setSafe(i, v) + } + root.setRowCount(ids.length) + root + } + + @Test + def emitsEachBatchInOrder(): Unit = { + val allocator = new RootAllocator(Long.MaxValue) + val schema = LanceArrowUtils.toArrowSchema(sparkSchema, "UTC", false) + val batches = Array( + newBatch(allocator, Seq(0L, 1L)), + newBatch(allocator, Seq(2L, 3L)), + newBatch(allocator, Seq(4L))) + val reader = new ChunkedArrowReader(allocator, schema, batches) + try { + val counts = ArrayBuffer[Int]() + val ids = ArrayBuffer[Long]() + while (reader.loadNextBatch()) { + val root = reader.getVectorSchemaRoot + counts += root.getRowCount + val idVec = root.getVector("id").asInstanceOf[BigIntVector] + var r = 0 + while (r < root.getRowCount) { + ids += idVec.get(r) + r += 1 + } + } + assertEquals(List(2, 2, 1), counts.toList) + assertEquals(List(0L, 1L, 2L, 3L, 4L), ids.toList) + assertFalse(reader.loadNextBatch(), "an exhausted reader stays exhausted") + } finally { + reader.close() + batches.foreach(_.close()) + allocator.close() + } + } + + @Test + def emptyReaderLoadsNoBatches(): Unit = { + val allocator = new RootAllocator(Long.MaxValue) + val schema = LanceArrowUtils.toArrowSchema(sparkSchema, "UTC", false) + val reader = new ChunkedArrowReader(allocator, schema, Array.empty[VectorSchemaRoot]) + try { + assertFalse(reader.loadNextBatch()) + } finally { + reader.close() + allocator.close() + } + } +}