Skip to content

Commit 845058e

Browse files
dfa1claude
dfa1
and
claude
committed
refactor: introduce EncodeContext with arena + registry; merge CompressorContext into it
EncodeContext now carries Arena, EncodingRegistry, and all cascade parameters (depth, exclusions, sampling), replacing the separate CompressorContext record. Key changes: - Encoding.encode(DType, Object) default removed — callers must supply EncodeContext; eliminates hidden EncodingRegistry.loadAll() / ServiceLoader scan on every call - Encoding.encodeCascade signature reduced to 3-arg (DType, Object, EncodeContext), symmetric with decode(DecodeContext) - CompressorContext deleted; EncodeContext.ofDepth() / .of() replace it - CascadingCompressor constructor drops the context param; cascade params come from the EncodeContext passed to encode() - Cross-encoder calls (DateTimePartsEncoding, AlpRdEncoding, ExtEncoding, etc.) use ctx.lookupEncoding(id) instead of new XyzEncoding(), avoiding per-call instantiation - VortexWriter uses Arena.ofConfined() scoped to each segment write, releasing encode buffers immediately after flush rather than waiting for GC - EncodingRegistry.of(List<Encoding>) factory added - Fixed DateTimePartsEncoding.encodeCascade override: was overriding 3-arg default instead of 4-arg, so CascadingCompressor never saw open children for timestamp columns - Fixed TestRegistry.withPrimitive duplicate-registration crash when sut is PrimitiveEncoding - Fixed missing DType import in RandomAccessTest Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 30f63c0 commit 845058e

75 files changed

Lines changed: 620 additions & 510 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

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

Lines changed: 21 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
import io.github.dfa1.vortex.core.array.DoubleArray;
1313
import io.github.dfa1.vortex.core.array.FloatArray;
1414

15-
import java.lang.foreign.Arena;
1615
import java.lang.foreign.MemorySegment;
1716
import java.lang.foreign.ValueLayout;
1817
import java.nio.ByteBuffer;
@@ -76,13 +75,13 @@ public boolean accepts(DType dtype) {
7675
}
7776

7877
@Override
79-
public EncodeResult encode(DType dtype, Object data) {
80-
return Encoder.encode(dtype, data);
78+
public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) {
79+
return Encoder.encode(dtype, data, ctx);
8180
}
8281

8382
@Override
84-
public CascadeStep encodeCascade(DType dtype, Object data, CompressorContext ctx) {
85-
return Encoder.encodeCascade(dtype, data);
83+
public CascadeStep encodeCascade(DType dtype, Object data, EncodeContext encodeCtx) {
84+
return Encoder.encodeCascade(dtype, data, encodeCtx);
8685
}
8786

8887
@Override
@@ -101,21 +100,21 @@ private record AlpF64Data(int expE, int expF, long[] encodedArr,
101100
List<Integer> patchIndices, List<Double> patchValues,
102101
byte[] statsMin, byte[] statsMax) {}
103102

104-
private static EncodeResult encode(DType dtype, Object data) {
103+
private static EncodeResult encode(DType dtype, Object data, EncodeContext ctx) {
105104
PType ptype = ((DType.Primitive) dtype).ptype();
106105
return switch (ptype) {
107-
case F64 -> encodeF64((double[]) data, dtype);
108-
case F32 -> encodeF32((float[]) data, dtype);
106+
case F64 -> encodeF64((double[]) data, dtype, ctx);
107+
case F32 -> encodeF32((float[]) data, dtype, ctx);
109108
default -> throw new UnsupportedOperationException("ALP encode not supported for " + ptype);
110109
};
111110
}
112111

113-
private static CascadeStep encodeCascade(DType dtype, Object data) {
112+
private static CascadeStep encodeCascade(DType dtype, Object data, EncodeContext ctx) {
114113
PType ptype = ((DType.Primitive) dtype).ptype();
115114
if (ptype == PType.F64) {
116-
return encodeCascadeF64((double[]) data);
115+
return encodeCascadeF64((double[]) data, ctx);
117116
}
118-
return CascadeStep.terminal(encode(dtype, data));
117+
return CascadeStep.terminal(encode(dtype, data, ctx));
119118
}
120119

121120
private static int[] findExponentsF64(double[] values) {
@@ -187,11 +186,11 @@ private static AlpF64Data computeF64(double[] values) {
187186
return new AlpF64Data(expE, expF, encodedArr, patchIndices, patchValues, statsMin, statsMax);
188187
}
189188

190-
private static EncodeResult encodeF64(double[] values, DType dtype) {
189+
private static EncodeResult encodeF64(double[] values, DType dtype, EncodeContext ctx) {
191190
AlpF64Data d = computeF64(values);
192191
int n = values.length;
193192

194-
MemorySegment encodedBuf = Arena.ofAuto().allocate((long) n * 8, 8);
193+
MemorySegment encodedBuf = ctx.arena().allocate((long) n * 8, 8);
195194
for (int i = 0; i < n; i++) {
196195
encodedBuf.setAtIndex(PTypeIO.LE_LONG, i, d.encodedArr()[i]);
197196
}
@@ -207,8 +206,8 @@ private static EncodeResult encodeF64(double[] values, DType dtype) {
207206
}
208207

209208
int numPatches = d.patchIndices().size();
210-
MemorySegment idxBuf = Arena.ofAuto().allocate((long) numPatches * 4, 4);
211-
MemorySegment valBuf = Arena.ofAuto().allocate((long) numPatches * 8, 8);
209+
MemorySegment idxBuf = ctx.arena().allocate((long) numPatches * 4, 4);
210+
MemorySegment valBuf = ctx.arena().allocate((long) numPatches * 8, 8);
212211
for (int i = 0; i < numPatches; i++) {
213212
idxBuf.setAtIndex(PTypeIO.LE_INT, i, d.patchIndices().get(i));
214213
valBuf.setAtIndex(PTypeIO.LE_DOUBLE, i, d.patchValues().get(i));
@@ -229,7 +228,7 @@ private static EncodeResult encodeF64(double[] values, DType dtype) {
229228

230229
/// Cascade-aware F64 encode: I64 child is an open ChildSlot so the compressor
231230
/// can bitpack it. Patch idx/val buffers (if any) stay in ownedBuffers.
232-
private static CascadeStep encodeCascadeF64(double[] values) {
231+
private static CascadeStep encodeCascadeF64(double[] values, EncodeContext ctx) {
233232
AlpF64Data d = computeF64(values);
234233
int n = values.length;
235234

@@ -244,8 +243,8 @@ private static CascadeStep encodeCascadeF64(double[] values) {
244243
}
245244

246245
int numPatches = d.patchIndices().size();
247-
MemorySegment idxBuf = Arena.ofAuto().allocate((long) numPatches * 4, 4);
248-
MemorySegment valBuf = Arena.ofAuto().allocate((long) numPatches * 8, 8);
246+
MemorySegment idxBuf = ctx.arena().allocate((long) numPatches * 4, 4);
247+
MemorySegment valBuf = ctx.arena().allocate((long) numPatches * 8, 8);
249248
for (int i = 0; i < numPatches; i++) {
250249
idxBuf.setAtIndex(PTypeIO.LE_INT, i, d.patchIndices().get(i));
251250
valBuf.setAtIndex(PTypeIO.LE_DOUBLE, i, d.patchValues().get(i));
@@ -295,7 +294,7 @@ private static int[] findExponentsF32(float[] values) {
295294
return new int[]{bestExpE, bestExpF};
296295
}
297296

298-
private static EncodeResult encodeF32(float[] values, DType dtype) {
297+
private static EncodeResult encodeF32(float[] values, DType dtype, EncodeContext ctx) {
299298
int n = values.length;
300299
int[] exps = findExponentsF32(values);
301300
int expE = exps[0], expF = exps[1];
@@ -331,7 +330,7 @@ private static EncodeResult encodeF32(float[] values, DType dtype) {
331330
byte[] statsMin = n > 0 ? scalarF32(min) : null;
332331
byte[] statsMax = n > 0 ? scalarF32(max) : null;
333332

334-
MemorySegment encodedBuf = Arena.ofAuto().allocate((long) n * 4, 4);
333+
MemorySegment encodedBuf = ctx.arena().allocate((long) n * 4, 4);
335334
for (int i = 0; i < n; i++) {
336335
encodedBuf.setAtIndex(PTypeIO.LE_INT, i, encodedArr[i]);
337336
}
@@ -347,8 +346,8 @@ private static EncodeResult encodeF32(float[] values, DType dtype) {
347346
}
348347

349348
int numPatches = patchIndices.size();
350-
MemorySegment idxBuf = Arena.ofAuto().allocate((long) numPatches * 4, 4);
351-
MemorySegment valBuf = Arena.ofAuto().allocate((long) numPatches * 4, 4);
349+
MemorySegment idxBuf = ctx.arena().allocate((long) numPatches * 4, 4);
350+
MemorySegment valBuf = ctx.arena().allocate((long) numPatches * 4, 4);
352351
for (int i = 0; i < numPatches; i++) {
353352
idxBuf.setAtIndex(PTypeIO.LE_INT, i, patchIndices.get(i));
354353
valBuf.setAtIndex(PTypeIO.LE_FLOAT, i, patchValues.get(i));

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

Lines changed: 21 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -60,8 +60,8 @@ public boolean accepts(DType dtype) {
6060
}
6161

6262
@Override
63-
public EncodeResult encode(DType dtype, Object data) {
64-
return Encoder.encode(dtype, data);
63+
public EncodeResult encode(DType dtype, Object data, EncodeContext ctx) {
64+
return Encoder.encode(dtype, data, ctx);
6565
}
6666

6767
@Override
@@ -79,21 +79,21 @@ private static final class Encoder {
7979
private static final int MAX_CUT = 16;
8080
private static final int MAX_DICT_SIZE = 8;
8181

82-
static EncodeResult encode(DType dtype, Object data) {
82+
static EncodeResult encode(DType dtype, Object data, EncodeContext ctx) {
8383
PType ptype = ((DType.Primitive) dtype).ptype();
8484
return switch (ptype) {
85-
case F64 -> encodeF64((double[]) data);
86-
case F32 -> encodeF32((float[]) data);
85+
case F64 -> encodeF64((double[]) data, ctx);
86+
case F32 -> encodeF32((float[]) data, ctx);
8787
default -> throw new UnsupportedOperationException("ALP-RD encode not supported for " + ptype);
8888
};
8989
}
9090

9191
// --- F64 ---
9292

93-
private static EncodeResult encodeF64(double[] values) {
93+
private static EncodeResult encodeF64(double[] values, EncodeContext ctx) {
9494
int n = values.length;
9595
if (n == 0) {
96-
return emptyResult(U64_DTYPE, U16_DTYPE);
96+
return emptyResult(U64_DTYPE, U16_DTYPE, ctx);
9797
}
9898

9999
int sampleLen = Math.min(SAMPLE_SIZE, n);
@@ -123,7 +123,7 @@ private static EncodeResult encodeF64(double[] values) {
123123

124124
return buildEncodeResult(
125125
best.dict, best.rightBitWidth, leftCodes, rightParts,
126-
U64_DTYPE, excPos, excVals);
126+
U64_DTYPE, excPos, excVals, ctx);
127127
}
128128

129129
private record Dictionary64(short[] dict, int rightBitWidth) {}
@@ -174,10 +174,10 @@ private static int countExceptions(double[] values, int sampleLen, short[] dict,
174174

175175
// --- F32 ---
176176

177-
private static EncodeResult encodeF32(float[] values) {
177+
private static EncodeResult encodeF32(float[] values, EncodeContext ctx) {
178178
int n = values.length;
179179
if (n == 0) {
180-
return emptyResult(U32_DTYPE, U16_DTYPE);
180+
return emptyResult(U32_DTYPE, U16_DTYPE, ctx);
181181
}
182182

183183
int sampleLen = Math.min(SAMPLE_SIZE, n);
@@ -207,7 +207,7 @@ private static EncodeResult encodeF32(float[] values) {
207207

208208
return buildEncodeResult(
209209
best.dict, best.rightBitWidth, leftCodes, rightParts,
210-
U32_DTYPE, excPos, excVals);
210+
U32_DTYPE, excPos, excVals, ctx);
211211
}
212212

213213
private record Dictionary32(short[] dict, int rightBitWidth) {}
@@ -280,11 +280,11 @@ private static Map<Short, Short> buildLookup(short[] dict) {
280280
private static EncodeResult buildEncodeResult(
281281
short[] dict, int rightBitWidth,
282282
short[] leftCodes, Object rightPartsData, DType rightDtype,
283-
List<Long> excPos, List<Short> excVals) {
283+
List<Long> excPos, List<Short> excVals, EncodeContext ctx) {
284284

285-
BitpackedEncoding bp = new BitpackedEncoding();
286-
EncodeResult leftResult = bp.encode(U16_DTYPE, leftCodes);
287-
EncodeResult rightResult = bp.encode(rightDtype, rightPartsData);
285+
Encoding bp = ctx.lookupEncoding(EncodingId.FASTLANES_BITPACKED);
286+
EncodeResult leftResult = bp.encode(U16_DTYPE, leftCodes, ctx);
287+
EncodeResult rightResult = bp.encode(rightDtype, rightPartsData, ctx);
288288

289289
List<MemorySegment> allBuffers = new ArrayList<>(leftResult.buffers());
290290
int leftBufCount = allBuffers.size();
@@ -312,8 +312,8 @@ private static EncodeResult buildEncodeResult(
312312
excValsArr[i] = excVals.get(i);
313313
}
314314

315-
EncodeResult idxResult = bp.encode(U64_DTYPE, excPosArr);
316-
EncodeResult valResult = bp.encode(U16_DTYPE, excValsArr);
315+
EncodeResult idxResult = bp.encode(U64_DTYPE, excPosArr, ctx);
316+
EncodeResult valResult = bp.encode(U16_DTYPE, excValsArr, ctx);
317317

318318
int idxOffset = allBuffers.size();
319319
allBuffers.addAll(idxResult.buffers());
@@ -338,11 +338,11 @@ private static EncodeResult buildEncodeResult(
338338
return new EncodeResult(root, List.copyOf(allBuffers), null, null);
339339
}
340340

341-
private static EncodeResult emptyResult(DType rightDtype, DType leftDtype) {
342-
BitpackedEncoding bp = new BitpackedEncoding();
343-
EncodeResult leftResult = bp.encode(leftDtype, new short[0]);
341+
private static EncodeResult emptyResult(DType rightDtype, DType leftDtype, EncodeContext ctx) {
342+
Encoding bp = ctx.lookupEncoding(EncodingId.FASTLANES_BITPACKED);
343+
EncodeResult leftResult = bp.encode(leftDtype, new short[0], ctx);
344344
EncodeResult rightResult = bp.encode(rightDtype,
345-
rightDtype.equals(U32_DTYPE) ? new int[0] : new long[0]);
345+
rightDtype.equals(U32_DTYPE) ? new int[0] : new long[0], ctx);
346346

347347
List<MemorySegment> allBuffers = new ArrayList<>(leftResult.buffers());
348348
int leftBufCount = allBuffers.size();

0 commit comments

Comments
 (0)