Skip to content

Commit 7a87c26

Browse files
dfa1claude
dfa1
and
claude
committed
feat(encoding): implement vortex.masked decoder
- Add MaskedArray to sealed Array hierarchy; buffer(i) delegates to child so existing consumers work transparently - Add DType.withNullable(boolean) helper via pattern-match switch - Implement MaskedEncoding.Decoder: validates 0 buffers + 1-2 children, decodes child with non-nullable dtype, optional validity as BoolArray - 7 unit tests covering all-valid, validity masking, dtype, buffer delegation, and validation error cases - Integration test: nullable I64 schema round-trips without error (JNI compressor uses fastlanes.bitpacked, not vortex.masked, for this data; S3 masked.vortex fixture test stays @disabled until published) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent de77b2d commit 7a87c26

6 files changed

Lines changed: 370 additions & 5 deletions

File tree

‎core/src/main/java/io/github/dfa1/vortex/core/DType.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,4 +64,19 @@ record Extension(
6464
boolean nullable
6565
) implements DType {
6666
}
67+
68+
default DType withNullable(boolean nullable) {
69+
return switch (this) {
70+
case Null _ -> new Null(nullable);
71+
case Bool _ -> new Bool(nullable);
72+
case Primitive(var pt, _) -> new Primitive(pt, nullable);
73+
case Decimal(var p, var s, _) -> new Decimal(p, s, nullable);
74+
case Utf8 _ -> new Utf8(nullable);
75+
case Binary _ -> new Binary(nullable);
76+
case Struct(var names, var types, _) -> new Struct(names, types, nullable);
77+
case List(var elem, _) -> new List(elem, nullable);
78+
case FixedSizeList(var elem, var size, _) -> new FixedSizeList(elem, size, nullable);
79+
case Extension(var id, var storage, var meta, _) -> new Extension(id, storage, meta, nullable);
80+
};
81+
}
6782
}

‎core/src/main/java/io/github/dfa1/vortex/core/array/Array.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@
1212
/// is tied to the `VortexFile`'s Arena.
1313
public sealed interface Array
1414
permits BoolArray, ByteArray, DoubleArray, EmptyArray, FixedSizeListArray, Float16Array,
15-
FloatArray, GenericArray, IntArray, ListArray, ListViewArray, LongArray, NullArray,
16-
ShortArray, StructArray, VarBinArray {
15+
FloatArray, GenericArray, IntArray, ListArray, ListViewArray, LongArray, MaskedArray,
16+
NullArray, ShortArray, StructArray, VarBinArray {
1717

1818
long length();
1919

‎core/src/main/java/io/github/dfa1/vortex/core/array/MaskedArray.java‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
package io.github.dfa1.vortex.core.array;
2+
3+
import io.github.dfa1.vortex.core.DType;
4+
5+
import java.lang.foreign.MemorySegment;
6+
7+
/// Decoded {@code vortex.masked} array: a non-nullable child paired with an optional validity bitmap.
8+
///
9+
/// <p>Invariant: {@code child} has no actual nulls — nullability is expressed solely via
10+
/// {@code validity}. When {@code validity} is {@code null} all positions are valid.
11+
///
12+
/// <p>{@link #buffer(int)} delegates to the child so that existing consumers that read raw
13+
/// value buffers continue to work transparently. Use {@link #isValid(long)} to check
14+
/// individual positions before trusting a value.
15+
public final class MaskedArray implements Array {
16+
17+
private final Array child;
18+
private final BoolArray validity;
19+
20+
public MaskedArray(Array child, BoolArray validity) {
21+
this.child = child;
22+
this.validity = validity;
23+
}
24+
25+
@Override
26+
public DType dtype() {
27+
return child.dtype().withNullable(true);
28+
}
29+
30+
@Override
31+
public long length() {
32+
return child.length();
33+
}
34+
35+
@Override
36+
public MemorySegment buffer(int i) {
37+
return child.buffer(i);
38+
}
39+
40+
@Override
41+
public Array child(int i) {
42+
return switch (i) {
43+
case 0 -> child;
44+
case 1 -> {
45+
if (validity == null) {
46+
throw new IndexOutOfBoundsException("no validity child (AllValid)");
47+
}
48+
yield validity;
49+
}
50+
default -> throw new IndexOutOfBoundsException(i);
51+
};
52+
}
53+
54+
public boolean isValid(long i) {
55+
return validity == null || validity.getBoolean(i);
56+
}
57+
}

‎core/src/main/java/io/github/dfa1/vortex/encoding/MaskedEncoding.java‎

Lines changed: 50 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,18 @@
33
import io.github.dfa1.vortex.core.DType;
44
import io.github.dfa1.vortex.core.VortexException;
55
import io.github.dfa1.vortex.core.array.Array;
6+
import io.github.dfa1.vortex.core.array.BoolArray;
7+
import io.github.dfa1.vortex.core.array.MaskedArray;
68

7-
/// Stub for {@code vortex.masked} — not yet implemented.
9+
/// Decoder for {@code vortex.masked}.
10+
///
11+
/// <p>Wire format:
12+
/// <ul>
13+
/// <li>Metadata: empty.</li>
14+
/// <li>Buffers: none.</li>
15+
/// <li>Child 0: payload array encoded with non-nullable dtype (no actual nulls by invariant).</li>
16+
/// <li>Child 1 (optional): validity bitmap, dtype {@code Bool(false)}. Absent means AllValid.</li>
17+
/// </ul>
818
public final class MaskedEncoding implements Encoding {
919

1020
@Override
@@ -14,11 +24,48 @@ public EncodingId encodingId() {
1424

1525
@Override
1626
public EncodeResult encode(DType dtype, Object data) {
17-
throw new VortexException(EncodingId.VORTEX_MASKED, "not yet implemented");
27+
throw new VortexException(EncodingId.VORTEX_MASKED, "encode not yet implemented");
1828
}
1929

2030
@Override
2131
public Array decode(DecodeContext ctx) {
22-
throw new VortexException(EncodingId.VORTEX_MASKED, "not yet implemented");
32+
return Decoder.decode(ctx);
33+
}
34+
35+
private static final class Decoder {
36+
37+
static Array decode(DecodeContext ctx) {
38+
if (ctx.node().bufferIndices().length != 0) {
39+
throw new VortexException(EncodingId.VORTEX_MASKED,
40+
"expected 0 buffers, got " + ctx.node().bufferIndices().length);
41+
}
42+
int numChildren = ctx.node().children().length;
43+
if (numChildren < 1 || numChildren > 2) {
44+
throw new VortexException(EncodingId.VORTEX_MASKED,
45+
"expected 1 or 2 children, got " + numChildren);
46+
}
47+
48+
Array child = decodeChild(ctx, 0, ctx.dtype().withNullable(false), ctx.rowCount());
49+
50+
BoolArray validity = null;
51+
if (numChildren == 2) {
52+
Array validityArray = decodeChild(ctx, 1, new DType.Bool(false), ctx.rowCount());
53+
if (!(validityArray instanceof BoolArray ba)) {
54+
throw new VortexException(EncodingId.VORTEX_MASKED,
55+
"validity child decoded to unexpected type: " + validityArray.getClass().getSimpleName());
56+
}
57+
validity = ba;
58+
}
59+
60+
return new MaskedArray(child, validity);
61+
}
62+
63+
private static Array decodeChild(DecodeContext parent, int idx, DType dtype, long rowCount) {
64+
ArrayNode childNode = parent.node().children()[idx];
65+
DecodeContext childCtx = new DecodeContext(
66+
childNode, dtype, rowCount,
67+
parent.segmentBuffers(), parent.registry(), parent.arena());
68+
return parent.registry().decode(childCtx);
69+
}
2370
}
2471
}

‎core/src/test/java/io/github/dfa1/vortex/encoding/MaskedEncodingTest.java‎

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
package io.github.dfa1.vortex.encoding;
2+
3+
import io.github.dfa1.vortex.core.DType;
4+
import io.github.dfa1.vortex.core.PType;
5+
import io.github.dfa1.vortex.core.VortexException;
6+
import io.github.dfa1.vortex.core.array.Array;
7+
import io.github.dfa1.vortex.core.array.MaskedArray;
8+
import org.junit.jupiter.api.Nested;
9+
import org.junit.jupiter.api.Test;
10+
11+
import java.lang.foreign.Arena;
12+
import java.lang.foreign.MemorySegment;
13+
import java.lang.foreign.ValueLayout;
14+
import java.nio.ByteOrder;
15+
import java.util.ArrayList;
16+
import java.util.List;
17+
18+
import static org.assertj.core.api.Assertions.assertThat;
19+
import static org.assertj.core.api.Assertions.assertThatThrownBy;
20+
21+
class MaskedEncodingTest {
22+
23+
private static final ValueLayout.OfInt LE_INT =
24+
ValueLayout.JAVA_INT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
25+
26+
@Nested
27+
class Decode {
28+
29+
@Test
30+
void oneChild_noValidity_allPositionsValid() {
31+
// Given
32+
var sut = new MaskedEncoding();
33+
EncodingRegistry registry = buildRegistry();
34+
DType i32Nullable = new DType.Primitive(PType.I32, true);
35+
EncodeResult ctx = maskedResult(new int[]{10, 20, 30}, null);
36+
37+
// When
38+
Array result = sut.decode(EncodeTestHelper.toDecodeContext(ctx, 3L, i32Nullable, registry));
39+
40+
// Then
41+
assertThat(result).isInstanceOf(MaskedArray.class);
42+
MaskedArray masked = (MaskedArray) result;
43+
assertThat(masked.length()).isEqualTo(3);
44+
assertThat(masked.isValid(0)).isTrue();
45+
assertThat(masked.isValid(1)).isTrue();
46+
assertThat(masked.isValid(2)).isTrue();
47+
}
48+
49+
@Test
50+
void twoChildren_withValidity_masksNulls() {
51+
// Given
52+
var sut = new MaskedEncoding();
53+
EncodingRegistry registry = buildRegistry();
54+
DType i32Nullable = new DType.Primitive(PType.I32, true);
55+
EncodeResult ctx = maskedResult(new int[]{1, 2, 3, 4, 5},
56+
new boolean[]{true, false, true, false, true});
57+
58+
// When
59+
Array result = sut.decode(EncodeTestHelper.toDecodeContext(ctx, 5L, i32Nullable, registry));
60+
61+
// Then
62+
assertThat(result).isInstanceOf(MaskedArray.class);
63+
MaskedArray masked = (MaskedArray) result;
64+
assertThat(masked.length()).isEqualTo(5);
65+
assertThat(masked.isValid(0)).isTrue();
66+
assertThat(masked.isValid(1)).isFalse();
67+
assertThat(masked.isValid(2)).isTrue();
68+
assertThat(masked.isValid(3)).isFalse();
69+
assertThat(masked.isValid(4)).isTrue();
70+
}
71+
72+
@Test
73+
void dtype_isNullable() {
74+
// Given
75+
var sut = new MaskedEncoding();
76+
EncodingRegistry registry = buildRegistry();
77+
DType i32Nullable = new DType.Primitive(PType.I32, true);
78+
EncodeResult ctx = maskedResult(new int[]{1, 2, 3}, null);
79+
80+
// When
81+
Array result = sut.decode(EncodeTestHelper.toDecodeContext(ctx, 3L, i32Nullable, registry));
82+
83+
// Then
84+
assertThat(result.dtype().nullable()).isTrue();
85+
}
86+
87+
@Test
88+
void buffer_delegatesToChild() {
89+
// Given
90+
var sut = new MaskedEncoding();
91+
EncodingRegistry registry = buildRegistry();
92+
DType i32Nullable = new DType.Primitive(PType.I32, true);
93+
EncodeResult ctx = maskedResult(new int[]{7, 8, 9}, null);
94+
95+
// When
96+
Array result = sut.decode(EncodeTestHelper.toDecodeContext(ctx, 3L, i32Nullable, registry));
97+
98+
// Then — buffer(0) works and contains child values
99+
assertThat(result.buffer(0).get(LE_INT, 0L)).isEqualTo(7);
100+
assertThat(result.buffer(0).get(LE_INT, 4L)).isEqualTo(8);
101+
assertThat(result.buffer(0).get(LE_INT, 8L)).isEqualTo(9);
102+
}
103+
104+
@Test
105+
void buffersPresentThrows() {
106+
// Given
107+
var sut = new MaskedEncoding();
108+
EncodingRegistry registry = buildRegistry();
109+
DType i32Nullable = new DType.Primitive(PType.I32, true);
110+
111+
// Build a node with an unexpected buffer index
112+
EncodeNode childNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 0);
113+
EncodeNode maskedNode = new EncodeNode(
114+
EncodingId.VORTEX_MASKED, null,
115+
new EncodeNode[]{childNode}, new int[]{1});
116+
MemorySegment dummyBuf = Arena.ofAuto().allocate(4);
117+
EncodeResult result = new EncodeResult(maskedNode, List.of(dummyBuf, dummyBuf), null, null);
118+
119+
// When / Then
120+
assertThatThrownBy(() ->
121+
sut.decode(EncodeTestHelper.toDecodeContext(result, 1L, i32Nullable, registry)))
122+
.isInstanceOf(VortexException.class)
123+
.hasMessageContaining("expected 0 buffers");
124+
}
125+
126+
@Test
127+
void zeroChildrenThrows() {
128+
// Given
129+
var sut = new MaskedEncoding();
130+
EncodingRegistry registry = buildRegistry();
131+
DType i32Nullable = new DType.Primitive(PType.I32, true);
132+
133+
EncodeNode maskedNode = new EncodeNode(
134+
EncodingId.VORTEX_MASKED, null, new EncodeNode[]{}, new int[]{});
135+
EncodeResult result = new EncodeResult(maskedNode, List.of(), null, null);
136+
137+
// When / Then
138+
assertThatThrownBy(() ->
139+
sut.decode(EncodeTestHelper.toDecodeContext(result, 0L, i32Nullable, registry)))
140+
.isInstanceOf(VortexException.class)
141+
.hasMessageContaining("expected 1 or 2 children");
142+
}
143+
144+
@Test
145+
void threeChildrenThrows() {
146+
// Given
147+
var sut = new MaskedEncoding();
148+
EncodingRegistry registry = buildRegistry();
149+
DType i32Nullable = new DType.Primitive(PType.I32, true);
150+
151+
PrimitiveEncoding primitiveEncoding = new PrimitiveEncoding();
152+
DType i32 = new DType.Primitive(PType.I32, false);
153+
EncodeResult childResult = primitiveEncoding.encode(i32, new int[]{1});
154+
EncodeNode childNode = childResult.rootNode();
155+
EncodeNode maskedNode = new EncodeNode(
156+
EncodingId.VORTEX_MASKED, null,
157+
new EncodeNode[]{childNode, childNode, childNode}, new int[]{});
158+
List<MemorySegment> bufs = new ArrayList<>(childResult.buffers());
159+
EncodeResult result = new EncodeResult(maskedNode, bufs, null, null);
160+
161+
// When / Then
162+
assertThatThrownBy(() ->
163+
sut.decode(EncodeTestHelper.toDecodeContext(result, 1L, i32Nullable, registry)))
164+
.isInstanceOf(VortexException.class)
165+
.hasMessageContaining("expected 1 or 2 children");
166+
}
167+
}
168+
169+
// ── helpers ───────────────────────────────────────────────────────────────
170+
171+
private static EncodingRegistry buildRegistry() {
172+
EncodingRegistry registry = EncodingRegistry.empty();
173+
registry.register(new MaskedEncoding());
174+
registry.register(new PrimitiveEncoding());
175+
registry.register(new BoolEncoding());
176+
return registry;
177+
}
178+
179+
/// Build a synthetic masked EncodeResult wrapping an I32 child and optional validity.
180+
private static EncodeResult maskedResult(int[] values, boolean[] validity) {
181+
PrimitiveEncoding primitiveEncoding = new PrimitiveEncoding();
182+
DType i32 = new DType.Primitive(PType.I32, false);
183+
EncodeResult childResult = primitiveEncoding.encode(i32, values);
184+
185+
List<MemorySegment> allBuffers = new ArrayList<>(childResult.buffers());
186+
EncodeNode[] children;
187+
188+
if (validity == null) {
189+
children = new EncodeNode[]{childResult.rootNode()};
190+
} else {
191+
BoolEncoding boolEncoding = new BoolEncoding();
192+
DType boolDtype = new DType.Bool(false);
193+
EncodeResult validityResult = boolEncoding.encode(boolDtype, validity);
194+
EncodeNode remapped = EncodeNode.remapBufferIndices(
195+
validityResult.rootNode(), childResult.buffers().size());
196+
allBuffers.addAll(validityResult.buffers());
197+
children = new EncodeNode[]{childResult.rootNode(), remapped};
198+
}
199+
200+
EncodeNode maskedNode = new EncodeNode(
201+
EncodingId.VORTEX_MASKED, null, children, new int[]{});
202+
return new EncodeResult(maskedNode, allBuffers, null, null);
203+
}
204+
}

0 commit comments

Comments
 (0)