fix(writer): write a nullable timestamp with nulls as datetimeparts

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Fixed
+- A nullable timestamp column holding nulls is written as `vortex.datetimeparts` with the validity on its days part, as Rust writes it, instead of `vortex.ext` over a masked i64. ([#627](https://github.com/dfa1/vortex-java/issues/627))
 - A column whose encoded bytes pass 2 GiB in one chunk, such as large audio blobs under FSST, writes and reads back instead of failing with an integer overflow. ([#629](https://github.com/dfa1/vortex-java/issues/629))
 - Parquet import reads a `DECIMAL` column as a Vortex decimal with its precision and scale, in all four physical encodings; `INT32`/`INT64` ones lost their scale and the byte-array ones failed. ([#628](https://github.com/dfa1/vortex-java/issues/628))
 - Parquet import keeps the nulls of a nullable `TIMESTAMP` column instead of failing as if the column were `REQUIRED`. ([#627](https://github.com/dfa1/vortex-java/issues/627))

integration/src/test/java/io/github/dfa1/vortex/integration/NullableTimestampInteropIntegrationTest.java

@@ -0,0 +1,187 @@
+package io.github.dfa1.vortex.integration;
+
+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.TimeUnit;
+import io.github.dfa1.vortex.core.model.TimestampDtype;
+import io.github.dfa1.vortex.parquet.ImportOptions;
+import io.github.dfa1.vortex.parquet.ParquetExporter;
+import io.github.dfa1.vortex.parquet.ParquetImporter;
+import io.github.dfa1.vortex.writer.WriteOptions;
+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.TimeStampMicroVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+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.Test;
+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.time.Instant;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.function.IntPredicate;
+import java.util.stream.Stream;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/// A nullable timestamp column is written as vortex-jni writes it: `vortex.datetimeparts`, with the
+/// validity on its days child as Rust's `split_temporal` keeps it, and no validity at all when every
+/// row is valid. The Java side used to decline datetimeparts for any chunk holding a null, writing
+/// `vortex.ext` over a masked i64, so the column's encoding changed with where its nulls fell.
+///
+/// Only the data column's root encoding and where its validity sits are compared: the cascade under
+/// the days child and the zone-stats table still differ from Rust with or without nulls (#634).
+class NullableTimestampInteropIntegrationTest {
+
+    private static final Session SESSION = Session.create();
+    private static final BufferAllocator ALLOCATOR = ArrowAllocation.rootAllocator();
+    private static final ColumnName TS = ColumnName.of("ts");
+    private static final DType.Extension TS_DTYPE = TimestampDtype.of(TimeUnit.Microseconds, null, true);
+    private static final int ROWS = 4 * 8192;
+    private static final long START_MICROS = 1_704_067_200_000_000L; // 2024-01-01T00:00:00Z
+    private static final long[] JITTER_MICROS = new Random(634).longs(ROWS, 0, 60_000_000L).toArray();
+
+    static {
+        NativeLoader.loadJni();
+    }
+
+    static Stream<Arguments> shapes() {
+        return Stream.of(
+                // nullable but every row valid: written without a validity, as Rust's AllValid
+                Arguments.of("no nulls", (IntPredicate) _ -> false),
+                Arguments.of("every 7th null", (IntPredicate) i -> i % 7 == 0),
+                Arguments.of("one null", (IntPredicate) i -> i == 5),
+                Arguments.of("nulls in the second half", (IntPredicate) i -> i >= ROWS / 2 && i % 3 == 0));
+    }
+
+    @ParameterizedTest(name = "{0}")
+    @MethodSource("shapes")
+    void javaWritesTheEncodingTreeJniWrites(String name, IntPredicate isNull, @TempDir Path tmp) throws Exception {
+        // Given
+        Path jniFile = tmp.resolve("jni.vortex");
+        Path javaFile = tmp.resolve("java.vortex");
+        writeWithJni(jniFile, isNull);
+        List<String> expected = EncodingTrees.of(jniFile);
+
+        // When
+        writeWithJava(javaFile, isNull);
+
+        // Then
+        assertThat(dataColumn(expected).getFirst()).contains("vortex.datetimeparts");
+        assertDataColumnMatches(EncodingTrees.of(javaFile), expected);
+    }
+
+    /// The Parquet import path hands the writer batches with and without nulls, which its
+    /// repartition coalesces into one chunk; that chunk is datetimeparts too.
+    @Test
+    void parquetImport_writesTheEncodingTreeJniWrites(@TempDir Path tmp) throws Exception {
+        // Given
+        IntPredicate isNull = i -> i >= ROWS / 2 && i % 3 == 0;
+        Path jniFile = tmp.resolve("jni.vortex");
+        writeWithJni(jniFile, isNull);
+        Path source = tmp.resolve("source.vortex");
+        writeWithJava(source, isNull);
+        Path parquet = tmp.resolve("source.parquet");
+        ParquetExporter.exportParquet(source, parquet);
+        Path imported = tmp.resolve("imported.vortex");
+
+        // When
+        ParquetImporter.importParquet(parquet, imported, ImportOptions.defaults().withChunkSize(4096));
+
+        // Then
+        assertDataColumnMatches(EncodingTrees.of(imported), EncodingTrees.of(jniFile));
+    }
+
+    /// Same root encoding for the data column, and a validity (`vortex.bool`) in its days child
+    /// exactly when Rust writes one.
+    private static void assertDataColumnMatches(List<String> actual, List<String> expected) {
+        List<String> actualColumn = dataColumn(actual);
+        List<String> expectedColumn = dataColumn(expected);
+        assertThat(actualColumn.getFirst()).isEqualTo(expectedColumn.getFirst());
+        assertThat(firstChild(actualColumn).stream().anyMatch(l -> l.contains("vortex.bool")))
+                .isEqualTo(firstChild(expectedColumn).stream().anyMatch(l -> l.contains("vortex.bool")));
+    }
+
+    /// The data column's encoding subtree: the lines under the zoned layout's first flat, up to the
+    /// stats table's flat.
+    private static List<String> dataColumn(List<String> tree) {
+        int start = tree.indexOf("    vortex.flat") + 1;
+        int end = tree.subList(start, tree.size()).indexOf("    vortex.flat");
+        return tree.subList(start, end < 0 ? tree.size() : start + end);
+    }
+
+    /// The first child (days) subtree of a data column subtree.
+    private static List<String> firstChild(List<String> column) {
+        int depth = indent(column.getFirst()) + 2;
+        int end = 2;
+        while (end < column.size() && indent(column.get(end)) > depth) {
+            end++;
+        }
+        return column.subList(1, end);
+    }
+
+    private static int indent(String line) {
+        return line.length() - line.stripLeading().length();
+    }
+
+    /// One row about every 90 s with a jitter of up to a minute: not an arithmetic sequence, which
+    /// Rust would write as `vortex.sequence` instead of splitting it (#634).
+    private static long micros(int i) {
+        return START_MICROS + i * 90_000_000L + JITTER_MICROS[i];
+    }
+
+    private static void writeWithJni(Path path, IntPredicate isNull) throws Exception {
+        Schema schema = new Schema(List.of(Field.nullable("ts",
+                new ArrowType.Timestamp(org.apache.arrow.vector.types.TimeUnit.MICROSECOND, null))));
+        try (VortexWriter writer = VortexWriter.builder(SESSION, path.toAbsolutePath().toUri().toString(), schema,
+                ALLOCATOR).build();
+             VectorSchemaRoot root = VectorSchemaRoot.create(schema, ALLOCATOR)) {
+            TimeStampMicroVector vector = (TimeStampMicroVector) root.getVector("ts");
+            vector.setInitialCapacity(ROWS);
+            vector.allocateNew();
+            for (int i = 0; i < ROWS; i++) {
+                if (isNull.test(i)) {
+                    vector.setNull(i);
+                } else {
+                    vector.setSafe(i, micros(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, IntPredicate isNull) throws IOException {
+        List<Instant> values = new ArrayList<>(ROWS);
+        for (int i = 0; i < ROWS; i++) {
+            values.add(isNull.test(i) ? null : Instant.EPOCH.plusNanos(micros(i) * 1_000L));
+        }
+        DType.Struct schema = new DType.Struct(List.of(TS), List.of(TS_DTYPE), false);
+        try (FileChannel channel = FileChannel.open(path, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
+             var writer = io.github.dfa1.vortex.writer.VortexWriter.create(channel, schema, WriteOptions.defaults())) {
+            writer.writeChunk(Map.of(TS, values));
+        }
+    }
+}

integration/src/test/java/io/github/dfa1/vortex/integration/TimestampFilterInteropIntegrationTest.java

@@ -93,9 +93,8 @@ static void writeFiles() throws IOException {
                 }
             }
         }
-        // As the Parquet importer writes a timestamp column: a nullable column handed over as
-        // DateTimePartsData, the only input the writer stores as datetimeparts (which has no null
-        // rows on this path; the unit tests cover nulls).
+        // As the Parquet importer writes a timestamp column without nulls: a nullable column
+        // handed over as DateTimePartsData, which the writer stores as datetimeparts.
         partsFile = write("timestamp-parts.vtx", WriteOptions.defaults(), ids,
                 b -> new DateTimePartsData(ts[b], true));
         // The same values as i64 storage under the extension: the default cascade would pick
@@ -154,12 +153,21 @@ static Stream<Case> cases() {
     void timestampColumnsAreStoredAsPartsAndAsFlatLongs() throws IOException {
         // Given the files written above
         // When
-        try (VortexReader parts = VortexReader.open(partsFile);
-             VortexReader flat = VortexReader.open(flatFile)) {
-            // Then the filters run over both storages the fast path reads
-            assertThat(parts.footer().arraySpecs()).contains(EncodingId.VORTEX_DATETIMEPARTS);
-            assertThat(flat.footer().arraySpecs()).doesNotContain(EncodingId.VORTEX_DATETIMEPARTS);
-        }
+        String partsData = tsDataEncoding(partsFile);
+        String flatData = tsDataEncoding(flatFile);
+
+        // Then the filters run over both storages the fast path reads. Only the data column is
+        // checked: the zone-stats table of a nullable timestamp is datetimeparts in both files.
+        assertThat(partsData).isEqualTo(EncodingId.VORTEX_DATETIMEPARTS.id());
+        assertThat(flatData).isEqualTo(EncodingId.VORTEX_EXT.id());
+    }
+
+    /// The root encoding of the ts column's data: the array under the first flat of its zoned layout.
+    private static String tsDataEncoding(Path file) throws IOException {
+        List<String> tree = EncodingTrees.of(file);
+        int tsZoned = tree.lastIndexOf("  vortex.zoned");
+        assertThat(tsZoned).as("the ts column's zoned layout in %s", tree).isPositive();
+        return tree.get(tree.subList(tsZoned, tree.size()).indexOf("    vortex.flat") + tsZoned + 1).strip();
     }
 
     @ParameterizedTest(name = "{0}: {1}")

writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java

@@ -357,11 +357,11 @@ private static WriteRegistry buildRegistry(List<EncodingEncoder> encoders) {
     private static List<EncodingEncoder> buildCascadeCodecs(WriteOptions options) {
         List<EncodingEncoder> codecs = new ArrayList<>();
         // FbsExtension-dtype dispatch order matters: findPrimitiveEncoding picks the first
-        // accepting codec. DateTimePartsEncoding goes first because it consumes
-        // pre-decomposed DateTimePartsData (Parquet importer path); when the data is
-        // raw primitive storage (JDBC's long[] via TimestampExtension.encodeAll) it
-        // returns notApplicable and spliceResult excludes it, falling back to
-        // ExtEncoding which cascades the storage child through FoR/Bitpacked/RLE/ALP.
+        // accepting codec. DateTimePartsEncoding goes first: it splits any vortex.timestamp
+        // storage (DateTimePartsData, long[], or NullableData over long[] with its validity on
+        // the days child); for other extensions, Days units or an all-null chunk it returns
+        // notApplicable and spliceResult excludes it, falling back to ExtEncoding, which
+        // cascades the storage child through FoR/Bitpacked/RLE/ALP.
         // FixedSizeListEncoding handles UUID-style fixed-size byte storage downstream
         // of ExtEncoding.
         codecs.add(new DateTimePartsEncodingEncoder());

writer/src/main/java/io/github/dfa1/vortex/writer/encode/CascadingCompressor.java

@@ -261,8 +261,7 @@ private EncodeResult encodeWithCtx(DType dtype, Object data, EncodeContext ctx)
         // splice through it so its cascaded children (e.g. datetimeparts → days/seconds/
         // subseconds) are recursively compressed rather than stored as raw primitives. Honor the
         // excluded set so spliceResult's notApplicable retry can rotate to the next accepting
-        // encoding (e.g. DateTimePartsEncoding → ExtEncoding when the input is raw storage rather
-        // than DateTimePartsData).
+        // encoding (e.g. DateTimePartsEncoding → ExtEncoding for a date, or an all-null chunk).
         if (!(dtype instanceof DType.Primitive)) {
             return spliceResult(findPrimitiveEncoding(dtype, ctx.excluded()), dtype, data, ctx);
         }