Skip to content

Commit 6b9d101

Browse files
dfa1claude
andcommitted
feat(writer): cascade the per-chunk dict Utf8 values pool (#299)
DictEncodingEncoder.encodeUtf8 hardcoded the distinct-values pool as a terminal raw vortex.varbin node, so per-chunk dict-encoded Utf8 columns never FSST-compressed their values. encodeCascade now exposes the values pool as an open Utf8 child slot instead of CascadeStep.terminal, so the cascade competes FSST/VarBin/Zstd on it. CascadingCompressor already excludes the winning vortex.dict from a child's competition, so the pool is never wrapped in a second dict the reader cannot unwrap. The non-cascade encode() path (cascade depth 0) still emits terminal varbin. Verified: encodeCascade exposes a single open values child (unit test); per-chunk dict Utf8 round-trips through the new path; 217 Java-writes-Rust-reads interop tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 3f0efa6 commit 6b9d101

3 files changed

Lines changed: 106 additions & 1 deletion

File tree

integration/src/test/java/io/github/dfa1/vortex/integration/DictValuesPoolCompressionIntegrationTest.java

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,43 @@ void utf8_cardinalityAbovePriorCap_admittedToGlobalDict(@TempDir Path tmp) throw
123123
}
124124
}
125125

126+
@Test
127+
void perChunkDictUtf8_valuesPool_isFsstCompressed(@TempDir Path tmp) throws IOException {
128+
// Given — global dict disabled, so a redundant low-cardinality Utf8 column is dict-encoded
129+
// per chunk (a vortex.dict *encoding*, not a layout). High row duplication makes per-chunk
130+
// dict win the cascade; the shared-prefix distinct values make FSST win the values child.
131+
Path file = tmp.resolve("perchunk_dict.vortex");
132+
int rowsPerChunk = 10_000;
133+
int chunkCount = 3;
134+
int distinct = 500;
135+
136+
List<String> expected = new ArrayList<>();
137+
try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
138+
var sut = VortexWriter.create(ch, SCHEMA, WriteOptions.cascading(3).withGlobalDict(false))) {
139+
// When
140+
for (int c = 0; c < chunkCount; c++) {
141+
String[] data = redundantChunk(rowsPerChunk, distinct, c * 53);
142+
expected.addAll(List.of(data));
143+
sut.writeChunk(Map.of(COL, data));
144+
}
145+
}
146+
147+
// Then — per-chunk dict still wins (the values pool is now a cascaded child rather than a
148+
// terminal varbin node; its encoding is nested inside the dict segment, so it is not visible
149+
// in usedEncodings — the mechanism is covered by DictEncodingEncoderTest, the round-trip
150+
// below guards correctness through the new cascade path).
151+
try (var vf = VortexReader.open(file, ReadRegistry.loadAll())) {
152+
assertThat(InspectorTree.build(vf).usedEncodings())
153+
.as("per-chunk dict encoding")
154+
.contains("vortex.dict");
155+
}
156+
157+
// And every value round-trips exactly.
158+
try (var vf = VortexReader.open(file, ReadRegistry.loadAll())) {
159+
assertThat(readAllStrings(vf)).isEqualTo(expected);
160+
}
161+
}
162+
126163
private static boolean hasDictLayout(InspectorTree.Node node) {
127164
if (node.layout().isDict()) {
128165
return true;

writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) {
8484
@Override
8585
public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) {
8686
if (dtype instanceof DType.Utf8) {
87-
return CascadeStep.terminal(encodeUtf8((String[]) data, ctx));
87+
return encodeUtf8Cascade((String[]) data, ctx);
8888
}
8989
DictData d = buildDictData(dtype, data, ctx);
9090
PType codePType = d.codePType();
@@ -101,6 +101,56 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) {
101101
return new CascadeStep(partialRoot, List.of(d.valuesBuf()), List.of(slot), null, null, true);
102102
}
103103

104+
/// Cascading Utf8 dict: emit the codes leaf but expose the distinct-values pool as an open Utf8
105+
/// child slot so the cascade competes FSST/VarBin/Zstd on it (#299) instead of hardcoding raw
106+
/// varbin. [CascadingCompressor] excludes the winning `vortex.dict` from that child's
107+
/// competition, so the pool is never wrapped in a second dict the reader cannot unwrap.
108+
///
109+
/// @param strings the column's string values
110+
/// @param ctx encoding context supplying the arena
111+
/// @return a cascade step whose values child (index 1) is left open for compression
112+
private static CascadeStep encodeUtf8Cascade(String[] strings, EncodeContext ctx) {
113+
int n = strings.length;
114+
115+
var valueMap = new LinkedHashMap<String, Integer>();
116+
for (String s : strings) {
117+
valueMap.computeIfAbsent(s, _ -> valueMap.size());
118+
}
119+
int dictSize = valueMap.size();
120+
PType codePType = codePType(dictSize);
121+
int codeBytes = codePType.byteSize();
122+
123+
MemorySegment codesBuf = ctx.arena().allocate((long) n * codeBytes);
124+
for (int i = 0; i < n; i++) {
125+
writeCodeToSeg(codesBuf, codePType, i, valueMap.get(strings[i]));
126+
}
127+
128+
byte[] metaBytes = new ProtoDictMetadata(
129+
dictSize,
130+
io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(codePType.ordinal()),
131+
null,
132+
null
133+
).encode();
134+
135+
// Child order matches the terminal encodeUtf8: [codes, values]. Codes is a raw primitive
136+
// leaf at buffer 0; values (index 1) is left null for the cascade to fill.
137+
EncodeNode codesNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 0);
138+
EncodeNode partialRoot = new EncodeNode(
139+
EncodingId.VORTEX_DICT, MemorySegment.ofArray(metaBytes),
140+
new EncodeNode[]{codesNode, null},
141+
new int[0]);
142+
143+
String[] distinct = valueMap.keySet().toArray(new String[0]);
144+
ChildSlot valuesSlot = new ChildSlot(DType.UTF8, distinct, 1);
145+
146+
String minStr = valueMap.keySet().stream().min(String::compareTo).orElse(null);
147+
String maxStr = valueMap.keySet().stream().max(String::compareTo).orElse(null);
148+
byte[] statsMin = minStr != null ? ProtoScalarValue.ofStringValue(minStr).encode() : null;
149+
byte[] statsMax = maxStr != null ? ProtoScalarValue.ofStringValue(maxStr).encode() : null;
150+
151+
return new CascadeStep(partialRoot, List.of(codesBuf), List.of(valuesSlot), statsMin, statsMax, true);
152+
}
153+
104154
private static EncodeResult encodeUtf8(String[] strings, EncodeContext ctx) {
105155
int n = strings.length;
106156

writer/src/test/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoderTest.java

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -178,4 +178,22 @@ void encode_utf8_metadata_valuesLen_matchesUniqueCount() throws Exception {
178178
// Then
179179
assertThat(meta.values_len()).isEqualTo(2);
180180
}
181+
182+
@Test
183+
void encodeCascade_utf8_exposesValuesPoolAsOpenChild() {
184+
// Given — a low-cardinality Utf8 column.
185+
String[] data = repeat(new String[]{"apple", "banana", "cherry"}, 100);
186+
187+
// When
188+
CascadeStep result = ENCODER.encodeCascade(DTypes.UTF8, data, EncodeTestHelper.testCtx());
189+
190+
// Then — the distinct-values pool (child index 1) is left open for the cascade to compress
191+
// (FSST/VarBin compete on it, #299), not hardcoded to a terminal raw-varbin node.
192+
assertThat(result.isTerminal()).isFalse();
193+
assertThat(result.openChildren()).singleElement().satisfies(slot -> {
194+
assertThat(slot.parentChildIdx()).isEqualTo(1);
195+
assertThat(slot.childDtype()).isEqualTo(DTypes.UTF8);
196+
assertThat((String[]) slot.childData()).containsExactly("apple", "banana", "cherry");
197+
});
198+
}
181199
}

0 commit comments

Comments
 (0)