Skip to content

Commit 4870e93

Browse files
dfa1claude
dfa1
and
claude
committed
refactor: add Arena to DecodeContext; eliminate long[] in BitpackedCodec
DecodeContext gains an Arena field (propagated via decodeChild/decodeChildAs) to enable future arena-backed allocations in the decode hot path. Currently ScanIterator passes Arena.global() as a pass-through. BitpackedCodec: eliminate the intermediate long[] + fromLongs byte[] copy by writing directly into a MemorySegment.ofArray(byte[]) via writeWordToSeg. Saves ~400KB per 50K-row chunk (the long[] allocation was the biggest GC waste in the decode chain). ForCodec / AlpCodec: replace ByteBuffer positional writes with MemorySegment.set and running-offset counters (same allocation, cleaner code, no position state). SparseCodec / DictCodec / RunEndCodec: propagate Arena in decodeChildAs. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent df6530c commit 4870e93

12 files changed

Lines changed: 83 additions & 88 deletions

File tree

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

Lines changed: 18 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -94,26 +94,26 @@ private Array decodeF64(DecodeContext ctx, EncodingProtos.ALPMetadata meta, int
9494
double factor = F10_F64[expF] * IF10_F64[expE];
9595

9696
byte[] outBytes = new byte[Math.toIntExact(n * 8)];
97-
ByteBuffer out = ByteBuffer.wrap(outBytes).order(ByteOrder.LITTLE_ENDIAN);
97+
MemorySegment out = MemorySegment.ofArray(outBytes);
9898
var srcLayout = ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
99+
var dstLayout = ValueLayout.JAVA_DOUBLE_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
99100
MemorySegment src = encoded.buffer(0);
100101
// Strength-reduce: running byte offset instead of i*8 per iteration.
101102
for (long off = 0, end = n * 8; off < end; off += 8) {
102103
long encVal = src.get(srcLayout, off);
103104
double dec = (double) encVal * factor;
104-
out.putDouble(dec);
105+
out.set(dstLayout, off, dec);
105106
}
106107

107108
if (meta.hasPatches()) {
108-
applyPatchesF64(ctx, meta.getPatches(), outBytes, n);
109+
applyPatchesF64(ctx, meta.getPatches(), out, n);
109110
}
110111

111-
return new Array(ctx.dtype(), n,
112-
new MemorySegment[]{MemorySegment.ofArray(outBytes)}, new Array[0], ArrayStats.empty());
112+
return new Array(ctx.dtype(), n, new MemorySegment[]{out}, new Array[0], ArrayStats.empty());
113113
}
114114

115115
private void applyPatchesF64(DecodeContext ctx, EncodingProtos.PatchesMetadata pm,
116-
byte[] outBytes, long n) {
116+
MemorySegment out, long n) {
117117
long numPatches = pm.getLen();
118118
long offset = pm.getOffset();
119119
PType idxPtype = ptypeFromProto(pm.getIndicesPtype());
@@ -124,17 +124,16 @@ private void applyPatchesF64(DecodeContext ctx, EncodingProtos.PatchesMetadata p
124124
Array idxArr = decodeChildAs(ctx, 1, idxDtype, numPatches);
125125
Array valArr = decodeChildAs(ctx, 2, valDtype, numPatches);
126126

127-
var idxLayout = unsignedLayout(idxPtype);
128127
var valLayout = ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
128+
var dstLayout = ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
129129

130-
ByteBuffer out = ByteBuffer.wrap(outBytes).order(ByteOrder.LITTLE_ENDIAN);
131130
MemorySegment idxSeg = idxArr.buffer(0);
132131
MemorySegment valSeg = valArr.buffer(0);
133132

134133
for (long i = 0; i < numPatches; i++) {
135-
long absIdx = readUnsigned(idxSeg, i, idxPtype) - offset;
134+
long absIdx = readUnsigned(idxSeg, i, idxPtype) - offset;
136135
long rawBits = valSeg.get(valLayout, i * 8);
137-
out.putLong((int) (absIdx * 8), rawBits);
136+
out.set(dstLayout, absIdx * 8, rawBits);
138137
}
139138
}
140139

@@ -147,25 +146,25 @@ private Array decodeF32(DecodeContext ctx, EncodingProtos.ALPMetadata meta, int
147146
float factor = F10_F32[expF] * IF10_F32[expE];
148147

149148
byte[] outBytes = new byte[Math.toIntExact(n * 4)];
150-
ByteBuffer out = ByteBuffer.wrap(outBytes).order(ByteOrder.LITTLE_ENDIAN);
149+
MemorySegment out = MemorySegment.ofArray(outBytes);
151150
var srcLayout = ValueLayout.JAVA_INT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
151+
var dstLayout = ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
152152
MemorySegment src = encoded.buffer(0);
153153
for (long off = 0, end = n * 4; off < end; off += 4) {
154154
int encVal = src.get(srcLayout, off);
155155
float dec = (float) encVal * factor;
156-
out.putFloat(dec);
156+
out.set(dstLayout, off, dec);
157157
}
158158

159159
if (meta.hasPatches()) {
160-
applyPatchesF32(ctx, meta.getPatches(), outBytes, n);
160+
applyPatchesF32(ctx, meta.getPatches(), out, n);
161161
}
162162

163-
return new Array(ctx.dtype(), n,
164-
new MemorySegment[]{MemorySegment.ofArray(outBytes)}, new Array[0], ArrayStats.empty());
163+
return new Array(ctx.dtype(), n, new MemorySegment[]{out}, new Array[0], ArrayStats.empty());
165164
}
166165

167166
private void applyPatchesF32(DecodeContext ctx, EncodingProtos.PatchesMetadata pm,
168-
byte[] outBytes, long n) {
167+
MemorySegment out, long n) {
169168
long numPatches = pm.getLen();
170169
long offset = pm.getOffset();
171170
PType idxPtype = ptypeFromProto(pm.getIndicesPtype());
@@ -177,15 +176,15 @@ private void applyPatchesF32(DecodeContext ctx, EncodingProtos.PatchesMetadata p
177176
Array valArr = decodeChildAs(ctx, 2, valDtype, numPatches);
178177

179178
var valLayout = ValueLayout.JAVA_INT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
179+
var dstLayout = ValueLayout.JAVA_INT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
180180

181-
ByteBuffer out = ByteBuffer.wrap(outBytes).order(ByteOrder.LITTLE_ENDIAN);
182181
MemorySegment idxSeg = idxArr.buffer(0);
183182
MemorySegment valSeg = valArr.buffer(0);
184183

185184
for (long i = 0; i < numPatches; i++) {
186185
long absIdx = readUnsigned(idxSeg, i, idxPtype) - offset;
187186
int rawBits = valSeg.get(valLayout, i * 4);
188-
out.putInt((int) (absIdx * 4), rawBits);
187+
out.set(dstLayout, absIdx * 4, rawBits);
189188
}
190189
}
191190

@@ -194,7 +193,7 @@ private void applyPatchesF32(DecodeContext ctx, EncodingProtos.PatchesMetadata p
194193
private static Array decodeChildAs(DecodeContext parent, int childIdx, DType dtype, long rowCount) {
195194
ArrayNode childNode = parent.node().children()[childIdx];
196195
DecodeContext childCtx = new DecodeContext(
197-
childNode, dtype, rowCount, parent.segmentBuffers(), parent.registry());
196+
childNode, dtype, rowCount, parent.segmentBuffers(), parent.registry(), parent.arena());
198197
return parent.registry().decode(childCtx);
199198
}
200199

@@ -211,10 +210,6 @@ private static long readUnsigned(MemorySegment seg, long i, PType ptype) {
211210
};
212211
}
213212

214-
private static ValueLayout.OfLong unsignedLayout(PType ptype) {
215-
return ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
216-
}
217-
218213
private static PType ptypeFromProto(DTypeProtos.PType proto) {
219214
return PType.values()[proto.getNumber()];
220215
}

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

Lines changed: 19 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -117,10 +117,12 @@ public Array decode(DecodeContext ctx) {
117117
long rowCount = ctx.rowCount();
118118

119119
MemorySegment packed = ctx.buffer(0);
120-
long[] longs = fastlanesUnpack(packed, bitWidth, offset, typeBits, rowCount);
120+
int wordBytes = typeBits / 8;
121+
MemorySegment output = MemorySegment.ofArray(new byte[(int) (rowCount * wordBytes)]);
122+
fastlanesUnpackToSeg(packed, bitWidth, offset, typeBits, rowCount, output);
121123

122124
return new Array(ctx.dtype(), rowCount,
123-
new MemorySegment[]{fromLongs(longs, ptype)}, new Array[0], ArrayStats.empty());
125+
new MemorySegment[]{output}, new Array[0], ArrayStats.empty());
124126
}
125127

126128
// ── FastLanes pack ────────────────────────────────────────────────────────
@@ -171,11 +173,11 @@ private static ByteBuffer packFastLanes(long[] values, int n, int bitWidth, int
171173

172174
// ── FastLanes unpack ──────────────────────────────────────────────────────
173175

174-
private static long[] fastlanesUnpack(
175-
MemorySegment buf, int bitWidth, int offset, int typeBits, long rowCount) {
176-
long[] output = new long[(int) rowCount];
176+
private static void fastlanesUnpackToSeg(
177+
MemorySegment buf, int bitWidth, int offset, int typeBits, long rowCount,
178+
MemorySegment output) {
177179
if (bitWidth == 0) {
178-
return output;
180+
return; // output already zero-initialized by arena.allocate()
179181
}
180182

181183
int LANES = 1024 / typeBits;
@@ -225,11 +227,10 @@ private static long[] fastlanesUnpack(
225227
value = (src >>> shift) & bitMask;
226228
}
227229

228-
output[logicalIdx] = value;
230+
writeWordToSeg(output, (long) logicalIdx * wordBytes, value, typeBits);
229231
}
230232
}
231233
}
232-
return output;
233234
}
234235

235236
// ── Buffer helpers ────────────────────────────────────────────────────────
@@ -294,6 +295,16 @@ private static long readWordFromSeg(MemorySegment seg, long off, int typeBits) {
294295
};
295296
}
296297

298+
private static void writeWordToSeg(MemorySegment seg, long off, long value, int typeBits) {
299+
switch (typeBits) {
300+
case 8 -> seg.set(ValueLayout.JAVA_BYTE, off, (byte) value);
301+
case 16 -> seg.set(ValueLayout.JAVA_SHORT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN), off, (short) value);
302+
case 32 -> seg.set(ValueLayout.JAVA_INT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN), off, (int) value);
303+
case 64 -> seg.set(ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN), off, value);
304+
default -> throw new IllegalArgumentException("unsupported typeBits: " + typeBits);
305+
}
306+
}
307+
297308
// ── Type conversion ───────────────────────────────────────────────────────
298309

299310
private static long[] toLongs(Object data, PType ptype) {
@@ -339,23 +350,6 @@ private static long[] toLongs(Object data, PType ptype) {
339350
};
340351
}
341352

342-
private static MemorySegment fromLongs(long[] longs, PType ptype) {
343-
int n = longs.length;
344-
int elemSize = ptype.byteSize();
345-
byte[] bytes = new byte[n * elemSize];
346-
ByteBuffer bb = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN);
347-
for (long v : longs) {
348-
switch (ptype) {
349-
case I8, U8 -> bb.put((byte) v);
350-
case I16, U16 -> bb.putShort((short) v);
351-
case I32, U32 -> bb.putInt((int) v);
352-
case I64, U64 -> bb.putLong(v);
353-
default -> throw new UnsupportedOperationException("unsupported ptype: " + ptype);
354-
}
355-
}
356-
return MemorySegment.ofArray(bytes);
357-
}
358-
359353
// ── Misc helpers ──────────────────────────────────────────────────────────
360354

361355
private static long typeMask(int typeBits) {

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,24 +3,28 @@
33
import io.github.dfa1.vortex.core.DType;
44

55
import java.io.IOException;
6+
import java.lang.foreign.Arena;
67
import java.lang.foreign.MemorySegment;
78
import java.nio.ByteBuffer;
89

910
/// Decoding context passed to each [Decoder].
1011
///
1112
/// Buffers are `MemorySegment` slices materialized from the file's segment table;
1213
/// children are decoded recursively via [#decodeChild(int)].
14+
/// The arena is scoped to one chunk epoch — all decode output allocated from it is
15+
/// valid until the next [io.github.dfa1.vortex.scan.ScanIterator#hasNext()] call.
1316
public record DecodeContext(
1417
ArrayNode node,
1518
DType dtype,
1619
long rowCount,
1720
MemorySegment[] segmentBuffers,
18-
DecoderRegistry registry
21+
DecoderRegistry registry,
22+
Arena arena
1923
) {
20-
/// Recursively decode child `i` using the same segment buffers and registry.
24+
/// Recursively decode child `i` using the same segment buffers, registry and arena.
2125
public Array decodeChild(int i) throws IOException {
2226
ArrayNode child = node.children()[i];
23-
var childCtx = new DecodeContext(child, dtype, rowCount, segmentBuffers, registry);
27+
var childCtx = new DecodeContext(child, dtype, rowCount, segmentBuffers, registry, arena);
2428
return registry.decode(childCtx);
2529
}
2630

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -156,7 +156,7 @@ private Array decodeRustProto(DecodeContext ctx, byte[] metaBytes) {
156156
private static Array decodeChildAs(DecodeContext parent, int childIdx, DType dtype, long rowCount) {
157157
ArrayNode childNode = parent.node().children()[childIdx];
158158
DecodeContext childCtx = new DecodeContext(
159-
childNode, dtype, rowCount, parent.segmentBuffers(), parent.registry());
159+
childNode, dtype, rowCount, parent.segmentBuffers(), parent.registry(), parent.arena());
160160
return parent.registry().decode(childCtx);
161161
}
162162

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

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -72,40 +72,40 @@ private static long referenceValue(ScalarProtos.ScalarValue scalar) {
7272
}
7373

7474
private static MemorySegment applyReference(MemorySegment src, long n, PType ptype, long ref) {
75-
int elemBytes = ptype.byteSize();
76-
byte[] bytes = new byte[(int) (n * elemBytes)];
77-
ByteBuffer dst = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN);
75+
int wordBytes = ptype.byteSize();
76+
byte[] bytes = new byte[(int) (n * wordBytes)];
77+
MemorySegment dst = MemorySegment.ofArray(bytes);
7878
switch (ptype) {
7979
case I8, U8 -> {
80-
for (long i = 0; i < n; i++) {
81-
byte v = src.get(ValueLayout.JAVA_BYTE, i);
82-
dst.put((byte) (v + (byte) ref));
80+
for (long off = 0, end = n; off < end; off++) {
81+
byte v = src.get(ValueLayout.JAVA_BYTE, off);
82+
dst.set(ValueLayout.JAVA_BYTE, off, (byte) (v + (byte) ref));
8383
}
8484
}
8585
case I16, U16 -> {
8686
var layout = ValueLayout.JAVA_SHORT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
87-
for (long i = 0; i < n; i++) {
88-
short v = src.get(layout, i * 2);
89-
dst.putShort((short) (v + (short) ref));
87+
for (long off = 0, end = n * 2; off < end; off += 2) {
88+
short v = src.get(layout, off);
89+
dst.set(layout, off, (short) (v + (short) ref));
9090
}
9191
}
9292
case I32, U32 -> {
9393
var layout = ValueLayout.JAVA_INT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
94-
for (long i = 0; i < n; i++) {
95-
int v = src.get(layout, i * 4);
96-
dst.putInt(v + (int) ref);
94+
for (long off = 0, end = n * 4; off < end; off += 4) {
95+
int v = src.get(layout, off);
96+
dst.set(layout, off, v + (int) ref);
9797
}
9898
}
9999
case I64, U64 -> {
100100
var layout = ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
101-
for (long i = 0; i < n; i++) {
102-
long v = src.get(layout, i * 8);
103-
dst.putLong(v + ref);
101+
for (long off = 0, end = n * 8; off < end; off += 8) {
102+
long v = src.get(layout, off);
103+
dst.set(layout, off, v + ref);
104104
}
105105
}
106106
default -> throw new UnsupportedOperationException(
107107
"fastlanes.for: unsupported ptype " + ptype);
108108
}
109-
return MemorySegment.ofArray(bytes);
109+
return dst;
110110
}
111111
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,7 @@ private static Array expand(
103103
private static Array decodeChildAs(DecodeContext parent, int childIdx, DType dtype, long rowCount) {
104104
ArrayNode childNode = parent.node().children()[childIdx];
105105
DecodeContext childCtx = new DecodeContext(
106-
childNode, dtype, rowCount, parent.segmentBuffers(), parent.registry());
106+
childNode, dtype, rowCount, parent.segmentBuffers(), parent.registry(), parent.arena());
107107
return parent.registry().decode(childCtx);
108108
}
109109

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

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -71,11 +71,12 @@ public Array decode(DecodeContext ctx) {
7171
// Allocate output buffer filled with the fill value
7272
int elemBytes = valuePtype.byteSize();
7373
byte[] outBytes = new byte[(int) (n * elemBytes)];
74-
fillBuffer(outBytes, n, valuePtype, fillScalar);
74+
MemorySegment out = MemorySegment.ofArray(outBytes);
75+
fillSegment(out, n, valuePtype, fillScalar);
7576

7677
if (numPatches == 0) {
7778
return new Array(ctx.dtype(), n,
78-
new MemorySegment[]{MemorySegment.ofArray(outBytes)}, new Array[0], ArrayStats.empty());
79+
new MemorySegment[]{out}, new Array[0], ArrayStats.empty());
7980
}
8081

8182
// Decode patch indices with their own dtype and rowCount
@@ -86,51 +87,51 @@ public Array decode(DecodeContext ctx) {
8687
Array valuesArray = decodeChildWithDtype(ctx, 1, ctx.dtype(), numPatches);
8788

8889
// Apply patches: output[index - offset] = value
89-
applyPatches(outBytes, n, valuePtype,
90+
applyPatches(out, n, valuePtype,
9091
indicesArray.buffer(0), valuesArray.buffer(0), indicesPtype, numPatches, offset);
9192

9293
return new Array(ctx.dtype(), n,
93-
new MemorySegment[]{MemorySegment.ofArray(outBytes)}, new Array[0], ArrayStats.empty());
94+
new MemorySegment[]{out}, new Array[0], ArrayStats.empty());
9495
}
9596

9697
// ── Helpers ───────────────────────────────────────────────────────────────
9798

9899
private static Array decodeChildWithDtype(DecodeContext parent, int childIdx, DType dtype, long rowCount) {
99100
ArrayNode childNode = parent.node().children()[childIdx];
100101
DecodeContext childCtx = new DecodeContext(
101-
childNode, dtype, rowCount, parent.segmentBuffers(), parent.registry());
102+
childNode, dtype, rowCount, parent.segmentBuffers(), parent.registry(), parent.arena());
102103
try {
103104
return parent.registry().decode(childCtx);
104105
} catch (Exception e) {
105106
throw new IllegalStateException("vortex.sparse: failed to decode child " + childIdx, e);
106107
}
107108
}
108109

109-
private static void fillBuffer(byte[] buf, long n, PType ptype, ScalarProtos.ScalarValue scalar) {
110+
private static void fillSegment(MemorySegment out, long n, PType ptype, ScalarProtos.ScalarValue scalar) {
110111
long fillLong = scalarToLong(scalar, ptype);
111112
int elemBytes = ptype.byteSize();
112-
ByteBuffer bb = ByteBuffer.wrap(buf).order(ByteOrder.LITTLE_ENDIAN);
113+
ByteBuffer bb = out.asByteBuffer().order(ByteOrder.LITTLE_ENDIAN);
113114
for (long i = 0; i < n; i++) {
114115
writeElem(bb, ptype, fillLong);
115116
}
116117
}
117118

118119
private static void applyPatches(
119-
byte[] outBytes, long n, PType valuePtype,
120+
MemorySegment out, long n, PType valuePtype,
120121
MemorySegment idxSeg, MemorySegment valSeg,
121122
PType idxPtype, long numPatches, long offset
122123
) {
123124
int elemBytes = valuePtype.byteSize();
124-
ByteBuffer out = ByteBuffer.wrap(outBytes).order(ByteOrder.LITTLE_ENDIAN);
125+
ByteBuffer outBuf = out.asByteBuffer().order(ByteOrder.LITTLE_ENDIAN);
125126
for (long i = 0; i < numPatches; i++) {
126127
long idx = readUnsignedIdx(idxSeg, i, idxPtype) - offset;
127128
if (idx < 0 || idx >= n) {
128129
throw new IllegalStateException(
129130
"vortex.sparse: patch index " + idx + " out of range [0," + n + ")");
130131
}
131132
long val = readElem(valSeg, i, valuePtype);
132-
out.position((int) (idx * elemBytes));
133-
writeElem(out, valuePtype, val);
133+
outBuf.position((int) (idx * elemBytes));
134+
writeElem(outBuf, valuePtype, val);
134135
}
135136
}
136137

0 commit comments

Comments
 (0)