diff --git a/CHANGELOG.md b/CHANGELOG.md index e5f44378f..e52682be0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,15 @@ 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)) +- 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)) +- 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/docs/compatibility.md b/docs/compatibility.md index d297b5d37..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; 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 30ebca470..64af5aaad 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 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. 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/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..9ce81f96a --- /dev/null +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/NullableFloatInteropIntegrationTest.java @@ -0,0 +1,244 @@ +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), + // 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)); + } + + @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 new file mode 100644 index 000000000..4043eed34 --- /dev/null +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/NullableIntegerInteropIntegrationTest.java @@ -0,0 +1,238 @@ +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.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.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.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 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 = 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, 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; + } +} 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 9ccea0cd6..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()); @@ -654,11 +659,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 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 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) { @@ -874,9 +878,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 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 && 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(); } @@ -1320,8 +1327,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() { 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..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 CascadeStep.terminal(encode(dtype, 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); @@ -367,7 +418,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 +458,53 @@ 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[], boolean[], EncodeContext)]: Rust's ALP scheme + /// compresses the encoded `i32`s through the cascade for either width. + 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 = encodedSlot(DType.I32, d.encodedArr(), validity); + 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 = 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); + } + + 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 +550,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/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..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,12 +320,18 @@ 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) { - 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. + 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, left, 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/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 bf6581b5d..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 @@ -48,10 +48,97 @@ 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, 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(); + 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 (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); + } + 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. + private static EncodeResult encodeNullable(DType dtype, NullableData data, EncodeContext ctx) { + PType ptype = ((DType.Primitive) dtype).ptype(); + 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 cbfb81d70..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 @@ -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,6 +183,36 @@ private static long primitiveBytes(DType dtype, int n) { return (long) n * p.ptype().byteSize(); } + /// 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) { + 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): every + /// primitive does. + /// + /// @param dtype the column's dtype + /// @return `true` for a primitive + public static boolean keepsValidity(DType dtype) { + return dtype instanceof DType.Primitive; + } + /// Entry point: encode `data` using the best cascading strategy. /// /// Cascade parameters (depth, sampling, exclusions) are taken from `ctx`. @@ -183,10 +224,19 @@ private static long primitiveBytes(DType dtype, int n) { /// @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) @@ -227,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()); } } @@ -259,22 +309,70 @@ 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); + } + 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(); } - // Run the winner on the full data - return spliceResult(choice.encoder(), dtype, data, ctx); + 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 /// @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. @@ -296,9 +394,10 @@ 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())) { + if (!eligible(enc, dtype, input.data(), ctx)) { continue; } Estimate est = competes(dtype) ? enc.expectedRatio(dtype, input, ctx) : Estimate.COMPLETE; @@ -310,13 +409,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(); @@ -331,9 +431,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 @@ -358,13 +459,16 @@ 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; } 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)); } @@ -422,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 6e61f1269..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 @@ -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). @@ -73,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; @@ -89,14 +97,43 @@ 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'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 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 afad4a9a4..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,22 +154,69 @@ 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); - 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 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(PType.U16, false), runs.indices(), 1, POSITIONS_EXCLUDED), - new ChildSlot(new DType.Primitive(PType.U64, false), runs.offsets(), 2, POSITIONS_EXCLUDED)); - MinMax stats = ZoneMapStats.of(dtype, data); + 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.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) @@ -177,12 +230,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..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,41 +224,70 @@ 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); 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 valuesSlot = new ChildSlot(dtype, PrimitiveArrays.fromLongsArray(values, ptype, EncodingId.VORTEX_RUNEND), 1, VALUES_EXCLUDED); - MinMax stats = ZoneMapStats.of(dtype, data); + ChildSlot endsSlot = new ChildSlot(new DType.Primitive(endsPtype, false), + PrimitiveArrays.fromIntsArray(ends, endsPtype, EncodingId.VORTEX_RUNEND), 0, ENDS_EXCLUDED); + // 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..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 @@ -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,8 @@ import java.lang.foreign.MemorySegment; 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`. @@ -22,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 @@ -48,6 +55,13 @@ public boolean accepts(DType dtype) { return dtype instanceof DType.Primitive; } + /// 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 accepts(dtype); + } + @Override public StatsOptions statsOptions() { return StatsOptions.defaults().withTrackMostFrequent(true); @@ -58,8 +72,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 (NullDominatedSparseScheme) only takes the null-dominated + // arrays above if (p.ptype().isFloating()) { return Estimate.SKIP; } @@ -83,12 +109,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 +148,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 +230,9 @@ 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); + // 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); } @@ -177,7 +265,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/AlpEncodingEncoderTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/AlpEncodingEncoderTest.java index e341153b0..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 @@ -363,6 +363,40 @@ void encodeCascade_f64_empty_doesNotThrow() { assertThat(result.statsMin()).isNull(); 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 + // 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 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). 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 6b99e1494..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,189 @@ 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. + @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 { 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); + } } } 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);