diff --git a/CHANGELOG.md b/CHANGELOG.md index adc7ae727..2493b968e 100644 --- a/CHANGELOG.md +++ b/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)) diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/FsstLargeChunkInteropIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/LargeChunkInteropIntegrationTest.java similarity index 64% rename from integration/src/test/java/io/github/dfa1/vortex/integration/FsstLargeChunkInteropIntegrationTest.java rename to integration/src/test/java/io/github/dfa1/vortex/integration/LargeChunkInteropIntegrationTest.java index c08a05bd6..1ac477c47 100644 --- a/integration/src/test/java/io/github/dfa1/vortex/integration/FsstLargeChunkInteropIntegrationTest.java +++ b/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 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 readTailWithJni(Path file, long begin, long end) throws IOException { ScanOptions opts = ScanOptions.builder() .projection(Expression.select(new String[]{"blob"}, Expression.root())) diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/array/VarBinOffsetArray.java b/reader/src/main/java/io/github/dfa1/vortex/reader/array/VarBinOffsetArray.java index 27b1683d5..45cbe496b 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/array/VarBinOffsetArray.java +++ b/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); diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/decode/ZstdEncodingDecoder.java b/reader/src/main/java/io/github/dfa1/vortex/reader/decode/ZstdEncodingDecoder.java index b17211e02..b00e0177b 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/decode/ZstdEncodingDecoder.java +++ b/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); + } } } diff --git a/reader/src/test/java/io/github/dfa1/vortex/reader/array/VarBinArrayTest.java b/reader/src/test/java/io/github/dfa1/vortex/reader/array/VarBinArrayTest.java index afed2a6c8..7f5a22e09 100644 --- a/reader/src/test/java/io/github/dfa1/vortex/reader/array/VarBinArrayTest.java +++ b/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 { diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoder.java index 1d3617274..a19d1aa02 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinEncodingEncoder.java +++ b/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(); diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinViewEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinViewEncodingEncoder.java index c9f769df5..def394904 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinViewEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/VarBinViewEncodingEncoder.java @@ -6,6 +6,7 @@ import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; +import java.util.ArrayList; import java.util.List; /// Write-only encoder for `vortex.varbinview`. @@ -13,9 +14,22 @@ public final class VarBinViewEncodingEncoder implements EncodingEncoder { private static final int MAX_INLINED_SIZE = 12; private static final int VIEW_SIZE = 16; + /// Largest data buffer: a view's offset into it stays a non-negative `int`, Rust's + /// `MAX_BUFFER_LEN` (`i32::MAX`). + private static final long MAX_DATA_BUFFER_BYTES = Integer.MAX_VALUE; + + private final long maxDataBufferBytes; /// Creates the encoder. public VarBinViewEncodingEncoder() { + this(MAX_DATA_BUFFER_BYTES); + } + + /// Creates the encoder with a smaller data buffer cap, so tests reach the split without 2 GiB. + /// + /// @param maxDataBufferBytes the largest data buffer + VarBinViewEncodingEncoder(long maxDataBufferBytes) { + this.maxDataBufferBytes = maxDataBufferBytes; } @Override @@ -33,19 +47,33 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { byte[][] bytes = VarBinBytes.toByteArrays(data); int n = bytes.length; - int totalDataBytes = 0; + // Out-of-line rows fill data buffers in order; a buffer is closed before a row would take + // it past 2 GiB, the most a view's offset reaches here, so a chunk of any size fits (#633) + List bufferSizes = new ArrayList<>(); + long current = 0; for (int i = 0; i < n; i++) { - if (bytes[i].length > MAX_INLINED_SIZE) { - totalDataBytes += bytes[i].length; + int len = bytes[i].length; + if (len > MAX_INLINED_SIZE) { + if (current > 0 && current + len > maxDataBufferBytes) { + bufferSizes.add(current); + current = 0; + } + current += len; } } + if (current > 0) { + bufferSizes.add(current); + } Arena arena = ctx.arena(); - boolean hasDataBuf = totalDataBytes > 0; - MemorySegment dataBuf = arena.allocate(hasDataBuf ? totalDataBytes : 1); + MemorySegment[] dataBufs = new MemorySegment[bufferSizes.size()]; + for (int b = 0; b < dataBufs.length; b++) { + dataBufs[b] = arena.allocate(bufferSizes.get(b)); + } MemorySegment viewsBuf = arena.allocate(n > 0 ? (long) n * VIEW_SIZE : 1); - int dataOffset = 0; + int bufferIndex = 0; + long dataOffset = 0; for (int i = 0; i < n; i++) { byte[] b = bytes[i]; long viewOff = (long) i * VIEW_SIZE; @@ -53,29 +81,33 @@ public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) { if (b.length <= MAX_INLINED_SIZE) { MemorySegment.copy(MemorySegment.ofArray(b), 0, viewsBuf, viewOff + 4, b.length); } else { + if (dataOffset + b.length > dataBufs[bufferIndex].byteSize()) { + bufferIndex++; + dataOffset = 0; + } MemorySegment.copy(MemorySegment.ofArray(b), 0, viewsBuf, viewOff + 4, 4); - viewsBuf.set(VortexFormat.LE_INT, viewOff + 8, 0); - viewsBuf.set(VortexFormat.LE_INT, viewOff + 12, dataOffset); - MemorySegment.copy(MemorySegment.ofArray(b), 0, dataBuf, dataOffset, b.length); + viewsBuf.set(VortexFormat.LE_INT, viewOff + 8, bufferIndex); + viewsBuf.set(VortexFormat.LE_INT, viewOff + 12, (int) dataOffset); + MemorySegment.copy(MemorySegment.ofArray(b), 0, dataBufs[bufferIndex], dataOffset, b.length); dataOffset += b.length; } } - int[] bufIndices; - List buffers; - // Views are Rust's `#[repr(C, align(16))] BinaryView`; the data buffers are plain bytes. - EncodedBuffer views = new EncodedBuffer(viewsBuf, VIEW_SIZE); - if (hasDataBuf) { - bufIndices = new int[]{0, 1}; - buffers = List.of(EncodedBuffer.bytes(dataBuf), views); - } else { - bufIndices = new int[]{0}; - buffers = List.of(views); + // Views are Rust's `#[repr(C, align(16))] BinaryView`; the data buffers are plain bytes, + // listed before the views. + List buffers = new ArrayList<>(dataBufs.length + 1); + for (MemorySegment dataBuf : dataBufs) { + buffers.add(EncodedBuffer.bytes(dataBuf)); + } + buffers.add(new EncodedBuffer(viewsBuf, VIEW_SIZE)); + int[] bufIndices = new int[buffers.size()]; + for (int b = 0; b < bufIndices.length; b++) { + bufIndices[b] = b; } EncodeNode root = new EncodeNode(EncodingId.VORTEX_VARBINVIEW, null, new EncodeNode[0], bufIndices); // Same lexicographic bounds vortex.varbin reports for the same column: which of the two // string layouts the compressor picks must not decide whether the column can be pruned. - return new EncodeResult(root, buffers, null, null).withStats(ZoneMapStats.of(dtype, data)); + return new EncodeResult(root, List.copyOf(buffers), null, null).withStats(ZoneMapStats.of(dtype, data)); } } diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZstdEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZstdEncodingEncoder.java index b6d2a4c45..943e6c822 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZstdEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ZstdEncodingEncoder.java @@ -59,6 +59,12 @@ static void requireBinding(boolean present) { } } + /// Largest uncompressed payload of one varbin frame: half the 2 GiB a reader takes per frame. + private static final long MAX_FRAME_BYTES = 1L << 30; + + /// The varbin frame cap in use: [#MAX_FRAME_BYTES], smaller in tests. + private final long maxFrameBytes; + /// Values per zstd frame; `0` (or any non-positive value) means a single frame for the whole array. private final long valuesPerFrame; @@ -87,8 +93,18 @@ public ZstdEncodingEncoder(long valuesPerFrame) { } private ZstdEncodingEncoder(long valuesPerFrame, boolean textOnly) { + this(valuesPerFrame, textOnly, MAX_FRAME_BYTES); + } + + /// Creates an encoder with a smaller varbin frame cap, so tests reach the split without 1 GiB. + /// + /// @param valuesPerFrame values per frame; non-positive means no count limit + /// @param textOnly whether only strings and binary are accepted + /// @param maxFrameBytes the largest uncompressed varbin frame payload + ZstdEncodingEncoder(long valuesPerFrame, boolean textOnly, long maxFrameBytes) { this.valuesPerFrame = valuesPerFrame; this.textOnly = textOnly; + this.maxFrameBytes = maxFrameBytes; } @Override @@ -276,27 +292,38 @@ private FrameLayout uniformLayout(long n, int byteWidth) { /// Frame layout for a length-prefixed varbin payload: `valuesPerFrame` values per frame, with /// each frame's byte span found by walking the 4-byte length prefixes to a value boundary. + /// + /// A frame is also closed before its payload passes [#MAX_FRAME_BYTES] (#633): the reader + /// caps a frame's declared size at 2 GiB until ADR 0004's resource limits exist, so a frame this + /// writer emits stays readable. A payload under the cap keeps the plain layout. A single value + /// larger than the cap is a frame of its own; past 2 GiB it would not read back, but no Java + /// array holds such a row. private FrameLayout varBinLayout(MemorySegment raw, long nValues) { - if (valuesPerFrame <= 0 || nValues <= valuesPerFrame) { + boolean fitsOneFrame = valuesPerFrame <= 0 || nValues <= valuesPerFrame; + if (fitsOneFrame && raw.byteSize() <= maxFrameBytes) { return new FrameLayout(new long[]{raw.byteSize()}, new long[]{nValues}); } - int frameCount = (int) ((nValues + valuesPerFrame - 1) / valuesPerFrame); - long[] byteLengths = new long[frameCount]; - long[] valueCounts = new long[frameCount]; + List byteLengths = new ArrayList<>(); + List valueCounts = new ArrayList<>(); long pos = 0; - long valueIdx = 0; - for (int f = 0; f < frameCount; f++) { - long count = Math.min(valuesPerFrame, nValues - valueIdx); - long start = pos; - for (long k = 0; k < count; k++) { - int len = raw.get(VortexFormat.LE_INT, pos); - pos += 4L + len; + long frameStart = 0; + long count = 0; + for (long v = 0; v < nValues; v++) { + long valueBytes = 4L + raw.get(VortexFormat.LE_INT, pos); + boolean full = valuesPerFrame > 0 && count == valuesPerFrame; + if (count > 0 && (full || pos - frameStart + valueBytes > maxFrameBytes)) { + byteLengths.add(pos - frameStart); + valueCounts.add(count); + frameStart = pos; + count = 0; } - byteLengths[f] = pos - start; - valueCounts[f] = count; - valueIdx += count; + pos += valueBytes; + count++; } - return new FrameLayout(byteLengths, valueCounts); + byteLengths.add(pos - frameStart); + valueCounts.add(count); + return new FrameLayout(byteLengths.stream().mapToLong(Long::longValue).toArray(), + valueCounts.stream().mapToLong(Long::longValue).toArray()); } private static Frames compressFrames(MemorySegment raw, FrameLayout layout, Arena arena) { @@ -371,7 +398,8 @@ private static int countNonNull(byte[][] rows) { } private static MemorySegment buildLengthPrefixed(byte[][] encoded, Arena arena) { - int total = 0; + // long: a chunk's rows may pass 2 GiB, which an int total wraps negative (#633) + long total = 0; for (byte[] bytes : encoded) { total += 4 + bytes.length; } diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/VarBinViewEncodingEncoderTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/VarBinViewEncodingEncoderTest.java index 47a3ff723..1aa6dc9c2 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/VarBinViewEncodingEncoderTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/VarBinViewEncodingEncoderTest.java @@ -10,6 +10,8 @@ import io.github.dfa1.vortex.reader.ReadRegistry; import io.github.dfa1.vortex.reader.decode.TestRegistry; import io.github.dfa1.vortex.reader.decode.VarBinViewEncodingDecoder; +import java.util.List; +import java.util.ArrayList; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -202,4 +204,79 @@ void decode_roundtrip_returnsAllStrings(String name, String[] values) { } } } + + /// Out-of-line rows fill data buffers up to a cap (`Integer.MAX_VALUE` by default, Rust's + /// `MAX_BUFFER_LEN`), so a chunk past 2 GiB still addresses every row (#633). A 26-byte cap + /// reaches the split here. + @Nested + class DataBufferCap { + + private final VarBinViewEncodingEncoder sut = new VarBinViewEncodingEncoder(26); + + @Test + void rowsPastTheCap_openANewBufferAndRoundTrip() { + // Given — 13-byte rows (out of line), an inlined row between them: a + b fill the first + // buffer exactly, c starts the second at offset 0 + String a = "aaaaaaaaaaaaa"; + String b = "bbbbbbbbbbbbb"; + String c = "ccccccccccccc"; + String[] values = {a, "short", b, c}; + + // When + EncodeResult result = sut.encode(DTypes.UTF8, values, EncodeTestHelper.testCtx()); + + // Then — two data buffers then the views, and every row decodes + List buffers = result.buffers(); + assertThat(buffers).hasSize(3); + assertThat(buffers.get(0).byteSize()).isEqualTo(26); + assertThat(buffers.get(1).byteSize()).isEqualTo(13); + MemorySegment views = buffers.get(2); + assertThat(views.get(VortexFormat.LE_INT, 3 * 16L + 8)).as("c's buffer index").isEqualTo(1); + assertThat(views.get(VortexFormat.LE_INT, 3 * 16L + 12)).as("c's offset").isZero(); + assertThat(decode(result, values.length)).containsExactly(values); + } + + @Test + void rowLargerThanTheCap_takesABufferOfItsOwn() { + // Given + String big = "x".repeat(40); + String[] values = {"aaaaaaaaaaaaa", big}; + + // When + EncodeResult result = sut.encode(DTypes.UTF8, values, EncodeTestHelper.testCtx()); + + // Then + assertThat(result.buffers()).hasSize(3); + assertThat(result.buffers().get(1).byteSize()).isEqualTo(40); + assertThat(decode(result, values.length)).containsExactly(values); + } + + @Test + void allInlinedOrEmpty_writesNoDataBuffer() { + // Given + String[] inlined = {"a", "bc", ""}; + + // When + EncodeResult inlinedResult = sut.encode(DTypes.UTF8, inlined, EncodeTestHelper.testCtx()); + EncodeResult emptyResult = sut.encode(DTypes.UTF8, new String[0], EncodeTestHelper.testCtx()); + + // Then — only the views buffer + assertThat(inlinedResult.buffers()).hasSize(1); + assertThat(decode(inlinedResult, inlined.length)).containsExactly(inlined); + assertThat(emptyResult.buffers()).hasSize(1); + } + + private static List decode(EncodeResult result, int n) { + MemorySegment[] bufs = result.buffers().toArray(MemorySegment[]::new); + ArrayNode node = new ArrayNode(EncodingId.VORTEX_VARBINVIEW, null, new ArrayNode[0], + result.rootNode().bufferIndices()); + DecodeContext ctx = new DecodeContext(node, DTypes.UTF8, n, bufs, REGISTRY, Arena.ofAuto()); + VarBinArray decoded = (VarBinArray) DECODER.decode(ctx); + List out = new ArrayList<>(); + for (int i = 0; i < n; i++) { + out.add(decoded.getString(i)); + } + return out; + } + } } diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/ZstdEncodingEncoderTest.java b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/ZstdEncodingEncoderTest.java index 2fbae4225..c9df4271a 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/ZstdEncodingEncoderTest.java +++ b/writer/src/test/java/io/github/dfa1/vortex/writer/encode/ZstdEncodingEncoderTest.java @@ -676,4 +676,74 @@ void encode_nullablePrimitive_framesOverValidValues_andRoundTrips() throws Excep assertThat(child.getInt(6)).isEqualTo(50); } } + + /// A varbin frame closes before its payload passes a cap (1 GiB by default), so every frame stays + /// within the 2 GiB a reader takes (#633). Each 10-byte value is a 14-byte length-prefixed + /// payload; small caps reach the split here. + @Nested + class FrameByteCap { + + private static final String[] VALUES = {"aaaaaaaaaa", "bbbbbbbbbb", "cccccccccc", "dddddddddd", "eeeeeeeeee"}; + + @Test + void payloadPastTheCap_splitsOnValueBoundariesAndRoundTrips() throws Exception { + // Given — 40 bytes hold two values (28), not three (42) + var sut = new ZstdEncodingEncoder(0, true, 40); + + // When + EncodeResult result = sut.encode(DTypes.UTF8, VALUES, EncodeTestHelper.testCtx()); + + // Then + assertThat(frameValueCounts(result)).containsExactly(2L, 2L, 1L); + assertThat(roundTrip(result)).containsExactly(VALUES); + } + + @Test + void valuesPerFrame_andByteCap_closeAtWhicheverComesFirst() throws Exception { + // Given — at most 3 values a frame, at most 40 bytes: the byte cap closes first + var sut = new ZstdEncodingEncoder(3, true, 40); + // at most 2 values a frame, a cap no frame reaches: the count closes first + var byCount = new ZstdEncodingEncoder(2, true, 1_000); + + // When + EncodeResult result = sut.encode(DTypes.UTF8, VALUES, EncodeTestHelper.testCtx()); + EncodeResult countResult = byCount.encode(DTypes.UTF8, VALUES, EncodeTestHelper.testCtx()); + + // Then + assertThat(frameValueCounts(result)).containsExactly(2L, 2L, 1L); + assertThat(frameValueCounts(countResult)).containsExactly(2L, 2L, 1L); + assertThat(roundTrip(result)).containsExactly(VALUES); + } + + @Test + void valueLargerThanTheCap_isAFrameOfItsOwn() throws Exception { + // Given + var sut = new ZstdEncodingEncoder(0, true, 10); + String[] values = {"x".repeat(20), "y"}; + + // When + EncodeResult result = sut.encode(DTypes.UTF8, values, EncodeTestHelper.testCtx()); + + // Then + assertThat(frameValueCounts(result)).containsExactly(1L, 1L); + assertThat(roundTrip(result)).containsExactly(values); + } + + private static List frameValueCounts(EncodeResult result) throws Exception { + var metaSeg = result.rootNode().metadata(); + return ProtoZstdMetadata.decode(metaSeg, 0, metaSeg.byteSize()).frames().stream() + .map(ProtoZstdFrameMetadata::n_values).toList(); + } + + private static List roundTrip(EncodeResult result) throws Exception { + int n = Math.toIntExact(frameValueCounts(result).stream().mapToLong(Long::longValue).sum()); + DecodeContext ctx = DecodeTestHelper.toDecodeContext(result, n, DTypes.UTF8, ReadRegistry.empty()); + VarBinArray decoded = (VarBinArray) DECODER.decode(ctx); + List out = new ArrayList<>(); + for (int i = 0; i < n; i++) { + out.add(decoded.getString(i)); + } + return out; + } + } }