From 4f91bc2cfa1a007d22380555b51800814cf45758 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sun, 25 Oct 2026 19:42:08 +0100 Subject: [PATCH 01/10] fix(writer): keep an encoding only when it is smaller than the raw array Rust's choose_and_compress (vortex-compressor compressor/cascade.rs) compresses the winning scheme on the full array and keeps the result only when its byte size is below the array's exact uncompressed size; otherwise the array stays canonical. Java ran the winner and kept whatever came out. The two disagree whenever the sample favors a scheme the full array does not. Bit-packing pads to 1024-value blocks, so a 1024-row sample of 1500 15-bit u16 values packs to 1920 of 2048 bytes and wins, while the full array pads to 2048 values: 3840 bytes over its 3000. Rust stores those as vortex.primitive. Found while comparing nullable columns against vortex-jni (#604), where it decides a zone table's four null counts (FoR wins their estimate, its bit-packed residuals pad past the 32 raw bytes) and a mostly-null column's sparse patch indices. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01SwhK4W5cgj7y7s4ZPCif5x --- CHANGELOG.md | 1 + .../writer/encode/CascadingCompressor.java | 35 ++++++++++++++-- .../encode/CascadingCompressorTest.java | 42 +++++++++++++++++++ 3 files changed, 74 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e5f44378f..0649588df 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **Breaking:** CSV import stores an empty cell in a number or boolean column as null instead of `0` or `false`, so inferred number and boolean columns are now nullable (read through `MaskedArray`), and an empty cell or a missing column fails the import where the schema declares the column non-nullable. ([#592](https://github.com/dfa1/vortex-java/issues/592)) ### Fixed +- The cascade keeps an encoding only when it comes out smaller than the raw array, as Rust does, so a short column no longer bit-packs into more bytes than it holds. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - An all-null chunk is written as one null `vortex.constant`, as Rust writes it, instead of a masked constant `0`. ([#594](https://github.com/dfa1/vortex-java/pull/594)) - An all-null chunk stored as a null constant (as vortex-jni writes it) reads as a masked array with every row null, like any nullable chunk, instead of a `NullArray` that the CLI, SQL and Parquet export failed on. ([#594](https://github.com/dfa1/vortex-java/pull/594)) - An all-null date, time, timestamp or uuid chunk is written as a bare null `vortex.constant`, as Rust writes it, without a `vortex.ext` around it. ([#595](https://github.com/dfa1/vortex-java/pull/595)) diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java index cbfb81d70..a112a82ab 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java @@ -172,6 +172,11 @@ private static long primitiveBytes(DType dtype, int n) { return (long) n * p.ptype().byteSize(); } + /// Rust's `exact_uncompressed_size` of a primitive array: its values. + private static long uncompressedBytes(DType dtype, Object data) { + return primitiveBytes(dtype, dataLength(data)); + } + /// Entry point: encode `data` using the best cascading strategy. /// /// Cascade parameters (depth, sampling, exclusions) are taken from `ctx`. @@ -259,13 +264,35 @@ private EncodeResult competeAndEncode( int sampleSize = sampleSize(n, ctx); Choice choice = chooseBest(dtype, input, ctx, baselineFn.applyAsLong(sampleSize), () -> sampleSize < n ? stratifiedSample(data, sampleSize, ctx.sampleSeed()) : data); + return encodeChoice(choice, dtype, data, ctx); + } + /// Runs the winner of a selection on the full data, or the first accepting encoder when no + /// candidate beat the baseline. + /// + /// A primitive winner is kept only when it comes out smaller than the array it encodes, as + /// Rust's `choose_and_compress` accepts it (`compressor/cascade.rs`): a sample can favor a + /// scheme the full array does not, such as bit-packing a short array that pads to a whole + /// 1024-value block. Otherwise the array stays canonical. + private EncodeResult encodeChoice(Choice choice, DType dtype, Object data, EncodeContext ctx) { + EncodingEncoder canonical = findPrimitiveEncoding(dtype, ctx.excluded()); if (choice.encoder() == null) { - // No encoding beats the baseline — fall back to the first accepting encoder - return spliceResult(findPrimitiveEncoding(dtype, ctx.excluded()), dtype, data, ctx); + return spliceResult(canonical, dtype, data, ctx); } - // Run the winner on the full data - return spliceResult(choice.encoder(), dtype, data, ctx); + EncodeResult result = spliceResult(choice.encoder(), dtype, data, ctx); + if (dtype instanceof DType.Primitive && choice.encoder() != canonical + && totalBytes(result) >= uncompressedBytes(dtype, data)) { + return spliceResult(canonical, dtype, data, ctx); + } + return result; + } + + private static long totalBytes(EncodeResult result) { + long total = 0; + for (EncodedBuffer buffer : result.encodedBuffers()) { + total += buffer.data().byteSize(); + } + return total; } /// The winner of a selection and the size it was scored at on the sample (or its estimate). diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/CascadingCompressorTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/CascadingCompressorTest.java index 6b99e1494..895a883e7 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/CascadingCompressorTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/CascadingCompressorTest.java @@ -413,6 +413,48 @@ static Stream floatArrays() { } } + /// A winner is kept only when it comes out smaller than the array, as Rust's + /// `choose_and_compress` accepts it: the sample alone can favor a scheme the full array does + /// not. + @Nested + class Acceptance { + + @Test + void bitPackingPaddedPastTheRawSize_keepsTheArrayCanonical() { + // Given — 1500 u16 values of 15 bits: a 1024-row sample packs to 1920 of 2048 bytes and + // picks bit-packing, but the full array pads to 2048 values, 3840 bytes over its 3000 + short[] values = new short[1500]; + for (int i = 0; i < values.length; i++) { + values[i] = (short) (i * 21); + } + DType u16 = new DType.Primitive(PType.U16, false); + CascadingCompressor sut = new CascadingCompressor( + List.of(new BitpackedEncodingEncoder(), new PrimitiveEncodingEncoder())); + + // When + EncodeResult result = sut.encode(u16, values, EncodeTestHelper.cascadingCtx()); + + // Then + assertThat(result.rootNode().encodingId()).isEqualTo(EncodingId.VORTEX_PRIMITIVE); + } + + @Test + void forOverAFewValues_keepsTheArrayCanonical() { + // Given — a zone table's null counts: FoR wins on its estimate, but bit-packing the + // residuals pads 4 values to a 1024-value block, past the 32 raw bytes + long[] values = {4681, 4681, 4681, 4682}; + DType u64 = new DType.Primitive(PType.U64, false); + CascadingCompressor sut = new CascadingCompressor(List.of(new FrameOfReferenceEncodingEncoder(), + new BitpackedEncodingEncoder(), new PrimitiveEncodingEncoder())); + + // When + EncodeResult result = sut.encode(u64, values, EncodeTestHelper.cascadingCtx()); + + // Then + assertThat(result.rootNode().encodingId()).isEqualTo(EncodingId.VORTEX_PRIMITIVE); + } + } + @Nested class NullableStructField { From 87d6723b0aa6de5539eb0d9188ce6caeae818368 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sun, 25 Oct 2026 19:43:25 +0100 Subject: [PATCH 02/10] fix(writer): cost a child encoding by what it encodes to Rust measures a sampled scheme by compressing the sample with it, children included (compress_child runs the full choose_and_compress on each child), so a child's cost is its winner's real output, and the child stays canonical when that is not smaller. Java picked the child's winner the same way, but when the winner won on a stats estimate it took the estimate as the cost. The estimate can be far off: a FoR ratio says nothing of the 1024-value block its bit-packing pads a short child to. On a 1024-row sample of runs of 20 wide values with null runs, FoR estimated 81 bytes for run-end's 65 ends, which encode to 1283 and are kept as 130 raw bytes. Run-end came out 49 bytes too cheap and beat RLE, which Rust picks (found comparing nullable columns against vortex-jni, #604). Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01SwhK4W5cgj7y7s4ZPCif5x --- CHANGELOG.md | 1 + .../writer/encode/CascadingCompressor.java | 21 ++++++++++++------- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0649588df..ad70488a6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **Breaking:** CSV import stores an empty cell in a number or boolean column as null instead of `0` or `false`, so inferred number and boolean columns are now nullable (read through `MaskedArray`), and an empty cell or a missing column fails the import where the schema declares the column non-nullable. ([#592](https://github.com/dfa1/vortex-java/issues/592)) ### Fixed +- The cascade costs a child encoding by what it encodes to, as Rust does, not by its estimate, which made it pick run-end where Rust picks RLE. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - The cascade keeps an encoding only when it comes out smaller than the raw array, as Rust does, so a short column no longer bit-packs into more bytes than it holds. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - An all-null chunk is written as one null `vortex.constant`, as Rust writes it, instead of a masked constant `0`. ([#594](https://github.com/dfa1/vortex-java/pull/594)) - An all-null chunk stored as a null constant (as vortex-jni writes it) reads as a masked array with every row null, like any nullable chunk, instead of a `NullArray` that the CLI, SQL and Parquet export failed on. ([#594](https://github.com/dfa1/vortex-java/pull/594)) diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java index a112a82ab..716ec009e 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java @@ -299,9 +299,10 @@ private static long totalBytes(EncodeResult result) { /// /// @param encoder the winning encoder, or `null` when none beat the baseline /// @param size the winner's sample size in bytes, or the baseline when there is none; - /// unmeasured when `alwaysUse` - /// @param alwaysUse whether the winner was settled by [Estimate#ALWAYS_USE] - private record Choice(EncodingEncoder encoder, double size, boolean alwaysUse) { + /// unmeasured when `estimated` + /// @param estimated whether the winner won without being encoded: by [Estimate#ALWAYS_USE] or + /// a stats [Estimate.Ratio], so `size` is not what it encodes to + private record Choice(EncodingEncoder encoder, double size, boolean estimated) { } /// Rust's `choose_best_scheme` (`compressor/select.rs`), for one array. @@ -323,6 +324,7 @@ private Choice chooseBest(DType dtype, ArrayAndStats input, EncodeContext ctx, l Supplier sample) { EncodingEncoder best = null; double bestSize = baseline; + boolean bestEstimated = false; List deferred = new ArrayList<>(); for (EncodingEncoder enc : encodings) { if (!enc.accepts(dtype) || ctx.excluded().contains(enc.encodingId())) { @@ -337,13 +339,14 @@ private Choice chooseBest(DType dtype, ArrayAndStats input, EncodeContext ctx, l if (size < bestSize) { best = enc; bestSize = size; + bestEstimated = true; } } else if (est == Estimate.COMPLETE) { deferred.add(enc); } } if (deferred.isEmpty()) { - return new Choice(best, bestSize, false); + return new Choice(best, bestSize, bestEstimated); } Object measured = sample.get(); @@ -358,9 +361,10 @@ private Choice chooseBest(DType dtype, ArrayAndStats input, EncodeContext ctx, l if (size < bestSize) { best = enc; bestSize = size; + bestEstimated = false; } } - return new Choice(best, bestSize, false); + return new Choice(best, bestSize, bestEstimated); } /// Derives the context a child slot is filled under: one cascade level deeper, with the slot's @@ -390,8 +394,11 @@ private long measureBestChild(DType dtype, Object data, EncodeContext ctx) { return baseline; } Choice choice = chooseBest(dtype, withStats(dtype, data, ctx), ctx, baseline, () -> data); - if (choice.alwaysUse()) { - // ALWAYS_USE: settled without a measurement, so measure the winner for its size + if (choice.estimated()) { + // Rust compresses a child for real even on a sample (`compress_child`), so an estimate + // picks the winner but its cost is what it encodes to, and the child stays canonical + // when that is not smaller: a FoR ratio says nothing of the 1024-value block its + // bit-packing pads a short child to CascadeStep step = choice.encoder().encodeCascade(dtype, data, ctx); return Math.min(baseline, measureStep(choice.encoder(), step, ctx)); } From c4b105f5f7cdc69e0a6d99a85af212381ec5cf10 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sun, 25 Oct 2026 19:45:40 +0100 Subject: [PATCH 03/10] fix(writer): FoR bit-packs its residuals instead of cascading them Rust's v1 FoRScheme declares no children: it subtracts the minimum and bit-packs the residuals at once with BITPACKING_V1 ("if any other scheme was preferable, it would be chosen instead of bitpacking", vortex-btrblocks schemes/integer/for_.rs), keeping them a primitive when the best width is the native one (compress_global returns None). Java opened the residuals as a cascade child, so any scheme could win them. On a zone table's null counts [4681, 4681, 4681, 4682] the child became a vortex.dict, giving for(dict) where vortex-jni writes a vortex.primitive (FoR wins the estimate, its bit-packing pads past the 32 raw bytes, and the array stays canonical). Found comparing nullable columns against vortex-jni (#604). Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01SwhK4W5cgj7y7s4ZPCif5x --- CHANGELOG.md | 1 + .../encode/BitpackedEncodingEncoder.java | 22 +++++++++++++ .../FrameOfReferenceEncodingEncoder.java | 23 +++++++++----- .../FrameOfReferenceEncodingEncoderTest.java | 31 +++++++++++++++++++ 4 files changed, 69 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ad70488a6..5ce9b7b24 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **Breaking:** CSV import stores an empty cell in a number or boolean column as null instead of `0` or `false`, so inferred number and boolean columns are now nullable (read through `MaskedArray`), and an empty cell or a missing column fails the import where the schema declares the column non-nullable. ([#592](https://github.com/dfa1/vortex-java/issues/592)) ### Fixed +- Frame-of-reference bit-packs its residuals directly, as Rust's FoR scheme does, instead of letting the cascade pick another encoding for them. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - The cascade costs a child encoding by what it encodes to, as Rust does, not by its estimate, which made it pick run-end where Rust picks RLE. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - The cascade keeps an encoding only when it comes out smaller than the raw array, as Rust does, so a short column no longer bit-packs into more bytes than it holds. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - An all-null chunk is written as one null `vortex.constant`, as Rust writes it, instead of a masked constant `0`. ([#594](https://github.com/dfa1/vortex-java/pull/594)) diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java index bf6581b5d..c999ddc3d 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java @@ -48,6 +48,28 @@ public boolean accepts(DType dtype) { return p.ptype().isInteger(); } + /// Rust's `BitPackingScheme::compress` run directly, as FoR runs it on its residuals: packed + /// at one global width, or kept a `vortex.primitive` when the best width is the native one + /// (`compress_global` returns `None`). + /// + /// @param dtype the non-negative integer dtype + /// @param data the values + /// @param ctx the encode context + /// @return the bit-packed or canonical array + static EncodeResult encodeLeaf(DType dtype, Object data, EncodeContext ctx) { + PType ptype = ((DType.Primitive) dtype).ptype(); + long[] longs = PrimitiveArrays.toLongs(data, ptype, EncodingId.FASTLANES_BITPACKED); + long typeMask = FastLanes.lowMask(ptype.bits()); + int[] bitWidthFreq = new int[ptype.bits() + 1]; + for (long v : longs) { + bitWidthFreq[Long.SIZE - Long.numberOfLeadingZeros(v & typeMask)]++; + } + if (bestBitWidth(bitWidthFreq, ptype.byteSize() + 4, longs.length) == ptype.bits()) { + return new PrimitiveEncodingEncoder().encode(dtype, data, ctx); + } + return new BitpackedEncodingEncoder().encode(dtype, data, ctx); + } + @Override public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { PType ptype = ((DType.Primitive) dtype).ptype(); diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FrameOfReferenceEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FrameOfReferenceEncodingEncoder.java index 6e61f1269..52859904d 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FrameOfReferenceEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FrameOfReferenceEncodingEncoder.java @@ -10,7 +10,6 @@ import io.github.dfa1.vortex.core.simd.SimdOperationsSupport; import java.lang.foreign.MemorySegment; -import java.util.Set; import java.util.List; /// Write-only encoder for `fastlanes.for` (Frame of Reference). @@ -89,14 +88,22 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext encodeC if (ref == 0L && ptype.isUnsigned()) { return CascadeStep.notApplicable(); } - MemorySegment meta = buildForMeta(ref, ptype); + return step(p, ref, residualsAsNativeArray(longs, ref, ptype), PrimitiveEncodingEncoder.minMaxStats(ptype, data), + encodeCtx); + } - EncodeNode partialRoot = new EncodeNode(EncodingId.FASTLANES_FOR, meta, new EncodeNode[1], new int[0]); - ChildSlot slot = new ChildSlot(dtype, residualsAsNativeArray(longs, ref, ptype), 0, - Set.of(EncodingId.FASTLANES_FOR)); - MinMax stats = PrimitiveEncodingEncoder.minMaxStats(ptype, data); - return new CascadeStep(partialRoot, List.of(), List.of(slot), - PrimitiveEncodingEncoder.minOf(stats), PrimitiveEncodingEncoder.maxOf(stats), true); + /// Rust's v1 `FoRScheme::compress`: the residuals are bit-packed at once, never cascaded + /// ("if any other scheme was preferable, it would be chosen instead of bitpacking", + /// `schemes/integer/for_.rs`), and stay canonical when bit-packing cannot narrow them. + private static CascadeStep step(DType.Primitive dtype, long ref, Object residuals, MinMax stats, + EncodeContext ctx) { + EncodeResult child = ctx.excluded().contains(EncodingId.FASTLANES_BITPACKED) + ? new PrimitiveEncodingEncoder().encode(dtype, residuals, ctx) + : BitpackedEncodingEncoder.encodeLeaf(dtype, residuals, ctx); + EncodeNode root = new EncodeNode(EncodingId.FASTLANES_FOR, buildForMeta(ref, dtype.ptype()), + new EncodeNode[]{child.rootNode()}, new int[0]); + return CascadeStep.terminal(new EncodeResult(root, child.encodedBuffers(), + PrimitiveEncodingEncoder.minOf(stats), PrimitiveEncodingEncoder.maxOf(stats))); } // Rust takes the reference from Stat::Min, an unsigned min for unsigned ptypes: a signed one diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/FrameOfReferenceEncodingEncoderTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/FrameOfReferenceEncodingEncoderTest.java index b96b664e6..99f687963 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/FrameOfReferenceEncodingEncoderTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/FrameOfReferenceEncodingEncoderTest.java @@ -296,5 +296,36 @@ void encodeCascade_signedI64_statsCarryMinAndMax() throws IOException { assertThat(result.statsMin().int64_value()).isEqualTo(990L); assertThat(result.statsMax().int64_value()).isEqualTo(1050L); } + + @Test + void encodeCascade_bitPacksTheResidualsWithoutCascading() { + // Given — Rust's v1 FoR has no cascaded child: it bit-packs its residuals at once, so + // no other scheme (a dict, a run-end) can win them + long[] data = new long[2048]; + for (int i = 0; i < data.length; i++) { + data[i] = 1_000_000L + i % 4; + } + + // When + CascadeStep result = ENCODER.encodeCascade(DTypes.I64, data, EncodeTestHelper.cascadingCtx()); + + // Then + assertThat(result.isTerminal()).isTrue(); + assertThat(result.partialRoot().children()[0].encodingId()).isEqualTo(EncodingId.FASTLANES_BITPACKED); + } + + @Test + void encodeCascade_residualsAtFullWidth_stayPrimitive() { + // Given — values spanning the whole i64 range, so most residuals need all 64 bits and + // bit-packing cannot narrow them (Rust's compress_global returns None): they stay a + // vortex.primitive + long[] data = {Long.MIN_VALUE, Long.MAX_VALUE, -1L, 0L}; + + // When + CascadeStep result = ENCODER.encodeCascade(DTypes.I64, data, EncodeTestHelper.cascadingCtx()); + + // Then + assertThat(result.partialRoot().children()[0].encodingId()).isEqualTo(EncodingId.VORTEX_PRIMITIVE); + } } } From 4f50db0e1368c611d9555849536229614ff93881 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sun, 25 Oct 2026 19:46:29 +0100 Subject: [PATCH 04/10] fix(writer): narrow run-end ends and RLE indices and offsets as Rust does Rust narrows these position arrays before compressing them (PrimitiveArray::narrow: the smallest unsigned type that holds the largest value): runend_encode narrows its u64 ends, and rle_compress its u16 indices and u64 offsets. Java always wrote u32 ends, u16 indices and u64 offsets. The width decides what the cascade does with them. A 32768-row column's ends fit u16, and bit-packing 15 of 16 bits pads past their raw size, so Rust keeps them a vortex.primitive where Java's u32 ends went through FoR. RLE indices within a chunk rarely pass 255, so Rust's u8 indices cost half of Java's u16, which tipped the choice between RLE and run-end (found comparing nullable columns against vortex-jni, #604). Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01SwhK4W5cgj7y7s4ZPCif5x --- CHANGELOG.md | 1 + .../writer/encode/RleEncodingEncoder.java | 29 +++++++++++++++---- .../writer/encode/RunEndEncodingEncoder.java | 7 +++-- .../writer/encode/RleEncodingEncoderTest.java | 13 +++++---- .../encode/RunEndEncodingEncoderTest.java | 5 ++-- 5 files changed, 41 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5ce9b7b24..7697166df 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **Breaking:** CSV import stores an empty cell in a number or boolean column as null instead of `0` or `false`, so inferred number and boolean columns are now nullable (read through `MaskedArray`), and an empty cell or a missing column fails the import where the schema declares the column non-nullable. ([#592](https://github.com/dfa1/vortex-java/issues/592)) ### Fixed +- Run-end ends and RLE indices and offsets are stored in the narrowest unsigned type that holds them, as Rust narrows them, instead of fixed `u32`, `u16` and `u64`. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - Frame-of-reference bit-packs its residuals directly, as Rust's FoR scheme does, instead of letting the cascade pick another encoding for them. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - The cascade costs a child encoding by what it encodes to, as Rust does, not by its estimate, which made it pick run-end where Rust picks RLE. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - The cascade keeps an encoding only when it comes out smaller than the raw array, as Rust does, so a short column no longer bit-packs into more bytes than it holds. ([#604](https://github.com/dfa1/vortex-java/issues/604)) 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 afad4a9a4..22bed2e5c 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 @@ -153,13 +153,28 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { return CascadeStep.notApplicable(); } Runs runs = runs(longs); - EncodeNode partialRoot = new EncodeNode(EncodingId.FASTLANES_RLE, MemorySegment.ofArray(runs.metadata()), + // Rust narrows the u16 indices and u64 offsets to the smallest unsigned type that holds + // their largest value before compressing them (`rle_compress`, `PrimitiveArray::narrow`) + int[] indices = new int[runs.paddedLen()]; + int maxIndex = 0; + for (int i = 0; i < indices.length; i++) { + indices[i] = Short.toUnsignedInt(runs.indices()[i]); + maxIndex = Math.max(maxIndex, indices[i]); + } + PType indicesPtype = PType.narrowestUnsigned(maxIndex); + long[] offsets = Arrays.copyOf(runs.offsets(), runs.numChunks()); + PType offsetsPtype = PType.narrowestUnsigned(offsets[offsets.length - 1]); + Object indicesData = PrimitiveArrays.fromIntsArray(indices, indicesPtype, EncodingId.FASTLANES_RLE); + EncodeNode partialRoot = new EncodeNode(EncodingId.FASTLANES_RLE, + MemorySegment.ofArray(runs.metadata(indicesPtype, offsetsPtype)), new EncodeNode[]{null, null, null}, new int[0]); long[] values = Arrays.copyOf(runs.values(), runs.valuesCount()); List slots = List.of( new ChildSlot(dtype, PrimitiveArrays.fromBitsArray(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)); + new ChildSlot(new DType.Primitive(indicesPtype, false), indicesData, 1, POSITIONS_EXCLUDED), + new ChildSlot(new DType.Primitive(offsetsPtype, false), + PrimitiveArrays.fromLongsArray(offsets, offsetsPtype, EncodingId.FASTLANES_RLE), 2, + POSITIONS_EXCLUDED)); MinMax stats = ZoneMapStats.of(dtype, data); return new CascadeStep(partialRoot, List.of(), slots, ZoneMapStats.minOf(stats), ZoneMapStats.maxOf(stats), true); } @@ -177,12 +192,16 @@ private record Runs(long[] values, int valuesCount, short[] indices, int paddedL int numChunks) { byte[] metadata() { + return metadata(PType.U16, PType.U64); + } + + byte[] metadata(PType indicesPtype, PType offsetsPtype) { return new ProtoRLEMetadata( valuesCount, paddedLen, - io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(PType.U16.ordinal()), + io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(indicesPtype.ordinal()), numChunks, - io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(PType.U64.ordinal()), + io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(offsetsPtype.ordinal()), 0L ).encode(); } diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/RunEndEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/RunEndEncodingEncoder.java index 26c5a26eb..55ff44712 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/RunEndEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/RunEndEncodingEncoder.java @@ -241,11 +241,14 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { ends[k] = n; values[k] = widened[n - 1]; + // Rust narrows the u64 ends to the smallest unsigned type that holds the last one (`n`) + PType endsPtype = PType.narrowestUnsigned(n); byte[] metaBytes = new ProtoRunEndMetadata( - io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(PType.U32.ordinal()), numRuns, 0L).encode(); + io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(endsPtype.ordinal()), numRuns, 0L).encode(); EncodeNode partialRoot = new EncodeNode(EncodingId.VORTEX_RUNEND, MemorySegment.ofArray(metaBytes), new EncodeNode[]{null, null}, new int[0]); - ChildSlot endsSlot = new ChildSlot(new DType.Primitive(PType.U32, false), ends, 0, ENDS_EXCLUDED); + ChildSlot endsSlot = new ChildSlot(new DType.Primitive(endsPtype, false), + PrimitiveArrays.fromIntsArray(ends, endsPtype, EncodingId.VORTEX_RUNEND), 0, ENDS_EXCLUDED); ChildSlot valuesSlot = new ChildSlot(dtype, PrimitiveArrays.fromLongsArray(values, ptype, EncodingId.VORTEX_RUNEND), 1, VALUES_EXCLUDED); MinMax stats = ZoneMapStats.of(dtype, data); return new CascadeStep(partialRoot, List.of(), List.of(endsSlot, valuesSlot), 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 b8b78d593..9718124aa 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 @@ -417,11 +417,14 @@ void encodeCascade_opensValuesIndicesAndOffsetsSlots() { 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); + // indices and offsets are narrowed to the smallest type holding their largest value + // (1 and 2), as Rust's rle_compress narrows them before compressing + assertThat(indices.childDtype()).isEqualTo(new DType.Primitive(io.github.dfa1.vortex.core.model.PType.U8, false)); + assertThat((byte[]) indices.childData()).hasSize(2_048); + assertThat(((byte[]) indices.childData())[599]).isZero(); + assertThat(((byte[]) indices.childData())[600]).isEqualTo((byte) 1); + assertThat(offsets.childDtype()).isEqualTo(new DType.Primitive(io.github.dfa1.vortex.core.model.PType.U8, false)); + assertThat((byte[]) offsets.childData()).containsExactly(0, 2); for (ChildSlot positions : List.of(indices, offsets)) { assertThat(positions.excluded()).containsExactlyInAnyOrder(EncodingId.FASTLANES_RLE, EncodingId.VORTEX_DICT, EncodingId.VORTEX_SPARSE); diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/RunEndEncodingEncoderTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/RunEndEncodingEncoderTest.java index 1f8a58ec8..6404f6104 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/RunEndEncodingEncoderTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/RunEndEncodingEncoderTest.java @@ -242,8 +242,9 @@ void encodeCascade_splitsRunsIntoEndsAndValuesSlots() { assertThat(result.ownedBuffers()).isEmpty(); ChildSlot ends = result.openChildren().get(0); ChildSlot values = result.openChildren().get(1); - assertThat((int[]) ends.childData()).containsExactly(3, 5, 6); - assertThat(ends.childDtype()).isEqualTo(new DType.Primitive(PType.U32, false)); + // ends are narrowed to the smallest type holding the last one, as Rust narrows them + assertThat((byte[]) ends.childData()).containsExactly(3, 5, 6); + assertThat(ends.childDtype()).isEqualTo(new DType.Primitive(PType.U8, false)); assertThat(ends.excluded()).containsExactlyInAnyOrder(EncodingId.VORTEX_DICT, EncodingId.VORTEX_RUNEND, EncodingId.FASTLANES_RLE, EncodingId.VORTEX_SPARSE); assertThat((long[]) values.childData()).containsExactly(7L, -2L, 7L); From 772bca6e76d38c77ebd56eeee1dfb5e3b4901c9f Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sun, 25 Oct 2026 19:47:24 +0100 Subject: [PATCH 05/10] fix(writer): compress zone stats tables with the full compressor Rust's file strategy builds two compressors from one builder: the data compressor excludes IntDictScheme (DictStrategy already dictionary-encodes suitable columns as a layout), while the stats compressor, which also compresses dict values, keeps every scheme (vortex-file strategy.rs). Java compressed stats tables with its data cascade codecs, so an integer dict never competed there. A zone table column with a repeated value differs: the per-zone null counts [1, 0, 0, 0] of a column with one null are a vortex.dict in vortex-jni's file (16 bytes of values and 4 of codes against 32 raw). The candidates are now the probe's codecs, which are the data cascade with the integer dict back in (found comparing nullable columns, #604). Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01SwhK4W5cgj7y7s4ZPCif5x --- CHANGELOG.md | 1 + .../java/io/github/dfa1/vortex/writer/VortexWriter.java | 6 +++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7697166df..12d7d2c2f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **Breaking:** CSV import stores an empty cell in a number or boolean column as null instead of `0` or `false`, so inferred number and boolean columns are now nullable (read through `MaskedArray`), and an empty cell or a missing column fails the import where the schema declares the column non-nullable. ([#592](https://github.com/dfa1/vortex-java/issues/592)) ### Fixed +- Zone stats tables are compressed with integer dictionaries among the candidates, as Rust's stats compressor has them. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - Run-end ends and RLE indices and offsets are stored in the narrowest unsigned type that holds them, as Rust narrows them, instead of fixed `u32`, `u16` and `u64`. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - Frame-of-reference bit-packs its residuals directly, as Rust's FoR scheme does, instead of letting the cascade pick another encoding for them. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - The cascade costs a child encoding by what it encodes to, as Rust does, not by its estimate, which made it pick run-end where Rust picks RLE. ([#604](https://github.com/dfa1/vortex-java/issues/604)) diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java b/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java index 9ccea0cd6..f8b73e872 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java @@ -1320,8 +1320,12 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { /// list, as Rust's stats strategy always compresses it: [CascadingCompressor] encodes the /// struct itself (including the nullable `bounded_max` state) and falls back to canonical /// encodings for any field no candidate accepts. + /// + /// The candidates are the probe's: Rust compresses stats tables with its full compressor, + /// integer dictionaries included, and keeps `IntDictScheme` out of the data compressor only + /// (`vortex-file` `strategy.rs`), so a zone's `null_count` column can be a `vortex.dict`. private EncodingEncoder zoneTableEncoder() { - List codecs = cascadeCodecs; + List codecs = probeCodecs; return new EncodingEncoder() { @Override public EncodingId encodingId() { From 6115747932fecd194752fc4827475d8688b9ce68 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sun, 25 Oct 2026 19:52:27 +0100 Subject: [PATCH 06/10] fix(writer): nullable integers keep their validity in the encoding (#604) Java wrote a nullable integer column as vortex.masked(values, validity) and ran the cascade over the values alone. Rust has no such wrapper: every array carries its own validity, so the leaf that stores the values holds it as a plain vortex.bool child (vortex.primitive, fastlanes.bitpacked after any patches), and the encodings above pass it down: FoR and ZigZag to their output, run-end to its run values (a run of nulls is one invalid run), RLE to its indices padded to a whole chunk, dict to its codes, Pco by compressing only the valid values. Rust never compresses that bitmap. The wrapper changed more than the tree. Rust's sampled schemes are measured with the bitmap, its estimated ones are not, so the two compressors ranked schemes differently on nullable data, and the dict layout could not take Rust's probe. Now: - NullableData flows through CascadingCompressor for integer dtypes. Stats skip the nulls as Rust's do (computed over the valid values: same min, max, distinct count and runs, since a null continues the run before it), the baseline adds the packed bitmap, and only encoders that keep validity compete (acceptsNullable). - The compressor's own checks come first, as in choose_and_compress: all null is a null constant, valid rows holding one value a constant masked by a plain bitmap (Rust's compress_constant). A column without nulls stores no bitmap, as Rust imports an all-valid Arrow array. - Sparse follows Sparse::encode: past 90% null the fill is null and the valid rows are the patches (with Rust's all-valid bitmap on their values); otherwise nulls are patches. - Bit-packing counts a null at width 0 and patches only valid values (bit_width_histogram, gather_patches); FoR takes its reference from the valid values and zeroes null residuals. - The dict layout probe now runs for nullable integers too, as Rust's DictStrategy does. Floats keep vortex.masked until ALP and ALP-RD carry a validity (keepsValidity). NullableIntegerInteropIntegrationTest writes 17 column shapes through vortex-jni and vortex-java and asserts the same encoding tree, data and zone table, then reads Java's file (default and compact) through vortex-jni, full, filtered and row-range. Not covered: patched bit-packing, whose chunk offsets Rust stores and Java does not yet. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01SwhK4W5cgj7y7s4ZPCif5x --- CHANGELOG.md | 1 + docs/compatibility.md | 2 +- docs/explanation.md | 4 +- .../DictAdmissionInteropIntegrationTest.java | 9 +- ...NullableIntegerInteropIntegrationTest.java | 283 ++++++++++++++++++ .../ParquetImportIntegrationTest.java | 7 +- .../vortex/parquet/ParquetImporterTest.java | 8 +- .../dfa1/vortex/writer/VortexWriter.java | 17 +- .../vortex/writer/encode/ArrayAndStats.java | 29 +- .../encode/BitpackedEncodingEncoder.java | 53 +++- .../writer/encode/CascadingCompressor.java | 89 +++++- .../writer/encode/DictEncodingEncoder.java | 46 ++- .../vortex/writer/encode/EncodeResult.java | 20 ++ .../FrameOfReferenceEncodingEncoder.java | 30 ++ .../writer/encode/MaskedEncodingEncoder.java | 26 ++ .../vortex/writer/encode/NullableData.java | 94 +++++- .../writer/encode/PcoEncodingEncoder.java | 12 + .../encode/PrimitiveEncodingEncoder.java | 12 + .../writer/encode/RleEncodingEncoder.java | 46 ++- .../writer/encode/RunEndEncodingEncoder.java | 51 +++- .../writer/encode/SparseEncodingEncoder.java | 108 ++++++- .../writer/encode/ZigZagEncodingEncoder.java | 29 +- .../writer/encode/ArrayAndStatsTest.java | 21 ++ .../encode/CascadingCompressorTest.java | 141 +++++++++ 24 files changed, 1051 insertions(+), 87 deletions(-) create mode 100644 integration/src/test/java/io/github/dfa1/vortex/integration/NullableIntegerInteropIntegrationTest.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 12d7d2c2f..08ec66c64 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **Breaking:** CSV import stores an empty cell in a number or boolean column as null instead of `0` or `false`, so inferred number and boolean columns are now nullable (read through `MaskedArray`), and an empty cell or a missing column fails the import where the schema declares the column non-nullable. ([#592](https://github.com/dfa1/vortex-java/issues/592)) ### Fixed +- A nullable integer column keeps its validity in the encoding that stores its values instead of a `vortex.masked` wrapper, as Rust writes it, so the cascade and the dict layout pick what Rust picks. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - Zone stats tables are compressed with integer dictionaries among the candidates, as Rust's stats compressor has them. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - Run-end ends and RLE indices and offsets are stored in the narrowest unsigned type that holds them, as Rust narrows them, instead of fixed `u32`, `u16` and `u64`. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - Frame-of-reference bit-packs its residuals directly, as Rust's FoR scheme does, instead of letting the cascade pick another encoding for them. ([#604](https://github.com/dfa1/vortex-java/issues/604)) diff --git a/docs/compatibility.md b/docs/compatibility.md index d297b5d37..48ddb880b 100644 --- a/docs/compatibility.md +++ b/docs/compatibility.md @@ -135,7 +135,7 @@ decimals ([#430](https://github.com/dfa1/vortex-java/pull/430)) and nulls in nul | `vortex.map` | `MapEncodingDecoder` | `MapEncodingEncoder` | ✅ | ✅ | Single `entries` child, a bare `vortex.listview` of non-nullable `{key, value}` structs (any other entries encoding is rejected, as in Rust); map validity is delegated to that child's validity slot. Single-flat columns only: neither `ChunkedArrayCombiner` nor `ScanIterator`'s row-window slicing has a `MapArray` arm yet, so a chunked map column or a split-window scan fails as `VortexException`. Pre-existing gap shared with `vortex.listview`; tracked separately | | `vortex.fixed_size_list` | `FixedSizeListEncodingDecoder` | `FixedSizeListEncodingEncoder` | ✅ | ✅ | | | `vortex.zstd` | `ZstdEncodingDecoder` | `ZstdEncodingEncoder` | ✅ | ✅ | Primitive, Utf8, Binary | -| `vortex.masked` | `MaskedEncodingDecoder` | `MaskedEncodingEncoder` | ✅ | ✅ | NullableData carrier; inner values cascade (Dict/FSST/ALP/…) when cascading is enabled; Utf8/Primitive still compete on measured size even at depth 0, else first-match FixedSizeList. A nullable Utf8 column whose per-chunk dict is smaller skips the mask entirely: null is one more dict entry, invalid in the pool's validity, with non-nullable codes, as Rust's dict builder writes it | +| `vortex.masked` | `MaskedEncodingDecoder` | `MaskedEncodingEncoder` | ✅ | ✅ | NullableData carrier. A nullable integer column is not masked when cascading: as in Rust, its validity is a plain `vortex.bool` child of the encoding that stores the values (`vortex.primitive`, `fastlanes.bitpacked`, the run values of `vortex.runend`, the indices of `fastlanes.rle`, the codes of `vortex.dict`, `vortex.pco` over the valid values only), passed down by `fastlanes.for` and `vortex.zigzag`, and one whose valid rows hold one value is a `vortex.constant` masked by a plain bitmap. Other nullable columns: inner values cascade (Dict/FSST/ALP/…) when cascading is enabled; Utf8/Primitive still compete on measured size even at depth 0, else first-match FixedSizeList. A nullable Utf8 column whose per-chunk dict is smaller skips the mask entirely: null is one more dict entry, invalid in the pool's validity, with non-nullable codes, as Rust's dict builder writes it | | `vortex.decimal` | `DecimalEncodingDecoder` | `DecimalEncodingEncoder` | ✅ | ✅ | Precision 1–76 (i8 … i256 storage by precision), negative scales. Written from `BigDecimal[]` | | `vortex.decimal_byte_parts` | `DecimalBytePartsEncodingDecoder`| `DecimalBytePartsEncodingEncoder`| ✅ | ✅ | Write: precision ≤ 18 only (one i64 mantissa, no low parts); competes in the cascade | | `vortex.datetimeparts` | `DateTimePartsEncodingDecoder` | `DateTimePartsEncodingEncoder` | ✅ | ✅ | | diff --git a/docs/explanation.md b/docs/explanation.md index 30ebca470..b461fd8f4 100644 --- a/docs/explanation.md +++ b/docs/explanation.md @@ -354,8 +354,8 @@ single `entries` child. `vortex.masked` + `vortex.list` as before. - **A null *entry value*** (`DType.Map(key, value.asNullable(), ...)`, e.g. `{a: 1, b: null}` inside an otherwise-present map row) is a completely different bit: it rides the entry - struct's own `value` field validity — ordinary `vortex.masked` wrapping, exactly like a - nullable field anywhere else inside a `vortex.struct`. It has nothing to do with the map row + struct's own `value` field validity — kept by the encoding of an integer value and wrapped in + `vortex.masked` otherwise, exactly like a nullable field anywhere else inside a `vortex.struct`. It has nothing to do with the map row being present or absent. See [how-to.md#write-and-read-a-map-column](how-to.md#write-and-read-a-map-column) for the diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/DictAdmissionInteropIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/DictAdmissionInteropIntegrationTest.java index 9abe7f2e6..6210a8d1b 100644 --- a/integration/src/test/java/io/github/dfa1/vortex/integration/DictAdmissionInteropIntegrationTest.java +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/DictAdmissionInteropIntegrationTest.java @@ -72,11 +72,12 @@ static Stream columns() { Arguments.of("f64, 5 distinct", Kind.F64, 5, 0), Arguments.of("utf8, 5 distinct", Kind.UTF8, 5, 0), Arguments.of("nullable i32, 10 distinct", Kind.I32, 10, 7), - // under half distinct, yet bit-packing 7 bits wins over a 100-value dictionary. - // A nullable column with 1000 distinct values still gets a dict in Java only: its - // admission keeps the cardinality check until the nullable cascade charges - // bit-packing the validity, as Rust's sampled sizes do. + // under half distinct, yet bit-packing 7 bits wins over a 100-value dictionary Arguments.of("i64, 100 distinct", Kind.I64, 100, 0), + // the probe charges bit-packing the validity, as Rust's sampled sizes do (#604): + // 10 bits a row plus a null bit still beat a 1000-value dictionary + Arguments.of("nullable i32, 1000 distinct", Kind.I32, 1000, 7), + Arguments.of("nullable i64, 100 distinct", Kind.I64, 100, 7), Arguments.of("i32, 1000 distinct", Kind.I32, 1000, 0), Arguments.of("utf8, 1000 distinct", Kind.UTF8, 1000, 0), Arguments.of("utf8, mostly distinct", Kind.UTF8, 1_000_000, 0)); diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/NullableIntegerInteropIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/NullableIntegerInteropIntegrationTest.java new file mode 100644 index 000000000..2e1034f18 --- /dev/null +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/NullableIntegerInteropIntegrationTest.java @@ -0,0 +1,283 @@ +package io.github.dfa1.vortex.integration; + +import dev.vortex.api.DataSource; +import dev.vortex.api.ImmutableScanOptions; +import dev.vortex.api.Expression; +import dev.vortex.api.Partition; +import dev.vortex.api.Scan; +import dev.vortex.api.ScanOptions; +import dev.vortex.api.Session; +import dev.vortex.api.VortexWriter; +import dev.vortex.arrow.ArrowAllocation; +import dev.vortex.jni.NativeLoader; +import io.github.dfa1.vortex.core.fbs.FbsArray; +import io.github.dfa1.vortex.core.fbs.FbsArrayNode; +import io.github.dfa1.vortex.core.model.ColumnName; +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.reader.SegmentSpec; +import io.github.dfa1.vortex.reader.VortexReader; +import io.github.dfa1.vortex.reader.layout.Layout; +import io.github.dfa1.vortex.writer.WriteOptions; +import io.github.dfa1.vortex.writer.encode.NullableData; +import org.apache.arrow.c.ArrowArray; +import org.apache.arrow.c.ArrowSchema; +import org.apache.arrow.c.Data; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import java.io.IOException; +import java.lang.foreign.MemorySegment; +import java.nio.channels.FileChannel; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.function.IntPredicate; +import java.util.function.IntToLongFunction; +import java.util.stream.Stream; + +import static io.github.dfa1.vortex.core.io.VortexFormat.LE_INT; +import static org.assertj.core.api.Assertions.assertThat; + +/// A nullable integer column keeps its validity in the encoding that stores its values, as Rust's +/// arrays do (#604): a `vortex.bool` child of the bit-packed or primitive leaf, passed down by +/// FoR, ZigZag, run-end, RLE and dict codes, never a `vortex.masked` wrapper. The cascade then +/// charges the validity to the schemes it samples, so it picks what Rust picks, and the dict +/// layout takes Rust's probe instead of a cardinality rule. +/// +/// Each column goes through vortex-jni and vortex-java, and the two files must hold the same +/// encoding tree, data and zone stats table alike. Each shape steers the cascade to a different +/// winner. Not covered: a column whose bit-packing needs patches, as Rust also stores the patches' +/// chunk offsets and Java does not yet. +class NullableIntegerInteropIntegrationTest { + + private static final Session SESSION = Session.create(); + private static final BufferAllocator ALLOCATOR = ArrowAllocation.rootAllocator(); + private static final ColumnName V = ColumnName.of("v"); + private static final int ROWS = 4 * 8192; + private static final long[] RANDOM = new Random(7).longs(ROWS).toArray(); + + static { + NativeLoader.loadJni(); + } + + static Stream shapes() { + IntPredicate every7 = i -> i % 7 == 0; + return Stream.of( + // no null: written without validity, as Rust writes an all-valid Arrow array + shape("no nulls", PType.I32, i -> Math.floorMod(RANDOM[i], 16), _ -> false), + // valid rows all 42: a constant masked by a plain bitmap + shape("constant with nulls", PType.I32, _ -> 42, every7), + // small values: bit-packed, the bitmap on the bit-packed array + shape("small values", PType.I32, i -> Math.floorMod(RANDOM[i], 16), every7), + // 10 distinct: a dict layout, its values pool keeping the null entry's bitmap + shape("10 distinct", PType.I32, i -> Math.floorMod(RANDOM[i], 10), every7), + // 1000 distinct: bit-packing 10 bits a row beats a dict once it pays for nulls + shape("1000 distinct", PType.I32, i -> Math.floorMod(RANDOM[i], 1000), every7), + // a large offset: frame of reference, then bit-packing with the bitmap + shape("offset values", PType.I64, i -> 1_000_000_000L + Math.floorMod(RANDOM[i], 1000), every7), + shape("negative values", PType.I64, i -> Math.floorMod(RANDOM[i], 1000) - 500, every7), + shape("sorted steps", PType.I64, i -> i * 1000L + Math.floorMod(RANDOM[i], 7), every7), + shape("sequence", PType.I32, i -> i, every7), + // run of nulls: a run-end whose run values keep a bitmap + shape("runs with null runs", PType.I32, i -> Math.floorMod(RANDOM[i / 6], 1000), i -> (i / 6) % 5 == 0), + shape("runs with scattered nulls", PType.I32, i -> Math.floorMod(RANDOM[i / 6], 1000), i -> i % 13 == 0), + shape("long runs", PType.I64, i -> (i / 100) * 37L, every7), + // wide runs: RLE, its narrowed indices keeping the bitmap padded to a whole chunk + shape("wide runs with null runs", PType.I64, i -> RANDOM[i / 20] >>> 20, i -> (i / 20) % 7 == 0), + shape("wide runs with scattered nulls", PType.I64, i -> RANDOM[i / 20] >>> 20, i -> i % 13 == 0), + // mostly null: sparse with a null fill, the valid rows as patches + shape("95% null", PType.I32, i -> Math.floorMod(RANDOM[i], 1000), i -> Math.floorMod(RANDOM[i] >>> 20, 100) < 95), + shape("clustered nulls", PType.I32, i -> Math.floorMod(RANDOM[i], 1000), i -> (i / 500) % 3 == 0), + // one null in one zone: that zone's null count makes the stats table's dict + shape("one null", PType.I32, i -> Math.floorMod(RANDOM[i], 1000), i -> i == 5)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("shapes") + void javaWritesTheEncodingTreeJniWrites(String name, PType ptype, IntToLongFunction values, IntPredicate isNull, + @TempDir Path tmp) throws Exception { + // Given + Path jniFile = tmp.resolve("jni.vortex"); + Path javaFile = tmp.resolve("java.vortex"); + writeWithJni(jniFile, ptype, values, isNull); + List expected = encodingTree(jniFile); + + // When + writeWithJava(javaFile, ptype, values, isNull, WriteOptions.defaults()); + + // Then + assertThat(encodingTree(javaFile)).isEqualTo(expected); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("shapes") + void jniReadsJavaFile(String name, PType ptype, IntToLongFunction values, IntPredicate isNull, + @TempDir Path tmp) throws Exception { + assertJniReads(ptype, values, isNull, WriteOptions.defaults(), tmp); + } + + /// Compact writes add Pco, which keeps only the valid values and the bitmap as its child. + /// vortex-jni cannot write compact files, so here only its read of Java's file checks them. + @ParameterizedTest(name = "{0}") + @MethodSource("shapes") + void jniReadsJavaCompactFile(String name, PType ptype, IntToLongFunction values, IntPredicate isNull, + @TempDir Path tmp) throws Exception { + assertJniReads(ptype, values, isNull, WriteOptions.defaults().withCompact(true), tmp); + } + + private static void assertJniReads(PType ptype, IntToLongFunction values, IntPredicate isNull, + WriteOptions options, Path tmp) throws Exception { + // Given + Path file = tmp.resolve("java.vortex"); + writeWithJava(file, ptype, values, isNull, options); + List expected = new ArrayList<>(); + for (int i = 0; i < ROWS; i++) { + expected.add(isNull.test(i) ? null : values.applyAsLong(i)); + } + long threshold = values.applyAsLong(ROWS / 2); + + // When + List result = readWithJni(file, ScanOptions.builder()); + // a filtered and a row-range read prune and slice, which a full scan never does + List filtered = readWithJni(file, ScanOptions.builder().filter(Expression.binary( + Expression.BinaryOp.GTE, Expression.column("v"), literal(ptype, threshold)))); + List range = readWithJni(file, ScanOptions.builder().rowRangeBegin(1000).rowRangeEnd(20_000)); + + // Then + assertThat(result).isEqualTo(expected); + assertThat(filtered).isEqualTo(expected.stream().filter(v -> v != null && v >= threshold).toList()); + assertThat(range).isEqualTo(expected.subList(1000, 20_000)); + } + + private static Arguments shape(String name, PType ptype, IntToLongFunction values, IntPredicate isNull) { + return Arguments.of(name, ptype, values, isNull); + } + + private static Expression literal(PType ptype, long value) { + if (ptype == PType.I32) { + return Expression.literal((int) value); + } + return Expression.literal(value); + } + + private static void writeWithJni(Path path, PType ptype, IntToLongFunction values, IntPredicate isNull) + throws Exception { + Schema schema = new Schema(List.of(Field.nullable("v", new ArrowType.Int(ptype.bits(), true)))); + try (VortexWriter writer = VortexWriter.builder(SESSION, path.toAbsolutePath().toUri().toString(), schema, + ALLOCATOR).build(); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, ALLOCATOR)) { + FieldVector vector = root.getVector("v"); + vector.setInitialCapacity(ROWS); + vector.allocateNew(); + for (int i = 0; i < ROWS; i++) { + if (isNull.test(i)) { + vector.setNull(i); + } else if (ptype == PType.I32) { + ((IntVector) vector).setSafe(i, (int) values.applyAsLong(i)); + } else { + ((BigIntVector) vector).setSafe(i, values.applyAsLong(i)); + } + } + vector.setValueCount(ROWS); + root.setRowCount(ROWS); + try (ArrowArray array = ArrowArray.allocateNew(ALLOCATOR); + ArrowSchema arrowSchema = ArrowSchema.allocateNew(ALLOCATOR)) { + Data.exportVectorSchemaRoot(ALLOCATOR, root, null, array, arrowSchema); + writer.writeBatch(array.memoryAddress(), arrowSchema.memoryAddress()); + } + } + } + + private static void writeWithJava(Path path, PType ptype, IntToLongFunction values, IntPredicate isNull, + WriteOptions options) throws IOException { + boolean[] valid = new boolean[ROWS]; + int[] ints = new int[ROWS]; + long[] longs = new long[ROWS]; + for (int i = 0; i < ROWS; i++) { + valid[i] = !isNull.test(i); + // a null slot holds 0, as an Arrow null does + long v = valid[i] ? values.applyAsLong(i) : 0; + ints[i] = (int) v; + longs[i] = v; + } + Object data = ptype == PType.I32 ? ints : longs; + DType.Struct schema = new DType.Struct(List.of(V), List.of(new DType.Primitive(ptype, true)), false); + try (var channel = FileChannel.open(path, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var writer = io.github.dfa1.vortex.writer.VortexWriter.create(channel, schema, options)) { + writer.writeChunk(Map.of(V, new NullableData(data, valid))); + } + } + + private static List readWithJni(Path file, ImmutableScanOptions.Builder options) throws IOException { + ScanOptions opts = options.projection(Expression.select(new String[]{"v"}, Expression.root())).build(); + List out = new ArrayList<>(); + DataSource ds = DataSource.open(SESSION, file.toAbsolutePath().toUri().toString()); + Scan scan = ds.scan(opts); + while (scan.hasNext()) { + Partition partition = scan.next(); + try (ArrowReader reader = partition.scanArrow(ALLOCATOR)) { + while (reader.loadNextBatch()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + FieldVector vec = root.getVector("v"); + for (int i = 0; i < root.getRowCount(); i++) { + out.add(vec.isNull(i) ? null : ((Number) vec.getObject(i)).longValue()); + } + } + } + } + return out; + } + + /// The file's layouts and array encodings, depth first, one indented id a line. A chunked + /// layout is left out: Java keeps one around a single chunk, where Rust writes the flat + /// layout directly, and that layout shape is not what these columns test. + private static List encodingTree(Path file) throws IOException { + List out = new ArrayList<>(); + try (VortexReader reader = VortexReader.open(file)) { + layoutTree(reader, reader.layout(), 0, out); + } + return out; + } + + private static void layoutTree(VortexReader reader, Layout layout, int depth, List out) { + if (layout.isChunked() && layout.children().size() == 1) { + layoutTree(reader, layout.children().getFirst(), depth, out); + return; + } + out.add(" ".repeat(depth) + layout.layoutId().id()); + if (layout.isFlat() && !layout.segments().isEmpty()) { + SegmentSpec spec = reader.footer().segmentSpecs().get(layout.segments().getFirst()); + MemorySegment seg = reader.rawSegment(spec); + int fbLen = seg.get(LE_INT, seg.byteSize() - 4L); + FbsArray array = FbsArray.getRootAsFbsArray(seg.asSlice(seg.byteSize() - 4L - fbLen, fbLen)); + arrayTree(array.root(), reader.footer().arraySpecs(), depth + 1, out); + return; + } + for (Layout child : layout.children()) { + layoutTree(reader, child, depth + 1, out); + } + } + + private static void arrayTree(FbsArrayNode node, List specs, int depth, List out) { + out.add(" ".repeat(depth) + specs.get(node.encoding()).id()); + for (int i = 0; i < node.childrenLength(); i++) { + arrayTree(node.children(i), specs, depth + 1, out); + } + } +} diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/ParquetImportIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/ParquetImportIntegrationTest.java index 50c111e6b..3763da4c1 100644 --- a/integration/src/test/java/io/github/dfa1/vortex/integration/ParquetImportIntegrationTest.java +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/ParquetImportIntegrationTest.java @@ -112,14 +112,13 @@ void longColumnValuesMatch(@TempDir Path tmp) throws Exception { // When ParquetImporter.importParquet(parquetFile, vortexFile); - // Then — c_customer_sk (INT64, nullable): first 3 values are 100, 99, 98. Nullable, so it - // round-trips as a MaskedArray (validity + Primitive values child). + // Then — c_customer_sk (INT64, nullable): first 3 values are 100, 99, 98. Nullable but + // without a null, so it is written without a validity, as Rust writes an all-valid array. try (VortexReader reader = VortexReader.open(vortexFile); ScanIterator iter = reader.scan(ScanOptions.all())) { assertThat(iter.hasNext()).isTrue(); try (Chunk first = iter.next()) { - MaskedArray col = first.column(ColumnName.of("c_customer_sk")); - LongArray colValues = (LongArray) col.inner(); + LongArray colValues = first.column(ColumnName.of("c_customer_sk")); assertThat(colValues.getLong(0)).isEqualTo(100L); assertThat(colValues.getLong(1)).isEqualTo(99L); assertThat(colValues.getLong(2)).isEqualTo(98L); diff --git a/parquet/src/test/java/io/github/dfa1/vortex/parquet/ParquetImporterTest.java b/parquet/src/test/java/io/github/dfa1/vortex/parquet/ParquetImporterTest.java index db013ff13..473cdaa00 100644 --- a/parquet/src/test/java/io/github/dfa1/vortex/parquet/ParquetImporterTest.java +++ b/parquet/src/test/java/io/github/dfa1/vortex/parquet/ParquetImporterTest.java @@ -400,11 +400,9 @@ void importsFixture_columnValuesRoundTrip(@TempDir Path tmp) throws Exception { ScanIterator iter = reader.scan(ScanOptions.all())) { assertThat(iter.hasNext()).isTrue(); try (Chunk first = iter.next()) { - // c_customer_sk is nullable I64, so it round-trips as a MaskedArray (validity - // + Primitive values child): a boxed Long[] buffer during import, so a null - // row is tracked instead of silently becoming 0. First three rows are non-null. - MaskedArray sk = first.column(ColumnName.of("c_customer_sk")); - LongArray skValues = (LongArray) sk.inner(); + // c_customer_sk is nullable I64 but holds no null, so it is written without a + // validity, as Rust writes an all-valid array, and reads as plain longs + LongArray skValues = first.column(ColumnName.of("c_customer_sk")); assertThat(skValues.getLong(0)).isEqualTo(100L); assertThat(skValues.getLong(1)).isEqualTo(99L); assertThat(skValues.getLong(2)).isEqualTo(98L); diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java b/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java index f8b73e872..c4f40a5eb 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java @@ -654,11 +654,13 @@ public void writeChunk(Map columns) throws IOException { // Rust's dict layout admits Primitive | Utf8 | Binary (dict_layout_supported). candidate = DictColumnState.isVarBinDictCandidate((Object[]) values, validity); } - // ponytail: only non-null primitives take Rust's probe. A nullable column's cascade - // runs under vortex.masked, which leaves the validity out of the sampled sizes Rust - // charges bit-packing with, and strings pick OnPair where Rust's StringDictScheme - // wins, so for those the cardinality check above still matches Rust more often. - if (candidate && colDtype instanceof DType.Primitive p && !p.nullable()) { + // ponytail: only primitives whose cascade sees their validity take Rust's probe + // (#604). Nullable floats still run under vortex.masked, which leaves the validity + // out of the sampled sizes Rust charges bit-packing with, and strings pick OnPair + // where Rust's StringDictScheme wins, so for those the cardinality check above + // still matches Rust more often. + if (candidate && colDtype instanceof DType.Primitive p + && (!p.nullable() || CascadingCompressor.keepsValidity(p))) { candidate = probeCompressesToDict(colDtype, data); } if (candidate) { @@ -874,9 +876,12 @@ private EncodingEncoder resolveOverride(DType dtype, Object data, EncodingEncode if (encodingOverride == null && dtype instanceof DType.Map) { return new io.github.dfa1.vortex.writer.encode.MapEncodingEncoder(); } + // With cascading on, a nullable integer column goes to the compressor, which keeps the + // validity in the encoding that stores the values, as Rust does (#604) + boolean keepsValidity = options.allowedCascading() > 0 && CascadingCompressor.keepsValidity(dtype); if (encodingOverride == null && data instanceof io.github.dfa1.vortex.writer.encode.NullableData - && !(dtype instanceof DType.Extension)) { + && !(dtype instanceof DType.Extension) && !keepsValidity) { EncodingEncoder nullableCapable = nullableCapableEncoder(dtype); return nullableCapable != null ? nullableCapable : new MaskedEncodingEncoder(); } diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ArrayAndStats.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ArrayAndStats.java index 75687044d..3cd46bffb 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ArrayAndStats.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ArrayAndStats.java @@ -2,6 +2,8 @@ import io.github.dfa1.vortex.core.model.DType; +import java.lang.reflect.Array; + /// An array bundled with its lazily computed [ArrayStats]: Rust's `ArrayAndStats` /// (`vortex-compressor/src/stats/cache.rs`). /// @@ -47,10 +49,31 @@ public Object data() { /// @return the stats under this bundle's [StatsOptions] public ArrayStats stats() { if (stats == null) { - stats = dtype instanceof DType.Primitive p - ? ArrayStats.compute(p.ptype(), data, options) - : ArrayStats.EMPTY; + stats = ArrayStats.EMPTY; + if (dtype instanceof DType.Primitive p) { + // Rust's stats skip nulls: min, max and the distinct values are the valid rows', + // and a null continues the run before it, so runs are counted over the valid rows + Object values = data instanceof NullableData nullable ? nullable.validValues() : data; + stats = ArrayStats.compute(p.ptype(), values, options); + } } return stats; } + + /// Returns the number of rows, null or not. + /// + /// @return the array length + public int length() { + if (data instanceof NullableData nullable) { + return nullable.length(); + } + return Array.getLength(data); + } + + /// Returns the number of null rows: Rust's `null_count`. + /// + /// @return how many rows are null; `0` for an array without validity + public int nullCount() { + return data instanceof NullableData nullable ? nullable.nullCount() : 0; + } } diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java index c999ddc3d..9834a21f7 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java @@ -48,21 +48,45 @@ public boolean accepts(DType dtype) { return p.ptype().isInteger(); } + /// A nullable array keeps its validity as a trailing `vortex.bool` child, after any patches, + /// as Rust's `BitPackedArray` does. + @Override + public boolean acceptsNullable(DType dtype) { + return accepts(dtype); + } + + @Override + public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { + if (data instanceof NullableData nullable) { + return encodeNullable(dtype, nullable, ctx); + } + PType ptype = ((DType.Primitive) dtype).ptype(); + long[] longs = PrimitiveArrays.toLongs(data, ptype, EncodingId.FASTLANES_BITPACKED); + return encodeLongs(ptype, longs, ctx); + } + /// Rust's `BitPackingScheme::compress` run directly, as FoR runs it on its residuals: packed /// at one global width, or kept a `vortex.primitive` when the best width is the native one /// (`compress_global` returns `None`). /// /// @param dtype the non-negative integer dtype - /// @param data the values + /// @param data the values, dense or a [NullableData] /// @param ctx the encode context /// @return the bit-packed or canonical array static EncodeResult encodeLeaf(DType dtype, Object data, EncodeContext ctx) { PType ptype = ((DType.Primitive) dtype).ptype(); - long[] longs = PrimitiveArrays.toLongs(data, ptype, EncodingId.FASTLANES_BITPACKED); + boolean[] validity = null; + Object values = data; + if (data instanceof NullableData nullable) { + validity = nullable.validity(); + values = nullable.values(); + } + long[] longs = PrimitiveArrays.toLongs(values, ptype, EncodingId.FASTLANES_BITPACKED); long typeMask = FastLanes.lowMask(ptype.bits()); int[] bitWidthFreq = new int[ptype.bits() + 1]; - for (long v : longs) { - bitWidthFreq[Long.SIZE - Long.numberOfLeadingZeros(v & typeMask)]++; + for (int i = 0; i < longs.length; i++) { + int width = validity == null || validity[i] ? Long.SIZE - Long.numberOfLeadingZeros(longs[i] & typeMask) : 0; + bitWidthFreq[width]++; } if (bestBitWidth(bitWidthFreq, ptype.byteSize() + 4, longs.length) == ptype.bits()) { return new PrimitiveEncodingEncoder().encode(dtype, data, ctx); @@ -70,10 +94,25 @@ static EncodeResult encodeLeaf(DType dtype, Object data, EncodeContext ctx) { return new BitpackedEncodingEncoder().encode(dtype, data, ctx); } - @Override - public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { + /// Rust's `bitpack_encode` on a nullable array: the bit-width histogram counts a null at width + /// 0 and only valid values become patches (`bit_width_histogram`, `gather_patches`). A null slot + /// is zeroed first, which gives both, and packs as Rust packs the zero an Arrow null holds. + private static EncodeResult encodeNullable(DType dtype, NullableData data, EncodeContext ctx) { PType ptype = ((DType.Primitive) dtype).ptype(); - long[] longs = PrimitiveArrays.toLongs(data, ptype, EncodingId.FASTLANES_BITPACKED); + long[] longs = PrimitiveArrays.toLongs(data.values(), ptype, EncodingId.FASTLANES_BITPACKED).clone(); + boolean[] validity = data.validity(); + for (int i = 0; i < longs.length; i++) { + if (!validity[i]) { + longs[i] = 0; + } + } + MinMax stats = ZoneMapStats.of(dtype.withNullable(false), data); + EncodeResult packed = encodeLongs(ptype, longs, ctx); + return new EncodeResult(packed.rootNode(), packed.encodedBuffers(), PrimitiveEncodingEncoder.minOf(stats), + PrimitiveEncodingEncoder.maxOf(stats)).withValidity(validity, ctx); + } + + private static EncodeResult encodeLongs(PType ptype, long[] longs, EncodeContext ctx) { int n = longs.length; int typeBits = ptype.bits(); long typeMask = FastLanes.lowMask(typeBits); diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java index 716ec009e..b550ce143 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java @@ -2,6 +2,8 @@ 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.simd.SimdOperationsSupport; import java.util.ArrayList; import java.util.List; @@ -36,6 +38,7 @@ public CascadingCompressor(List encodings) { private static int dataLength(Object data) { return switch (data) { case StructData(var fieldArrays) -> fieldArrays.isEmpty() ? 0 : dataLength(fieldArrays.getFirst()); + case NullableData nullable -> nullable.length(); case byte[] a -> a.length; case short[] a -> a.length; case int[] a -> a.length; @@ -80,6 +83,14 @@ case StructData(var fieldArrays) -> { .map(f -> stratifiedSample(f, sampleSize, seed)).toList(); yield new StructData(sliced); } + case NullableData(var values, var validity) -> new NullableData( + stratifiedSample(values, sampleSize, seed), (boolean[]) stratifiedSample(validity, sampleSize, seed)); + case boolean[] a -> { + boolean[] out = new boolean[sampleSize]; + forEachStride(a.length, sampleSize, seed, (srcOff, dstOff, len) -> + System.arraycopy(a, srcOff, out, dstOff, len)); + yield out; + } case byte[] a -> { byte[] out = new byte[sampleSize]; forEachStride(a.length, sampleSize, seed, (srcOff, dstOff, len) -> @@ -172,9 +183,34 @@ private static long primitiveBytes(DType dtype, int n) { return (long) n * p.ptype().byteSize(); } - /// Rust's `exact_uncompressed_size` of a primitive array: its values. + /// Rust's `exact_uncompressed_size` of a primitive array: the values, plus a packed bitmap + /// when some row is null (`validity_uncompressed_size_in_bytes`, nothing when all are valid). private static long uncompressedBytes(DType dtype, Object data) { - return primitiveBytes(dtype, dataLength(data)); + int n = dataLength(data); + long values = primitiveBytes(dtype, n); + if (data instanceof NullableData nullable && nullable.nullCount() > 0) { + return values + (n + 7) / 8; + } + return values; + } + + /// Whether `enc` may compete for `data`: it accepts the dtype, is not excluded, and, for a + /// nullable primitive, keeps the validity itself, as every Rust scheme does on its array. + private static boolean eligible(EncodingEncoder enc, DType dtype, Object data, EncodeContext ctx) { + if (!enc.accepts(dtype) || ctx.excluded().contains(enc.encodingId())) { + return false; + } + return !(data instanceof NullableData) || enc.acceptsNullable(dtype); + } + + /// Whether a nullable column of `dtype` keeps its validity in the encoding that stores its + /// values, as every Rust array does, rather than in a `vortex.masked` wrapper (#604). Integers + /// do; floats still go through `vortex.masked` until ALP and ALP-RD carry a validity. + /// + /// @param dtype the column's dtype + /// @return `true` for an integer primitive + public static boolean keepsValidity(DType dtype) { + return dtype instanceof DType.Primitive p && p.ptype().isInteger(); } /// Entry point: encode `data` using the best cascading strategy. @@ -188,10 +224,19 @@ private static long uncompressedBytes(DType dtype, Object data) { /// @param ctx encoding context supplying the arena, encoder map, and cascade parameters /// @return the [EncodeResult] produced by the winning encoding public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { + // A column without nulls has no validity to store, as Rust imports an Arrow array whose + // null count is zero (`Validity::AllValid`). Deeper down a NullableData keeps its bitmap + // even when every bit is set, as a filtered or sampled Rust array does. + if (keepsValidity(dtype) && data instanceof NullableData nullable && nullable.nullCount() == 0) { + return encodeWithCtx(dtype, nullable.values(), ctx); + } return encodeWithCtx(dtype, data, ctx); } private EncodeResult encodeWithCtx(DType dtype, Object data, EncodeContext ctx) { + if (keepsValidity(dtype) && data instanceof NullableData nullable) { + return encodeNullablePrimitive(dtype, nullable, ctx); + } if (dtype instanceof DType.Struct structDtype) { return data instanceof NullableData(StructData fields, boolean[] validity) ? encodeStruct(structDtype, fields, validity, ctx) @@ -232,7 +277,7 @@ private EncodeResult encodeWithCtx(DType dtype, Object data, EncodeContext ctx) private ArrayAndStats withStats(DType dtype, Object data, EncodeContext ctx) { StatsOptions merged = StatsOptions.defaults(); for (EncodingEncoder enc : encodings) { - if (enc.accepts(dtype) && !ctx.excluded().contains(enc.encodingId())) { + if (eligible(enc, dtype, data, ctx)) { merged = StatsOptions.merge(merged, enc.statsOptions()); } } @@ -295,6 +340,31 @@ private static long totalBytes(EncodeResult result) { return total; } + /// Rust's `choose_and_compress` for a primitive array with nulls: an all-null array is a null + /// constant and one whose valid rows hold a single value is a constant masked by the validity, + /// both decided by the compressor before any scheme (`compressor/cascade.rs`). Otherwise the + /// schemes compete with the validity in their arrays: a sampled scheme is measured with its + /// validity child, against the values plus a packed bitmap. + private EncodeResult encodeNullablePrimitive(DType dtype, NullableData data, EncodeContext ctx) { + int n = data.length(); + if (MaskedEncodingEncoder.isNullConstant(data.validity(), ctx)) { + return MaskedEncodingEncoder.nullConstant(); + } + if (n > 0 && ctx.allowedCascading() > 0 && !ctx.sample() + && !ctx.excluded().contains(EncodingId.VORTEX_CONSTANT)) { + Object valid = data.validValues(); + PType ptype = ((DType.Primitive) dtype).ptype(); + if (SimdOperationsSupport.preferred().allEqual(valid, ptype)) { + return MaskedEncodingEncoder.maskedConstant(dtype, data, valid, ctx); + } + } + ArrayAndStats input = withStats(dtype, data, ctx); + int sampleSize = sampleSize(n, ctx); + Object sample = sampleSize < n ? stratifiedSample(data, sampleSize, ctx.sampleSeed()) : data; + Choice choice = chooseBest(dtype, input, ctx, uncompressedBytes(dtype, sample), () -> sample); + return encodeChoice(choice, dtype, data, ctx); + } + /// The winner of a selection and the size it was scored at on the sample (or its estimate). /// /// @param encoder the winning encoder, or `null` when none beat the baseline @@ -327,7 +397,7 @@ private Choice chooseBest(DType dtype, ArrayAndStats input, EncodeContext ctx, l boolean bestEstimated = false; List deferred = new ArrayList<>(); for (EncodingEncoder enc : encodings) { - if (!enc.accepts(dtype) || ctx.excluded().contains(enc.encodingId())) { + if (!eligible(enc, dtype, input.data(), ctx)) { continue; } Estimate est = competes(dtype) ? enc.expectedRatio(dtype, input, ctx) : Estimate.COMPLETE; @@ -389,7 +459,7 @@ private long measureStep(EncodingEncoder enc, CascadeStep step, EncodeContext ct /// (`compressor/cascade.rs`), so it costs nothing. private long measureBestChild(DType dtype, Object data, EncodeContext ctx) { int n = dataLength(data); - long baseline = primitiveBytes(dtype, n); + long baseline = uncompressedBytes(dtype, data); if (n == 0) { return baseline; } @@ -456,12 +526,13 @@ private EncodeResult encodeStruct(DType.Struct dtype, StructData data, boolean[] // Mirrors StructEncodingEncoder's own field loop: a nullable field arrives as // NullableData(values, validity), not the dense array (String[], byte[][], ...) // encodeWithCtx's per-dtype dispatch expects. A nullable struct carries its validity - // as its own child instead, so it recurses rather than being masked. - boolean masked = fieldData instanceof NullableData - && !(fieldDtype instanceof DType.Extension) && !(fieldDtype instanceof DType.Struct); + // as its own child instead, so it recurses rather than being masked, and a nullable + // primitive keeps its validity in the encoding that stores its values. + boolean masked = fieldData instanceof NullableData && !(fieldDtype instanceof DType.Extension) + && !(fieldDtype instanceof DType.Struct) && !keepsValidity(fieldDtype); EncodeResult fieldResult = masked ? new MaskedEncodingEncoder().encode(fieldDtype, fieldData, ctx) - : encodeWithCtx(fieldDtype, fieldData, ctx); + : encode(fieldDtype, fieldData, ctx); int bufOffset = allBuffers.size(); children[fieldOffset + i] = EncodeNode.remapBufferIndices(fieldResult.rootNode(), bufOffset); allBuffers.addAll(fieldResult.encodedBuffers()); diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java index a39512ba4..46f72dbf7 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java @@ -70,6 +70,12 @@ public boolean accepts(DType dtype) { return dtype instanceof DType.Utf8; } + /// A nullable primitive keeps its validity on the codes, as Rust's dict schemes do. + @Override + public boolean acceptsNullable(DType dtype) { + return dtype instanceof DType.Primitive && accepts(dtype); + } + @Override public StatsOptions statsOptions() { return StatsOptions.defaults().withCountDistinct(true); @@ -117,7 +123,7 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { if (dtype instanceof DType.Utf8) { return encodeUtf8(data, ctx); } - DictData d = buildDictData(dtype, data); + DictData d = buildDictData(dtype, data, null); PType codePType = d.codePType(); int codeBytes = codePType.byteSize(); @@ -150,7 +156,13 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { if (dtype instanceof DType.Utf8) { return encodeUtf8Cascade(data); } - DictData d = buildDictData(dtype, data); + boolean[] validity = null; + Object values = data; + if (data instanceof NullableData nullable) { + validity = nullable.validity(); + values = nullable.values(); + } + DictData d = buildDictData(dtype, values, validity); PType codePType = d.codePType(); // Same [codes, values] wire shape as the terminal path above, with both children left open: @@ -163,10 +175,13 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { new EncodeNode[]{null, null}, new int[0]); - DType codesDtype = new DType.Primitive(codePType, false); - MinMax stats = PrimitiveEncodingEncoder.minMaxStats(((DType.Primitive) dtype).ptype(), data); + // Rust's dict schemes keep a nullable array's validity on the codes and give the values an + // all-valid one (`builtins/dict/integer.rs`), so only the codes carry a bitmap + DType codesDtype = new DType.Primitive(codePType, validity != null); + Object codes = validity == null ? d.codesArr() : new NullableData(d.codesArr(), validity); + MinMax stats = ZoneMapStats.of(dtype.withNullable(false), data); return new CascadeStep(partialRoot, List.of(), - List.of(new ChildSlot(codesDtype, d.codesArr(), 0, CODES_EXCLUDED), + List.of(new ChildSlot(codesDtype, codes, 0, CODES_EXCLUDED), new ChildSlot(dtype, d.valuesArr(), 1, Set.of(EncodingId.VORTEX_DICT))), PrimitiveEncodingEncoder.minOf(stats), PrimitiveEncodingEncoder.maxOf(stats), true); } @@ -413,7 +428,10 @@ private static EncodeResult encodeUtf8(Object data, EncodeContext ctx) { EncodedBuffer.of(codesBuf, codePType)), statsMin, statsMax); } - private static DictData buildDictData(DType dtype, Object data) { + /// A null row is not a dictionary entry: its code is that of the value in its slot when that + /// value is in the dictionary, else 0, as Rust's `DictEncoder::encode` looks it up + /// (`unwrap_or_default`). Its validity keeps the code from being read. + private static DictData buildDictData(DType dtype, Object data, boolean[] validity) { PType ptype = ((DType.Primitive) dtype).ptype(); int len = Array.getLength(data); @@ -432,6 +450,9 @@ private static DictData buildDictData(DType dtype, Object data) { long[] firstSeenRaw = new long[16]; int dictSize = 0; for (int i = 0; i < len; i++) { + if (validity != null && !validity[i]) { + continue; + } long raw = rawBits(data, ptype, i); int code = index.get(canonicalBits(raw, ptype)); if (code < 0) { @@ -446,13 +467,20 @@ private static DictData buildDictData(DType dtype, Object data) { codes[i] = code; } + if (validity != null) { + for (int i = 0; i < len; i++) { + if (!validity[i]) { + codes[i] = Math.max(0, index.get(canonicalBits(rawBits(data, ptype, i), ptype))); + } + } + } PType codePType = codePType(dictSize); Object uniqueArray = buildUniqueArray(ptype, firstSeenRaw, dictSize); MemorySegment valuesBuf = PTypeIO.copyArray(ptype, uniqueArray, dictSize); Object codesArr = PrimitiveArrays.fromIntsArray(codes, codePType, EncodingId.VORTEX_DICT); - return new DictData(valuesBuf, uniqueArray, codesArr, codePType, len, dictSize); + return new DictData(valuesBuf, uniqueArray, codesArr, codePType, len, dictSize, validity != null); } private static PType codePType(int dictSize) { @@ -531,12 +559,12 @@ private static byte[] dictMetadata(DictData d) { return new ProtoDictMetadata( d.dictSize(), io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(d.codePType().ordinal()), - null, + d.nullableCodes() ? Boolean.TRUE : null, null ).encode(); } private record DictData(MemorySegment valuesBuf, Object valuesArr, Object codesArr, PType codePType, int len, - int dictSize) { + int dictSize, boolean nullableCodes) { } } diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/EncodeResult.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/EncodeResult.java index b9b4ed5a4..1dd53d1cf 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/EncodeResult.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/EncodeResult.java @@ -1,9 +1,12 @@ package io.github.dfa1.vortex.writer.encode; import io.github.dfa1.vortex.core.proto.ProtoScalarValue; +import io.github.dfa1.vortex.core.model.DType; import io.github.dfa1.vortex.core.model.EncodingId; import java.lang.foreign.MemorySegment; +import java.util.ArrayList; +import java.util.Arrays; import java.util.List; /// Output of encoding an array to bytes for one flat segment. @@ -59,6 +62,23 @@ public EncodeResult withStats(MinMax stats) { : new EncodeResult(rootNode, encodedBuffers, stats.min(), stats.max()); } + /// This result with `validity` appended to the root as its last child, a plain `vortex.bool`: + /// where Rust's primitive and bit-packed arrays keep the validity of a nullable array, which + /// the encodings stacked above them (FoR, ZigZag, dict codes) pass down unchanged. + /// + /// @param validity per-row validity + /// @param ctx the encode context + /// @return a copy of this result with the validity child and its buffer + EncodeResult withValidity(boolean[] validity, EncodeContext ctx) { + EncodeResult bool = new BoolEncodingEncoder().encode(DType.BOOL, validity, ctx); + EncodeNode[] children = Arrays.copyOf(rootNode.children(), rootNode.children().length + 1); + children[children.length - 1] = EncodeNode.remapBufferIndices(bool.rootNode(), encodedBuffers.size()); + List buffers = new ArrayList<>(encodedBuffers); + buffers.addAll(bool.encodedBuffers()); + EncodeNode root = new EncodeNode(rootNode.encodingId(), rootNode.metadata(), children, rootNode.bufferIndices()); + return new EncodeResult(root, List.copyOf(buffers), statsMin, statsMax); + } + /// Returns `true` if both `statsMin` and `statsMax` are present. /// /// @return `true` if zone-map statistics are available for this result diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FrameOfReferenceEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FrameOfReferenceEncodingEncoder.java index 52859904d..6b9ddd2ef 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FrameOfReferenceEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/FrameOfReferenceEncodingEncoder.java @@ -72,11 +72,20 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { PrimitiveEncodingEncoder.minOf(stats), PrimitiveEncodingEncoder.maxOf(stats)); } + /// A nullable array passes its validity down to the residuals, as Rust's `FoR::encode` does. + @Override + public boolean acceptsNullable(DType dtype) { + return accepts(dtype); + } + @Override public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext encodeCtx) { if (!(dtype instanceof DType.Primitive p)) { throw new VortexException(EncodingId.FASTLANES_FOR, "expected primitive dtype, got " + dtype); } + if (data instanceof NullableData nullable) { + return encodeNullableCascade(p, nullable, encodeCtx); + } PType ptype = p.ptype(); long[] longs = PrimitiveArrays.toLongs(data, ptype, EncodingId.FASTLANES_FOR); int n = longs.length; @@ -106,6 +115,27 @@ private static CascadeStep step(DType.Primitive dtype, long ref, Object residual PrimitiveEncodingEncoder.minOf(stats), PrimitiveEncodingEncoder.maxOf(stats))); } + /// Rust's `FoR::encode` on a nullable array: the reference is the smallest valid value, and a + /// null slot's residual is zero, so it costs no bits in the bit-packing below + /// (`for_compress.rs`). The residuals keep the validity. + private static CascadeStep encodeNullableCascade(DType.Primitive dtype, NullableData data, EncodeContext ctx) { + PType ptype = dtype.ptype(); + long[] valid = PrimitiveArrays.toLongs(data.validValues(), ptype, EncodingId.FASTLANES_FOR); + long ref = computeRef(valid, valid.length, ptype); + if (ref == 0L && ptype.isUnsigned()) { + return CascadeStep.notApplicable(); + } + long[] longs = PrimitiveArrays.toLongs(data.values(), ptype, EncodingId.FASTLANES_FOR).clone(); + boolean[] validity = data.validity(); + for (int i = 0; i < longs.length; i++) { + if (!validity[i]) { + longs[i] = ref; + } + } + return step(dtype, ref, new NullableData(residualsAsNativeArray(longs, ref, ptype), validity), + ZoneMapStats.of(dtype.withNullable(false), data), ctx); + } + // Rust takes the reference from Stat::Min, an unsigned min for unsigned ptypes: a signed one // picks a u64 above 2^63 over a smaller value below it, so a column straddling 2^63 got full-width // residuals. Widened U8..U32 values are zero-extended, so the U64 order serves every unsigned width. diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/MaskedEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/MaskedEncodingEncoder.java index de5bbd696..303ef8d75 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/MaskedEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/MaskedEncodingEncoder.java @@ -6,6 +6,7 @@ import io.github.dfa1.vortex.core.error.VortexException; import io.github.dfa1.vortex.core.model.EncodingId; +import io.github.dfa1.vortex.core.model.PType; import io.github.dfa1.vortex.core.simd.SimdOperationsSupport; import java.lang.foreign.MemorySegment; @@ -180,6 +181,31 @@ static boolean isNullConstant(boolean[] validity, EncodeContext ctx) { && SimdOperationsSupport.preferred().allEqual(validity); } + /// Rust's `compress_constant` for a nullable primitive whose valid rows all hold one value: a + /// `vortex.constant` of that value masked by the validity, stored as a plain `vortex.bool`. + /// As `compress_as_constant`, the canonical array is kept when the constant is not smaller. + /// + /// @param dtype the column's nullable primitive dtype + /// @param data the column, with at least one null + /// @param valid the valid rows' values, all equal and at least one + /// @param ctx the encode context + /// @return the masked constant, or the canonical primitive with its validity + static EncodeResult maskedConstant(DType dtype, NullableData data, Object valid, EncodeContext ctx) { + PType ptype = ((DType.Primitive) dtype).ptype(); + boolean[] validity = data.validity(); + EncodeResult constant = new ConstantEncodingEncoder().encode(dtype, valid, ctx); + if (constant.buffers().getFirst().byteSize() >= (long) validity.length * ptype.byteSize()) { + return new PrimitiveEncodingEncoder().encode(dtype, data, ctx); + } + EncodeResult bool = new BoolEncodingEncoder().encode(DType.BOOL, validity, ctx); + EncodeNode validityNode = EncodeNode.remapBufferIndices(bool.rootNode(), constant.buffers().size()); + List buffers = new ArrayList<>(constant.encodedBuffers()); + buffers.addAll(bool.encodedBuffers()); + EncodeNode root = new EncodeNode(EncodingId.VORTEX_MASKED, null, + new EncodeNode[]{constant.rootNode(), validityNode}, new int[0]); + return new EncodeResult(root, buffers, constant.statsMin(), constant.statsMax()); + } + /// A `vortex.constant` whose scalar is null: every row null, and no min/max. /// /// @return the encoded constant diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/NullableData.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/NullableData.java index 59b76a0b9..f99d55b56 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/NullableData.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/NullableData.java @@ -9,17 +9,101 @@ /// alongside real data. `validity` has the same logical length: `true` /// for valid rows, `false` for nulls. /// -/// The writer recognizes this shape and emits the `vortex.masked` -/// wire layout: a non-nullable child (the storage) plus an optional Bool -/// validity child. Readers reconstruct a `MaskedArray` — named as text, not linked, because -/// the writer module does not depend on the reader. +/// A nullable primitive column keeps its validity on the encoding that stores the values, as +/// Rust's arrays do (a `vortex.primitive` or `fastlanes.bitpacked` with a `vortex.bool` child); +/// other nullable columns are written as `vortex.masked`: a non-nullable child (the storage) +/// plus a Bool validity child. Readers reconstruct a `MaskedArray` — named as text, not linked, +/// because the writer module does not depend on the reader. /// /// Invariant: `values` length (or storage row count) matches -/// `validity.length`; placeholders at null positions must be zero so the +/// `validity.length`; placeholders at null positions must be zero so /// the validity child contains no semantic nulls. /// /// @param values packed values array; concrete type depends on the column dtype /// @param validity per-row validity bitmap, parallel to `values` @SuppressWarnings("java:S6218") // internal data carrier; record components are arrays of immutable primitives or refs that flow through pipelines without ever being compared. public record NullableData(Object values, boolean[] validity) { + + /// Returns the number of rows. + /// + /// @return the validity length + public int length() { + return validity.length; + } + + /// Returns the number of null rows. + /// + /// @return how many rows are invalid + public int nullCount() { + int valid = 0; + for (boolean v : validity) { + valid += v ? 1 : 0; + } + return validity.length - valid; + } + + /// Returns the values of the valid rows, in row order: the array Rust's stats see, since + /// they skip nulls (`vortex-compressor` `stats/integer.rs`). + /// + /// @return a new primitive array of the same type as `values`, one element per valid row + public Object validValues() { + int count = validity.length - nullCount(); + return switch (values) { + case byte[] a -> { + byte[] out = new byte[count]; + for (int i = 0, j = 0; j < count; i++) { + if (validity[i]) { + out[j++] = a[i]; + } + } + yield out; + } + case short[] a -> { + short[] out = new short[count]; + for (int i = 0, j = 0; j < count; i++) { + if (validity[i]) { + out[j++] = a[i]; + } + } + yield out; + } + case int[] a -> { + int[] out = new int[count]; + for (int i = 0, j = 0; j < count; i++) { + if (validity[i]) { + out[j++] = a[i]; + } + } + yield out; + } + case long[] a -> { + long[] out = new long[count]; + for (int i = 0, j = 0; j < count; i++) { + if (validity[i]) { + out[j++] = a[i]; + } + } + yield out; + } + case float[] a -> { + float[] out = new float[count]; + for (int i = 0, j = 0; j < count; i++) { + if (validity[i]) { + out[j++] = a[i]; + } + } + yield out; + } + case double[] a -> { + double[] out = new double[count]; + for (int i = 0, j = 0; j < count; i++) { + if (validity[i]) { + out[j++] = a[i]; + } + } + yield out; + } + default -> throw new IllegalArgumentException("not a primitive array: " + values.getClass()); + }; + } } diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/PcoEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/PcoEncodingEncoder.java index d6728e76a..5b179fee5 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/PcoEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/PcoEncodingEncoder.java @@ -63,8 +63,20 @@ public boolean accepts(DType dtype) { }; } + /// A nullable array compresses only its valid values and keeps the validity as child 0, as + /// Rust's `Pco::from_primitive` does (`collect_valid`). + @Override + public boolean acceptsNullable(DType dtype) { + return accepts(dtype); + } + @Override public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { + if (data instanceof NullableData nullable) { + return Encoder.encode(dtype, nullable.validValues(), ctx) + .withStats(ZoneMapStats.of(dtype.withNullable(false), nullable)) + .withValidity(nullable.validity(), ctx); + } // Pco is a lossless numeric codec: it reorders nothing and drops nothing, so the column's // bounds are the input's. Attached here rather than inside Encoder, which assembles the // result several frames deep from the values. diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/PrimitiveEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/PrimitiveEncodingEncoder.java index be2a09a69..5dfb91455 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/PrimitiveEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/PrimitiveEncodingEncoder.java @@ -29,8 +29,20 @@ public boolean accepts(DType dtype) { return dtype instanceof DType.Primitive; } + /// A nullable primitive keeps its validity as a `vortex.bool` child, as Rust's canonical + /// `PrimitiveArray` does. + @Override + public boolean acceptsNullable(DType dtype) { + return dtype instanceof DType.Primitive; + } + @Override public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { + if (data instanceof NullableData nullable) { + return encode(dtype, nullable.values(), ctx) + .withStats(ZoneMapStats.of(dtype.withNullable(false), nullable)) + .withValidity(nullable.validity(), ctx); + } PType ptype = ((DType.Primitive) dtype).ptype(); // byte[] needs no byte swap, so it is wrapped in place rather than copied off-heap. MemorySegment seg = data instanceof byte[] bytes 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 22bed2e5c..3b82eb12b 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 @@ -44,6 +44,12 @@ public Estimate expectedRatio(DType dtype, ArrayAndStats data, EncodeContext ctx return data.stats().averageRunLength() < RUN_LENGTH_THRESHOLD ? Estimate.SKIP : Estimate.COMPLETE; } + /// A nullable array keeps its validity on the indices, as Rust's `RLEData::encode` does. + @Override + public boolean acceptsNullable(DType dtype) { + return accepts(dtype); + } + @Override public boolean accepts(DType dtype) { // every primitive, floats included (Rust's FloatRLEScheme): runs compare raw bits, so -0.0, @@ -148,10 +154,19 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { @Override public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { PType ptype = ((DType.Primitive) dtype).ptype(); - long[] longs = toLongs(data, ptype); + boolean[] validity = null; + Object raw = data; + if (data instanceof NullableData nullable) { + validity = nullable.validity(); + raw = nullable.values(); + } + long[] longs = toLongs(raw, ptype); if (longs.length == 0) { return CascadeStep.notApplicable(); } + if (validity != null) { + longs = fillForwardNulls(longs, validity); + } Runs runs = runs(longs); // Rust narrows the u16 indices and u64 offsets to the smallest unsigned type that holds // their largest value before compressing them (`rle_compress`, `PrimitiveArray::narrow`) @@ -164,21 +179,44 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { PType indicesPtype = PType.narrowestUnsigned(maxIndex); long[] offsets = Arrays.copyOf(runs.offsets(), runs.numChunks()); PType offsetsPtype = PType.narrowestUnsigned(offsets[offsets.length - 1]); - Object indicesData = PrimitiveArrays.fromIntsArray(indices, indicesPtype, EncodingId.FASTLANES_RLE); + Object narrowIndices = PrimitiveArrays.fromIntsArray(indices, indicesPtype, EncodingId.FASTLANES_RLE); + Object indicesData = narrowIndices; + if (validity != null) { + // the rows' validity, padded invalid to a whole chunk (Rust's `padded_validity`) + indicesData = new NullableData(narrowIndices, Arrays.copyOf(validity, runs.paddedLen())); + } EncodeNode partialRoot = new EncodeNode(EncodingId.FASTLANES_RLE, MemorySegment.ofArray(runs.metadata(indicesPtype, offsetsPtype)), new EncodeNode[]{null, null, null}, new int[0]); long[] values = Arrays.copyOf(runs.values(), runs.valuesCount()); List slots = List.of( new ChildSlot(dtype, PrimitiveArrays.fromBitsArray(values, ptype, EncodingId.FASTLANES_RLE), 0, VALUES_EXCLUDED), - new ChildSlot(new DType.Primitive(indicesPtype, false), indicesData, 1, POSITIONS_EXCLUDED), + new ChildSlot(new DType.Primitive(indicesPtype, validity != null), indicesData, 1, POSITIONS_EXCLUDED), new ChildSlot(new DType.Primitive(offsetsPtype, false), PrimitiveArrays.fromLongsArray(offsets, offsetsPtype, EncodingId.FASTLANES_RLE), 2, POSITIONS_EXCLUDED)); - MinMax stats = ZoneMapStats.of(dtype, data); + MinMax stats = ZoneMapStats.of(dtype.withNullable(false), data); return new CascadeStep(partialRoot, List.of(), slots, ZoneMapStats.minOf(stats), ZoneMapStats.maxOf(stats), true); } + /// Rust's `fill_forward_nulls` with `ChunkBoundary::Reset` (`encodings/fastlanes/src/lib.rs`): + /// a null takes the last valid value, or 0 when it opens a chunk, so nulls add no runs. + private static long[] fillForwardNulls(long[] longs, boolean[] validity) { + long[] filled = longs.clone(); + long lastValid = 0; + for (int i = 0; i < filled.length; i++) { + if (validity[i]) { + lastValid = filled[i]; + } else { + if ((i & (FL_CHUNK_SIZE - 1)) == 0) { + lastValid = 0; + } + filled[i] = lastValid; + } + } + return filled; + } + /// One column run-length encoded in FastLanes 1024-row chunks. /// /// @param values the run values, chunk after chunk (first `valuesCount` are used) diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/RunEndEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/RunEndEncodingEncoder.java index 55ff44712..0350b6cc4 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/RunEndEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/RunEndEncodingEncoder.java @@ -36,6 +36,13 @@ public boolean accepts(DType dtype) { return dtype instanceof DType.Primitive p && p.ptype().isInteger(); } + /// A nullable array keeps its nulls as runs, invalid in the run values' validity, as Rust's + /// `runend_encode` does. + @Override + public boolean acceptsNullable(DType dtype) { + return accepts(dtype); + } + /// Rust's `RunEndScheme` verdict: skip unless runs average at least [#RUN_END_THRESHOLD] /// values, else defer to the sample. Needs no distinct count, so it declares no /// [StatsOptions], as Rust's scheme declares none. @@ -217,29 +224,48 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { return CascadeStep.notApplicable(); } PType ptype = p.ptype(); - int n = Array.getLength(data); + boolean[] validity = null; + Object raw = data; + if (data instanceof NullableData nullable) { + validity = nullable.validity(); + raw = nullable.values(); + } + int n = Array.getLength(raw); if (n == 0) { return CascadeStep.notApplicable(); } - long[] widened = PrimitiveArrays.toLongs(data, ptype, EncodingId.VORTEX_RUNEND); + long[] widened = PrimitiveArrays.toLongs(raw, ptype, EncodingId.VORTEX_RUNEND); + if (validity != null) { + // Rust's runend_encode_nullable_primitive compares Option values: a null ends a valid + // run and a run of nulls is one run, stored as an invalid 0. Zeroing the null slots + // and comparing validity too gives exactly those runs. + widened = widened.clone(); + for (int i = 0; i < n; i++) { + if (!validity[i]) { + widened[i] = 0; + } + } + } int numRuns = 1; for (int i = 1; i < n; i++) { - if (widened[i] != widened[i - 1]) { + if (newRun(widened, validity, i)) { numRuns++; } } int[] ends = new int[numRuns]; long[] values = new long[numRuns]; + boolean[] runValidity = new boolean[numRuns]; + boolean anyNullRun = false; int k = 0; - for (int i = 1; i < n; i++) { - if (widened[i] != widened[i - 1]) { + for (int i = 1; i <= n; i++) { + if (i == n || newRun(widened, validity, i)) { ends[k] = i; values[k] = widened[i - 1]; + runValidity[k] = validity == null || validity[i - 1]; + anyNullRun |= !runValidity[k]; k++; } } - ends[k] = n; - values[k] = widened[n - 1]; // Rust narrows the u64 ends to the smallest unsigned type that holds the last one (`n`) PType endsPtype = PType.narrowestUnsigned(n); @@ -249,12 +275,19 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { new EncodeNode[]{null, null}, new int[0]); ChildSlot endsSlot = new ChildSlot(new DType.Primitive(endsPtype, false), PrimitiveArrays.fromIntsArray(ends, endsPtype, EncodingId.VORTEX_RUNEND), 0, ENDS_EXCLUDED); - ChildSlot valuesSlot = new ChildSlot(dtype, PrimitiveArrays.fromLongsArray(values, ptype, EncodingId.VORTEX_RUNEND), 1, VALUES_EXCLUDED); - MinMax stats = ZoneMapStats.of(dtype, data); + // The run values' validity is all-valid unless a run is null, as Rust builds it from bits + Object runValues = PrimitiveArrays.fromLongsArray(values, ptype, EncodingId.VORTEX_RUNEND); + Object valuesData = anyNullRun ? new NullableData(runValues, runValidity) : runValues; + ChildSlot valuesSlot = new ChildSlot(dtype, valuesData, 1, VALUES_EXCLUDED); + MinMax stats = ZoneMapStats.of(dtype.withNullable(false), data); return new CascadeStep(partialRoot, List.of(), List.of(endsSlot, valuesSlot), ZoneMapStats.minOf(stats), ZoneMapStats.maxOf(stats), true); } + private static boolean newRun(long[] values, boolean[] validity, int i) { + return values[i] != values[i - 1] || validity != null && validity[i] != validity[i - 1]; + } + private static ProtoScalarValue statsBytes(PType ptype, long value) { if (ptype.isUnsigned()) { return ProtoScalarValue.ofUint64Value(value); diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/SparseEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/SparseEncodingEncoder.java index ecc6cddf4..92a6f2d04 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/SparseEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/SparseEncodingEncoder.java @@ -6,6 +6,7 @@ import io.github.dfa1.vortex.core.error.VortexException; import io.github.dfa1.vortex.core.model.EncodingId; import io.github.dfa1.vortex.core.io.PTypeIO; +import io.github.dfa1.vortex.core.proto.ProtoNullValue; import io.github.dfa1.vortex.core.proto.ProtoPatchesMetadata; import io.github.dfa1.vortex.core.proto.ProtoScalarValue; import io.github.dfa1.vortex.core.proto.ProtoSparseMetadata; @@ -13,6 +14,7 @@ import java.lang.foreign.MemorySegment; import java.util.Set; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; /// Write-only encoder for `vortex.sparse`. @@ -48,6 +50,14 @@ public boolean accepts(DType dtype) { return dtype instanceof DType.Primitive; } + /// A nullable integer array keeps its nulls in the fill or the patch values, as Rust's + /// `Sparse::encode` does. Nullable floats are not taken yet: they still go through + /// `vortex.masked`. + @Override + public boolean acceptsNullable(DType dtype) { + return dtype instanceof DType.Primitive p && p.ptype().isInteger(); + } + @Override public StatsOptions statsOptions() { return StatsOptions.defaults().withTrackMostFrequent(true); @@ -58,8 +68,20 @@ public Estimate expectedRatio(DType dtype, ArrayAndStats data, EncodeContext ctx if (!(dtype instanceof DType.Primitive p)) { return Estimate.COMPLETE; } - // Rust's float sparse scheme only takes null-dominated arrays; nulls never reach a - // primitive here (validity is split off by vortex.masked), so floats never qualify. + // Rust's SparseScheme takes a mostly-null array whatever its values: only the valid rows + // are stored. An all-null array is a null constant, decided before any scheme. + int len = data.length(); + int nulls = data.nullCount(); + if (nulls > 0) { + if (nulls == len) { + return Estimate.SKIP; + } + if ((double) nulls / len > 0.9) { + return Estimate.ratio((double) len / (len - nulls)); + } + } + // Rust's float sparse scheme only takes null-dominated arrays; nullable floats are still + // split off by vortex.masked before they reach a primitive here, so floats never qualify. if (p.ptype().isFloating()) { return Estimate.SKIP; } @@ -83,12 +105,15 @@ public Estimate expectedRatio(DType dtype, ArrayAndStats data, EncodeContext ctx /// Airport_fee / congestion_surcharge. @Override public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { - // `data` is the whole logical column -- the fill value and patch positions are derived - // from it below -- so the column's bounds are just its bounds, patches included. - MinMax stats = ZoneMapStats.of(dtype, data); if (!(dtype instanceof DType.Primitive p)) { return CascadeStep.notApplicable(); } + if (data instanceof NullableData nullable) { + return encodeNullableCascade(p, nullable, ctx); + } + // `data` is the whole logical column -- the fill value and patch positions are derived + // from it below -- so the column's bounds are just its bounds, patches included. + MinMax stats = ZoneMapStats.of(dtype, data); PType ptype = p.ptype(); int n = arrayLength(data, ptype); if (n == 0) { @@ -119,16 +144,73 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { return CascadeStep.notApplicable(); } - // Owned buffers: fill scalar. idx and val moved to ChildSlot so cascade can bitpack. - ProtoScalarValue fillScalar = scalar(ptype, fill); + Object valArr = PrimitiveArrays.fromBitsArray( + patchBits.stream().mapToLong(Long::longValue).toArray(), ptype, EncodingId.VORTEX_SPARSE); + return cascadeStep(dtype, scalar(ptype, fill), patchIdx.stream().mapToInt(Integer::intValue).toArray(), + idxPtype, valArr, stats, ctx); + } + + /// Rust's `Sparse::encode` on a nullable array (`encodings/sparse`). When more than 90% of + /// the rows are null the fill is null and the valid rows are the patches, their values keeping + /// an all-valid bitmap as Rust's filtered array does. Otherwise the most frequent valid value is + /// the fill and every other row, null or not, is a patch, the values keeping the validity. + private static CascadeStep encodeNullableCascade(DType.Primitive dtype, NullableData data, EncodeContext ctx) { + PType ptype = dtype.ptype(); + boolean[] validity = data.validity(); + int n = validity.length; + int nulls = data.nullCount(); + MinMax stats = ZoneMapStats.of(dtype.withNullable(false), data); + PType idxPtype = PType.narrowestUnsigned(n); + if (nulls > 0.9 * n) { + int[] idx = new int[n - nulls]; + for (int i = 0, p = 0; i < n; i++) { + if (validity[i]) { + idx[p++] = i; + } + } + boolean[] allValid = new boolean[idx.length]; + Arrays.fill(allValid, true); + ProtoScalarValue nullFill = ProtoScalarValue.ofNullValue(ProtoNullValue.NULL_VALUE); + return cascadeStep(dtype, nullFill, idx, idxPtype, new NullableData(data.validValues(), allValid), + stats, ctx); + } + Object values = data.values(); + long fill = fillBits(ptype, data.validValues()); + int count = 0; + for (int i = 0; i < n; i++) { + if (!validity[i] || readBits(values, ptype, i) != fill) { + count++; + } + } + if (count == 0) { + return CascadeStep.notApplicable(); + } + int[] idx = new int[count]; + long[] bits = new long[count]; + boolean[] patchValidity = new boolean[count]; + for (int i = 0, p = 0; i < n; i++) { + if (!validity[i] || readBits(values, ptype, i) != fill) { + idx[p] = i; + bits[p] = validity[i] ? readBits(values, ptype, i) : 0; + patchValidity[p] = validity[i]; + p++; + } + } + Object patchValues = PrimitiveArrays.fromBitsArray(bits, ptype, EncodingId.VORTEX_SPARSE); + return cascadeStep(dtype, scalar(ptype, fill), idx, idxPtype, new NullableData(patchValues, patchValidity), + stats, ctx); + } + + /// The sparse node with its fill scalar as the one owned buffer and the patch indices and + /// values as open children, so the cascade compresses both. + private static CascadeStep cascadeStep(DType dtype, ProtoScalarValue fillScalar, int[] patchIdx, PType idxPtype, + Object patchValues, MinMax stats, EncodeContext ctx) { + int numPatches = patchIdx.length; byte[] fillBytes = fillScalar.encode(); MemorySegment fillBuf = ctx.arena().allocate(fillBytes.length); MemorySegment.copy(MemorySegment.ofArray(fillBytes), 0, fillBuf, 0, fillBytes.length); + Object idxArr = PrimitiveArrays.fromIntsArray(patchIdx, idxPtype, EncodingId.VORTEX_SPARSE); - Object idxArr = PrimitiveArrays.fromIntsArray( - patchIdx.stream().mapToInt(Integer::intValue).toArray(), idxPtype, EncodingId.VORTEX_SPARSE); - Object valArr = PrimitiveArrays.fromBitsArray( - patchBits.stream().mapToLong(Long::longValue).toArray(), ptype, EncodingId.VORTEX_SPARSE); ProtoPatchesMetadata patchesMeta = new ProtoPatchesMetadata( numPatches, @@ -144,7 +226,7 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { new EncodeNode[]{null, null}, new int[]{0}); DType idxDtype = new DType.Primitive(idxPtype, false); ChildSlot idxSlot = new ChildSlot(idxDtype, idxArr, 0, SELF); - ChildSlot valSlot = new ChildSlot(dtype, valArr, 1, SELF); + ChildSlot valSlot = new ChildSlot(dtype, patchValues, 1, SELF); return new CascadeStep(partialRoot, List.of(EncodedBuffer.bytes(fillBuf)), List.of(idxSlot, valSlot), ZoneMapStats.minOf(stats), ZoneMapStats.maxOf(stats), true); } @@ -177,7 +259,7 @@ static EncodeResult encodeBool(boolean[] validity, EncodeContext ctx) { } } boolean[] patchVals = new boolean[numPatches]; - java.util.Arrays.fill(patchVals, !fillValue); + Arrays.fill(patchVals, !fillValue); PType idxPtype = PType.narrowestUnsigned(n); Object idxArr = PrimitiveArrays.fromIntsArray(patchIdx, idxPtype, EncodingId.VORTEX_SPARSE); diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZigZagEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZigZagEncodingEncoder.java index e7261fc9d..6fca6b7dd 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZigZagEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZigZagEncodingEncoder.java @@ -155,29 +155,46 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { /// On size it never beats frame-of-reference, which competes for the same columns: FoR needs /// `bits(max - min)`, zigzag `bits(2 * max|v|)`, never fewer (they tie when the data is /// symmetric around zero). It is kept for parity with the reference compressor. + /// A nullable array passes its validity down to the encoded values, as Rust's `zigzag_encode` + /// does; a null slot's zero encodes as zero. + @Override + public boolean acceptsNullable(DType dtype) { + return accepts(dtype); + } + @Override public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { PType signed = ((DType.Primitive) dtype).ptype(); - long[] values = PrimitiveArrays.toLongs(data, signed, EncodingId.VORTEX_ZIGZAG); + boolean[] validity = null; + Object raw = data; + Object valid = data; + if (data instanceof NullableData nullable) { + validity = nullable.validity(); + raw = nullable.values(); + valid = nullable.validValues(); + } + long[] values = PrimitiveArrays.toLongs(raw, signed, EncodingId.VORTEX_ZIGZAG); int n = values.length; // Rust skips on stats before sampling; ours carry no minimum, so find it with a read-only // scan before allocating anything — the common all-non-negative column stops here - if (n == 0 || SimdOperationsSupport.preferred().minMax(values, PType.I64)[0] >= 0L) { + long[] validLongs = PrimitiveArrays.toLongs(valid, signed, EncodingId.VORTEX_ZIGZAG); + if (validLongs.length == 0 || SimdOperationsSupport.preferred().minMax(validLongs, PType.I64)[0] >= 0L) { return CascadeStep.notApplicable(); } long[] encoded = new long[n]; for (int i = 0; i < n; i++) { - long v = values[i]; + long v = validity == null || validity[i] ? values[i] : 0; // computed at 64 bits: for a value sign-extended from a narrower width, the low bits // of the result equal that width's zigzag, and fromLongsArray truncates to them encoded[i] = (v << 1) ^ (v >> 63); } PType unsigned = toUnsigned(signed); EncodeNode partialRoot = new EncodeNode(EncodingId.VORTEX_ZIGZAG, null, new EncodeNode[]{null}, new int[0]); - ChildSlot slot = new ChildSlot(new DType.Primitive(unsigned, false), - PrimitiveArrays.fromLongsArray(encoded, unsigned, EncodingId.VORTEX_ZIGZAG), 0, ENCODED_EXCLUDED); + Object encodedValues = PrimitiveArrays.fromLongsArray(encoded, unsigned, EncodingId.VORTEX_ZIGZAG); + ChildSlot slot = new ChildSlot(new DType.Primitive(unsigned, validity != null), + validity == null ? encodedValues : new NullableData(encodedValues, validity), 0, ENCODED_EXCLUDED); // zigzag is not order-preserving, so the bounds come from the original signed values - MinMax stats = ZoneMapStats.of(dtype, data); + MinMax stats = ZoneMapStats.of(dtype.withNullable(false), data); return new CascadeStep(partialRoot, List.of(), List.of(slot), ZoneMapStats.minOf(stats), ZoneMapStats.maxOf(stats), true); } diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/ArrayAndStatsTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/ArrayAndStatsTest.java index c53bca72e..3a7fe20fc 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/ArrayAndStatsTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/ArrayAndStatsTest.java @@ -27,6 +27,27 @@ void stats_computedOnce_andSharedByEveryCaller() { assertThat(result.topFrequency()).isEqualTo(2); } + @Test + void stats_nullable_skipNulls() { + // Given: the nulls hold 0, below every valid value, and sit between equal values, as + // Rust's stats see them only as gaps: a null continues the run before it + long[] values = {5, 0, 5, 9, 0, 9, 9}; + boolean[] validity = {true, false, true, true, false, true, true}; + ArrayAndStats sut = new ArrayAndStats(DType.I64, new NullableData(values, validity), ALL_STATS); + + // When + ArrayStats result = sut.stats(); + + // Then: 5 valid values in 2 runs, 0 never counted + assertThat(result.valueCount()).isEqualTo(5); + assertThat(result.distinctCount()).isEqualTo(2); + assertThat(result.min()).isEqualTo(5L); + assertThat(result.max()).isEqualTo(9L); + assertThat(result.averageRunLength()).isEqualTo(2); + assertThat(sut.length()).isEqualTo(7); + assertThat(sut.nullCount()).isEqualTo(2); + } + @Test void stats_nonPrimitive_isEmpty() { // Given: strings go through the same competition, but no stats consumer reads them diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/CascadingCompressorTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/CascadingCompressorTest.java index 895a883e7..c7aadeed9 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/CascadingCompressorTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/CascadingCompressorTest.java @@ -413,6 +413,147 @@ static Stream floatArrays() { } } + /// A nullable integer keeps its validity in the encoding that stores the values, as Rust's + /// arrays do (#604), so each case checks the tree Rust writes and that the nulls decode. + @Nested + class NullableInteger { + + private static final DType I32 = new DType.Primitive(PType.I32, true); + + /// The integer schemes of Rust's default compressor (no Pco or Zstd). + private final CascadingCompressor sut = new CascadingCompressor(List.of( + new ConstantEncodingEncoder(), new FrameOfReferenceEncodingEncoder(), new ZigZagEncodingEncoder(), + new RunEndEncodingEncoder(), new RleEncodingEncoder(), new SparseEncodingEncoder(), + new SequenceEncodingEncoder(), new BitpackedEncodingEncoder(), new PrimitiveEncodingEncoder())); + + @Test + void smallValues_bitPackedWithTheBitmapAsLastChild() { + // Given — 4 bits a row, a null every 7th: Rust bit-packs and keeps the bitmap on it + NullableData data = nullable(4096, i -> i % 16, i -> i % 7 == 0); + + // When + EncodeResult result = sut.encode(I32, data, EncodeTestHelper.cascadingCtx()); + + // Then + EncodeNode root = result.rootNode(); + assertThat(root.encodingId()).isEqualTo(EncodingId.FASTLANES_BITPACKED); + assertThat(root.children()[root.children().length - 1].encodingId()).isEqualTo(EncodingId.VORTEX_BOOL); + assertDecodes(result, data); + } + + @Test + void noNull_writtenWithoutValidity() { + // Given — nullable, but no row is null: Rust imports it as all-valid and stores no bitmap + NullableData data = nullable(4096, i -> i % 16, _ -> false); + + // When + EncodeResult result = sut.encode(I32, data, EncodeTestHelper.cascadingCtx()); + + // Then + assertThat(result.rootNode().encodingId()).isEqualTo(EncodingId.FASTLANES_BITPACKED); + assertThat(result.rootNode().children()).isEmpty(); + } + + @Test + void validRowsOneValue_constantMaskedByAPlainBitmap() { + // Given — the compressor's own constant check, before any scheme (Rust compress_constant) + NullableData data = nullable(4096, _ -> 42, i -> i % 7 == 0); + + // When + EncodeResult result = sut.encode(I32, data, EncodeTestHelper.cascadingCtx()); + + // Then + EncodeNode root = result.rootNode(); + assertThat(root.encodingId()).isEqualTo(EncodingId.VORTEX_MASKED); + assertThat(root.children()).extracting(EncodeNode::encodingId) + .containsExactly(EncodingId.VORTEX_CONSTANT, EncodingId.VORTEX_BOOL); + assertDecodes(result, data); + } + + @Test + void everyRowNull_nullConstant() { + // Given + NullableData data = nullable(4096, i -> i, _ -> true); + + // When + EncodeResult result = sut.encode(I32, data, EncodeTestHelper.cascadingCtx()); + + // Then + assertThat(result.rootNode().encodingId()).isEqualTo(EncodingId.VORTEX_CONSTANT); + assertThat(result.rootNode().children()).isEmpty(); + } + + @Test + void largeOffset_forPassesTheBitmapToTheBitPackedResiduals() { + // Given — 10-bit residuals over a 1e9 reference + NullableData data = nullable(4096, i -> 1_000_000_000 + i % 1000, i -> i % 7 == 0); + + // When + EncodeResult result = sut.encode(I32, data, EncodeTestHelper.cascadingCtx()); + + // Then + EncodeNode root = result.rootNode(); + assertThat(root.encodingId()).isEqualTo(EncodingId.FASTLANES_FOR); + EncodeNode packed = root.children()[0]; + assertThat(packed.encodingId()).isEqualTo(EncodingId.FASTLANES_BITPACKED); + assertThat(packed.children()[packed.children().length - 1].encodingId()) + .isEqualTo(EncodingId.VORTEX_BOOL); + assertDecodes(result, data); + } + + @Test + void nullRuns_runEndValuesKeepTheBitmap() { + // Given — runs of 200, every 4th run null: a null run is one more run, invalid + NullableData data = nullable(4096, i -> (i / 200) * 1_000_003, i -> (i / 200) % 4 == 0); + + // When + EncodeResult result = sut.encode(I32, data, EncodeTestHelper.cascadingCtx()); + + // Then + assertThat(result.rootNode().encodingId()).isEqualTo(EncodingId.VORTEX_RUNEND); + assertDecodes(result, data); + } + + @Test + void mostlyNull_sparseWithANullFill() { + // Given — 95% null: Rust's sparse stores only the valid rows + NullableData data = nullable(4096, i -> i * 7, i -> i % 20 != 0); + + // When + EncodeResult result = sut.encode(I32, data, EncodeTestHelper.cascadingCtx()); + + // Then + assertThat(result.rootNode().encodingId()).isEqualTo(EncodingId.VORTEX_SPARSE); + assertDecodes(result, data); + } + + private static NullableData nullable(int n, java.util.function.IntUnaryOperator value, + java.util.function.IntPredicate isNull) { + int[] values = new int[n]; + boolean[] validity = new boolean[n]; + for (int i = 0; i < n; i++) { + validity[i] = !isNull.test(i); + values[i] = validity[i] ? value.applyAsInt(i) : 0; + } + return new NullableData(values, validity); + } + + private static void assertDecodes(EncodeResult result, NullableData expected) { + ReadRegistry readRegistry = ReadRegistry.loadAll(); + int n = expected.length(); + DecodeContext decodeCtx = DecodeTestHelper.toDecodeContext(result, n, I32, readRegistry); + MaskedArray decoded = (MaskedArray) readRegistry.decode(decodeCtx); + IntArray values = (IntArray) decoded.inner(); + int[] expectedValues = (int[]) expected.values(); + for (int i = 0; i < n; i++) { + assertThat(decoded.isValid(i)).as("validity of row %d", i).isEqualTo(expected.validity()[i]); + if (expected.validity()[i]) { + assertThat(values.getInt(i)).as("row %d", i).isEqualTo(expectedValues[i]); + } + } + } + } + /// A winner is kept only when it comes out smaller than the array, as Rust's /// `choose_and_compress` accepts it: the sample alone can favor a scheme the full array does /// not. From 5a5ae30b70179bcb2dfa3af77dee95979eb8e146 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sun, 25 Oct 2026 19:55:42 +0100 Subject: [PATCH 07/10] fix(writer): ALP cascades the encoded ints of an F32 column Rust's ALPScheme compresses the encoded integers through compress_child for either float width (vortex-btrblocks schemes/float/alp.rs). Java's encodeCascade only opened that child for F64 and fell back to the terminal encode for F32, which writes the i32s as a raw vortex.primitive, so an F32 column ALP won never got them bit-packed. vortex-jni writes an F32 column of halves as alp(bitpacked), and one with one decimal as alp(bitpacked) with patches; Java wrote alp(primitive) for both. The F32 path now mirrors the F64 one: computeF32 shared by both entry points, the patches owned by the node and the encoded ints left open (found comparing nullable floats against vortex-jni, #604). Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01SwhK4W5cgj7y7s4ZPCif5x --- CHANGELOG.md | 1 + .../writer/encode/AlpEncodingEncoder.java | 56 ++++++++++++++++++- .../writer/encode/AlpEncodingEncoderTest.java | 15 +++++ 3 files changed, 70 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 08ec66c64..85009e6ce 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **Breaking:** CSV import stores an empty cell in a number or boolean column as null instead of `0` or `false`, so inferred number and boolean columns are now nullable (read through `MaskedArray`), and an empty cell or a missing column fails the import where the schema declares the column non-nullable. ([#592](https://github.com/dfa1/vortex-java/issues/592)) ### Fixed +- ALP on an `F32` column compresses its encoded integers through the cascade, as Rust does for both float widths, instead of storing them raw. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - A nullable integer column keeps its validity in the encoding that stores its values instead of a `vortex.masked` wrapper, as Rust writes it, so the cascade and the dict layout pick what Rust picks. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - Zone stats tables are compressed with integer dictionaries among the candidates, as Rust's stats compressor has them. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - Run-end ends and RLE indices and offsets are stored in the narrowest unsigned type that holds them, as Rust narrows them, instead of fixed `u32`, `u16` and `u64`. ([#604](https://github.com/dfa1/vortex-java/issues/604)) diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoder.java index e14847403..edd06a4ea 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoder.java @@ -82,7 +82,7 @@ public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { if (ptype == PType.F64) { return encodeCascadeF64((double[]) data, ctx); } - return CascadeStep.terminal(encode(dtype, data, ctx)); + return encodeCascadeF32((float[]) data, ctx); } // Rust's `encode_single_unchecked`: `(value * 10^e * 10^-f).fast_round() as i64`, the cast @@ -367,7 +367,7 @@ private static long estimateEncodedSizeF32(float[] values, int e, int f, long li return ((long) values.length * bits + 7) / 8 + (long) patches * PATCH_BYTES_F32; } - private static EncodeResult encodeF32(float[] values, EncodeContext ctx) { + private static AlpF32Data computeF32(float[] values) { int n = values.length; int[] exps = findExponentsF32(values); int expE = exps[0]; @@ -407,6 +407,52 @@ private static EncodeResult encodeF32(float[] values, EncodeContext ctx) { ProtoScalarValue statsMin = n > 0 ? scalarF32(min) : null; ProtoScalarValue statsMax = n > 0 ? scalarF32(max) : null; + return new AlpF32Data(expE, expF, encodedArr, patchIndices, patchValues, statsMin, statsMax); + } + + /// The F32 counterpart of [#encodeCascadeF64(double[], EncodeContext)]: Rust's ALP scheme + /// compresses the encoded `i32`s through the cascade for either width. + private static CascadeStep encodeCascadeF32(float[] values, EncodeContext ctx) { + AlpF32Data d = computeF32(values); + if (d.patchIndices().isEmpty()) { + byte[] metaBytes = new ProtoALPMetadata(d.expE(), d.expF(), null).encode(); + EncodeNode partialRoot = new EncodeNode(EncodingId.VORTEX_ALP, + MemorySegment.ofArray(metaBytes), new EncodeNode[1], new int[0]); + ChildSlot slot = new ChildSlot(DType.I32, d.encodedArr(), 0, Set.of(EncodingId.VORTEX_ALP)); + return new CascadeStep(partialRoot, List.of(), List.of(slot), d.statsMin(), d.statsMax(), true); + } + + int numPatches = d.patchIndices().size(); + PatchBuffers patchBuffers = PatchBuffers.of(d.patchIndices(), values.length, ctx); + MemorySegment valBuf = ctx.arena().allocate((long) numPatches * 4, 4); + for (int i = 0; i < numPatches; i++) { + valBuf.setAtIndex(VortexFormat.LE_FLOAT, i, d.patchValues().get(i)); + } + + ProtoPatchesMetadata patches = patchBuffers.meta(numPatches); + byte[] metaBytes = new ProtoALPMetadata(d.expE(), d.expF(), patches).encode(); + + EncodeNode idxNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 0); + EncodeNode valNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 1); + EncodeNode offsetsNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 2); + EncodeNode partialRoot = new EncodeNode(EncodingId.VORTEX_ALP, + MemorySegment.ofArray(metaBytes), new EncodeNode[]{null, idxNode, valNode, offsetsNode}, new int[0]); + ChildSlot slot = new ChildSlot(DType.I32, d.encodedArr(), 0, Set.of(EncodingId.VORTEX_ALP)); + return new CascadeStep(partialRoot, List.of(EncodedBuffer.of(patchBuffers.indices(), patchBuffers.indexType()), + EncodedBuffer.of(valBuf, PType.F32), EncodedBuffer.of(patchBuffers.offsets(), patchBuffers.offsetType())), + List.of(slot), d.statsMin(), d.statsMax(), true); + } + + private static EncodeResult encodeF32(float[] values, EncodeContext ctx) { + AlpF32Data d = computeF32(values); + int n = values.length; + int expE = d.expE(); + int expF = d.expF(); + int[] encodedArr = d.encodedArr(); + List patchIndices = d.patchIndices(); + List patchValues = d.patchValues(); + ProtoScalarValue statsMin = d.statsMin(); + ProtoScalarValue statsMax = d.statsMax(); MemorySegment encodedBuf = ctx.arena().allocate((long) n * 4, 4); for (int i = 0; i < n; i++) { @@ -452,6 +498,12 @@ private static ProtoScalarValue scalarF32(float v) { return ProtoScalarValue.ofF32Value(v); } + @SuppressWarnings("java:S6218") // internal data carrier; record components are arrays of immutable primitives or refs that flow through pipelines without ever being compared. + private record AlpF32Data(int expE, int expF, int[] encodedArr, + List patchIndices, List patchValues, + ProtoScalarValue statsMin, ProtoScalarValue statsMax) { + } + @SuppressWarnings("java:S6218") // internal data carrier; record components are arrays of immutable primitives or refs that flow through pipelines without ever being compared. private record AlpF64Data(int expE, int expF, long[] encodedArr, List patchIndices, List patchValues, diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoderTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoderTest.java index e341153b0..69b784d99 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoderTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoderTest.java @@ -363,6 +363,21 @@ void encodeCascade_f64_empty_doesNotThrow() { assertThat(result.statsMin()).isNull(); assertThat(result.statsMax()).isNull(); } + + @Test + void encodeCascade_f32_opensTheEncodedIntsForTheCascade() { + // Given — Rust's ALP scheme compresses the encoded ints of either width; f32 wrote + // them as a raw i32 buffer, which the cascade could never bit-pack + float[] values = {1.5f, 2.5f, 3.5f, 1.5f}; + + // When + CascadeStep result = ENCODER.encodeCascade(DTypes.F32, values, EncodeTestHelper.cascadingCtx()); + + // Then + assertThat(result.openChildren()).hasSize(1); + assertThat(result.openChildren().getFirst().childDtype()).isEqualTo(DTypes.I32); + assertThat((int[]) result.openChildren().getFirst().childData()).containsExactly(15, 25, 35, 15); + } } /// Property-based round-trip. ALP is a **lossless** codec: every value either fits the From 55f986dca5cc9697ca7131303bd9a53f969d8b92 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sun, 25 Oct 2026 20:08:38 +0100 Subject: [PATCH 08/10] fix(writer): ALP-RD packs its parts at the dictionary's widths Rust's encode_generic (encodings/alp/src/alp_rd/mod.rs) packs the dictionary-coded left parts at the encoder's left_bit_width (the width of the largest code, at least 1 bit: bit_width(dict_len - 1) in the alp crate) and the right parts at right_bit_width, both with bitpack_encode_unchecked. Java ran them through the bit-packing encoder, whose histogram finds narrower widths when the values allow: on quarters the right parts are all zero and packed in 0 bits. That made ALP-RD measure smaller in Java than in Rust, so it won columns Rust gives to a dictionary: a float column with five distinct values, or taxi's Airport_fee shape, which vortex-jni writes with a dict layout (found comparing nullable floats, #604). The parts now go through BitpackedEncodingEncoder.encodeAtWidth, Rust's unchecked pack. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01SwhK4W5cgj7y7s4ZPCif5x --- CHANGELOG.md | 1 + .../writer/encode/AlpRdEncodingEncoder.java | 11 +++++--- .../encode/BitpackedEncodingEncoder.java | 26 +++++++++++++++++++ .../encode/AlpRdEncodingEncoderTest.java | 22 ++++++++++++++++ 4 files changed, 57 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 85009e6ce..f71692323 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **Breaking:** CSV import stores an empty cell in a number or boolean column as null instead of `0` or `false`, so inferred number and boolean columns are now nullable (read through `MaskedArray`), and an empty cell or a missing column fails the import where the schema declares the column non-nullable. ([#592](https://github.com/dfa1/vortex-java/issues/592)) ### Fixed +- ALP-RD packs its left and right parts at the dictionary's widths, as Rust does, instead of narrower widths that made it win where Rust picks a dictionary. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - ALP on an `F32` column compresses its encoded integers through the cascade, as Rust does for both float widths, instead of storing them raw. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - A nullable integer column keeps its validity in the encoding that stores its values instead of a `vortex.masked` wrapper, as Rust writes it, so the cascade and the dict layout pick what Rust picks. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - Zone stats tables are compressed with integer dictionaries among the candidates, as Rust's stats compressor has them. ([#604](https://github.com/dfa1/vortex-java/issues/604)) diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoder.java index f48514710..6f2b23b3f 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoder.java @@ -308,9 +308,14 @@ private static EncodeResult buildEncodeResult( short[] leftCodes, Object rightPartsData, DType rightDtype, Exceptions exceptions, ProtoScalarValue statsMin, ProtoScalarValue statsMax, EncodeContext ctx) { - EncodingEncoder bp = ctx.lookupEncoder(EncodingId.FASTLANES_BITPACKED); - EncodeResult leftResult = bp.encode(DType.U16, leftCodes, ctx); - EncodeResult rightResult = bp.encode(rightDtype, rightPartsData, ctx); + // Rust packs both parts at the dictionary's widths (`bitpack_encode_unchecked` in + // `encode_generic`), never at narrower ones a histogram would find: the left codes at the + // width of the largest code, at least 1 bit (`alp` crate `bit_width`), the right parts at + // the right bit width. A narrower pack measured ALP-RD smaller than Rust measures it. + int leftBitWidth = Math.max(1, Integer.SIZE - Integer.numberOfLeadingZeros(dict.length - 1)); + EncodeResult leftResult = BitpackedEncodingEncoder.encodeAtWidth(DType.U16, leftCodes, leftBitWidth, ctx); + EncodeResult rightResult = BitpackedEncodingEncoder.encodeAtWidth(rightDtype, rightPartsData, rightBitWidth, + ctx); List allBuffers = new ArrayList<>(leftResult.encodedBuffers()); int leftBufCount = allBuffers.size(); diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java index 9834a21f7..3d1040cc0 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/BitpackedEncodingEncoder.java @@ -94,6 +94,32 @@ static EncodeResult encodeLeaf(DType dtype, Object data, EncodeContext ctx) { return new BitpackedEncodingEncoder().encode(dtype, data, ctx); } + /// Rust's `bitpack_encode_unchecked`: every value packed at `bitWidth`, without patches or + /// stats, and a nullable array's validity as the last child. The caller guarantees the values + /// fit, as ALP-RD's dictionary codes and right parts do by construction. + /// + /// @param dtype the unsigned integer dtype of the values + /// @param data the values, dense or a [NullableData] + /// @param bitWidth the width every value is packed at + /// @param ctx the encode context + /// @return the bit-packed array + static EncodeResult encodeAtWidth(DType dtype, Object data, int bitWidth, EncodeContext ctx) { + PType ptype = ((DType.Primitive) dtype).ptype(); + boolean[] validity = null; + Object values = data; + if (data instanceof NullableData nullable) { + validity = nullable.validity(); + values = nullable.values(); + } + long[] longs = PrimitiveArrays.toLongs(values, ptype, EncodingId.FASTLANES_BITPACKED); + MemorySegment packed = packFastLanes(longs, longs.length, bitWidth, ptype.bits(), ctx.arena()); + byte[] metaBytes = new ProtoBitPackedMetadata(bitWidth, 0, null).encode(); + EncodeNode root = new EncodeNode(EncodingId.FASTLANES_BITPACKED, MemorySegment.ofArray(metaBytes), + new EncodeNode[0], new int[]{0}); + EncodeResult result = new EncodeResult(root, List.of(EncodedBuffer.of(packed, ptype)), null, null); + return validity == null ? result : result.withValidity(validity, ctx); + } + /// Rust's `bitpack_encode` on a nullable array: the bit-width histogram counts a null at width /// 0 and only valid values become patches (`bit_width_histogram`, `gather_patches`). A null slot /// is zeroed first, which gives both, and packs as Rust packs the zero an Arrow null holds. diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoderTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoderTest.java index 4721e70fa..20153ba8b 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoderTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoderTest.java @@ -85,6 +85,28 @@ void encode_f64_metadata_rightBitWidth_isNonZero() throws Exception { assertThat(meta.right_bit_width()).isGreaterThan(0); } + @Test + void encode_f64_packsBothPartsAtTheDictionaryWidths() throws IOException { + // Given — quarters, whose right parts are all zero: a histogram would pack them in 0 bits, + // but Rust packs them at the right bit width, and the left codes at the width of the + // largest code, never narrower, so the cascade measures ALP-RD as Rust does + double[] values = new double[2048]; + for (int i = 0; i < values.length; i++) { + values[i] = (i % 5) * 0.25; + } + var sut = new AlpRdEncodingEncoder(); + + // When + EncodeResult result = sut.encode(DTypes.F64, values, EncodeTestHelper.testCtx()); + + // Then + var metaSeg = result.rootNode().metadata(); + ProtoALPRDMetadata meta = ProtoALPRDMetadata.decode(metaSeg, 0, metaSeg.byteSize()); + int leftWidth = Math.max(1, Integer.SIZE - Integer.numberOfLeadingZeros(meta.dict_len() - 1)); + assertThat(result.buffers().get(0).byteSize()).isEqualTo(2048L * leftWidth / 8); + assertThat(result.buffers().get(1).byteSize()).isEqualTo(2048L * meta.right_bit_width() / 8); + } + @Test void encode_f64_tiedPrefixesGetDeterministicDictionaryCodes() throws IOException { // Given: the `alp` crate's own tie test (rd_encoder_uses_deterministic_codes_for_tied_prefixes). From 5e0c7d1f986e98a633ce256c073b83b416da4128 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sun, 25 Oct 2026 20:10:53 +0100 Subject: [PATCH 09/10] fix(writer): nullable floats keep their validity in the encoding (#604) The integer change extended to every primitive: a nullable float column no longer goes through vortex.masked. As in Rust: - ALP encodes every slot, keeps the validity on the encoded ints and drops the exceptions at null rows from its patches ("we remove null exceptions in order to keep the Patches small", alp_encode_components_typed). Its terminal encode, used without a cascade, keeps the encoded ints a primitive with the validity. - ALP-RD splits every slot and keeps the validity on its bit-packed left parts (encode_generic). - Sparse takes a nullable float only when it is more than 90% null (Rust's NullDominatedSparseScheme), and compresses only its indices: the patch values stay as they are, with the all-valid bitmap Rust's filter leaves on them. - The dict layout probe now runs for every primitive. NullableFloatInteropIntegrationTest writes 11 F32/F64 shapes through vortex-jni and vortex-java and asserts the same encoding tree, then reads Java's file (default and compact) through vortex-jni, full, filtered and row-range. EncodingTrees is the tree walk both nullable cross-checks share. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01SwhK4W5cgj7y7s4ZPCif5x --- CHANGELOG.md | 2 +- docs/compatibility.md | 2 +- docs/explanation.md | 2 +- .../vortex/integration/EncodingTrees.java | 60 +++++ .../NullableFloatInteropIntegrationTest.java | 240 ++++++++++++++++++ ...NullableIntegerInteropIntegrationTest.java | 49 +--- .../dfa1/vortex/writer/VortexWriter.java | 13 +- .../writer/encode/AlpEncodingEncoder.java | 70 ++++- .../writer/encode/AlpRdEncodingEncoder.java | 34 ++- .../writer/encode/CascadingCompressor.java | 8 +- .../writer/encode/SparseEncodingEncoder.java | 20 +- .../writer/encode/AlpEncodingEncoderTest.java | 19 ++ 12 files changed, 432 insertions(+), 87 deletions(-) create mode 100644 integration/src/test/java/io/github/dfa1/vortex/integration/EncodingTrees.java create mode 100644 integration/src/test/java/io/github/dfa1/vortex/integration/NullableFloatInteropIntegrationTest.java diff --git a/CHANGELOG.md b/CHANGELOG.md index f71692323..6f59c51b2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,7 +26,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - ALP-RD packs its left and right parts at the dictionary's widths, as Rust does, instead of narrower widths that made it win where Rust picks a dictionary. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - ALP on an `F32` column compresses its encoded integers through the cascade, as Rust does for both float widths, instead of storing them raw. ([#604](https://github.com/dfa1/vortex-java/issues/604)) -- A nullable integer column keeps its validity in the encoding that stores its values instead of a `vortex.masked` wrapper, as Rust writes it, so the cascade and the dict layout pick what Rust picks. ([#604](https://github.com/dfa1/vortex-java/issues/604)) +- A nullable integer or float column keeps its validity in the encoding that stores its values instead of a `vortex.masked` wrapper, as Rust writes it, so the cascade and the dict layout pick what Rust picks. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - Zone stats tables are compressed with integer dictionaries among the candidates, as Rust's stats compressor has them. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - Run-end ends and RLE indices and offsets are stored in the narrowest unsigned type that holds them, as Rust narrows them, instead of fixed `u32`, `u16` and `u64`. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - Frame-of-reference bit-packs its residuals directly, as Rust's FoR scheme does, instead of letting the cascade pick another encoding for them. ([#604](https://github.com/dfa1/vortex-java/issues/604)) diff --git a/docs/compatibility.md b/docs/compatibility.md index 48ddb880b..968fe5d38 100644 --- a/docs/compatibility.md +++ b/docs/compatibility.md @@ -135,7 +135,7 @@ decimals ([#430](https://github.com/dfa1/vortex-java/pull/430)) and nulls in nul | `vortex.map` | `MapEncodingDecoder` | `MapEncodingEncoder` | ✅ | ✅ | Single `entries` child, a bare `vortex.listview` of non-nullable `{key, value}` structs (any other entries encoding is rejected, as in Rust); map validity is delegated to that child's validity slot. Single-flat columns only: neither `ChunkedArrayCombiner` nor `ScanIterator`'s row-window slicing has a `MapArray` arm yet, so a chunked map column or a split-window scan fails as `VortexException`. Pre-existing gap shared with `vortex.listview`; tracked separately | | `vortex.fixed_size_list` | `FixedSizeListEncodingDecoder` | `FixedSizeListEncodingEncoder` | ✅ | ✅ | | | `vortex.zstd` | `ZstdEncodingDecoder` | `ZstdEncodingEncoder` | ✅ | ✅ | Primitive, Utf8, Binary | -| `vortex.masked` | `MaskedEncodingDecoder` | `MaskedEncodingEncoder` | ✅ | ✅ | NullableData carrier. A nullable integer column is not masked when cascading: as in Rust, its validity is a plain `vortex.bool` child of the encoding that stores the values (`vortex.primitive`, `fastlanes.bitpacked`, the run values of `vortex.runend`, the indices of `fastlanes.rle`, the codes of `vortex.dict`, `vortex.pco` over the valid values only), passed down by `fastlanes.for` and `vortex.zigzag`, and one whose valid rows hold one value is a `vortex.constant` masked by a plain bitmap. Other nullable columns: inner values cascade (Dict/FSST/ALP/…) when cascading is enabled; Utf8/Primitive still compete on measured size even at depth 0, else first-match FixedSizeList. A nullable Utf8 column whose per-chunk dict is smaller skips the mask entirely: null is one more dict entry, invalid in the pool's validity, with non-nullable codes, as Rust's dict builder writes it | +| `vortex.masked` | `MaskedEncodingDecoder` | `MaskedEncodingEncoder` | ✅ | ✅ | NullableData carrier. A nullable primitive column is not masked when cascading: as in Rust, its validity is a plain `vortex.bool` child of the encoding that stores the values (`vortex.primitive`, `fastlanes.bitpacked`, the run values of `vortex.runend`, the indices of `fastlanes.rle`, the codes of `vortex.dict`, the left parts of `vortex.alprd`, `vortex.pco` over the valid values only), passed down by `fastlanes.for`, `vortex.zigzag` and `vortex.alp` (which patches only valid exceptions), and one whose valid rows hold one value is a `vortex.constant` masked by a plain bitmap. Other nullable columns: inner values cascade (Dict/FSST/ALP/…) when cascading is enabled; Utf8/Primitive still compete on measured size even at depth 0, else first-match FixedSizeList. A nullable Utf8 column whose per-chunk dict is smaller skips the mask entirely: null is one more dict entry, invalid in the pool's validity, with non-nullable codes, as Rust's dict builder writes it | | `vortex.decimal` | `DecimalEncodingDecoder` | `DecimalEncodingEncoder` | ✅ | ✅ | Precision 1–76 (i8 … i256 storage by precision), negative scales. Written from `BigDecimal[]` | | `vortex.decimal_byte_parts` | `DecimalBytePartsEncodingDecoder`| `DecimalBytePartsEncodingEncoder`| ✅ | ✅ | Write: precision ≤ 18 only (one i64 mantissa, no low parts); competes in the cascade | | `vortex.datetimeparts` | `DateTimePartsEncodingDecoder` | `DateTimePartsEncodingEncoder` | ✅ | ✅ | | diff --git a/docs/explanation.md b/docs/explanation.md index b461fd8f4..64af5aaad 100644 --- a/docs/explanation.md +++ b/docs/explanation.md @@ -354,7 +354,7 @@ single `entries` child. `vortex.masked` + `vortex.list` as before. - **A null *entry value*** (`DType.Map(key, value.asNullable(), ...)`, e.g. `{a: 1, b: null}` inside an otherwise-present map row) is a completely different bit: it rides the entry - struct's own `value` field validity — kept by the encoding of an integer value and wrapped in + struct's own `value` field validity — kept by the encoding of a primitive value and wrapped in `vortex.masked` otherwise, exactly like a nullable field anywhere else inside a `vortex.struct`. It has nothing to do with the map row being present or absent. diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/EncodingTrees.java b/integration/src/test/java/io/github/dfa1/vortex/integration/EncodingTrees.java new file mode 100644 index 000000000..2f4ada076 --- /dev/null +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/EncodingTrees.java @@ -0,0 +1,60 @@ +package io.github.dfa1.vortex.integration; + +import io.github.dfa1.vortex.core.fbs.FbsArray; +import io.github.dfa1.vortex.core.fbs.FbsArrayNode; +import io.github.dfa1.vortex.core.model.EncodingId; +import io.github.dfa1.vortex.reader.SegmentSpec; +import io.github.dfa1.vortex.reader.VortexReader; +import io.github.dfa1.vortex.reader.layout.Layout; + +import java.io.IOException; +import java.lang.foreign.MemorySegment; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; + +import static io.github.dfa1.vortex.core.io.VortexFormat.LE_INT; + +/// A file's layouts and array encodings, for comparing what vortex-jni and vortex-java write. +final class EncodingTrees { + + private EncodingTrees() { + } + + /// The file's layouts and array encodings, depth first, one indented id a line. A chunked + /// layout around a single chunk is left out: Java keeps one where Rust writes the flat layout + /// directly, and that layout shape is not what the encoding comparisons test. + static List of(Path file) throws IOException { + List out = new ArrayList<>(); + try (VortexReader reader = VortexReader.open(file)) { + layoutTree(reader, reader.layout(), 0, out); + } + return out; + } + + private static void layoutTree(VortexReader reader, Layout layout, int depth, List out) { + if (layout.isChunked() && layout.children().size() == 1) { + layoutTree(reader, layout.children().getFirst(), depth, out); + return; + } + out.add(" ".repeat(depth) + layout.layoutId().id()); + if (layout.isFlat() && !layout.segments().isEmpty()) { + SegmentSpec spec = reader.footer().segmentSpecs().get(layout.segments().getFirst()); + MemorySegment seg = reader.rawSegment(spec); + int fbLen = seg.get(LE_INT, seg.byteSize() - 4L); + FbsArray array = FbsArray.getRootAsFbsArray(seg.asSlice(seg.byteSize() - 4L - fbLen, fbLen)); + arrayTree(array.root(), reader.footer().arraySpecs(), depth + 1, out); + return; + } + for (Layout child : layout.children()) { + layoutTree(reader, child, depth + 1, out); + } + } + + private static void arrayTree(FbsArrayNode node, List specs, int depth, List out) { + out.add(" ".repeat(depth) + specs.get(node.encoding()).id()); + for (int i = 0; i < node.childrenLength(); i++) { + arrayTree(node.children(i), specs, depth + 1, out); + } + } +} diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/NullableFloatInteropIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/NullableFloatInteropIntegrationTest.java new file mode 100644 index 000000000..fb184f41c --- /dev/null +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/NullableFloatInteropIntegrationTest.java @@ -0,0 +1,240 @@ +package io.github.dfa1.vortex.integration; + +import dev.vortex.api.DataSource; +import dev.vortex.api.Expression; +import dev.vortex.api.ImmutableScanOptions; +import dev.vortex.api.Partition; +import dev.vortex.api.Scan; +import dev.vortex.api.ScanOptions; +import dev.vortex.api.Session; +import dev.vortex.api.VortexWriter; +import dev.vortex.arrow.ArrowAllocation; +import dev.vortex.jni.NativeLoader; +import io.github.dfa1.vortex.core.model.ColumnName; +import io.github.dfa1.vortex.core.model.DType; +import io.github.dfa1.vortex.core.model.PType; +import io.github.dfa1.vortex.writer.WriteOptions; +import io.github.dfa1.vortex.writer.encode.NullableData; +import org.apache.arrow.c.ArrowArray; +import org.apache.arrow.c.ArrowSchema; +import org.apache.arrow.c.Data; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import java.io.IOException; +import java.nio.channels.FileChannel; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.function.IntPredicate; +import java.util.function.IntToDoubleFunction; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; + +/// A nullable float column keeps its validity in the encoding that stores its values, as Rust's +/// arrays do (#604): ALP passes it to its encoded integers and patches only valid exceptions, +/// ALP-RD keeps it on its bit-packed left parts, a mostly-null column is sparse with a null fill, +/// and a column whose valid rows hold one value is a constant masked by a plain bitmap. +/// +/// Each column goes through vortex-jni and vortex-java, and the two files must hold the same +/// encoding tree, data and zone stats table alike; vortex-jni then reads Java's file back. +class NullableFloatInteropIntegrationTest { + + private static final Session SESSION = Session.create(); + private static final BufferAllocator ALLOCATOR = ArrowAllocation.rootAllocator(); + private static final ColumnName V = ColumnName.of("v"); + private static final int ROWS = 4 * 8192; + private static final long[] RANDOM = new Random(7).longs(ROWS).toArray(); + + static { + NativeLoader.loadJni(); + } + + static Stream shapes() { + IntPredicate every7 = i -> i % 7 == 0; + IntToDoubleFunction twoDecimals = i -> Math.floorMod(RANDOM[i], 100_000) / 100.0; + return Stream.of( + // no null: written without validity, so the plain cascade (ALP, F32 included) + shape("f64 two decimals, no nulls", PType.F64, twoDecimals, _ -> false), + shape("f32 halves, no nulls", PType.F32, i -> Math.floorMod(RANDOM[i], 1000) * 0.5, _ -> false), + shape("f32 one decimal, no nulls", PType.F32, i -> Math.floorMod(RANDOM[i], 10_000) / 10.0, _ -> false), + shape("f32 two decimals, no nulls", PType.F32, twoDecimals, _ -> false), + // ALP: the encoded ints bit-packed with the bitmap, patches only at valid rows + shape("f64 two decimals", PType.F64, twoDecimals, every7), + shape("f32 two decimals", PType.F32, twoDecimals, every7), + // ALP-RD: the bitmap on the left parts + shape("f64 random", PType.F64, i -> new Random(RANDOM[i]).nextDouble() * 1000, every7), + // mostly null: sparse, its patch values left raw with their all-valid bitmap + shape("f64 95% null", PType.F64, twoDecimals, i -> Math.floorMod(RANDOM[i] >>> 20, 100) < 95), + shape("f64 constant with nulls", PType.F64, _ -> 1.5, every7), + // few distinct: a flat dict whose codes keep the bitmap + shape("f64 five distinct", PType.F64, i -> Math.floorMod(RANDOM[i], 5) * 0.25, every7), + // taxi Airport_fee: one dominant value, two rare ones, nulls: a dict layout + shape("f64 airport fee", PType.F64, i -> i % 17 == 0 ? 1.75 : i % 101 == 0 ? -1.75 : 0.0, + i -> i % 10 == 3)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("shapes") + void javaWritesTheEncodingTreeJniWrites(String name, PType ptype, IntToDoubleFunction values, + IntPredicate isNull, @TempDir Path tmp) throws Exception { + // Given + Path jniFile = tmp.resolve("jni.vortex"); + Path javaFile = tmp.resolve("java.vortex"); + writeWithJni(jniFile, ptype, values, isNull); + List expected = EncodingTrees.of(jniFile); + + // When + writeWithJava(javaFile, ptype, values, isNull, WriteOptions.defaults()); + + // Then + assertThat(EncodingTrees.of(javaFile)).isEqualTo(expected); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("shapes") + void jniReadsJavaFile(String name, PType ptype, IntToDoubleFunction values, IntPredicate isNull, + @TempDir Path tmp) throws Exception { + assertJniReads(ptype, values, isNull, WriteOptions.defaults(), tmp); + } + + /// Compact writes add Pco, which keeps only the valid values and the bitmap as its child. + /// vortex-jni cannot write compact files, so here only its read of Java's file checks them. + @ParameterizedTest(name = "{0}") + @MethodSource("shapes") + void jniReadsJavaCompactFile(String name, PType ptype, IntToDoubleFunction values, IntPredicate isNull, + @TempDir Path tmp) throws Exception { + assertJniReads(ptype, values, isNull, WriteOptions.defaults().withCompact(true), tmp); + } + + private static void assertJniReads(PType ptype, IntToDoubleFunction values, IntPredicate isNull, + WriteOptions options, Path tmp) throws Exception { + // Given + Path file = tmp.resolve("java.vortex"); + writeWithJava(file, ptype, values, isNull, options); + List expected = new ArrayList<>(); + for (int i = 0; i < ROWS; i++) { + expected.add(isNull.test(i) ? null : stored(ptype, values.applyAsDouble(i))); + } + double threshold = stored(ptype, values.applyAsDouble(ROWS / 2)); + + // When + List result = readWithJni(file, ScanOptions.builder()); + // a filtered and a row-range read prune and slice, which a full scan never does + List filtered = readWithJni(file, ScanOptions.builder().filter(Expression.binary( + Expression.BinaryOp.GTE, Expression.column("v"), literal(ptype, threshold)))); + List range = readWithJni(file, ScanOptions.builder().rowRangeBegin(1000).rowRangeEnd(20_000)); + + // Then + assertThat(result).isEqualTo(expected); + assertThat(filtered).isEqualTo(expected.stream().filter(v -> v != null && v >= threshold).toList()); + assertThat(range).isEqualTo(expected.subList(1000, 20_000)); + } + + private static Arguments shape(String name, PType ptype, IntToDoubleFunction values, IntPredicate isNull) { + return Arguments.of(name, ptype, values, isNull); + } + + /// The value as the column stores it: an F32 column rounds it to a float. + private static double stored(PType ptype, double value) { + if (ptype == PType.F32) { + return (float) value; + } + return value; + } + + private static Expression literal(PType ptype, double value) { + if (ptype == PType.F32) { + return Expression.literal((float) value); + } + return Expression.literal(value); + } + + private static void writeWithJni(Path path, PType ptype, IntToDoubleFunction values, IntPredicate isNull) + throws Exception { + FloatingPointPrecision precision = FloatingPointPrecision.DOUBLE; + if (ptype == PType.F32) { + precision = FloatingPointPrecision.SINGLE; + } + Schema schema = new Schema(List.of(Field.nullable("v", new ArrowType.FloatingPoint(precision)))); + try (VortexWriter writer = VortexWriter.builder(SESSION, path.toAbsolutePath().toUri().toString(), schema, + ALLOCATOR).build(); + VectorSchemaRoot root = VectorSchemaRoot.create(schema, ALLOCATOR)) { + FieldVector vector = root.getVector("v"); + vector.setInitialCapacity(ROWS); + vector.allocateNew(); + for (int i = 0; i < ROWS; i++) { + if (isNull.test(i)) { + vector.setNull(i); + } else if (ptype == PType.F32) { + ((Float4Vector) vector).setSafe(i, (float) values.applyAsDouble(i)); + } else { + ((Float8Vector) vector).setSafe(i, values.applyAsDouble(i)); + } + } + vector.setValueCount(ROWS); + root.setRowCount(ROWS); + try (ArrowArray array = ArrowArray.allocateNew(ALLOCATOR); + ArrowSchema arrowSchema = ArrowSchema.allocateNew(ALLOCATOR)) { + Data.exportVectorSchemaRoot(ALLOCATOR, root, null, array, arrowSchema); + writer.writeBatch(array.memoryAddress(), arrowSchema.memoryAddress()); + } + } + } + + private static void writeWithJava(Path path, PType ptype, IntToDoubleFunction values, IntPredicate isNull, + WriteOptions options) throws IOException { + boolean[] valid = new boolean[ROWS]; + float[] floats = new float[ROWS]; + double[] doubles = new double[ROWS]; + for (int i = 0; i < ROWS; i++) { + valid[i] = !isNull.test(i); + // a null slot holds 0, as an Arrow null does + double v = valid[i] ? values.applyAsDouble(i) : 0; + floats[i] = (float) v; + doubles[i] = v; + } + Object data = ptype == PType.F32 ? floats : doubles; + DType.Struct schema = new DType.Struct(List.of(V), List.of(new DType.Primitive(ptype, true)), false); + try (var channel = FileChannel.open(path, StandardOpenOption.CREATE, StandardOpenOption.WRITE); + var writer = io.github.dfa1.vortex.writer.VortexWriter.create(channel, schema, options)) { + writer.writeChunk(Map.of(V, new NullableData(data, valid))); + } + } + + private static List readWithJni(Path file, ImmutableScanOptions.Builder options) throws IOException { + ScanOptions opts = options.projection(Expression.select(new String[]{"v"}, Expression.root())).build(); + List out = new ArrayList<>(); + DataSource ds = DataSource.open(SESSION, file.toAbsolutePath().toUri().toString()); + Scan scan = ds.scan(opts); + while (scan.hasNext()) { + Partition partition = scan.next(); + try (ArrowReader reader = partition.scanArrow(ALLOCATOR)) { + while (reader.loadNextBatch()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + FieldVector vec = root.getVector("v"); + for (int i = 0; i < root.getRowCount(); i++) { + out.add(vec.isNull(i) ? null : ((Number) vec.getObject(i)).doubleValue()); + } + } + } + } + return out; + } +} diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/NullableIntegerInteropIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/NullableIntegerInteropIntegrationTest.java index 2e1034f18..4043eed34 100644 --- a/integration/src/test/java/io/github/dfa1/vortex/integration/NullableIntegerInteropIntegrationTest.java +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/NullableIntegerInteropIntegrationTest.java @@ -10,15 +10,9 @@ import dev.vortex.api.VortexWriter; import dev.vortex.arrow.ArrowAllocation; import dev.vortex.jni.NativeLoader; -import io.github.dfa1.vortex.core.fbs.FbsArray; -import io.github.dfa1.vortex.core.fbs.FbsArrayNode; import io.github.dfa1.vortex.core.model.ColumnName; 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.reader.SegmentSpec; -import io.github.dfa1.vortex.reader.VortexReader; -import io.github.dfa1.vortex.reader.layout.Layout; import io.github.dfa1.vortex.writer.WriteOptions; import io.github.dfa1.vortex.writer.encode.NullableData; import org.apache.arrow.c.ArrowArray; @@ -39,7 +33,6 @@ import org.junit.jupiter.params.provider.MethodSource; import java.io.IOException; -import java.lang.foreign.MemorySegment; import java.nio.channels.FileChannel; import java.nio.file.Path; import java.nio.file.StandardOpenOption; @@ -51,7 +44,6 @@ import java.util.function.IntToLongFunction; import java.util.stream.Stream; -import static io.github.dfa1.vortex.core.io.VortexFormat.LE_INT; import static org.assertj.core.api.Assertions.assertThat; /// A nullable integer column keeps its validity in the encoding that stores its values, as Rust's @@ -116,13 +108,13 @@ void javaWritesTheEncodingTreeJniWrites(String name, PType ptype, IntToLongFunct Path jniFile = tmp.resolve("jni.vortex"); Path javaFile = tmp.resolve("java.vortex"); writeWithJni(jniFile, ptype, values, isNull); - List expected = encodingTree(jniFile); + List expected = EncodingTrees.of(jniFile); // When writeWithJava(javaFile, ptype, values, isNull, WriteOptions.defaults()); // Then - assertThat(encodingTree(javaFile)).isEqualTo(expected); + assertThat(EncodingTrees.of(javaFile)).isEqualTo(expected); } @ParameterizedTest(name = "{0}") @@ -243,41 +235,4 @@ private static List readWithJni(Path file, ImmutableScanOptions.Builder op } return out; } - - /// The file's layouts and array encodings, depth first, one indented id a line. A chunked - /// layout is left out: Java keeps one around a single chunk, where Rust writes the flat - /// layout directly, and that layout shape is not what these columns test. - private static List encodingTree(Path file) throws IOException { - List out = new ArrayList<>(); - try (VortexReader reader = VortexReader.open(file)) { - layoutTree(reader, reader.layout(), 0, out); - } - return out; - } - - private static void layoutTree(VortexReader reader, Layout layout, int depth, List out) { - if (layout.isChunked() && layout.children().size() == 1) { - layoutTree(reader, layout.children().getFirst(), depth, out); - return; - } - out.add(" ".repeat(depth) + layout.layoutId().id()); - if (layout.isFlat() && !layout.segments().isEmpty()) { - SegmentSpec spec = reader.footer().segmentSpecs().get(layout.segments().getFirst()); - MemorySegment seg = reader.rawSegment(spec); - int fbLen = seg.get(LE_INT, seg.byteSize() - 4L); - FbsArray array = FbsArray.getRootAsFbsArray(seg.asSlice(seg.byteSize() - 4L - fbLen, fbLen)); - arrayTree(array.root(), reader.footer().arraySpecs(), depth + 1, out); - return; - } - for (Layout child : layout.children()) { - layoutTree(reader, child, depth + 1, out); - } - } - - private static void arrayTree(FbsArrayNode node, List specs, int depth, List out) { - out.add(" ".repeat(depth) + specs.get(node.encoding()).id()); - for (int i = 0; i < node.childrenLength(); i++) { - arrayTree(node.children(i), specs, depth + 1, out); - } - } } diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java b/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java index c4f40a5eb..ea16298f1 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java @@ -654,13 +654,10 @@ public void writeChunk(Map columns) throws IOException { // Rust's dict layout admits Primitive | Utf8 | Binary (dict_layout_supported). candidate = DictColumnState.isVarBinDictCandidate((Object[]) values, validity); } - // ponytail: only primitives whose cascade sees their validity take Rust's probe - // (#604). Nullable floats still run under vortex.masked, which leaves the validity - // out of the sampled sizes Rust charges bit-packing with, and strings pick OnPair - // where Rust's StringDictScheme wins, so for those the cardinality check above - // still matches Rust more often. - if (candidate && colDtype instanceof DType.Primitive p - && (!p.nullable() || CascadingCompressor.keepsValidity(p))) { + // ponytail: only primitives take Rust's probe. Strings pick OnPair where Rust's + // StringDictScheme wins, so for them the cardinality check above still matches + // Rust more often. + if (candidate && colDtype instanceof DType.Primitive) { candidate = probeCompressesToDict(colDtype, data); } if (candidate) { @@ -876,7 +873,7 @@ private EncodingEncoder resolveOverride(DType dtype, Object data, EncodingEncode if (encodingOverride == null && dtype instanceof DType.Map) { return new io.github.dfa1.vortex.writer.encode.MapEncodingEncoder(); } - // With cascading on, a nullable integer column goes to the compressor, which keeps the + // With cascading on, a nullable primitive column goes to the compressor, which keeps the // validity in the encoding that stores the values, as Rust does (#604) boolean keepsValidity = options.allowedCascading() > 0 && CascadingCompressor.keepsValidity(dtype); if (encodingOverride == null diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoder.java index edd06a4ea..3b4a66482 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoder.java @@ -66,8 +66,28 @@ public boolean accepts(DType dtype) { return p.ptype() == PType.F64 || p.ptype() == PType.F32; } + /// A nullable array encodes every slot, keeps the validity on the encoded ints and patches + /// only valid exceptions, as Rust's `alp_encode` does. + @Override + public boolean acceptsNullable(DType dtype) { + return accepts(dtype); + } + @Override public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { + if (data instanceof NullableData) { + // No cascade to fill the encoded ints with: they stay a primitive keeping the validity + CascadeStep step = encodeCascade(dtype, data, ctx); + ChildSlot slot = step.openChildren().getFirst(); + EncodeResult child = new PrimitiveEncodingEncoder().encode(slot.childDtype(), slot.childData(), ctx); + List buffers = new ArrayList<>(step.ownedBuffers()); + EncodeNode[] children = step.partialRoot().children().clone(); + children[slot.parentChildIdx()] = EncodeNode.remapBufferIndices(child.rootNode(), buffers.size()); + buffers.addAll(child.encodedBuffers()); + EncodeNode root = new EncodeNode(EncodingId.VORTEX_ALP, step.partialRoot().metadata(), children, + new int[0]); + return new EncodeResult(root, List.copyOf(buffers), step.statsMin(), step.statsMax()); + } PType ptype = ((DType.Primitive) dtype).ptype(); return switch (ptype) { case F64 -> encodeF64((double[]) data, ctx); @@ -79,10 +99,40 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { @Override public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) { PType ptype = ((DType.Primitive) dtype).ptype(); + boolean[] validity = null; + Object values = data; + if (data instanceof NullableData nullable) { + validity = nullable.validity(); + values = nullable.values(); + } + CascadeStep step; if (ptype == PType.F64) { - return encodeCascadeF64((double[]) data, ctx); + step = encodeCascadeF64((double[]) values, validity, ctx); + } else { + step = encodeCascadeF32((float[]) values, validity, ctx); + } + // the bounds are the valid values', never a null slot's placeholder + return validity == null ? step : step.withStats(ZoneMapStats.of(dtype.withNullable(false), data)); + } + + /// Rust drops the exceptions at null rows: "we remove null exceptions in order to keep the + /// Patches small" (`alp_encode_components_typed`). Their encoded slots keep the fill value. + private static void dropNullExceptions(List indices, List values, boolean[] validity) { + if (validity == null) { + return; + } + for (int i = indices.size() - 1; i >= 0; i--) { + if (!validity[indices.get(i)]) { + indices.remove(i); + values.remove(i); + } } - return encodeCascadeF32((float[]) data, ctx); + } + + /// The encoded ints, keeping the column's validity when it has one. + private static ChildSlot encodedSlot(DType dtype, Object encoded, boolean[] validity) { + Object data = validity == null ? encoded : new NullableData(encoded, validity); + return new ChildSlot(dtype.withNullable(validity != null), data, 0, Set.of(EncodingId.VORTEX_ALP)); } // Rust's `encode_single_unchecked`: `(value * 10^e * 10^-f).fast_round() as i64`, the cast @@ -265,13 +315,14 @@ private static EncodeResult encodeF64(double[] values, EncodeContext ctx) { EncodedBuffer.of(patchBuffers.offsets(), patchBuffers.offsetType())), d.statsMin(), d.statsMax()); } - private static CascadeStep encodeCascadeF64(double[] values, EncodeContext ctx) { + private static CascadeStep encodeCascadeF64(double[] values, boolean[] validity, EncodeContext ctx) { AlpF64Data d = computeF64(values); + dropNullExceptions(d.patchIndices(), d.patchValues(), validity); if (d.patchIndices().isEmpty()) { byte[] metaBytes = new ProtoALPMetadata(d.expE(), d.expF(), null).encode(); EncodeNode partialRoot = new EncodeNode(EncodingId.VORTEX_ALP, MemorySegment.ofArray(metaBytes), new EncodeNode[1], new int[0]); - ChildSlot slot = new ChildSlot(DType.I64, d.encodedArr(), 0, Set.of(EncodingId.VORTEX_ALP)); + ChildSlot slot = encodedSlot(DType.I64, d.encodedArr(), validity); return new CascadeStep(partialRoot, List.of(), List.of(slot), d.statsMin(), d.statsMax(), true); } @@ -290,7 +341,7 @@ private static CascadeStep encodeCascadeF64(double[] values, EncodeContext ctx) EncodeNode offsetsNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 2); EncodeNode partialRoot = new EncodeNode(EncodingId.VORTEX_ALP, MemorySegment.ofArray(metaBytes), new EncodeNode[]{null, idxNode, valNode, offsetsNode}, new int[0]); - ChildSlot slot = new ChildSlot(DType.I64, d.encodedArr(), 0, Set.of(EncodingId.VORTEX_ALP)); + ChildSlot slot = encodedSlot(DType.I64, d.encodedArr(), validity); return new CascadeStep(partialRoot, List.of(EncodedBuffer.of(patchBuffers.indices(), patchBuffers.indexType()), EncodedBuffer.of(valBuf, PType.F64), EncodedBuffer.of(patchBuffers.offsets(), patchBuffers.offsetType())), List.of(slot), d.statsMin(), d.statsMax(), true); @@ -410,15 +461,16 @@ private static AlpF32Data computeF32(float[] values) { return new AlpF32Data(expE, expF, encodedArr, patchIndices, patchValues, statsMin, statsMax); } - /// The F32 counterpart of [#encodeCascadeF64(double[], EncodeContext)]: Rust's ALP scheme + /// The F32 counterpart of [#encodeCascadeF64(double[], boolean[], EncodeContext)]: Rust's ALP scheme /// compresses the encoded `i32`s through the cascade for either width. - private static CascadeStep encodeCascadeF32(float[] values, EncodeContext ctx) { + private static CascadeStep encodeCascadeF32(float[] values, boolean[] validity, EncodeContext ctx) { AlpF32Data d = computeF32(values); + dropNullExceptions(d.patchIndices(), d.patchValues(), validity); if (d.patchIndices().isEmpty()) { byte[] metaBytes = new ProtoALPMetadata(d.expE(), d.expF(), null).encode(); EncodeNode partialRoot = new EncodeNode(EncodingId.VORTEX_ALP, MemorySegment.ofArray(metaBytes), new EncodeNode[1], new int[0]); - ChildSlot slot = new ChildSlot(DType.I32, d.encodedArr(), 0, Set.of(EncodingId.VORTEX_ALP)); + ChildSlot slot = encodedSlot(DType.I32, d.encodedArr(), validity); return new CascadeStep(partialRoot, List.of(), List.of(slot), d.statsMin(), d.statsMax(), true); } @@ -437,7 +489,7 @@ private static CascadeStep encodeCascadeF32(float[] values, EncodeContext ctx) { EncodeNode offsetsNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 2); EncodeNode partialRoot = new EncodeNode(EncodingId.VORTEX_ALP, MemorySegment.ofArray(metaBytes), new EncodeNode[]{null, idxNode, valNode, offsetsNode}, new int[0]); - ChildSlot slot = new ChildSlot(DType.I32, d.encodedArr(), 0, Set.of(EncodingId.VORTEX_ALP)); + ChildSlot slot = encodedSlot(DType.I32, d.encodedArr(), validity); return new CascadeStep(partialRoot, List.of(EncodedBuffer.of(patchBuffers.indices(), patchBuffers.indexType()), EncodedBuffer.of(valBuf, PType.F32), EncodedBuffer.of(patchBuffers.offsets(), patchBuffers.offsetType())), List.of(slot), d.statsMin(), d.statsMax(), true); diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoder.java index 6f2b23b3f..fc381eebc 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoder.java @@ -42,17 +42,32 @@ public boolean accepts(DType dtype) { return p.ptype() == PType.F32 || p.ptype() == PType.F64; } + /// A nullable array splits every slot and keeps the validity on the bit-packed left parts, as + /// Rust's `RDEncoder::encode` does. + @Override + public boolean acceptsNullable(DType dtype) { + return accepts(dtype); + } + @Override public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { PType ptype = ((DType.Primitive) dtype).ptype(); - return switch (ptype) { - case F64 -> encodeF64((double[]) data, ctx); - case F32 -> encodeF32((float[]) data, ctx); + boolean[] validity = null; + Object values = data; + if (data instanceof NullableData nullable) { + validity = nullable.validity(); + values = nullable.values(); + } + EncodeResult result = switch (ptype) { + case F64 -> encodeF64((double[]) values, validity, ctx); + case F32 -> encodeF32((float[]) values, validity, ctx); default -> throw new UnsupportedOperationException("ALP-RD encode not supported for " + ptype); }; + // the bounds are the valid values', never a null slot's placeholder + return validity == null ? result : result.withStats(ZoneMapStats.of(dtype.withNullable(false), data)); } - private static EncodeResult encodeF64(double[] values, EncodeContext ctx) { + private static EncodeResult encodeF64(double[] values, boolean[] validity, EncodeContext ctx) { int n = values.length; if (n == 0) { return emptyResult(DType.U64, ctx); @@ -94,7 +109,7 @@ private static EncodeResult encodeF64(double[] values, EncodeContext ctx) { } return buildEncodeResult( - best.dict, best.rightBitWidth, leftCodes, rightParts, + best.dict, best.rightBitWidth, leftCodes, validity, rightParts, DType.U64, exceptions, scalarF64(min), scalarF64(max), ctx); } @@ -182,7 +197,7 @@ private static int countExceptionsF64(double[] values, int sampleLen, short[] di return count; } - private static EncodeResult encodeF32(float[] values, EncodeContext ctx) { + private static EncodeResult encodeF32(float[] values, boolean[] validity, EncodeContext ctx) { int n = values.length; if (n == 0) { return emptyResult(DType.U32, ctx); @@ -222,7 +237,7 @@ private static EncodeResult encodeF32(float[] values, EncodeContext ctx) { } return buildEncodeResult( - best.dict, best.rightBitWidth, leftCodes, rightParts, + best.dict, best.rightBitWidth, leftCodes, validity, rightParts, DType.U32, exceptions, scalarF32(min), scalarF32(max), ctx); } @@ -305,15 +320,16 @@ private static int indexOf(short[] dict, short value) { private static EncodeResult buildEncodeResult( short[] dict, int rightBitWidth, - short[] leftCodes, Object rightPartsData, DType rightDtype, + short[] leftCodes, boolean[] validity, Object rightPartsData, DType rightDtype, Exceptions exceptions, ProtoScalarValue statsMin, ProtoScalarValue statsMax, EncodeContext ctx) { // Rust packs both parts at the dictionary's widths (`bitpack_encode_unchecked` in // `encode_generic`), never at narrower ones a histogram would find: the left codes at the // width of the largest code, at least 1 bit (`alp` crate `bit_width`), the right parts at // the right bit width. A narrower pack measured ALP-RD smaller than Rust measures it. + Object left = validity == null ? leftCodes : new NullableData(leftCodes, validity); int leftBitWidth = Math.max(1, Integer.SIZE - Integer.numberOfLeadingZeros(dict.length - 1)); - EncodeResult leftResult = BitpackedEncodingEncoder.encodeAtWidth(DType.U16, leftCodes, leftBitWidth, ctx); + EncodeResult leftResult = BitpackedEncodingEncoder.encodeAtWidth(DType.U16, left, leftBitWidth, ctx); EncodeResult rightResult = BitpackedEncodingEncoder.encodeAtWidth(rightDtype, rightPartsData, rightBitWidth, ctx); diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java index b550ce143..000e04797 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java @@ -204,13 +204,13 @@ private static boolean eligible(EncodingEncoder enc, DType dtype, Object data, E } /// Whether a nullable column of `dtype` keeps its validity in the encoding that stores its - /// values, as every Rust array does, rather than in a `vortex.masked` wrapper (#604). Integers - /// do; floats still go through `vortex.masked` until ALP and ALP-RD carry a validity. + /// values, as every Rust array does, rather than in a `vortex.masked` wrapper (#604): every + /// primitive does. /// /// @param dtype the column's dtype - /// @return `true` for an integer primitive + /// @return `true` for a primitive public static boolean keepsValidity(DType dtype) { - return dtype instanceof DType.Primitive p && p.ptype().isInteger(); + return dtype instanceof DType.Primitive; } /// Entry point: encode `data` using the best cascading strategy. diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/SparseEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/SparseEncodingEncoder.java index 92a6f2d04..123fa1bcf 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/SparseEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/SparseEncodingEncoder.java @@ -15,6 +15,7 @@ import java.util.Set; import java.util.ArrayList; import java.util.Arrays; +import java.util.EnumSet; import java.util.List; /// Write-only encoder for `vortex.sparse`. @@ -24,6 +25,10 @@ public final class SparseEncodingEncoder implements EncodingEncoder { /// without this the cascade could sparse-encode the sparse values. private static final Set SELF = Set.of(EncodingId.VORTEX_SPARSE); + /// Every encoding but `vortex.primitive`, so a child stays canonical. + private static final Set ALL_BUT_PRIMITIVE = + Set.copyOf(EnumSet.complementOf(EnumSet.of(EncodingId.WellKnown.VORTEX_PRIMITIVE))); + /// Candidates for compressing a boolean mask's patch-index array ([#encodeBool]) — a sorted, /// often-regular sequence (e.g. a periodic null pattern), so `vortex.sequence`/`fastlanes.delta` /// are tried first. Deliberately separate from [io.github.dfa1.vortex.writer.VortexWriter]'s @@ -50,12 +55,11 @@ public boolean accepts(DType dtype) { return dtype instanceof DType.Primitive; } - /// A nullable integer array keeps its nulls in the fill or the patch values, as Rust's - /// `Sparse::encode` does. Nullable floats are not taken yet: they still go through - /// `vortex.masked`. + /// A nullable array keeps its nulls in the fill or the patch values, as Rust's + /// `Sparse::encode` does. @Override public boolean acceptsNullable(DType dtype) { - return dtype instanceof DType.Primitive p && p.ptype().isInteger(); + return accepts(dtype); } @Override @@ -80,8 +84,8 @@ public Estimate expectedRatio(DType dtype, ArrayAndStats data, EncodeContext ctx return Estimate.ratio((double) len / (len - nulls)); } } - // Rust's float sparse scheme only takes null-dominated arrays; nullable floats are still - // split off by vortex.masked before they reach a primitive here, so floats never qualify. + // Rust's float sparse scheme (NullDominatedSparseScheme) only takes the null-dominated + // arrays above if (p.ptype().isFloating()) { return Estimate.SKIP; } @@ -226,7 +230,9 @@ private static CascadeStep cascadeStep(DType dtype, ProtoScalarValue fillScalar, new EncodeNode[]{null, null}, new int[]{0}); DType idxDtype = new DType.Primitive(idxPtype, false); ChildSlot idxSlot = new ChildSlot(idxDtype, idxArr, 0, SELF); - ChildSlot valSlot = new ChildSlot(dtype, patchValues, 1, SELF); + // Rust's float sparse scheme compresses only the indices: the values stay as they are + Set valuesExcluded = ((DType.Primitive) dtype).ptype().isFloating() ? ALL_BUT_PRIMITIVE : SELF; + ChildSlot valSlot = new ChildSlot(dtype, patchValues, 1, valuesExcluded); return new CascadeStep(partialRoot, List.of(EncodedBuffer.bytes(fillBuf)), List.of(idxSlot, valSlot), ZoneMapStats.minOf(stats), ZoneMapStats.maxOf(stats), true); } diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoderTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoderTest.java index 69b784d99..dbfc2c83f 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoderTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoderTest.java @@ -364,6 +364,25 @@ void encodeCascade_f64_empty_doesNotThrow() { assertThat(result.statsMax()).isNull(); } + @Test + void encodeCascade_nullable_patchesOnlyValidExceptions() { + // Given — PI does not fit two decimals, so it is an exception at row 1 (valid) and at + // row 3 (null): Rust keeps only the valid one, and the encoded ints keep the validity + double[] values = {1.25, Math.PI, 2.5, Math.PI, 3.75}; + boolean[] validity = {true, true, true, false, true}; + NullableData data = new NullableData(values, validity); + + // When + CascadeStep result = ENCODER.encodeCascade(DTypes.F64_N, data, EncodeTestHelper.cascadingCtx()); + + // Then + ChildSlot encoded = result.openChildren().getFirst(); + assertThat(encoded.childData()).isInstanceOf(NullableData.class); + assertThat(((NullableData) encoded.childData()).validity()).isEqualTo(validity); + assertThat(result.ownedBuffers().getFirst().data().byteSize()) + .as("one patch index").isEqualTo(1); + } + @Test void encodeCascade_f32_opensTheEncodedIntsForTheCascade() { // Given — Rust's ALP scheme compresses the encoded ints of either width; f32 wrote From 1f8c3db7c51184ffb7125d29bcc98230086a9cf7 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sun, 25 Oct 2026 20:14:12 +0100 Subject: [PATCH 10/10] fix(writer): try encodings in Rust's registration order Rust breaks a tie between two equal scores by registration order within each pass (choose_best_scheme), and registers integers as FoR, ZigZag, BitPacking, Sparse, IntDict, RunEnd, Sequence, RLE, Delta and floats as ALP, ALP-RD, FloatDict, sparse, RLE (vortex-btrblocks session.rs). Java's list had RLE and run-end ahead of dict and bit-packing last, so a tie went the other way. Ties are rare but real: on a float column of runs with null runs, RLE and dict both measure 644 bytes on the dict layout probe's sample. Rust takes the dict and writes a dict layout; Java took RLE and wrote none (found comparing nullable floats, #604). One list serves both families, so Dict sits before Sparse as floats have them; for integers the two only compete in the probe and zone tables, where a tie of their two estimates is the one case still broken the other way. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01SwhK4W5cgj7y7s4ZPCif5x --- CHANGELOG.md | 1 + .../NullableFloatInteropIntegrationTest.java | 4 +++ .../dfa1/vortex/writer/VortexWriter.java | 33 +++++++++++-------- 3 files changed, 24 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6f59c51b2..e52682be0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **Breaking:** CSV import stores an empty cell in a number or boolean column as null instead of `0` or `false`, so inferred number and boolean columns are now nullable (read through `MaskedArray`), and an empty cell or a missing column fails the import where the schema declares the column non-nullable. ([#592](https://github.com/dfa1/vortex-java/issues/592)) ### Fixed +- The cascade tries encodings in Rust's registration order, so a tie between two equally small encodings goes the way Rust breaks it. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - ALP-RD packs its left and right parts at the dictionary's widths, as Rust does, instead of narrower widths that made it win where Rust picks a dictionary. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - ALP on an `F32` column compresses its encoded integers through the cascade, as Rust does for both float widths, instead of storing them raw. ([#604](https://github.com/dfa1/vortex-java/issues/604)) - A nullable integer or float column keeps its validity in the encoding that stores its values instead of a `vortex.masked` wrapper, as Rust writes it, so the cascade and the dict layout pick what Rust picks. ([#604](https://github.com/dfa1/vortex-java/issues/604)) diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/NullableFloatInteropIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/NullableFloatInteropIntegrationTest.java index fb184f41c..9ce81f96a 100644 --- a/integration/src/test/java/io/github/dfa1/vortex/integration/NullableFloatInteropIntegrationTest.java +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/NullableFloatInteropIntegrationTest.java @@ -85,6 +85,10 @@ static Stream shapes() { shape("f64 constant with nulls", PType.F64, _ -> 1.5, every7), // few distinct: a flat dict whose codes keep the bitmap shape("f64 five distinct", PType.F64, i -> Math.floorMod(RANDOM[i], 5) * 0.25, every7), + // runs with null runs: RLE and dict tie on the probe's sample, and Rust's + // registration order (dict first) gives the dict layout + shape("f64 runs with null runs", PType.F64, i -> Math.floorMod(RANDOM[i / 20], 100_000) / 7.0, + i -> (i / 20) % 7 == 0), // taxi Airport_fee: one dominant value, two rare ones, nulls: a dict layout shape("f64 airport fee", PType.F64, i -> i % 17 == 0 ? 1.75 : i % 101 == 0 ? -1.75 : 0.0, i -> i % 10 == 3)); diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java b/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java index ea16298f1..6747e1019 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java @@ -359,18 +359,32 @@ private static List buildCascadeCodecs(WriteOptions options) { codecs.add(new FixedSizeListEncodingEncoder()); codecs.add(new ListEncodingEncoder()); codecs.add(new ConstantEncodingEncoder()); + // The primitive schemes in Rust's registration order (vortex-btrblocks session.rs), which + // breaks a tie between equal scores: integers FoR, ZigZag, BitPacking, Sparse, IntDict, + // RunEnd, Sequence, RLE, Delta; floats ALP, ALP-RD, FloatDict, sparse, RLE. One list + // serves both, so Dict sits before Sparse as the floats have it; for integers the two only + // meet in the probe and zone tables, on an exact tie of two estimates. + codecs.add(new FrameOfReferenceEncodingEncoder()); + // ZigZag folds signed values to small unsigned magnitudes and cascades them into + // bit-packing, as Rust's ZigZagScheme does (issue #410). It competes with FoR, which + // always fits in as few bits, so it is here for parity rather than for size. + codecs.add(new ZigZagEncodingEncoder()); + codecs.add(new BitpackedEncodingEncoder()); codecs.add(new AlpEncodingEncoder()); // ALP-RD competes for high-precision F64/F32 that plain ALP can't fit without too many // exceptions (e.g. nyc-311 Latitude): without it in the competition such columns fall back // to raw vortex.primitive or dict, matching neither the Rust reference (which uses alprd) nor // its size (#304). Registered on WriteRegistry already; this adds it as a top-level candidate. codecs.add(new AlpRdEncodingEncoder()); - codecs.add(new FrameOfReferenceEncodingEncoder()); - // ZigZag folds signed values to small unsigned magnitudes and cascades them into - // bit-packing, as Rust's ZigZagScheme does (issue #410). It competes with FoR, which - // always fits in as few bits, so it is here for parity rather than for size. - codecs.add(new ZigZagEncodingEncoder()); + // Rust's file writer keeps integers out of the data cascade because its DictStrategy + // dictionary-encodes suitable columns as a layout. That is our global dictionary: with it + // off (a streaming CSV import cannot buffer a whole column), an integer chunk keeps its + // per-chunk dictionary, or a low-cardinality integer column loses its dictionary entirely. + codecs.add(options.globalDict() ? DictEncodingEncoder.forDataCascade() : new DictEncodingEncoder()); + codecs.add(new SparseEncodingEncoder()); codecs.add(new RunEndEncodingEncoder()); + // Rust's SequenceScheme: an exact integer arithmetic sequence costs two scalars (#507). + codecs.add(new SequenceEncodingEncoder()); codecs.add(new RleEncodingEncoder()); // Delta is in no edition ("no edition includes fastlanes.delta yet"), so like Rust's // DeltaScheme it competes only when the edition guard is off (WriteOptions#withoutEditions). @@ -378,15 +392,6 @@ private static List buildCascadeCodecs(WriteOptions options) { if (allowed.isEmpty() || allowed.contains(EncodingId.FASTLANES_DELTA)) { codecs.add(new DeltaEncodingEncoder()); } - codecs.add(new SparseEncodingEncoder()); - // Rust's SequenceScheme: an exact integer arithmetic sequence costs two scalars (#507). - codecs.add(new SequenceEncodingEncoder()); - // Rust's file writer keeps integers out of the data cascade because its DictStrategy - // dictionary-encodes suitable columns as a layout. That is our global dictionary: with it - // off (a streaming CSV import cannot buffer a whole column), an integer chunk keeps its - // per-chunk dictionary, or a low-cardinality integer column loses its dictionary entirely. - codecs.add(options.globalDict() ? DictEncodingEncoder.forDataCascade() : new DictEncodingEncoder()); - codecs.add(new BitpackedEncodingEncoder()); // Decimals: byte-parts (precision <= 18) cascades its i64 mantissa through FoR/bit-packing; // plain vortex.decimal covers every width. The competition keeps the smaller. codecs.add(new DecimalBytePartsEncodingEncoder());