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 50314526..0bdf5f42 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 @@ -54,7 +54,78 @@ public boolean accepts(DType dtype) { @Override public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { - String[] strings = (String[]) data; + Arena arena = ctx.arena(); + Fsst c = compress((String[]) data, arena); + + // Terminal layout: the per-row length and cumulative-offset children are raw primitive + // segments (buffers 3 and 4). The cascading path (encodeCascade) instead exposes them as + // open child slots so they can be bitpacked/constant-folded. + long uncompLenBytes = c.uncompLenPType().byteSize(); + MemorySegment uncompLenBuf = arena.allocate(Math.max((long) c.n() * uncompLenBytes, 1)); + for (int i = 0; i < c.n(); i++) { + PTypeIO.set(uncompLenBuf, i * uncompLenBytes, c.uncompLenPType(), c.uncompLens()[i]); + } + long codesOffBytes = c.codesOffPType().byteSize(); + MemorySegment codesOffBuf = arena.allocate((long) (c.n() + 1) * codesOffBytes); + for (int i = 0; i <= c.n(); i++) { + PTypeIO.set(codesOffBuf, i * codesOffBytes, c.codesOffPType(), c.codesOffsets()[i]); + } + + EncodeNode uncompLensNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 3); + EncodeNode codesOffNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 4); + EncodeNode root = new EncodeNode( + EncodingId.VORTEX_FSST, + MemorySegment.ofArray(c.metaBytes()), + new EncodeNode[]{uncompLensNode, codesOffNode}, + new int[]{0, 1, 2}); + + return new EncodeResult(root, + List.of(c.symBuf(), c.symLenBuf(), c.compBuf(), uncompLenBuf, codesOffBuf), + null, null); + } + + /// Cascading FSST: expose the per-row uncompressed-length and code-offset children as open + /// primitive slots so the cascade can compress them — bitpacking the monotonic offsets and + /// constant-folding the lengths of fixed-width columns (dates, coordinates). Matches the Rust + /// reference, which stores these children as `fastlanes.bitpacked` / `vortex.constant` rather + /// than raw primitives; leaving them raw cost ~2-4 bytes/row and was the largest remaining gap + /// on the nyc-311 date columns (#299). At cascade depth 0 there is no competition, so fall back + /// to the terminal raw-primitive layout. + /// + /// @param dtype the Utf8/Binary type being encoded + /// @param data the string values + /// @param ctx encoding context supplying the arena and cascade depth + /// @return a cascade step with the two offset children left open, or a terminal step at depth 0 + @Override + public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { + if (ctx.allowedCascading() == 0) { + return CascadeStep.terminal(encode(dtype, data, ctx)); + } + Fsst c = compress((String[]) data, ctx.arena()); + Object uncompLens = typedUnsigned(c.uncompLenPType(), c.uncompLens()); + Object codesOffsets = typedUnsigned(c.codesOffPType(), c.codesOffsets()); + EncodeNode partialRoot = new EncodeNode( + EncodingId.VORTEX_FSST, + MemorySegment.ofArray(c.metaBytes()), + new EncodeNode[]{null, null}, + new int[]{0, 1, 2}); + return new CascadeStep(partialRoot, + List.of(c.symBuf(), c.symLenBuf(), c.compBuf()), + List.of(new ChildSlot(new DType.Primitive(c.uncompLenPType(), false), uncompLens, 0), + new ChildSlot(new DType.Primitive(c.codesOffPType(), false), codesOffsets, 1)), + null, null, true); + } + + /// The FSST-specific product of compression: the symbol-table buffers, the wire code stream, the + /// [ProtoFSSTMetadata] bytes, and the per-row uncompressed lengths / cumulative code offsets as + /// plain `int[]` (the two offset children, in narrowest-unsigned ptypes). + private record Fsst( + MemorySegment symBuf, MemorySegment symLenBuf, MemorySegment compBuf, + byte[] metaBytes, int[] uncompLens, int[] codesOffsets, + PType uncompLenPType, PType codesOffPType, int n) { + } + + private static Fsst compress(String[] strings, Arena arena) { int n = strings.length; byte[][] byteArrays = new byte[n][]; @@ -81,8 +152,6 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { internalToWire[wireOrder[i]] = i; } - Arena arena = ctx.arena(); - MemorySegment symBuf = arena.allocate(Math.max(numSymbols * 8L, 1), 8); MemorySegment symLenBuf = arena.allocate(Math.max(numSymbols, 1)); for (int i = 0; i < numSymbols; i++) { @@ -116,17 +185,13 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { PType uncompLenPType = PType.narrowestUnsigned(maxUncompLen); PType codesOffPType = PType.narrowestUnsigned(totalCompressed); - long uncompLenBytes = uncompLenPType.byteSize(); - MemorySegment uncompLenBuf = arena.allocate(Math.max((long) n * uncompLenBytes, 1)); + int[] uncompLens = new int[n]; for (int i = 0; i < n; i++) { - PTypeIO.set(uncompLenBuf, i * uncompLenBytes, uncompLenPType, byteArrays[i].length); + uncompLens[i] = byteArrays[i].length; } - - long codesOffBytes = codesOffPType.byteSize(); - MemorySegment codesOffBuf = arena.allocate((long) (n + 1) * codesOffBytes); - PTypeIO.set(codesOffBuf, 0, codesOffPType, 0); + int[] codesOffsets = new int[n + 1]; for (int i = 0; i < n; i++) { - PTypeIO.set(codesOffBuf, (i + 1) * codesOffBytes, codesOffPType, rowEnds[i]); + codesOffsets[i + 1] = rowEnds[i]; } byte[] metaBytes = new ProtoFSSTMetadata( @@ -134,17 +199,34 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(codesOffPType.ordinal()) ).encode(); - EncodeNode uncompLensNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 3); - EncodeNode codesOffNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 4); - EncodeNode root = new EncodeNode( - EncodingId.VORTEX_FSST, - MemorySegment.ofArray(metaBytes), - new EncodeNode[]{uncompLensNode, codesOffNode}, - new int[]{0, 1, 2}); + return new Fsst(symBuf, symLenBuf, compBuf, metaBytes, uncompLens, codesOffsets, + uncompLenPType, codesOffPType, n); + } - return new EncodeResult(root, - List.of(symBuf, symLenBuf, compBuf, uncompLenBuf, codesOffBuf), - null, null); + /// Copies unsigned values into the narrowest Java array matching `ptype` (U8→`byte[]`, + /// U16→`short[]`, else `int[]`) so a [ChildSlot] can hand them to the cascade's primitive codecs. + /// + /// @param ptype the child's primitive type + /// @param vals the values (already within `ptype`'s unsigned range) + /// @return a `byte[]`, `short[]`, or `int[]` holding `vals` + private static Object typedUnsigned(PType ptype, int[] vals) { + return switch (ptype) { + case U8, I8 -> { + byte[] a = new byte[vals.length]; + for (int i = 0; i < vals.length; i++) { + a[i] = (byte) vals[i]; + } + yield a; + } + case U16, I16 -> { + short[] a = new short[vals.length]; + for (int i = 0; i < vals.length; i++) { + a[i] = (short) vals[i]; + } + yield a; + } + default -> vals; + }; } /// Remaps every code byte in `stream[0..length)` from the compressor's internal (gain-descending) diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoderTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoderTest.java index c57b0785..636fcba4 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoderTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/FsstEncodingEncoderTest.java @@ -1,7 +1,9 @@ package io.github.dfa1.vortex.writer.encode; +import io.github.dfa1.vortex.core.model.DType; import io.github.dfa1.vortex.core.model.PType; import io.github.dfa1.vortex.core.error.VortexException; +import io.github.dfa1.vortex.writer.WriteRegistry; import io.github.dfa1.vortex.reader.array.VarBinArray; import io.github.dfa1.vortex.reader.decode.ArrayNode; import io.github.dfa1.vortex.core.testing.DTypes; @@ -428,6 +430,34 @@ void encode_metadata_bothPTypes_escalate_forU16RoundtripCase() throws Exception assertThat(meta.codes_offsets_ptype().value()).isEqualTo(PType.U16.ordinal()); } + @Test + void encodeCascade_exposesOffsetChildrenAsOpenSlots() { + // Given — a cascading context (depth 3). + var ctx = EncodeContext.ofDepth(3, Arena.ofAuto(), WriteRegistry.loadAll()); + + // When + CascadeStep result = ENCODER.encodeCascade( + DTypes.UTF8, java.util.Collections.nCopies(100, "2020").toArray(new String[0]), ctx); + + // Then — the per-row length (child 0) and code-offset (child 1) children are left open + // for the cascade to bitpack/constant-fold, not emitted as raw primitive leaves (#299). + assertThat(result.isTerminal()).isFalse(); + assertThat(result.openChildren()).extracting(ChildSlot::parentChildIdx).containsExactly(0, 1); + assertThat(result.openChildren()) + .allSatisfy(slot -> assertThat(slot.childDtype()).isInstanceOf(DType.Primitive.class)); + } + + @Test + void encodeCascade_atDepthZero_isTerminal() { + // Given / When — no cascade depth, so there is no competition to run. + CascadeStep result = ENCODER.encodeCascade( + DTypes.UTF8, java.util.Collections.nCopies(20, "ab").toArray(new String[0]), + EncodeTestHelper.testCtx()); + + // Then — falls back to the terminal raw-primitive layout. + assertThat(result.isTerminal()).isTrue(); + } + // Mirror of Encode.u16TierRows() — the nested classes cannot share a private static helper // without exposing it, and the intent (deterministic U16-on-both data) is small enough to // restate here next to the assertion that depends on it.