Skip to content

Commit 0622b5a

Browse files
dfa1claude
andcommitted
feat(chunked): implement vortex.chunked encode via ChunkedData
Encodes chunk offsets (U64) as first child, each chunk via fallback encoding. ChunkedData carries chunk data + explicit row counts. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 70a46c4 commit 0622b5a

4 files changed

Lines changed: 136 additions & 2 deletions

File tree

TODO.md

Lines changed: 1 addition & 1 deletion
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` | `ChunkedEncoding` | ✅ | ❌ stub | medium | decode: primitive + struct concat; encode stub |
96+
| `vortex.chunked` | `ChunkedEncoding` | ✅ | | medium | decode: primitive + struct concat; encode via ChunkedData |
9797
| `fastlanes.rle` | — | ❌ | ❌ | medium | unblocks `rle.vortex` |
9898
| `vortex.alprd` | — | ❌ | ❌ | medium | unblocks `alprd.vortex` |
9999
| `vortex.decimal` | `DecimalEncoding` | ✅ | ✅ | — | Decimal |
Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
package io.github.dfa1.vortex.encoding;
2+
3+
import java.util.List;
4+
5+
/// Input data for encoding a chunked array.
6+
/// Each element of {@code chunks} is the raw data for one chunk, in the same format
7+
/// the inner encoding expects (e.g. {@code long[]} for I64, {@link StructData} for Struct).
8+
/// {@code chunkLengths[i]} is the row count of {@code chunks.get(i)}.
9+
public record ChunkedData(List<Object> chunks, long[] chunkLengths) {
10+
public ChunkedData {
11+
if (chunks.size() != chunkLengths.length) {
12+
throw new IllegalArgumentException(
13+
"chunks.size() %d != chunkLengths.length %d".formatted(chunks.size(), chunkLengths.length));
14+
}
15+
chunks = List.copyOf(chunks);
16+
chunkLengths = chunkLengths.clone();
17+
}
18+
}

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

Lines changed: 58 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import java.lang.foreign.MemorySegment;
1717
import java.lang.foreign.SegmentAllocator;
1818
import java.lang.foreign.ValueLayout;
19+
import java.nio.ByteBuffer;
1920
import java.nio.ByteOrder;
2021
import java.util.ArrayList;
2122
import java.util.List;
@@ -43,14 +44,70 @@ public boolean accepts(DType dtype) {
4344

4445
@Override
4546
public EncodeResult encode(DType dtype, Object data) {
46-
throw new UnsupportedOperationException("encode not yet implemented for " + encodingId());
47+
return Encoder.encode(dtype, (ChunkedData) data);
4748
}
4849

4950
@Override
5051
public Array decode(DecodeContext ctx) {
5152
return Decoder.decode(ctx);
5253
}
5354

55+
private static final class Encoder {
56+
57+
private static final List<Encoding> FALLBACK = List.of(
58+
new PrimitiveEncoding(), new VarBinEncoding(), new BoolEncoding(),
59+
new NullEncoding(), new ByteBoolEncoding(), new StructEncoding());
60+
61+
static EncodeResult encode(DType dtype, ChunkedData data) {
62+
List<Object> chunks = data.chunks();
63+
long[] chunkLengths = data.chunkLengths();
64+
int nchunks = chunks.size();
65+
if (nchunks == 0) {
66+
throw new VortexException(EncodingId.VORTEX_CHUNKED, "at least one chunk required");
67+
}
68+
69+
// Build cumulative offsets: [0, len0, len0+len1, ...]
70+
long[] offsets = new long[nchunks + 1];
71+
offsets[0] = 0;
72+
for (int i = 0; i < nchunks; i++) {
73+
offsets[i + 1] = offsets[i] + chunkLengths[i];
74+
}
75+
76+
// Encode offsets child (U64 primitive)
77+
DType u64 = new DType.Primitive(PType.U64, false);
78+
EncodeResult offsetsResult = new PrimitiveEncoding().encode(u64, offsets);
79+
80+
List<MemorySegment> allBuffers = new ArrayList<>(offsetsResult.buffers());
81+
EncodeNode[] children = new EncodeNode[nchunks + 1];
82+
children[0] = offsetsResult.rootNode();
83+
84+
// Encode each chunk
85+
Encoding inner = findEncoding(dtype);
86+
for (int i = 0; i < nchunks; i++) {
87+
EncodeResult chunkResult = inner.encode(dtype, chunks.get(i));
88+
int bufOffset = allBuffers.size();
89+
children[i + 1] = EncodeNode.remapBufferIndices(chunkResult.rootNode(), bufOffset);
90+
allBuffers.addAll(chunkResult.buffers());
91+
}
92+
93+
EncodeNode root = new EncodeNode(
94+
EncodingId.VORTEX_CHUNKED,
95+
ByteBuffer.wrap(new byte[0]),
96+
children,
97+
new int[]{});
98+
return new EncodeResult(root, List.copyOf(allBuffers), null, null);
99+
}
100+
101+
private static Encoding findEncoding(DType dtype) {
102+
for (Encoding enc : FALLBACK) {
103+
if (enc.accepts(dtype)) {
104+
return enc;
105+
}
106+
}
107+
throw new UnsupportedOperationException("no fallback encoding for dtype: " + dtype);
108+
}
109+
}
110+
54111
private static final class Decoder {
55112

56113
private static final ValueLayout.OfLong LE_LONG =

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

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import java.lang.foreign.MemorySegment;
1212
import java.lang.foreign.ValueLayout;
1313
import java.nio.ByteOrder;
14+
import java.util.List;
1415

1516
import static org.assertj.core.api.Assertions.assertThat;
1617
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -20,6 +21,64 @@ class ChunkedEncodingTest {
2021
private static final ValueLayout.OfLong LE_LONG =
2122
ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
2223

24+
@Nested
25+
class Encode {
26+
27+
@Test
28+
void roundTrip_twoChunks_i64_preservesValues() {
29+
// Given
30+
long[] chunk0 = {10L, 20L, 30L};
31+
long[] chunk1 = {40L, 50L};
32+
DType i64 = new DType.Primitive(PType.I64, false);
33+
var sut = new ChunkedEncoding();
34+
EncodingRegistry registry = EncodingRegistry.empty();
35+
registry.register(sut);
36+
registry.register(new PrimitiveEncoding());
37+
ChunkedData data = new ChunkedData(List.of(chunk0, chunk1), new long[]{3, 2});
38+
39+
// When
40+
EncodeResult encoded = sut.encode(i64, data);
41+
DecodeContext ctx = EncodeTestHelper.toDecodeContext(encoded, 5L, i64, registry);
42+
Array result = sut.decode(ctx);
43+
44+
// Then
45+
assertThat(result.length()).isEqualTo(5);
46+
assertThat(result.buffer(0).get(LE_LONG, 0L)).isEqualTo(10L);
47+
assertThat(result.buffer(0).get(LE_LONG, 8L)).isEqualTo(20L);
48+
assertThat(result.buffer(0).get(LE_LONG, 16L)).isEqualTo(30L);
49+
assertThat(result.buffer(0).get(LE_LONG, 24L)).isEqualTo(40L);
50+
assertThat(result.buffer(0).get(LE_LONG, 32L)).isEqualTo(50L);
51+
}
52+
53+
@Test
54+
void encodeNode_hasNoDirectBuffers_offsetsAsFirstChild() {
55+
// Given
56+
long[] chunk0 = {1L, 2L};
57+
DType i64 = new DType.Primitive(PType.I64, false);
58+
var sut = new ChunkedEncoding();
59+
ChunkedData data = new ChunkedData(List.of(chunk0), new long[]{2});
60+
61+
// When
62+
EncodeResult result = sut.encode(i64, data);
63+
64+
// Then
65+
assertThat(result.rootNode().bufferIndices()).isEmpty();
66+
assertThat(result.rootNode().children()).hasSize(2); // offsets + 1 chunk
67+
assertThat(result.buffers()).hasSize(2); // offsets buf + chunk buf
68+
}
69+
70+
@Test
71+
void mismatchedLengths_throws() {
72+
// Given
73+
var sut = new ChunkedEncoding();
74+
DType i64 = new DType.Primitive(PType.I64, false);
75+
76+
// When / Then
77+
assertThatThrownBy(() -> new ChunkedData(List.of(new long[]{1L}), new long[]{1, 2}))
78+
.isInstanceOf(IllegalArgumentException.class);
79+
}
80+
}
81+
2382
@Nested
2483
class Decode {
2584

0 commit comments

Comments
 (0)