diff --git a/CHANGELOG.md b/CHANGELOG.md index 39f97d23..801268ab 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)) +- A non-null integer or float column gets a dict layout only when the compressor picks a dictionary for its first 8192 rows, as Rust's writer decides, so a column that bit-packs better is no longer dictionary-encoded. ([#599](https://github.com/dfa1/vortex-java/pull/599)) - 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)) diff --git a/docs/compatibility.md b/docs/compatibility.md index 183aa630..d297b5d3 100644 --- a/docs/compatibility.md +++ b/docs/compatibility.md @@ -124,7 +124,7 @@ decimals ([#430](https://github.com/dfa1/vortex-java/pull/430)) and nulls in nul | `vortex.varbinview` | `VarBinViewEncodingDecoder` | `VarBinViewEncodingEncoder` | ✅ | ✅ | Utf8, Binary | | `vortex.alp` | `AlpEncodingDecoder` | `AlpEncodingEncoder` | ✅ | ✅ | F64, F32 | | `vortex.alprd` | `AlpRdEncodingDecoder` | `AlpRdEncodingEncoder` | ✅ | ✅ | F64, F32 | -| `vortex.dict` | `DictEncodingDecoder` | `DictEncodingEncoder` | ✅ | ✅ | Primitive, Utf8/Binary. Writes `DictMetadata` + children `[codes, values]`, both cascaded (a float pool goes through ALP). With `globalDict` on (the default) the writer's data cascade never dictionary-encodes integers: Rust's file writer excludes `IntDictScheme` there and dictionary-encodes suitable columns once, as a layout. With it off (the CSV and Parquet importers stream), integer chunks keep their per-chunk dictionary. Decode also accepts the pre-0.14.5 Java shape (1-byte metadata, `[values, codes]`) so older Java-written files stay readable; that shape was never readable by Rust | +| `vortex.dict` | `DictEncodingDecoder` | `DictEncodingEncoder` | ✅ | ✅ | Primitive, Utf8/Binary. Writes `DictMetadata` + children `[codes, values]`, both cascaded (a float pool goes through ALP). With `globalDict` on (the default) the writer's data cascade never dictionary-encodes integers: Rust's file writer excludes `IntDictScheme` there and dictionary-encodes suitable columns once, as a layout. A column first needs under half distinct values in its first batch, at most 2,048 of them (32,768 for strings), and a type the dict layout reads (no `I8`/`U8`/`I16`/`U16`/`F16`/`F32`); Rust has no such cap and skips only above half. A non-null primitive column that passes then gets the layout only when the cascade, with the integer dictionary back in, picks `vortex.dict` for its first rows: Rust's `DictStrategy` probe. Rust probes exactly the first 8192 rows across batches; Java probes the first batch's first 8192, so a first batch under 8192 rows can decide differently. Nullable and string columns skip the probe and decide by the cardinality check alone, which admits some nullable integer columns Rust bit-packs. With it off (the CSV and Parquet importers stream), integer chunks keep their per-chunk dictionary. Decode also accepts the pre-0.14.5 Java shape (1-byte metadata, `[values, codes]`) so older Java-written files stay readable; that shape was never readable by Rust | | `vortex.sparse` | `SparseEncodingDecoder` | `SparseEncodingEncoder` | ✅ | ✅ | Primitive, Bool, Utf8/Binary (decode); Primitive (encode) | | `vortex.sequence` | `SequenceEncodingDecoder` | `SequenceEncodingEncoder` | ✅ | ✅ | Primitive (encode: integer only, as Rust; decode also float) | | `vortex.struct` | `StructEncodingDecoder` | `StructEncodingEncoder` | ✅ | ✅ | Struct | diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/DictAdmissionInteropIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/DictAdmissionInteropIntegrationTest.java new file mode 100644 index 00000000..9abe7f2e --- /dev/null +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/DictAdmissionInteropIntegrationTest.java @@ -0,0 +1,198 @@ +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.inspect.InspectorTree; +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.BigIntVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.FloatingPointPrecision; +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.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.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; + +/// Rust's dict layout admits a column only when its probe compressor, the full cascade, turns the +/// column's first 8192-row block into a `vortex.dict` (`DictStrategy`); a cardinality rule of thumb +/// admits columns Rust leaves to bit-packing. Each column goes through vortex-jni and vortex-java, +/// and both must agree on whether it gets a dict layout. +class DictAdmissionInteropIntegrationTest { + + private static final Session SESSION = Session.create(); + private static final BufferAllocator ALLOCATOR = ArrowAllocation.rootAllocator(); + private static final ColumnName V = ColumnName.of("v"); + private static final int ROWS = 20_000; + + static { + NativeLoader.loadJni(); + } + + /// The column type. The wide kinds spread their values over the whole range, where a + /// dictionary beats bit-packing; the others keep them small. + private enum Kind { I32, I64, WIDE_I64, F64, WIDE_F64, UTF8 } + + static Stream columns() { + return Stream.of( + // wide values: a dictionary beats the values + Arguments.of("wide i64, 100 distinct", Kind.WIDE_I64, 100, 0), + Arguments.of("wide f64, 10 distinct", Kind.WIDE_F64, 10, 0), + // few small values: bit-packing a few bits a row beats a dictionary + Arguments.of("i64, 4 distinct", Kind.I64, 4, 0), + Arguments.of("f64, 5 distinct", Kind.F64, 5, 0), + Arguments.of("utf8, 5 distinct", Kind.UTF8, 5, 0), + Arguments.of("nullable i32, 10 distinct", Kind.I32, 10, 7), + // under half distinct, yet bit-packing 7 bits wins over a 100-value dictionary. + // A nullable column with 1000 distinct values still gets a dict in Java only: its + // admission keeps the cardinality check until the nullable cascade charges + // bit-packing the validity, as Rust's sampled sizes do. + Arguments.of("i64, 100 distinct", Kind.I64, 100, 0), + Arguments.of("i32, 1000 distinct", Kind.I32, 1000, 0), + Arguments.of("utf8, 1000 distinct", Kind.UTF8, 1000, 0), + Arguments.of("utf8, mostly distinct", Kind.UTF8, 1_000_000, 0)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("columns") + void javaGivesADictLayoutWhereJniDoes(String name, Kind kind, int distinct, int nullEvery, @TempDir Path tmp) + throws Exception { + // Given + long[] draws = new Random(3).longs(ROWS, 0, distinct).toArray(); + Path jniFile = tmp.resolve("jni.vortex"); + Path javaFile = tmp.resolve("java.vortex"); + writeWithJni(jniFile, kind, draws, nullEvery); + boolean expected = hasDictLayout(jniFile); + + // When + writeWithJava(javaFile, kind, draws, nullEvery); + + // Then + assertThat(hasDictLayout(javaFile)).as("vortex-jni dict layout: %s", expected).isEqualTo(expected); + } + + private static boolean isNull(int row, int nullEvery) { + return nullEvery > 0 && row % nullEvery == 0; + } + + private static long i64(Kind kind, long draw) { + return kind == Kind.WIDE_I64 ? draw * 0x9E3779B97F4A7C15L : draw; + } + + private static double f64(Kind kind, long draw) { + return kind == Kind.WIDE_F64 ? new Random(draw).nextDouble() * 1000 : draw * 0.25; + } + + private static String string(long draw) { + return "value-" + draw; + } + + private static void writeWithJni(Path path, Kind kind, long[] draws, int nullEvery) throws Exception { + ArrowType type = switch (kind) { + case I32 -> new ArrowType.Int(32, true); + case I64, WIDE_I64 -> new ArrowType.Int(64, true); + case F64, WIDE_F64 -> new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); + case UTF8 -> new ArrowType.Utf8(); + }; + Field field = nullEvery > 0 ? Field.nullable("v", type) : Field.notNullable("v", type); + Schema schema = new Schema(List.of(field)); + try (VortexWriter writer = VortexWriter.builder(SESSION, path.toAbsolutePath().toUri().toString(), schema, + ALLOCATOR).build(); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, ALLOCATOR)) { + FieldVector vector = root.getVector("v"); + vector.setInitialCapacity(ROWS); + vector.allocateNew(); + for (int i = 0; i < ROWS; i++) { + if (isNull(i, nullEvery)) { + vector.setNull(i); + continue; + } + switch (kind) { + case I32 -> ((IntVector) vector).setSafe(i, (int) draws[i]); + case I64, WIDE_I64 -> ((BigIntVector) vector).setSafe(i, i64(kind, draws[i])); + case F64, WIDE_F64 -> ((Float8Vector) vector).setSafe(i, f64(kind, draws[i])); + case UTF8 -> ((VarCharVector) vector).setSafe(i, string(draws[i]).getBytes(StandardCharsets.UTF_8)); + } + } + vector.setValueCount(ROWS); + root.setRowCount(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, Kind kind, long[] draws, int nullEvery) throws IOException { + boolean nullable = nullEvery > 0; + DType dtype = switch (kind) { + case I32 -> new DType.Primitive(PType.I32, nullable); + case I64, WIDE_I64 -> new DType.Primitive(PType.I64, nullable); + case F64, WIDE_F64 -> new DType.Primitive(PType.F64, nullable); + case UTF8 -> new DType.Utf8(nullable); + }; + Object values = switch (kind) { + case I32 -> Arrays.stream(draws).mapToInt(d -> (int) d).toArray(); + case I64, WIDE_I64 -> Arrays.stream(draws).map(d -> i64(kind, d)).toArray(); + case F64, WIDE_F64 -> Arrays.stream(draws).mapToDouble(d -> f64(kind, d)).toArray(); + case UTF8 -> Arrays.stream(draws).mapToObj(DictAdmissionInteropIntegrationTest::string) + .toArray(String[]::new); + }; + Object data = values; + if (nullable) { + boolean[] valid = new boolean[ROWS]; + for (int i = 0; i < ROWS; i++) { + valid[i] = !isNull(i, nullEvery); + } + data = new NullableData(values, valid); + } + DType.Struct schema = new DType.Struct(List.of(V), List.of(dtype), false); + try (var channel = FileChannel.open(path, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var writer = io.github.dfa1.vortex.writer.VortexWriter.create(channel, schema, WriteOptions.defaults())) { + writer.writeChunk(Map.of(V, data)); + } + } + + private static boolean hasDictLayout(Path file) throws IOException { + try (VortexReader reader = VortexReader.open(file)) { + return hasDictLayout(InspectorTree.build(reader).root()); + } + } + + private static boolean hasDictLayout(InspectorTree.Node node) { + if (node.layout().isDict()) { + return true; + } + return node.children().stream().anyMatch(DictAdmissionInteropIntegrationTest::hasDictLayout); + } +} 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 4856fa4e..2a229f21 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 @@ -82,6 +82,15 @@ static long[] sizePieces(Object data, long rows, long firstRow) { return sizes; } + /// The first `rows` rows of a batch, in the batch's carrier shape. + /// + /// @param data a batch of a column [#supports(DType)] accepts + /// @param rows how many leading rows to keep, at most the batch's row count + /// @return the leading rows + static Object head(Object data, int rows) { + return concat(List.of(new Piece(data, 0, rows, 0))); + } + /// Whether sizing a batch of `data` reads its values, as for strings and binaries -- the /// only batches worth [#sizePieces(Object,long,long)] ahead of time. /// 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 8aad3682..9ccea0cd 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 @@ -133,6 +133,10 @@ public final class VortexWriter implements Closeable { private final WriteRegistry extensionRegistry; private final List cascadeCodecs; private final WriteRegistry cascadeRegistry; + // The dict-layout probe's codecs: the cascade with the integer dictionary back in, as Rust's + // probe is its full compressor while its data compressor leaves IntDictScheme out. + private final List probeCodecs; + private final WriteRegistry probeRegistry; private final List segs = new ArrayList<>(); private final Map> colChunks = new LinkedHashMap<>(); private final Map encodingIdx = new LinkedHashMap<>(); @@ -226,6 +230,10 @@ private VortexWriter( // the caller left out would produce a file their own read registry may not decode. this.cascadeCodecs = encodings == DEFAULT_CODECS ? buildCascadeCodecs(options) : List.copyOf(encodings); this.cascadeRegistry = buildRegistry(this.cascadeCodecs); + this.probeCodecs = this.cascadeCodecs.stream() + .map(e -> e instanceof DictEncodingEncoder ? new DictEncodingEncoder() : e) + .toList(); + this.probeRegistry = buildRegistry(this.probeCodecs); this.editionAllowed = editionAllowed(options.editions()); this.zonedLayout = options.enableZoneMaps() && editionHasZonedLayout(options.editions()); for (Map.Entry e : options.columnEncodings().entrySet()) { @@ -646,6 +654,13 @@ public void writeChunk(Map columns) throws IOException { // Rust's dict layout admits Primitive | Utf8 | Binary (dict_layout_supported). candidate = DictColumnState.isVarBinDictCandidate((Object[]) values, validity); } + // ponytail: only non-null primitives take Rust's probe. A nullable column's cascade + // runs under vortex.masked, which leaves the validity out of the sampled sizes Rust + // charges bit-packing with, and strings pick OnPair where Rust's StringDictScheme + // wins, so for those the cardinality check above still matches Rust more often. + if (candidate && colDtype instanceof DType.Primitive p && !p.nullable()) { + candidate = probeCompressesToDict(colDtype, data); + } if (candidate) { dictCandidates.add(colName); dictStates.put(colName, new DictColumnState(colDtype)); @@ -684,6 +699,22 @@ public void writeChunk(Map columns) throws IOException { firstChunkSeen = true; } + /// Rust's `DictStrategy` admission: the column gets a dict layout only when the full cascade, + /// its probe compressor, turns the first 8192-row block into a `vortex.dict` + /// (`vortex-layout/src/layouts/dict/writer.rs`). Without a cascade there is no compressor to + /// ask, and the cardinality check alone decides. + private boolean probeCompressesToDict(DType colDtype, Object data) { + if (options.allowedCascading() == 0) { + return true; + } + Object block = Repartitioner.head(data, (int) Math.min(Repartitioner.BLOCK_ROWS, arrayLength(data))); + try (Arena arena = Arena.ofConfined()) { + EncodeContext ctx = EncodeContext.ofDepth(options.allowedCascading(), arena, probeRegistry, editionExcluded); + EncodeResult result = new CascadingCompressor(probeCodecs).encode(colDtype, block, ctx); + return result.rootNode().encodingId().equals(EncodingId.VORTEX_DICT); + } + } + /// Hands one batch of a column to its [Repartitioner], writing every chunk it completes; a /// column whose carriers cannot be coalesced is written one chunk per batch. private void appendChunk(ColumnName colName, DType colDtype, Object data) throws IOException { diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/GlobalDictPrimitiveTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/GlobalDictPrimitiveTest.java index 20d0cebf..8fea5106 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/GlobalDictPrimitiveTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/GlobalDictPrimitiveTest.java @@ -114,7 +114,7 @@ void cardinalityJustOverU8_i64_roundTrips(@TempDir Path tmp) throws IOException int cardinality = 300; long[] data = new long[rows]; for (int i = 0; i < rows; i++) { - data[i] = i % cardinality; + data[i] = wide(i % cardinality); } Path file = tmp.resolve("u16_i64.vortex"); @@ -133,7 +133,7 @@ void cardinalityExceededAcrossChunks_i64_fallsBackToPerChunk(@TempDir Path tmp) // full key set only at flush, so it trips the per-chunk fallback in writeGlobalDictColumn. long[] c0 = new long[1_000]; for (int i = 0; i < c0.length; i++) { - c0[i] = i % 400; + c0[i] = wide(i % 400); } long[] c1 = new long[1_700]; for (int i = 0; i < c1.length; i++) { @@ -174,7 +174,7 @@ void cardinalityGrowsPastCapMidFile_i64_demotesImmediatelyToChunkedLayout(@TempD for (int c = 0; c < chunkCount; c++) { long[] data = new long[rowsPerChunk]; for (int i = 0; i < rowsPerChunk; i++) { - long value = (c == 0) ? (i % 200) : (1_000_000L + next++); + long value = (c == 0) ? wide(i % 200) : (1_000_000L + next++); data[i] = value; expected[c * rowsPerChunk + i] = value; } @@ -290,6 +290,37 @@ void dominantValueNotFirstSeen_codesStillCompress(@TempDir Path tmp) throws IOEx } } + /// The `k`th of a set of 64-bit-wide values. The dict layout takes a non-null integer column + /// only when the cascade picks a dictionary for its first block, as Rust's probe does; small + /// values like `i % 200` bit-pack better and never get one, so the global-dict paths below need + /// values a dictionary actually wins on. + private static long wide(long k) { + return k * 0x9E3779B97F4A7C15L; + } + + @Test + void wideLowCardinality_i64_getsTheGlobalDictLayout(@TempDir Path tmp) throws IOException { + // Given — the first chunk the demotion and cap tests above rely on being admitted + long[] data = new long[1_000]; + for (int i = 0; i < data.length; i++) { + data[i] = wide(i % 200); + } + var schema = new DType.Struct(List.of(ColumnName.of("v")), List.of(DType.I64), false); + Path file = tmp.resolve("wide_dict_i64.vortex"); + try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var sut = VortexWriter.create(ch, schema, WriteOptions.cascading(3))) { + sut.writeChunk(Map.of(ColumnName.of("v"), data)); + } + + // When + try (var vf = VortexReader.open(file, ReadRegistry.loadAll())) { + var result = unwrapZoned(vf.layout().children().getFirst()); + + // Then + assertThat(result.isDict()).isTrue(); + } + } + /// Unwraps a column's Zoned/Stats wrapper (the writer wraps every column in one for zone-map /// pruning) to reach the encoding layout beneath. private static io.github.dfa1.vortex.reader.layout.Layout unwrapZoned(