Skip to content

Commit ff66e7a

Browse files
dfa1claude
dfa1
and
claude
committed
fix(nullable): propagate validity bitmaps in primitive, FOR, and RLE decoders
PrimitiveEncoding: decode validity child (if present) and wrap in MaskedArray. FrameOfReferenceEncoding: unwrap MaskedArray from child, apply reference, re-wrap. RleEncoding: extract validity from nullable indices (child[1]); propagate to output. StructEncoding: scalar nullable wrapper (2 children) now returns MaskedArray. Fixes incorrect numeric sums for nullable_i32 (for.vortex) and nullable_runs (rle.vortex) in the JNI vs Java comparison integration test — all 25 fixtures pass. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 64538cb commit ff66e7a

10 files changed

Lines changed: 412 additions & 42 deletions

File tree

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

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,12 @@
66
import io.github.dfa1.vortex.core.ArrayStats;
77
import io.github.dfa1.vortex.core.DType;
88
import io.github.dfa1.vortex.core.PType;
9+
import io.github.dfa1.vortex.core.array.BoolArray;
910
import io.github.dfa1.vortex.core.array.ByteArray;
1011
import io.github.dfa1.vortex.core.array.DoubleArray;
1112
import io.github.dfa1.vortex.core.array.IntArray;
1213
import io.github.dfa1.vortex.core.array.LongArray;
14+
import io.github.dfa1.vortex.core.array.MaskedArray;
1315
import io.github.dfa1.vortex.core.array.ShortArray;
1416
import io.github.dfa1.vortex.core.VortexException;
1517

@@ -199,26 +201,35 @@ private static Array decode(DecodeContext ctx) {
199201

200202
Array encoded = ctx.decodeChild(0);
201203

204+
// Nullable primitive: child decodes as MaskedArray; extract values and propagate validity.
205+
BoolArray validity = null;
206+
Array rawEncoded = encoded;
207+
if (encoded instanceof MaskedArray masked) {
208+
rawEncoded = masked.child(0);
209+
validity = (BoolArray) masked.child(1);
210+
}
211+
202212
if (!(ctx.dtype() instanceof DType.Primitive p)) {
203213
throw new VortexException(EncodingId.FASTLANES_FOR, "expected primitive dtype, got " + ctx.dtype());
204214
}
205215

206216
long ref = referenceValue(scalar);
207217
if (ref == 0L) {
208-
return encoded;
218+
return validity != null ? new MaskedArray(rawEncoded, validity) : rawEncoded;
209219
}
210220

211-
MemorySegment src = encoded.buffer(0);
221+
MemorySegment src = rawEncoded.buffer(0);
212222
long n = ctx.rowCount();
213223
MemorySegment dst = applyReference(src, n, p.ptype(), ref, ctx.arena());
214-
return switch (p.ptype()) {
224+
Array result = switch (p.ptype()) {
215225
case I64, U64 -> new LongArray(ctx.dtype(), n, dst, ArrayStats.empty());
216226
case I32, U32 -> new IntArray(ctx.dtype(), n, dst, ArrayStats.empty());
217227
case F64 -> new DoubleArray(ctx.dtype(), n, dst, ArrayStats.empty());
218228
case I16, U16 -> new ShortArray(ctx.dtype(), n, dst, ArrayStats.empty());
219229
case I8, U8 -> new ByteArray(ctx.dtype(), n, dst, ArrayStats.empty());
220230
default -> throw new VortexException(EncodingId.FASTLANES_FOR, "unsupported ptype " + p.ptype());
221231
};
232+
return validity != null ? new MaskedArray(result, validity) : result;
222233
}
223234

224235
private static long referenceValue(ScalarProtos.ScalarValue scalar) {

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,15 @@
44
import io.github.dfa1.vortex.core.array.Array;
55
import io.github.dfa1.vortex.core.DType;
66
import io.github.dfa1.vortex.core.PType;
7+
import io.github.dfa1.vortex.core.VortexException;
8+
import io.github.dfa1.vortex.core.array.BoolArray;
79
import io.github.dfa1.vortex.core.array.ByteArray;
810
import io.github.dfa1.vortex.core.array.DoubleArray;
911
import io.github.dfa1.vortex.core.array.Float16Array;
1012
import io.github.dfa1.vortex.core.array.FloatArray;
1113
import io.github.dfa1.vortex.core.array.IntArray;
1214
import io.github.dfa1.vortex.core.array.LongArray;
15+
import io.github.dfa1.vortex.core.array.MaskedArray;
1316
import io.github.dfa1.vortex.core.array.ShortArray;
1417
import java.lang.foreign.Arena;
1518
import java.lang.foreign.MemorySegment;
@@ -309,7 +312,7 @@ private static Array decode(DecodeContext ctx) {
309312
DType dt = ctx.dtype();
310313
PType ptype = ((DType.Primitive) dt).ptype();
311314
var stats = ctx.node().stats();
312-
return switch (ptype) {
315+
Array values = switch (ptype) {
313316
case I64, U64 -> new LongArray(dt, n, buf, stats);
314317
case I32, U32 -> new IntArray(dt, n, buf, stats);
315318
case F64 -> new DoubleArray(dt, n, buf, stats);
@@ -318,6 +321,19 @@ private static Array decode(DecodeContext ctx) {
318321
case I8, U8 -> new ByteArray(dt, n, buf, stats);
319322
case F16 -> new Float16Array(dt, n, buf, stats);
320323
};
324+
if (ctx.node().children().length == 1) {
325+
ArrayNode validityNode = ctx.node().children()[0];
326+
var validityCtx = new DecodeContext(
327+
validityNode, new DType.Bool(false), n,
328+
ctx.segmentBuffers(), ctx.registry(), ctx.arena());
329+
Array va = ctx.registry().decode(validityCtx);
330+
if (!(va instanceof BoolArray validity)) {
331+
throw new VortexException(EncodingId.VORTEX_PRIMITIVE,
332+
"validity child decoded to unexpected type: " + va.getClass().getSimpleName());
333+
}
334+
return new MaskedArray(values, validity);
335+
}
336+
return values;
321337
}
322338
}
323339
}

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

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,14 @@
66
import io.github.dfa1.vortex.core.PType;
77
import io.github.dfa1.vortex.core.VortexException;
88
import io.github.dfa1.vortex.core.array.Array;
9+
import io.github.dfa1.vortex.core.array.BoolArray;
910
import io.github.dfa1.vortex.core.array.ByteArray;
1011
import io.github.dfa1.vortex.core.array.DoubleArray;
1112
import io.github.dfa1.vortex.core.array.Float16Array;
1213
import io.github.dfa1.vortex.core.array.FloatArray;
1314
import io.github.dfa1.vortex.core.array.IntArray;
1415
import io.github.dfa1.vortex.core.array.LongArray;
16+
import io.github.dfa1.vortex.core.array.MaskedArray;
1517
import io.github.dfa1.vortex.core.array.ShortArray;
1618
import io.github.dfa1.vortex.proto.DTypeProtos;
1719
import io.github.dfa1.vortex.proto.EncodingProtos;
@@ -329,9 +331,17 @@ static Array decode(DecodeContext ctx) {
329331
DType offsetsDtype = new DType.Primitive(offsetsPtype, false);
330332

331333
Array valuesArr = decodeChildAs(ctx, 0, valuesDtype, valuesLen);
332-
Array indicesArr = decodeChildAs(ctx, 1, indicesDtype, indicesLen);
334+
Array indicesRaw = decodeChildAs(ctx, 1, indicesDtype, indicesLen);
333335
Array offsetsArr = decodeChildAs(ctx, 2, offsetsDtype, offsetsLen);
334336

337+
// Indices may carry a validity bitmap when the output column is nullable.
338+
BoolArray indicesValidity = null;
339+
Array indicesArr = indicesRaw;
340+
if (indicesRaw instanceof MaskedArray masked) {
341+
indicesArr = masked.child(0);
342+
indicesValidity = (BoolArray) masked.child(1);
343+
}
344+
335345
long[] values = readLongs(valuesArr.buffer(0), (int) valuesLen, ptype);
336346
int[] indices = readIndices(indicesArr.buffer(0), (int) indicesLen, indicesPtype);
337347
long[] valuesIdxOffsets = readUnsignedLongs(offsetsArr.buffer(0), (int) offsetsLen, offsetsPtype);
@@ -370,7 +380,22 @@ static Array decode(DecodeContext ctx) {
370380
}
371381

372382
MemorySegment seg = fromLongs(decoded, offset, (int) rowCount, ptype, ctx.arena());
373-
return toArray(ctx.dtype(), rowCount, seg, ptype);
383+
Array result = toArray(ctx.dtype(), rowCount, seg, ptype);
384+
if (indicesValidity == null) {
385+
return result;
386+
}
387+
// Propagate indices validity to output: output[j] is null iff indices[offset+j] is null.
388+
int validityBytes = (int) ((rowCount + 7) / 8);
389+
MemorySegment validityBuf = ctx.arena().allocate(validityBytes);
390+
for (long j = 0; j < rowCount; j++) {
391+
if (indicesValidity.getBoolean(offset + j)) {
392+
int byteIdx = (int) (j >>> 3);
393+
byte current = validityBuf.get(ValueLayout.JAVA_BYTE, (long) byteIdx);
394+
validityBuf.set(ValueLayout.JAVA_BYTE, (long) byteIdx, (byte) (current | (1 << (j & 7))));
395+
}
396+
}
397+
BoolArray outputValidity = new BoolArray(new DType.Bool(false), rowCount, validityBuf, ArrayStats.empty());
398+
return new MaskedArray(result, outputValidity);
374399
}
375400

376401
private static Array emptyArray(DecodeContext ctx) {

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

Lines changed: 49 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
package io.github.dfa1.vortex.encoding;
22

3-
import io.github.dfa1.vortex.core.array.Array;
4-
import io.github.dfa1.vortex.core.array.StructArray;
53
import io.github.dfa1.vortex.core.DType;
64
import io.github.dfa1.vortex.core.VortexException;
5+
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;
8+
import io.github.dfa1.vortex.core.array.StructArray;
79

810
import java.lang.foreign.MemorySegment;
911
import java.util.ArrayList;
@@ -97,39 +99,67 @@ private static Array decode(DecodeContext ctx) {
9799
boolean hasValidity = (numChildren == nfields + 1);
98100
int fieldOffset = hasValidity ? 1 : 0;
99101

102+
BoolArray structValidity = null;
103+
if (hasValidity) {
104+
ArrayNode validityNode = ctx.node().children()[0];
105+
var validityCtx = new DecodeContext(validityNode, new DType.Bool(false),
106+
ctx.rowCount(), ctx.segmentBuffers(), ctx.registry(), ctx.arena());
107+
Array va = ctx.registry().decode(validityCtx);
108+
if (!(va instanceof BoolArray ba)) {
109+
throw new VortexException(EncodingId.VORTEX_STRUCT,
110+
"struct validity decoded to unexpected type: " + va.getClass().getSimpleName());
111+
}
112+
structValidity = ba;
113+
}
114+
100115
if (nfields == 1) {
101116
DType fieldDtype = structDtype.fieldTypes().getFirst();
102117
ArrayNode fieldNode = ctx.node().children()[fieldOffset];
103-
var fieldCtx = new DecodeContext(fieldNode, fieldDtype, ctx.rowCount(),
104-
ctx.segmentBuffers(), ctx.registry(), ctx.arena());
105-
return ctx.registry().decode(fieldCtx);
118+
var fieldCtx = new DecodeContext(fieldNode, fieldDtype.withNullable(false),
119+
ctx.rowCount(), ctx.segmentBuffers(), ctx.registry(), ctx.arena());
120+
Array field = ctx.registry().decode(fieldCtx);
121+
return structValidity != null ? new MaskedArray(field, structValidity) : field;
106122
}
107123

108124
List<Array> fieldArrays = new ArrayList<>(nfields);
109125
for (int i = 0; i < nfields; i++) {
110126
ArrayNode fieldNode = ctx.node().children()[fieldOffset + i];
111127
DType fieldDtype = structDtype.fieldTypes().get(i);
112-
var fieldCtx = new DecodeContext(fieldNode, fieldDtype, ctx.rowCount(),
113-
ctx.segmentBuffers(), ctx.registry(), ctx.arena());
114-
fieldArrays.add(ctx.registry().decode(fieldCtx));
128+
var fieldCtx = new DecodeContext(fieldNode, fieldDtype.withNullable(false),
129+
ctx.rowCount(), ctx.segmentBuffers(), ctx.registry(), ctx.arena());
130+
Array field = ctx.registry().decode(fieldCtx);
131+
fieldArrays.add(structValidity != null ? new MaskedArray(field, structValidity) : field);
115132
}
116133
return new StructArray(structDtype, ctx.rowCount(), fieldArrays);
117134
}
118135

119136
// Scalar nullable wrapper: nfields == 1
120137
// children = [values] (non-nullable) or [validity, values] (nullable)
121-
int valuesIdx = switch (numChildren) {
122-
case 1 -> 0;
123-
case 2 -> 1;
124-
default -> throw new VortexException(EncodingId.VORTEX_STRUCT,
138+
if (numChildren == 1) {
139+
ArrayNode valuesNode = ctx.node().children()[0];
140+
var valuesCtx = new DecodeContext(
141+
valuesNode, ctx.dtype(), ctx.rowCount(),
142+
ctx.segmentBuffers(), ctx.registry(), ctx.arena());
143+
return ctx.registry().decode(valuesCtx);
144+
} else if (numChildren == 2) {
145+
ArrayNode validityNode = ctx.node().children()[0];
146+
var validityCtx = new DecodeContext(validityNode, new DType.Bool(false),
147+
ctx.rowCount(), ctx.segmentBuffers(), ctx.registry(), ctx.arena());
148+
Array va = ctx.registry().decode(validityCtx);
149+
if (!(va instanceof BoolArray validity)) {
150+
throw new VortexException(EncodingId.VORTEX_STRUCT,
151+
"scalar wrapper validity decoded to unexpected type: " + va.getClass().getSimpleName());
152+
}
153+
ArrayNode valuesNode = ctx.node().children()[1];
154+
var valuesCtx = new DecodeContext(
155+
valuesNode, ctx.dtype().withNullable(false), ctx.rowCount(),
156+
ctx.segmentBuffers(), ctx.registry(), ctx.arena());
157+
Array values = ctx.registry().decode(valuesCtx);
158+
return new MaskedArray(values, validity);
159+
} else {
160+
throw new VortexException(EncodingId.VORTEX_STRUCT,
125161
"unexpected child count " + numChildren + " for scalar wrapper");
126-
};
127-
128-
ArrayNode valuesNode = ctx.node().children()[valuesIdx];
129-
var valuesCtx = new DecodeContext(
130-
valuesNode, ctx.dtype(), ctx.rowCount(),
131-
ctx.segmentBuffers(), ctx.registry(), ctx.arena());
132-
return ctx.registry().decode(valuesCtx);
162+
}
133163
}
134164
}
135165
}

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

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import io.github.dfa1.vortex.core.ArrayStats;
66
import io.github.dfa1.vortex.core.DType;
77
import io.github.dfa1.vortex.core.PType;
8+
import io.github.dfa1.vortex.core.array.MaskedArray;
89
import org.junit.jupiter.api.Nested;
910
import org.junit.jupiter.api.Test;
1011
import org.junit.jupiter.params.ParameterizedTest;
@@ -107,6 +108,53 @@ void decode_wrappingAdd_i64(long reference) {
107108
assertThat(got).isEqualTo(residuals[0] + reference);
108109
}
109110

111+
@Test
112+
void decode_nullableResiduals_returnsMaskedArrayWithCorrectValues() {
113+
// Given — 4 I32 residuals; positions 1 and 3 are null (validity: 0b00000101 = 0x05)
114+
// Residuals: [0, 0, 5, 0], reference: 100 → valid outputs: [100, ?, 105, ?]
115+
long reference = 100L;
116+
long[] residuals = {0, 0, 5, 0};
117+
MemorySegment validitySeg = MemorySegment.ofArray(new byte[]{0x05}); // bits 0,2
118+
119+
byte[] residualBytes = new byte[residuals.length * 4];
120+
ByteBuffer bb = ByteBuffer.wrap(residualBytes).order(ByteOrder.LITTLE_ENDIAN);
121+
for (long v : residuals) {
122+
bb.putInt((int) v);
123+
}
124+
125+
ArrayNode validityNode = new ArrayNode(
126+
EncodingId.VORTEX_BOOL, null, new ArrayNode[0], new int[]{1}, ArrayStats.empty());
127+
ArrayNode primNode = new ArrayNode(
128+
EncodingId.VORTEX_PRIMITIVE, null, new ArrayNode[]{validityNode}, new int[]{0}, ArrayStats.empty());
129+
byte[] metaBytes = ScalarProtos.ScalarValue.newBuilder().setInt64Value(reference).build().toByteArray();
130+
ArrayNode forNode = new ArrayNode(
131+
EncodingId.FASTLANES_FOR, ByteBuffer.wrap(metaBytes), new ArrayNode[]{primNode}, new int[0], ArrayStats.empty());
132+
133+
EncodingRegistry registry = EncodingRegistry.empty();
134+
registry.register(new FrameOfReferenceEncoding());
135+
registry.register(new PrimitiveEncoding());
136+
registry.register(new BoolEncoding());
137+
138+
MemorySegment[] segments = {MemorySegment.ofArray(residualBytes), validitySeg};
139+
DecodeContext ctx = new DecodeContext(
140+
forNode, I32_DTYPE, residuals.length, segments, registry, java.lang.foreign.Arena.global());
141+
FrameOfReferenceEncoding sut = new FrameOfReferenceEncoding();
142+
143+
// When
144+
Array result = sut.decode(ctx);
145+
146+
// Then — MaskedArray; reference added to valid positions only
147+
assertThat(result).isInstanceOf(MaskedArray.class);
148+
MaskedArray masked = (MaskedArray) result;
149+
assertThat(masked.isValid(0)).isTrue();
150+
assertThat(masked.isValid(1)).isFalse();
151+
assertThat(masked.isValid(2)).isTrue();
152+
assertThat(masked.isValid(3)).isFalse();
153+
var layout = ValueLayout.JAVA_INT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
154+
assertThat(masked.child(0).buffer(0).get(layout, 0L)).isEqualTo(100);
155+
assertThat(masked.child(0).buffer(0).get(layout, 8L)).isEqualTo(105);
156+
}
157+
110158
private static DecodeContext buildForContext(
111159
DType dtype, long reference, long[] residuals, PType ptype
112160
) {

0 commit comments

Comments
 (0)