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 new file mode 100644 index 00000000..a39e9432 --- /dev/null +++ b/fsst/src/main/java/io/github/dfa1/vortex/fsst/Compressor.java @@ -0,0 +1,173 @@ +package io.github.dfa1.vortex.fsst; + +import java.util.ArrayList; +import java.util.List; + +/// An immutable, trained FSST symbol table together with the branch-free [Matcher] built from it. +/// +/// A compressor is produced by [CompressorBuilder#train(byte[][])] and holds up to 255 symbols +/// (codes `0..254`; `0xFF` is the escape). The symbols are retained in gain-descending order, which +/// is load-bearing: [Matcher#of(List)] relies on that order for its lossy-hash first-writer-wins +/// collision rule, so a symbol's code equals its index in the retained list. +/// +/// The table can be handed straight to a [Decompressor] via [#toDecompressor()], and reordered into +/// the length-sorted permutation the `vortex.fsst` wire format wants via [#codesSortedByLength()]. +public final class Compressor { + + /// Escape code: emitted as `0xFF` followed by one literal byte. Real symbol codes are `0..254`. + static final int ESCAPE = Decompressor.ESCAPE; + + private final List symbolsByGainDescending; + private final Matcher matcher; + + private Compressor(List symbolsByGainDescending, Matcher matcher) { + this.symbolsByGainDescending = symbolsByGainDescending; + this.matcher = matcher; + } + + /// Creates a compressor over the given symbols, which must be in gain-descending order (code = + /// list index). The list is copied defensively so the compressor stays immutable. + /// + /// @param symbolsByGainDescending the trained symbols, code = list index, gain-descending + /// @return a compressor bound to that table + static Compressor of(List symbolsByGainDescending) { + List copy = List.copyOf(symbolsByGainDescending); + return new Compressor(copy, Matcher.of(copy)); + } + + /// Returns the branch-free matcher over this table, used by training's compress-count pass. + /// + /// @return the matcher built from this compressor's symbols + Matcher matcher() { + return matcher; + } + + /// Returns the number of symbols in the table (0-255). + /// + /// @return the symbol count + public int symbolCount() { + return symbolsByGainDescending.size(); + } + + /// Returns the bytes of the symbol with the given code, packed LSB-first into a `long` (the + /// [Symbol] convention). + /// + /// @param code the symbol code, in `0 .. symbolCount() - 1` + /// @return the symbol's packed bytes + public long packedSymbol(int code) { + return symbolsByGainDescending.get(code).packedBytes(); + } + + /// Returns the length in bytes (1-8) of the symbol with the given code. + /// + /// @param code the symbol code, in `0 .. symbolCount() - 1` + /// @return the symbol's length in bytes + public int symbolLength(int code) { + return symbolsByGainDescending.get(code).length(); + } + + /// Returns a permutation of the codes ordered by symbol length: all multi-byte symbols (length + /// 2-8) first in non-decreasing length order, then all length-1 symbols last. + /// + /// This is pure, wire-format-agnostic reordering logic — "sort by length, single-byte-last" — + /// with no `vortex.fsst`-specific naming or semantics baked in. The `vortex.fsst` wire adapter + /// (a separate module) happens to interpret this exact permutation as its on-wire symbol order, + /// but that interpretation lives entirely in the adapter; the algorithm module only knows how to + /// sort. Ties within a length bucket keep their gain-descending relative order (a stable sort), + /// so the permutation is deterministic. + /// + /// @return an `int[]` permutation of `0 .. symbolCount() - 1` in length-sorted, single-byte-last + /// order + public int[] codesSortedByLength() { + List codes = new ArrayList<>(symbolCount()); + for (int code = 0; code < symbolCount(); code++) { + codes.add(code); + } + codes.sort((a, b) -> { + int lengthA = symbolLength(a); + int lengthB = symbolLength(b); + boolean singleA = lengthA == 1; + boolean singleB = lengthB == 1; + if (singleA != singleB) { + return singleA ? 1 : -1; + } + return Integer.compare(lengthA, lengthB); + }); + int[] result = new int[codes.size()]; + for (int i = 0; i < result.length; i++) { + result[i] = codes.get(i); + } + return result; + } + + /// Builds a [Decompressor] over this compressor's table so the same trained symbols round-trip + /// without re-deriving the parallel arrays at the call site. + /// + /// @return a decompressor bound to this table + public Decompressor toDecompressor() { + int count = symbolCount(); + long[] packed = new long[count]; + int[] lengths = new int[count]; + for (int code = 0; code < count; code++) { + packed[code] = packedSymbol(code); + lengths[code] = symbolLength(code); + } + return Decompressor.of(packed, lengths); + } + + /// Compresses the byte range `[start, end)` of `input` into `out` using longest-match-first + /// greedy parsing, writing the code stream starting at `outPos`. + /// + /// At each position the [Matcher] resolves the longest matching symbol; a match emits its + /// one-byte code and advances by the symbol length, while a position matching no symbol emits an + /// escape (`0xFF` followed by the literal byte). The caller must size `out` so every code fits: + /// a range of `end - start` input bytes expands to at most `2 * (end - start)` output bytes (a + /// full run of escapes). This is an internal algorithm boundary, not an untrusted-input + /// boundary, so no defensive bounds checks are performed here. + /// + /// @param input the raw bytes to compress + /// @param start the index of the first byte to compress, inclusive + /// @param end the index one past the last byte to compress, exclusive + /// @param out the destination buffer for the code stream + /// @param outPos the index in `out` at which to start writing + /// @return the index in `out` one past the last byte written + public long compress(byte[] input, int start, int end, byte[] out, long outPos) { + int pos = start; + int outIndex = (int) outPos; + while (pos < end) { + int packedMatch = matcher.longestMatch(loadWord(input, pos, end)); + 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 + // satisfy the masked compare against the padding and report a match longer than the real + // remaining input. Trusting it would emit one code that decodes to more bytes than were + // compressed, silently corrupting the tail of any row (NUL-containing or binary data + // especially). Falling back to an escape keeps the byte count exact. + if (length > 0 && pos + length <= end) { + out[outIndex++] = (byte) Matcher.codeOf(packedMatch); + pos += length; + } else { + out[outIndex++] = (byte) ESCAPE; + out[outIndex++] = input[pos]; + pos++; + } + } + return outIndex; + } + + /// Loads an 8-byte little-endian word from `input` starting at `pos`, zero-padding any bytes at + /// or past `end` so the [Matcher]'s masked compare never reads across the range boundary. + /// + /// @param input the source bytes + /// @param pos the first byte to load + /// @param end the exclusive range end; bytes at or after it read as zero + /// @return the 8-byte little-endian word at `pos` + static long loadWord(byte[] input, int pos, int end) { + long word = 0; + int available = Math.min(8, end - pos); + for (int k = 0; k < available; k++) { + word |= Byte.toUnsignedLong(input[pos + k]) << (k * 8); + } + return word; + } +} diff --git a/fsst/src/main/java/io/github/dfa1/vortex/fsst/CompressorBuilder.java b/fsst/src/main/java/io/github/dfa1/vortex/fsst/CompressorBuilder.java new file mode 100644 index 00000000..85df7e62 --- /dev/null +++ b/fsst/src/main/java/io/github/dfa1/vortex/fsst/CompressorBuilder.java @@ -0,0 +1,74 @@ +package io.github.dfa1.vortex.fsst; + +import java.util.List; + +/// Trains an FSST [Compressor] from a corpus of rows using the FSST paper's bottom-up algorithm +/// (Algorithm 3) plus the Rust reference's four engineering refinements (issue #287). +/// +/// Training runs [#GENERATIONS] generations starting from an empty symbol table. Each generation +/// compress-counts a growing fraction of one fixed [Sample] with the current table, then rebuilds +/// the table from the top-gain candidates (see [TrainingGeneration]). The refinements — adaptive +/// per-generation sample fraction, per-generation min-count pruning, the 8x single-byte gain boost, +/// and a final cost-based prune — are all applied inside that per-generation step; this class only +/// drives the loop and threads the seed through. +/// +/// Training is deterministic: the same rows and seed always produce a byte-identical [Compressor], +/// so an encoder built on top of it produces byte-identical output for a given input. Use the +/// builder to override the seed: +/// +/// ```java +/// Compressor compressor = new CompressorBuilder().seed(42L).train(rows); +/// ``` +public final class CompressorBuilder { + + /// Number of training generations. The FSST paper reports the table stabilizing after roughly + /// five passes; each pass can pair adjacent symbols to roughly double their length, so five + /// passes grow symbols to the full 8-byte length from an empty start. + private static final int GENERATIONS = Sample.SAMPLE_FRACTION_NUMERATORS.length; + + /// Default PRNG seed when [#seed(long)] is not called. A fixed constant (not wall-clock derived) + /// keeps training reproducible out of the box. + private static final long DEFAULT_SEED = 0x5EEDL; + + private long seed = DEFAULT_SEED; + + /// Creates a builder using the default seed until [#seed(long)] is called. + public CompressorBuilder() { + // Explicit no-arg constructor: fields default to the reproducible DEFAULT_SEED. + } + + /// Sets the PRNG seed used to draw the training sample. The same rows and seed always train a + /// byte-identical compressor. + /// + /// @param seed the sample-drawing seed + /// @return this builder, for chaining + public CompressorBuilder seed(long seed) { + this.seed = seed; + return this; + } + + /// Trains a compressor over `rows`, running the full bottom-up generation loop. + /// + /// Empty input (no rows, or every row empty) trains an empty compressor (`symbolCount() == 0`), + /// so encoding an empty column degrades to all escapes rather than failing. + /// + /// @param rows the corpus rows, each a raw byte array (never `null`; individual rows may be + /// empty) + /// @return the trained compressor + public Compressor train(byte[][] rows) { + Sample sample = Sample.draw(rows, seed); + Compressor compressor = Compressor.of(List.of()); + if (sample.chunkCount() == 0) { + return compressor; + } + for (int gen = 0; gen < GENERATIONS; gen++) { + boolean finalGeneration = gen == GENERATIONS - 1; + int chunkLimit = sample.chunkCountForGeneration(gen); + int fractionNumerator = Sample.SAMPLE_FRACTION_NUMERATORS[gen]; + List symbols = TrainingGeneration.run( + compressor, sample, chunkLimit, fractionNumerator, finalGeneration); + compressor = Compressor.of(symbols); + } + return compressor; + } +} diff --git a/fsst/src/main/java/io/github/dfa1/vortex/fsst/Sample.java b/fsst/src/main/java/io/github/dfa1/vortex/fsst/Sample.java new file mode 100644 index 00000000..ad912c28 --- /dev/null +++ b/fsst/src/main/java/io/github/dfa1/vortex/fsst/Sample.java @@ -0,0 +1,166 @@ +package io.github.dfa1.vortex.fsst; + +import java.util.ArrayList; +import java.util.List; +import java.util.Random; + +/// A fixed, byte-size-bounded training sample drawn deterministically from a corpus of rows. +/// +/// FSST training does not scan the whole input on every generation. Instead it draws one bounded +/// sample up front — targeting [#TARGET_SAMPLE_BYTES] bytes assembled from contiguous chunks of at +/// most [#MAX_CHUNK_BYTES] bytes, each taken from a randomly chosen row — and then each of the five +/// training generations replays only a growing fraction of that same sample. This replaces the old +/// `FsstEncodingEncoder`'s string-count-bounded, always-full-sample scheme, which on long-string +/// columns pulled in far more training bytes than the reference intends for no compression benefit +/// (issue #287). Sampling by bytes keeps training cost proportional to the sample size regardless of +/// how long individual rows are. +/// +/// The sample is held as a single concatenated byte array with explicit chunk boundaries, because a +/// symbol must never be counted across a chunk boundary: chunks come from unrelated rows, so a +/// "symbol" spanning two chunks would be an artifact of sampling, not a real substring of any row. +/// +/// Drawing is deterministic given a seed: the same rows and seed always produce a byte-identical +/// sample, which is what makes the whole trained table reproducible across runs (the encoder's +/// output must be byte-identical for a given input). +final class Sample { + + /// Target total sample size in bytes. The Rust reference `spiraldb/fsst` targets ~16KB; a + /// sample this size is large enough to learn a stable table while keeping every generation's + /// compress-count pass cheap. + static final int TARGET_SAMPLE_BYTES = 16 * 1024; + + /// Maximum length of a single contiguous chunk drawn from one row. Long rows contribute several + /// chunks; short rows contribute one (truncated to the row length). Bounding chunk length keeps + /// the sample spread across many rows rather than dominated by a single long one. + static final int MAX_CHUNK_BYTES = 512; + + /// Denominator for the per-generation growing sample fraction. Each generation replays + /// `SAMPLE_FRACTION_NUMERATORS[gen] / SAMPLE_FRACTION_DENOMINATOR` of the sample's chunks. + static final int SAMPLE_FRACTION_DENOMINATOR = 128; + + /// Per-generation numerators over [#SAMPLE_FRACTION_DENOMINATOR]. Early generations see only a + /// small fraction of the sample (the table is still small and low-quality, so compressing more + /// bytes buys nothing), ramping to the full sample on the last generation once the table has + /// mostly converged. These are the Rust reference's fractions: ~6%, ~30%, ~53%, ~77%, 100%. + static final int[] SAMPLE_FRACTION_NUMERATORS = {8, 38, 68, 98, 128}; + + private final byte[] bytes; + private final int[] chunkStarts; + private final int[] chunkEnds; + + private Sample(byte[] bytes, int[] chunkStarts, int[] chunkEnds) { + this.bytes = bytes; + this.chunkStarts = chunkStarts; + this.chunkEnds = chunkEnds; + } + + /// Draws a deterministic byte-size-bounded sample from `rows`. + /// + /// Rows are selected at random; each selection contributes one contiguous chunk of at most + /// [#MAX_CHUNK_BYTES] bytes, starting at a random offset within the chosen row. Empty rows are + /// skipped. Drawing stops once [#TARGET_SAMPLE_BYTES] bytes have been collected, so a corpus + /// smaller than the target is used in full. When every row is empty the sample is empty. + /// + /// @param rows the corpus rows, each a raw byte array (never `null`, but individually may be + /// empty) + /// @param seed the PRNG seed; the same rows and seed yield a byte-identical sample + /// @return a sample of up to [#TARGET_SAMPLE_BYTES] bytes split into per-row chunks + @SuppressWarnings("java:S2245") // Deterministic PRNG is the contract: training samples must be + // reproducible across builds. No security boundary here. + static Sample draw(byte[][] rows, long seed) { + long totalBytes = 0; + int nonEmptyRows = 0; + for (byte[] row : rows) { + totalBytes += row.length; + if (row.length > 0) { + nonEmptyRows++; + } + } + if (nonEmptyRows == 0) { + return new Sample(new byte[0], new int[0], new int[0]); + } + + int target = (int) Math.min(TARGET_SAMPLE_BYTES, totalBytes); + byte[] out = new byte[target]; + List starts = new ArrayList<>(); + List ends = new ArrayList<>(); + Random rng = new Random(seed); + int outPos = 0; + + // Bound the number of empty-row skips so an all-but-one-empty corpus cannot spin forever; + // once we have drawn from the one non-empty row enough times to fill the target, we stop. + while (outPos < target) { + int row = rng.nextInt(rows.length); + byte[] source = rows[row]; + if (source.length == 0) { + continue; + } + int remaining = target - outPos; + int chunkLen = Math.min(Math.min(MAX_CHUNK_BYTES, source.length), remaining); + int maxStart = source.length - chunkLen; + int start = maxStart == 0 ? 0 : rng.nextInt(maxStart + 1); + System.arraycopy(source, start, out, outPos, chunkLen); + starts.add(outPos); + ends.add(outPos + chunkLen); + outPos += chunkLen; + } + + return new Sample(out, toIntArray(starts), toIntArray(ends)); + } + + /// Returns the concatenated sample bytes. Chunk boundaries in [#chunkStart(int)] / + /// [#chunkEnd(int)] delimit the parts that came from distinct rows. + /// + /// @return the sample bytes; never `null`, possibly empty + byte[] bytes() { + return bytes; + } + + /// Returns the number of chunks the sample is split into. Generation `gen` replays only the + /// first [#chunkCountForGeneration(int)] of these. + /// + /// @return the total chunk count + int chunkCount() { + return chunkStarts.length; + } + + /// Returns the start offset (inclusive) of chunk `i` within [#bytes()]. + /// + /// @param i the chunk index, in `0 .. chunkCount() - 1` + /// @return the chunk's start offset into the sample bytes + int chunkStart(int i) { + return chunkStarts[i]; + } + + /// Returns the end offset (exclusive) of chunk `i` within [#bytes()]. + /// + /// @param i the chunk index, in `0 .. chunkCount() - 1` + /// @return the chunk's end offset into the sample bytes + int chunkEnd(int i) { + return chunkEnds[i]; + } + + /// Returns how many leading chunks generation `gen` (0-based) should replay, applying that + /// generation's growing sample fraction. Generation 0 sees roughly 6% of the chunks, the last + /// generation sees them all. At least one chunk is always replayed when the sample is non-empty, + /// so no generation trains on nothing. + /// + /// @param gen the 0-based generation index, in `0 .. SAMPLE_FRACTION_NUMERATORS.length - 1` + /// @return the number of leading chunks to replay for that generation + int chunkCountForGeneration(int gen) { + int total = chunkCount(); + if (total == 0) { + return 0; + } + long scaled = (long) total * SAMPLE_FRACTION_NUMERATORS[gen] / SAMPLE_FRACTION_DENOMINATOR; + return (int) Math.max(1, Math.min(total, scaled)); + } + + private static int[] toIntArray(List values) { + int[] array = new int[values.size()]; + for (int i = 0; i < array.length; i++) { + array[i] = values.get(i); + } + return array; + } +} diff --git a/fsst/src/main/java/io/github/dfa1/vortex/fsst/TrainingGeneration.java b/fsst/src/main/java/io/github/dfa1/vortex/fsst/TrainingGeneration.java new file mode 100644 index 00000000..84ee1b19 --- /dev/null +++ b/fsst/src/main/java/io/github/dfa1/vortex/fsst/TrainingGeneration.java @@ -0,0 +1,268 @@ +package io.github.dfa1.vortex.fsst; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.PriorityQueue; + +/// One generation of FSST bottom-up training (the FSST paper's Algorithm 3): compress-count the +/// sample with the current symbol table, then build the next table from the top-gain candidates. +/// +/// Compress-count parses each sample chunk with the current [Compressor] using longest-match-first +/// greedy parsing, tallying how often each matched symbol occurs (`count1`) and how often each pair +/// of adjacent matches occurs (`count2`). Make-table then proposes a candidate for every counted +/// symbol (its own bytes) and every counted adjacent pair (their concatenation, truncated to 8 +/// bytes), ranks them by gain, and keeps the top 255. +/// +/// Beyond the bare paper this ports the Rust reference's four refinements (issue #287), each guarded +/// and documented at its site below: per-generation min-count pruning, the 8x single-byte gain +/// boost, and — on the final generation only — a cost-based prune pass. The growing sample fraction +/// (the fourth refinement) is applied by the caller via [Sample#chunkCountForGeneration(int)]. +final class TrainingGeneration { + + /// Maximum number of real symbols kept per generation (codes `0..254`; `0xFF` is the escape). + private static final int MAX_SYMBOLS = 255; + + /// Maximum symbol length in bytes, bounded by what fits in one `long`. + private static final int MAX_SYMBOL_LENGTH = 8; + + /// Single-byte gain multiplier. WHY: multiplying every length-1 candidate's gain by 8 before + /// ranking deliberately over-values frequent single bytes relative to their raw `count*length`. + /// A single byte in the table replaces a 2-byte escape on every occurrence, so keeping frequent + /// bytes suppresses escapes; without the boost a marginally-higher-gain multi-byte candidate can + /// crowd a genuinely more valuable single byte out of the 255 slots. The old encoder used plain + /// gain and had exactly this gap. + private static final int SINGLE_BYTE_GAIN_BOOST = 8; + + private TrainingGeneration() { + } + + /// Runs one training generation and returns the next generation's symbols in gain-descending + /// order (code = list index), ready to build the next [Compressor] / [Matcher]. + /// + /// @param current the current generation's compressor (an empty-table compressor for generation + /// 0, so every position escapes and bootstraps single-byte candidates) + /// @param sample the fixed training sample + /// @param chunkLimit the number of leading sample chunks this generation replays (its growing + /// fraction), from [Sample#chunkCountForGeneration(int)] + /// @param sampleFractionNumerator this generation's sample-fraction numerator over + /// [Sample#SAMPLE_FRACTION_DENOMINATOR], used for the min-count + /// floor + /// @param finalGeneration whether this is the last generation, which relaxes the min-count floor + /// to 1 and enables the cost-based final prune + /// @return the next generation's symbols, gain-descending, at most [#MAX_SYMBOLS] of them + static List run(Compressor current, Sample sample, int chunkLimit, + int sampleFractionNumerator, boolean finalGeneration) { + Counts counts = compressCount(current, sample, chunkLimit); + return makeTable(counts, sampleFractionNumerator, finalGeneration); + } + + /// Compress-counts the first `chunkLimit` chunks of `sample`: greedily parses each chunk with + /// `current`, tallying single-match counts and adjacent-pair counts. Candidates are keyed by + /// packed symbol so both the parsed symbols and the input bytes they cover are captured directly + /// from the sample (generation 0's empty table escapes everything, seeding single bytes). + private static Counts compressCount(Compressor current, Sample sample, int chunkLimit) { + Counts counts = new Counts(); + byte[] bytes = sample.bytes(); + for (int c = 0; c < chunkLimit; c++) { + int start = sample.chunkStart(c); + int end = sample.chunkEnd(c); + long previous = -1L; // No previous match within this chunk yet. + int pos = start; + while (pos < end) { + int matchLength = matchLengthAt(current, bytes, pos, end); + long packed = Compressor.loadWord(bytes, pos, pos + matchLength) & lengthMask(matchLength); + counts.bumpSingle(packed, matchLength); + if (previous >= 0) { + counts.bumpPair(previous, packed); + } + previous = packed; + pos += matchLength; + } + } + return counts; + } + + /// Returns the length of the current table's longest match at `pos`, or 1 when nothing matches + /// (an escaped single byte still advances one position and counts as a length-1 candidate). + /// + /// An over-long match is rejected the same way [Compressor#compress(byte[], int, int, byte[], + /// long)] rejects it: loadWord zero-pads past `end`, so a symbol with trailing zero bytes can + /// spuriously match the padding beyond the chunk. Counting such a match would tally a candidate + /// spanning bytes that are not really in the sample, and could advance `pos` past `end`. + private static int matchLengthAt(Compressor current, byte[] bytes, int pos, int end) { + int packedMatch = current.matcher().longestMatch(Compressor.loadWord(bytes, pos, end)); + int length = Matcher.lengthOf(packedMatch); + return length > 0 && pos + length <= end ? length : 1; + } + + /// Builds the next table from the counted candidates: proposes each single symbol and each + /// adjacent-pair concatenation, prunes by min-count, boosts single-byte gain, ranks, applies the + /// final cost prune on the last generation, and keeps the top [#MAX_SYMBOLS] by gain. + private static List makeTable(Counts counts, int sampleFractionNumerator, + boolean finalGeneration) { + // Refinement 2 — per-generation min-count pruning. A candidate seen fewer times than this + // floor is noise at this generation's sample coverage and is dropped before its gain is even + // computed. The floor scales with the sample fraction so early, low-coverage generations + // prune more aggressively; the final generation relaxes it to 1 so rare-but-real candidates + // survive the last, authoritative ranking pass. + int minCount = finalGeneration + ? 1 + : Math.max(1, 5 * sampleFractionNumerator / Sample.SAMPLE_FRACTION_DENOMINATOR); + + List candidates = new ArrayList<>(); + counts.forEachSingle((packed, length, count) -> { + if (count >= minCount) { + candidates.add(new Candidate(packed, length, gainOf(count, length))); + } + }); + counts.forEachPair((first, second, count) -> { + if (count < minCount) { + return; + } + int firstLength = counts.lengthOf(first); + int combinedLength = Math.min(firstLength + counts.lengthOf(second), MAX_SYMBOL_LENGTH); + long combined = concatenate(first, firstLength, second); + long packed = combined & lengthMask(combinedLength); + candidates.add(new Candidate(packed, combinedLength, gainOf(count, combinedLength))); + }); + + return selectTop(candidates, finalGeneration); + } + + /// Ranks candidates by gain-descending (length-descending on ties) and keeps the top + /// [#MAX_SYMBOLS], applying the final cost prune on the last generation. Uses a bounded min-heap + /// so a large candidate set is not fully sorted — only the surviving 255 are. + private static List selectTop(List candidates, boolean finalGeneration) { + // A bounded min-heap of size MAX_SYMBOLS: the weakest survivor sits at the head, so a new + // candidate either loses to it (discarded) or evicts it. This is top-K in O(n log K) rather + // than an O(n log n) full sort of every candidate. + // Min-heap on ranking gain: the weakest survivor sits at the head. + PriorityQueue heap = new PriorityQueue<>(TrainingGeneration::compareGain); + for (Candidate candidate : candidates) { + if (heap.size() < MAX_SYMBOLS) { + heap.add(candidate); + } else if (compareGain(candidate, heap.peek()) > 0) { + heap.poll(); + heap.add(candidate); + } + } + + List survivors = new ArrayList<>(heap); + // Restore gain-descending order (the heap only guarantees the head is weakest). + survivors.sort((a, b) -> compareGain(b, a)); + + List result = new ArrayList<>(survivors.size()); + for (Candidate candidate : survivors) { + if (result.size() >= MAX_SYMBOLS) { + break; + } + // Refinement 4 — final cost-based prune. Only on the last generation, a symbol earns its + // code slot only if its real (un-boosted) gain exceeds length+1: it must save at least + // that many bytes over escaping every occurrence, otherwise the slot is better left for a + // candidate that does. Provisional earlier generations skip this so their symbols stay + // available to be re-evaluated next pass. + if (finalGeneration && realGain(candidate) <= candidate.length() + 1L) { + continue; + } + result.add(new Symbol(candidate.packed(), candidate.length())); + } + return result; + } + + /// Gain used for ranking: `count * length`, with the single-byte boost applied to length-1 + /// candidates (refinement 3). + private static long gainOf(long count, int length) { + long gain = count * length; + return length == 1 ? gain * SINGLE_BYTE_GAIN_BOOST : gain; + } + + /// Real (un-boosted) gain of a candidate, dividing the length-1 boost back out — used only by + /// the final cost prune, which must judge a symbol on its true byte savings, not its boosted + /// ranking gain. + private static long realGain(Candidate candidate) { + return candidate.length() == 1 + ? candidate.gain() / SINGLE_BYTE_GAIN_BOOST + : candidate.gain(); + } + + private static int compareGain(Candidate a, Candidate b) { + int byGain = Long.compare(a.gain(), b.gain()); + if (byGain != 0) { + return byGain; + } + return Integer.compare(a.length(), b.length()); + } + + /// Concatenates two packed symbols LSB-first: the first symbol's bytes stay in the low bytes and + /// the second symbol's bytes are shifted up by the first symbol's byte length, so the combined + /// word reads the first symbol's bytes then the second's. A shift of 64 bits (first already 8 + /// bytes long) drops the second symbol entirely, which is correct — the result is then truncated + /// to 8 bytes anyway. + private static long concatenate(long first, int firstLength, long second) { + if (firstLength >= MAX_SYMBOL_LENGTH) { + return first; + } + return first | (second << (firstLength * 8)); + } + + private static long lengthMask(int length) { + return length >= MAX_SYMBOL_LENGTH ? ~0L : (1L << (length * 8)) - 1; + } + + /// A ranked candidate symbol: packed bytes, length, and ranking gain (boosted for length 1). + private record Candidate(long packed, int length, long gain) { + } + + /// Tallies of single-match and adjacent-pair counts during compress-count. Single matches are + /// keyed by packed symbol (its length is stored alongside); pairs are keyed by the ordered + /// `(firstPacked, secondPacked)` pair. + private static final class Counts { + + private final Map singles = new HashMap<>(); // packed -> {count, length} + private final Map pairs = new HashMap<>(); + + void bumpSingle(long packed, int length) { + singles.merge(packed, new long[]{1, length}, (existing, added) -> { + existing[0]++; + return existing; + }); + } + + void bumpPair(long first, long second) { + pairs.merge(new Pair(first, second), 1L, Long::sum); + } + + int lengthOf(long packed) { + return (int) singles.get(packed)[1]; + } + + void forEachSingle(SingleConsumer consumer) { + for (Map.Entry entry : singles.entrySet()) { + long[] value = entry.getValue(); + consumer.accept(entry.getKey(), (int) value[1], value[0]); + } + } + + void forEachPair(PairConsumer consumer) { + for (Map.Entry entry : pairs.entrySet()) { + Pair pair = entry.getKey(); + consumer.accept(pair.first(), pair.second(), entry.getValue()); + } + } + } + + private record Pair(long first, long second) { + } + + @FunctionalInterface + private interface SingleConsumer { + void accept(long packed, int length, long count); + } + + @FunctionalInterface + private interface PairConsumer { + void accept(long first, long second, long count); + } +} diff --git a/fsst/src/test/java/io/github/dfa1/vortex/fsst/CompressorBuilderTest.java b/fsst/src/test/java/io/github/dfa1/vortex/fsst/CompressorBuilderTest.java new file mode 100644 index 00000000..85f7b7fc --- /dev/null +++ b/fsst/src/test/java/io/github/dfa1/vortex/fsst/CompressorBuilderTest.java @@ -0,0 +1,269 @@ +package io.github.dfa1.vortex.fsst; + +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.Random; + +import static org.assertj.core.api.Assertions.assertThat; + +class CompressorBuilderTest { + + @Test + void train_emptyInput_producesEmptyTable() { + // Given — no rows at all. + CompressorBuilder sut = new CompressorBuilder(); + + // When + Compressor result = sut.train(new byte[0][]); + + // Then + assertThat(result.symbolCount()).isZero(); + } + + @Test + void train_allEmptyRows_producesEmptyTable() { + // Given — rows exist but every one is empty, so there is nothing to learn. + CompressorBuilder sut = new CompressorBuilder(); + + // When + Compressor result = sut.train(new byte[][]{new byte[0], new byte[0]}); + + // Then + assertThat(result.symbolCount()).isZero(); + } + + @Test + void train_repetitiveText_learnsMultiByteSymbols() { + // Given — realistic log-like lines with heavily repeated substrings ("ERROR", timestamps, + // "host=example.com"), which FSST should collapse into multi-byte symbols. + byte[][] rows = logLikeRows(2_000); + CompressorBuilder sut = new CompressorBuilder(); + + // When + Compressor result = sut.train(rows); + + // Then — at least one learned symbol is longer than a bigram, i.e. training pairs adjacent + // matches into longer symbols rather than stopping at length 2. + assertThat(hasSymbolLongerThan(result, 2)).isTrue(); + } + + @Test + void train_repetitiveText_compressesAndRoundTrips() { + // Given + byte[][] rows = logLikeRows(2_000); + Compressor sut = new CompressorBuilder().train(rows); + + // When — compress then decompress row 0 through the trained table. + byte[] input = rows[0]; + byte[] compressed = new byte[input.length * 2]; + long compressedLength = sut.compress(input, 0, input.length, compressed, 0); + byte[] decompressed = new byte[input.length]; + long decompressedLength = + sut.toDecompressor().decompress(compressed, 0, (int) compressedLength, decompressed, 0); + + // Then — round-trips exactly and the compressed form is smaller than the raw input. + assertThat(decompressedLength).isEqualTo(input.length); + assertThat(decompressed).isEqualTo(input); + assertThat(compressedLength).isLessThan(input.length); + } + + @Nested + class EndOfInput { + + @Test + void compress_symbolWithTrailingZeros_doesNotMatchPastEndOfInput() { + // Given — a table with one length-8 symbol "AB\0\0\0\0\0\0". Compressing the 3-byte + // input "AB\0" must NOT match that symbol: loadWord zero-pads the 5 bytes past the + // input, so the branch-free matcher's masked compare would spuriously satisfy the full + // 8-byte compare. Trusting it emits one code that decodes to 8 bytes, silently + // appending 5 garbage bytes past the true length — a real corruption on any NUL-tailed + // or binary row, which this encoder also accepts (DType.Binary). The caller must reject + // any match reaching past the input end and escape instead. + long packed = ('A' & 0xFFL) | (('B' & 0xFFL) << 8); // "AB" then six zero bytes. + Compressor sut = Compressor.of(List.of(new Symbol(packed, 8))); + byte[] input = {'A', 'B', 0}; + + // When — compress then decompress the shorter input. + byte[] compressed = new byte[input.length * 2]; + long compressedLength = sut.compress(input, 0, input.length, compressed, 0); + byte[] decompressed = new byte[input.length]; + long decompressedLength = sut.toDecompressor() + .decompress(compressed, 0, (int) compressedLength, decompressed, 0); + + // Then — the round-trip preserves the exact original length and bytes. + assertThat(decompressedLength).isEqualTo(input.length); + assertThat(decompressed).isEqualTo(input); + } + } + + @Nested + class Determinism { + + @Test + void train_sameInputAndSeed_producesByteIdenticalTable() { + // Given — the same corpus trained twice with the same explicit seed. + byte[][] rows = logLikeRows(2_000); + + // When + Compressor first = new CompressorBuilder().seed(1234L).train(rows); + Compressor second = new CompressorBuilder().seed(1234L).train(rows); + + // Then — identical symbol count and identical (packed, length) for every code. + assertThat(second.symbolCount()).isEqualTo(first.symbolCount()); + for (int code = 0; code < first.symbolCount(); code++) { + assertThat(second.packedSymbol(code)).isEqualTo(first.packedSymbol(code)); + assertThat(second.symbolLength(code)).isEqualTo(first.symbolLength(code)); + } + } + + @Test + void train_defaultSeed_isAlsoDeterministic() { + // Given — no explicit seed: the default seed must be a fixed constant, not wall-clock. + byte[][] rows = logLikeRows(1_000); + + // When + Compressor first = new CompressorBuilder().train(rows); + Compressor second = new CompressorBuilder().train(rows); + + // Then + assertThat(second.symbolCount()).isEqualTo(first.symbolCount()); + for (int code = 0; code < first.symbolCount(); code++) { + assertThat(second.packedSymbol(code)).isEqualTo(first.packedSymbol(code)); + } + } + } + + @Nested + class WireOrder { + + @Test + void codesSortedByLength_multiByteAscendingThenSingleByteLast() { + // Given — a trained table with a spread of symbol lengths. + Compressor sut = new CompressorBuilder().train(logLikeRows(2_000)); + + // When + int[] result = sut.codesSortedByLength(); + + // Then — multi-byte symbols come first in non-decreasing length order, then every + // length-1 symbol, mirroring the vortex.fsst wire contract this permutation feeds. + assertThat(result).hasSize(sut.symbolCount()); + int previousMultiByteLength = 0; + boolean seenSingleByte = false; + for (int code : result) { + int length = sut.symbolLength(code); + if (length == 1) { + seenSingleByte = true; + } else { + assertThat(seenSingleByte) + .as("no multi-byte symbol may follow a length-1 symbol") + .isFalse(); + assertThat(length) + .as("multi-byte symbols are in non-decreasing length order") + .isGreaterThanOrEqualTo(previousMultiByteLength); + previousMultiByteLength = length; + } + } + } + } + + @Nested + class MinCountPruning { + + @Test + void run_nonFinalGeneration_prunesCandidateBelowFloor() { + // Given — "ab" repeated a thousand times (very frequent) with a single "cd" appended + // (count 1). A non-final generation at full sample fraction (numerator 128) uses a + // min-count floor of 5, so the rare "cd" pair — and its bytes 'c'/'d' — must be pruned + // before ranking, unlike a naive keep-everything scheme. + StringBuilder text = new StringBuilder(); + for (int i = 0; i < 1_000; i++) { + text.append("ab"); + } + text.append("cd"); + Sample sample = Sample.draw(new byte[][]{utf8(text.toString())}, 1L); + Compressor emptyTable = Compressor.of(List.of()); + + // When — one non-final generation over the whole sample. + List result = + TrainingGeneration.run(emptyTable, sample, sample.chunkCount(), 128, false); + + // Then — the rare bytes never make it into the table. + assertThat(containsByte(result, (byte) 'c')).isFalse(); + assertThat(containsByte(result, (byte) 'd')).isFalse(); + // Sanity: the frequent 'a'/'b' content is retained. + assertThat(containsByte(result, (byte) 'a')).isTrue(); + } + } + + @Nested + class SingleByteGainBoost { + + @Test + void train_frequentSingleByte_survivesAmidManyMultiByteCompetitors() { + // Given — a frequent single byte 'x' separates many distinct 3-byte groups. Under plain + // count*length gain a single byte is easily out-ranked by longer candidates; the 8x + // single-byte boost is exactly what keeps a frequent byte in the table, so it is not + // crowded out and its occurrences do not degrade into 2-byte escapes. + String[] groups = {"abc", "def", "ghi", "jkl", "mno", "pqr", "stu", "vwx"}; + Random random = new Random(3); + StringBuilder text = new StringBuilder(); + for (int i = 0; i < 4_000; i++) { + text.append(groups[random.nextInt(groups.length)]).append('x'); + } + Compressor sut = new CompressorBuilder().train(new byte[][]{utf8(text.toString())}); + + // When + boolean result = hasSingleByteSymbol(sut, (byte) 'x'); + + // Then + assertThat(result).isTrue(); + } + } + + private static byte[][] logLikeRows(int count) { + // Real repeated substrings: a shared timestamp/level/host prefix plus a varying tail, the + // kind of column FSST is meant to compress well. + byte[][] rows = new byte[count][]; + String prefix = "2026-07-20T12:34:56Z ERROR request failed host=example.com id="; + for (int i = 0; i < count; i++) { + rows[i] = utf8(prefix + i); + } + return rows; + } + + private static boolean hasSymbolLongerThan(Compressor compressor, int length) { + for (int code = 0; code < compressor.symbolCount(); code++) { + if (compressor.symbolLength(code) > length) { + return true; + } + } + return false; + } + + private static boolean hasSingleByteSymbol(Compressor compressor, byte value) { + for (int code = 0; code < compressor.symbolCount(); code++) { + if (compressor.symbolLength(code) == 1 && (byte) compressor.packedSymbol(code) == value) { + return true; + } + } + return false; + } + + private static boolean containsByte(List symbols, byte value) { + for (Symbol symbol : symbols) { + for (int k = 0; k < symbol.length(); k++) { + if (symbol.byteAt(k) == value) { + return true; + } + } + } + return false; + } + + private static byte[] utf8(String s) { + return s.getBytes(StandardCharsets.UTF_8); + } +} diff --git a/fsst/src/test/java/io/github/dfa1/vortex/fsst/SampleTest.java b/fsst/src/test/java/io/github/dfa1/vortex/fsst/SampleTest.java new file mode 100644 index 00000000..d613c60f --- /dev/null +++ b/fsst/src/test/java/io/github/dfa1/vortex/fsst/SampleTest.java @@ -0,0 +1,97 @@ +package io.github.dfa1.vortex.fsst; + +import org.junit.jupiter.api.Test; + +import java.nio.charset.StandardCharsets; + +import static org.assertj.core.api.Assertions.assertThat; + +class SampleTest { + + @Test + void draw_largeInput_totalSampledBytesNearTarget() { + // Given — a corpus far larger than the 16KB target, so sampling must cap the total. + byte[][] rows = repeatedRows("the quick brown fox jumps over the lazy dog ", 5_000); + + // When + Sample result = Sample.draw(rows, 1L); + + // Then — the total lands at the target exactly here, since every draw chunk is full-length; + // in general it is bounded by TARGET_SAMPLE_BYTES and never overshoots (draws stop once the + // target is reached). Allow a small tolerance for the last, possibly-truncated chunk. + assertThat(totalBytes(result)) + .isLessThanOrEqualTo(Sample.TARGET_SAMPLE_BYTES) + .isGreaterThan(Sample.TARGET_SAMPLE_BYTES - Sample.MAX_CHUNK_BYTES); + } + + @Test + void draw_smallInput_usesFullCorpus() { + // Given — a corpus well under the target: no under-sampling should occur. + byte[][] rows = repeatedRows("abc", 4); // 12 bytes total, << 16KB. + + // When + Sample result = Sample.draw(rows, 1L); + + // Then — every input byte is represented. + assertThat(totalBytes(result)).isEqualTo(12); + } + + @Test + void draw_allEmptyRows_producesEmptySample() { + // Given — every row empty, so there is nothing to sample. + byte[][] rows = {new byte[0], new byte[0], new byte[0]}; + + // When + Sample result = Sample.draw(rows, 1L); + + // Then + assertThat(result.chunkCount()).isZero(); + assertThat(result.bytes()).isEmpty(); + } + + @Test + void draw_sameSeed_producesByteIdenticalSample() { + // Given + byte[][] rows = repeatedRows("deterministic sampling must reproduce ", 2_000); + + // When + Sample first = Sample.draw(rows, 99L); + Sample second = Sample.draw(rows, 99L); + + // Then + assertThat(second.bytes()).isEqualTo(first.bytes()); + assertThat(second.chunkCount()).isEqualTo(first.chunkCount()); + } + + @Test + void chunkCountForGeneration_growsMonotonicallyToFullSample() { + // Given + byte[][] rows = repeatedRows("grow the fraction over generations ", 2_000); + Sample sut = Sample.draw(rows, 7L); + + // When + int gen0 = sut.chunkCountForGeneration(0); + int gen4 = sut.chunkCountForGeneration(4); + + // Then — the first generation sees only a fraction; the last sees every chunk. + assertThat(gen0).isGreaterThanOrEqualTo(1).isLessThan(sut.chunkCount()); + assertThat(gen4).isEqualTo(sut.chunkCount()); + } + + private static byte[][] repeatedRows(String text, int count) { + byte[] row = text.getBytes(StandardCharsets.UTF_8); + byte[][] rows = new byte[count][]; + for (int i = 0; i < count; i++) { + rows[i] = row; + } + return rows; + } + + private static int totalBytes(Sample sample) { + int total = 0; + for (int i = 0; i < sample.chunkCount(); i++) { + total += sample.chunkEnd(i) - sample.chunkStart(i); + } + return total; + } +}