Skip to content

Commit 70a46c4

Browse files
dfa1claude
dfa1
and
claude
committed
feat(chunked): implement vortex.chunked segment-level decode
Supports primitive and struct dtype concatenation. Fixes chunked.vortex fixture (23/35). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent b65e479 commit 70a46c4

5 files changed

Lines changed: 293 additions & 2 deletions

File tree

‎TODO.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@
9393
| `vortex.fsst` | `FsstEncoding` | ✅ | ❌ stub | — | Utf8, Binary |
9494
| `vortex.varbinview` | `VarBinViewEncoding` | ✅ | ❌ stub | — | Utf8, Binary |
9595
| `vortex.pco` | `PcoEncoding` (stub) | ❌ | ❌ | very hard | ANS + bin tokenization not ported; unblocks `pco.vortex`, tpch/clickbench fixtures |
96-
| `vortex.chunked` | — | ❌ | ❌ | medium | unblocks `chunked.vortex` (segment-level chunked array) |
96+
| `vortex.chunked` | `ChunkedEncoding` | ✅ | ❌ stub | medium | decode: primitive + struct concat; encode stub |
9797
| `fastlanes.rle` | — | ❌ | ❌ | medium | unblocks `rle.vortex` |
9898
| `vortex.alprd` | — | ❌ | ❌ | medium | unblocks `alprd.vortex` |
9999
| `vortex.decimal` | `DecimalEncoding` | ✅ | ✅ | — | Decimal |
@@ -140,7 +140,7 @@
140140
| `dict.vortex` | ✅ | |
141141
| `sparse.vortex` | ✅ | |
142142
| `varbinview.vortex` | ✅ | |
143-
| `chunked.vortex` | ❌ | `vortex.chunked` at segment level |
143+
| `chunked.vortex` | ✅ | |
144144
| `rle.vortex` | ❌ | `fastlanes.rle` missing |
145145
| `alprd.vortex` | ❌ | `vortex.alprd` missing |
146146
| `decimal.vortex` | ✅ | |

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

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
package io.github.dfa1.vortex.encoding;
2+
3+
import io.github.dfa1.vortex.core.ArrayStats;
4+
import io.github.dfa1.vortex.core.DType;
5+
import io.github.dfa1.vortex.core.PType;
6+
import io.github.dfa1.vortex.core.VortexException;
7+
import io.github.dfa1.vortex.core.array.Array;
8+
import io.github.dfa1.vortex.core.array.ByteArray;
9+
import io.github.dfa1.vortex.core.array.DoubleArray;
10+
import io.github.dfa1.vortex.core.array.FloatArray;
11+
import io.github.dfa1.vortex.core.array.IntArray;
12+
import io.github.dfa1.vortex.core.array.LongArray;
13+
import io.github.dfa1.vortex.core.array.ShortArray;
14+
import io.github.dfa1.vortex.core.array.StructArray;
15+
16+
import java.lang.foreign.MemorySegment;
17+
import java.lang.foreign.SegmentAllocator;
18+
import java.lang.foreign.ValueLayout;
19+
import java.nio.ByteOrder;
20+
import java.util.ArrayList;
21+
import java.util.List;
22+
23+
/// Encoding for {@code vortex.chunked} — a segment-level chunked array.
24+
///
25+
/// <p>Wire format (per Rust vtable):
26+
/// <ul>
27+
/// <li>Metadata: empty bytes
28+
/// <li>Buffers: 0
29+
/// <li>Children: N+1 — {@code children[0]} is the chunk offsets (U64, non-nullable,
30+
/// length = nchunks+1, cumulative from 0); {@code children[1..N]} are the chunk arrays.
31+
/// </ul>
32+
public final class ChunkedEncoding implements Encoding {
33+
34+
@Override
35+
public EncodingId encodingId() {
36+
return EncodingId.VORTEX_CHUNKED;
37+
}
38+
39+
@Override
40+
public boolean accepts(DType dtype) {
41+
return dtype instanceof DType.Primitive || dtype instanceof DType.Struct;
42+
}
43+
44+
@Override
45+
public EncodeResult encode(DType dtype, Object data) {
46+
throw new UnsupportedOperationException("encode not yet implemented for " + encodingId());
47+
}
48+
49+
@Override
50+
public Array decode(DecodeContext ctx) {
51+
return Decoder.decode(ctx);
52+
}
53+
54+
private static final class Decoder {
55+
56+
private static final ValueLayout.OfLong LE_LONG =
57+
ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
58+
59+
static Array decode(DecodeContext ctx) {
60+
int nchildren = ctx.node().children().length;
61+
if (nchildren < 1) {
62+
throw new VortexException(EncodingId.VORTEX_CHUNKED,
63+
"needs at least one child (chunk offsets)");
64+
}
65+
int nchunks = nchildren - 1;
66+
long[] offsets = readOffsets(ctx, nchunks);
67+
68+
DType dtype = ctx.dtype();
69+
List<Array> chunks = new ArrayList<>(nchunks);
70+
for (int i = 0; i < nchunks; i++) {
71+
long chunkLen = offsets[i + 1] - offsets[i];
72+
ArrayNode chunkNode = ctx.node().children()[i + 1];
73+
var chunkCtx = new DecodeContext(
74+
chunkNode, dtype, chunkLen, ctx.segmentBuffers(), ctx.registry(), ctx.arena());
75+
chunks.add(ctx.registry().decode(chunkCtx));
76+
}
77+
78+
return concat(chunks, dtype, ctx.rowCount(), ctx.arena());
79+
}
80+
81+
private static long[] readOffsets(DecodeContext ctx, int nchunks) {
82+
ArrayNode offsetsNode = ctx.node().children()[0];
83+
DType u64 = new DType.Primitive(PType.U64, false);
84+
var offsetsCtx = new DecodeContext(
85+
offsetsNode, u64, nchunks + 1L, ctx.segmentBuffers(), ctx.registry(), ctx.arena());
86+
Array offsetsArray = ctx.registry().decode(offsetsCtx);
87+
MemorySegment offsetsBuf = offsetsArray.buffer(0);
88+
long[] offsets = new long[nchunks + 1];
89+
for (int i = 0; i <= nchunks; i++) {
90+
offsets[i] = offsetsBuf.get(LE_LONG, (long) i * 8);
91+
}
92+
return offsets;
93+
}
94+
95+
private static Array concat(List<Array> chunks, DType dtype, long totalRows, SegmentAllocator arena) {
96+
if (dtype instanceof DType.Primitive pt) {
97+
return concatPrimitive(chunks, pt, dtype, totalRows, arena);
98+
}
99+
if (dtype instanceof DType.Struct struct) {
100+
return concatStruct(chunks, struct, totalRows, arena);
101+
}
102+
throw new VortexException(EncodingId.VORTEX_CHUNKED,
103+
"concat not supported for dtype: " + dtype);
104+
}
105+
106+
private static Array concatPrimitive(
107+
List<Array> chunks, DType.Primitive pt, DType dtype, long totalRows, SegmentAllocator arena
108+
) {
109+
PType ptype = pt.ptype();
110+
MemorySegment combined = arena.allocate(totalRows * ptype.byteSize());
111+
long byteOffset = 0;
112+
for (Array chunk : chunks) {
113+
MemorySegment src = chunk.buffer(0);
114+
MemorySegment.copy(src, 0, combined, byteOffset, src.byteSize());
115+
byteOffset += src.byteSize();
116+
}
117+
MemorySegment ro = combined.asReadOnly();
118+
return switch (ptype) {
119+
case I64, U64 -> new LongArray(dtype, totalRows, ro, ArrayStats.empty());
120+
case I32, U32 -> new IntArray(dtype, totalRows, ro, ArrayStats.empty());
121+
case F64 -> new DoubleArray(dtype, totalRows, ro, ArrayStats.empty());
122+
case F32 -> new FloatArray(dtype, totalRows, ro, ArrayStats.empty());
123+
case I16, U16 -> new ShortArray(dtype, totalRows, ro, ArrayStats.empty());
124+
case I8, U8 -> new ByteArray(dtype, totalRows, ro, ArrayStats.empty());
125+
default -> throw new VortexException(EncodingId.VORTEX_CHUNKED,
126+
"unsupported ptype for concat: " + ptype);
127+
};
128+
}
129+
130+
private static StructArray concatStruct(
131+
List<Array> chunks, DType.Struct struct, long totalRows, SegmentAllocator arena
132+
) {
133+
int nfields = struct.fieldTypes().size();
134+
List<Array> concatFields = new ArrayList<>(nfields);
135+
for (int f = 0; f < nfields; f++) {
136+
DType fieldDtype = struct.fieldTypes().get(f);
137+
List<Array> fieldChunks = new ArrayList<>(chunks.size());
138+
for (Array chunk : chunks) {
139+
fieldChunks.add(((StructArray) chunk).field(f));
140+
}
141+
concatFields.add(concat(fieldChunks, fieldDtype, totalRows, arena));
142+
}
143+
return new StructArray(struct, totalRows, concatFields);
144+
}
145+
}
146+
}

‎core/src/main/resources/META-INF/services/io.github.dfa1.vortex.encoding.Encoding‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,3 +22,4 @@ io.github.dfa1.vortex.encoding.DecimalEncoding
2222
io.github.dfa1.vortex.encoding.DecimalBytePartsEncoding
2323
io.github.dfa1.vortex.encoding.DateTimePartsEncoding
2424
io.github.dfa1.vortex.encoding.ZstdEncoding
25+
io.github.dfa1.vortex.encoding.ChunkedEncoding

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

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
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 org.junit.jupiter.api.Nested;
8+
import org.junit.jupiter.api.Test;
9+
10+
import java.lang.foreign.Arena;
11+
import java.lang.foreign.MemorySegment;
12+
import java.lang.foreign.ValueLayout;
13+
import java.nio.ByteOrder;
14+
15+
import static org.assertj.core.api.Assertions.assertThat;
16+
import static org.assertj.core.api.Assertions.assertThatThrownBy;
17+
18+
class ChunkedEncodingTest {
19+
20+
private static final ValueLayout.OfLong LE_LONG =
21+
ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
22+
23+
@Nested
24+
class Decode {
25+
26+
@Test
27+
void roundTrip_twoChunks_concatenatesValues() {
28+
// Given
29+
long[] chunk0 = {10L, 20L, 30L};
30+
long[] chunk1 = {40L, 50L};
31+
DType i64 = new DType.Primitive(PType.I64, false);
32+
DType u64 = new DType.Primitive(PType.U64, false);
33+
34+
EncodingRegistry registry = EncodingRegistry.empty();
35+
var sut = new ChunkedEncoding();
36+
registry.register(sut);
37+
registry.register(new PrimitiveEncoding());
38+
39+
// Build chunk_offsets segment: [0, 3, 5] as U64 LE
40+
EncodeResult offsetsResult = new PrimitiveEncoding().encode(u64, new long[]{0L, 3L, 5L});
41+
// Build chunk0 segment
42+
EncodeResult chunk0Result = new PrimitiveEncoding().encode(i64, chunk0);
43+
// Build chunk1 segment
44+
EncodeResult chunk1Result = new PrimitiveEncoding().encode(i64, chunk1);
45+
46+
// Collect all buffers: [offsets_buf, chunk0_buf, chunk1_buf]
47+
MemorySegment[] allBufs = {
48+
offsetsResult.buffers().getFirst(),
49+
chunk0Result.buffers().getFirst(),
50+
chunk1Result.buffers().getFirst()
51+
};
52+
53+
// Build ArrayNode tree:
54+
// root: ChunkedEncoding, children=[offsetsNode, chunk0Node, chunk1Node], buffers=[]
55+
// offsetsNode: PrimitiveEncoding, bufferIndices=[0]
56+
// chunk0Node: PrimitiveEncoding, bufferIndices=[1]
57+
// chunk1Node: PrimitiveEncoding, bufferIndices=[2]
58+
ArrayNode offsetsNode = toArrayNode(offsetsResult.rootNode());
59+
ArrayNode chunk0Node = toArrayNode(remapped(chunk0Result.rootNode(), 1));
60+
ArrayNode chunk1Node = toArrayNode(remapped(chunk1Result.rootNode(), 2));
61+
ArrayNode root = new ArrayNode(
62+
EncodingId.VORTEX_CHUNKED, null,
63+
new ArrayNode[]{offsetsNode, chunk0Node, chunk1Node},
64+
new int[]{}, null);
65+
66+
DecodeContext ctx = new DecodeContext(root, i64, 5L, allBufs, registry, Arena.ofAuto());
67+
68+
// When
69+
Array result = sut.decode(ctx);
70+
71+
// Then
72+
assertThat(result.length()).isEqualTo(5);
73+
assertThat(result.buffer(0).get(LE_LONG, 0L)).isEqualTo(10L);
74+
assertThat(result.buffer(0).get(LE_LONG, 8L)).isEqualTo(20L);
75+
assertThat(result.buffer(0).get(LE_LONG, 16L)).isEqualTo(30L);
76+
assertThat(result.buffer(0).get(LE_LONG, 24L)).isEqualTo(40L);
77+
assertThat(result.buffer(0).get(LE_LONG, 32L)).isEqualTo(50L);
78+
}
79+
80+
@Test
81+
void singleChunk_returnsSameValues() {
82+
// Given
83+
long[] data = {1L, 2L, 3L};
84+
DType i64 = new DType.Primitive(PType.I64, false);
85+
DType u64 = new DType.Primitive(PType.U64, false);
86+
87+
EncodingRegistry registry = EncodingRegistry.empty();
88+
registry.register(new ChunkedEncoding());
89+
registry.register(new PrimitiveEncoding());
90+
91+
EncodeResult offsetsResult = new PrimitiveEncoding().encode(u64, new long[]{0L, 3L});
92+
EncodeResult chunkResult = new PrimitiveEncoding().encode(i64, data);
93+
94+
MemorySegment[] allBufs = {
95+
offsetsResult.buffers().getFirst(),
96+
chunkResult.buffers().getFirst()
97+
};
98+
99+
ArrayNode root = new ArrayNode(
100+
EncodingId.VORTEX_CHUNKED, null,
101+
new ArrayNode[]{toArrayNode(offsetsResult.rootNode()), toArrayNode(remapped(chunkResult.rootNode(), 1))},
102+
new int[]{}, null);
103+
104+
DecodeContext ctx = new DecodeContext(root, i64, 3L, allBufs, registry, Arena.ofAuto());
105+
106+
// When
107+
Array result = new ChunkedEncoding().decode(ctx);
108+
109+
// Then
110+
assertThat(result.length()).isEqualTo(3);
111+
for (int i = 0; i < 3; i++) {
112+
assertThat(result.buffer(0).get(LE_LONG, (long) i * 8)).isEqualTo(data[i]);
113+
}
114+
}
115+
116+
@Test
117+
void noChildren_throws() {
118+
// Given
119+
DType i64 = new DType.Primitive(PType.I64, false);
120+
EncodingRegistry registry = EncodingRegistry.empty();
121+
registry.register(new ChunkedEncoding());
122+
ArrayNode root = new ArrayNode(EncodingId.VORTEX_CHUNKED, null, new ArrayNode[]{}, new int[]{}, null);
123+
DecodeContext ctx = new DecodeContext(root, i64, 0L, new MemorySegment[]{}, registry, Arena.ofAuto());
124+
125+
// When / Then
126+
assertThatThrownBy(() -> new ChunkedEncoding().decode(ctx))
127+
.isInstanceOf(VortexException.class)
128+
.hasMessageContaining("at least one child");
129+
}
130+
}
131+
132+
private static ArrayNode toArrayNode(EncodeNode enc) {
133+
ArrayNode[] children = new ArrayNode[enc.children().length];
134+
for (int i = 0; i < children.length; i++) {
135+
children[i] = toArrayNode(enc.children()[i]);
136+
}
137+
return new ArrayNode(enc.encodingId(), enc.metadata(), children, enc.bufferIndices(), null);
138+
}
139+
140+
private static EncodeNode remapped(EncodeNode node, int offset) {
141+
return EncodeNode.remapBufferIndices(node, offset);
142+
}
143+
}

‎reader/src/test/java/io/github/dfa1/vortex/io/VortexHttpReaderIT.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@ void scan_forVortex_decodesAllRows() throws Exception {
9797
"decimal_byte_parts.vortex",
9898
"datetimeparts.vortex",
9999
"zstd.vortex",
100+
"chunked.vortex",
100101
})
101102
void scan_fixture_decodesAllRows(String fixture) throws Exception {
102103
// Given

0 commit comments

Comments
 (0)