fix: read and write string and binary chunks past 2 GiB in every encoding

CHANGELOG.md

@@ -9,7 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### 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))
+- A string or binary column whose FSST codes pass 2 GiB in one chunk, such as large audio blobs, writes and reads back instead of failing with an integer overflow. ([#629](https://github.com/dfa1/vortex-java/issues/629))
+- A string or binary chunk past 2 GiB writes and reads back as `vortex.varbin`, `vortex.varbinview` and compact `vortex.zstd` too, and Rust-written `varbin` columns with `u32` offsets past 2 GiB read instead of failing as out of range. ([#633](https://github.com/dfa1/vortex-java/issues/633))
 - 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/FsstLargeChunkInteropIntegrationTest.java → integration/src/test/java/io/github/dfa1/vortex/integration/LargeChunkInteropIntegrationTest.java

@@ -23,8 +23,9 @@
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.ipc.ArrowReader;
 import org.junit.jupiter.api.Tag;
-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.ValueSource;
 
 import java.io.IOException;
 import java.nio.channels.FileChannel;
@@ -39,17 +40,25 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
-/// An FSST column whose compressed codes pass 2 GiB in one chunk writes and reads back (#629): the
-/// code offsets go 64-bit, as Rust's `fsst_compress` switches to `i64` offsets past `i32::MAX`; the
-/// segment, whose length is a `u32` on the wire, is written past a `ByteBuffer`'s 2 GiB and read
-/// back as one. Found on Raincloud's `peoples-speech-clean-validation`, whose audio column is
-/// ~2.3 GB in one 8192-row block.
+/// A binary column whose bytes pass 2 GiB in one chunk writes and reads back in every encoding that
+/// can hold it: FSST's code offsets go 64-bit as Rust's `fsst_compress` does past `i32::MAX`, and the
+/// segment, a `u32` length on the wire, is written and read past a `ByteBuffer`'s 2 GiB (#629);
+/// `vortex.varbin` copies its rows per row past 2 GiB, with `i64` offsets; `vortex.varbinview`
+/// spreads its rows over data buffers below 2 GiB each; `vortex.zstd` sums its length-prefixed
+/// payload in a `long` and closes a frame before 1 GiB (#633). Found on Raincloud's
+/// `peoples-speech-clean-validation` (2.3 GB of audio in one 8192-row block) and `mmmu` (2.95 GB of
+/// images).
 ///
-/// FSST only wins on bytes it can shrink, so the rows compress slightly (a repeated word in a tenth
-/// of their 8-byte slots, the rest bytes FSST codes one for one) and ~2.46 GB of input is needed;
-/// every row shares one array to keep the heap small. Tagged `large`: several GB of memory.
+/// Not covered here: reading Rust's `u32` varbin offsets past 2 GiB (`mmmu`). Java writes `i64`
+/// offsets, and a Rust-written fixture that size is too big to commit, so `VarBinArrayTest` pins the
+/// unsigned read and Raincloud's `mmmu` is its cross-check.
+///
+/// Each case forces one encoding on ~2.46 GB of slightly compressible rows (a repeated word in a tenth
+/// of their 8-byte slots, the rest bytes FSST codes one for one), so FSST and zstd both win and keep
+/// their encoding; every row shares one array to keep the heap small. Tagged `large`: several GB of
+/// memory.
 @Tag("large")
-class FsstLargeChunkInteropIntegrationTest {
+class LargeChunkInteropIntegrationTest {
 
     private static final Session SESSION = Session.create();
     private static final BufferAllocator ALLOCATOR = ArrowAllocation.rootAllocator();
@@ -62,15 +71,16 @@ class FsstLargeChunkInteropIntegrationTest {
         NativeLoader.loadJni();
     }
 
-    @Test
-    void codesPastTwoGibibytes_roundTripAndJniReads(@TempDir Path tmp) throws Exception {
+    @ParameterizedTest
+    @ValueSource(strings = {"vortex.fsst", "vortex.varbin", "vortex.varbinview", "vortex.zstd"})
+    void chunkPastTwoGibibytes_roundTripsAndJniReads(String encoding, @TempDir Path tmp) throws Exception {
         // Given
         byte[] row = slightlyCompressibleRow();
         byte[][] rows = new byte[ROWS][];
         Arrays.fill(rows, row);
-        Path file = tmp.resolve("fsst.vortex");
-        WriteOptions options = WriteOptions.defaults()
-                .withColumnEncoding(BLOB, ColumnEncoding.candidates(EncodingId.VORTEX_FSST));
+        Path file = tmp.resolve("large.vortex");
+        EncodingId id = EncodingId.parse(encoding);
+        WriteOptions options = WriteOptions.defaults().withColumnEncoding(BLOB, ColumnEncoding.candidates(id));
 
         // When
         DType.Struct schema = new DType.Struct(List.of(BLOB), List.of(new DType.Binary(false)), false);
@@ -79,9 +89,13 @@ void codesPastTwoGibibytes_roundTripAndJniReads(@TempDir Path tmp) throws Except
             writer.writeChunk(Map.of(BLOB, rows));
         }
 
-        // Then — FSST was kept and its codes passed 2 GiB; the last rows, whose codes start past it,
+        // Then — the forced encoding was kept, and the last rows, whose bytes start past 2 GiB,
         // decode to the input in vortex-jni and vortex-java
-        assertThat(file.toFile().length()).isGreaterThan(Integer.MAX_VALUE).isLessThan((long) ROWS * ROW_BYTES);
+        assertThat(dataEncoding(file)).isEqualTo(encoding);
+        if (id != EncodingId.VORTEX_ZSTD) {
+            // zstd shrinks the repeated row to almost nothing; its 2 GiB is the payload it compresses
+            assertThat(file.toFile().length()).isGreaterThan(Integer.MAX_VALUE);
+        }
         assertThat(readTailWithJni(file, ROWS - 4, ROWS)).hasSize(4).allSatisfy(b -> assertThat(b).isEqualTo(row));
         assertThat(readLastWithJava(file)).isEqualTo(row);
     }
@@ -102,6 +116,14 @@ private static byte[] slightlyCompressibleRow() {
         return row;
     }
 
+    /// The root encoding of the blob column's data: the array under its zoned layout's first flat.
+    private static String dataEncoding(Path file) throws IOException {
+        List<String> tree = EncodingTrees.of(file);
+        int flat = tree.indexOf("    vortex.flat");
+        assertThat(flat).as("the blob column's data flat in %s", tree).isPositive();
+        return tree.get(flat + 1).strip();
+    }
+
     private static List<byte[]> readTailWithJni(Path file, long begin, long end) throws IOException {
         ScanOptions opts = ScanOptions.builder()
                 .projection(Expression.select(new String[]{"blob"}, Expression.root()))

reader/src/main/java/io/github/dfa1/vortex/reader/array/VarBinOffsetArray.java

@@ -109,10 +109,14 @@ private int offsetWidth() {
     /// `IndexOutOfBoundsException` (ADR 0003).
     ///
     /// @param i zero-based offset index, in `[0, length]`
-    /// @return the offset value widened to a signed long
+    /// @return the offset value widened to a long, unsigned for `U32`
     private long readOffset(long i) {
         try {
-            if (offsetsPtype == PType.I32 || offsetsPtype == PType.U32) {
+            if (offsetsPtype == PType.U32) {
+                // a u32 offset passes Integer.MAX_VALUE once the data buffer does (#633)
+                return Integer.toUnsignedLong(offsetsSegment.getAtIndex(VortexFormat.LE_INT, i));
+            }
+            if (offsetsPtype == PType.I32) {
                 return offsetsSegment.getAtIndex(VortexFormat.LE_INT, i);
             }
             return offsetsSegment.getAtIndex(VortexFormat.LE_LONG, i);

reader/src/main/java/io/github/dfa1/vortex/reader/decode/ZstdEncodingDecoder.java

@@ -161,8 +161,8 @@ private static VarBinArray buildScatteredVarBin(
         }
 
         MemorySegment values = ctx.arena().allocate(totalDataBytes > 0 ? totalDataBytes : 1);
-        MemorySegment offsets = ctx.arena().allocate((rowCount + 1) * 4L, 4);
-        offsets.setAtIndex(VortexFormat.LE_INT, 0, 0);
+        PType offsetsPtype = offsetsPtype(totalDataBytes);
+        MemorySegment offsets = ctx.arena().allocate((rowCount + 1) * offsetsPtype.byteSize(), offsetsPtype.byteSize());
 
         // Second pass reads the same positions the first pass already bounds-checked via
         // readVarBinLen, so a raw get/copy here cannot overrun; values is sized to the validated
@@ -177,10 +177,10 @@ private static VarBinArray buildScatteredVarBin(
                 readPos += len;
                 dataPos += len;
             }
-            offsets.setAtIndex(VortexFormat.LE_INT, i + 1, (int) dataPos);
+            setOffset(offsets, offsetsPtype, i + 1, dataPos);
         }
 
-        return new VarBinOffsetArray(dtype.withNullable(false), rowCount, values, offsets, PType.I32);
+        return new VarBinOffsetArray(dtype.withNullable(false), rowCount, values, offsets, offsetsPtype);
     }
 
     private static MemorySegment decompressFrames(
@@ -313,8 +313,8 @@ private static VarBinArray buildVarBin(DType dtype, long n, MemorySegment decomp
         }
 
         MemorySegment values = ctx.arena().allocate(totalDataBytes);
-        MemorySegment offsets = ctx.arena().allocate((n + 1) * 4L, 4);
-        offsets.setAtIndex(VortexFormat.LE_INT, 0, 0);
+        PType offsetsPtype = offsetsPtype(totalDataBytes);
+        MemorySegment offsets = ctx.arena().allocate((n + 1) * offsetsPtype.byteSize(), offsetsPtype.byteSize());
 
         // Second pass reads the same positions the first pass already bounds-checked via
         // readVarBinLen, so a raw get/copy here cannot overrun; values is sized to the validated
@@ -327,9 +327,22 @@ private static VarBinArray buildVarBin(DType dtype, long n, MemorySegment decomp
             MemorySegment.copy(decompressed, pos, values, dataPos, len);
             pos += len;
             dataPos += len;
-            offsets.setAtIndex(VortexFormat.LE_INT, i + 1, (int) dataPos);
+            setOffset(offsets, offsetsPtype, i + 1, dataPos);
         }
 
-        return new VarBinOffsetArray(dtype, n, values, offsets, PType.I32);
+        return new VarBinOffsetArray(dtype, n, values, offsets, offsetsPtype);
+    }
+
+    /// `I32` offsets, as before, unless the rows' bytes pass 2 GiB, which only `I64` reaches (#633).
+    private static PType offsetsPtype(long totalDataBytes) {
+        return totalDataBytes > Integer.MAX_VALUE ? PType.I64 : PType.I32;
+    }
+
+    private static void setOffset(MemorySegment offsets, PType ptype, long index, long value) {
+        if (ptype == PType.I64) {
+            offsets.setAtIndex(VortexFormat.LE_LONG, index, value);
+        } else {
+            offsets.setAtIndex(VortexFormat.LE_INT, index, (int) value);
+        }
     }
 }

reader/src/test/java/io/github/dfa1/vortex/reader/array/VarBinArrayTest.java

@@ -42,6 +42,48 @@ private static VarBinOffsetArray of(String... values) {
         return new VarBinOffsetArray(UTF8, values.length, bytes, offsetsSeg, PType.I32);
     }
 
+    /// `U32` offsets reach past 2 GiB once the data buffer does, as in Rust-written files (#633).
+    @Nested
+    class U32Offsets {
+
+        @Test
+        @SuppressWarnings("restricted") // reinterpret sizes a segment over no memory; never read
+        void getByteLength_offsetsPastTwoGibibytes_readUnsigned() {
+            // Given — offsets straddling 2^31; only the data buffer's size is read, so it is a
+            // 3 GiB view over no memory rather than an allocation (reading its bytes would crash
+            // the JVM: keep this to getByteLength)
+            long start = Integer.MAX_VALUE - 9L;
+            long end = Integer.MAX_VALUE + 11L;
+            ByteBuffer bb = ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN);
+            bb.putInt((int) start).putInt((int) end);
+            MemorySegment bytes = MemorySegment.NULL.reinterpret(3L << 30);
+            var sut = new VarBinOffsetArray(DTypes.BINARY, 1, bytes, MemorySegment.ofArray(bb.array()), PType.U32);
+
+            // When
+            int result = sut.getByteLength(0);
+
+            // Then
+            assertThat(result).isEqualTo(20);
+        }
+
+        @Test
+        void getByteLength_unsignedOffsetPastTheDataBuffer_throwsVortexException() {
+            // Given — 0xFFFFFFF0 used to read as negative; unsigned it is ~4 GiB, past a 16-byte buffer
+            ByteBuffer bb = ByteBuffer.allocate(12).order(ByteOrder.LITTLE_ENDIAN);
+            bb.putInt(0).putInt(0xFFFF_FFF0).putInt(0xFFFF_FFF0);
+            MemorySegment bytes = MemorySegment.ofArray(new byte[16]);
+            var sut = new VarBinOffsetArray(DTypes.BINARY, 2, bytes, MemorySegment.ofArray(bb.array()), PType.U32);
+
+            // When / Then — the row read and the row-range slice both reject it
+            assertThatThrownBy(() -> sut.getByteLength(0))
+                    .isInstanceOf(VortexException.class)
+                    .hasMessageContaining("out of range");
+            assertThatThrownBy(() -> sut.limited(1))
+                    .isInstanceOf(VortexException.class)
+                    .hasMessageContaining("out of range");
+        }
+    }
+
     @Nested
     class Regular {
 

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

@@ -1,7 +1,6 @@
 package io.github.dfa1.vortex.writer.encode;
 
 import io.github.dfa1.vortex.core.compute.Utf8Order;
-import io.github.dfa1.vortex.core.error.VortexException;
 import io.github.dfa1.vortex.core.model.DType;
 import io.github.dfa1.vortex.core.model.PType;
 import io.github.dfa1.vortex.core.model.EncodingId;
@@ -40,18 +39,34 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) {
         // values child of a masked/nullable layout, where validity carries nullity) contributes a
         // zero-length slot.
         VarBinBytes.Rows rows = VarBinBytes.toContiguous(data);
-        if (rows == null) {
-            throw new VortexException(EncodingId.VORTEX_VARBIN, "chunk exceeds 2 GB of row bytes");
-        }
-        int n = rows.count();
-        int totalBytes = rows.offsets()[n];
-
         Arena arena = ctx.arena();
-        MemorySegment bytesBuf = arena.allocate(totalBytes > 0 ? totalBytes : 1);
-        MemorySegment.copy(rows.bytes(), 0, bytesBuf, ValueLayout.JAVA_BYTE, 0, totalBytes);
-        MemorySegment offsetsBuf = arena.allocate((long) (n + 1) * Long.BYTES, Long.BYTES);
-        for (int i = 0; i <= n; i++) {
-            offsetsBuf.setAtIndex(VortexFormat.LE_LONG, i, rows.offsets()[i]);
+        MemorySegment bytesBuf;
+        MemorySegment offsetsBuf;
+        if (rows != null) {
+            int n = rows.count();
+            int totalBytes = rows.offsets()[n];
+            bytesBuf = arena.allocate(totalBytes > 0 ? totalBytes : 1);
+            MemorySegment.copy(rows.bytes(), 0, bytesBuf, ValueLayout.JAVA_BYTE, 0, totalBytes);
+            offsetsBuf = arena.allocate((long) (n + 1) * Long.BYTES, Long.BYTES);
+            for (int i = 0; i <= n; i++) {
+                offsetsBuf.setAtIndex(VortexFormat.LE_LONG, i, rows.offsets()[i]);
+            }
+        } else {
+            // Past 2 GB of row bytes no byte[] holds them contiguously, but the off-heap buffer and
+            // its i64 offsets do (#633): copy the rows straight into it
+            byte[][] perRow = VarBinBytes.toByteArrays(data);
+            long totalBytes = 0;
+            for (byte[] row : perRow) {
+                totalBytes += row.length;
+            }
+            bytesBuf = arena.allocate(totalBytes);
+            offsetsBuf = arena.allocate((long) (perRow.length + 1) * Long.BYTES, Long.BYTES);
+            long at = 0;
+            for (int i = 0; i < perRow.length; i++) {
+                MemorySegment.copy(perRow[i], 0, bytesBuf, ValueLayout.JAVA_BYTE, at, perRow[i].length);
+                at += perRow[i].length;
+                offsetsBuf.setAtIndex(VortexFormat.LE_LONG, i + 1, at);
+            }
         }
 
         byte[] metaBytes = new ProtoVarBinMetadata(io.github.dfa1.vortex.core.proto.ProtoPType.fromValue(PType.I64.ordinal())).encode();