Skip to content

feat(writer): cascade RLE like Rust's IntRLEScheme (issue #410) #437

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dfa1 merged 1 commit into from
Oct 3, 2026
Merged

feat(writer): cascade RLE like Rust's IntRLEScheme (issue #410) #437

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
164 changes: 90 additions & 74 deletions writer/src/main/java/io/github/dfa1/vortex/writer/encode/RleEncodingEncoder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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<EncodingId> 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<EncodingId> 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<ChildSlot> 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;

Expand Down Expand Up @@ -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) {
Expand Down
45 changes: 45 additions & 0 deletions writer/src/test/java/io/github/dfa1/vortex/writer/encode/RleEncodingEncoderTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
}
Loading