Skip to content

Commit f852bef

Browse files
dfa1claude
andcommitted
feat(encoding): implement pco FloatMult mode (mode=2)
Port intFloatFromLatent + toLatentOrdered helpers for F32/F64; combine step: rawLatents[i] = toLatentOrdered(intFloatFromLatent(mult) * baseFloat) + adj. Unblocks s3_pcoVortex_javaDecodeMatchesJni (removed @disabled). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 7a87c26 commit f852bef

2 files changed

Lines changed: 92 additions & 19 deletions

File tree

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

Lines changed: 92 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -34,14 +34,13 @@
3434
/// [0–7b alignment]</li>
3535
/// <li>Classic page: [primary: deltaOrder×dtypeSize b moments, 4×ansSizeLog b ANS states]
3636
/// [0–7b alignment] [per 256-batch: ANS bits, offset bits]</li>
37-
/// <li>IntMult page: [primary header][secondary header][0–7b alignment]
37+
/// <li>IntMult/FloatMult page: [primary header][secondary header][0–7b alignment]
3838
/// [per 256-batch: primary ANS+offsets, secondary ANS+offsets]</li>
3939
/// <li>All bit packing little-endian (LSB first)</li>
4040
/// </ul>
4141
///
42-
/// <p>Supported: Classic and IntMult modes, None+Consecutive delta (order ≤ 7), non-null,
43-
/// all integer/float ptypes except F16. Other modes (FloatMult, FloatQuant, Dict) throw
44-
/// "not implemented".
42+
/// <p>Supported: Classic, IntMult, FloatMult modes, None+Consecutive delta (order ≤ 7), non-null,
43+
/// all integer/float ptypes except F16. Other modes (FloatQuant, Dict) throw "not implemented".
4544
public final class PcoEncoding implements Encoding {
4645

4746
static final byte PCO_FORMAT_MAJOR = 0x04;
@@ -120,7 +119,7 @@ static Array decode(DecodeContext ctx) {
120119
pageBuf, pageN, rawLatents, rawByteOffset);
121120
}
122121
} else {
123-
// IntMult mode: two latent variables (mult, adj); num = base * mult + adj.
122+
// IntMult (1) / FloatMult (2): two latent variables; page decode is identical.
124123
long base = chunkMeta.base();
125124
int primaryAnsSizeLog = chunkMeta.ansSizeLog();
126125
int secondaryAnsSizeLog = chunkMeta.secondaryAnsSizeLog();
@@ -151,13 +150,17 @@ static Array decode(DecodeContext ctx) {
151150
adjByteOffset += (long) pageN * Long.BYTES;
152151
}
153152

154-
// Combine: rawLatents[i] = (mult[i] * base + adj[i]) & mask
155-
long mask = typeMask(dtypeSize);
156-
for (int i = 0; i < chunkN; i++) {
157-
long off = chunkMultsOffset + (long) i * Long.BYTES;
158-
long mult = rawLatents.get(LE_LONG, off);
159-
long adj = rawAdjs.get(LE_LONG, (long) i * Long.BYTES);
160-
rawLatents.set(LE_LONG, off, (mult * base + adj) & mask);
153+
// Combine — mode 1 = IntMult, mode 2 = FloatMult.
154+
if (chunkMeta.mode() == 1) {
155+
long mask = typeMask(dtypeSize);
156+
for (int i = 0; i < chunkN; i++) {
157+
long off = chunkMultsOffset + (long) i * Long.BYTES;
158+
long mult = rawLatents.get(LE_LONG, off);
159+
long adj = rawAdjs.get(LE_LONG, (long) i * Long.BYTES);
160+
rawLatents.set(LE_LONG, off, (mult * base + adj) & mask);
161+
}
162+
} else {
163+
combineFloatMult(ptype, base, chunkN, rawLatents, chunkMultsOffset, rawAdjs);
161164
}
162165
}
163166
}
@@ -372,17 +375,89 @@ private static Array toArray(DType dtype, long n, MemorySegment out) {
372375
};
373376
}
374377

378+
/// Int-float latent → integer-valued float for F32.
379+
///
380+
/// mid = 2^31; negative when l &lt; mid; GPI = 2^24 (f32 mantissa bits).
381+
/// Inverse of pcodec {@code int_float_to_latent}.
382+
private static float intFloatFromLatentF32(long l) {
383+
long mid = 0x80000000L;
384+
boolean negative = (l < mid);
385+
long absInt = negative ? (0x7FFFFFFFL - l) : (l ^ 0x80000000L);
386+
long gpi = 1L << 24;
387+
float absFloat = (absInt < gpi) ? (float) absInt
388+
: Float.intBitsToFloat(0x4B800000 + (int) (absInt - gpi));
389+
return negative ? -absFloat : absFloat;
390+
}
391+
392+
/// Int-float latent → integer-valued float for F64.
393+
///
394+
/// mid = 2^63; l &lt; 2^63 (unsigned, i.e. l ≥ 0 signed) → negative float; GPI = 2^53.
395+
/// Inverse of pcodec {@code int_float_to_latent}.
396+
private static double intFloatFromLatentF64(long l) {
397+
// l >= 0 signed ↔ l < 2^63 unsigned → original float was negative
398+
boolean negative = (l >= 0);
399+
long absInt = negative ? (Long.MAX_VALUE - l) : (l ^ Long.MIN_VALUE);
400+
long gpi = 1L << 53;
401+
double absFloat = (absInt < gpi) ? (double) absInt
402+
: Double.longBitsToDouble(0x4340000000000000L + (absInt - gpi));
403+
return negative ? -absFloat : absFloat;
404+
}
405+
406+
/// Inverse of {@link #fromLatentOrdered} for F32: float → ordered latent.
407+
private static long toLatentOrderedF32(float f) {
408+
int bits = Float.floatToRawIntBits(f);
409+
if ((bits & 0x80000000) != 0) {
410+
return (~bits) & 0xFFFFFFFFL;
411+
} else {
412+
return (bits ^ 0x80000000) & 0xFFFFFFFFL;
413+
}
414+
}
415+
416+
/// Inverse of {@link #fromLatentOrdered} for F64: float → ordered latent.
417+
private static long toLatentOrderedF64(double d) {
418+
long bits = Double.doubleToRawLongBits(d);
419+
if ((bits & Long.MIN_VALUE) != 0) {
420+
return ~bits;
421+
} else {
422+
return bits ^ Long.MIN_VALUE;
423+
}
424+
}
425+
426+
/// FloatMult combine: rawLatents[i] = toLatentOrdered(intFloatFromLatent(mult) * baseFloat) + adj.
427+
private static void combineFloatMult(PType ptype, long baseLatent, int chunkN,
428+
MemorySegment rawLatents, long multsOffset, MemorySegment rawAdjs) {
429+
if (ptype == PType.F32) {
430+
float baseFloat = Float.intBitsToFloat((int) fromLatentOrdered(baseLatent, PType.F32));
431+
for (int i = 0; i < chunkN; i++) {
432+
long off = multsOffset + (long) i * Long.BYTES;
433+
long mult = rawLatents.get(LE_LONG, off);
434+
long adj = rawAdjs.get(LE_LONG, (long) i * Long.BYTES);
435+
long unadjusted = toLatentOrderedF32(intFloatFromLatentF32(mult) * baseFloat);
436+
rawLatents.set(LE_LONG, off, (unadjusted + adj) & 0xFFFFFFFFL);
437+
}
438+
} else {
439+
double baseDouble = Double.longBitsToDouble(fromLatentOrdered(baseLatent, PType.F64));
440+
for (int i = 0; i < chunkN; i++) {
441+
long off = multsOffset + (long) i * Long.BYTES;
442+
long mult = rawLatents.get(LE_LONG, off);
443+
long adj = rawAdjs.get(LE_LONG, (long) i * Long.BYTES);
444+
long unadjusted = toLatentOrderedF64(intFloatFromLatentF64(mult) * baseDouble);
445+
rawLatents.set(LE_LONG, off, unadjusted + adj);
446+
}
447+
}
448+
}
449+
375450
private static PcoChunkMeta readChunkMeta(MemorySegment buf, int dtypeSize) {
376451
LeBitReader r = new LeBitReader(buf);
377452

378453
int modeNibble = (int) r.readBits(4);
379454
long base = 0L;
380-
if (modeNibble == 1) {
381-
// IntMult: read base value (dtypeSize bits, uncompressed).
455+
if (modeNibble == 1 || modeNibble == 2) {
456+
// IntMult / FloatMult: read base value (dtypeSize bits, uncompressed latent).
382457
base = r.readBits(dtypeSize);
383458
} else if (modeNibble != 0) {
384459
throw new VortexException(EncodingId.VORTEX_PCO,
385-
"pco mode " + modeNibble + " not yet implemented (Classic=0, IntMult=1 supported)");
460+
"pco mode " + modeNibble + " not yet implemented (Classic=0, IntMult=1, FloatMult=2 supported)");
386461
}
387462

388463
// delta_encoding_variant: 4 bits. 0=NoOp, 1=Consecutive (3b order + 1b secondary_uses_delta).
@@ -404,10 +479,10 @@ private static PcoChunkMeta readChunkMeta(MemorySegment buf, int dtypeSize) {
404479
int nBins = (int) r.readBits(15);
405480
PcoBin[] bins = readBins(r, nBins, ansSizeLog, dtypeSize);
406481

407-
// Secondary latent variable (IntMult only).
482+
// Secondary latent variable (IntMult and FloatMult).
408483
int secondaryAnsSizeLog = 0;
409484
PcoBin[] secondaryBins = new PcoBin[0];
410-
if (modeNibble == 1) {
485+
if (modeNibble == 1 || modeNibble == 2) {
411486
secondaryAnsSizeLog = (int) r.readBits(4);
412487
int nSecondaryBins = (int) r.readBits(15);
413488
secondaryBins = readBins(r, nSecondaryBins, secondaryAnsSizeLog, dtypeSize);

integration/src/test/java/io/github/dfa1/vortex/integration/RustWritesJavaReadsIntegrationTest.java

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
import org.apache.arrow.vector.types.pojo.ArrowType;
2828
import org.apache.arrow.vector.types.pojo.Field;
2929
import org.apache.arrow.vector.types.pojo.Schema;
30-
import org.junit.jupiter.api.Disabled;
3130
import org.junit.jupiter.api.Test;
3231
import org.junit.jupiter.api.io.TempDir;
3332

@@ -259,7 +258,6 @@ void jniWriter_nullableColumn_decodesWithoutError(@TempDir Path tmp) throws IOEx
259258

260259
// ── S3 fixture round-trip: Rust-written pco → Java reader ────────────────
261260

262-
@Disabled("pco.vortex has FloatMult columns (mode=2) — blocked on Phase 9")
263261
@Test
264262
void s3_pcoVortex_javaDecodeMatchesJni(@TempDir Path tmp) throws Exception {
265263
// Given — pco.vortex: synthetic file with all pco ptypes; Classic+Consecutive+IntMult+FloatMult modes

0 commit comments

Comments
 (0)