From 9283daec9f76ddbc2e37f4dbe26c17f67756d18c Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sat, 3 Oct 2026 14:45:47 +0200 Subject: [PATCH] feat(writer): global-dict Binary columns like Rust's dict layout Rust's dict layout admits Primitive | Utf8 | Binary (dict_layout_supported); the Java writer left Binary out. Binary values are keyed by a ByteBuffer view so equal bytes in distinct byte[] instances dedup by content. Adds DictAllTypesInteropIntegrationTest: vortex-jni writes a low-cardinality column of every Arrow type, Java rebuilds every value, and the test pins which types Rust dict-encodes. Rust never dicts uuid/date/FSL/bool/decimal, so compatibility.md downgrades that dict-layout row from a gap to unreachable. Co-Authored-By: Claude Sonnet 5 --- CHANGELOG.md | 1 + docs/compatibility.md | 2 +- .../DictAllTypesInteropIntegrationTest.java | 331 ++++++++++++++++++ .../JavaWritesRustReadsIntegrationTest.java | 63 ++++ .../dfa1/vortex/writer/DictColumnState.java | 98 ++++-- .../dfa1/vortex/writer/VortexWriter.java | 19 +- .../vortex/writer/DictColumnStateTest.java | 53 ++- 7 files changed, 519 insertions(+), 48 deletions(-) create mode 100644 integration/src/test/java/io/github/dfa1/vortex/integration/DictAllTypesInteropIntegrationTest.java diff --git a/CHANGELOG.md b/CHANGELOG.md index f32efe20..6f9ffcf4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Read and write `vortex.onpair` (unstable edition `unstable2026.06.0`): `tpch_orders.regular` and `clickbench_hits_5k.regular` from the v0.86.1 fixtures now scan, and `WriteOptions.withEdition(Editions.UNSTABLE_2026_06_0)` lets the cascade pick OnPair for string columns ([#425](https://github.com/dfa1/vortex-java/issues/425)). ### Changed +- Low-cardinality `Binary` columns are now written as a global dictionary, as Rust does for `Primitive`, `Utf8` and `Binary`; before, only primitive and `Utf8` columns were ([#443](https://github.com/dfa1/vortex-java/pull/443)). - With an unstable edition enabled (`WriteOptions.withEdition(Editions.UNSTABLE_2025_05_0)` or later), cascading writes can pick `fastlanes.delta`, its bases and deltas cascaded as in Rust; default writes never emit it ([#439](https://github.com/dfa1/vortex-java/pull/439)). - The writer run-length encodes float columns too (`fastlanes.rle`, as Rust's float RLE scheme does), losslessly: `-0.0` and NaN payloads round-trip ([#438](https://github.com/dfa1/vortex-java/pull/438)). - Cascading writes (`WriteOptions.cascading(n)`) compress run-end columns further: the run ends and values now go through the cascade (bit-packing, frame-of-reference, …) instead of being stored raw, as in Rust; about 1% smaller on the mixed-column size benchmark ([#435](https://github.com/dfa1/vortex-java/pull/435)). diff --git a/docs/compatibility.md b/docs/compatibility.md index ad56a889..8730c149 100644 --- a/docs/compatibility.md +++ b/docs/compatibility.md @@ -40,7 +40,7 @@ only the built-in decoders in `reader`; no encoder class is loaded. | `vortex.onpair` experimental string encoding | Rust 0.74.0 | ✅ Read and written. The writer emits it only when `WriteOptions.withEdition(Editions.UNSTABLE_2026_06_0)` is set; the trained dictionary is valid for Rust but not byte-identical to Rust's (Rust's sampling RNG is not portable). | | `vortex.variant` arbitrary nested objects | Rust (`vortex.parquet.variant`) | ⚠️ Java encodes/decodes variant columns of **typed scalar** values (constant / chunked-of-constants core, optional shredded child); Java↔Rust round-trip verified. Arbitrary nested JSON objects and real path-based shredding need the `vortex.parquet.variant` physical encoding — deferred ([ADR 0014](../adr/0014-variant-encoding-strategy.md)). | | Arrow extension array import affecting Variant shape | Rust 0.74.0 (#8125) | Untested against the currently pinned v0.85.0 fixtures; #8125 not yet re-verified. | -| `vortex.dict` **layout** over a values pool that is neither VarBin- nor primitive-shaped (e.g. a dict-encoded `vortex.uuid`, whose storage is `FixedSizeList(U8, 16)`) | Rust's dict layout accepts any dtype | ❌ No lazy dict carrier exists for that pool shape, so decode throws `VortexException("unsupported dict values shape: …")`. The `vortex.dict` *encoding* is unaffected. | +| `vortex.dict` **layout** over a values pool that is neither VarBin- nor primitive-shaped (e.g. a dict-encoded `vortex.uuid`, whose storage is `FixedSizeList(U8, 16)`) | Not written by Rust: its dict layout admits only `Primitive \| Utf8 \| Binary` (`dict_layout_supported`), and its dict compressor schemes only integers, floats and strings | ⚠️ Unreachable from Rust- or Java-written files. Every type vortex-jni writes reads back exactly, dict-encoded or not (`DictAllTypesInteropIntegrationTest`, which also fails if a Rust bump starts dict-encoding another type). A foreign file with such a pool fails with `VortexException("unsupported dict values shape: …")`. | | Duplicate struct field names | Rust writer rejects ("StructLayout must have unique field names"); Rust reader tolerates foreign files (first-match access) | ⚠️ Deliberate divergence on read: Java rejects such files with `VortexException("duplicate field name in file schema")` instead of tolerating them — the name-keyed `Chunk` API cannot represent both columns, and silent column loss is worse than a loud failure on a file the reference writer refuses to produce. Java's writer mirrors the Rust writer's rejection. | | Blank / control-character field names | Wire-legal; the Rust writer produces `""` and whitespace-only names. NUL (`U+0000`) additionally aborts the Rust toolchain: Arrow FFI schema export hits a panic-cannot-unwind in `arrow-rs` (`ffi_stream::get_schema`) and SIGABRTs the process (measured against vortex-jni 0.75.0) | ⚠️ Deliberate strictness BOTH ways: vortex-java's writer refuses blank and control-character field names (`IllegalArgumentException`), and its reader rejects files carrying them (`VortexException` naming the producing pipeline as the likely bug) — the JSON-`""`-key principle: wire-legal is a floor, not a policy. Printable names of any shape (`$`-runs, spaces inside, emoji) are legal and round-trip intact both directions (measured; pinned by `ColumnNameEdgeCasesIntegrationTest`). | diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/DictAllTypesInteropIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/DictAllTypesInteropIntegrationTest.java new file mode 100644 index 00000000..22a47ee8 --- /dev/null +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/DictAllTypesInteropIntegrationTest.java @@ -0,0 +1,331 @@ +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.csv.CsvExporter; +import io.github.dfa1.vortex.csv.ExportOptions; +import io.github.dfa1.vortex.inspect.InspectorTree; +import io.github.dfa1.vortex.reader.ReadRegistry; +import io.github.dfa1.vortex.reader.VortexReader; +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.BigIntVector; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.DateDayVector; +import org.apache.arrow.vector.DecimalVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.SmallIntVector; +import org.apache.arrow.vector.TimeMicroVector; +import org.apache.arrow.vector.TimeStampMicroVector; +import org.apache.arrow.vector.TinyIntVector; +import org.apache.arrow.vector.UInt1Vector; +import org.apache.arrow.vector.UInt2Vector; +import org.apache.arrow.vector.UInt4Vector; +import org.apache.arrow.vector.UInt8Vector; +import org.apache.arrow.vector.UuidVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.FixedSizeListVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.StructVector; +import org.apache.arrow.vector.extension.UuidType; +import org.apache.arrow.vector.types.DateUnit; +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.TimeUnit; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + +import java.io.IOException; +import java.io.StringReader; +import java.io.StringWriter; +import java.math.BigDecimal; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; +import java.util.UUID; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; + +/// vortex-jni writes a low-cardinality column of every Arrow type it accepts, dict-encoding the ones +/// Rust supports; vortex-java must read each one back and rebuild the exact values. +/// +/// The Java reader's `vortex.dict` support grew per type as bugs surfaced (Utf8, then primitives, +/// then I16, then VarBin extensions), leaving shapes like dict-encoded `vortex.uuid` (a +/// FixedSizeList pool) or `vortex.date` (an extension over a primitive pool) rejected with +/// "unsupported dict values shape". This sweeps every type at once so a missing pool shape shows +/// up as one failing case instead of a Raincloud surprise. +/// +/// Values are compared through [CsvExporter]'s per-type rendering so one generic comparison +/// covers primitive, nested and extension columns alike (extensions render their storage value). +class DictAllTypesInteropIntegrationTest { + + private static final Session SESSION = Session.create(); + private static final BufferAllocator ALLOCATOR = ArrowAllocation.rootAllocator(); + /// Far more rows than distinct values so the Rust compressor's dict scheme wins wherever it applies. + private static final int ROWS = 20_000; + + static { + NativeLoader.loadJni(); + } + + /// One Arrow column type: how to write distinct value `j` at row `i`, how each distinct value + /// renders after the Java read (an empty string is a null row), and whether Rust dict-encodes + /// it at all. + /// + /// Rust's dict layout only admits `Primitive | Utf8 | Binary` (`dict_layout_supported` in + /// `vortex-layout/src/layouts/dict/writer.rs`) and its compressor's dict schemes are + /// integer/float/string only, so the other types never reach a dictionary. The flag is asserted + /// both ways: a vortex-jni bump that starts dict-encoding, say, dates fails here and points at + /// the reader's still-missing dict-over-extension pool path. + private record Case(String name, Field field, Setter setter, List expected, boolean rustDicts) { + + Case notDictedByRust() { + return new Case(name, field, setter, expected, false); + } + + @Override + public String toString() { + return name; + } + } + + @FunctionalInterface + private interface Setter { + void set(FieldVector vector, int row, int distinct); + } + + private static final UUID[] UUIDS = { + UUID.fromString("00000000-0000-0000-0000-000000000001"), + UUID.fromString("123e4567-e89b-12d3-a456-426614174000"), + UUID.fromString("ffffffff-ffff-ffff-ffff-ffffffffffff"), + UUID.fromString("a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a11"), + }; + + static Stream cases() { + return Stream.of( + primitive("i8", new ArrowType.Int(8, true), + (v, i, j) -> ((TinyIntVector) v).setSafe(i, new byte[]{-128, 0, 7, 127}[j]), + "-128", "0", "7", "127"), + primitive("u8", new ArrowType.Int(8, false), + (v, i, j) -> ((UInt1Vector) v).setSafe(i, new int[]{0, 1, 128, 255}[j]), + "0", "1", "128", "255"), + primitive("i16", new ArrowType.Int(16, true), + (v, i, j) -> ((SmallIntVector) v).setSafe(i, new short[]{-32768, 0, 300, 32767}[j]), + "-32768", "0", "300", "32767"), + primitive("u16", new ArrowType.Int(16, false), + (v, i, j) -> ((UInt2Vector) v).setSafe(i, new int[]{0, 1, 40_000, 65_535}[j]), + "0", "1", "40000", "65535"), + primitive("i32", new ArrowType.Int(32, true), + (v, i, j) -> ((IntVector) v).setSafe(i, new int[]{Integer.MIN_VALUE, 0, 42, Integer.MAX_VALUE}[j]), + "-2147483648", "0", "42", "2147483647"), + primitive("u32", new ArrowType.Int(32, false), + (v, i, j) -> ((UInt4Vector) v).setSafe(i, new int[]{0, 1, 0x8000_0000, -1}[j]), + "0", "1", "2147483648", "4294967295"), + primitive("i64", new ArrowType.Int(64, true), + (v, i, j) -> ((BigIntVector) v).setSafe(i, new long[]{Long.MIN_VALUE, 0, 1L << 40, Long.MAX_VALUE}[j]), + "-9223372036854775808", "0", "1099511627776", "9223372036854775807"), + primitive("u64", new ArrowType.Int(64, false), + (v, i, j) -> ((UInt8Vector) v).setSafe(i, new long[]{0, 1, Long.MIN_VALUE, -1}[j]), + "0", "1", "9223372036854775808", "18446744073709551615"), + primitive("f32", new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE), + (v, i, j) -> ((Float4Vector) v).setSafe(i, new float[]{-1.5f, 0f, 3.25f, 1e30f}[j]), + "-1.5", "0.0", "3.25", "1.0E30"), + primitive("f64", new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE), + (v, i, j) -> ((Float8Vector) v).setSafe(i, new double[]{-1.1, 0.0, 2.2, 1e300}[j]), + "-1.1", "0.0", "2.2", "1.0E300"), + primitive("bool", ArrowType.Bool.INSTANCE, + (v, i, j) -> ((BitVector) v).setSafe(i, j & 1), + "false", "true", "false", "true").notDictedByRust(), + primitive("utf8", ArrowType.Utf8.INSTANCE, + (v, i, j) -> ((VarCharVector) v).setSafe(i, utf8(new String[]{"", "alpha", "βeta", "gamma gamma"}[j])), + "", "alpha", "βeta", "gamma gamma"), + primitive("binary", ArrowType.Binary.INSTANCE, + (v, i, j) -> ((VarBinaryVector) v).setSafe(i, utf8(new String[]{"x", "yy", "zzz", "wwww"}[j])), + "x", "yy", "zzz", "wwww"), + primitive("decimal(10,2)", new ArrowType.Decimal(10, 2, 128), + (v, i, j) -> ((DecimalVector) v).setSafe(i, new BigDecimal(new String[]{"-99999999.99", "0.00", "12.34", "99999999.99"}[j])), + "-99999999.99", "0.00", "12.34", "99999999.99").notDictedByRust(), + // Extensions render their storage value: days / micros since epoch / midnight. + primitive("date", new ArrowType.Date(DateUnit.DAY), + (v, i, j) -> ((DateDayVector) v).setSafe(i, new int[]{-1, 0, 19_000, 20_000}[j]), + "-1", "0", "19000", "20000").notDictedByRust(), + primitive("time", new ArrowType.Time(TimeUnit.MICROSECOND, 64), + (v, i, j) -> ((TimeMicroVector) v).setSafe(i, new long[]{0, 1, 3_600_000_000L, 86_399_999_999L}[j]), + "0", "1", "3600000000", "86399999999").notDictedByRust(), + primitive("timestamp", new ArrowType.Timestamp(TimeUnit.MICROSECOND, null), + (v, i, j) -> ((TimeStampMicroVector) v).setSafe(i, new long[]{-1, 0, 1_700_000_000_000_000L, 1_800_000_000_000_000L}[j]), + "-1", "0", "1700000000000000", "1800000000000000").notDictedByRust(), + primitive("uuid", UuidType.INSTANCE, + (v, i, j) -> ((UuidVector) v).setSafe(i, UUIDS[j]), + uuidCells()).notDictedByRust(), + nested("fixed_size_list", + new Field("v", FieldType.notNullable(new ArrowType.FixedSizeList(2)), + List.of(Field.notNullable("item", new ArrowType.Int(32, true)))), + (v, i, j) -> { + var list = (FixedSizeListVector) v; + var items = (IntVector) list.getDataVector(); + list.setNotNull(i); + items.setSafe(i * 2, j); + items.setSafe(i * 2 + 1, -j); + }, + "[0,0]", "[1,-1]", "[2,-2]", "[3,-3]").notDictedByRust(), + nested("list", + new Field("v", FieldType.notNullable(ArrowType.List.INSTANCE), + List.of(Field.notNullable("item", new ArrowType.Int(32, true)))), + (v, i, j) -> { + var list = (ListVector) v; + var items = (IntVector) list.getDataVector(); + int start = list.startNewValue(i); + for (int k = 0; k < j; k++) { + items.setSafe(start + k, k); + } + list.endValue(i, j); + }, + "[]", "[0]", "[0,1]", "[0,1,2]").notDictedByRust(), + nested("struct{a:i32}", + new Field("v", FieldType.notNullable(ArrowType.Struct.INSTANCE), + List.of(Field.notNullable("a", new ArrowType.Int(32, true)))), + (v, i, j) -> { + var struct = (StructVector) v; + struct.setIndexDefined(i); + ((IntVector) struct.getChild("a")).setSafe(i, j * 10); + }, + "{\"a\":0}", "{\"a\":10}", "{\"a\":20}", "{\"a\":30}"), + // Nullable: a null row rides on the codes, the pool itself holds no nulls. + nested("nullable i32", + Field.nullable("v", new ArrowType.Int(32, true)), + (v, i, j) -> { + if (j == 0) { + ((IntVector) v).setNull(i); + } else { + ((IntVector) v).setSafe(i, j); + } + }, + "", "1", "2", "3"), + nested("nullable utf8", + Field.nullable("v", ArrowType.Utf8.INSTANCE), + (v, i, j) -> { + if (j == 0) { + ((VarCharVector) v).setNull(i); + } else { + ((VarCharVector) v).setSafe(i, utf8("s" + j)); + } + }, + "", "s1", "s2", "s3")); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("cases") + void jniDictColumn_javaRebuildsValues(Case c, @TempDir Path tmp) throws IOException { + // Given — ROWS rows cycling through 4 distinct values, written by vortex-jni + Path file = tmp.resolve("dict.vortex"); + writeJni(file, c); + List expected = new ArrayList<>(ROWS); + for (int i = 0; i < ROWS; i++) { + expected.add(c.expected().get(i % c.expected().size())); + } + + // When + List result = readColumnAsCsv(file); + + // Then — Rust dict-encoded exactly the types it supports, and every row matches + try (var vf = VortexReader.open(file, ReadRegistry.loadAll())) { + InspectorTree tree = InspectorTree.build(vf); + assertThat(hasDictLayout(tree.root()) || tree.usedEncodings().contains("vortex.dict")) + .as("vortex-jni dict-encoded %s (encodings %s)", c.field().getType(), tree.usedEncodings()) + .isEqualTo(c.rustDicts()); + } + assertThat(result).isEqualTo(expected); + } + + private static void writeJni(Path file, Case c) throws IOException { + Schema schema = new Schema(List.of(c.field())); + String uri = file.toAbsolutePath().toUri().toString(); + try (VortexWriter writer = VortexWriter.builder(SESSION, uri, schema, ALLOCATOR).build(); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, ALLOCATOR)) { + FieldVector vector = root.getVector("v"); + vector.setInitialCapacity(ROWS); + vector.allocateNew(); + int distinct = c.expected().size(); + for (int i = 0; i < ROWS; i++) { + c.setter().set(vector, i, i % distinct); + } + vector.setValueCount(ROWS); + root.setRowCount(ROWS); + try (ArrowArray arr = ArrowArray.allocateNew(ALLOCATOR); + ArrowSchema arrowSchema = ArrowSchema.allocateNew(ALLOCATOR)) { + Data.exportVectorSchemaRoot(ALLOCATOR, root, null, arr, arrowSchema); + writer.writeBatch(arr.memoryAddress(), arrowSchema.memoryAddress()); + } + } + } + + private static List readColumnAsCsv(Path file) throws IOException { + var csv = new StringWriter(); + CsvExporter.exportCsv(file, csv, ExportOptions.defaults()); + var out = new ArrayList(ROWS); + try (var reader = de.siegmar.fastcsv.reader.CsvReader.builder().skipEmptyLines(false).ofCsvRecord(new StringReader(csv.toString()))) { + var rows = reader.iterator(); + rows.next(); // header + rows.forEachRemaining(row -> out.add(row.getField(0))); + } + return out; + } + + private static boolean hasDictLayout(InspectorTree.Node node) { + if (node.layout().isDict()) { + return true; + } + for (InspectorTree.Node child : node.children()) { + if (hasDictLayout(child)) { + return true; + } + } + return false; + } + + private static Case primitive(String name, ArrowType type, Setter setter, String... expected) { + return nested(name, Field.notNullable("v", type), setter, expected); + } + + private static Case nested(String name, Field field, Setter setter, String... expected) { + return new Case(name, field, setter, List.of(expected), true); + } + + private static byte[] utf8(String s) { + return s.getBytes(StandardCharsets.UTF_8); + } + + /// `vortex.uuid` stores 16 big-endian bytes as a FixedSizeList of u8, which the CSV exporter + /// renders as a JSON array of the unsigned byte values. + private static String[] uuidCells() { + String[] cells = new String[UUIDS.length]; + for (int u = 0; u < UUIDS.length; u++) { + var sb = new StringBuilder("["); + long[] halves = {UUIDS[u].getMostSignificantBits(), UUIDS[u].getLeastSignificantBits()}; + for (int b = 0; b < 16; b++) { + if (b > 0) { + sb.append(','); + } + sb.append((halves[b / 8] >>> (56 - 8 * (b % 8))) & 0xFF); + } + cells[u] = sb.append(']').toString(); + } + return cells; + } +} diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/JavaWritesRustReadsIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/JavaWritesRustReadsIntegrationTest.java index 884cb090..ca017de6 100644 --- a/integration/src/test/java/io/github/dfa1/vortex/integration/JavaWritesRustReadsIntegrationTest.java +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/JavaWritesRustReadsIntegrationTest.java @@ -1882,6 +1882,69 @@ void javaWriter_rustReader_globalDict_f64(@TempDir Path tmp) throws IOException assertThat(decoded).containsExactly(all); } + @Test + void javaWriter_rustReader_globalDict_nullableBinary(@TempDir Path tmp) throws IOException { + // Given — a low-cardinality nullable Binary column over two chunks. Rust's dict layout admits + // Primitive | Utf8 | Binary; the Java writer left Binary out, so this column was never + // dict-encoded. Equal values arrive as distinct byte[] instances, so the dedup must compare + // content; the empty value and nulls sit on the edges of the VarBin pool and the codes mask. + Path file = tmp.resolve("java_globaldict_binary.vtx"); + ColumnName id = ColumnName.of("id"); + ColumnName b = ColumnName.of("b"); + DType.Struct schema = new DType.Struct(List.of(id, b), List.of(DType.I64, new DType.Binary(true)), false); + int rowsPerChunk = 1_000; + List expected = new ArrayList<>(); + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var sut = VortexWriter.create(ch, schema, WriteOptions.defaults())) { + // When + for (int c = 0; c < 2; c++) { + long[] ids = new long[rowsPerChunk]; + byte[][] values = new byte[rowsPerChunk][]; + for (int i = 0; i < rowsPerChunk; i++) { + int row = c * rowsPerChunk + i; + ids[i] = row; + values[i] = switch (row % 4) { + case 0 -> null; + case 1 -> new byte[0]; + case 2 -> new byte[]{1, 2, 3}; + default -> new byte[]{(byte) 0xFF}; + }; + expected.add(values[i]); + } + sut.writeChunk(Map.of(id, ids, b, values)); + } + } + + // Then — the column is a global dict layout, and vortex-jni reads every row back, full scan + // and filtered (zone pruning over the dict column's per-chunk stats) + try (var vf = io.github.dfa1.vortex.reader.VortexReader.open(file, + io.github.dfa1.vortex.reader.ReadRegistry.loadAll())) { + var tree = io.github.dfa1.vortex.inspect.InspectorTree.build(vf); + assertThat(tree.root().children()) + .filteredOn(n -> n.fieldName().equals(java.util.Optional.of("b"))) + .singleElement() + .satisfies(n -> assertThat(hasDictLayout(n)).as("Binary column is a dict layout").isTrue()); + } + List full = readColumnFiltered(file, "b", Expression.binary(Expression.BinaryOp.GTE, + Expression.column("id"), Expression.literal(0L))); + assertThat(full).containsExactlyElementsOf(expected); + List filtered = readColumnFiltered(file, "b", Expression.binary(Expression.BinaryOp.GTE, + Expression.column("id"), Expression.literal(1_500L))); + assertThat(filtered).containsExactlyElementsOf(expected.subList(1_500, expected.size())); + } + + private static boolean hasDictLayout(io.github.dfa1.vortex.inspect.InspectorTree.Node node) { + if (node.layout().isDict()) { + return true; + } + for (var child : node.children()) { + if (hasDictLayout(child)) { + return true; + } + } + return false; + } + @Test void javaWriter_rustReader_nullable_date(@TempDir Path tmp) throws IOException { // Given — nullable vortex.date column; row 1 is null. diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/DictColumnState.java b/writer/src/main/java/io/github/dfa1/vortex/writer/DictColumnState.java index e4395776..8f1b3f1f 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/DictColumnState.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/DictColumnState.java @@ -7,6 +7,7 @@ import io.github.dfa1.vortex.writer.encode.PrimitiveEncodingEncoder; import io.github.dfa1.vortex.writer.encode.VarBinEncodingEncoder; +import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.HashSet; import java.util.LinkedHashMap; @@ -27,7 +28,7 @@ final class DictColumnState { // Columns with global cardinality below this threshold are dict-encoded across all chunks. // The cap is type-aware. Numeric stays low: a global dict hurts high-cardinality F64/I64 - // columns (ALP/bitpacked codes beat U16 dict codes). Utf8 is raised far higher — text columns + // columns (ALP/bitpacked codes beat U16 dict codes). Utf8/Binary is raised far higher — text columns // with thousands of repeated distinct values (street/place names) dictionary-compress well // (#299), and the per-chunk short[] code buffer holds codes 0..32767 (up to 32768 distinct) // with no wider buffer; codePTypeForSize already emits U16 codes above 256. @@ -37,15 +38,18 @@ final class DictColumnState { private static final int INDEX_MIN_CAPACITY = 64; private final DType dtype; - private final boolean utf8; + // Utf8 or Binary: keyed by value (String, or a ByteBuffer wrapping the bytes, whose + // equals/hashCode compare content), coded in first-seen order with no frequency remap. + private final boolean varBin; + private final boolean binary; private final PType ptype; private final boolean nullable; - // First-seen value -> code map (keys are boxed primitives or String, matching readPrimitiveElement). + // First-seen value -> code map (keys are boxed primitives, String, or ByteBuffer — see varBinKey). private final Map valueToCode = new LinkedHashMap<>(); // Hot-path side index for the primitive path: raw value bits -> code + 1 (0 == empty), open // addressing with a power-of-two capacity so probing masks instead of taking a modulo (CLAUDE.md // hot-loop rule). valueToCode stays the authoritative store — it carries first-seen order, the - // Utf8 keys, and everything the demotion and flush paths read — but probing it needs a boxed key, + // VarBin keys, and everything the demotion and flush paths read — but probing it needs a boxed key, // and one Long per row of every candidate column profiled as the writer's hottest single frame. private LongIntMap bitsIndex = new LongIntMap(INDEX_MIN_CAPACITY); // Occurrence count per code, indexed by code; grows in lockstep with valueToCode. A primitive @@ -64,7 +68,8 @@ final class DictColumnState { DictColumnState(DType dtype) { this.dtype = dtype; - this.utf8 = dtype instanceof DType.Utf8; + this.binary = dtype instanceof DType.Binary; + this.varBin = binary || dtype instanceof DType.Utf8; this.ptype = dtype instanceof DType.Primitive p ? p.ptype() : null; this.nullable = dtype.nullable(); } @@ -73,8 +78,8 @@ DType dtype() { return dtype; } - boolean utf8() { - return utf8; + boolean varBin() { + return varBin; } PType ptype() { @@ -130,9 +135,18 @@ List chunkStatsSum() { return chunkStatsSum; } - /// The distinct Utf8 values seen so far, in first-seen order. Only valid when [#utf8()]. - String[] utf8Uniques() { - return valueToCode.keySet().toArray(new String[0]); + /// The distinct Utf8/Binary values seen so far, in first-seen order: a `String[]` for Utf8, a + /// `byte[][]` for Binary. Only valid when [#varBin()]. + Object varBinUniques() { + if (!binary) { + return valueToCode.keySet().toArray(new String[0]); + } + byte[][] out = new byte[valueToCode.size()][]; + int i = 0; + for (Object key : valueToCode.keySet()) { + out[i++] = ((ByteBuffer) key).array(); + } + return out; } /// Ingests one chunk into this candidate column's cardinality-bounded dict state (ADR 0021): dedups @@ -145,17 +159,17 @@ String[] utf8Uniques() { /// encoding. This moves the cap check from `close()` to a continuous, mid-file guard so a column /// whose distinct set grows past the cap never accumulates unbounded memory first. /// - /// @param data the chunk data (primitive array, `String[]`, or a [NullableData] wrapper) + /// @param data the chunk data (primitive array, `String[]`, `byte[][]`, or a [NullableData] wrapper) /// @return `true` if the chunk was ingested within the cardinality cap; `false` if the column /// must be demoted boolean ingestDictChunk(Object data) { boolean nullableData = data instanceof NullableData; Object values = nullableData ? ((NullableData) data).values() : data; boolean[] validity = nullableData ? ((NullableData) data).validity() : null; - int len = utf8 ? ((String[]) values).length : primitiveArrayLen(values, ptype); - int cap = dictMaxCardinality(utf8); + int len = varBin ? ((Object[]) values).length : primitiveArrayLen(values, ptype); + int cap = dictMaxCardinality(varBin); int startSize = valueToCode.size(); - String[] strings = utf8 ? (String[]) values : null; + Object[] strings = varBin ? (Object[]) values : null; // One pass: insert new values and build the per-chunk code array. Ingest stays // all-or-nothing — a chunk that would breach the cap rolls back the entries it added (the @@ -168,12 +182,12 @@ boolean ingestDictChunk(Object data) { if (validity != null && !validity[i]) { continue; } - // Nullable Utf8 keeps a real null at invalid positions (ChunkImpl.adaptUtf8); treat - // it as a null slot (code 0), never as a dictionary entry. - String v = strings[i]; - if (v == null) { + // Nullable Utf8/Binary keeps a real null at invalid positions (ChunkImpl.adaptUtf8, + // adaptBinary); treat it as a null slot (code 0), never as a dictionary entry. + if (strings[i] == null) { continue; } + Object v = varBinKey(strings[i]); Integer code = valueToCode.get(v); if (code == null) { if (valueToCode.size() == cap) { @@ -226,7 +240,12 @@ boolean ingestDictChunk(Object data) { chunkValidity.add(validity); chunkRowCounts.add((long) len); chunkNullCounts.add(validity != null ? VortexWriter.countNulls(validity) : 0L); - if (utf8) { + if (binary) { + // No min/max for Binary, matching the per-chunk VarBin path (VarBinEncodingEncoder). + chunkStatsMin.add(null); + chunkStatsMax.add(null); + chunkStatsSum.add(null); + } else if (varBin) { byte[][] mm = VarBinEncodingEncoder.minMaxStats((String[]) values); chunkStatsMin.add(mm != null ? mm[0] : null); chunkStatsMax.add(mm != null ? mm[1] : null); @@ -246,7 +265,7 @@ boolean ingestDictChunk(Object data) { /// through [#buildInverseMap] and [#reconstructChunk], which must not see this chunk's values. private void rollbackTo(int startSize) { valueToCode.values().removeIf(code -> code >= startSize); - if (utf8) { + if (varBin) { return; } bitsIndex = new LongIntMap(INDEX_MIN_CAPACITY); @@ -308,7 +327,7 @@ Object[] buildInverseMap() { return inverse; } - /// Reconstructs demoted chunk `c`'s raw array (a typed primitive array or `String[]`, wrapped in + /// Reconstructs demoted chunk `c`'s raw array (a typed primitive array, `String[]` or `byte[][]`, wrapped in /// [NullableData] when the chunk carried validity) from its buffered `short[]` codes and the /// inverse code-to-value map. Null slots restore a zero/`null` placeholder — exactly what the /// per-chunk encoders expect from [NullableData]. @@ -317,7 +336,15 @@ Object reconstructChunk(Object[] inverse, int c) { boolean[] validity = chunkValidity.get(c); int len = codes.length; Object values; - if (utf8) { + if (binary) { + byte[][] arr = new byte[len][]; + for (int i = 0; i < len; i++) { + if (validity == null || validity[i]) { + arr[i] = ((ByteBuffer) inverse[codes[i] & 0xFFFF]).array(); + } + } + values = arr; + } else if (varBin) { String[] arr = new String[len]; for (int i = 0; i < len; i++) { if (validity == null || validity[i]) { @@ -465,28 +492,29 @@ static Object emitCodes(short[] buffered, int[] remap, boolean[] validity, PType }; } - static boolean isUtf8DictCandidate(String[] data) { - return isUtf8DictCandidate(data, null); + static boolean isVarBinDictCandidate(Object[] data) { + return isVarBinDictCandidate(data, null); } - /// Like [#isUtf8DictCandidate(String[])] but ignores null (invalid) rows when counting distinct + /// Like [#isVarBinDictCandidate(Object[])] but ignores null (invalid) rows when counting distinct /// values, so a nullable low-cardinality column still qualifies for the shared global dictionary. /// The ratio denominator stays the total row count (not the valid-row count), matching the /// per-chunk encoders' convention that null placeholders occupy a row like any other value. /// - /// @param data the string values; null elements at invalid positions are skipped + /// @param data the Utf8 (`String[]`) or Binary (`byte[][]`) values; null elements at invalid + /// positions are skipped /// @param validity per-row validity bitmap, or `null` meaning every row is valid /// @return `true` if the column's distinct valid-value count is low enough to dictionary-encode - static boolean isUtf8DictCandidate(String[] data, boolean[] validity) { + static boolean isVarBinDictCandidate(Object[] data, boolean[] validity) { if (data.length == 0) { return false; } - var seen = HashSet.newHashSet(Math.min(GLOBAL_DICT_MAX_CARDINALITY_UTF8, data.length)); + var seen = HashSet.newHashSet(Math.min(GLOBAL_DICT_MAX_CARDINALITY_UTF8, data.length)); for (int i = 0; i < data.length; i++) { if ((validity != null && !validity[i]) || data[i] == null) { continue; } - seen.add(data[i]); + seen.add(varBinKey(data[i])); if (seen.size() > GLOBAL_DICT_MAX_CARDINALITY_UTF8) { return false; } @@ -605,8 +633,14 @@ static PType codePTypeForSize(int dictSize) { return PType.U32; } - // The global-dict cardinality cap for a column, by whether it is Utf8 (see the constants above). - private static int dictMaxCardinality(boolean utf8) { - return utf8 ? GLOBAL_DICT_MAX_CARDINALITY_UTF8 : GLOBAL_DICT_MAX_CARDINALITY; + // The global-dict cardinality cap for a column, by whether it is Utf8/Binary (see the constants above). + private static int dictMaxCardinality(boolean varBin) { + return varBin ? GLOBAL_DICT_MAX_CARDINALITY_UTF8 : GLOBAL_DICT_MAX_CARDINALITY; + } + + // A byte[] compares by identity, so Binary values are keyed by a ByteBuffer view, whose + // equals/hashCode compare content. The array is never mutated after ChunkImpl hands it over. + private static Object varBinKey(Object value) { + return value instanceof byte[] bytes ? ByteBuffer.wrap(bytes) : value; } } 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 d7235bd2..87de70c1 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 @@ -512,7 +512,7 @@ public void writeChunk(Map columns) throws IOException { } if (!firstChunkSeen && options.globalDict()) { - // Global dict candidate detection inspects raw primitive/String arrays. Nullable + // Global dict candidate detection inspects raw primitive/String/byte[] arrays. Nullable // columns (carried as NullableData) run the same cardinality/ratio check against // their values, skipping null positions per the validity bitmap; the reader's dict // lazy-decode already handles masked (nullable) codes children. @@ -524,8 +524,9 @@ public void writeChunk(Map columns) throws IOException { boolean candidate = false; if (colDtype instanceof DType.Primitive p) { candidate = DictColumnState.isDictCandidate(p.ptype(), values, validity); - } else if (colDtype instanceof DType.Utf8) { - candidate = DictColumnState.isUtf8DictCandidate((String[]) values, validity); + } else if (colDtype instanceof DType.Utf8 || colDtype instanceof DType.Binary) { + // Rust's dict layout admits Primitive | Utf8 | Binary (dict_layout_supported). + candidate = DictColumnState.isVarBinDictCandidate((Object[]) values, validity); } if (candidate) { dictCandidates.add(colName); @@ -1194,8 +1195,8 @@ private void flushDictColumns() throws IOException { if (state == null || state.chunkCount() == 0 || state.cardinality() == 0) { continue; } - if (state.utf8()) { - writeGlobalDictUtf8Column(colName, state); + if (state.varBin()) { + writeGlobalDictVarBinColumn(colName, state); } else { writeGlobalDictColumn(colName, state); } @@ -1234,17 +1235,17 @@ private void writeGlobalDictColumn(ColumnName colName, DictColumnState state) th state.chunkStatsMin(), state.chunkStatsMax(), state.chunkStatsSum())); } - private void writeGlobalDictUtf8Column(ColumnName colName, DictColumnState state) throws IOException { + private void writeGlobalDictVarBinColumn(ColumnName colName, DictColumnState state) throws IOException { int dictSize = state.cardinality(); PType codePType = DictColumnState.codePTypeForSize(dictSize); - // Utf8 assigns codes in first-seen order with no frequency sort, so the incremental map's + // Utf8/Binary assigns codes in first-seen order with no frequency sort, so the incremental map's // order already matches — no remap pass (ADR 0021). Compress the distinct-values pool - // through the normal Utf8 competition (FSST/VarBin/Zstd) so it captures substring + // through the normal Utf8/Binary competition (FSST/VarBin/Zstd) so it captures substring // redundancy across dictionary entries (#299), but exclude Dict so the cascade never wraps // the (all-unique-by-construction) dictionary in another dict the reader cannot unwrap. At // cascade depth 0 there is no competition to run, so force flat VarBin as before. - String[] uniques = state.utf8Uniques(); + Object uniques = state.varBinUniques(); int valuesSegIdx = options.allowedCascading() > 0 ? writeSegment(state.dtype(), uniques, null, Set.of(EncodingId.VORTEX_DICT)) : writeSegment(state.dtype(), uniques, new VarBinEncodingEncoder()); diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/DictColumnStateTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/DictColumnStateTest.java index 7b28bcf6..4a2d8ea8 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/DictColumnStateTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/DictColumnStateTest.java @@ -1,6 +1,9 @@ package io.github.dfa1.vortex.writer; +import io.github.dfa1.vortex.core.model.DType; 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; @@ -90,7 +93,7 @@ void isDictCandidate_nullable(String name, long[] data, boolean[] validity, bool assertThat(result).isEqualTo(expected); } - // ── isUtf8DictCandidate ────────────────────────────────────────────────────── + // ── isVarBinDictCandidate ────────────────────────────────────────────────────── static Stream utf8DictCandidateCases() { return Stream.of( @@ -107,17 +110,17 @@ static Stream utf8DictCandidateCases() { @ParameterizedTest(name = "{0}") @MethodSource("utf8DictCandidateCases") - void isUtf8DictCandidate(String name, String[] data, boolean expected) { + void isVarBinDictCandidate(String name, String[] data, boolean expected) { // Given — a string column with the case's data // When - boolean result = DictColumnState.isUtf8DictCandidate(data); + boolean result = DictColumnState.isVarBinDictCandidate(data); // Then assertThat(result).isEqualTo(expected); } - // ── isUtf8DictCandidate (nullable) ─────────────────────────────────────────── + // ── isVarBinDictCandidate (nullable) ─────────────────────────────────────────── static Stream nullableUtf8DictCandidateCases() { // Nullable Utf8 keeps real null array elements at invalid positions (ChunkImpl.adaptUtf8), @@ -138,16 +141,54 @@ static Stream nullableUtf8DictCandidateCases() { @ParameterizedTest(name = "{0}") @MethodSource("nullableUtf8DictCandidateCases") - void isUtf8DictCandidate_nullable(String name, String[] data, boolean[] validity, boolean expected) { + void isVarBinDictCandidate_nullable(String name, String[] data, boolean[] validity, boolean expected) { // Given — a string column with the case's values and validity // When - boolean result = DictColumnState.isUtf8DictCandidate(data, validity); + boolean result = DictColumnState.isVarBinDictCandidate(data, validity); // Then assertThat(result).isEqualTo(expected); } + @Test + void isVarBinDictCandidate_binaryComparesContentNotIdentity() { + // Given — 5 rows of 2 distinct byte strings, each row its own byte[] instance. Keyed by + // array identity (byte[] has no content equals) every row would look distinct and the + // column would never qualify. + byte[][] data = new byte[5][]; + for (int i = 0; i < data.length; i++) { + data[i] = new byte[]{(byte) (i % 2), 42}; + } + + // When + boolean result = DictColumnState.isVarBinDictCandidate(data); + + // Then + assertThat(result).isTrue(); + } + + @Test + void ingestDictChunk_binary_dedupsByContentAndReconstructs() { + // Given — a nullable Binary column whose equal values arrive as distinct byte[] instances + var sut = new DictColumnState(new DType.Binary(true)); + byte[][] values = {{1, 2}, {3}, null, {1, 2}, {3}}; + boolean[] validity = {true, true, false, true, true}; + + // When + boolean admitted = sut.ingestDictChunk(new NullableData(values, validity)); + Object result = sut.reconstructChunk(sut.buildInverseMap(), 0); + + // Then — two dictionary entries, and demotion rebuilds the original rows (null kept null) + assertThat(admitted).isTrue(); + assertThat(sut.cardinality()).isEqualTo(2); + assertThat((byte[][]) sut.varBinUniques()).isDeepEqualTo(new byte[][]{{1, 2}, {3}}); + assertThat(result).isInstanceOfSatisfying(NullableData.class, nd -> { + assertThat((byte[][]) nd.values()).isDeepEqualTo(values); + assertThat(nd.validity()).containsExactly(validity); + }); + } + // ── codePTypeForSize ───────────────────────────────────────────────────────── static Stream codePTypeCases() {