feat(variant): decode variant columns Java-side (constant + chunked)

ConstantEncodingDecoder now handles DType.Variant by unwrapping the variant
scalar to its typed inner scalar and materialising the inner-typed constant
array; the dispatch was refactored to thread the dtype explicitly so Extension
and Variant can recurse without re-reading the buffer. ChunkedEncodingDecoder
wraps Variant chunks under their inner dtype. VariantEncodingDecoder.accepts now
reports Variant.

Together with VariantEncodingDecoder this makes the Java reader symmetric with
the writer: a Java-encoded variant column round-trips through Java decode, with
core storage exposing each row's inner value (verified for constant and
row-varying columns).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: dfa1

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

@@ -86,6 +86,23 @@ private static Array wrap(List<Array> chunks, DType dtype, long totalRows) {
         if (dtype instanceof DType.Struct struct) {
             return wrapStruct(chunks, struct, totalRows);
         }
+        if (dtype instanceof DType.Variant) {
+            // Each chunk decoded as Variant materialises to its inner-typed constant array
+            // (see ConstantEncodingDecoder). Wrap the chunks under that inner dtype; the
+            // VariantArray container re-applies the logical Variant dtype.
+            if (chunks.isEmpty()) {
+                throw new VortexException(EncodingId.VORTEX_CHUNKED, "chunked variant has no chunks");
+            }
+            DType innerDtype = chunks.get(0).dtype();
+            if (innerDtype instanceof DType.Primitive innerPt) {
+                return wrapPrimitive(chunks, innerPt, innerDtype, totalRows);
+            }
+            if (innerDtype instanceof DType.Bool) {
+                return ChunkedBoolArray.of(innerDtype, totalRows, chunks);
+            }
+            throw new VortexException(EncodingId.VORTEX_CHUNKED,
+                    "chunked variant inner dtype not supported: " + innerDtype);
+        }
         throw new VortexException(EncodingId.VORTEX_CHUNKED,
                 "chunked not supported for dtype: " + dtype);
     }

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

@@ -52,69 +52,77 @@ public Array decode(DecodeContext ctx) {
             throw new VortexException(EncodingId.VORTEX_CONSTANT, "invalid scalar value", e);
         }
 
-        long n = ctx.rowCount();
+        return arrayFromScalar(ctx, scalar, ctx.dtype(), ctx.rowCount());
+    }
 
-        if (ctx.dtype() instanceof DType.Null) {
-            return new NullArray(ctx.dtype(), n);
+    /// Builds the constant array for `scalar` interpreted as `dtype`, broadcast to `n` rows.
+    /// Recurses for Extension (its storage dtype) and Variant (the wrapped inner scalar).
+    private static Array arrayFromScalar(DecodeContext ctx, ScalarValue scalar, DType dtype, long n) {
+        if (dtype instanceof DType.Null) {
+            return new NullArray(dtype, n);
         }
-
-        if (ctx.dtype() instanceof DType.Utf8 || ctx.dtype() instanceof DType.Binary) {
-            return decodeString(ctx, scalar, n);
+        if (dtype instanceof DType.Variant) {
+            // A constant variant wraps a typed inner scalar (Scalar::variant(inner)); the
+            // physical storage is the inner-typed constant array. The VariantArray wrapper
+            // re-applies the logical Variant dtype.
+            io.github.dfa1.vortex.proto.Scalar inner = scalar.variant_value();
+            if (inner == null || inner.value() == null) {
+                throw new VortexException(EncodingId.VORTEX_CONSTANT, "constant variant missing variant_value");
+            }
+            DType innerDtype = VariantEncodingDecoder.dtypeFromProto(inner.dtype());
+            return arrayFromScalar(ctx, inner.value(), innerDtype, n);
         }
-
-        if (ctx.dtype() instanceof DType.Bool) {
-            return decodeBool(ctx, scalar, n);
+        if (dtype instanceof DType.Utf8 || dtype instanceof DType.Binary) {
+            return decodeString(ctx, scalar, dtype, n);
         }
-
-        if (ctx.dtype() instanceof DType.Decimal) {
-            return decodeDecimal(ctx, scalar, n);
+        if (dtype instanceof DType.Bool) {
+            return decodeBool(dtype, scalar, n);
         }
-
-        if (ctx.dtype() instanceof DType.Extension ext) {
-            var storageCtx = new DecodeContext(ctx.node(), ext.storageDType(), ctx.rowCount(),
-                    ctx.segmentBuffers(), ctx.registry(), ctx.arena());
-            Array storage = decode(storageCtx);
-            // GenericArray needs a backing buffer; the recursive call now returns a
-            // metadata-only LazyConstantXxxArray. Materialise once into the chunk arena
-            // so downstream extension consumers that read via ArraySegments.of(arr) still
-            // find a segment. Extension-on-constant is rare enough that the small alloc
-            // doesn't matter — the bare primitive path stays buffer-free.
-            return new GenericArray(ctx.dtype(), n, ArraySegments.of(storage, ctx.arena()));
+        if (dtype instanceof DType.Decimal) {
+            return decodeDecimal(dtype, scalar, n);
         }
-
-        if (!(ctx.dtype() instanceof DType.Primitive p)) {
-            throw new VortexException(EncodingId.VORTEX_CONSTANT, "unsupported dtype " + ctx.dtype());
+        if (dtype instanceof DType.Extension ext) {
+            Array storage = arrayFromScalar(ctx, scalar, ext.storageDType(), n);
+            // GenericArray needs a backing buffer; the recursive call returns a metadata-only
+            // LazyConstantXxxArray. Materialise once into the chunk arena so downstream
+            // extension consumers that read via ArraySegments.of(arr) still find a segment.
+            // Extension-on-constant is rare enough that the small alloc doesn't matter — the
+            // bare primitive path stays buffer-free.
+            return new GenericArray(dtype, n, ArraySegments.of(storage, ctx.arena()));
+        }
+        if (!(dtype instanceof DType.Primitive p)) {
+            throw new VortexException(EncodingId.VORTEX_CONSTANT, "unsupported dtype " + dtype);
         }
 
         PType ptype = p.ptype();
         long rawBits = scalarToRawBits(scalar, ptype);
 
         return switch (ptype) {
-            case I64, U64 -> new LazyConstantLongArray(ctx.dtype(), n, rawBits);
-            case I32, U32 -> new LazyConstantIntArray(ctx.dtype(), n, (int) rawBits);
-            case F64 -> new LazyConstantDoubleArray(ctx.dtype(), n, Double.longBitsToDouble(rawBits));
-            case F32 -> new LazyConstantFloatArray(ctx.dtype(), n, Float.intBitsToFloat((int) rawBits));
-            case I16, U16 -> new LazyConstantShortArray(ctx.dtype(), n, (short) rawBits);
-            case I8, U8 -> new LazyConstantByteArray(ctx.dtype(), n, (byte) rawBits);
+            case I64, U64 -> new LazyConstantLongArray(dtype, n, rawBits);
+            case I32, U32 -> new LazyConstantIntArray(dtype, n, (int) rawBits);
+            case F64 -> new LazyConstantDoubleArray(dtype, n, Double.longBitsToDouble(rawBits));
+            case F32 -> new LazyConstantFloatArray(dtype, n, Float.intBitsToFloat((int) rawBits));
+            case I16, U16 -> new LazyConstantShortArray(dtype, n, (short) rawBits);
+            case I8, U8 -> new LazyConstantByteArray(dtype, n, (byte) rawBits);
             default -> throw new VortexException(EncodingId.VORTEX_CONSTANT, "unsupported ptype " + ptype);
         };
     }
 
-    private static Array decodeDecimal(DecodeContext ctx, ScalarValue scalar, long n) {
+    private static Array decodeDecimal(DType dtype, ScalarValue scalar, long n) {
         byte[] elemBytes = scalar.bytes_value();
         int elemLen = elemBytes.length;
         // Decode the single scalar value via LazyDecimalArray (reuses its LE byte-order logic),
         // then wrap in a constant array — O(1) allocation regardless of row count.
-        var value = new LazyDecimalArray(ctx.dtype(), 1, MemorySegment.ofArray(elemBytes), elemLen).getDecimal(0);
-        return new LazyConstantDecimalArray(ctx.dtype(), n, value, elemLen);
+        var value = new LazyDecimalArray(dtype, 1, MemorySegment.ofArray(elemBytes), elemLen).getDecimal(0);
+        return new LazyConstantDecimalArray(dtype, n, value, elemLen);
     }
 
-    private static Array decodeBool(DecodeContext ctx, ScalarValue scalar, long n) {
+    private static Array decodeBool(DType dtype, ScalarValue scalar, long n) {
         boolean value = scalar.bool_value() != null && scalar.bool_value();
-        return new LazyConstantBoolArray(ctx.dtype(), n, value);
+        return new LazyConstantBoolArray(dtype, n, value);
     }
 
-    private static Array decodeString(DecodeContext ctx, ScalarValue scalar, long n) {
+    private static Array decodeString(DecodeContext ctx, ScalarValue scalar, DType dtype, long n) {
         byte[] strBytes = scalar.string_value() != null
                               ? scalar.string_value().getBytes(StandardCharsets.UTF_8)
                               : (scalar.bytes_value() != null ? scalar.bytes_value() : new byte[0]);
@@ -131,7 +139,7 @@ private static Array decodeString(DecodeContext ctx, ScalarValue scalar, long n)
             offsetsSeg.setAtIndex(PTypeIO.LE_INT, i, (int) (i * strLen));
         }
 
-        return new VarBinArray.OffsetMode(ctx.dtype(), n, bytesSeg.asReadOnly(), offsetsSeg.asReadOnly(), PType.I32);
+        return new VarBinArray.OffsetMode(dtype, n, bytesSeg.asReadOnly(), offsetsSeg.asReadOnly(), PType.I32);
     }
 
     private static long scalarToRawBits(ScalarValue scalar, PType ptype) {

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

@@ -28,7 +28,7 @@ public EncodingId encodingId() {
 
     @Override
     public boolean accepts(DType dtype) {
-        return false;
+        return dtype instanceof DType.Variant;
     }
 
     @Override

writer/src/test/java/io/github/dfa1/vortex/writer/encode/VariantEncodingEncoderTest.java

@@ -141,6 +141,63 @@ void adjacentEqualValues_coalesceIntoOneRun() throws Exception {
         }
     }
 
+    @Nested
+    class RoundTrip {
+
+        private static final io.github.dfa1.vortex.reader.ReadRegistry REGISTRY =
+                io.github.dfa1.vortex.reader.decode.TestRegistry.ofDecoders(
+                        new io.github.dfa1.vortex.reader.decode.VariantEncodingDecoder(),
+                        new io.github.dfa1.vortex.reader.decode.ConstantEncodingDecoder(),
+                        new io.github.dfa1.vortex.reader.decode.ChunkedEncodingDecoder(),
+                        new io.github.dfa1.vortex.reader.decode.PrimitiveEncodingDecoder());
+
+        private static io.github.dfa1.vortex.reader.decode.ArrayNode toArrayNode(EncodeNode node) {
+            io.github.dfa1.vortex.reader.decode.ArrayNode[] children =
+                    new io.github.dfa1.vortex.reader.decode.ArrayNode[node.children().length];
+            for (int i = 0; i < children.length; i++) {
+                children[i] = toArrayNode(node.children()[i]);
+            }
+            return io.github.dfa1.vortex.reader.decode.ArrayNode.of(
+                    node.encodingId(), node.metadata(), children, node.bufferIndices());
+        }
+
+        private static io.github.dfa1.vortex.reader.array.VariantArray decode(EncodeResult result, long rows) {
+            MemorySegment[] bufs = result.buffers().toArray(MemorySegment[]::new);
+            var ctx = new io.github.dfa1.vortex.reader.decode.DecodeContext(
+                    toArrayNode(result.rootNode()), VARIANT, rows, bufs, REGISTRY, java.lang.foreign.Arena.global());
+            return (io.github.dfa1.vortex.reader.array.VariantArray) new io.github.dfa1.vortex.reader.decode.VariantEncodingDecoder().decode(ctx);
+        }
+
+        @Test
+        void constantColumn_decodesToBroadcastInnerValues() {
+            // Given/When a constant column is encoded then decoded back
+            var result = SUT.encode(VARIANT, VariantData.constant(4, i32Scalar(7L)), EncodeTestHelper.testCtx());
+            var variant = decode(result, 4);
+
+            // Then core storage is the inner i32 value broadcast to every row
+            assertThat(variant.length()).isEqualTo(4);
+            assertThat(variant.shredded()).isNull();
+            var core = (io.github.dfa1.vortex.reader.array.IntArray) variant.coreStorage();
+            assertThat(core.dtype()).isEqualTo(new DType.Primitive(io.github.dfa1.vortex.core.PType.I32, false));
+            for (long i = 0; i < 4; i++) {
+                assertThat(core.getInt(i)).isEqualTo(7);
+            }
+        }
+
+        @Test
+        void varyingColumn_decodesPerRowValuesInOrder() {
+            // Given/When distinct per-row values are encoded (chunked) then decoded back
+            var data = new VariantData(List.of(i32Scalar(10L), i32Scalar(20L), i32Scalar(30L)));
+            var variant = decode(SUT.encode(VARIANT, data, EncodeTestHelper.testCtx()), 3);
+
+            // Then the chunked core storage yields each row's inner value in order
+            var core = (io.github.dfa1.vortex.reader.array.IntArray) variant.coreStorage();
+            assertThat(core.getInt(0)).isEqualTo(10);
+            assertThat(core.getInt(1)).isEqualTo(20);
+            assertThat(core.getInt(2)).isEqualTo(30);
+        }
+    }
+
     @Nested
     class Errors {