Skip to content

Commit 0397b7b

Browse files
dfa1claude
dfa1
and
claude
committed
feat(encoding): implement ALP encode for F32
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent a8bd874 commit 0397b7b

2 files changed

Lines changed: 184 additions & 4 deletions

File tree

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

Lines changed: 114 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -99,10 +99,11 @@ public boolean accepts(DType dtype) {
9999
@Override
100100
public EncodeResult encode(DType dtype, Object data) {
101101
PType ptype = ((DType.Primitive) dtype).ptype();
102-
if (ptype == PType.F64) {
103-
return encodeF64((double[]) data, dtype);
104-
}
105-
throw new UnsupportedOperationException("ALP encode not yet supported for " + ptype);
102+
return switch (ptype) {
103+
case F64 -> encodeF64((double[]) data, dtype);
104+
case F32 -> encodeF32((float[]) data, dtype);
105+
default -> throw new UnsupportedOperationException("ALP encode not supported for " + ptype);
106+
};
106107
}
107108

108109
// ── F64 encode ────────────────────────────────────────────────────────────
@@ -214,6 +215,115 @@ private static byte[] scalarF64(double v) {
214215
return ScalarProtos.ScalarValue.newBuilder().setF64Value(v).build().toByteArray();
215216
}
216217

218+
// ── F32 encode ────────────────────────────────────────────────────────────
219+
220+
private static int[] findExponentsF32(float[] values) {
221+
int sampleLen = Math.min(SAMPLE_SIZE, values.length);
222+
int bestExpE = 0, bestExpF = 0, bestExceptions = sampleLen + 1;
223+
224+
outer:
225+
for (int expE = 0; expE < F10_F32.length; expE++) {
226+
for (int expF = 0; expF < F10_F32.length; expF++) {
227+
float encFactor = F10_F32[expE] * IF10_F32[expF];
228+
float decFactor = F10_F32[expF] * IF10_F32[expE];
229+
int exceptions = 0;
230+
for (int i = 0; i < sampleLen; i++) {
231+
float enc = values[i] * encFactor;
232+
if (!Float.isFinite(enc) || (float) Math.round(enc) * decFactor != values[i]) {
233+
exceptions++;
234+
}
235+
}
236+
if (exceptions < bestExceptions) {
237+
bestExceptions = exceptions;
238+
bestExpE = expE;
239+
bestExpF = expF;
240+
if (bestExceptions == 0) {
241+
break outer;
242+
}
243+
}
244+
}
245+
}
246+
return new int[]{bestExpE, bestExpF};
247+
}
248+
249+
private static EncodeResult encodeF32(float[] values, DType dtype) {
250+
int n = values.length;
251+
int[] exps = findExponentsF32(values);
252+
int expE = exps[0], expF = exps[1];
253+
float encFactor = F10_F32[expE] * IF10_F32[expF];
254+
float decFactor = F10_F32[expF] * IF10_F32[expE];
255+
256+
int[] encodedArr = new int[n];
257+
var patchIndices = new ArrayList<Integer>();
258+
var patchValues = new ArrayList<Float>();
259+
260+
float min = Float.MAX_VALUE, max = -Float.MAX_VALUE;
261+
for (int i = 0; i < n; i++) {
262+
float v = values[i];
263+
float enc = v * encFactor;
264+
int encoded;
265+
if (Float.isFinite(enc) && (float) (encoded = Math.round(enc)) * decFactor == v) {
266+
encodedArr[i] = encoded;
267+
} else {
268+
encodedArr[i] = 0;
269+
patchIndices.add(i);
270+
patchValues.add(v);
271+
}
272+
if (v < min) {
273+
min = v;
274+
}
275+
if (v > max) {
276+
max = v;
277+
}
278+
}
279+
280+
byte[] statsMin = n > 0 ? scalarF32(min) : null;
281+
byte[] statsMax = n > 0 ? scalarF32(max) : null;
282+
283+
MemorySegment encodedBuf = Arena.ofAuto().allocate((long) n * 4, 4);
284+
for (int i = 0; i < n; i++) {
285+
encodedBuf.setAtIndex(PTypeIO.LE_INT, i, encodedArr[i]);
286+
}
287+
288+
EncodeNode encodedNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 0);
289+
290+
if (patchIndices.isEmpty()) {
291+
byte[] metaBytes = EncodingProtos.ALPMetadata.newBuilder()
292+
.setExpE(expE).setExpF(expF).build().toByteArray();
293+
EncodeNode root = new EncodeNode(EncodingId.VORTEX_ALP,
294+
ByteBuffer.wrap(metaBytes), new EncodeNode[]{encodedNode}, new int[0]);
295+
return new EncodeResult(root, List.of(encodedBuf), statsMin, statsMax);
296+
}
297+
298+
int numPatches = patchIndices.size();
299+
MemorySegment idxBuf = Arena.ofAuto().allocate((long) numPatches * 4, 4);
300+
MemorySegment valBuf = Arena.ofAuto().allocate((long) numPatches * 4, 4);
301+
for (int i = 0; i < numPatches; i++) {
302+
idxBuf.setAtIndex(PTypeIO.LE_INT, i, patchIndices.get(i));
303+
valBuf.setAtIndex(PTypeIO.LE_FLOAT, i, patchValues.get(i));
304+
}
305+
306+
EncodingProtos.PatchesMetadata patches = EncodingProtos.PatchesMetadata.newBuilder()
307+
.setLen(numPatches)
308+
.setOffset(0)
309+
.setIndicesPtype(DTypeProtos.PType.forNumber(PType.U32.ordinal()))
310+
.build();
311+
byte[] metaBytes = EncodingProtos.ALPMetadata.newBuilder()
312+
.setExpE(expE).setExpF(expF).setPatches(patches).build().toByteArray();
313+
314+
EncodeNode idxNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 1);
315+
EncodeNode valNode = EncodeNode.leaf(EncodingId.VORTEX_PRIMITIVE, 2);
316+
EncodeNode root = new EncodeNode(EncodingId.VORTEX_ALP,
317+
ByteBuffer.wrap(metaBytes),
318+
new EncodeNode[]{encodedNode, idxNode, valNode},
319+
new int[0]);
320+
return new EncodeResult(root, List.of(encodedBuf, idxBuf, valBuf), statsMin, statsMax);
321+
}
322+
323+
private static byte[] scalarF32(float v) {
324+
return ScalarProtos.ScalarValue.newBuilder().setF32Value(v).build().toByteArray();
325+
}
326+
217327
@Override
218328
public Array decode(DecodeContext ctx) {
219329
ByteBuffer rawMeta = ctx.metadata();

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -220,4 +220,74 @@ void decode_f32_noPatches() {
220220
.as("index %d", i).isCloseTo(expected[i], within(1e-6f));
221221
}
222222
}
223+
224+
@Test
225+
void encode_f32_roundTrip_noPatches() {
226+
// Given
227+
float[] values = {1.0f, 2.5f, 3.75f, 10.0f, 0.1f};
228+
AlpEncoding sut = new AlpEncoding();
229+
230+
EncodingRegistry registry = EncodingRegistry.empty();
231+
registry.register(sut);
232+
registry.register(new PrimitiveEncoding());
233+
234+
// When
235+
EncodeResult encoded = sut.encode(F32_DTYPE, values);
236+
DecodeContext ctx = EncodeTestHelper.toDecodeContext(encoded, values.length, F32_DTYPE, registry);
237+
Array result = sut.decode(ctx);
238+
239+
// Then
240+
var layout = ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
241+
for (int i = 0; i < values.length; i++) {
242+
assertThat(result.buffer(0).get(layout, (long) i * 4))
243+
.as("index %d", i).isCloseTo(values[i], within(1e-6f));
244+
}
245+
}
246+
247+
@Test
248+
void encode_f32_roundTrip_withPatches() {
249+
// Given — Float.NaN and Float.POSITIVE_INFINITY can't be ALP-encoded; must become patches
250+
float[] values = {1.0f, Float.NaN, 2.5f, Float.POSITIVE_INFINITY, 3.0f};
251+
AlpEncoding sut = new AlpEncoding();
252+
253+
EncodingRegistry registry = EncodingRegistry.empty();
254+
registry.register(sut);
255+
registry.register(new PrimitiveEncoding());
256+
257+
// When
258+
EncodeResult encoded = sut.encode(F32_DTYPE, values);
259+
DecodeContext ctx = EncodeTestHelper.toDecodeContext(encoded, values.length, F32_DTYPE, registry);
260+
Array result = sut.decode(ctx);
261+
262+
// Then
263+
var layout = ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
264+
assertThat(result.buffer(0).get(layout, 0L)).isCloseTo(1.0f, within(1e-6f));
265+
assertThat(result.buffer(0).get(layout, 4L)).isNaN();
266+
assertThat(result.buffer(0).get(layout, 8L)).isCloseTo(2.5f, within(1e-6f));
267+
assertThat(result.buffer(0).get(layout, 12L)).isInfinite();
268+
assertThat(result.buffer(0).get(layout, 16L)).isCloseTo(3.0f, within(1e-6f));
269+
}
270+
271+
@Test
272+
void encode_f64_roundTrip_noPatches() {
273+
// Given
274+
double[] values = {1.23, 4.56, 7.89, 0.001, 100.0};
275+
AlpEncoding sut = new AlpEncoding();
276+
277+
EncodingRegistry registry = EncodingRegistry.empty();
278+
registry.register(sut);
279+
registry.register(new PrimitiveEncoding());
280+
281+
// When
282+
EncodeResult encoded = sut.encode(F64_DTYPE, values);
283+
DecodeContext ctx = EncodeTestHelper.toDecodeContext(encoded, values.length, F64_DTYPE, registry);
284+
Array result = sut.decode(ctx);
285+
286+
// Then
287+
var layout = ValueLayout.JAVA_DOUBLE_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
288+
for (int i = 0; i < values.length; i++) {
289+
assertThat(result.buffer(0).get(layout, (long) i * 8))
290+
.as("index %d", i).isCloseTo(values[i], within(1e-9));
291+
}
292+
}
223293
}

0 commit comments

Comments
 (0)