fix(writer): admit an integer column to the dict layout as Rust's probe does

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))
 

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                                                                |

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<Arguments> 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);
+    }
+}

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.
     ///

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<EncodingEncoder> 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<EncodingEncoder> probeCodecs;
+    private final WriteRegistry probeRegistry;
     private final List<SegRef> segs = new ArrayList<>();
     private final Map<ColumnName, List<ChunkRef>> colChunks = new LinkedHashMap<>();
     private final Map<EncodingId, Integer> 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<ColumnName, ColumnEncoding> e : options.columnEncodings().entrySet()) {
@@ -646,6 +654,13 @@ public void writeChunk(Map<ColumnName, Object> 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<ColumnName, Object> 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 {