From 3e186afb9770c3a42a7c92d39dac0975dc97fee8 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Wed, 22 Jul 2026 08:29:08 +0200 Subject: [PATCH] perf(fsst): close the vortex-jni gap on encode and decode hot paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Compressor: single unaligned 8-byte word loads on the fast path (byte-by-byte loadWord kept for the <8-byte tail), escape literal taken from the loaded word. LossyPerfectHashTable: packed-int lookup over a two-longs-per-slot table (the paper's C layout) — one cache line per lookup instead of three parallel-array reads; empty slots self-answer no-match with no occupancy flag or sentinel. ShortCodeTable/Matcher: pre-packed no-match slots, one array read per fallback. FsstEncodingEncoder: all rows compress into one shared scratch (was two heap allocations plus an extra copy per row), one-pass wire-code remap. FsstEncodingDecoder: prefix-sum output offsets from the lengths child and 256-row batched decompression (per-row switch/modulo removed; whole-chunk single call measured 1.8x slower via OSR-compiled mega-loop), plus new malformed-input guards (empty children, invalid offsets, row-count and decoded-length overflow, decoded-vs-claimed mismatch) with negative tests. JavaVsJniFsstBenchmark -f 3: encode 1.848 -> 3.059 +- 0.430 ops/s (jni 2.854, parity), decode 27.137 -> 32.924 +- 0.824 ops/s (jni 28.353, 1.16x faster). Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 9 + TODO.md | 11 +- .../github/dfa1/vortex/fsst/Compressor.java | 57 +++- .../vortex/fsst/LossyPerfectHashTable.java | 93 ++---- .../io/github/dfa1/vortex/fsst/Matcher.java | 12 +- .../dfa1/vortex/fsst/ShortCodeTable.java | 30 +- .../fsst/LossyPerfectHashTableTest.java | 36 +-- .../reader/decode/FsstEncodingDecoder.java | 133 +++++++- .../decode/FsstEncodingDecoderTest.java | 305 ++++++++++++++++++ .../writer/encode/FsstEncodingEncoder.java | 35 +- 10 files changed, 592 insertions(+), 129 deletions(-) create mode 100644 reader/src/test/java/io/github/dfa1/vortex/reader/decode/FsstEncodingDecoderTest.java diff --git a/CHANGELOG.md b/CHANGELOG.md index c2db3c57..b117b044 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 Measured with `JavaVsJniFsstBenchmark` ([1003a673](https://github.com/dfa1/vortex-java/commit/1003a673)), unmodified before and after. +- `vortex.fsst` hot paths close the remaining `vortex-jni` gap on both sides: single 8-byte word loads and a one-cache-line packed hash-slot layout in the compressor, one shared scratch buffer in the encoder (was two heap allocations per row), and prefix-sum output offsets plus 256-row batched decompression in the decoder (was a per-row offset `switch` and two broadcast modulos per row). ([#300](https://github.com/dfa1/vortex-java/pull/300)) + + | Benchmark | Before | After | `vortex-java` vs `vortex-jni` | + |---|---|---|---| + | `javaFsstEncode` | 1.848 ops/s | 3.059 ± 0.430 ops/s | 1.6x slower → parity (1.07x, overlapping error) | + | `javaFsstDecode` | 27.137 ops/s | 32.924 ± 0.824 ops/s | 1.3x slower → 1.16x faster | + + Measured with `JavaVsJniFsstBenchmark` `-f 3` (jni: encode 2.854 ± 0.289, decode 28.353 ± 5.283 ops/s on the same machine and run). + ### Fixed - `ParquetImporter` detects duplicate column names in the source Parquet schema and throws a clear message naming the duplicate(s) and the source file, instead of a confusing `VortexWriter` internal-invariant error several frames removed from the actual cause. ([f2c05e57](https://github.com/dfa1/vortex-java/commit/f2c05e57)) diff --git a/TODO.md b/TODO.md index 58ca09b4..efec5196 100644 --- a/TODO.md +++ b/TODO.md @@ -15,13 +15,14 @@ ### FSST follow-ups -Out of scope for the #287 rewrite (which closed most of the `vortex-jni` gap with a scalar, -branch-free algorithm — see [ADR-0022](adr/0022-fsst-module-extraction.md)): +Out of scope for the #287 rewrite (which, together with the follow-up hot-path pass, closed the +`vortex-jni` gap with a scalar, branch-free algorithm — encode at parity, decode ~1.16x faster on +`JavaVsJniFsstBenchmark` — see [ADR-0022](adr/0022-fsst-module-extraction.md)): - [ ] **AVX512/SIMD compression kernel** — the paper measures a SIMD kernel as the fastest known - string compressor at that tier, but this project's scalar rewrite already narrowed the encode gap - to `vortex-jni` from 36x to 1.6x. Needs a Vector-API/JDK-incubator decision and its own ADR (cf. - [ADR-0005](adr/0005-vector-api-adoption.md)). + string compressor at that tier. The scalar path now matches `vortex-jni` (itself scalar), so a + SIMD kernel is about beating the reference, not catching it. Needs a Vector-API/JDK-incubator + decision and its own ADR (cf. [ADR-0005](adr/0005-vector-api-adoption.md)). - [ ] **True per-row lazy/random-access decompression** exploiting FSST's headline random-access property — today `FsstEncodingDecoder.decode()` eagerly materializes the whole column up front regardless of what is queried. Connects to [ADR-0010](adr/0010-lazy-decode.md) (Lazy decode) but diff --git a/fsst/src/main/java/io/github/dfa1/vortex/fsst/Compressor.java b/fsst/src/main/java/io/github/dfa1/vortex/fsst/Compressor.java index 192020af..2aa2e47a 100644 --- a/fsst/src/main/java/io/github/dfa1/vortex/fsst/Compressor.java +++ b/fsst/src/main/java/io/github/dfa1/vortex/fsst/Compressor.java @@ -2,6 +2,9 @@ import java.lang.foreign.MemorySegment; import java.lang.foreign.ValueLayout; +import java.lang.invoke.MethodHandles; +import java.lang.invoke.VarHandle; +import java.nio.ByteOrder; import java.util.ArrayList; import java.util.List; @@ -23,7 +26,13 @@ public final class Compressor { /// `core`'s `VortexFormat.LE_LONG` so this module keeps its zero dependency on `core`; the FSST /// wire format packs symbol bytes LSB-first, so all `MemorySegment` reads/writes are little-endian. static final ValueLayout.OfLong LE_LONG = - ValueLayout.JAVA_LONG_UNALIGNED.withOrder(java.nio.ByteOrder.LITTLE_ENDIAN); + ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN); + + /// Little-endian unaligned `long` view over a `byte[]`, the JIT-intrinsified way to load an + /// 8-byte input word in one instruction on the `byte[]` compress fast path (the byte-by-byte + /// [#loadWord(byte[], int, int)] assembly is 8 loads + shifts and is kept for the tail only). + private static final VarHandle LONG_LE_BYTES = + MethodHandles.byteArrayViewVarHandle(long[].class, ByteOrder.LITTLE_ENDIAN); private final List symbolsByGainDescending; private final Matcher matcher; @@ -142,8 +151,27 @@ public Decompressor toDecompressor() { public long compress(byte[] input, int start, int end, byte[] out, long outPos) { int pos = start; int outIndex = (int) outPos; + // Fast path: while 8 real input bytes remain, load the word in one intrinsified read. No + // over-long-match guard is needed here — a match is at most 8 bytes, so it can never reach + // past end while pos <= end - 8. + int fastEnd = end - 8; + while (pos <= fastEnd) { + long word = (long) LONG_LE_BYTES.get(input, pos); + int packedMatch = matcher.longestMatch(word); + int length = Matcher.lengthOf(packedMatch); + if (length > 0) { + out[outIndex++] = (byte) Matcher.codeOf(packedMatch); + pos += length; + } else { + // The escaped literal is the word's low byte — already loaded, no re-read. + out[outIndex++] = (byte) ESCAPE; + out[outIndex++] = (byte) word; + pos++; + } + } while (pos < end) { - int packedMatch = matcher.longestMatch(loadWord(input, pos, end)); + long word = loadWord(input, pos, end); + int packedMatch = matcher.longestMatch(word); int length = Matcher.lengthOf(packedMatch); // Reject an over-long match: loadWord zero-pads bytes past end, and the branch-free // Matcher has no notion of end, so a symbol whose trailing bytes are zero can spuriously @@ -156,7 +184,7 @@ public long compress(byte[] input, int start, int end, byte[] out, long outPos) pos += length; } else { out[outIndex++] = (byte) ESCAPE; - out[outIndex++] = input[pos]; + out[outIndex++] = (byte) word; pos++; } } @@ -187,8 +215,27 @@ public long compress(byte[] input, int start, int end, byte[] out, long outPos) public long compress(MemorySegment input, long start, long end, MemorySegment out, long outPos) { long pos = start; long outIndex = outPos; + // Fast path: while 8 readable bytes remain (bounded by both end and the segment's own + // size — the segment may end exactly at end with no trailing slack), load the word in one + // unaligned read. No over-long-match guard is needed here — a match is at most 8 bytes. + long fastEnd = Math.min(end, input.byteSize()) - 8; + while (pos <= fastEnd) { + long word = input.get(LE_LONG, pos); + int packedMatch = matcher.longestMatch(word); + int length = Matcher.lengthOf(packedMatch); + if (length > 0) { + out.set(ValueLayout.JAVA_BYTE, outIndex++, (byte) Matcher.codeOf(packedMatch)); + pos += length; + } else { + // The escaped literal is the word's low byte — already loaded, no re-read. + out.set(ValueLayout.JAVA_BYTE, outIndex++, (byte) ESCAPE); + out.set(ValueLayout.JAVA_BYTE, outIndex++, (byte) word); + pos++; + } + } while (pos < end) { - int packedMatch = matcher.longestMatch(loadWord(input, pos, end)); + long word = loadWord(input, pos, end); + int packedMatch = matcher.longestMatch(word); int length = Matcher.lengthOf(packedMatch); // Same boundary guard as the byte[] overload: reject a match that reaches past the real // input end, since loadWord zero-pads and the branch-free Matcher has no notion of end. @@ -197,7 +244,7 @@ public long compress(MemorySegment input, long start, long end, MemorySegment ou pos += length; } else { out.set(ValueLayout.JAVA_BYTE, outIndex++, (byte) ESCAPE); - out.set(ValueLayout.JAVA_BYTE, outIndex++, input.get(ValueLayout.JAVA_BYTE, pos)); + out.set(ValueLayout.JAVA_BYTE, outIndex++, (byte) word); pos++; } } diff --git a/fsst/src/main/java/io/github/dfa1/vortex/fsst/LossyPerfectHashTable.java b/fsst/src/main/java/io/github/dfa1/vortex/fsst/LossyPerfectHashTable.java index bcdf104b..91c8917d 100644 --- a/fsst/src/main/java/io/github/dfa1/vortex/fsst/LossyPerfectHashTable.java +++ b/fsst/src/main/java/io/github/dfa1/vortex/fsst/LossyPerfectHashTable.java @@ -35,30 +35,21 @@ final class LossyPerfectHashTable { /// Low three bytes of the input word — the only bytes the hash keys on. private static final long PREFIX_MASK = 0x00FF_FFFFL; - /// Per-slot packed symbol bytes, LSB-first ([Symbol] convention). Meaningful only where the - /// corresponding [#occupied] entry is set. - private final long[] packedBytes; - - /// Per-slot mask `~0L >>> ignoredBits` (`ignoredBits = 64 - 8 * length`) applied to an input - /// word before comparing, clearing the high bytes past the candidate's length. - private final long[] keepMask; - - /// Per-slot symbol length in bytes, 3-8. - private final int[] lengths; - - /// Per-slot code (the symbol's list index), meaningful only where [#occupied] is set. - private final int[] codes; - - /// Whether each slot holds a candidate. - private final boolean[] occupied; - - private LossyPerfectHashTable(long[] packedBytes, long[] keepMask, int[] lengths, int[] codes, - boolean[] occupied) { - this.packedBytes = packedBytes; - this.keepMask = keepMask; - this.lengths = lengths; - this.codes = codes; - this.occupied = occupied; + /// Slot table with two adjacent longs per slot (the FSST paper's C layout): `slots[2 * s]` is + /// the candidate's packed symbol bytes (LSB-first, [Symbol] convention) and `slots[2 * s + 1]` + /// is its metadata `ignoredBits << 16 | code << 8 | length` (`ignoredBits = 64 - 8 * length`). + /// A lookup derives the keep-mask from the ignored-bits with one shift instead of loading a + /// separate mask array, so the whole 16-byte slot lands in a single cache line — one memory + /// touch per lookup instead of three parallel-array reads across three lines. + /// + /// Empty slots are all-zero: metadata 0 makes the keep-mask `~0L` and the symbol 0, so an + /// input word of exactly 0 can "hit" an empty slot — harmlessly, because the returned low 16 + /// bits (`code << 8 | length`) are then 0, which is precisely the "no match" answer. No + /// occupancy flag or sentinel is needed. + private final long[] slots; + + private LossyPerfectHashTable(long[] slots) { + this.slots = slots; } /// Builds the table from the trained symbols in descending-gain order, keeping only those of @@ -76,43 +67,41 @@ private LossyPerfectHashTable(long[] packedBytes, long[] keepMask, int[] lengths /// @param symbolsByGainDescending the trained symbols, code = list index, gain-descending /// @return a hash table resolving 3-8 byte matches with first-writer-wins on collision static LossyPerfectHashTable of(List symbolsByGainDescending) { - long[] packedBytes = new long[SLOTS]; - long[] keepMask = new long[SLOTS]; - int[] lengths = new int[SLOTS]; - int[] codes = new int[SLOTS]; - boolean[] occupied = new boolean[SLOTS]; + long[] slots = new long[2 * SLOTS]; for (int code = 0; code < symbolsByGainDescending.size(); code++) { Symbol symbol = symbolsByGainDescending.get(code); if (symbol.length() < 3) { continue; // Length 1-2 belongs to ShortCodeTable. } int slot = slotFor(symbol.packedBytes()); - if (occupied[slot]) { + if (slots[2 * slot + 1] != 0) { continue; // First writer (higher gain) wins; skip the collision. } - occupied[slot] = true; - packedBytes[slot] = symbol.packedBytes(); - keepMask[slot] = keepMaskFor(symbol.length()); - lengths[slot] = symbol.length(); - codes[slot] = code; + long ignoredBits = 64L - 8 * symbol.length(); + slots[2 * slot] = symbol.packedBytes(); + slots[2 * slot + 1] = ignoredBits << 16 | (long) code << 8 | symbol.length(); } - return new LossyPerfectHashTable(packedBytes, keepMask, lengths, codes, occupied); + return new LossyPerfectHashTable(slots); } - /// Looks up `word` and reports whether a stored 3-8 byte symbol really matches it. + /// Looks up `word` and returns the matched symbol as `code << 8 | length`, or 0 when no stored + /// 3-8 byte symbol matches. /// /// The input word's first three bytes select one slot; the candidate there matches only if the - /// masked compare passes (`(word & keepMask) == candidate.packedBytes`), which rules out both - /// hash collisions with unrelated bytes and empty slots. On a miss the returned [HashMatch] has - /// `hit == false` and the caller falls back to [ShortCodeTable]. + /// masked compare passes (`(word & (~0L >>> ignoredBits)) == candidate.packedBytes`), which + /// rules out hash collisions with unrelated bytes. An empty slot can only "hit" an all-zero + /// input word, and then still answers 0 (its metadata is 0), which is the no-match result. A + /// stored symbol's length is 3-8, so a real hit is never 0 and the caller can test the result + /// directly. On a miss the caller falls back to [ShortCodeTable]. /// /// @param word an 8-byte little-endian input word starting at the current match position, with /// any bytes past the remaining input already zero-padded by the caller - /// @return the match result: `hit`, and when hit the matched `code` and `length` - HashMatch lookup(long word) { - int slot = slotFor(word); - boolean hit = occupied[slot] && (word & keepMask[slot]) == packedBytes[slot]; - return new HashMatch(hit, codes[slot], lengths[slot]); + /// @return the match as `code << 8 | length`, or 0 when there is no match + int lookup(long word) { + int slot = slotFor(word) << 1; + long symbol = slots[slot]; + long meta = slots[slot + 1]; + return (word & (~0L >>> (int) (meta >>> 16))) == symbol ? (int) (meta & 0xFFFF) : 0; } /// Computes the slot index a word hashes to, keyed on its first three bytes. Package-visible so @@ -126,18 +115,4 @@ static int slotFor(long word) { return (int) (mixed >>> 32) & SLOT_MASK; } - private static long keepMaskFor(int length) { - int ignoredBits = 64 - 8 * length; - return ~0L >>> ignoredBits; - } - - /// Result of a [LossyPerfectHashTable#lookup(long)]: whether a 3-8 byte symbol matched and, if - /// so, its code and length. When `hit` is false, `code` and `length` are unspecified and the - /// caller must ignore them. - /// - /// @param hit whether a stored symbol really matched the input word - /// @param code the matched symbol code, valid only when `hit` - /// @param length the matched symbol length in bytes (3-8), valid only when `hit` - record HashMatch(boolean hit, int code, int length) { - } } diff --git a/fsst/src/main/java/io/github/dfa1/vortex/fsst/Matcher.java b/fsst/src/main/java/io/github/dfa1/vortex/fsst/Matcher.java index 21885bbd..03b035c4 100644 --- a/fsst/src/main/java/io/github/dfa1/vortex/fsst/Matcher.java +++ b/fsst/src/main/java/io/github/dfa1/vortex/fsst/Matcher.java @@ -55,11 +55,13 @@ public static Matcher of(List symbolsByGainDescending) { /// any bytes past the remaining input already zero-padded by the caller /// @return the longest match as `code << 8 | length`; length 0 signals "no match, escape" public int longestMatch(long word) { - LossyPerfectHashTable.HashMatch hashMatch = hashTable.lookup(word); - if (hashMatch.hit()) { - return hashMatch.code() << 8 | hashMatch.length(); - } - return shortCodes.codeFor(word) << 8 | shortCodes.lengthFor(word); + // Both tables are read unconditionally so the select below has no side to skip — the JIT + // can lower it to a conditional move instead of a data-dependent branch, which matters + // because hash hit/miss alternates unpredictably across input positions. The extra + // short-code read on a hash hit is one L1 load into a 256 KB table. + int hashMatch = hashTable.lookup(word); + int shortMatch = shortCodes.packedFor(word); + return hashMatch != 0 ? hashMatch : shortMatch; } /// Extracts the code from a packed [#longestMatch(long)] result. diff --git a/fsst/src/main/java/io/github/dfa1/vortex/fsst/ShortCodeTable.java b/fsst/src/main/java/io/github/dfa1/vortex/fsst/ShortCodeTable.java index dff51251..173439fa 100644 --- a/fsst/src/main/java/io/github/dfa1/vortex/fsst/ShortCodeTable.java +++ b/fsst/src/main/java/io/github/dfa1/vortex/fsst/ShortCodeTable.java @@ -1,5 +1,6 @@ package io.github.dfa1.vortex.fsst; +import java.util.Arrays; import java.util.List; /// Direct-indexed table resolving the shortest FSST matches — length 0 (no match), 1, or 2 — @@ -25,7 +26,12 @@ final class ShortCodeTable { /// codes are `0..254` (`0xFF` is the escape), so `-1` can never collide with a real code. static final int NO_CODE = -1; - /// Packed `code << 8 | length` per 16-bit key. A zero length marks "no match". + /// Packed no-match value stored in unpopulated slots: `NO_CODE << 8 | 0`. Storing the sentinel + /// pre-packed lets [#packedFor(long)] return the slot verbatim — one array read, no branch — + /// and an arithmetic `>> 8` recovers [#NO_CODE] while `& 0xFF` recovers length 0. + private static final int NO_MATCH = NO_CODE << 8; + + /// Packed `code << 8 | length` per 16-bit key. A zero length marks "no match" ([#NO_MATCH]). private final int[] slots; private ShortCodeTable(int[] slots) { @@ -46,6 +52,7 @@ private ShortCodeTable(int[] slots) { /// @return a table resolving 0/1/2-byte matches for any two-byte input prefix static ShortCodeTable of(List symbolsByGainDescending) { int[] slots = new int[SLOTS]; + Arrays.fill(slots, NO_MATCH); for (int code = 0; code < symbolsByGainDescending.size(); code++) { Symbol symbol = symbolsByGainDescending.get(code); if (symbol.length() == 1) { @@ -69,23 +76,32 @@ static ShortCodeTable of(List symbolsByGainDescending) { return new ShortCodeTable(slots); } - /// Returns the code matched by the low two bytes of `word`, or [#NO_CODE] if none. One array - /// read, no branching. + /// Returns the match for the low two bytes of `word` as `code << 8 | length`, or + /// `NO_CODE << 8` (length 0) when there is no length-1 or length-2 match. One array read, no + /// branching — the slot value is returned verbatim, which is what makes this the hot path's + /// fallback of choice. + /// + /// @param word an input word; only its low 16 bits (first two input bytes) are consulted + /// @return the match as `code << 8 | length`; length 0 (and code [#NO_CODE]) means no match + int packedFor(long word) { + return slots[(int) (word & 0xFFFF)]; + } + + /// Returns the code matched by the low two bytes of `word`, or [#NO_CODE] if none. /// /// @param word an input word; only its low 16 bits (first two input bytes) are consulted /// @return the matched symbol code, or [#NO_CODE] when there is no length-1 or length-2 match int codeFor(long word) { - int packed = slots[(int) (word & 0xFFFF)]; - return packed == 0 ? NO_CODE : packed >>> 8; + return packedFor(word) >> 8; } /// Returns the length of the symbol matched by the low two bytes of `word`: 2, 1, or 0 for no - /// match. One array read, no branching. + /// match. /// /// @param word an input word; only its low 16 bits (first two input bytes) are consulted /// @return the matched symbol length in bytes, or 0 when there is no match int lengthFor(long word) { - return length(slots[(int) (word & 0xFFFF)]); + return length(packedFor(word)); } private static int length(int packed) { diff --git a/fsst/src/test/java/io/github/dfa1/vortex/fsst/LossyPerfectHashTableTest.java b/fsst/src/test/java/io/github/dfa1/vortex/fsst/LossyPerfectHashTableTest.java index 6cb59c72..29eb9f99 100644 --- a/fsst/src/test/java/io/github/dfa1/vortex/fsst/LossyPerfectHashTableTest.java +++ b/fsst/src/test/java/io/github/dfa1/vortex/fsst/LossyPerfectHashTableTest.java @@ -17,12 +17,12 @@ void lookup_realThreeByteSymbol_resolvesCode() { // When — the input word starts with "abc" and carries an unrelated trailing byte that the // masked compare must ignore (the symbol is only 3 bytes). - LossyPerfectHashTable.HashMatch result = sut.lookup(wordOf("abcZZZZZ")); + int result = sut.lookup(wordOf("abcZZZZZ")); - // Then - assertThat(result.hit()).isTrue(); - assertThat(result.code()).isEqualTo(0); - assertThat(result.length()).isEqualTo(3); + // Then — packed as code << 8 | length; a hit is never 0. + assertThat(result).isNotZero(); + assertThat(result >>> 8).isEqualTo(0); + assertThat(result & 0xFF).isEqualTo(3); } @Test @@ -32,11 +32,11 @@ void lookup_eightByteSymbol_resolvesFullWord() { LossyPerfectHashTable sut = LossyPerfectHashTable.of(List.of(full)); // When - LossyPerfectHashTable.HashMatch result = sut.lookup(wordOf("ABCDEFGH")); + int result = sut.lookup(wordOf("ABCDEFGH")); // Then - assertThat(result.hit()).isTrue(); - assertThat(result.length()).isEqualTo(8); + assertThat(result).isNotZero(); + assertThat(result & 0xFF).isEqualTo(8); } @Test @@ -48,10 +48,10 @@ void lookup_lengthOneAndTwoSymbolsSkipped() { LossyPerfectHashTable sut = LossyPerfectHashTable.of(List.of(a, ab)); // When - LossyPerfectHashTable.HashMatch result = sut.lookup(wordOf("ab")); + int result = sut.lookup(wordOf("ab")); - // Then - assertThat(result.hit()).isFalse(); + // Then — 0 is the packed "no match". + assertThat(result).isZero(); } @Test @@ -65,14 +65,14 @@ void lookup_hashCollisionInGainOrder_higherGainWinsSlotLowerGainMisses() { LossyPerfectHashTable sut = LossyPerfectHashTable.of(List.of(higherGain, lowerGain)); // When — look up each symbol's own bytes. - LossyPerfectHashTable.HashMatch higher = sut.lookup(higherGain.packedBytes()); - LossyPerfectHashTable.HashMatch lower = sut.lookup(lowerGain.packedBytes()); + int higher = sut.lookup(higherGain.packedBytes()); + int lower = sut.lookup(lowerGain.packedBytes()); // Then — the higher-gain (first-inserted) symbol wins; the lower-gain one correctly misses // and does not corrupt or masquerade as the winner. - assertThat(higher.hit()).isTrue(); - assertThat(higher.code()).isEqualTo(0); - assertThat(lower.hit()).isFalse(); + assertThat(higher).isNotZero(); + assertThat(higher >>> 8).isEqualTo(0); + assertThat(lower).isZero(); } @Test @@ -85,10 +85,10 @@ void lookup_hashCollisionButByteMismatch_reportsNoMatch() { long adversarial = findWordCollidingWithButNotEqualTo(stored.packedBytes()); // When - LossyPerfectHashTable.HashMatch result = sut.lookup(adversarial); + int result = sut.lookup(adversarial); // Then - assertThat(result.hit()).isFalse(); + assertThat(result).isZero(); } /// Finds two distinct 3-byte symbols whose first three bytes hash to the same slot, so the test diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/decode/FsstEncodingDecoder.java b/reader/src/main/java/io/github/dfa1/vortex/reader/decode/FsstEncodingDecoder.java index 6f3c6ae8..0b12d7c3 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/decode/FsstEncodingDecoder.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/decode/FsstEncodingDecoder.java @@ -23,6 +23,13 @@ /// which runs the FSST paper's Algorithm 1 decode. public final class FsstEncodingDecoder implements EncodingDecoder { + /// Rows decoded per [Decompressor#decompress] call on the fast path. Batching (rather than one + /// whole-chunk call) keeps each call's loop short enough that execution stays in the cleanly + /// compiled method entry instead of an OSR-compiled mega-loop — a single 50k-row call measured + /// ~1.8x slower per byte than short calls on the same data — while still amortizing the per-row + /// offset read down to one read per batch. + private static final int ROWS_PER_DECODE_BATCH = 256; + @Override public EncodingId encodingId() { return EncodingId.VORTEX_FSST; @@ -81,9 +88,24 @@ public Array decode(DecodeContext ctx) { } Decompressor decompressor = Decompressor.of(packedSymbols, symbolLengths); - long totalUncompressed = 0L; - for (long i = 0; i < n; i++) { - totalUncompressed += readUnsigned(uncompLensSeg, i % uncompLensCap, uncompLenPType); + if (n >= Integer.MAX_VALUE) { + throw new VortexException(EncodingId.VORTEX_FSST, "row count too large: " + n); + } + + // The decoded row boundaries are fully determined by the uncompressed-lengths child, so the + // output offsets are its prefix sums — computed in one per-ptype loop (no switch, no + // modulo in the body; hot-loop rule) that also yields the total for sizing the output. + MemorySegment outOffsets = ctx.arena().allocate((n + 1) * 4L, 4); + outOffsets.setAtIndex(VortexFormat.LE_INT, 0, 0); + long totalUncompressed = writeDecodedOffsets(uncompLensSeg, outOffsets, n, uncompLenPType, uncompLensCap); + + // The output offsets are I32 prefix sums, so a decoded chunk past 2 GB would silently wrap + // (and adversarial I32/I64-typed lengths can make the running sum go negative). Reject both + // before sizing/allocating outBytes, so a hostile length child can never drive a wrapped or + // enormous allocation. + if (totalUncompressed > Integer.MAX_VALUE || totalUncompressed < 0) { + throw new VortexException(EncodingId.VORTEX_FSST, + "decoded length too large: " + totalUncompressed); } // Allocate 7 bytes of slack past the true logical length: the decompressor's unconditional @@ -91,21 +113,110 @@ public Array decode(DecodeContext ctx) { // few as 1 real byte. The slack is sliced off before the buffer is exposed, so callers still // see an exactly-sized buffer. MemorySegment outBytes = ctx.arena().allocate(totalUncompressed + 7); - MemorySegment outOffsets = ctx.arena().allocate((n + 1) * 4L, 4); - outOffsets.setAtIndex(VortexFormat.LE_INT, 0, 0); - long outPos = 0L; - for (long i = 0; i < n; i++) { - long cStart = readUnsigned(codesOffsetsSeg, i % codesOffCap, codesOffPType); - long cEnd = readUnsigned(codesOffsetsSeg, (i + 1) % codesOffCap, codesOffPType); - outPos = decompressor.decompress(compressedBytes, cStart, cEnd, outBytes, outPos); - outOffsets.setAtIndex(VortexFormat.LE_INT, i + 1, (int) outPos); + if (codesOffCap == 0) { + // No physical offsets at all: only valid for a zero-row chunk, where there is nothing + // to decode. + if (n > 0) { + throw new VortexException(EncodingId.VORTEX_FSST, "empty codes-offsets child"); + } + } else if (codesOffCap >= n + 1 || codesOffCap == 1) { + // Fast path. Row i's code range is [offsets[i], offsets[i+1]) out of ONE shared offsets + // array, so consecutive rows are contiguous by construction and the code stream decodes + // in row batches — only every ROWS_PER_DECODE_BATCH-th offset is read (no per-row + // offset reads, no per-row switch/modulo). codesOffCap == 1 is the constant-broadcast + // child: every row's range is the same empty [off, off). + long firstOffset = readUnsigned(codesOffsetsSeg, 0, codesOffPType); + long lastOffset = readUnsigned(codesOffsetsSeg, Math.min(n, codesOffCap - 1), codesOffPType); + if (firstOffset > lastOffset || lastOffset > compressedBytes.byteSize()) { + throw new VortexException(EncodingId.VORTEX_FSST, "invalid code offsets: [" + + firstOffset + ", " + lastOffset + ") of " + compressedBytes.byteSize()); + } + long outPos = 0L; + long batchStart = firstOffset; + for (long i = 0; i < n; i += ROWS_PER_DECODE_BATCH) { + long batchEndRow = Math.min(i + ROWS_PER_DECODE_BATCH, n); + long batchEnd = batchEndRow == n + ? lastOffset + : readUnsigned(codesOffsetsSeg, batchEndRow, codesOffPType); + outPos = decompressor.decompress(compressedBytes, batchStart, batchEnd, outBytes, outPos); + batchStart = batchEnd; + } + if (outPos != totalUncompressed) { + throw new VortexException(EncodingId.VORTEX_FSST, "decoded " + outPos + + " bytes but uncompressed lengths claim " + totalUncompressed); + } + } else { + // Defensive broadcast path (1 < physical offsets < n + 1): ranges wrap around the + // physical elements, so decode row by row and record the offsets the decode actually + // produced, overwriting the claimed prefix sums. + long outPos = 0L; + for (long i = 0; i < n; i++) { + long cStart = readUnsigned(codesOffsetsSeg, i % codesOffCap, codesOffPType); + long cEnd = readUnsigned(codesOffsetsSeg, (i + 1) % codesOffCap, codesOffPType); + outPos = decompressor.decompress(compressedBytes, cStart, cEnd, outBytes, outPos); + outOffsets.setAtIndex(VortexFormat.LE_INT, i + 1, (int) outPos); + } } return new VarBinArray.OffsetMode(ctx.dtype(), n, outBytes.asSlice(0, totalUncompressed).asReadOnly(), outOffsets.asReadOnly(), PType.I32); } + /// Writes the prefix sums of the `count` unsigned per-row lengths in `seg` into `outOffsets` + /// (I32 slots `1 .. count`; slot 0 is the caller's) and returns the total. Branch-split per + /// ptype and on the broadcast case so the per-element loop bodies carry no switch and no + /// modulo (hot-loop rule). + private static long writeDecodedOffsets(MemorySegment seg, MemorySegment outOffsets, long count, + PType ptype, long cap) { + if (cap == 0 && count > 0) { + throw new VortexException(EncodingId.VORTEX_FSST, "empty uncompressed-lengths child"); + } + long sum = 0L; + if (cap >= count) { + switch (ptype) { + case U8 -> { + for (long i = 0; i < count; i++) { + sum += Byte.toUnsignedLong(seg.get(ValueLayout.JAVA_BYTE, i)); + outOffsets.setAtIndex(VortexFormat.LE_INT, i + 1, (int) sum); + } + } + case U16 -> { + for (long i = 0; i < count; i++) { + sum += Short.toUnsignedLong(seg.get(VortexFormat.LE_SHORT, i * 2)); + outOffsets.setAtIndex(VortexFormat.LE_INT, i + 1, (int) sum); + } + } + case U32 -> { + for (long i = 0; i < count; i++) { + sum += Integer.toUnsignedLong(seg.getAtIndex(VortexFormat.LE_INT, i)); + outOffsets.setAtIndex(VortexFormat.LE_INT, i + 1, (int) sum); + } + } + case I32 -> { + for (long i = 0; i < count; i++) { + sum += seg.getAtIndex(VortexFormat.LE_INT, i); + outOffsets.setAtIndex(VortexFormat.LE_INT, i + 1, (int) sum); + } + } + case I64, U64 -> { + for (long i = 0; i < count; i++) { + sum += seg.getAtIndex(VortexFormat.LE_LONG, i); + outOffsets.setAtIndex(VortexFormat.LE_INT, i + 1, (int) sum); + } + } + default -> throw new VortexException(EncodingId.VORTEX_FSST, "unsupported ptype " + ptype); + } + return sum; + } + // Broadcast slow path (constant-encoded child): only ever a handful of physical elements. + for (long i = 0; i < count; i++) { + sum += readUnsigned(seg, i % cap, ptype); + outOffsets.setAtIndex(VortexFormat.LE_INT, i + 1, (int) sum); + } + return sum; + } + private static long readUnsigned(MemorySegment seg, long idx, PType ptype) { return switch (ptype) { case U8 -> Byte.toUnsignedLong(seg.get(ValueLayout.JAVA_BYTE, idx)); diff --git a/reader/src/test/java/io/github/dfa1/vortex/reader/decode/FsstEncodingDecoderTest.java b/reader/src/test/java/io/github/dfa1/vortex/reader/decode/FsstEncodingDecoderTest.java new file mode 100644 index 00000000..5fff005b --- /dev/null +++ b/reader/src/test/java/io/github/dfa1/vortex/reader/decode/FsstEncodingDecoderTest.java @@ -0,0 +1,305 @@ +package io.github.dfa1.vortex.reader.decode; + +import io.github.dfa1.vortex.core.error.VortexException; +import io.github.dfa1.vortex.core.io.VortexFormat; +import io.github.dfa1.vortex.core.model.DType; +import io.github.dfa1.vortex.core.model.EncodingId; +import io.github.dfa1.vortex.core.model.PType; +import io.github.dfa1.vortex.core.proto.ProtoFSSTMetadata; +import io.github.dfa1.vortex.core.proto.ProtoPType; +import io.github.dfa1.vortex.reader.ReadRegistry; +import io.github.dfa1.vortex.reader.array.VarBinArray; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; +import java.nio.charset.StandardCharsets; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatExceptionOfType; + +/// Unit tests for [FsstEncodingDecoder], focused on the untrusted-input guards that turn a +/// malformed `vortex.fsst` node into a [VortexException] rather than an +/// `ArrayIndexOutOfBoundsException`, `NegativeArraySizeException`, or a wrapped/oversized +/// allocation (the reader parses untrusted binary input; see the security contract in CLAUDE.md). +/// +/// The wire node shape mirrors [io.github.dfa1.vortex.writer.encode.FsstEncodingEncoder]: 3 buffers +/// (symbols — one LE long per symbol; symbol lengths — one byte per symbol; compressed code stream) +/// plus 2 primitive children (uncompressed lengths, `n` elements; code offsets, `n + 1` elements) +/// plus a [ProtoFSSTMetadata] naming the two child ptypes. The children are `vortex.primitive` +/// leaves, so the test registry registers [PrimitiveEncodingDecoder] alongside the SUT. +class FsstEncodingDecoderTest { + + private static final FsstEncodingDecoder SUT = new FsstEncodingDecoder(); + private static final ReadRegistry REGISTRY = + TestRegistry.ofDecoders(SUT, new PrimitiveEncodingDecoder()); + + /// Escape code copied verbatim by the decompressor (see [io.github.dfa1.vortex.fsst.Decompressor]). + private static final int ESCAPE = 0xFF; + + @Test + void encodingId_isVortexFsst() { + // Given / When / Then + assertThat(SUT.encodingId()).isEqualTo(EncodingId.VORTEX_FSST); + } + + @Nested + class HappyPath { + + @Test + void decodesTinyNode_provingHarnessIsSound() { + // Given — a symbol table of one symbol "ab" (wire code 0) plus a stream that uses both + // that symbol and the escape path, so the round-trip exercises the fast decode loop, the + // symbol store, and the escape branch at once. Row 0 = "ab" (code [0]); row 1 = "abZ" + // (codes [0, ESCAPE, 'Z']) — 'Z' is not a symbol, so it must round-trip via the escape. + long[] symbols = {packSymbol("ab")}; + byte[] symbolLengths = {2}; + byte[] compressed = {0x00, 0x00, (byte) ESCAPE, 'Z'}; + long[] uncompLengths = {2, 3}; // "ab" is 2 bytes, "abZ" is 3 + long[] codeOffsets = {0, 1, 4}; // row 0 = [0,1), row 1 = [1,4) + + // When + VarBinArray result = decodeFsst(2, symbols, symbolLengths, compressed, + PType.U8, uncompLengths, PType.U8, codeOffsets); + + // Then + assertThat(result.length()).isEqualTo(2); + assertThat(result.getString(0)).isEqualTo("ab"); + assertThat(result.getString(1)).isEqualTo("abZ"); + } + } + + @Nested + class Guards { + + @Test + void rowCountAtIntMax_throws() { + // Given — n == Integer.MAX_VALUE overflows the I32 offset array sizing; the guard rejects + // it before any allocation. Empty buffers/children suffice because the check fires first. + long[] symbols = {}; + byte[] symbolLengths = {}; + byte[] compressed = {}; + // Broadcast (single-element) children keep the fixtures tiny while claiming a huge n. + long[] uncompLengths = {0}; + long[] codeOffsets = {0}; + + // When / Then + assertThatExceptionOfType(VortexException.class) + .isThrownBy(() -> decodeFsst(Integer.MAX_VALUE, symbols, symbolLengths, compressed, + PType.U8, uncompLengths, PType.U8, codeOffsets)) + .withMessageContaining("row count too large"); + } + + @Test + void emptyUncompressedLengthsChild_throws() { + // Given — n > 0 but the uncompressed-lengths child has zero physical elements. There is + // no per-row length to read, so the prefix-sum loop cannot proceed. + long[] symbols = {packSymbol("ab")}; + byte[] symbolLengths = {2}; + byte[] compressed = {0x00}; + long[] uncompLengths = {}; // empty length child with n == 1 + long[] codeOffsets = {0, 1}; + + // When / Then + assertThatExceptionOfType(VortexException.class) + .isThrownBy(() -> decodeFsst(1, symbols, symbolLengths, compressed, + PType.U8, uncompLengths, PType.U8, codeOffsets)) + .withMessageContaining("empty uncompressed-lengths child"); + } + + @Test + void emptyCodesOffsetsChild_throws() { + // Given — n > 0 but the code-offsets child is empty. Without offsets there is no code + // range for any row, so the fast path's codesOffCap == 0 branch must reject it. + long[] symbols = {packSymbol("ab")}; + byte[] symbolLengths = {2}; + byte[] compressed = {0x00}; + long[] uncompLengths = {2}; + long[] codeOffsets = {}; // empty offsets child with n == 1 + + // When / Then + assertThatExceptionOfType(VortexException.class) + .isThrownBy(() -> decodeFsst(1, symbols, symbolLengths, compressed, + PType.U8, uncompLengths, PType.U8, codeOffsets)) + .withMessageContaining("empty codes-offsets child"); + } + + @Test + void codeOffsetsDescending_throws() { + // Given — firstOffset > lastOffset. A descending offset pair would make the code range + // negative-length; the guard must reject it rather than let the decompressor read wild. + long[] symbols = {packSymbol("ab")}; + byte[] symbolLengths = {2}; + byte[] compressed = {0x00, 0x00}; + long[] uncompLengths = {2}; + long[] codeOffsets = {2, 0}; // first (2) > last (0) + + // When / Then + assertThatExceptionOfType(VortexException.class) + .isThrownBy(() -> decodeFsst(1, symbols, symbolLengths, compressed, + PType.U8, uncompLengths, PType.U8, codeOffsets)) + .withMessageContaining("invalid code offsets"); + } + + @Test + void codeOffsetsPastCompressedBuffer_throws() { + // Given — lastOffset exceeds the compressed buffer's size, so the code range points past + // the end of the stream. The guard must reject it before the decompressor reads OOB. + long[] symbols = {packSymbol("ab")}; + byte[] symbolLengths = {2}; + byte[] compressed = {0x00}; // only 1 code byte physically present + long[] uncompLengths = {2}; + long[] codeOffsets = {0, 5}; // claims 5 code bytes, buffer has 1 + + // When / Then + assertThatExceptionOfType(VortexException.class) + .isThrownBy(() -> decodeFsst(1, symbols, symbolLengths, compressed, + PType.U8, uncompLengths, PType.U8, codeOffsets)) + .withMessageContaining("invalid code offsets"); + } + + @Test + void decodedLengthMismatchesClaim_throws() { + // Given — a valid symbol table plus a code stream that decodes to FEWER bytes than the + // uncompressed-lengths child claims: the stream is a single "ab" symbol (2 bytes) but the + // length child claims 5. The fast-path post-decode check must catch the disagreement. + long[] symbols = {packSymbol("ab")}; + byte[] symbolLengths = {2}; + byte[] compressed = {0x00}; // decodes to exactly "ab" (2 bytes) + long[] uncompLengths = {5}; // but the child claims 5 bytes + long[] codeOffsets = {0, 1}; + + // When / Then + assertThatExceptionOfType(VortexException.class) + .isThrownBy(() -> decodeFsst(1, symbols, symbolLengths, compressed, + PType.U8, uncompLengths, PType.U8, codeOffsets)) + .withMessageContaining("uncompressed lengths claim"); + } + + @Test + void totalUncompressedOverflowsIntMax_throws() { + // Given — an I64 length child whose few huge values sum past Integer.MAX_VALUE. The + // output offsets are I32 prefix sums, so this would silently wrap; the step-1 guard must + // reject it. Placed BEFORE the outBytes allocation, so the test must fail without ever + // attempting a multi-gigabyte allocation. Two rows of ~1.5 GB each overflow I32. + long huge = 1_500_000_000L; + long[] symbols = {packSymbol("ab")}; + byte[] symbolLengths = {2}; + byte[] compressed = {0x00, 0x00}; + long[] uncompLengths = {huge, huge}; // sum == 3e9 > Integer.MAX_VALUE + long[] codeOffsets = {0, 1, 2}; + + // When / Then + assertThatExceptionOfType(VortexException.class) + .isThrownBy(() -> decodeFsst(2, symbols, symbolLengths, compressed, + PType.I64, uncompLengths, PType.U8, codeOffsets)) + .withMessageContaining("decoded length too large"); + } + + @Test + void negativeTotalUncompressed_throws() { + // Given — an adversarial I32 length whose value is negative; the running prefix sum goes + // negative, which the step-1 guard must reject (a negative size would otherwise flow into + // the allocation and throw a raw exception instead of a sanitized VortexException). + long[] symbols = {packSymbol("ab")}; + byte[] symbolLengths = {2}; + byte[] compressed = {0x00}; + long[] uncompLengths = {-1}; // I32 -1 => sum == -1 + long[] codeOffsets = {0, 1}; + + // When / Then + assertThatExceptionOfType(VortexException.class) + .isThrownBy(() -> decodeFsst(1, symbols, symbolLengths, compressed, + PType.I32, uncompLengths, PType.U8, codeOffsets)) + .withMessageContaining("decoded length too large"); + } + } + + // ── decode harness ───────────────────────────────────────────────────────── + + /// Builds and decodes a `vortex.fsst` node from the raw component arrays. + /// + /// Segment layout matches the encoder: buffers 0-2 are the FSST node's own (symbols, symbol + /// lengths, compressed stream); buffers 3-4 back the two primitive children (uncompressed + /// lengths, code offsets). The children's ptypes are carried in the [ProtoFSSTMetadata]. + private static VarBinArray decodeFsst(long rowCount, long[] symbols, byte[] symbolLengths, + byte[] compressed, PType uncompLenPType, long[] uncompLengths, + PType codesOffPType, long[] codeOffsets) { + MemorySegment symbolsSeg = leLongs(symbols); + MemorySegment symbolLensSeg = bytes(symbolLengths); + MemorySegment compressedSeg = bytes(compressed); + MemorySegment uncompLensSeg = typedSegment(uncompLenPType, uncompLengths); + MemorySegment codeOffsetsSeg = typedSegment(codesOffPType, codeOffsets); + MemorySegment[] segs = {symbolsSeg, symbolLensSeg, compressedSeg, uncompLensSeg, codeOffsetsSeg}; + + MemorySegment meta = MemorySegment.ofArray(new ProtoFSSTMetadata( + ProtoPType.fromValue(uncompLenPType.ordinal()), + ProtoPType.fromValue(codesOffPType.ordinal())).encode()); + + ArrayNode uncompLensNode = primitiveNode(3); + ArrayNode codeOffsetsNode = primitiveNode(4); + ArrayNode fsstNode = new ArrayNode(EncodingId.VORTEX_FSST, meta, + new ArrayNode[]{uncompLensNode, codeOffsetsNode}, new int[]{0, 1, 2}); + + DecodeContext ctx = new DecodeContext(fsstNode, DType.UTF8, rowCount, segs, REGISTRY, + Arena.ofAuto()); + return (VarBinArray) SUT.decode(ctx); + } + + private static ArrayNode primitiveNode(int bufferIndex) { + return new ArrayNode(EncodingId.VORTEX_PRIMITIVE, null, new ArrayNode[0], new int[]{bufferIndex}); + } + + // ── segment builders ─────────────────────────────────────────────────────── + + /// Packs a symbol string's UTF-8 bytes LSB-first into a `long` (byte 0 in the low byte), the + /// [io.github.dfa1.vortex.fsst.Symbol] wire convention. + private static long packSymbol(String s) { + byte[] b = s.getBytes(StandardCharsets.UTF_8); + long value = 0; + for (int k = 0; k < b.length; k++) { + value |= Byte.toUnsignedLong(b[k]) << (k * 8); + } + return value; + } + + private static MemorySegment leLongs(long[] values) { + MemorySegment seg = Arena.ofAuto().allocate(Math.max(values.length * 8L, 1), 8); + for (int i = 0; i < values.length; i++) { + seg.setAtIndex(VortexFormat.LE_LONG, i, values[i]); + } + return seg; + } + + private static MemorySegment bytes(byte[] values) { + MemorySegment seg = Arena.ofAuto().allocate(Math.max(values.length, 1)); + for (int i = 0; i < values.length; i++) { + seg.set(ValueLayout.JAVA_BYTE, i, values[i]); + } + return seg; + } + + /// Writes `values` into a segment at the stride of `ptype`. Only the ptypes the FSST children + /// can carry (U8, I32, I64) are needed by these tests. An empty `values` yields a zero-length + /// segment (not a 1-byte floor) so the child decodes with capacity 0 — the shape the + /// empty-child guards inspect. + private static MemorySegment typedSegment(PType ptype, long[] values) { + long stride = ptype.byteSize(); + if (values.length == 0) { + return MemorySegment.ofArray(new byte[0]); + } + MemorySegment seg = Arena.ofAuto().allocate(values.length * stride, stride); + for (int i = 0; i < values.length; i++) { + switch (ptype) { + case U8, I8 -> seg.set(ValueLayout.JAVA_BYTE, i, (byte) values[i]); + case U16, I16 -> seg.set(VortexFormat.LE_SHORT, i * 2, (short) values[i]); + case U32, I32 -> seg.setAtIndex(VortexFormat.LE_INT, i, (int) values[i]); + case I64, U64 -> seg.setAtIndex(VortexFormat.LE_LONG, i, values[i]); + default -> throw new IllegalArgumentException("unsupported ptype: " + ptype); + } + } + return seg; + } +} diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoder.java index fac11207..50314526 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoder.java @@ -58,8 +58,12 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { int n = strings.length; byte[][] byteArrays = new byte[n][]; + long totalInput = 0; + int maxUncompLen = 0; for (int i = 0; i < n; i++) { byteArrays[i] = strings[i].getBytes(StandardCharsets.UTF_8); + totalInput += byteArrays[i].length; + maxUncompLen = Math.max(maxUncompLen, byteArrays[i].length); } Compressor compressor = new CompressorBuilder().seed(TRAINING_SAMPLE_SEED).train(byteArrays); @@ -87,28 +91,23 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { symLenBuf.set(ValueLayout.JAVA_BYTE, i, (byte) compressor.symbolLength(internalCode)); } - // Compress each row into a scratch buffer holding the compressor's internal codes, then - // remap only the code bytes (never the literal byte following an escape) to wire codes. - byte[][] compressed = new byte[n][]; + // Compress every row back-to-back into one shared scratch buffer (a per-row scratch plus a + // per-row exact-size copy costs two heap allocations and an extra copy per row — millions + // per chunk), then remap the code bytes (never the literal byte following an escape) to + // wire codes in a single pass over the whole stream. Worst case each input byte escapes to + // 2 output bytes, so 2 * totalInput bounds the entire stream. + byte[] scratch = new byte[Math.toIntExact(2 * totalInput)]; + int[] rowEnds = new int[n]; int totalCompressed = 0; - int maxUncompLen = 0; for (int i = 0; i < n; i++) { byte[] row = byteArrays[i]; - byte[] scratch = new byte[row.length * 2]; - int len = (int) compressor.compress(row, 0, row.length, scratch, 0); - remapCodesToWire(scratch, len, internalToWire); - compressed[i] = new byte[len]; - System.arraycopy(scratch, 0, compressed[i], 0, len); - totalCompressed += len; - maxUncompLen = Math.max(maxUncompLen, row.length); + totalCompressed = (int) compressor.compress(row, 0, row.length, scratch, totalCompressed); + rowEnds[i] = totalCompressed; } + remapCodesToWire(scratch, totalCompressed, internalToWire); MemorySegment compBuf = arena.allocate(Math.max(totalCompressed, 1)); - long pos = 0; - for (byte[] c : compressed) { - MemorySegment.copy(MemorySegment.ofArray(c), 0, compBuf, pos, c.length); - pos += c.length; - } + MemorySegment.copy(scratch, 0, compBuf, ValueLayout.JAVA_BYTE, 0, totalCompressed); // Narrowest ptype that fits every value: row lengths and cumulative offsets are // typically far below the 4-byte ceiling this always used to pay (e.g. a 6-byte string @@ -125,11 +124,9 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { long codesOffBytes = codesOffPType.byteSize(); MemorySegment codesOffBuf = arena.allocate((long) (n + 1) * codesOffBytes); - long off = 0; PTypeIO.set(codesOffBuf, 0, codesOffPType, 0); for (int i = 0; i < n; i++) { - off += compressed[i].length; - PTypeIO.set(codesOffBuf, (i + 1) * codesOffBytes, codesOffPType, off); + PTypeIO.set(codesOffBuf, (i + 1) * codesOffBytes, codesOffPType, rowEnds[i]); } byte[] metaBytes = new ProtoFSSTMetadata(