diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/RleEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/RleEncodingEncoder.java index 86d85b49..b6850a33 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/RleEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/RleEncodingEncoder.java @@ -12,7 +12,9 @@ import java.lang.foreign.MemorySegment; import java.lang.foreign.SegmentAllocator; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; +import java.util.Set; /// Write-only encoder for `fastlanes.rle`. public final class RleEncodingEncoder implements EncodingEncoder { @@ -44,58 +46,18 @@ static EncodeResult encodeBool(boolean[] validity, EncodeContext ctx) { for (int i = 0; i < n; i++) { longs[i] = validity[i] ? 1L : 0L; } - - int numChunks = (n + FL_CHUNK_SIZE - 1) / FL_CHUNK_SIZE; - int paddedLen = numChunks * FL_CHUNK_SIZE; - - long[] globalValues = new long[paddedLen]; - short[] globalIndices = new short[paddedLen]; - long[] valuesIdxOffsets = new long[numChunks]; - - long[] chunkInput = new long[FL_CHUNK_SIZE]; - long[] chunkValues = new long[FL_CHUNK_SIZE]; - short[] chunkIndices = new short[FL_CHUNK_SIZE]; - - int globalValuesCount = 0; - for (int chunk = 0; chunk < numChunks; chunk++) { - int chunkStart = chunk * FL_CHUNK_SIZE; - int chunkEnd = Math.min(chunkStart + FL_CHUNK_SIZE, n); - int chunkLen = chunkEnd - chunkStart; - - System.arraycopy(longs, chunkStart, chunkInput, 0, chunkLen); - long lastVal = longs[chunkEnd - 1]; - for (int i = chunkLen; i < FL_CHUNK_SIZE; i++) { - chunkInput[i] = lastVal; - } - - int numChunkValues = rleEncode(chunkInput, chunkValues, chunkIndices); - - valuesIdxOffsets[chunk] = globalValuesCount; - System.arraycopy(chunkValues, 0, globalValues, globalValuesCount, numChunkValues); - globalValuesCount += numChunkValues; - - System.arraycopy(chunkIndices, 0, globalIndices, chunkStart, FL_CHUNK_SIZE); - } - - boolean[] valuesArr = new boolean[globalValuesCount]; - for (int i = 0; i < globalValuesCount; i++) { - valuesArr[i] = globalValues[i] != 0L; + Runs runs = runs(longs); + boolean[] valuesArr = new boolean[runs.valuesCount()]; + for (int i = 0; i < valuesArr.length; i++) { + valuesArr[i] = runs.values()[i] != 0L; } EncodeResult valuesResult = new BoolEncodingEncoder().encode(DType.BOOL, valuesArr, ctx); - MemorySegment indicesSeg = toIndicesSeg(globalIndices, paddedLen, ctx.arena()); - MemorySegment offsetsSeg = fromLongsU64(valuesIdxOffsets, numChunks, ctx.arena()); + MemorySegment indicesSeg = toIndicesSeg(runs.indices(), runs.paddedLen(), ctx.arena()); + MemorySegment offsetsSeg = fromLongsU64(runs.offsets(), runs.numChunks(), ctx.arena()); PType indicesPtype = PType.U16; PType offsetsPtype = PType.U64; - - byte[] metaBytes = new ProtoRLEMetadata( - globalValuesCount, - paddedLen, - io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(indicesPtype.ordinal()), - numChunks, - io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(offsetsPtype.ordinal()), - 0L - ).encode(); + byte[] metaBytes = runs.metadata(); int indicesBufIdx = valuesResult.buffers().size(); EncodeNode indicesNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, indicesBufIdx); @@ -129,6 +91,86 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { // raw values do. byte[][] stats = ZoneMapStats.of(dtype, data); + Runs runs = runs(longs); + MemorySegment valuesSeg = fromLongs(runs.values(), runs.valuesCount(), ptype, ctx.arena()); + MemorySegment indicesSeg = toIndicesSeg(runs.indices(), runs.paddedLen(), ctx.arena()); + MemorySegment offsetsSeg = fromLongsU64(runs.offsets(), runs.numChunks(), ctx.arena()); + + PType indicesPtype = PType.U16; + PType offsetsPtype = PType.U64; + byte[] metaBytes = runs.metadata(); + + EncodeNode valuesNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 0); + EncodeNode indicesNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 1); + EncodeNode offsetsNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 2); + EncodeNode root = new EncodeNode( + EncodingId.FASTLANES_RLE, + MemorySegment.ofArray(metaBytes), + new EncodeNode[]{valuesNode, indicesNode, offsetsNode}, + new int[0]); + return new EncodeResult(root, List.of(EncodedBuffer.of(valuesSeg, ptype), EncodedBuffer.of(indicesSeg, indicesPtype), + EncodedBuffer.of(offsetsSeg, offsetsPtype)), null, null).withStats(stats); + } + + /// Barred from the indices and offsets children (issue #410, Rust's `rle_descendant_exclusions`): + /// Dict and Sparse cannot pay off on per-chunk run positions or monotone offsets. Rust keeps a + /// RunEnd rule there commented out as unsound, so RunEnd is not barred. RLE names itself, as + /// Rust's compressor never repeats a scheme within one chain. + private static final Set POSITIONS_EXCLUDED = + Set.of(EncodingId.FASTLANES_RLE, EncodingId.VORTEX_DICT, EncodingId.VORTEX_SPARSE); + + /// Barred from the values child: only RLE itself (no rule upstream). + private static final Set VALUES_EXCLUDED = Set.of(EncodingId.FASTLANES_RLE); + + /// Cascading RLE, mirroring Rust's `IntRLEScheme`: the values, indices and offsets become open + /// children (Rust ids values=0, indices=1, offsets=2, also our wire order) for the compressor + /// to bit-pack / FoR, instead of raw buffers. Same layout and metadata as [#encode]. + @Override + public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { + PType ptype = ((DType.Primitive) dtype).ptype(); + long[] longs = toLongs(data, ptype); + if (longs.length == 0) { + return CascadeStep.notApplicable(); + } + Runs runs = runs(longs); + EncodeNode partialRoot = new EncodeNode(EncodingId.FASTLANES_RLE, MemorySegment.ofArray(runs.metadata()), + new EncodeNode[]{null, null, null}, new int[0]); + long[] values = Arrays.copyOf(runs.values(), runs.valuesCount()); + List slots = List.of( + new ChildSlot(dtype, PrimitiveArrays.fromLongsArray(values, ptype, EncodingId.FASTLANES_RLE), 0, + VALUES_EXCLUDED), + new ChildSlot(new DType.Primitive(PType.U16, false), runs.indices(), 1, POSITIONS_EXCLUDED), + new ChildSlot(new DType.Primitive(PType.U64, false), runs.offsets(), 2, POSITIONS_EXCLUDED)); + byte[][] stats = ZoneMapStats.of(dtype, data); + return new CascadeStep(partialRoot, List.of(), slots, ZoneMapStats.minOf(stats), ZoneMapStats.maxOf(stats), true); + } + + /// One column run-length encoded in FastLanes 1024-row chunks. + /// + /// @param values the run values, chunk after chunk (first `valuesCount` are used) + /// @param valuesCount number of run values + /// @param indices per padded row, the index of its run within its chunk + /// @param paddedLen row count rounded up to a whole chunk + /// @param offsets per chunk, the index of its first run value + /// @param numChunks chunk count + private record Runs(long[] values, int valuesCount, short[] indices, int paddedLen, long[] offsets, + int numChunks) { + + byte[] metadata() { + return new ProtoRLEMetadata( + valuesCount, + paddedLen, + io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(PType.U16.ordinal()), + numChunks, + io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(PType.U64.ordinal()), + 0L + ).encode(); + } + } + + /// Run-length encodes `longs` chunk by chunk; the last chunk is padded with its final value. + private static Runs runs(long[] longs) { + int n = longs.length; int numChunks = (n + FL_CHUNK_SIZE - 1) / FL_CHUNK_SIZE; int paddedLen = numChunks * FL_CHUNK_SIZE; @@ -161,33 +203,7 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { System.arraycopy(chunkIndices, 0, globalIndices, chunkStart, FL_CHUNK_SIZE); } - - MemorySegment valuesSeg = fromLongs(globalValues, globalValuesCount, ptype, ctx.arena()); - MemorySegment indicesSeg = toIndicesSeg(globalIndices, paddedLen, ctx.arena()); - MemorySegment offsetsSeg = fromLongsU64(valuesIdxOffsets, numChunks, ctx.arena()); - - PType indicesPtype = PType.U16; - PType offsetsPtype = PType.U64; - - byte[] metaBytes = new ProtoRLEMetadata( - globalValuesCount, - paddedLen, - io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(indicesPtype.ordinal()), - numChunks, - io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(offsetsPtype.ordinal()), - 0L - ).encode(); - - EncodeNode valuesNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 0); - EncodeNode indicesNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 1); - EncodeNode offsetsNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 2); - EncodeNode root = new EncodeNode( - EncodingId.FASTLANES_RLE, - MemorySegment.ofArray(metaBytes), - new EncodeNode[]{valuesNode, indicesNode, offsetsNode}, - new int[0]); - return new EncodeResult(root, List.of(EncodedBuffer.of(valuesSeg, ptype), EncodedBuffer.of(indicesSeg, indicesPtype), - EncodedBuffer.of(offsetsSeg, offsetsPtype)), null, null).withStats(stats); + return new Runs(globalValues, globalValuesCount, globalIndices, paddedLen, valuesIdxOffsets, numChunks); } private static int rleEncode(long[] input, long[] chunkValues, short[] chunkIndices) { diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/RleEncodingEncoderTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/RleEncodingEncoderTest.java index 7b3fe37a..4810e211 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/RleEncodingEncoderTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/RleEncodingEncoderTest.java @@ -392,4 +392,49 @@ void encode_i32_metadata_valuesLen_matchesRunCount() throws Exception { assertThat(meta.indices_ptype().value()).isEqualTo(1); } } + + /// Cascading RLE (issue #410) hands values, indices and offsets to the compressor as open + /// children instead of raw buffers, with Rust's `rle_descendant_exclusions` on each. + @Nested + class Cascade { + + @Test + void encodeCascade_opensValuesIndicesAndOffsetsSlots() { + // Given — 1500 rows: two FastLanes chunks; chunk 0 has runs 5,9 and chunk 1 a single 9 + int[] data = new int[1_500]; + java.util.Arrays.fill(data, 0, 600, 5); + java.util.Arrays.fill(data, 600, 1_500, 9); + var sut = new RleEncodingEncoder(); + + // When + CascadeStep result = sut.encodeCascade(DTypes.I32, data, EncodeTestHelper.testCtx()); + + // Then + assertThat(result.applicable()).isTrue(); + assertThat(result.ownedBuffers()).isEmpty(); + ChildSlot values = result.openChildren().get(0); + ChildSlot indices = result.openChildren().get(1); + ChildSlot offsets = result.openChildren().get(2); + assertThat((int[]) values.childData()).containsExactly(5, 9, 9); + assertThat(values.excluded()).containsExactly(EncodingId.FASTLANES_RLE); + assertThat(indices.childDtype()).isEqualTo(new DType.Primitive(io.github.dfa1.vortex.core.model.PType.U16, false)); + assertThat((short[]) indices.childData()).hasSize(2_048); + assertThat(((short[]) indices.childData())[599]).isZero(); + assertThat(((short[]) indices.childData())[600]).isEqualTo((short) 1); + assertThat((long[]) offsets.childData()).containsExactly(0L, 2L); + for (ChildSlot positions : List.of(indices, offsets)) { + assertThat(positions.excluded()).containsExactlyInAnyOrder(EncodingId.FASTLANES_RLE, + EncodingId.VORTEX_DICT, EncodingId.VORTEX_SPARSE); + } + } + + @Test + void encodeCascade_empty_notApplicable() { + // When + CascadeStep result = new RleEncodingEncoder().encodeCascade(DTypes.I32, new int[0], EncodeTestHelper.testCtx()); + + // Then + assertThat(result.applicable()).isFalse(); + } + } }