Skip to content

Commit 9c70f94

Browse files
dfa1claude
dfa1
and
claude
committed
refactor(reader): polymorphic Array.truncate, drop the type switch
Replace ScanIterator's ~135-line instanceof switch + 7 truncateChunked* helpers with a polymorphic Array.truncate(long rows): - Array: default throws; static truncated(arr, rows) holds the single length<=rows no-op guard. - numeric family interfaces (Long/Int/Double/Float/Short/Byte/Bool) default to a zero-copy length-capping view (OffsetXxxArray, offset 0) instead of materialising — fixes the wasteful full-column copy the old catch-all did for lazy arrays. - Chunked* override with the structural keep-prefix-children truncate (shared ChunkedArrays.truncateChildren) to preserve batch iteration. - Null/Generic/Masked/Decimal* override directly; VarBin already had it. Dispatch is now on the runtime type, eliminating the fragile "Chunked/Dict/Constant must precede the family catch-all" ordering that the switch required. No materialization on any truncation path. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent a6126f2 commit 9c70f94

23 files changed

Lines changed: 206 additions & 162 deletions

‎reader/src/main/java/io/github/dfa1/vortex/reader/ScanIterator.java‎

Lines changed: 3 additions & 161 deletions
Original file line numberDiff line numberDiff line change
@@ -20,39 +20,17 @@
2020
import io.github.dfa1.vortex.reader.array.DictIntArray;
2121
import io.github.dfa1.vortex.reader.array.DictLongArray;
2222
import io.github.dfa1.vortex.reader.array.DoubleArray;
23-
import io.github.dfa1.vortex.reader.array.EmptyArray;
2423
import io.github.dfa1.vortex.reader.array.FloatArray;
25-
import io.github.dfa1.vortex.reader.array.GenericArray;
2624
import io.github.dfa1.vortex.reader.array.IntArray;
27-
import io.github.dfa1.vortex.reader.array.LazyConstantBoolArray;
28-
import io.github.dfa1.vortex.reader.array.LazyConstantByteArray;
29-
import io.github.dfa1.vortex.reader.array.LazyConstantDoubleArray;
30-
import io.github.dfa1.vortex.reader.array.LazyConstantFloatArray;
31-
import io.github.dfa1.vortex.reader.array.LazyConstantIntArray;
32-
import io.github.dfa1.vortex.reader.array.LazyConstantLongArray;
33-
import io.github.dfa1.vortex.reader.array.LazyConstantShortArray;
34-
import io.github.dfa1.vortex.reader.array.LazyConstantDecimalArray;
35-
import io.github.dfa1.vortex.reader.array.LazyDecimalArray;
36-
import io.github.dfa1.vortex.reader.array.LazyDecimalBytePartsArray;
37-
import io.github.dfa1.vortex.reader.array.LazyRunEndBoolArray;
38-
import io.github.dfa1.vortex.reader.array.LazySparseBoolArray;
3925
import io.github.dfa1.vortex.reader.array.LongArray;
4026
import io.github.dfa1.vortex.reader.array.MaskedArray;
41-
import io.github.dfa1.vortex.reader.array.MaterializedBoolArray;
42-
import io.github.dfa1.vortex.reader.array.MaterializedByteArray;
43-
import io.github.dfa1.vortex.reader.array.MaterializedDoubleArray;
44-
import io.github.dfa1.vortex.reader.array.MaterializedFloatArray;
45-
import io.github.dfa1.vortex.reader.array.MaterializedIntArray;
46-
import io.github.dfa1.vortex.reader.array.MaterializedLongArray;
47-
import io.github.dfa1.vortex.reader.array.MaterializedShortArray;
4827
import io.github.dfa1.vortex.reader.array.OffsetBoolArray;
4928
import io.github.dfa1.vortex.reader.array.OffsetByteArray;
5029
import io.github.dfa1.vortex.reader.array.OffsetDoubleArray;
5130
import io.github.dfa1.vortex.reader.array.OffsetFloatArray;
5231
import io.github.dfa1.vortex.reader.array.OffsetIntArray;
5332
import io.github.dfa1.vortex.reader.array.OffsetLongArray;
5433
import io.github.dfa1.vortex.reader.array.OffsetShortArray;
55-
import io.github.dfa1.vortex.reader.array.NullArray;
5634
import io.github.dfa1.vortex.reader.array.ShortArray;
5735
import io.github.dfa1.vortex.reader.array.StructArray;
5836
import io.github.dfa1.vortex.reader.array.VarBinArray;
@@ -279,150 +257,14 @@ private static long readUnsigned(MemorySegment seg, long idx, PType ptype) {
279257

280258
// ── Zone-map pruning ──────────────────────────────────────────────────────
281259

282-
private static Map<String, Array> truncateColumns(Map<String, Array> columns, long rows,
283-
SegmentAllocator arena) {
260+
private static Map<String, Array> truncateColumns(Map<String, Array> columns, long rows) {
284261
var result = new LinkedHashMap<String, Array>(columns.size());
285262
for (var entry : columns.entrySet()) {
286-
result.put(entry.getKey(), truncateArray(entry.getValue(), rows, arena));
263+
result.put(entry.getKey(), Array.truncated(entry.getValue(), rows));
287264
}
288265
return Map.copyOf(result);
289266
}
290267

291-
private static Array truncateArray(Array arr, long rows, SegmentAllocator arena) {
292-
if (arr.length() <= rows) {
293-
return arr;
294-
}
295-
// Chunked* cases must precede the LongArray/IntArray/etc catch-alls below — a
296-
// ChunkedLongArray IS a LongArray, but slicing it via ArraySegments.of would
297-
// materialise the entire column before slicing (defeating the zero-copy win
298-
// we just added). Instead, keep the prefix children intact and recursively
299-
// truncate the boundary child.
300-
return switch (arr) {
301-
case ChunkedLongArray a -> truncateChunkedLong(a, rows, arena);
302-
case ChunkedIntArray a -> truncateChunkedInt(a, rows, arena);
303-
case ChunkedDoubleArray a -> truncateChunkedDouble(a, rows, arena);
304-
case ChunkedFloatArray a -> truncateChunkedFloat(a, rows, arena);
305-
case ChunkedShortArray a -> truncateChunkedShort(a, rows, arena);
306-
case ChunkedByteArray a -> truncateChunkedByte(a, rows, arena);
307-
case ChunkedBoolArray a -> truncateChunkedBool(a, rows, arena);
308-
// Dict* cases must precede the LongArray/etc catch-alls below: a DictLongArray
309-
// IS a LongArray, but the catch-all materialises via ArraySegments.of which
310-
// would scatter the entire column to truncate. Instead keep the values
311-
// dictionary intact and just truncate the codes — codes are a primitive Array
312-
// that recursively flows through this same switch.
313-
case DictLongArray a ->
314-
DictLongArray.of(a.dtype(), rows, a.values(), truncateArray(a.codes(), rows, arena));
315-
case DictIntArray a ->
316-
DictIntArray.of(a.dtype(), rows, a.values(), truncateArray(a.codes(), rows, arena));
317-
case DictDoubleArray a ->
318-
DictDoubleArray.of(a.dtype(), rows, a.values(), truncateArray(a.codes(), rows, arena));
319-
case DictFloatArray a ->
320-
DictFloatArray.of(a.dtype(), rows, a.values(), truncateArray(a.codes(), rows, arena));
321-
// LazyConstant* cases must precede the LongArray / IntArray / etc catch-alls:
322-
// each LazyConstantXxxArray IS a typed Array, but the catch-all would
323-
// materialise a length-sized buffer just to slice it. Truncating a constant
324-
// array is a no-buffer length swap.
325-
case LazyConstantLongArray a -> new LazyConstantLongArray(a.dtype(), rows, a.value());
326-
case LazyConstantIntArray a -> new LazyConstantIntArray(a.dtype(), rows, a.value());
327-
case LazyConstantDoubleArray a -> new LazyConstantDoubleArray(a.dtype(), rows, a.value());
328-
case LazyConstantFloatArray a -> new LazyConstantFloatArray(a.dtype(), rows, a.value());
329-
case LazyConstantShortArray a -> new LazyConstantShortArray(a.dtype(), rows, a.value());
330-
case LazyConstantByteArray a -> new LazyConstantByteArray(a.dtype(), rows, a.value());
331-
case LazyConstantBoolArray a -> new LazyConstantBoolArray(a.dtype(), rows, a.value());
332-
case LongArray a ->
333-
new MaterializedLongArray(a.dtype(), rows, ArraySegments.of(a, arena).asSlice(0, rows * Long.BYTES));
334-
case IntArray a ->
335-
new MaterializedIntArray(a.dtype(), rows, ArraySegments.of(a, arena).asSlice(0, rows * Integer.BYTES));
336-
case DoubleArray a ->
337-
new MaterializedDoubleArray(a.dtype(), rows, ArraySegments.of(a, arena).asSlice(0, rows * Double.BYTES));
338-
case FloatArray a ->
339-
new MaterializedFloatArray(a.dtype(), rows, ArraySegments.of(a, arena).asSlice(0, rows * Float.BYTES));
340-
case ShortArray a ->
341-
new MaterializedShortArray(a.dtype(), rows, ArraySegments.of(a, arena).asSlice(0, rows * Short.BYTES));
342-
case ByteArray a -> new MaterializedByteArray(a.dtype(), rows, ArraySegments.of(a, arena).asSlice(0, rows));
343-
case LazyRunEndBoolArray a ->
344-
new LazyRunEndBoolArray(a.dtype(), rows, a.values(), a.runEnds(), a.offset());
345-
case LazySparseBoolArray a ->
346-
new LazySparseBoolArray(a.dtype(), rows, a.fillValue(), a.patchValues(), a.patchIndices(), a.offset());
347-
case BoolArray a ->
348-
new MaterializedBoolArray(a.dtype(), rows, ArraySegments.of(a, arena).asSlice(0, (rows + 7) / 8));
349-
case NullArray a -> new NullArray(a.dtype(), rows);
350-
case VarBinArray a -> a.truncate(rows);
351-
case MaskedArray a -> {
352-
Array truncChild = truncateArray(a.inner(), rows, arena);
353-
BoolArray v = a.validity();
354-
BoolArray truncValidity = (v != null) ? (BoolArray) truncateArray(v, rows, arena) : null;
355-
yield new MaskedArray(truncChild, truncValidity);
356-
}
357-
case EmptyArray a -> a;
358-
case GenericArray a -> a.withLength(rows);
359-
case LazyDecimalArray a ->
360-
new LazyDecimalArray(a.dtype(), rows, a.buf().asSlice(0, rows * (long) a.byteWidth()), a.byteWidth());
361-
case LazyDecimalBytePartsArray a ->
362-
new LazyDecimalBytePartsArray(a.dtype(), rows, truncateArray(a.msp(), rows, arena));
363-
case LazyConstantDecimalArray a ->
364-
new LazyConstantDecimalArray(a.dtype(), rows, a.value(), a.byteWidth());
365-
default ->
366-
throw new VortexException("limit: truncation not supported for " + arr.getClass().getSimpleName());
367-
};
368-
}
369-
370-
/// Truncates a `ChunkedXxxArray` by keeping full children that fit within
371-
/// `rows` and recursively truncating the boundary child. Avoids the
372-
/// full-column materialisation that the {@link LongArray}/{@link IntArray}/etc.
373-
/// catch-all cases would trigger via {@link ArraySegments#of(Array, SegmentAllocator)}.
374-
private static Array truncateChunkedLong(ChunkedLongArray arr, long rows, SegmentAllocator arena) {
375-
List<Array> kept = collectTruncatedChildren(arr.children(), arr.offsets(), rows, arena);
376-
return ChunkedLongArray.of(arr.dtype(), rows, kept);
377-
}
378-
379-
private static Array truncateChunkedInt(ChunkedIntArray arr, long rows, SegmentAllocator arena) {
380-
List<Array> kept = collectTruncatedChildren(arr.children(), arr.offsets(), rows, arena);
381-
return ChunkedIntArray.of(arr.dtype(), rows, kept);
382-
}
383-
384-
private static Array truncateChunkedDouble(ChunkedDoubleArray arr, long rows, SegmentAllocator arena) {
385-
List<Array> kept = collectTruncatedChildren(arr.children(), arr.offsets(), rows, arena);
386-
return ChunkedDoubleArray.of(arr.dtype(), rows, kept);
387-
}
388-
389-
private static Array truncateChunkedFloat(ChunkedFloatArray arr, long rows, SegmentAllocator arena) {
390-
List<Array> kept = collectTruncatedChildren(arr.children(), arr.offsets(), rows, arena);
391-
return ChunkedFloatArray.of(arr.dtype(), rows, kept);
392-
}
393-
394-
private static Array truncateChunkedShort(ChunkedShortArray arr, long rows, SegmentAllocator arena) {
395-
List<Array> kept = collectTruncatedChildren(arr.children(), arr.offsets(), rows, arena);
396-
return ChunkedShortArray.of(arr.dtype(), rows, kept);
397-
}
398-
399-
private static Array truncateChunkedByte(ChunkedByteArray arr, long rows, SegmentAllocator arena) {
400-
List<Array> kept = collectTruncatedChildren(arr.children(), arr.offsets(), rows, arena);
401-
return ChunkedByteArray.of(arr.dtype(), rows, kept);
402-
}
403-
404-
private static Array truncateChunkedBool(ChunkedBoolArray arr, long rows, SegmentAllocator arena) {
405-
List<Array> kept = collectTruncatedChildren(arr.children(), arr.offsets(), rows, arena);
406-
return ChunkedBoolArray.of(arr.dtype(), rows, kept);
407-
}
408-
409-
private static List<Array> collectTruncatedChildren(Array[] children, long[] offsets,
410-
long rows, SegmentAllocator arena) {
411-
var kept = new ArrayList<Array>(children.length);
412-
for (int i = 0; i < children.length; i++) {
413-
long start = offsets[i];
414-
long end = offsets[i + 1];
415-
if (start >= rows) {
416-
break;
417-
}
418-
if (end <= rows) {
419-
kept.add(children[i]);
420-
} else {
421-
kept.add(truncateArray(children[i], rows - start, arena));
422-
}
423-
}
424-
return kept;
425-
}
426268

427269
@Override
428270
public boolean hasNext() {
@@ -470,7 +312,7 @@ public Chunk next() {
470312
try {
471313
Map<String, Array> columns = buildColumnMap(spec, arena);
472314
if (chunkRows < spec.rowCount()) {
473-
columns = truncateColumns(columns, chunkRows, arena);
315+
columns = truncateColumns(columns, chunkRows);
474316
}
475317
rowsReturned += chunkRows;
476318
Map<String, DType> chunkDtypes = new java.util.LinkedHashMap<>();

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package io.github.dfa1.vortex.reader.array;
22

33
import io.github.dfa1.vortex.core.DType;
4+
import io.github.dfa1.vortex.core.VortexException;
45

56
/// Decoded columnar data. Concrete subtypes specialise element access for the JIT;
67
/// each covers a specific dtype family.
@@ -23,5 +24,35 @@ public sealed interface Array
2324
/// @return dtype
2425
DType dtype();
2526

27+
/// Returns an array holding only the first `rows` elements of this array.
28+
///
29+
/// Used by the scan layer to honour [io.github.dfa1.vortex.reader.ScanOptions]
30+
/// row limits. Callers must guarantee `rows < length()` — use the
31+
/// [#truncated(Array, long)] guard, which short-circuits the no-op case.
32+
///
33+
/// The primitive family interfaces ([LongArray], [IntArray], …) provide a
34+
/// zero-copy default: a length-capping view over `this` (no buffer is copied).
35+
/// Subtypes that can do better override it — chunked arrays keep their whole
36+
/// prefix children and only truncate the boundary child, preserving batch
37+
/// iteration. The base default fails fast for types with no truncation support
38+
/// (struct, list, …).
39+
///
40+
/// @param rows number of leading elements to keep; must be `< length()`
41+
/// @return a new array of length `rows`
42+
/// @throws VortexException if truncation is not supported for this array type
43+
default Array truncate(long rows) {
44+
throw new VortexException("limit: truncation not supported for " + getClass().getSimpleName());
45+
}
2646

47+
/// Truncates `arr` to `rows` elements, returning it unchanged when it already
48+
/// fits. Single guard shared by the scan layer and the composite subtypes that
49+
/// recurse into children, so the `rows >= length()` no-op is handled in exactly
50+
/// one place.
51+
///
52+
/// @param arr array to truncate
53+
/// @param rows desired length
54+
/// @return `arr` when `arr.length() <= rows`, otherwise `arr.truncate(rows)`
55+
static Array truncated(Array arr, long rows) {
56+
return arr.length() <= rows ? arr : arr.truncate(rows);
57+
}
2758
}

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,4 +23,13 @@ default void forEachBoolean(BooleanConsumer c) {
2323
c.accept(getBoolean(i));
2424
}
2525
}
26+
27+
/// Zero-copy truncation: a length-capping view over this array.
28+
///
29+
/// @param rows number of leading elements to keep
30+
/// @return a length-`rows` bool array view
31+
@Override
32+
default Array truncate(long rows) {
33+
return new OffsetBoolArray(dtype(), rows, this, 0);
34+
}
2635
}

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,4 +43,14 @@ default void forEachByte(ByteConsumer c) {
4343
c.accept(getByte(i));
4444
}
4545
}
46+
47+
/// Zero-copy truncation: a length-capping view over this array.
48+
/// [ChunkedByteArray] overrides to keep batch iteration.
49+
///
50+
/// @param rows number of leading elements to keep
51+
/// @return a length-`rows` byte array view
52+
@Override
53+
default Array truncate(long rows) {
54+
return new OffsetByteArray(dtype(), rows, this, 0);
55+
}
4656
}

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

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
package io.github.dfa1.vortex.reader.array;
2+
3+
import java.util.ArrayList;
4+
import java.util.List;
5+
6+
/// Shared helper for the `ChunkedXxxArray.truncate` overrides. Keeps whole prefix
7+
/// children that fit within `rows` and recursively truncates the single boundary
8+
/// child — preserving each chunk's batch iteration / zero-copy backing instead of
9+
/// the per-element view a plain length-cap would impose.
10+
final class ChunkedArrays {
11+
12+
private ChunkedArrays() {
13+
}
14+
15+
/// Selects the children of a chunked array that cover the first `rows` rows.
16+
///
17+
/// @param children chunk arrays in scan order
18+
/// @param offsets cumulative row offsets (length = `children.length + 1`)
19+
/// @param rows number of leading rows to keep
20+
/// @return prefix children unchanged, with the boundary child truncated
21+
static List<Array> truncateChildren(Array[] children, long[] offsets, long rows) {
22+
var kept = new ArrayList<Array>(children.length);
23+
for (int i = 0; i < children.length; i++) {
24+
long start = offsets[i];
25+
if (start >= rows) {
26+
break;
27+
}
28+
long end = offsets[i + 1];
29+
if (end <= rows) {
30+
kept.add(children[i]);
31+
} else {
32+
kept.add(Array.truncated(children[i], rows - start));
33+
}
34+
}
35+
return kept;
36+
}
37+
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,4 +70,9 @@ public void forEachBoolean(BooleanConsumer c) {
7070
child.forEachBoolean(c);
7171
}
7272
}
73+
74+
@Override
75+
public Array truncate(long rows) {
76+
return ChunkedBoolArray.of(dtype, rows, ChunkedArrays.truncateChildren(children, offsets, rows));
77+
}
7378
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,4 +82,9 @@ public void forEachByte(ByteConsumer c) {
8282
child.forEachByte(c);
8383
}
8484
}
85+
86+
@Override
87+
public Array truncate(long rows) {
88+
return ChunkedByteArray.of(dtype, rows, ChunkedArrays.truncateChildren(children, offsets, rows));
89+
}
8590
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,4 +79,9 @@ public double fold(double identity, DoubleBinaryOperator op) {
7979
}
8080
return result;
8181
}
82+
83+
@Override
84+
public Array truncate(long rows) {
85+
return ChunkedDoubleArray.of(dtype, rows, ChunkedArrays.truncateChildren(children, offsets, rows));
86+
}
8287
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,4 +69,9 @@ public double fold(double identity, DoubleBinaryOperator op) {
6969
}
7070
return result;
7171
}
72+
73+
@Override
74+
public Array truncate(long rows) {
75+
return ChunkedFloatArray.of(dtype, rows, ChunkedArrays.truncateChildren(children, offsets, rows));
76+
}
7277
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,4 +77,9 @@ public int fold(int identity, IntBinaryOperator op) {
7777
}
7878
return result;
7979
}
80+
81+
@Override
82+
public Array truncate(long rows) {
83+
return ChunkedIntArray.of(dtype, rows, ChunkedArrays.truncateChildren(children, offsets, rows));
84+
}
8085
}

0 commit comments

Comments
 (0)