feat: implement vortex.zstd encode (Primitive + Utf8/Binary)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: dfa1

TODO.md

@@ -102,7 +102,21 @@
 | `vortex.list`                | —                          | ❌       | ❌       | hard      | needs list array model; unblocks `list.vortex` |
 | `vortex.listview`            | —                          | ❌       | ❌       | hard      | unblocks `listview.vortex` |
 | `vortex.fixed_size_list`     | —                          | ❌       | ❌       | hard      | unblocks `fixed_size_list.vortex` |
-| `vortex.zstd`                | `ZstdEncoding`             | ✅       | ❌ stub  | —         | Primitive, Utf8, Binary (no dict, no nullable); uses airlift/aircompressor |
+| `vortex.zstd`                | `ZstdEncoding`             | ✅       | ✅       | —         | Primitive, Utf8, Binary (no dict, no nullable); uses airlift/aircompressor |
+
+### `vortex.zstd` known limitations
+
+- [ ] **Nullable arrays (decode)** — `ZstdEncoding.Decoder` throws when `node.children().length > 0` (validity child present).
+  Fix: if child[0] exists, decode it as a validity bitmap (Bool array). Zstd stores only the _valid_ values compactly — after decompressing, scatter them back into a full-length array using the validity positions. For strings, reconstruct null slots as zero-length entries or handle at the `VarBinArray` level.
+
+- [ ] **Dictionary support (decode)** — `ZstdEncoding.Decoder` throws when `metadata.dictionary_size != 0`.
+  Fix: when `dictionary_size > 0`, buffer[0] is the dictionary bytes and frames start at buffer[1]. Pass dict bytes to `new ZstdDecompressor()` — check if aircompressor supports dictionary-mode decompression; if not, switch to `com.github.luben:zstd-jni` for this path only (native, dictionary-trained zstd is hard to replicate in pure Java).
+
+- [ ] **Multi-frame encode** — `ZstdEncoding.Encoder` always produces a single frame for the whole array.
+  Fix: accept a `valuesPerFrame` parameter (default: all values in one frame). Split the raw byte buffer at frame boundaries (`valuesPerFrame * byteWidth`), compress each slice independently, emit one `ZstdFrameMetadata` per frame. Enables partial decompression during slice scans.
+
+- [ ] **Nullable arrays (encode)** — `ZstdEncoding.Encoder` has no null handling.
+  Fix: accept nullable input (e.g. `Integer[]` or a validity mask alongside the data array). Strip null positions before compression. Encode the validity bitmap as a Bool child (child[0]) in the `EncodeNode`. Mirrors what Rust does: only valid values go into the compressed payload.
 
 ### S3 Fixture Status (`v0.72.0/arrays/`)
 

core/src/main/java/io/github/dfa1/vortex/encoding/ZstdEncoding.java

@@ -1,6 +1,7 @@
 package io.github.dfa1.vortex.encoding;
 
 import com.google.protobuf.InvalidProtocolBufferException;
+import io.airlift.compress.zstd.ZstdCompressor;
 import io.airlift.compress.zstd.ZstdDecompressor;
 import io.github.dfa1.vortex.core.ArrayStats;
 import io.github.dfa1.vortex.core.DType;
@@ -17,11 +18,15 @@
 import io.github.dfa1.vortex.core.array.VarBinArray;
 import io.github.dfa1.vortex.proto.EncodingProtos;
 
+import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.ValueLayout;
 import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.List;
 
-/// Decoder for {@code vortex.zstd} — Zstandard-compressed columnar array.
+/// Encoder/decoder for {@code vortex.zstd} — Zstandard-compressed columnar array.
 ///
 /// <p>Wire format:
 /// <ul>
@@ -30,27 +35,152 @@
 ///   <li>Child 0: validity bitmap (optional; this decoder rejects nullable arrays)</li>
 /// </ul>
 ///
-/// <p>Primitive dtype: decompressed bytes are raw LE values → returned as typed primitive array.
-/// <p>Utf8/Binary dtype: decompressed bytes interleave {@code [u32-LE length][data]} per valid string.
+/// <p>Primitive dtype: raw LE values compressed into a single frame.
+/// <p>Utf8/Binary dtype: {@code [u32-LE length][data]} per string, compressed into a single frame.
 ///
-/// <p>Scope: decode only, no dictionary, non-nullable.
+/// <p>No dictionary, no nullable (validity child not supported).
 public final class ZstdEncoding implements Encoding {
 
     @Override
     public EncodingId encodingId() {
         return EncodingId.VORTEX_ZSTD;
     }
 
+    @Override
+    public boolean accepts(DType dtype) {
+        return dtype instanceof DType.Primitive || dtype instanceof DType.Utf8 || dtype instanceof DType.Binary;
+    }
+
     @Override
     public EncodeResult encode(DType dtype, Object data) {
-        throw new UnsupportedOperationException("encode not supported by " + encodingId());
+        return Encoder.encode(dtype, data);
     }
 
     @Override
     public Array decode(DecodeContext ctx) {
         return Decoder.decode(ctx);
     }
 
+    private static final class Encoder {
+
+        private static EncodeResult encode(DType dtype, Object data) {
+            if (dtype instanceof DType.Primitive dt) {
+                return encodePrimitive(dt, data);
+            }
+            if (dtype instanceof DType.Utf8 || dtype instanceof DType.Binary) {
+                return encodeVarBin(dtype, (String[]) data);
+            }
+            throw new VortexException(EncodingId.VORTEX_ZSTD, "unsupported dtype: " + dtype);
+        }
+
+        private static EncodeResult encodePrimitive(DType.Primitive dt, Object data) {
+            MemorySegment raw = primitiveToLeBytes(dt.ptype(), data);
+            long n = primitiveLength(dt.ptype(), data);
+            byte[] rawBytes = raw.toArray(ValueLayout.JAVA_BYTE);
+            return buildResult(rawBytes, n);
+        }
+
+        private static EncodeResult encodeVarBin(DType dtype, String[] strings) {
+            byte[] raw = buildLengthPrefixed(strings);
+            return buildResult(raw, strings.length);
+        }
+
+        private static EncodeResult buildResult(byte[] raw, long n) {
+            byte[] compressed = compress(raw);
+            byte[] meta = EncodingProtos.ZstdMetadata.newBuilder()
+                    .setDictionarySize(0)
+                    .addFrames(EncodingProtos.ZstdFrameMetadata.newBuilder()
+                            .setUncompressedSize(raw.length)
+                            .setNValues(n))
+                    .build().toByteArray();
+            EncodeNode root = new EncodeNode(EncodingId.VORTEX_ZSTD, ByteBuffer.wrap(meta),
+                    new EncodeNode[0], new int[]{0});
+            return new EncodeResult(root, List.of(MemorySegment.ofArray(compressed)), null, null);
+        }
+
+        private static byte[] compress(byte[] input) {
+            ZstdCompressor compressor = new ZstdCompressor();
+            byte[] out = new byte[compressor.maxCompressedLength(input.length)];
+            int len = compressor.compress(input, 0, input.length, out, 0, out.length);
+            return Arrays.copyOf(out, len);
+        }
+
+        private static MemorySegment primitiveToLeBytes(PType ptype, Object data) {
+            return switch (ptype) {
+                case I8, U8 -> MemorySegment.ofArray((byte[]) data);
+                case I16, U16, F16 -> {
+                    short[] arr = (short[]) data;
+                    MemorySegment seg = Arena.ofAuto().allocate((long) arr.length * 2, 2);
+                    for (int i = 0; i < arr.length; i++) {
+                        seg.setAtIndex(PTypeIO.LE_SHORT, i, arr[i]);
+                    }
+                    yield seg;
+                }
+                case I32, U32 -> {
+                    int[] arr = (int[]) data;
+                    MemorySegment seg = Arena.ofAuto().allocate((long) arr.length * 4, 4);
+                    for (int i = 0; i < arr.length; i++) {
+                        seg.setAtIndex(PTypeIO.LE_INT, i, arr[i]);
+                    }
+                    yield seg;
+                }
+                case I64, U64 -> {
+                    long[] arr = (long[]) data;
+                    MemorySegment seg = Arena.ofAuto().allocate((long) arr.length * 8, 8);
+                    for (int i = 0; i < arr.length; i++) {
+                        seg.setAtIndex(PTypeIO.LE_LONG, i, arr[i]);
+                    }
+                    yield seg;
+                }
+                case F32 -> {
+                    float[] arr = (float[]) data;
+                    MemorySegment seg = Arena.ofAuto().allocate((long) arr.length * 4, 4);
+                    for (int i = 0; i < arr.length; i++) {
+                        seg.setAtIndex(PTypeIO.LE_FLOAT, i, arr[i]);
+                    }
+                    yield seg;
+                }
+                case F64 -> {
+                    double[] arr = (double[]) data;
+                    MemorySegment seg = Arena.ofAuto().allocate((long) arr.length * 8, 8);
+                    for (int i = 0; i < arr.length; i++) {
+                        seg.setAtIndex(PTypeIO.LE_DOUBLE, i, arr[i]);
+                    }
+                    yield seg;
+                }
+            };
+        }
+
+        private static long primitiveLength(PType ptype, Object data) {
+            return switch (ptype) {
+                case I8, U8 -> ((byte[]) data).length;
+                case I16, U16, F16 -> ((short[]) data).length;
+                case I32, U32 -> ((int[]) data).length;
+                case F32 -> ((float[]) data).length;
+                case I64, U64 -> ((long[]) data).length;
+                case F64 -> ((double[]) data).length;
+            };
+        }
+
+        private static byte[] buildLengthPrefixed(String[] strings) {
+            int total = 0;
+            byte[][] encoded = new byte[strings.length][];
+            for (int i = 0; i < strings.length; i++) {
+                encoded[i] = strings[i].getBytes(StandardCharsets.UTF_8);
+                total += 4 + encoded[i].length;
+            }
+            MemorySegment seg = Arena.ofAuto().allocate(total > 0 ? total : 1);
+            long pos = 0;
+            for (byte[] bytes : encoded) {
+                seg.set(PTypeIO.LE_INT, pos, bytes.length);
+                pos += 4;
+                MemorySegment.copy(MemorySegment.ofArray(bytes), 0, seg, pos, bytes.length);
+                pos += bytes.length;
+            }
+            return seg.asSlice(0, total).toArray(ValueLayout.JAVA_BYTE);
+        }
+    }
+
     private static final class Decoder {
 
         private static Array decode(DecodeContext ctx) {

core/src/test/java/io/github/dfa1/vortex/encoding/ZstdEncodingTest.java

@@ -31,13 +31,82 @@ class ZstdEncodingTest {
     class Encode {
 
         @Test
-        void encode_throwsUnsupportedOperationException() {
+        void encode_i32_roundTrips() {
+            // Given
+            var sut = new ZstdEncoding();
+            int[] data = {10, 20, 30, 40};
+
+            // When
+            EncodeResult result = sut.encode(I32, data);
+            DecodeContext ctx = EncodeTestHelper.toDecodeContext(result, data.length, I32, EncodingRegistry.empty());
+            IntArray decoded = (IntArray) sut.decode(ctx);
+
+            // Then
+            assertThat(decoded.length()).isEqualTo(data.length);
+            for (int i = 0; i < data.length; i++) {
+                assertThat(decoded.getInt(i)).as("index %d", i).isEqualTo(data[i]);
+            }
+        }
+
+        @Test
+        void encode_i64_roundTrips() {
+            // Given
+            var sut = new ZstdEncoding();
+            long[] data = {100L, 200L, 300L};
+
+            // When
+            EncodeResult result = sut.encode(I64, data);
+            DecodeContext ctx = EncodeTestHelper.toDecodeContext(result, data.length, I64, EncodingRegistry.empty());
+            LongArray decoded = (LongArray) sut.decode(ctx);
+
+            // Then
+            assertThat(decoded.length()).isEqualTo(data.length);
+            for (int i = 0; i < data.length; i++) {
+                assertThat(decoded.getLong(i)).as("index %d", i).isEqualTo(data[i]);
+            }
+        }
+
+        @Test
+        void encode_utf8_roundTrips() {
+            // Given
+            var sut = new ZstdEncoding();
+            String[] data = {"hello", "world", "zstd"};
+
+            // When
+            EncodeResult result = sut.encode(UTF8, data);
+            DecodeContext ctx = EncodeTestHelper.toDecodeContext(result, data.length, UTF8, EncodingRegistry.empty());
+            VarBinArray decoded = (VarBinArray) sut.decode(ctx);
+
+            // Then
+            assertThat(decoded.length()).isEqualTo(data.length);
+            for (int i = 0; i < data.length; i++) {
+                assertThat(decoded.getString(i)).as("index %d", i).isEqualTo(data[i]);
+            }
+        }
+
+        @Test
+        void encode_emptyArray_roundTrips() {
+            // Given
+            var sut = new ZstdEncoding();
+            int[] data = {};
+
+            // When
+            EncodeResult result = sut.encode(I32, data);
+            DecodeContext ctx = EncodeTestHelper.toDecodeContext(result, data.length, I32, EncodingRegistry.empty());
+            IntArray decoded = (IntArray) sut.decode(ctx);
+
+            // Then
+            assertThat(decoded.length()).isZero();
+        }
+
+        @Test
+        void encode_unsupportedDtype_throwsVortexException() {
             // Given
             var sut = new ZstdEncoding();
 
             // When / Then
-            assertThatThrownBy(() -> sut.encode(I32, new int[]{1, 2, 3}))
-                    .isInstanceOf(UnsupportedOperationException.class);
+            assertThatThrownBy(() -> sut.encode(new DType.Null(false), null))
+                    .isInstanceOf(VortexException.class);
         }
     }