Skip to content

Commit 59302db

Browse files
dfa1claude
andcommitted
feat(encoding): implement vortex.fixed_size_list encode+decode
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 94db18f commit 59302db

8 files changed

Lines changed: 315 additions & 5 deletions

File tree

TODO.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@
101101
| `vortex.datetimeparts` | `DateTimePartsEncoding` | ✅ | ❌ stub | — | Timestamp parts |
102102
| `vortex.list` | — | ❌ | ❌ | hard | two children: offsets (i32/i64, len N+1) + values; needs `ListArray`; offsets are cascadable (use `decodeChildAs`); chunked requires offset re-basing; unblocks `list.vortex` |
103103
| `vortex.listview` | — | ❌ | ❌ | hard | unblocks `listview.vortex` |
104-
| `vortex.fixed_size_list` | | ❌ | | medium | one child: values only, len = N*fixedSize; no offsets, no re-basing; needs `FixedSizeListArray` + child dispatch; easier than `vortex.list`; unblocks `fixed_size_list.vortex` |
104+
| `vortex.fixed_size_list` | `FixedSizeListEncoding` | ✅ | ✅ | | one child: flat elements; no offsets |
105105
| `vortex.zstd` | `ZstdEncoding` | ✅ | ✅ | — | Primitive, Utf8, Binary (no dict, no nullable); uses airlift/aircompressor |
106106
107107
### `vortex.zstd` known limitations
@@ -148,7 +148,7 @@
148148
| `datetimeparts.vortex` | ✅ | |
149149
| `list.vortex` | ❌ | `vortex.list` + list array model |
150150
| `listview.vortex` | ❌ | `vortex.listview` missing |
151-
| `fixed_size_list.vortex` | | `vortex.fixed_size_list` missing |
151+
| `fixed_size_list.vortex` | | |
152152
| `zstd.vortex` | ✅ | |
153153
| `tpch_lineitem.compact.vortex` | ❌ | `vortex.pco` |
154154
| `tpch_lineitem.regular.vortex` | ❌ | `vortex.pco` |

core/src/main/java/io/github/dfa1/vortex/core/array/Array.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@
1111
/// Buffers are `MemorySegment` slices backed by the memory-mapped file; lifetime
1212
/// is tied to the `VortexFile`'s Arena.
1313
public sealed interface Array
14-
permits BoolArray, ByteArray, DoubleArray, EmptyArray, Float16Array, FloatArray,
15-
GenericArray, IntArray, LongArray, NullArray, ShortArray, StructArray,
14+
permits BoolArray, ByteArray, DoubleArray, EmptyArray, FixedSizeListArray, Float16Array,
15+
FloatArray, GenericArray, IntArray, LongArray, NullArray, ShortArray, StructArray,
1616
VarBinArray {
1717

1818
long length();
Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
package io.github.dfa1.vortex.core.array;
2+
3+
import io.github.dfa1.vortex.core.ArrayStats;
4+
import io.github.dfa1.vortex.core.DType;
5+
6+
/// Decoded fixed-size list array: holds a flat elements [Array] of length {@code outerLen * fixedSize}.
7+
///
8+
/// <p>List {@code i} covers elements {@code [i*fixedSize, (i+1)*fixedSize)}.
9+
public final class FixedSizeListArray implements Array {
10+
11+
private final DType.FixedSizeList dtype;
12+
private final long outerLen;
13+
private final Array elements;
14+
15+
public FixedSizeListArray(DType.FixedSizeList dtype, long outerLen, Array elements) {
16+
this.dtype = dtype;
17+
this.outerLen = outerLen;
18+
this.elements = elements;
19+
}
20+
21+
@Override
22+
public long length() {
23+
return outerLen;
24+
}
25+
26+
@Override
27+
public DType dtype() {
28+
return dtype;
29+
}
30+
31+
public ArrayStats stats() {
32+
return ArrayStats.empty();
33+
}
34+
35+
public Array elements() {
36+
return elements;
37+
}
38+
39+
public int fixedSize() {
40+
return dtype.fixedSize();
41+
}
42+
43+
@Override
44+
public Array child(int i) {
45+
if (i != 0) {
46+
throw new ArrayIndexOutOfBoundsException("FixedSizeListArray child index " + i);
47+
}
48+
return elements;
49+
}
50+
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,8 @@ public enum EncodingId {
4040
VORTEX_DECIMAL("vortex.decimal"),
4141
VORTEX_DECIMAL_BYTE_PARTS("vortex.decimal_byte_parts"),
4242
VORTEX_DATETIMEPARTS("vortex.datetimeparts"),
43-
VORTEX_ZSTD("vortex.zstd");
43+
VORTEX_ZSTD("vortex.zstd"),
44+
VORTEX_FIXED_SIZE_LIST("vortex.fixed_size_list");
4445

4546
// O(1) access to EncodingId by its string representation
4647
private static final Map<String, EncodingId> LOOKUP = Stream.of(EncodingId.values())
Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
package io.github.dfa1.vortex.encoding;
2+
3+
/// Input data for encoding a {@code vortex.fixed_size_list} column.
4+
///
5+
/// <p>{@code elements} is the flat array of inner values; its length must equal
6+
/// {@code outerLen * fixedSize}. The element type is determined by the DType passed to
7+
/// {@link FixedSizeListEncoding#encode}.
8+
public record FixedSizeListData(Object elements, long outerLen) {
9+
}
Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
package io.github.dfa1.vortex.encoding;
2+
3+
import io.github.dfa1.vortex.core.DType;
4+
import io.github.dfa1.vortex.core.VortexException;
5+
import io.github.dfa1.vortex.core.array.Array;
6+
import io.github.dfa1.vortex.core.array.FixedSizeListArray;
7+
8+
import java.lang.foreign.MemorySegment;
9+
import java.nio.ByteBuffer;
10+
import java.util.ArrayList;
11+
import java.util.List;
12+
13+
/// Encoder/decoder for {@code vortex.fixed_size_list}.
14+
///
15+
/// <p>Wire format (per Rust vtable):
16+
/// <ul>
17+
/// <li>Buffers: 0
18+
/// <li>Metadata: empty byte array
19+
/// <li>Children: 1 or 2.
20+
/// {@code children[0]} is the flat elements array (len = outerLen * fixedSize, dtype = elementType).
21+
/// {@code children[1]}, when present, is the validity (bool) array (len = outerLen).
22+
/// </ul>
23+
///
24+
/// <p>The {@code fixedSize} is read directly from the {@link DType.FixedSizeList} dtype —
25+
/// it is not stored in metadata.
26+
public final class FixedSizeListEncoding implements Encoding {
27+
28+
@Override
29+
public EncodingId encodingId() {
30+
return EncodingId.VORTEX_FIXED_SIZE_LIST;
31+
}
32+
33+
@Override
34+
public boolean accepts(DType dtype) {
35+
return dtype instanceof DType.FixedSizeList;
36+
}
37+
38+
@Override
39+
public EncodeResult encode(DType dtype, Object data) {
40+
return Encoder.encode((DType.FixedSizeList) dtype, (FixedSizeListData) data);
41+
}
42+
43+
@Override
44+
public Array decode(DecodeContext ctx) {
45+
return Decoder.decode(ctx);
46+
}
47+
48+
private static final class Encoder {
49+
50+
private static final List<Encoding> FALLBACK = List.of(
51+
new PrimitiveEncoding(), new VarBinEncoding(), new BoolEncoding(),
52+
new NullEncoding(), new ByteBoolEncoding());
53+
54+
static EncodeResult encode(DType.FixedSizeList dtype, FixedSizeListData data) {
55+
DType elementType = dtype.elementType();
56+
Encoding inner = findEncoding(elementType);
57+
58+
EncodeResult elemResult = inner.encode(elementType, data.elements());
59+
60+
List<MemorySegment> allBuffers = new ArrayList<>(elemResult.buffers());
61+
EncodeNode elemNode = EncodeNode.remapBufferIndices(elemResult.rootNode(), 0);
62+
63+
EncodeNode root = new EncodeNode(
64+
EncodingId.VORTEX_FIXED_SIZE_LIST,
65+
ByteBuffer.wrap(new byte[0]),
66+
new EncodeNode[]{elemNode},
67+
new int[]{});
68+
return new EncodeResult(root, List.copyOf(allBuffers), null, null);
69+
}
70+
71+
private static Encoding findEncoding(DType dtype) {
72+
for (Encoding enc : FALLBACK) {
73+
if (enc.accepts(dtype)) {
74+
return enc;
75+
}
76+
}
77+
throw new UnsupportedOperationException("no fallback encoding for dtype: " + dtype);
78+
}
79+
}
80+
81+
private static final class Decoder {
82+
83+
static Array decode(DecodeContext ctx) {
84+
if (!(ctx.dtype() instanceof DType.FixedSizeList fsl)) {
85+
throw new VortexException(EncodingId.VORTEX_FIXED_SIZE_LIST,
86+
"expected DType.FixedSizeList, got " + ctx.dtype());
87+
}
88+
89+
int nchildren = ctx.node().children().length;
90+
if (nchildren < 1 || nchildren > 2) {
91+
throw new VortexException(EncodingId.VORTEX_FIXED_SIZE_LIST,
92+
"expected 1 or 2 children, got " + nchildren);
93+
}
94+
95+
long outerLen = ctx.rowCount();
96+
long elemLen = outerLen * fsl.fixedSize();
97+
DType elementType = fsl.elementType();
98+
99+
ArrayNode elemNode = ctx.node().children()[0];
100+
var elemCtx = new DecodeContext(
101+
elemNode, elementType, elemLen,
102+
ctx.segmentBuffers(), ctx.registry(), ctx.arena());
103+
Array elements = ctx.registry().decode(elemCtx);
104+
105+
return new FixedSizeListArray(fsl, outerLen, elements);
106+
}
107+
}
108+
}

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
@@ -23,3 +23,4 @@ io.github.dfa1.vortex.encoding.DecimalBytePartsEncoding
2323
io.github.dfa1.vortex.encoding.DateTimePartsEncoding
2424
io.github.dfa1.vortex.encoding.ZstdEncoding
2525
io.github.dfa1.vortex.encoding.ChunkedEncoding
26+
io.github.dfa1.vortex.encoding.FixedSizeListEncoding
Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
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.array.FixedSizeListArray;
7+
import io.github.dfa1.vortex.core.array.IntArray;
8+
import org.junit.jupiter.api.Nested;
9+
import org.junit.jupiter.api.Test;
10+
11+
import java.lang.foreign.Arena;
12+
import java.lang.foreign.MemorySegment;
13+
14+
import static org.assertj.core.api.Assertions.assertThat;
15+
import static org.assertj.core.api.Assertions.assertThatThrownBy;
16+
17+
class FixedSizeListEncodingTest {
18+
19+
private static final DType I32 = new DType.Primitive(PType.I32, false);
20+
21+
private static ArrayNode toArrayNode(EncodeNode node) {
22+
ArrayNode[] children = new ArrayNode[node.children().length];
23+
for (int i = 0; i < children.length; i++) {
24+
children[i] = toArrayNode(node.children()[i]);
25+
}
26+
return new ArrayNode(node.encodingId(), node.metadata(), children, node.bufferIndices(), ArrayStats.empty());
27+
}
28+
29+
private static EncodingRegistry registry() {
30+
EncodingRegistry r = EncodingRegistry.empty();
31+
r.register(new FixedSizeListEncoding());
32+
r.register(new PrimitiveEncoding());
33+
return r;
34+
}
35+
36+
@Nested
37+
class Encode {
38+
39+
@Test
40+
void accepts_fixedSizeListDtype_true() {
41+
// Given
42+
FixedSizeListEncoding sut = new FixedSizeListEncoding();
43+
DType.FixedSizeList dtype = new DType.FixedSizeList(I32, 3, false);
44+
45+
// When / Then
46+
assertThat(sut.accepts(dtype)).isTrue();
47+
}
48+
49+
@Test
50+
void accepts_primitiveDtype_false() {
51+
// Given
52+
FixedSizeListEncoding sut = new FixedSizeListEncoding();
53+
54+
// When / Then
55+
assertThat(sut.accepts(I32)).isFalse();
56+
}
57+
58+
@Test
59+
void encode_producesOneChild_noBuffers() {
60+
// Given
61+
DType.FixedSizeList dtype = new DType.FixedSizeList(I32, 2, false);
62+
int[] elements = {1, 2, 3, 4};
63+
FixedSizeListData data = new FixedSizeListData(elements, 2);
64+
FixedSizeListEncoding sut = new FixedSizeListEncoding();
65+
66+
// When
67+
EncodeResult result = sut.encode(dtype, data);
68+
69+
// Then
70+
assertThat(result.rootNode().encodingId()).isEqualTo(EncodingId.VORTEX_FIXED_SIZE_LIST);
71+
assertThat(result.rootNode().bufferIndices()).isEmpty();
72+
assertThat(result.rootNode().children()).hasSize(1);
73+
}
74+
}
75+
76+
@Nested
77+
class Decode {
78+
79+
@Test
80+
void roundTrip_i32Elements_preservesValues() {
81+
// Given
82+
DType.FixedSizeList dtype = new DType.FixedSizeList(I32, 3, false);
83+
int[] elements = {10, 20, 30, 40, 50, 60};
84+
FixedSizeListData data = new FixedSizeListData(elements, 2);
85+
FixedSizeListEncoding sut = new FixedSizeListEncoding();
86+
87+
// When
88+
EncodeResult result = sut.encode(dtype, data);
89+
MemorySegment[] bufs = result.buffers().toArray(MemorySegment[]::new);
90+
DecodeContext ctx = new DecodeContext(
91+
toArrayNode(result.rootNode()), dtype, 2, bufs, registry(), Arena.global());
92+
FixedSizeListArray decoded = (FixedSizeListArray) sut.decode(ctx);
93+
94+
// Then
95+
assertThat(decoded.length()).isEqualTo(2);
96+
assertThat(decoded.fixedSize()).isEqualTo(3);
97+
IntArray elems = (IntArray) decoded.elements();
98+
assertThat(elems.length()).isEqualTo(6);
99+
for (int i = 0; i < elements.length; i++) {
100+
assertThat(elems.getInt(i)).isEqualTo(elements[i]);
101+
}
102+
}
103+
104+
@Test
105+
void roundTrip_fixedSizeOne_preservesValues() {
106+
// Given
107+
DType.FixedSizeList dtype = new DType.FixedSizeList(I32, 1, false);
108+
int[] elements = {7, 8, 9};
109+
FixedSizeListData data = new FixedSizeListData(elements, 3);
110+
FixedSizeListEncoding sut = new FixedSizeListEncoding();
111+
112+
// When
113+
EncodeResult result = sut.encode(dtype, data);
114+
MemorySegment[] bufs = result.buffers().toArray(MemorySegment[]::new);
115+
DecodeContext ctx = new DecodeContext(
116+
toArrayNode(result.rootNode()), dtype, 3, bufs, registry(), Arena.global());
117+
FixedSizeListArray decoded = (FixedSizeListArray) sut.decode(ctx);
118+
119+
// Then
120+
assertThat(decoded.length()).isEqualTo(3);
121+
assertThat(decoded.fixedSize()).isEqualTo(1);
122+
IntArray elems = (IntArray) decoded.elements();
123+
for (int i = 0; i < elements.length; i++) {
124+
assertThat(elems.getInt(i)).isEqualTo(elements[i]);
125+
}
126+
}
127+
128+
@Test
129+
void decode_wrongDtype_throws() {
130+
// Given
131+
FixedSizeListEncoding sut = new FixedSizeListEncoding();
132+
ArrayNode node = new ArrayNode(EncodingId.VORTEX_FIXED_SIZE_LIST, null,
133+
new ArrayNode[0], new int[0], ArrayStats.empty());
134+
DecodeContext ctx = new DecodeContext(node, I32, 0, new MemorySegment[0], registry(), Arena.global());
135+
136+
// When / Then
137+
assertThatThrownBy(() -> sut.decode(ctx))
138+
.hasMessageContaining("DType.FixedSizeList");
139+
}
140+
}
141+
}

0 commit comments

Comments
 (0)