Skip to content

Commit d320d42

Browse files
dfa1claude
dfa1
and
claude
committed
fix: chunked List columns and NullArray chunks in ChunkedLayoutDecoder
ChunkedLayoutDecoder.decodeChunkedLayout blindly cast dtype to DType.Primitive for anything that wasn't Bool/Utf8/Binary, throwing a raw ClassCastException instead of VortexException for chunked List columns (#268). VarBinArray.ChunkedMode.of separately threw on a chunk that decoded to NullArray (an all-null vortex.null/vortex.constant chunk) instead of a VarBinArray (#269). Extracts the per-chunk dispatch into ChunkedArrayCombiner: primitives and Bool keep their existing ChunkedXxxArray shapes, Utf8/Binary route through VarBinArray.ChunkedMode (now NullArray-aware, materializing an all-null chunk as a zero-length OffsetMode run), List columns stitch into one ListArray via a recursive combine of their element arrays plus a rebuilt cumulative offsets table, and any other dtype now fails with a VortexException instead of a raw cast. combineChunkValidity also folds NullArray chunks into the row-validity bitmap it reconstructs across chunks. Found via the Raincloud conformance corpus (squad-v2, oasst1, mnist). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
1 parent 70053ea commit d320d42

6 files changed

Lines changed: 414 additions & 80 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
99

1010
### Fixed
1111

12+
- A chunked `List` column spanning several flat chunks now decodes into a single stitched array instead of throwing a raw `ClassCastException`; any other unhandled dtype now fails with `VortexException`. ([#268](https://github.com/dfa1/vortex-java/issues/268))
13+
- A chunked `Utf8`/`Binary` column with an entirely-null chunk (`NullArray`) no longer throws `chunk is not a VarBinArray`; the chunk materializes as an all-null run. ([#269](https://github.com/dfa1/vortex-java/issues/269))
1214
- Nullable low-cardinality columns now share one global dictionary across chunks instead of re-emitting a per-chunk dictionary. ([5fe8b544](https://github.com/dfa1/vortex-java/commit/5fe8b544))
1315
- `MaskedEncodingEncoder` encodes an all-valid or all-invalid validity bitmap as `vortex.constant` instead of a raw per-row bitmap. ([ecd47ead](https://github.com/dfa1/vortex-java/commit/ecd47ead))
1416
- `MaskedEncodingEncoder` tries `vortex.sparse` for a mixed-validity bitmap and keeps it over the raw bitmap when smaller. ([506d036f](https://github.com/dfa1/vortex-java/commit/506d036f))

‎reader/src/main/java/io/github/dfa1/vortex/reader/array/ChunkedArrayCombiner.java‎

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,195 @@
1+
package io.github.dfa1.vortex.reader.array;
2+
3+
import io.github.dfa1.vortex.core.error.VortexException;
4+
import io.github.dfa1.vortex.core.io.VortexFormat;
5+
import io.github.dfa1.vortex.core.model.DType;
6+
import io.github.dfa1.vortex.core.model.PType;
7+
8+
import java.lang.foreign.MemorySegment;
9+
import java.lang.foreign.SegmentAllocator;
10+
import java.lang.foreign.ValueLayout;
11+
import java.util.ArrayList;
12+
import java.util.List;
13+
14+
/// Stitches the per-chunk arrays of one logical column into a single view, dispatching on the
15+
/// column's [DType]. Each family gets its zero-copy composite shape (ADR 0012): primitive and
16+
/// boolean chunks fold into the `ChunkedXxxArray` records, variable-length chunks into
17+
/// [VarBinArray.ChunkedMode], and list chunks into a stitched [ListArray] whose bulk element data
18+
/// stays zero-copy (the child element arrays are themselves combined recursively) while only the
19+
/// small outer offsets table is rebuilt so per-chunk offsets — which each reset to zero — become one
20+
/// cumulative table.
21+
///
22+
/// Extracted from `ChunkedLayoutDecoder` so the same dispatch drives both the top-level chunked
23+
/// column and the recursive combine of a chunked list's elements. Every unsupported dtype fails with
24+
/// a [VortexException] rather than a raw `ClassCastException` — the reader parses untrusted input and
25+
/// must never leak a JDK runtime exception (issue #268).
26+
public final class ChunkedArrayCombiner {
27+
28+
private ChunkedArrayCombiner() {
29+
}
30+
31+
/// Combines the per-chunk arrays of one column into a single logical array, choosing the shape
32+
/// from `dtype`. Masked chunks keep their per-chunk validity: it is concatenated into one
33+
/// row-level bitmap and the result re-wrapped in a [MaskedArray], so a nullable column that
34+
/// spans more than one chunk retains its nulls.
35+
///
36+
/// @param dtype the column's logical type
37+
/// @param totalRows the total logical row count across all chunks
38+
/// @param chunks the decoded per-chunk arrays, in row order (non-empty)
39+
/// @param arena allocator for the combined validity bitmap and, for lists, the rebuilt offsets
40+
/// @return the combined [Array], masked when any chunk contributes nulls
41+
/// @throws VortexException on an empty chunk list or a dtype with no chunked shape
42+
public static Array combine(DType dtype, long totalRows, List<Array> chunks,
43+
SegmentAllocator arena) {
44+
if (chunks.isEmpty()) {
45+
throw new VortexException("chunked combine: empty chunk list");
46+
}
47+
Array data = switch (dtype) {
48+
case DType.Bool ignored -> ChunkedBoolArray.of(dtype, totalRows, chunks);
49+
case DType.Utf8 ignored -> VarBinArray.ChunkedMode.of(dtype, totalRows, chunks, arena);
50+
case DType.Binary ignored -> VarBinArray.ChunkedMode.of(dtype, totalRows, chunks, arena);
51+
case DType.List list -> combineLists(list, totalRows, chunks, arena);
52+
case DType.Primitive prim -> combinePrimitive(prim.ptype(), dtype, totalRows, chunks);
53+
default -> throw new VortexException("unsupported dtype for chunked layout: " + dtype);
54+
};
55+
BoolArray validity = combineChunkValidity(chunks, totalRows, arena);
56+
return validity != null ? new MaskedArray(data, validity) : data;
57+
}
58+
59+
private static Array combinePrimitive(PType ptype, DType dtype, long totalRows,
60+
List<Array> chunks) {
61+
return switch (ptype) {
62+
case I64, U64 -> ChunkedLongArray.of(dtype, totalRows, chunks);
63+
case I32, U32 -> ChunkedIntArray.of(dtype, totalRows, chunks);
64+
case F64 -> ChunkedDoubleArray.of(dtype, totalRows, chunks);
65+
case F32 -> ChunkedFloatArray.of(dtype, totalRows, chunks);
66+
case I16, U16 -> ChunkedShortArray.of(dtype, totalRows, chunks);
67+
case I8, U8 -> ChunkedByteArray.of(dtype, totalRows, chunks);
68+
default -> throw new VortexException("unsupported ptype for chunked layout: " + ptype);
69+
};
70+
}
71+
72+
/// Stitches per-chunk [ListArray] chunks into one [ListArray]. The element arrays of every chunk
73+
/// are combined recursively (staying zero-copy), and one cumulative I64 offsets table of
74+
/// `totalRows + 1` entries is built in `arena`: each chunk's local offsets are shifted by the
75+
/// number of elements contributed by earlier chunks, since per-chunk offsets each restart at
76+
/// zero.
77+
///
78+
/// @param dtype the list's logical type
79+
/// @param totalRows the total outer-list row count across all chunks
80+
/// @param chunks the per-chunk list arrays (each a [ListArray], possibly wrapped [MaskedArray])
81+
/// @param arena allocator for the rebuilt offsets segment
82+
/// @return a single [ListArray] over the combined chunks
83+
private static ListArray combineLists(DType.List dtype, long totalRows, List<Array> chunks,
84+
SegmentAllocator arena) {
85+
var listChunks = new ArrayList<ListArray>(chunks.size());
86+
for (Array chunk : chunks) {
87+
Array unwrapped = chunk instanceof MaskedArray m ? m.inner() : chunk;
88+
if (unwrapped instanceof ListArray la) {
89+
listChunks.add(la);
90+
} else {
91+
throw new VortexException("chunked list: chunk is not a ListArray: "
92+
+ unwrapped.getClass().getSimpleName());
93+
}
94+
}
95+
long outerRows = 0;
96+
for (ListArray la : listChunks) {
97+
outerRows += la.length();
98+
}
99+
if (outerRows != totalRows) {
100+
throw new VortexException("chunked list: chunk rows sum to " + outerRows
101+
+ ", expected " + totalRows);
102+
}
103+
104+
var elementChunks = new ArrayList<Array>(listChunks.size());
105+
for (ListArray la : listChunks) {
106+
elementChunks.add(la.elements());
107+
}
108+
Array combinedElements = combine(dtype.elementType(), sumLengths(elementChunks),
109+
elementChunks, arena);
110+
111+
MemorySegment offsets = arena.allocate((totalRows + 1) * Long.BYTES, Long.BYTES);
112+
offsets.setAtIndex(VortexFormat.LE_LONG, 0, 0L);
113+
long outRow = 0;
114+
long elementBase = 0;
115+
for (ListArray la : listChunks) {
116+
long localRows = la.length();
117+
Array localOffsets = la.offsets();
118+
for (long i = 0; i < localRows; i++) {
119+
long localEnd = readOffset(localOffsets, i + 1);
120+
offsets.setAtIndex(VortexFormat.LE_LONG, outRow + i + 1, elementBase + localEnd);
121+
}
122+
outRow += localRows;
123+
elementBase += readOffset(localOffsets, localRows);
124+
}
125+
Array offsetsArray = new MaterializedLongArray(DType.I64, totalRows + 1, offsets.asReadOnly());
126+
return new ListArray(dtype, totalRows, combinedElements, offsetsArray);
127+
}
128+
129+
private static long sumLengths(List<Array> arrays) {
130+
long total = 0;
131+
for (Array a : arrays) {
132+
total += a.length();
133+
}
134+
return total;
135+
}
136+
137+
/// Reads offset `idx` from a list offsets array as a non-negative long, widening whatever integer
138+
/// ptype the encoder chose (it picks the narrowest that fits the max offset).
139+
///
140+
/// @param offsets the offsets array
141+
/// @param idx the offset index to read
142+
/// @return the offset value as a long
143+
private static long readOffset(Array offsets, long idx) {
144+
return switch (offsets) {
145+
case LongArray la -> la.getLong(idx);
146+
case IntArray ia -> Integer.toUnsignedLong(ia.getInt(idx));
147+
case ShortArray sa -> sa.getInt(idx);
148+
case ByteArray ba -> ba.getInt(idx);
149+
default -> throw new VortexException("unexpected list offsets type: "
150+
+ offsets.getClass().getSimpleName());
151+
};
152+
}
153+
154+
/// Concatenates the per-chunk validity of masked chunks into one row-level bitmap, or returns
155+
/// `null` when no chunk is nullable. Unmasked chunks (and masked chunks with a `null` validity,
156+
/// which mean all-valid) contribute all-valid rows; an entirely-null [NullArray] chunk (#269)
157+
/// contributes all-invalid rows.
158+
///
159+
/// @param chunkArrays the decoded per-chunk arrays, in row order
160+
/// @param totalRows the total logical row count across all chunks
161+
/// @param arena allocator for the combined bitmap
162+
/// @return a bit-packed row-validity [BoolArray] of `totalRows` bits, or `null` if all valid
163+
private static BoolArray combineChunkValidity(List<Array> chunkArrays, long totalRows,
164+
SegmentAllocator arena) {
165+
boolean anyNullable = false;
166+
for (Array chunk : chunkArrays) {
167+
if ((chunk instanceof MaskedArray m && m.validity() != null) || chunk instanceof NullArray) {
168+
anyNullable = true;
169+
break;
170+
}
171+
}
172+
if (!anyNullable) {
173+
return null;
174+
}
175+
MemorySegment bits = arena.allocate((totalRows + 7) >>> 3);
176+
long row = 0;
177+
for (Array chunk : chunkArrays) {
178+
long chunkLen = chunk.length();
179+
if (!(chunk instanceof NullArray)) {
180+
BoolArray validity = chunk instanceof MaskedArray m ? m.validity() : null;
181+
for (long i = 0; i < chunkLen; i++) {
182+
if (validity == null || validity.getBoolean(i)) {
183+
long globalRow = row + i;
184+
long byteIdx = globalRow >>> 3;
185+
byte cur = bits.get(ValueLayout.JAVA_BYTE, byteIdx);
186+
bits.set(ValueLayout.JAVA_BYTE, byteIdx,
187+
(byte) ((cur & 0xff) | (1 << (globalRow & 7))));
188+
}
189+
}
190+
}
191+
row += chunkLen;
192+
}
193+
return new MaterializedBoolArray(DType.BOOL, totalRows, bits.asReadOnly());
194+
}
195+
}

‎reader/src/main/java/io/github/dfa1/vortex/reader/array/VarBinArray.java‎

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -442,7 +442,9 @@ private long dictReadOff(long i) {
442442
record ChunkedMode(DType dtype, long length, VarBinArray[] children, long[] offsets)
443443
implements VarBinArray {
444444

445-
/// Builds a `ChunkedMode` from a list of chunk arrays.
445+
/// Builds a `ChunkedMode` from a list of chunk arrays that are already all typed
446+
/// [VarBinArray]s (no all-null [NullArray] chunks). Used by [#limited(long)],
447+
/// whose children are always concrete `VarBinArray`s.
446448
///
447449
/// @param dtype logical element type
448450
/// @param totalRows expected total row count
@@ -451,6 +453,28 @@ record ChunkedMode(DType dtype, long length, VarBinArray[] children, long[] offs
451453
/// @throws VortexException on empty input, non-[VarBinArray] chunks, or row-count mismatch
452454
public static ChunkedMode of(DType dtype, long totalRows,
453455
java.util.List<? extends Array> chunks) {
456+
return of(dtype, totalRows, chunks, null);
457+
}
458+
459+
/// Builds a `ChunkedMode` from a list of chunk arrays.
460+
///
461+
/// An entirely-null chunk decodes to a [NullArray] rather than a [VarBinArray]
462+
/// (e.g. a `vortex.null` flat, or `vortex.constant` with a null scalar, #269).
463+
/// Such a chunk is materialized into an all-null [OffsetMode] of the same row
464+
/// count — every row zero-length, all offsets zero — so the chunked column keeps
465+
/// a uniform `VarBinArray` shape. Row-level nullability is preserved separately by
466+
/// the caller's validity bitmap.
467+
///
468+
/// @param dtype logical element type
469+
/// @param totalRows expected total row count
470+
/// @param chunks non-empty list of chunk arrays; each a [VarBinArray] or [NullArray]
471+
/// @param arena allocator for the offsets segment of a materialized null chunk;
472+
/// may be `null` only when no chunk is a [NullArray]
473+
/// @return a new `ChunkedMode`
474+
/// @throws VortexException on empty input, non-`VarBinArray`/`NullArray` chunks,
475+
/// or row-count mismatch
476+
public static ChunkedMode of(DType dtype, long totalRows,
477+
java.util.List<? extends Array> chunks, SegmentAllocator arena) {
454478
if (chunks.isEmpty()) {
455479
throw new VortexException("VarBinArray.ChunkedMode: empty chunk list");
456480
}
@@ -461,6 +485,12 @@ public static ChunkedMode of(DType dtype, long totalRows,
461485
java.util.Collections.addAll(typed, nested.children);
462486
} else if (data instanceof VarBinArray vb) {
463487
typed.add(vb);
488+
} else if (data instanceof NullArray na) {
489+
if (arena == null) {
490+
throw new VortexException(
491+
"VarBinArray.ChunkedMode: null chunk requires an allocator");
492+
}
493+
typed.add(allNull(dtype, na.length(), arena));
464494
} else {
465495
throw new VortexException("VarBinArray.ChunkedMode: chunk is not a VarBinArray: "
466496
+ data.getClass().getSimpleName());
@@ -477,6 +507,19 @@ public static ChunkedMode of(DType dtype, long totalRows,
477507
return new ChunkedMode(dtype, totalRows, typed.toArray(VarBinArray[]::new), off);
478508
}
479509

510+
/// Builds an all-null [OffsetMode] of `n` rows: an empty bytes segment and an
511+
/// offsets segment of `n + 1` zeros, so every row is zero-length. Row nullability
512+
/// is carried by the caller's validity bitmap, not the byte data.
513+
///
514+
/// @param dtype logical element type (Utf8 or Binary)
515+
/// @param n number of all-null rows
516+
/// @param arena allocator for the offsets segment
517+
/// @return an [OffsetMode] with `n` zero-length rows
518+
private static OffsetMode allNull(DType dtype, long n, SegmentAllocator arena) {
519+
MemorySegment offsets = arena.allocate((n + 1) * Long.BYTES, Long.BYTES);
520+
return new OffsetMode(dtype, n, MemorySegment.NULL, offsets, PType.I64);
521+
}
522+
480523
private int findChunk(long i) {
481524
int hit = java.util.Arrays.binarySearch(offsets, i);
482525
int idx = hit >= 0 ? hit : -hit - 2;

0 commit comments

Comments
 (0)