diff --git a/java/lance-jni/src/blocking_scanner.rs b/java/lance-jni/src/blocking_scanner.rs index 335cb2a4fa3..c94b51ebe7e 100644 --- a/java/lance-jni/src/blocking_scanner.rs +++ b/java/lance-jni/src/blocking_scanner.rs @@ -589,6 +589,37 @@ pub extern "system" fn Java_org_lance_ipc_LanceScanner_openStream( } fn inner_open_stream(env: &mut JNIEnv, j_scanner: JObject, stream_addr: jlong) -> Result<()> { + if stream_addr == 0 { + return Err(Error::input_error( + "ArrowArrayStream address must not be null".to_string(), + )); + } + + // Reject a stream that already holds a producer. We write the C struct in place below with + // `ptr::write_unaligned`, which does not run any destructor on the previous contents. If the + // caller passed a stream whose `release` callback is already set (e.g. it was populated by an + // earlier export and not yet released), overwriting it would drop that callback and leak the + // first producer's resources. A freshly-allocated `ArrowArrayStream` has a null `release`, per + // the Arrow C Data Interface, so requiring `release == None` is the contract for "empty". + // + // The struct is allocated by Arrow Java inside an ArrowBuf and is not guaranteed to be aligned + // (hence `write_unaligned` below), so we must not form a reference to it. We read only the + // `release` field through an unaligned read: `addr_of!` computes the field address without + // creating an intermediate (mis)aligned reference, and the field is an `Option` which is + // `Copy` with no destructor, so reading a copy of it leaves the caller's stream untouched. + let release_is_set = unsafe { + let stream_ptr = stream_addr as *const FFI_ArrowArrayStream; + let release = std::ptr::read_unaligned(std::ptr::addr_of!((*stream_ptr).release)); + release.is_some() + }; + if release_is_set { + return Err(Error::input_error( + "ArrowArrayStream is already populated; exporting into it would leak the existing \ + producer. Pass a freshly-allocated, empty stream." + .to_string(), + )); + } + let record_batch_stream = { let scanner_guard = unsafe { env.get_rust_field::<_, _, BlockingScanner>(j_scanner, NATIVE_SCANNER) }?; diff --git a/java/src/main/java/org/lance/ipc/LanceScanner.java b/java/src/main/java/org/lance/ipc/LanceScanner.java index 3a413e0ccfd..fbc271e85f4 100644 --- a/java/src/main/java/org/lance/ipc/LanceScanner.java +++ b/java/src/main/java/org/lance/ipc/LanceScanner.java @@ -146,6 +146,66 @@ public ArrowReader scanBatches() { } } + /** + * Export this scan's results into a caller-owned Arrow C stream identified by its memory address, + * using the Arrow C Data Interface release callback to transfer ownership. + * + *

This method intentionally takes a raw {@code streamAddress} (an {@code ArrowArrayStream} + * memory address) rather than a Java {@link ArrowArrayStream} object. A typed parameter would be + * an {@code org.apache.arrow.c.ArrowArrayStream} loaded by Lance's classloader / Arrow + * version; a caller running a different Arrow version (or under a different classloader, e.g. + * Spark + a native engine bundling its own Arrow) cannot construct that exact type and would hit + * a {@code ClassCastException}/{@code NoSuchMethodError} at the very boundary this method exists + * to cross. The C Data Interface ABI is stable across Arrow versions, so passing the C struct's + * address keeps the two sides fully decoupled: the caller allocates the stream with its + * own Arrow runtime and only the {@code long} address crosses into Lance. See gluten#12263 + * for the cross-Arrow-version integration that motivated this. + * + *

Unlike {@link #scanBatches()}, no Java Arrow {@link ArrowReader} is created on Lance's side: + * Lance writes the C struct directly at {@code streamAddress} and the caller drives the read loop + * with its own Arrow runtime. + * + *

The {@code streamAddress} must point to a freshly-allocated, empty {@code ArrowArrayStream} + * (its {@code release} callback must be null). Exporting into a stream that already holds a + * producer is rejected with an {@link IllegalArgumentException}, because overwriting the struct + * would drop the existing {@code release} callback and leak the first producer. The caller owns + * the stream and is responsible for closing it; the release callback installed by this call + * routes back through Lance's native side. + * + *

The provided stream must not be shared across concurrent exports. An {@code + * ArrowArrayStream} is a plain C struct in caller-owned memory with no internal synchronization, + * so a single stream must be exported into, then drained, by one thread at a time. The + * already-populated check above guards the sequential "export twice" mistake, but it cannot make + * two concurrent exports into the same struct safe — that is a caller-side data race on + * caller-owned memory, the same contract as Arrow's C Data Interface itself. Use a separate + * stream per concurrent export. + * + *

Example (caller on its own Arrow version / allocator): + * + * 

{@code
+   * try (ArrowArrayStream stream = ArrowArrayStream.allocateNew(callerAllocator)) {
+   *   scanner.exportArrowStream(stream.memoryAddress());
+   *   try (ArrowReader reader = Data.importArrayStream(callerAllocator, stream)) {
+   *     while (reader.loadNextBatch()) {
+   *       VectorSchemaRoot batch = reader.getVectorSchemaRoot();
+   *       // ...
+   *     }
+   *   }
+   * }
+   * }
+ * + * @param streamAddress the memory address of a freshly-allocated, empty {@code ArrowArrayStream} + * to populate + * @throws IllegalArgumentException if the scanner is closed or the stream is already populated + * @throws IOException if the native scan fails to start + */ + public void exportArrowStream(long streamAddress) throws IOException { + try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) { + Preconditions.checkArgument(nativeScannerHandle != 0, "Scanner is closed"); + openStream(streamAddress); + } + } + private native void openStream(long streamAddress) throws IOException; @Override diff --git a/java/src/test/java/org/lance/ScannerTest.java b/java/src/test/java/org/lance/ScannerTest.java index 00434034b64..d824d8bb13f 100644 --- a/java/src/test/java/org/lance/ScannerTest.java +++ b/java/src/test/java/org/lance/ScannerTest.java @@ -22,6 +22,8 @@ import org.lance.ipc.ScanOptions; import org.lance.ipc.ScanStats; +import org.apache.arrow.c.ArrowArrayStream; +import org.apache.arrow.c.Data; import org.apache.arrow.dataset.scanner.Scanner; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; @@ -39,6 +41,7 @@ import org.junit.jupiter.api.io.TempDir; import java.io.IOException; +import java.nio.charset.StandardCharsets; import java.nio.file.Path; import java.util.ArrayList; import java.util.Arrays; @@ -49,6 +52,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; public class ScannerTest { @@ -158,6 +162,450 @@ void testDatasetScannerSchema(@TempDir Path tempDir) throws Exception { } } + /** + * Imports a caller-owned C stream populated by {@link LanceScanner#exportArrowStream(long)} and + * returns the {@code id} values in the order the stream produced them. + * + *

The projected schema is asserted to be exactly a single {@code id: int32} field, and the + * assertion is made on the imported reader before the first {@code loadNextBatch()} call + * so that it still runs for an empty (zero-batch) result — a regression that exported the wrong + * schema for an empty scan would otherwise slip through. See {@code + * org.apache.arrow.vector.ipc.ArrowReader#getVectorSchemaRoot()}, which exposes the schema as + * soon as the stream is imported. + * + *

This helper intentionally makes no assertion about per-batch row counts. The + * scanner's {@code batchSize} is only a hint unless {@code strictBatchSize(true)} is set, so the + * number of batches and the rows per batch are not part of the contract being tested here; that + * dimension is covered separately by {@link #testExportArrowStreamStrictBatchSize}. Row ordering + * and exact values are asserted by the callers against the returned list. + */ + private static List drainIdStream(BufferAllocator allocator, ArrowArrayStream stream) + throws IOException { + List ids = new ArrayList<>(); + try (ArrowReader reader = Data.importArrayStream(allocator, stream)) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + List fields = root.getSchema().getFields(); + assertEquals(1, fields.size()); + assertEquals("id", fields.get(0).getName()); + assertEquals(ArrowType.ArrowTypeID.Int, fields.get(0).getType().getTypeID()); + while (reader.loadNextBatch()) { + IntVector vector = (IntVector) root.getVector("id"); + int rowsInBatch = vector.getValueCount(); + for (int i = 0; i < rowsInBatch; i++) { + ids.add(vector.get(i)); + } + } + } + return ids; + } + + /** + * Happy path: a single-fragment ordered scan exported through a caller-owned C stream returns + * every row exactly once, in scan order. The caller allocates the {@link ArrowArrayStream} from + * its own allocator and passes only the memory address; the scanner fills the C struct in place. + * This is the cross-Arrow-version / cross-classloader boundary the API exists to serve. + */ + @Test + void testExportArrowStream(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("export_stream_basic").toString(); + try (BufferAllocator allocator = new RootAllocator()) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + testDataset.createEmptyDataset().close(); + int totalRows = 40; + int batchRows = 20; + try (Dataset dataset = testDataset.write(1, totalRows)) { + try (LanceScanner scanner = + dataset.newScan( + new ScanOptions.Builder() + .batchSize(batchRows) + .columns(Arrays.asList("id")) + .build())) { + try (ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { + scanner.exportArrowStream(stream.memoryAddress()); + // SimpleTestDataset writes id = 0..totalRows-1; an ordered scan must return them in + // exactly that sequence, so assert the exact ordering (no sort). + List ids = drainIdStream(allocator, stream); + assertEquals(totalRows, ids.size()); + for (int i = 0; i < totalRows; i++) { + assertEquals(i, ids.get(i)); + } + } + } + } + } + } + + /** + * A scan that spans multiple fragments is exported as a single C stream that concatenates the + * fragments in fragment order. {@code createNewFragment(40, 10)} produces 4 fragments of 10 rows + * (ids 0-9, 10-19, 20-29, 30-39), and an ordered scan must return 0..39 in exactly that order. + * + *

The expected ids are asserted in stream order without sorting: sorting would mask a + * regression that returned fragments out of order, which is exactly the kind of bug this test + * exists to catch. A non-divisor batch size (7) is used so batch boundaries do not line up with + * fragment boundaries, exercising the stream's batch stitching across fragments. + */ + @Test + void testExportArrowStreamMultipleFragments(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("export_stream_multi_fragment").toString(); + try (BufferAllocator allocator = new RootAllocator()) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + testDataset.createEmptyDataset().close(); + int totalRows = 40; + // maxRowsPerFile < totalRows forces multiple fragments (4 fragments of 10 rows). + List fragments = testDataset.createNewFragment(totalRows, 10); + assertEquals(4, fragments.size()); + FragmentOperation.Append appendOp = new FragmentOperation.Append(fragments); + try (Dataset dataset = Dataset.commit(allocator, datasetPath, appendOp, Optional.of(1L))) { + int batchRows = 7; // deliberately not a divisor of any fragment size + try (LanceScanner scanner = + dataset.newScan( + new ScanOptions.Builder() + .batchSize(batchRows) + .columns(Arrays.asList("id")) + .build())) { + try (ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { + scanner.exportArrowStream(stream.memoryAddress()); + List ids = drainIdStream(allocator, stream); + assertEquals(totalRows, ids.size()); + // Assert exact scan order (no sort) so out-of-order fragments would fail. + for (int i = 0; i < totalRows; i++) { + assertEquals(i, ids.get(i), "row " + i + " out of expected scan order"); + } + } + } + } + } + } + + /** + * A pushed-down filter is honored by the exported stream: only matching rows cross the C-data + * boundary. {@code id < 20} over ids 0..39 must yield exactly 0..19 in order. Asserted in scan + * order without sorting so a filter/ordering regression cannot hide behind a sort. + */ + @Test + void testExportArrowStreamWithFilter(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("export_stream_filter").toString(); + try (BufferAllocator allocator = new RootAllocator()) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + testDataset.createEmptyDataset().close(); + try (Dataset dataset = testDataset.write(1, 40)) { + try (LanceScanner scanner = + dataset.newScan( + new ScanOptions.Builder() + .batchSize(50) + .columns(Arrays.asList("id")) + .filter("id < 20") + .build())) { + try (ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { + scanner.exportArrowStream(stream.memoryAddress()); + List ids = drainIdStream(allocator, stream); + assertEquals(20, ids.size()); + for (int i = 0; i < 20; i++) { + assertEquals(i, ids.get(i)); + } + } + } + } + } + } + + /** + * Pushed-down limit and offset are honored by the exported stream. Over ids 0..39, {@code + * offset(10).limit(5)} must yield exactly [10, 11, 12, 13, 14] in order — asserted as an exact + * ordered list so both the window bounds and the ordering are checked. + */ + @Test + void testExportArrowStreamWithLimitOffset(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("export_stream_limit_offset").toString(); + try (BufferAllocator allocator = new RootAllocator()) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + testDataset.createEmptyDataset().close(); + try (Dataset dataset = testDataset.write(1, 40)) { + try (LanceScanner scanner = + dataset.newScan( + new ScanOptions.Builder() + .batchSize(50) + .columns(Arrays.asList("id")) + .limit(5) + .offset(10) + .build())) { + try (ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { + scanner.exportArrowStream(stream.memoryAddress()); + List ids = drainIdStream(allocator, stream); + assertEquals(Arrays.asList(10, 11, 12, 13, 14), ids); + } + } + } + } + } + + /** + * Column projection is reflected in the exported stream's schema. {@code SimpleTestDataset} has + * columns {@code (id, name)}; projecting only {@code name} must produce a stream whose schema is + * exactly that one column. The schema is checked on the imported reader before draining, and the + * full row count is verified after. + */ + @Test + void testExportArrowStreamProjectsRequestedColumnsOnly(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("export_stream_projection").toString(); + try (BufferAllocator allocator = new RootAllocator()) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + testDataset.createEmptyDataset().close(); + try (Dataset dataset = testDataset.write(1, 10)) { + // Project only "name"; the exported stream's schema must contain exactly that column. + try (LanceScanner scanner = + dataset.newScan(new ScanOptions.Builder().columns(Arrays.asList("name")).build())) { + try (ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { + scanner.exportArrowStream(stream.memoryAddress()); + try (ArrowReader reader = Data.importArrayStream(allocator, stream)) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + assertEquals(1, root.getSchema().getFields().size()); + assertEquals("name", root.getSchema().getFields().get(0).getName()); + int rows = 0; + while (reader.loadNextBatch()) { + rows += root.getRowCount(); + } + assertEquals(10, rows); + } + } + } + } + } + } + + /** + * A scan that matches no rows ({@code id < 0}) still exports a valid, well-formed stream that + * yields zero rows. {@link #drainIdStream} asserts the projected schema ({@code id: int32}) on + * the imported reader before any {@code loadNextBatch()}, so this case also guards the empty-scan + * schema — a regression that exported a wrong or absent schema for zero-row results would fail + * here even though no batch is ever produced. + */ + @Test + void testExportArrowStreamEmptyResult(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("export_stream_empty").toString(); + try (BufferAllocator allocator = new RootAllocator()) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + testDataset.createEmptyDataset().close(); + try (Dataset dataset = testDataset.write(1, 40)) { + try (LanceScanner scanner = + dataset.newScan( + new ScanOptions.Builder().columns(Arrays.asList("id")).filter("id < 0").build())) { + try (ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { + scanner.exportArrowStream(stream.memoryAddress()); + List ids = drainIdStream(allocator, stream); + assertTrue(ids.isEmpty()); + } + } + } + } + } + + /** + * Guards against the sequential "export twice into the same stream" mistake. After the first + * export installs a producer (non-null {@code release} callback), a second export into the same + * stream must be rejected with {@link IllegalArgumentException} rather than overwriting the C + * struct in place — overwriting would drop the first producer's release callback and leak it. + * + *

The test also verifies the rejection is non-destructive: the first producer is still intact + * and fully drainable (all 40 rows) after the rejected second call. This is the single-threaded + * misuse case; concurrent exports into one caller-owned stream are the caller's responsibility, + * as documented on {@link LanceScanner#exportArrowStream(long)}. + */ + @Test + void testExportArrowStreamRejectsPopulatedStream(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("export_stream_reject_populated").toString(); + try (BufferAllocator allocator = new RootAllocator()) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + testDataset.createEmptyDataset().close(); + try (Dataset dataset = testDataset.write(1, 40)) { + try (LanceScanner scanner = + dataset.newScan(new ScanOptions.Builder().columns(Arrays.asList("id")).build())) { + try (ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { + // First export populates the stream and installs a release callback. + scanner.exportArrowStream(stream.memoryAddress()); + // Exporting again into the same (already-populated) stream must be rejected rather + // than silently overwriting and leaking the first producer's release callback. + IllegalArgumentException ex = + assertThrows( + IllegalArgumentException.class, + () -> scanner.exportArrowStream(stream.memoryAddress())); + assertTrue(ex.getMessage().toLowerCase().contains("already populated")); + // The first producer is still intact and drainable. + try (ArrowReader reader = Data.importArrayStream(allocator, stream)) { + int rows = 0; + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + while (reader.loadNextBatch()) { + rows += root.getRowCount(); + } + assertEquals(40, rows); + } + } + } + } + } + } + + /** + * A null (0) stream address is rejected with {@link IllegalArgumentException} before any native + * dereference, so a caller mistake cannot turn into a native null-pointer write. + */ + @Test + void testExportArrowStreamRejectsNullAddress(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("export_stream_reject_null").toString(); + try (BufferAllocator allocator = new RootAllocator()) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + testDataset.createEmptyDataset().close(); + try (Dataset dataset = testDataset.write(1, 10)) { + try (LanceScanner scanner = + dataset.newScan(new ScanOptions.Builder().columns(Arrays.asList("id")).build())) { + assertThrows(IllegalArgumentException.class, () -> scanner.exportArrowStream(0L)); + } + } + } + } + + /** + * Exporting from a closed scanner is rejected with {@link IllegalArgumentException} (the native + * scanner handle is zero after {@code close()}), rather than dereferencing a freed handle. The + * scanner is closed explicitly here, so it is intentionally not in a try-with-resources. + */ + @Test + void testExportArrowStreamRejectsClosedScanner(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("export_stream_reject_closed").toString(); + try (BufferAllocator allocator = new RootAllocator()) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + testDataset.createEmptyDataset().close(); + try (Dataset dataset = testDataset.write(1, 10)) { + LanceScanner scanner = + dataset.newScan(new ScanOptions.Builder().columns(Arrays.asList("id")).build()); + scanner.close(); + try (ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { + assertThrows( + IllegalArgumentException.class, + () -> scanner.exportArrowStream(stream.memoryAddress())); + } + } + } + } + + /** + * Null values survive the C-data export round-trip. {@code writeSortByDataset} writes 10 rows + * (insertion order) in which {@code id} is null at rows 2 and 5 and {@code name} is null at rows + * 0 and 6. An unordered scan returns rows in insertion order, so the exported stream must + * reproduce both the non-null values and the null positions exactly — null/validity bitmaps are a + * common casualty of an incorrect C-data export, so this guards them explicitly. + */ + @Test + void testExportArrowStreamPreservesNulls(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("export_stream_nulls").toString(); + try (BufferAllocator allocator = new RootAllocator()) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + testDataset.createEmptyDataset().close(); + try (Dataset dataset = testDataset.writeSortByDataset(1)) { + // Insertion order, row -> (id, name): + // 0 -> (0, null) 3 -> (2, "P2") 6 -> (3, null) 9 -> (5, "P5") + // 1 -> (1, "P0") 4 -> (2, "P3") 7 -> (4, "P4") + // 2 -> (null,"P1") 5 -> (null,"P3") 8 -> (4, "P5") + Integer[] expectedIds = {0, 1, null, 2, 2, null, 3, 4, 4, 5}; + String[] expectedNames = {null, "P0", "P1", "P2", "P3", "P3", null, "P4", "P5", "P5"}; + try (LanceScanner scanner = + dataset.newScan( + new ScanOptions.Builder().columns(Arrays.asList("id", "name")).build())) { + try (ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { + scanner.exportArrowStream(stream.memoryAddress()); + try (ArrowReader reader = Data.importArrayStream(allocator, stream)) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + assertEquals(2, root.getSchema().getFields().size()); + int row = 0; + while (reader.loadNextBatch()) { + IntVector idVector = (IntVector) root.getVector("id"); + VarCharVector nameVector = (VarCharVector) root.getVector("name"); + for (int i = 0; i < root.getRowCount(); i++, row++) { + if (expectedIds[row] == null) { + assertTrue(idVector.isNull(i), "id should be null at row " + row); + } else { + assertEquals( + expectedIds[row].intValue(), idVector.get(i), "id mismatch at row " + row); + } + if (expectedNames[row] == null) { + assertTrue(nameVector.isNull(i), "name should be null at row " + row); + } else { + assertEquals( + expectedNames[row], + new String(nameVector.get(i), StandardCharsets.UTF_8), + "name mismatch at row " + row); + } + } + } + assertEquals(expectedIds.length, row); + } + } + } + } + } + } + + /** + * With {@code strictBatchSize(true)}, the exported stream must split into batches no larger than + * the requested batch size, and still reproduce every row in order. This is the one place the + * per-batch size is part of the contract; the other export tests deliberately leave batch sizing + * unasserted because it is only a hint by default. Mirrors {@link #testStrictBatchSize} but over + * the C-data export path. A batch size of 10 over 25 rows yields batches of at most 10. + */ + @Test + void testExportArrowStreamStrictBatchSize(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("export_stream_strict_batch").toString(); + try (BufferAllocator allocator = new RootAllocator()) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + testDataset.createEmptyDataset().close(); + int totalRows = 25; + int batchSize = 10; + try (Dataset dataset = testDataset.write(1, totalRows)) { + try (LanceScanner scanner = + dataset.newScan( + new ScanOptions.Builder() + .batchSize(batchSize) + .strictBatchSize(true) + .columns(Arrays.asList("id")) + .build())) { + try (ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator)) { + scanner.exportArrowStream(stream.memoryAddress()); + try (ArrowReader reader = Data.importArrayStream(allocator, stream)) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + List ids = new ArrayList<>(); + while (reader.loadNextBatch()) { + int rowsInBatch = root.getRowCount(); + assertTrue( + rowsInBatch <= batchSize, + "strict: batch of " + rowsInBatch + " should be <= " + batchSize); + IntVector idVector = (IntVector) root.getVector("id"); + for (int i = 0; i < rowsInBatch; i++) { + ids.add(idVector.get(i)); + } + } + assertEquals(totalRows, ids.size()); + for (int i = 0; i < totalRows; i++) { + assertEquals(i, ids.get(i)); + } + } + } + } + } + } + } + @Test void testDatasetScannerCountRows(@TempDir Path tempDir) throws Exception { String datasetPath = tempDir.resolve("dataset_scanner_count").toString();