From 3c0cc7ee0c18f2c409ee4ef4b6078c3cc31e5b8d Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 20:34:29 +0000 Subject: [PATCH 1/2] build(deps): Bump dev.vortex:vortex-jni from 0.88.0 to 0.89.0 Bumps [dev.vortex:vortex-jni](https://github.com/spiraldb/vortex) from 0.88.0 to 0.89.0. - [Release notes](https://github.com/spiraldb/vortex/releases) - [Commits](https://github.com/spiraldb/vortex/compare/0.88.0...0.89.0) --- updated-dependencies: - dependency-name: dev.vortex:vortex-jni dependency-version: 0.89.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 00319bf8a..b2bcc499b 100644 --- a/pom.xml +++ b/pom.xml @@ -78,7 +78,7 @@ 0.30.0 3.27.7 5.24.0 - 0.88.0 + 0.89.0 19.0.0 2.0.20 From 841c3c99b67dc4714801b0e0fc2a74a343e90b40 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sat, 24 Oct 2026 13:02:55 +0200 Subject: [PATCH 2/2] fix(writer): cut chunks where vortex-jni 0.89 cuts them Rust 0.89 measures the 8192-row blocks its coalescing repartition buffers by the bytes their rows reference (vortex-data/vortex#10439, #10445): a string of at most 12 bytes costs only its view, and a block with no null carries no validity. Its outer repartition also re-cuts the stream into blocks aligned to the column's first row before they are measured, so a block spanning two batches is measured as one array. The Repartitioner now cuts at absolute block boundaries and measures each whole block that way; a partial last block that tops the pending blocks up to 1 MiB is flushed as its own chunk, as Rust's ChunksBuffer does. RepartitionInteropIntegrationTest pins identical chunk boundaries against vortex-jni for six string column shapes and five null patterns of a nullable I16 column, where a block's validity moves the cut. CalciteInteropIntegrationTest's string fixture no longer reached 1 MiB per column under the exact size, so it grows to 50k-row batches to keep several chunks. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01SwhK4W5cgj7y7s4ZPCif5x --- CHANGELOG.md | 1 + docs/compatibility.md | 4 +- docs/reference.md | 4 +- .../CalciteInteropIntegrationTest.java | 9 +- .../RepartitionInteropIntegrationTest.java | 222 ++++++++++++++++ .../dfa1/vortex/writer/Repartitioner.java | 251 ++++++++++-------- .../dfa1/vortex/writer/VortexWriter.java | 15 +- .../dfa1/vortex/writer/RepartitionerTest.java | 135 +++++++++- 8 files changed, 511 insertions(+), 130 deletions(-) create mode 100644 integration/src/test/java/io/github/dfa1/vortex/integration/RepartitionInteropIntegrationTest.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 2adda6f53..39f97d23b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - An all-null chunk stored as a null constant (as vortex-jni writes it) reads as a masked array with every row null, like any nullable chunk, instead of a `NullArray` that the CLI, SQL and Parquet export failed on. ([#594](https://github.com/dfa1/vortex-java/pull/594)) - An all-null date, time, timestamp or uuid chunk is written as a bare null `vortex.constant`, as Rust writes it, without a `vortex.ext` around it. ([#595](https://github.com/dfa1/vortex-java/pull/595)) - An all-null decimal chunk reads as a masked array too, and a filter selects every row of an all-null uuid or list chunk for `IS NULL`, where it selected none. ([#595](https://github.com/dfa1/vortex-java/pull/595)) +- The writer cuts string and nullable columns into the same chunks as vortex-jni 0.89, which measures its 8192-row blocks by the bytes their rows reference. ([#596](https://github.com/dfa1/vortex-java/pull/596)) - `filter` and `select --where` on the command line work on a nullable column, where they failed; a null row matches no comparison, as in SQL. ([#593](https://github.com/dfa1/vortex-java/pull/593)) ## [0.18.1] — 2026-10-22 diff --git a/docs/compatibility.md b/docs/compatibility.md index dad1a1d3b..183aa630a 100644 --- a/docs/compatibility.md +++ b/docs/compatibility.md @@ -1,6 +1,6 @@ # Compatibility -Tested against the [Rust reference implementation](https://github.com/vortex-data/vortex) v0.88.0. +Tested against the [Rust reference implementation](https://github.com/vortex-data/vortex) v0.89.0. For the rest of the API surface (reader, writer, scan, CLI), see [reference.md](reference.md). ## Read-only deployment @@ -256,7 +256,7 @@ guarantee once frozen (ADR 0023) — a write-time/read-time policy, not part of | `preview2026.08.0` | `preview` | nothing yet | | `zstd2026.02.0` | `zstd` | `vortex.zstd_buffers` (read only; buffer-level Zstd for GPU decode; declared by Rust's `vortex-zstd` plugin, opt-in) | -Mirrors Rust's `vortex-edition` declarations at the pinned vortex-jni release (0.88.0); +Mirrors Rust's `vortex-edition` declarations at the pinned vortex-jni release (0.89.0); `EditionCatalogParityIntegrationTest` fails the build on drift. `core` editions are frozen with a forever read-compatibility guarantee; `preview` holds opt-in components not yet adopted by `core`; `zstd` is declared by Rust's `vortex-zstd` plugin crate rather than the core declarations, and is diff --git a/docs/reference.md b/docs/reference.md index 49c2748fc..30ac2caba 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -148,7 +148,9 @@ Writes a Vortex file. Implements `Closeable`. The file is complete and readable Batches are not chunks: as Rust's writer does, each primitive, boolean, decimal, string and binary column is repartitioned into chunks of whole 8192-row blocks reaching 1 MB of uncompressed data -(e.g. 131 072 rows of I64, 524 288 of U16 dict codes), the rest becoming the last chunk. Nested +(e.g. 131 072 rows of I64, 524 288 of U16 dict codes), the rest becoming the last chunk. Blocks +count from the column's first row, whatever the batches, and are measured as Rust measures them: a +string of at most 12 bytes costs only its 16-byte view, and a block without nulls no validity. Nested columns, and targets before `core2026.08.0` (legacy `vortex.stats` zones), keep one chunk per batch. ### `Chunk` (`io.github.dfa1.vortex.writer.Chunk`) diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/CalciteInteropIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/CalciteInteropIntegrationTest.java index 1d05b9a70..44b6e1e58 100644 --- a/integration/src/test/java/io/github/dfa1/vortex/integration/CalciteInteropIntegrationTest.java +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/CalciteInteropIntegrationTest.java @@ -57,8 +57,8 @@ class CalciteInteropIntegrationTest { // several batches, so the row count is a sum over chunks, not a single chunk's length private static final int BATCHES = 3; private static final int BATCH_ROWS = 1_000; - // over 1 MiB a batch, the least Rust writes as a chunk of its own - private static final int EVENT_BATCH_ROWS = 20_000; + // three batches past 1 MiB per column, so Rust writes several chunks (see writeEventsWithJni) + private static final int EVENT_BATCH_ROWS = 50_000; private static final LocalDate DATE = LocalDate.of(2024, 2, 29); private static final LocalTime TIME = LocalTime.of(13, 45, 30); // pre-epoch on purpose: millisecond truncation must floor, not round toward zero @@ -760,8 +760,9 @@ void jniWriter_whereOnNarrowIntegers_selectsTheMatchingRows(String where, long e /// Writes `id`, `e` (an ascending event id, "é"-suffixed every 13th row, null every 97th), `t` /// (an ascending value repeated ten times) and `r` (a scrambled value with repeats) in [#BATCHES] - /// batches of [#EVENT_BATCH_ROWS], over 1 MiB each, so Rust writes several chunks: Rust coalesces - /// smaller batches into one. Its chunks hold a multiple of 8,192 rows, not of ten, so a run of + /// batches of [#EVENT_BATCH_ROWS], past 1 MiB per column by Rust's exact size (a string of at + /// most 12 bytes costs only its 16-byte view), so Rust writes several chunks: Rust coalesces + /// smaller columns into one. Its chunks hold a multiple of 8,192 rows, not of ten, so a run of /// equal `t` straddles each chunk edge. private static void writeEventsWithJni(Path file) throws Exception { Schema schema = new Schema(List.of( diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/RepartitionInteropIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/RepartitionInteropIntegrationTest.java new file mode 100644 index 000000000..0d68d9394 --- /dev/null +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/RepartitionInteropIntegrationTest.java @@ -0,0 +1,222 @@ +package io.github.dfa1.vortex.integration; + +import dev.vortex.api.Session; +import dev.vortex.api.VortexWriter; +import dev.vortex.arrow.ArrowAllocation; +import dev.vortex.jni.NativeLoader; +import io.github.dfa1.vortex.core.model.ColumnName; +import io.github.dfa1.vortex.core.model.DType; +import io.github.dfa1.vortex.core.model.PType; +import io.github.dfa1.vortex.reader.ScanIterator; +import io.github.dfa1.vortex.reader.ScanOptions; +import io.github.dfa1.vortex.reader.VortexReader; +import io.github.dfa1.vortex.writer.WriteOptions; +import io.github.dfa1.vortex.writer.encode.NullableData; +import org.apache.arrow.c.ArrowArray; +import org.apache.arrow.c.ArrowSchema; +import org.apache.arrow.c.Data; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.SmallIntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import java.io.IOException; +import java.nio.channels.FileChannel; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.function.IntFunction; +import java.util.function.IntPredicate; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; + +/// The writers coalesce `writeChunk` batches into chunks of the fewest 8192-row blocks reaching +/// 1 MiB, measured by Rust's exact uncompressed size (vortex-data/vortex#10445): a string of at +/// most 12 bytes costs only its view, and a block with no null carries no validity. The same +/// column through vortex-jni and vortex-java must cut the same chunks. +@Tag("slow") // about 1 s a case: each writes the column with both writers +class RepartitionInteropIntegrationTest { + + private static final Session SESSION = Session.create(); + private static final BufferAllocator ALLOCATOR = ArrowAllocation.rootAllocator(); + private static final ColumnName S = ColumnName.of("s"); + private static final int BATCHES = 4; + private static final int BATCH_ROWS = 60_000; + // 600k rows: past the 64 blocks of I16 that reach 1 MiB + private static final int SHORT_BATCHES = 10; + + static { + NativeLoader.loadJni(); + } + + static Stream columns() { + Random random = new Random(7); + // distinct values, so neither writer dictionary-encodes the column: dict codes chunk by + // their own size + int[] lengths = random.ints(BATCHES * BATCH_ROWS, 0, 24).toArray(); + return Stream.of( + Arguments.of("inlined strings", false, (IntFunction) row -> "v%07d".formatted(row)), + Arguments.of("out-of-line strings", false, (IntFunction) row -> "value-%010d".formatted(row)), + Arguments.of("mixed lengths", false, (IntFunction) row -> "%06d".formatted(row) + "x".repeat(lengths[row])), + Arguments.of("nullable, no nulls", true, (IntFunction) row -> "v%07d".formatted(row)), + Arguments.of("nullable, sparse nulls", true, + (IntFunction) row -> row % 50_000 == 7 ? null : "v%07d".formatted(row)), + Arguments.of("nullable, null runs", true, + (IntFunction) row -> row / 8192 % 3 == 0 ? null : "value-%010d".formatted(row))); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("columns") + void javaCutsTheChunksJniCuts(String name, boolean nullable, IntFunction valueAt, @TempDir Path tmp) + throws Exception { + // Given + Path jniFile = tmp.resolve("jni.vortex"); + Path javaFile = tmp.resolve("java.vortex"); + writeWithJni(jniFile, nullable, valueAt); + + // When + writeWithJava(javaFile, nullable, valueAt); + + // Then + long[] expected = chunkRowCounts(jniFile); + assertThat(expected).as("Rust cut more than one chunk").hasSizeGreaterThan(1); + assertThat(chunkRowCounts(javaFile)).isEqualTo(expected); + } + + // At 2 B a row of I16 values, a block's validity moves the cut: 64 blocks reach 1 MiB without a + // bitmap, 61 with one. 5000-row null runs straddle every block, so a block measured per batch + // instead of from the column's first row would carry no bitmap. + static Stream nullPatterns() { + return Stream.of( + Arguments.of("no nulls", (IntPredicate) row -> false), + Arguments.of("sparse nulls", (IntPredicate) row -> row % 50_000 == 7), + Arguments.of("null blocks", (IntPredicate) row -> row / 8192 % 3 == 0), + Arguments.of("null runs across blocks", (IntPredicate) row -> row / 5000 % 2 == 0), + Arguments.of("all null", (IntPredicate) row -> true)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("nullPatterns") + void javaCutsTheChunksJniCuts_nullableI16(String name, IntPredicate isNull, @TempDir Path tmp) throws Exception { + // Given + Path jniFile = tmp.resolve("jni.vortex"); + Path javaFile = tmp.resolve("java.vortex"); + writeShortsWithJni(jniFile, isNull); + + // When + writeShortsWithJava(javaFile, isNull); + + // Then + long[] expected = chunkRowCounts(jniFile); + assertThat(expected).as("Rust cut more than one chunk").hasSizeGreaterThan(1); + assertThat(chunkRowCounts(javaFile)).isEqualTo(expected); + } + + private static void writeShortsWithJni(Path path, IntPredicate isNull) throws Exception { + Schema schema = new Schema(List.of(Field.nullable("s", new ArrowType.Int(16, true)))); + try (VortexWriter writer = VortexWriter.builder(SESSION, path.toAbsolutePath().toUri().toString(), schema, + ALLOCATOR).build()) { + for (int batch = 0; batch < SHORT_BATCHES; batch++) { + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, ALLOCATOR)) { + SmallIntVector s = (SmallIntVector) root.getVector("s"); + s.allocateNew(BATCH_ROWS); + for (int i = 0; i < BATCH_ROWS; i++) { + int row = batch * BATCH_ROWS + i; + if (isNull.test(row)) { + s.setNull(i); + } else { + s.set(i, (short) row); + } + } + root.setRowCount(BATCH_ROWS); + try (ArrowArray array = ArrowArray.allocateNew(ALLOCATOR); + ArrowSchema arrowSchema = ArrowSchema.allocateNew(ALLOCATOR)) { + Data.exportVectorSchemaRoot(ALLOCATOR, root, null, array, arrowSchema); + writer.writeBatch(array.memoryAddress(), arrowSchema.memoryAddress()); + } + } + } + } + } + + private static void writeShortsWithJava(Path path, IntPredicate isNull) throws Exception { + DType.Struct schema = new DType.Struct(List.of(S), List.of(new DType.Primitive(PType.I16, true)), false); + try (var channel = FileChannel.open(path, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var writer = io.github.dfa1.vortex.writer.VortexWriter.create(channel, schema, WriteOptions.defaults())) { + for (int batch = 0; batch < SHORT_BATCHES; batch++) { + short[] s = new short[BATCH_ROWS]; + boolean[] valid = new boolean[BATCH_ROWS]; + for (int i = 0; i < BATCH_ROWS; i++) { + int row = batch * BATCH_ROWS + i; + valid[i] = !isNull.test(row); + s[i] = valid[i] ? (short) row : 0; + } + writer.writeChunk(Map.of(S, new NullableData(s, valid))); + } + } + } + + private static long[] chunkRowCounts(Path file) throws IOException { + try (VortexReader reader = VortexReader.open(file); + ScanIterator scan = reader.scan(ScanOptions.all())) { + return scan.chunkRowCounts(); + } + } + + private static void writeWithJni(Path path, boolean nullable, IntFunction valueAt) throws Exception { + Field field = nullable ? Field.nullable("s", new ArrowType.Utf8()) : Field.notNullable("s", new ArrowType.Utf8()); + Schema schema = new Schema(List.of(field)); + try (VortexWriter writer = VortexWriter.builder(SESSION, path.toAbsolutePath().toUri().toString(), schema, + ALLOCATOR).build()) { + for (int batch = 0; batch < BATCHES; batch++) { + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, ALLOCATOR)) { + VarCharVector s = (VarCharVector) root.getVector("s"); + s.allocateNew(BATCH_ROWS); + for (int i = 0; i < BATCH_ROWS; i++) { + String value = valueAt.apply(batch * BATCH_ROWS + i); + if (value == null) { + s.setNull(i); + } else { + s.setSafe(i, value.getBytes(StandardCharsets.UTF_8)); + } + } + root.setRowCount(BATCH_ROWS); + try (ArrowArray array = ArrowArray.allocateNew(ALLOCATOR); + ArrowSchema arrowSchema = ArrowSchema.allocateNew(ALLOCATOR)) { + Data.exportVectorSchemaRoot(ALLOCATOR, root, null, array, arrowSchema); + writer.writeBatch(array.memoryAddress(), arrowSchema.memoryAddress()); + } + } + } + } + } + + private static void writeWithJava(Path path, boolean nullable, IntFunction valueAt) throws Exception { + DType.Struct schema = new DType.Struct(List.of(S), List.of(new DType.Utf8(nullable)), false); + try (var channel = FileChannel.open(path, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var writer = io.github.dfa1.vortex.writer.VortexWriter.create(channel, schema, WriteOptions.defaults())) { + for (int batch = 0; batch < BATCHES; batch++) { + String[] s = new String[BATCH_ROWS]; + boolean[] valid = new boolean[BATCH_ROWS]; + for (int i = 0; i < BATCH_ROWS; i++) { + String value = valueAt.apply(batch * BATCH_ROWS + i); + s[i] = value == null ? "" : value; + valid[i] = value != null; + } + writer.writeChunk(Map.of(S, nullable ? new NullableData(s, valid) : s)); + } + } + } +} diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/Repartitioner.java b/writer/src/main/java/io/github/dfa1/vortex/writer/Repartitioner.java index 803ac1850..4856fa4ee 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/Repartitioner.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/Repartitioner.java @@ -7,18 +7,18 @@ import java.lang.reflect.Array; import java.nio.charset.StandardCharsets; -import java.util.ArrayDeque; import java.util.ArrayList; import java.util.Arrays; import java.util.List; /// Coalesces one column's `writeChunk` batches into the chunks Rust's default write strategy -/// emits (`vortex-layout/src/layouts/repartition.rs`, configured in `vortex-file/src/strategy.rs`): -/// each batch is cut into [#BLOCK_ROWS]-row pieces, and the moment at least [#BLOCK_BYTES] of -/// uncompressed data and [#BLOCK_ROWS] rows are pending, every whole block pending is emitted as -/// one chunk (a piece straddling the last block boundary is split). Whatever is pending at the end -/// becomes the last chunk. So a fixed-width column gets chunks of the fewest whole blocks reaching -/// 1 MB: 131 072 rows of 8-byte values, 524 288 of U16 dict codes. +/// emits (`vortex-file/src/strategy.rs`). Its outer `RepartitionStrategy` first re-cuts the stream +/// into [#BLOCK_ROWS]-row blocks aligned to the column's first row, whatever the batches; the inner +/// one (`vortex-layout/src/layouts/repartition.rs`) emits every pending block as one chunk the +/// moment they reach [#BLOCK_BYTES], each block measured as a whole by Rust's exact uncompressed +/// size. The partial block at the end becomes the last chunk, after the pending whole blocks when +/// it tops them up to [#BLOCK_BYTES]. So a fixed-width column gets chunks of the fewest whole blocks +/// reaching 1 MB: 131 072 rows of 8-byte values, 524 288 of U16 dict codes. /// /// Chunk size drives both segment granularity and compression: a ~1 MB chunk gives the cascade's /// sample enough rows to pick what Rust picks (#458, #470). @@ -30,14 +30,16 @@ final class Repartitioner { /// Rust's `data_block_target_bytes` (`ONE_MEG`), the coalescing `block_size_minimum`. static final long BLOCK_BYTES = 1L << 20; - // A row range of one batch, with its Rust nbytes (Rust's ChunksBuffer entry). - private record Piece(Object data, int offset, int length, long nbytes) { + // A row range of one batch inside one block, with the bytes its valid out-of-line strings add. + private record Piece(Object data, int offset, int length, long outOfLineBytes) { } private final DType dtype; - private final ArrayDeque pending = new ArrayDeque<>(); - private long pendingRows; + private final List block = new ArrayList<>(); + private int blockRows; + private final List pending = new ArrayList<>(); private long pendingBytes; + private long rows; Repartitioner(DType dtype) { this.dtype = dtype; @@ -52,34 +54,36 @@ static boolean supports(DType dtype) { || storage instanceof DType.Binary; } - /// Adds one batch and returns the chunks now ready, in order. - /// - /// @param data the batch, as handed to the column's encoder - /// @param rows the batch's row count - /// @return the chunks to write, possibly none - List add(Object data, long rows) { - return add(data, rows, null); + /// The rows added so far: where the next batch starts. + long rows() { + return rows; } - /// Sizes the [#BLOCK_ROWS]-row pieces of a batch ahead of [#add(Object,long,long[])], so a caller - /// can measure the columns of a batch in parallel: columns are independent, and measuring a - /// string column reads every one of its strings. + /// Measures the pieces of a batch ahead of [#add(Object,long,long[])], so a caller can measure + /// the columns of a batch in parallel: columns are independent, and measuring a string column + /// reads every one of its strings. /// - /// @param data the batch, as it will be handed to [#add(Object,long,long[])] - /// @param rows the batch's row count - /// @return the byte count of each piece, in order - long[] sizePieces(Object data, long rows) { + /// @param data the batch, as it will be handed to [#add(Object,long,long[])] + /// @param rows the batch's row count + /// @param firstRow the column's [#rows()] before the batch + /// @return the out-of-line string bytes of each piece, in order + static long[] sizePieces(Object data, long rows, long firstRow) { int len = (int) rows; - long[] sizes = new long[(len + BLOCK_ROWS - 1) / BLOCK_ROWS]; - for (int i = 0; i < sizes.length; i++) { - int off = i * BLOCK_ROWS; - sizes[i] = nbytes(data, off, Math.min(BLOCK_ROWS, len - off)); + int first = pieceRows(len, firstRow); + long[] sizes = new long[len == 0 ? 0 : 1 + (len - first + BLOCK_ROWS - 1) / BLOCK_ROWS]; + int piece = 0; + int off = 0; + while (off < len) { + int length = pieceRows(len - off, firstRow + off); + sizes[piece] = outOfLineBytes(data, off, length); + piece++; + off += length; } return sizes; } /// Whether sizing a batch of `data` reads its values, as for strings and binaries -- the - /// only batches worth [#sizePieces(Object,long)] ahead of time. + /// only batches worth [#sizePieces(Object,long,long)] ahead of time. /// /// @param data the batch /// @return `true` for a variable-width batch @@ -87,114 +91,153 @@ static boolean sizeReadsValues(Object data) { return values(data) instanceof Object[]; } - /// Adds one batch whose pieces were already sized by [#sizePieces(Object,long)]. + /// Adds one batch and returns the chunks now ready, in order. + /// + /// @param data the batch, as handed to the column's encoder + /// @param rows the batch's row count + /// @return the chunks to write, possibly none + List add(Object data, long rows) { + return add(data, rows, null); + } + + /// Adds one batch whose pieces were already measured by [#sizePieces(Object,long,long)]. /// /// @param data the batch, as handed to the column's encoder /// @param rows the batch's row count - /// @param pieceSizes the batch's [#sizePieces(Object,long)], or `null` to size it here + /// @param pieceSizes the batch's [#sizePieces(Object,long,long)], or `null` to measure it here /// @return the chunks to write, possibly none List add(Object data, long rows, long[] pieceSizes) { List out = new ArrayList<>(); int len = (int) rows; - for (int off = 0; off < len; off += BLOCK_ROWS) { - int length = Math.min(BLOCK_ROWS, len - off); - pushBack(pieceSizes != null - ? new Piece(data, off, length, pieceSizes[off / BLOCK_ROWS]) - : piece(data, off, length)); - if (pendingBytes >= BLOCK_BYTES && pendingRows >= BLOCK_ROWS) { - out.add(collectExactBlocks()); + int piece = 0; + int off = 0; + while (off < len) { + int length = pieceRows(len - off, this.rows); + long outOfLine = pieceSizes != null ? pieceSizes[piece] : outOfLineBytes(data, off, length); + piece++; + block.add(new Piece(data, off, length, outOfLine)); + blockRows += length; + this.rows += length; + off += length; + if (blockRows == BLOCK_ROWS) { + pendingBytes += blockBytes(); + pending.addAll(block); + block.clear(); + blockRows = 0; + if (pendingBytes >= BLOCK_BYTES) { + out.add(drainPending()); + } } } return out; } - /// The rows still pending, as the column's last chunk, or `null` when none are. - Object finish() { - if (pending.isEmpty()) { - return null; + /// The rows still pending, as the column's last chunks: the whole blocks, when the partial + /// block tops them up to [#BLOCK_BYTES], then the partial block, as Rust's `ChunksBuffer` + /// does when the outer repartition flushes it; otherwise all of them as one chunk. + List finish() { + List out = new ArrayList<>(); + if (!pending.isEmpty() && pendingBytes + blockBytes() >= BLOCK_BYTES) { + out.add(drainPending()); } - List rest = new ArrayList<>(pending); - pending.clear(); - pendingRows = 0; - pendingBytes = 0; - return concat(rest); - } - - // Rust's ChunksBuffer::collect_exact_blocks: pops whole blocks' worth of rows, splitting the - // piece that straddles the boundary and keeping its tail pending. - private Object collectExactBlocks() { - long remaining = pendingRows / BLOCK_ROWS * BLOCK_ROWS; - List taken = new ArrayList<>(); - while (remaining > 0) { - Piece p = popFront(); - if (p.length() > remaining) { - int head = (int) remaining; - taken.add(piece(p.data(), p.offset(), head)); - pushFront(piece(p.data(), p.offset() + head, p.length() - head)); - remaining = 0; - } else { - taken.add(p); - remaining -= p.length(); - } + pending.addAll(block); + if (!pending.isEmpty()) { + out.add(drainPending()); } - return concat(taken); + block.clear(); + blockRows = 0; + return out; } - private Piece piece(Object data, int offset, int length) { - return new Piece(data, offset, length, nbytes(data, offset, length)); + // A piece starting at column row `at` runs to the end of that row's block, or of the batch. + private static int pieceRows(int remaining, long at) { + return (int) Math.min(remaining, BLOCK_ROWS - at % BLOCK_ROWS); } - private void pushBack(Piece p) { - pending.addLast(p); - pendingRows += p.length(); - pendingBytes += p.nbytes(); + private Object drainPending() { + Object chunk = concat(pending); + pending.clear(); + pendingBytes = 0; + return chunk; } - private void pushFront(Piece p) { - pending.addFirst(p); - pendingRows += p.length(); - pendingBytes += p.nbytes(); + /// Rust's exact `uncompressed_size_in_bytes` of the block being filled, as one canonical + /// array (`vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes`): fixed-width values + /// at their width, Utf8/Binary as a 16-byte view per row plus each valid value too long to + /// inline in its view, and the validity per [#validityBytes(long,long)]. + private long blockBytes() { + long outOfLine = 0; + long valid = 0; + for (Piece p : block) { + outOfLine += p.outOfLineBytes(); + valid += validRows(p); + } + DType storage = dtype instanceof DType.Extension ext ? ext.storageDType() : dtype; + long values = switch (storage) { + case DType.Utf8 _, DType.Binary _ -> 16L * blockRows + outOfLine; + case DType.Bool _ -> (blockRows + 7L) / 8; + default -> (long) blockRows * elementBytes(storage); + }; + return values + validityBytes(valid, blockRows); } - private Piece popFront() { - Piece p = pending.removeFirst(); - pendingRows -= p.length(); - pendingBytes -= p.nbytes(); - return p; + /// Rust's `validity_uncompressed_size_in_bytes` of a block's mask: nothing when every row is + /// valid, a bitmap when some are, and when none are, a `ConstantArray(false)`'s one buffer -- + /// the 2-byte proto `ScalarValue { bool_value: false }`. + private static long validityBytes(long valid, long rows) { + if (valid == rows) { + return 0; + } + if (valid == 0) { + return 2; + } + return (rows + 7) / 8; } - /// Rust's `ArrayRef::nbytes` of the canonical array: fixed-width values at their width plus a - /// validity bitmap when nullable, and Utf8/Binary as Rust's canonical `VarBinView` -- a 16-byte - /// view per row plus its bytes. - // ponytail: a sliced Rust VarBinView counts its whole shared data buffers, so string chunk - // boundaries can differ from Rust's; exact only for fixed-width columns. - private long nbytes(Object data, int offset, int rows) { - long validity = data instanceof NullableData ? (rows + 7L) / 8 : 0; - Object values = values(data); - return validity + switch (values) { + private static long validRows(Piece p) { + if (!(p.data() instanceof NullableData nd)) { + return p.length(); + } + long valid = 0; + for (int i = p.offset(); i < p.offset() + p.length(); i++) { + if (nd.validity()[i]) { + valid++; + } + } + return valid; + } + + // The bytes of the valid values longer than 12, which a VarBinView keeps out of its views + // (BinaryView::MAX_INLINED_SIZE); zero for any other carrier. + private static long outOfLineBytes(Object data, int offset, int rows) { + boolean[] validity = data instanceof NullableData nd ? nd.validity() : null; + long n = 0; + switch (values(data)) { case String[] a -> { - long n = 16L * rows; for (int i = offset; i < offset + rows; i++) { - // Encoding to count beats a char-by-char count by 2x: the JDK encodes ASCII - // with a vectorized copy, and the throwaway array dies young. - n += a[i] == null ? 0 : a[i].getBytes(StandardCharsets.UTF_8).length; + if (a[i] != null && (validity == null || validity[i])) { + // Encoding to count beats a char-by-char count by 2x: the JDK encodes ASCII + // with a vectorized copy, and the throwaway array dies young. + int length = a[i].getBytes(StandardCharsets.UTF_8).length; + n += length > 12 ? length : 0; + } } - yield n; } case byte[][] a -> { - long n = 16L * rows; for (int i = offset; i < offset + rows; i++) { - n += a[i] == null ? 0 : a[i].length; + if (a[i] != null && (validity == null || validity[i])) { + n += a[i].length > 12 ? a[i].length : 0; + } } - yield n; } - case boolean[] _ -> (rows + 7L) / 8; - default -> (long) rows * elementBytes(); - }; + default -> { + // fixed width: blockBytes sizes it from the dtype + } + } + return n; } - private long elementBytes() { - DType storage = dtype instanceof DType.Extension ext ? ext.storageDType() : dtype; + private static long elementBytes(DType storage) { return switch (storage) { case DType.Primitive p -> p.ptype().byteSize(); case DType.Decimal d -> d.byteWidth(); diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java b/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java index a0ae26496..8aad36820 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java @@ -607,12 +607,13 @@ public void writeChunk(Map columns) throws IOException { // of an import. Columns are independent, so the pool sizes them all at once. if (zonedLayout && Repartitioner.supports(colDtype) && Repartitioner.sizeReadsValues(data) && !dictCandidates.contains(colName)) { - // Not the column's own Repartitioner: creating that here would change the order + // Only read the column's own Repartitioner: creating it here would change the order // flushRepartitioners writes the last chunks in. - Repartitioner sizer = new Repartitioner(colDtype); + Repartitioner own = repartitioners.get(colName); + long firstRow = own == null ? 0 : own.rows(); Object batch = data; pieceSizes.put(colName, CompletableFuture.supplyAsync( - () -> sizer.sizePieces(batch, arrayLength(batch)), options.executor())); + () -> Repartitioner.sizePieces(batch, arrayLength(batch), firstRow), options.executor())); } } @@ -704,11 +705,10 @@ private void writeDataChunk(ColumnName colName, DType colDtype, Object data) thr colChunks.get(colName).add(new ChunkRef(segIdx, rowCount)); } - /// Writes every column's pending rows as its last chunk. + /// Writes every column's pending rows as its last chunk or chunks. private void flushRepartitioners() throws IOException { for (Map.Entry e : repartitioners.entrySet()) { - Object last = e.getValue().finish(); - if (last != null) { + for (Object last : e.getValue().finish()) { writeDataChunk(e.getKey(), columnDtype(e.getKey()), last); } } @@ -1580,8 +1580,7 @@ private void writeGlobalDictCodes(ColumnName colName, DictColumnState state, int codesRowCounts.add((long) ((short[]) chunk).length); } } - Object last = repartitioner != null ? repartitioner.finish() : null; - if (last != null) { + for (Object last : repartitioner != null ? repartitioner.finish() : List.of()) { codesSegIdxes.add(writeSegment(codesDtype, last)); codesRowCounts.add((long) ((short[]) last).length); } diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/RepartitionerTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/RepartitionerTest.java index 49cee3f7b..84104cd01 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/RepartitionerTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/RepartitionerTest.java @@ -4,10 +4,15 @@ import io.github.dfa1.vortex.core.model.PType; import io.github.dfa1.vortex.writer.encode.NullableData; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import java.util.stream.LongStream; +import java.util.stream.Stream; import static org.assertj.core.api.Assertions.assertThat; @@ -26,7 +31,7 @@ void smallBatches_coalesceIntoOneMegabyteChunks() { for (int b = 0; b < 5; b++) { chunks.addAll(sut.add(sequence(b * 65_536L, 65_536), 65_536)); } - chunks.add(sut.finish()); + chunks.addAll(sut.finish()); // Then — 2 full chunks, then the odd batch left over at the end, rows in order assertThat(chunks).extracting(c -> ((long[]) c).length).containsExactly(131_072, 131_072, 65_536); @@ -41,7 +46,7 @@ void largeBatch_splitsAtOneMegabyteOfWholeBlocks() { // When List chunks = new ArrayList<>(sut.add(sequence(0, 300_000), 300_000)); - chunks.add(sut.finish()); + chunks.addAll(sut.finish()); // Then assertThat(chunks).extracting(c -> ((long[]) c).length).containsExactly(131_072, 131_072, 37_856); @@ -59,7 +64,7 @@ void unalignedBatches_splitThePieceStraddlingTheBlockBoundary() { for (int b = 0; b < 3; b++) { chunks.addAll(sut.add(sequence(b * 100_000L, 100_000), 100_000)); } - chunks.add(sut.finish()); + chunks.addAll(sut.finish()); // Then assertThat(chunks).extracting(c -> ((long[]) c).length) @@ -76,7 +81,7 @@ void nullableBatches_keepValidityAlignedWithValues() { sut.add(new long[]{4, 5}, 2); // When - Object result = sut.finish(); + Object result = sut.finish().getFirst(); // Then assertThat(result).isInstanceOfSatisfying(NullableData.class, nd -> { @@ -85,22 +90,119 @@ void nullableBatches_keepValidityAlignedWithValues() { }); } + // Each batch's first chunk is the fewest 8192-row blocks reaching 1 MiB by Rust 0.89's exact size + // (vortex-data/vortex#10445). The inlined, all-valid and all-null cases cut elsewhere under the + // rule before it, which counted every string byte and a bitmap on every nullable block. + static Stream exactSizeCases() { + String[] inlined = new String[600_000]; + Arrays.fill(inlined, "twelve bytes"); + String[] outOfLine = new String[600_000]; + Arrays.fill(outOfLine, "thirteen byte"); + boolean[] allValid = new boolean[600_000]; + Arrays.fill(allValid, true); + boolean[] nullPerBlock = allValid.clone(); + for (int i = 0; i < nullPerBlock.length; i += 8192) { + nullPerBlock[i] = false; + } + byte[][] inlinedBinary = new byte[600_000][]; + Arrays.fill(inlinedBinary, new byte[12]); + byte[][] outOfLineBinary = new byte[600_000][]; + Arrays.fill(outOfLineBinary, new byte[13]); + DType nullableI16 = new DType.Primitive(PType.I16, true); + short[] i16 = new short[600_000]; + // 1024 B of bool values a block: the validity's cost decides the cut. All null at 2 B a + // block needs 1023 blocks; 0 B would need 1024, a bitmap 512. + boolean[] bools = new boolean[1024 * 8192 + 1]; + return Stream.of( + // a value of at most 12 bytes lives in its 16-byte view: 16 B/row + Arguments.of("inlined strings", DType.UTF8, inlined, 65_536), + // 16 + 13 = 29 B/row: 5 blocks are 1 187 840 B, 4 fall short + Arguments.of("out-of-line strings", DType.UTF8, outOfLine, 40_960), + Arguments.of("inlined binaries", new DType.Binary(false), inlinedBinary, 65_536), + Arguments.of("out-of-line binaries", new DType.Binary(false), outOfLineBinary, 40_960), + // an all-valid block has no bitmap: 2 B/row + Arguments.of("all valid", nullableI16, new NullableData(i16, allValid), 524_288), + // a block with a null keeps its bitmap: 2.125 B/row, 61 blocks + Arguments.of("a null per block", nullableI16, new NullableData(i16, nullPerBlock), 499_712), + // an all-null block costs 2 bytes, a constant false's scalar + Arguments.of("all null", new DType.Bool(true), new NullableData(bools, bools), 1023 * 8192)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("exactSizeCases") + void blocksAreMeasuredByRustsExactSize(String name, DType dtype, Object batch, int expectedRows) { + // Given + var sut = new Repartitioner(dtype); + + // When + List result = sut.add(batch, rows(batch)); + + // Then + assertThat(result).first().satisfies(chunk -> assertThat(rows(chunk)).isEqualTo(expectedRows)); + } + @Test - void finish_withNothingPending_isNull() { + void finish_withNothingPending_isEmpty() { // Given var sut = new Repartitioner(I64); // When - Object result = sut.finish(); + List result = sut.finish(); + + // Then + assertThat(result).isEmpty(); + } + + @Test + void finish_partialBlockToppingUpToOneMegabyte_isItsOwnChunk() { + // Given — 3 whole blocks of inlined strings pending (384 KiB) and 100 rows of 10 000-byte + // strings: the partial block tips the total over 1 MiB, so Rust's coalescer emits the whole + // blocks, then flushes the rest on its own + var sut = new Repartitioner(DType.UTF8); + String[] blocks = new String[3 * 8192]; + Arrays.fill(blocks, "short"); + String[] tail = new String[100]; + Arrays.fill(tail, "x".repeat(10_000)); + sut.add(blocks, blocks.length); + sut.add(tail, tail.length); + + // When + List result = sut.finish(); + + // Then + assertThat(result).extracting(RepartitionerTest::rows).containsExactly(3 * 8192, 100); + } + + @Test + void blocksAlignToTheColumnsFirstRow_notToEachBatch() { + // Given — 5000-row batches, alternately all null and all valid. Rust re-cuts the stream + // into 8192-row blocks before measuring them, so every block mixes both and carries a + // bitmap (2.125 B/row: 61 blocks reach 1 MiB). Measured per batch, the all-null and + // all-valid pieces would carry none, and 64 blocks would be needed. + var sut = new Repartitioner(new DType.Primitive(PType.I16, true)); + boolean[] allNull = new boolean[5000]; + boolean[] allValid = new boolean[5000]; + Arrays.fill(allValid, true); + List result = new ArrayList<>(); + + // When + for (int b = 0; b < 120; b++) { + result.addAll(sut.add(new NullableData(new short[5000], b % 2 == 0 ? allNull : allValid), 5000)); + } // Then - assertThat(result).isNull(); + assertThat(result).first().satisfies(chunk -> assertThat(rows(chunk)).isEqualTo(61 * 8192)); } private static long[] sequence(long from, int n) { return LongStream.range(from, from + n).toArray(); } + private static int rows(Object chunk) { + Object values = chunk instanceof NullableData nd ? nd.values() : chunk; + return java.lang.reflect.Array.getLength(values); + } + private static long[] concat(List chunks) { return chunks.stream().flatMapToLong(c -> LongStream.of((long[]) c)).toArray(); } @@ -112,10 +214,21 @@ void addWithPresizedPiecesCutsTheSameChunks() { for (int i = 0; i < batch.length; i++) { batch[i] = "x".repeat(10 + i % 40) + (i % 7 == 0 ? "\u20ac" : ""); } - Repartitioner sizer = new Repartitioner(DType.UTF8); + // and a 3000-row batch ahead of it, so its pieces start mid-block + String[] head = new String[3000]; + Arrays.fill(head, "a string longer than twelve bytes"); Repartitioner direct = new Repartitioner(DType.UTF8); Repartitioner presized = new Repartitioner(DType.UTF8); - long[] sizes = sizer.sizePieces(batch, batch.length); + direct.add(head, head.length); + presized.add(head, head.length); + long[] sizes = Repartitioner.sizePieces(batch, batch.length, presized.rows()); + // the first piece only fills the head's block: rows 0..5191 + long firstPiece = 0; + for (int i = 0; i < 8192 - head.length; i++) { + int length = batch[i].getBytes(java.nio.charset.StandardCharsets.UTF_8).length; + firstPiece += length > 12 ? length : 0; + } + assertThat(sizes[0]).isEqualTo(firstPiece); // When List expected = direct.add(batch, batch.length); @@ -124,8 +237,8 @@ void addWithPresizedPiecesCutsTheSameChunks() { // Then assertThat(result).hasSameSizeAs(expected); for (int i = 0; i < result.size(); i++) { - assertThat((String[]) result.get(i)).containsExactly((String[]) expected.get(i)); + assertThat((String[]) result.get(i)).isEqualTo((String[]) expected.get(i)); } - assertThat((String[]) presized.finish()).containsExactly((String[]) direct.finish()); + assertThat(presized.finish().getFirst()).isEqualTo(direct.finish().getFirst()); } }