Skip to content

Commit 4f9ca0b

Browse files
dfa1claude
andcommitted
fix(writer): admit an integer column to the dict layout as Rust's probe does
Rust's DictStrategy gives a column a dict layout only when its probe compressor, the full cascade with IntDictScheme, turns the column's first 8192-row block into a vortex.dict. Java admitted any column with under half distinct values and at most 2,048 of them, so a column that bit-packs better (100 small distinct I64 values: 7 bits a row against a 100-value dictionary plus codes) got a dict layout vortex-jni never writes; the taxi payment_type divergence was this. A non-null primitive column now takes the same probe: the cascade with the per-chunk integer dictionary back in, over the column's first block. Nullable and string columns keep the cardinality check: Java cascades a nullable column's values under vortex.masked, without the validity cost Rust's sampled bit-packing size carries, and picks OnPair for strings where Rust's StringDictScheme wins, so the probe would diverge more there than the check does. DictAdmissionInteropIntegrationTest writes ten column shapes with vortex-jni and vortex-java and asserts both agree on the dict layout. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SwhK4W5cgj7y7s4ZPCif5x
1 parent 9aebb3b commit 4f9ca0b

6 files changed

Lines changed: 274 additions & 4 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
2121
- 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))
2222
- 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))
2323
- 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))
24+
- 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))
2425
- 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))
2526
- `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))
2627

‎docs/compatibility.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ decimals ([#430](https://github.com/dfa1/vortex-java/pull/430)) and nulls in nul
124124
| `vortex.varbinview` | `VarBinViewEncodingDecoder` | `VarBinViewEncodingEncoder` | ✅ | ✅ | Utf8, Binary |
125125
| `vortex.alp` | `AlpEncodingDecoder` | `AlpEncodingEncoder` | ✅ | ✅ | F64, F32 |
126126
| `vortex.alprd` | `AlpRdEncodingDecoder` | `AlpRdEncodingEncoder` | ✅ | ✅ | F64, F32 |
127-
| `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 |
127+
| `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 |
128128
| `vortex.sparse` | `SparseEncodingDecoder` | `SparseEncodingEncoder` | ✅ | ✅ | Primitive, Bool, Utf8/Binary (decode); Primitive (encode) |
129129
| `vortex.sequence` | `SequenceEncodingDecoder` | `SequenceEncodingEncoder` | ✅ | ✅ | Primitive (encode: integer only, as Rust; decode also float) |
130130
| `vortex.struct` | `StructEncodingDecoder` | `StructEncodingEncoder` | ✅ | ✅ | Struct |
Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
1+
package io.github.dfa1.vortex.integration;
2+
3+
import dev.vortex.api.Session;
4+
import dev.vortex.api.VortexWriter;
5+
import dev.vortex.arrow.ArrowAllocation;
6+
import dev.vortex.jni.NativeLoader;
7+
import io.github.dfa1.vortex.core.model.ColumnName;
8+
import io.github.dfa1.vortex.core.model.DType;
9+
import io.github.dfa1.vortex.core.model.PType;
10+
import io.github.dfa1.vortex.inspect.InspectorTree;
11+
import io.github.dfa1.vortex.reader.VortexReader;
12+
import io.github.dfa1.vortex.writer.WriteOptions;
13+
import io.github.dfa1.vortex.writer.encode.NullableData;
14+
import org.apache.arrow.c.ArrowArray;
15+
import org.apache.arrow.c.ArrowSchema;
16+
import org.apache.arrow.c.Data;
17+
import org.apache.arrow.memory.BufferAllocator;
18+
import org.apache.arrow.vector.BigIntVector;
19+
import org.apache.arrow.vector.FieldVector;
20+
import org.apache.arrow.vector.Float8Vector;
21+
import org.apache.arrow.vector.IntVector;
22+
import org.apache.arrow.vector.VarCharVector;
23+
import org.apache.arrow.vector.VectorSchemaRoot;
24+
import org.apache.arrow.vector.types.FloatingPointPrecision;
25+
import org.apache.arrow.vector.types.pojo.ArrowType;
26+
import org.apache.arrow.vector.types.pojo.Field;
27+
import org.apache.arrow.vector.types.pojo.Schema;
28+
import org.junit.jupiter.api.io.TempDir;
29+
import org.junit.jupiter.params.ParameterizedTest;
30+
import org.junit.jupiter.params.provider.Arguments;
31+
import org.junit.jupiter.params.provider.MethodSource;
32+
33+
import java.io.IOException;
34+
import java.nio.channels.FileChannel;
35+
import java.nio.charset.StandardCharsets;
36+
import java.nio.file.Path;
37+
import java.nio.file.StandardOpenOption;
38+
import java.util.Arrays;
39+
import java.util.List;
40+
import java.util.Map;
41+
import java.util.Random;
42+
import java.util.stream.Stream;
43+
44+
import static org.assertj.core.api.Assertions.assertThat;
45+
46+
/// Rust's dict layout admits a column only when its probe compressor, the full cascade, turns the
47+
/// column's first 8192-row block into a `vortex.dict` (`DictStrategy`); a cardinality rule of thumb
48+
/// admits columns Rust leaves to bit-packing. Each column goes through vortex-jni and vortex-java,
49+
/// and both must agree on whether it gets a dict layout.
50+
class DictAdmissionInteropIntegrationTest {
51+
52+
private static final Session SESSION = Session.create();
53+
private static final BufferAllocator ALLOCATOR = ArrowAllocation.rootAllocator();
54+
private static final ColumnName V = ColumnName.of("v");
55+
private static final int ROWS = 20_000;
56+
57+
static {
58+
NativeLoader.loadJni();
59+
}
60+
61+
/// The column type. The wide kinds spread their values over the whole range, where a
62+
/// dictionary beats bit-packing; the others keep them small.
63+
private enum Kind { I32, I64, WIDE_I64, F64, WIDE_F64, UTF8 }
64+
65+
static Stream<Arguments> columns() {
66+
return Stream.of(
67+
// wide values: a dictionary beats the values
68+
Arguments.of("wide i64, 100 distinct", Kind.WIDE_I64, 100, 0),
69+
Arguments.of("wide f64, 10 distinct", Kind.WIDE_F64, 10, 0),
70+
// few small values: bit-packing a few bits a row beats a dictionary
71+
Arguments.of("i64, 4 distinct", Kind.I64, 4, 0),
72+
Arguments.of("f64, 5 distinct", Kind.F64, 5, 0),
73+
Arguments.of("utf8, 5 distinct", Kind.UTF8, 5, 0),
74+
Arguments.of("nullable i32, 10 distinct", Kind.I32, 10, 7),
75+
// under half distinct, yet bit-packing 7 bits wins over a 100-value dictionary.
76+
// A nullable column with 1000 distinct values still gets a dict in Java only: its
77+
// admission keeps the cardinality check until the nullable cascade charges
78+
// bit-packing the validity, as Rust's sampled sizes do.
79+
Arguments.of("i64, 100 distinct", Kind.I64, 100, 0),
80+
Arguments.of("i32, 1000 distinct", Kind.I32, 1000, 0),
81+
Arguments.of("utf8, 1000 distinct", Kind.UTF8, 1000, 0),
82+
Arguments.of("utf8, mostly distinct", Kind.UTF8, 1_000_000, 0));
83+
}
84+
85+
@ParameterizedTest(name = "{0}")
86+
@MethodSource("columns")
87+
void javaGivesADictLayoutWhereJniDoes(String name, Kind kind, int distinct, int nullEvery, @TempDir Path tmp)
88+
throws Exception {
89+
// Given
90+
long[] draws = new Random(3).longs(ROWS, 0, distinct).toArray();
91+
Path jniFile = tmp.resolve("jni.vortex");
92+
Path javaFile = tmp.resolve("java.vortex");
93+
writeWithJni(jniFile, kind, draws, nullEvery);
94+
boolean expected = hasDictLayout(jniFile);
95+
96+
// When
97+
writeWithJava(javaFile, kind, draws, nullEvery);
98+
99+
// Then
100+
assertThat(hasDictLayout(javaFile)).as("vortex-jni dict layout: %s", expected).isEqualTo(expected);
101+
}
102+
103+
private static boolean isNull(int row, int nullEvery) {
104+
return nullEvery > 0 && row % nullEvery == 0;
105+
}
106+
107+
private static long i64(Kind kind, long draw) {
108+
return kind == Kind.WIDE_I64 ? draw * 0x9E3779B97F4A7C15L : draw;
109+
}
110+
111+
private static double f64(Kind kind, long draw) {
112+
return kind == Kind.WIDE_F64 ? new Random(draw).nextDouble() * 1000 : draw * 0.25;
113+
}
114+
115+
private static String string(long draw) {
116+
return "value-" + draw;
117+
}
118+
119+
private static void writeWithJni(Path path, Kind kind, long[] draws, int nullEvery) throws Exception {
120+
ArrowType type = switch (kind) {
121+
case I32 -> new ArrowType.Int(32, true);
122+
case I64, WIDE_I64 -> new ArrowType.Int(64, true);
123+
case F64, WIDE_F64 -> new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE);
124+
case UTF8 -> new ArrowType.Utf8();
125+
};
126+
Field field = nullEvery > 0 ? Field.nullable("v", type) : Field.notNullable("v", type);
127+
Schema schema = new Schema(List.of(field));
128+
try (VortexWriter writer = VortexWriter.builder(SESSION, path.toAbsolutePath().toUri().toString(), schema,
129+
ALLOCATOR).build();
130+
VectorSchemaRoot root = VectorSchemaRoot.create(schema, ALLOCATOR)) {
131+
FieldVector vector = root.getVector("v");
132+
vector.setInitialCapacity(ROWS);
133+
vector.allocateNew();
134+
for (int i = 0; i < ROWS; i++) {
135+
if (isNull(i, nullEvery)) {
136+
vector.setNull(i);
137+
continue;
138+
}
139+
switch (kind) {
140+
case I32 -> ((IntVector) vector).setSafe(i, (int) draws[i]);
141+
case I64, WIDE_I64 -> ((BigIntVector) vector).setSafe(i, i64(kind, draws[i]));
142+
case F64, WIDE_F64 -> ((Float8Vector) vector).setSafe(i, f64(kind, draws[i]));
143+
case UTF8 -> ((VarCharVector) vector).setSafe(i, string(draws[i]).getBytes(StandardCharsets.UTF_8));
144+
}
145+
}
146+
vector.setValueCount(ROWS);
147+
root.setRowCount(ROWS);
148+
try (ArrowArray array = ArrowArray.allocateNew(ALLOCATOR);
149+
ArrowSchema arrowSchema = ArrowSchema.allocateNew(ALLOCATOR)) {
150+
Data.exportVectorSchemaRoot(ALLOCATOR, root, null, array, arrowSchema);
151+
writer.writeBatch(array.memoryAddress(), arrowSchema.memoryAddress());
152+
}
153+
}
154+
}
155+
156+
private static void writeWithJava(Path path, Kind kind, long[] draws, int nullEvery) throws IOException {
157+
boolean nullable = nullEvery > 0;
158+
DType dtype = switch (kind) {
159+
case I32 -> new DType.Primitive(PType.I32, nullable);
160+
case I64, WIDE_I64 -> new DType.Primitive(PType.I64, nullable);
161+
case F64, WIDE_F64 -> new DType.Primitive(PType.F64, nullable);
162+
case UTF8 -> new DType.Utf8(nullable);
163+
};
164+
Object values = switch (kind) {
165+
case I32 -> Arrays.stream(draws).mapToInt(d -> (int) d).toArray();
166+
case I64, WIDE_I64 -> Arrays.stream(draws).map(d -> i64(kind, d)).toArray();
167+
case F64, WIDE_F64 -> Arrays.stream(draws).mapToDouble(d -> f64(kind, d)).toArray();
168+
case UTF8 -> Arrays.stream(draws).mapToObj(DictAdmissionInteropIntegrationTest::string)
169+
.toArray(String[]::new);
170+
};
171+
Object data = values;
172+
if (nullable) {
173+
boolean[] valid = new boolean[ROWS];
174+
for (int i = 0; i < ROWS; i++) {
175+
valid[i] = !isNull(i, nullEvery);
176+
}
177+
data = new NullableData(values, valid);
178+
}
179+
DType.Struct schema = new DType.Struct(List.of(V), List.of(dtype), false);
180+
try (var channel = FileChannel.open(path, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
181+
var writer = io.github.dfa1.vortex.writer.VortexWriter.create(channel, schema, WriteOptions.defaults())) {
182+
writer.writeChunk(Map.of(V, data));
183+
}
184+
}
185+
186+
private static boolean hasDictLayout(Path file) throws IOException {
187+
try (VortexReader reader = VortexReader.open(file)) {
188+
return hasDictLayout(InspectorTree.build(reader).root());
189+
}
190+
}
191+
192+
private static boolean hasDictLayout(InspectorTree.Node node) {
193+
if (node.layout().isDict()) {
194+
return true;
195+
}
196+
return node.children().stream().anyMatch(DictAdmissionInteropIntegrationTest::hasDictLayout);
197+
}
198+
}

‎writer/src/main/java/io/github/dfa1/vortex/writer/Repartitioner.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,15 @@ static long[] sizePieces(Object data, long rows, long firstRow) {
8282
return sizes;
8383
}
8484

85+
/// The first `rows` rows of a batch, in the batch's carrier shape.
86+
///
87+
/// @param data a batch of a column [#supports(DType)] accepts
88+
/// @param rows how many leading rows to keep, at most the batch's row count
89+
/// @return the leading rows
90+
static Object head(Object data, int rows) {
91+
return concat(List.of(new Piece(data, 0, rows, 0)));
92+
}
93+
8594
/// Whether sizing a batch of `data` reads its values, as for strings and binaries -- the
8695
/// only batches worth [#sizePieces(Object,long,long)] ahead of time.
8796
///

‎writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,10 @@ public final class VortexWriter implements Closeable {
133133
private final WriteRegistry extensionRegistry;
134134
private final List<EncodingEncoder> cascadeCodecs;
135135
private final WriteRegistry cascadeRegistry;
136+
// The dict-layout probe's codecs: the cascade with the integer dictionary back in, as Rust's
137+
// probe is its full compressor while its data compressor leaves IntDictScheme out.
138+
private final List<EncodingEncoder> probeCodecs;
139+
private final WriteRegistry probeRegistry;
136140
private final List<SegRef> segs = new ArrayList<>();
137141
private final Map<ColumnName, List<ChunkRef>> colChunks = new LinkedHashMap<>();
138142
private final Map<EncodingId, Integer> encodingIdx = new LinkedHashMap<>();
@@ -226,6 +230,10 @@ private VortexWriter(
226230
// the caller left out would produce a file their own read registry may not decode.
227231
this.cascadeCodecs = encodings == DEFAULT_CODECS ? buildCascadeCodecs(options) : List.copyOf(encodings);
228232
this.cascadeRegistry = buildRegistry(this.cascadeCodecs);
233+
this.probeCodecs = this.cascadeCodecs.stream()
234+
.map(e -> e instanceof DictEncodingEncoder ? new DictEncodingEncoder() : e)
235+
.toList();
236+
this.probeRegistry = buildRegistry(this.probeCodecs);
229237
this.editionAllowed = editionAllowed(options.editions());
230238
this.zonedLayout = options.enableZoneMaps() && editionHasZonedLayout(options.editions());
231239
for (Map.Entry<ColumnName, ColumnEncoding> e : options.columnEncodings().entrySet()) {
@@ -646,6 +654,13 @@ public void writeChunk(Map<ColumnName, Object> columns) throws IOException {
646654
// Rust's dict layout admits Primitive | Utf8 | Binary (dict_layout_supported).
647655
candidate = DictColumnState.isVarBinDictCandidate((Object[]) values, validity);
648656
}
657+
// ponytail: only non-null primitives take Rust's probe. A nullable column's cascade
658+
// runs under vortex.masked, which leaves the validity out of the sampled sizes Rust
659+
// charges bit-packing with, and strings pick OnPair where Rust's StringDictScheme
660+
// wins, so for those the cardinality check above still matches Rust more often.
661+
if (candidate && colDtype instanceof DType.Primitive p && !p.nullable()) {
662+
candidate = probeCompressesToDict(colDtype, data);
663+
}
649664
if (candidate) {
650665
dictCandidates.add(colName);
651666
dictStates.put(colName, new DictColumnState(colDtype));
@@ -684,6 +699,22 @@ public void writeChunk(Map<ColumnName, Object> columns) throws IOException {
684699
firstChunkSeen = true;
685700
}
686701

702+
/// Rust's `DictStrategy` admission: the column gets a dict layout only when the full cascade,
703+
/// its probe compressor, turns the first 8192-row block into a `vortex.dict`
704+
/// (`vortex-layout/src/layouts/dict/writer.rs`). Without a cascade there is no compressor to
705+
/// ask, and the cardinality check alone decides.
706+
private boolean probeCompressesToDict(DType colDtype, Object data) {
707+
if (options.allowedCascading() == 0) {
708+
return true;
709+
}
710+
Object block = Repartitioner.head(data, (int) Math.min(Repartitioner.BLOCK_ROWS, arrayLength(data)));
711+
try (Arena arena = Arena.ofConfined()) {
712+
EncodeContext ctx = EncodeContext.ofDepth(options.allowedCascading(), arena, probeRegistry, editionExcluded);
713+
EncodeResult result = new CascadingCompressor(probeCodecs).encode(colDtype, block, ctx);
714+
return result.rootNode().encodingId().equals(EncodingId.VORTEX_DICT);
715+
}
716+
}
717+
687718
/// Hands one batch of a column to its [Repartitioner], writing every chunk it completes; a
688719
/// column whose carriers cannot be coalesced is written one chunk per batch.
689720
private void appendChunk(ColumnName colName, DType colDtype, Object data) throws IOException {

0 commit comments

Comments
 (0)