Skip to content

Commit 8a62ee9

Browse files
dfa1claude
dfa1
and
claude
committed
feat(encoding): spec-compliant fastlanes.bitpacked rewrite
Replace 9-byte custom / 2-byte JNI metadata with protobuf BitPackedMetadata {bit_width, offset}. Implement unified FastLanes FL_ORDER block-transposed pack/unpack; correct LANES per element type (1024/T). Drops embedded FOR shift — signed compression deferred to fastlanes.for (#7b). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 0b13654 commit 8a62ee9

4 files changed

Lines changed: 221 additions & 101 deletions

File tree

‎core/pom.xml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@
9595
<argument>--proto_path=${project.basedir}/src/main/proto</argument>
9696
<argument>${project.basedir}/src/main/proto/dtype.proto</argument>
9797
<argument>${project.basedir}/src/main/proto/scalar.proto</argument>
98+
<argument>${project.basedir}/src/main/proto/encodings.proto</argument>
9899
</arguments>
99100
</configuration>
100101
</execution>

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

Lines changed: 207 additions & 100 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package io.github.dfa1.vortex.encoding;
22

3+
import com.google.protobuf.InvalidProtocolBufferException;
4+
import dev.vortex.proto.EncodingProtos;
35
import dev.vortex.proto.ScalarProtos;
46
import io.github.dfa1.vortex.core.DType;
57
import io.github.dfa1.vortex.core.PType;
@@ -10,14 +12,19 @@
1012
import java.nio.ByteOrder;
1113
import java.util.List;
1214

13-
/// Codec for `fastlanes.bitpacked` — bit-packing for integer columns.
15+
/// Codec for {@code fastlanes.bitpacked} — spec-compliant FastLanes bit-packing.
1416
///
15-
/// Metadata (9 bytes): bit_width (u8) | frame_of_reference (i64 LE).
16-
/// Values are shifted by the column minimum (frame-of-reference), then packed
17-
/// bit_width bits each, LSB-first, contiguous. bit_width=0 means all values
18-
/// equal the frame-of-reference.
17+
/// <p>Metadata: protobuf {@code BitPackedMetadata} — {@code bit_width u32} (tag 1),
18+
/// {@code offset u32} (tag 2, element offset within the first 1024-element block).
19+
///
20+
/// <p>Buffer layout: {@code ceil((len + offset) / 1024)} blocks, each block {@code 128 * bit_width}
21+
/// bytes. Within each block the values are transposed using the FastLanes FL_ORDER permutation so
22+
/// that adjacent bit-planes are contiguous (enables SIMD-friendly decompression).
1923
public final class BitpackedCodec implements Codec {
2024

25+
// FL_ORDER permutation from the FastLanes paper / spiraldb/fastlanes-rs.
26+
private static final int[] FL_ORDER = {0, 4, 2, 6, 1, 5, 3, 7};
27+
2128
@Override
2229
public String encodingId() {
2330
return "fastlanes.bitpacked";
@@ -41,148 +48,244 @@ public EncodeResult encode(DType dtype, Object data) {
4148
PType ptype = ((DType.Primitive) dtype).ptype();
4249
long[] longs = toLongs(data, ptype);
4350
int n = longs.length;
51+
int typeBits = ptype.byteSize() * 8;
52+
long typeMask = typeMask(typeBits);
4453
boolean unsign = isUnsigned(ptype);
4554

46-
long frameOfRef = 0L;
47-
long maxVal = 0L;
55+
long signedMin = 0L;
56+
long signedMax = 0L;
57+
long maxUnsigned = 0L;
4858
int bitWidth = 0;
4959

5060
if (n > 0) {
51-
long min = longs[0];
52-
long max = longs[0];
53-
for (int i = 1; i < n; i++) {
61+
signedMin = longs[0];
62+
signedMax = longs[0];
63+
for (int i = 0; i < n; i++) {
5464
long v = longs[i];
55-
if (unsign ? Long.compareUnsigned(v, min) < 0 : v < min) {
56-
min = v;
65+
if (unsign ? Long.compareUnsigned(v, signedMin) < 0 : v < signedMin) {
66+
signedMin = v;
67+
}
68+
if (unsign ? Long.compareUnsigned(v, signedMax) > 0 : v > signedMax) {
69+
signedMax = v;
5770
}
58-
if (unsign ? Long.compareUnsigned(v, max) > 0 : v > max) {
59-
max = v;
71+
long uv = v & typeMask;
72+
if (Long.compareUnsigned(uv, maxUnsigned) > 0) {
73+
maxUnsigned = uv;
6074
}
6175
}
62-
frameOfRef = min;
63-
maxVal = max;
64-
long range = max - min;
65-
bitWidth = (range == 0L) ? 0 : (64 - Long.numberOfLeadingZeros(range));
76+
bitWidth = maxUnsigned == 0L ? 0 : (Long.SIZE - Long.numberOfLeadingZeros(maxUnsigned));
6677
}
6778

68-
ByteBuffer packed = pack(longs, frameOfRef, bitWidth);
79+
ByteBuffer packed = packFastLanes(longs, n, bitWidth, typeBits);
6980

70-
ByteBuffer meta = ByteBuffer.allocate(9).order(ByteOrder.LITTLE_ENDIAN);
71-
meta.put((byte) bitWidth);
72-
meta.putLong(frameOfRef);
73-
meta.flip();
81+
byte[] metaBytes = EncodingProtos.BitPackedMetadata.newBuilder()
82+
.setBitWidth(bitWidth)
83+
.setOffset(0)
84+
.build()
85+
.toByteArray();
7486

75-
byte[] statsMin = n > 0 ? statsBytes(ptype, frameOfRef) : null;
76-
byte[] statsMax = n > 0 ? statsBytes(ptype, maxVal) : null;
87+
byte[] statsMin = n > 0 ? statsBytes(ptype, signedMin) : null;
88+
byte[] statsMax = n > 0 ? statsBytes(ptype, signedMax) : null;
7789

78-
EncodeNode root = new EncodeNode(encodingId(), meta, new EncodeNode[0], new int[]{0});
90+
EncodeNode root = new EncodeNode(encodingId(), ByteBuffer.wrap(metaBytes),
91+
new EncodeNode[0], new int[]{0});
7992
return new EncodeResult(root, List.of(packed), statsMin, statsMax);
8093
}
8194

8295
// ── Decode ────────────────────────────────────────────────────────────────
8396

8497
@Override
85-
public Array decode(DecodeContext ctx) {
98+
public Array decode(DecodeContext ctx) {
8699
ByteBuffer rawMeta = ctx.metadata();
87-
if (rawMeta == null || rawMeta.capacity() < 2) {
88-
throw new IllegalStateException("fastlanes.bitpacked: missing or truncated metadata");
100+
if (rawMeta == null) {
101+
throw new IllegalStateException("fastlanes.bitpacked: missing metadata");
89102
}
90-
if (rawMeta.capacity() == 9) {
91-
return decodeJava(ctx, rawMeta);
92-
}
93-
// 2-byte JNI format: [constant:u8, bit_width:u8] — FastLanes block packing
94-
return decodeJni(ctx, rawMeta);
95-
}
96103

97-
private Array decodeJava(DecodeContext ctx, ByteBuffer rawMeta) {
98-
ByteBuffer meta = rawMeta.duplicate().order(ByteOrder.LITTLE_ENDIAN);
99-
int bitWidth = Byte.toUnsignedInt(meta.get(0));
100-
long frameOfRef = meta.getLong(1);
101-
102-
PType ptype = ((DType.Primitive) ctx.dtype()).ptype();
103-
long rowCount = ctx.rowCount();
104+
EncodingProtos.BitPackedMetadata meta;
105+
try {
106+
byte[] bytes = new byte[rawMeta.remaining()];
107+
rawMeta.duplicate().get(bytes);
108+
meta = EncodingProtos.BitPackedMetadata.parseFrom(bytes);
109+
} catch (InvalidProtocolBufferException e) {
110+
throw new IllegalStateException("fastlanes.bitpacked: invalid metadata", e);
111+
}
104112

105-
MemorySegment packed = ctx.buffer(0);
106-
byte[] packedBytes = new byte[(int) packed.byteSize()];
107-
MemorySegment.copy(packed, ValueLayout.JAVA_BYTE, 0, packedBytes, 0, packedBytes.length);
113+
int bitWidth = meta.getBitWidth();
114+
int offset = meta.getOffset();
115+
PType ptype = ((DType.Primitive) ctx.dtype()).ptype();
116+
int typeBits = ptype.byteSize() * 8;
117+
long rowCount = ctx.rowCount();
108118

109-
long[] longs = unpack(packedBytes, (int) rowCount, bitWidth, frameOfRef);
119+
MemorySegment packed = ctx.buffer(0);
120+
long[] longs = fastlanesUnpack(packed, bitWidth, offset, typeBits, rowCount);
110121

111122
return new Array(ctx.dtype(), rowCount,
112123
new MemorySegment[]{fromLongs(longs, ptype)}, new Array[0], ArrayStats.empty());
113124
}
114125

115-
// Decode FastLanes block-transposed format (JNI/Rust Vortex writer).
116-
// Layout: ceil(rowCount/1024) blocks, each block = bit_width * 16 * 8 bytes.
117-
// Word index = bitPos * 16 + lane; bit j = bit bitPos of value at index (lane + j*16) in block.
118-
private Array decodeJni(DecodeContext ctx, ByteBuffer rawMeta) {
119-
int bitWidth = Byte.toUnsignedInt(rawMeta.get(1));
120-
121-
PType ptype = ((DType.Primitive) ctx.dtype()).ptype();
122-
long rowCount = ctx.rowCount();
126+
// ── FastLanes pack ────────────────────────────────────────────────────────
123127

124-
MemorySegment packed = ctx.buffer(0);
125-
long[] output = new long[(int) rowCount];
128+
private static ByteBuffer packFastLanes(long[] values, int n, int bitWidth, int typeBits) {
129+
if (bitWidth == 0 || n == 0) {
130+
return ByteBuffer.wrap(new byte[0]);
131+
}
132+
int LANES = 1024 / typeBits;
133+
int wordBytes = typeBits / 8;
134+
int blockCount = (n + 1023) / 1024;
135+
long typeMask = typeMask(typeBits);
136+
byte[] buf = new byte[blockCount * 128 * bitWidth];
126137

127-
int blockCount = (int) ((rowCount + 1023) / 1024);
128138
for (int block = 0; block < blockCount; block++) {
129-
int blockStart = block * 1024;
130-
long blockByteOff = (long) block * bitWidth * 16L * 8L;
131-
for (int bitPos = 0; bitPos < bitWidth; bitPos++) {
132-
for (int lane = 0; lane < 16; lane++) {
133-
long wordByteOff = blockByteOff + ((long) bitPos * 16 + lane) * 8L;
134-
long word = packed.get(
135-
ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN),
136-
wordByteOff);
137-
for (int j = 0; j < 64; j++) {
138-
int vi = blockStart + lane + j * 16;
139-
if (vi >= rowCount) {
140-
break;
141-
}
142-
output[vi] |= ((word >>> j) & 1L) << bitPos;
139+
int blockByteOff = block * 128 * bitWidth;
140+
int blockStart = block * 1024;
141+
142+
for (int row = 0; row < typeBits; row++) {
143+
int currWord = (row * bitWidth) / typeBits;
144+
int nextWord = ((row + 1) * bitWidth) / typeBits;
145+
int shift = (row * bitWidth) % typeBits;
146+
int remainingBits = (nextWord > currWord) ? ((row + 1) * bitWidth) % typeBits : 0;
147+
int currentBits = bitWidth - remainingBits;
148+
149+
for (int lane = 0; lane < LANES; lane++) {
150+
int o = row / 8;
151+
int s = row % 8;
152+
int logicalIdx = blockStart + FL_ORDER[o] * 16 + s * 128 + lane;
153+
long value = (logicalIdx < n) ? (values[logicalIdx] & typeMask) : 0L;
154+
155+
int wordOff = blockByteOff + (LANES * currWord + lane) * wordBytes;
156+
long existing = readWordFromBuf(buf, wordOff, typeBits);
157+
existing |= (value << shift) & typeMask;
158+
writeWordToBuf(buf, wordOff, existing, typeBits);
159+
160+
if (remainingBits > 0) {
161+
int hiWordOff = blockByteOff + (LANES * nextWord + lane) * wordBytes;
162+
long existingHi = readWordFromBuf(buf, hiWordOff, typeBits);
163+
existingHi |= (value >>> currentBits) & typeMask;
164+
writeWordToBuf(buf, hiWordOff, existingHi, typeBits);
143165
}
144166
}
145167
}
146168
}
147-
148-
return new Array(ctx.dtype(), rowCount,
149-
new MemorySegment[]{fromLongs(output, ptype)}, new Array[0], ArrayStats.empty());
169+
return ByteBuffer.wrap(buf);
150170
}
151171

152-
// ── Bit packing ───────────────────────────────────────────────────────────
172+
// ── FastLanes unpack ──────────────────────────────────────────────────────
173+
174+
private static long[] fastlanesUnpack(
175+
MemorySegment buf, int bitWidth, int offset, int typeBits, long rowCount) {
176+
long[] output = new long[(int) rowCount];
177+
if (bitWidth == 0) {
178+
return output;
179+
}
180+
181+
int LANES = 1024 / typeBits;
182+
int wordBytes = typeBits / 8;
183+
long totalElems = rowCount + offset;
184+
int blockCount = (int) ((totalElems + 1023) / 1024);
185+
long bitMask = bitWidth == 64 ? -1L : (1L << bitWidth) - 1L;
153186

154-
private static ByteBuffer pack(long[] values, long frameOfRef, int bitWidth) {
155-
int n = values.length;
156-
int totalBits = n * bitWidth;
157-
byte[] buf = new byte[(totalBits + 7) / 8];
187+
for (int block = 0; block < blockCount; block++) {
188+
long blockByteOff = (long) block * 128L * bitWidth;
189+
int blockLogicStart = block * 1024 - offset;
190+
191+
for (int row = 0; row < typeBits; row++) {
192+
int currWord = (row * bitWidth) / typeBits;
193+
int nextWord = ((row + 1) * bitWidth) / typeBits;
194+
int shift = (row * bitWidth) % typeBits;
195+
int remainingBits = (nextWord > currWord) ? ((row + 1) * bitWidth) % typeBits : 0;
196+
int currentBits = bitWidth - remainingBits;
197+
198+
for (int lane = 0; lane < LANES; lane++) {
199+
int o = row / 8;
200+
int s = row % 8;
201+
int logicalIdx = blockLogicStart + FL_ORDER[o] * 16 + s * 128 + lane;
202+
203+
if (logicalIdx < 0 || logicalIdx >= rowCount) {
204+
continue;
205+
}
206+
207+
long wordOff = blockByteOff + (long) (LANES * currWord + lane) * wordBytes;
208+
long src = readWordFromSeg(buf, wordOff, typeBits);
209+
210+
long value;
211+
if (remainingBits > 0) {
212+
long loMask = (1L << currentBits) - 1L;
213+
long lo = (src >>> shift) & loMask;
214+
long hiOff = blockByteOff + (long) (LANES * nextWord + lane) * wordBytes;
215+
long hi = readWordFromSeg(buf, hiOff, typeBits)
216+
& ((1L << remainingBits) - 1L);
217+
value = lo | (hi << currentBits);
218+
} else {
219+
value = (src >>> shift) & bitMask;
220+
}
158221

159-
for (int i = 0; i < n; i++) {
160-
long shifted = values[i] - frameOfRef;
161-
int bitPos = i * bitWidth;
162-
for (int b = 0; b < bitWidth; b++) {
163-
if (((shifted >>> b) & 1L) == 1L) {
164-
buf[(bitPos + b) >>> 3] |= (byte) (1 << ((bitPos + b) & 7));
222+
output[logicalIdx] = value;
165223
}
166224
}
167225
}
168-
return ByteBuffer.wrap(buf);
226+
return output;
169227
}
170228

171-
private static long[] unpack(byte[] packed, int rowCount, int bitWidth, long frameOfRef) {
172-
long[] out = new long[rowCount];
173-
for (int i = 0; i < rowCount; i++) {
174-
long val = 0L;
175-
int bitPos = i * bitWidth;
176-
for (int b = 0; b < bitWidth; b++) {
177-
int byteIdx = (bitPos + b) >>> 3;
178-
int bitIdx = (bitPos + b) & 7;
179-
if (byteIdx < packed.length && ((packed[byteIdx] >>> bitIdx) & 1) == 1) {
180-
val |= 1L << b;
181-
}
229+
// ── Buffer helpers ────────────────────────────────────────────────────────
230+
231+
private static long readWordFromBuf(byte[] buf, int off, int typeBits) {
232+
return switch (typeBits) {
233+
case 8 -> buf[off] & 0xFFL;
234+
case 16 -> (buf[off] & 0xFFL) | ((buf[off + 1] & 0xFFL) << 8);
235+
case 32 -> Integer.toUnsignedLong(
236+
(buf[off] & 0xFF) | ((buf[off + 1] & 0xFF) << 8)
237+
| ((buf[off + 2] & 0xFF) << 16) | ((buf[off + 3] & 0xFF) << 24));
238+
case 64 -> {
239+
long lo = Integer.toUnsignedLong(
240+
(buf[off] & 0xFF) | ((buf[off + 1] & 0xFF) << 8)
241+
| ((buf[off + 2] & 0xFF) << 16) | ((buf[off + 3] & 0xFF) << 24));
242+
long hi = Integer.toUnsignedLong(
243+
(buf[off + 4] & 0xFF) | ((buf[off + 5] & 0xFF) << 8)
244+
| ((buf[off + 6] & 0xFF) << 16) | ((buf[off + 7] & 0xFF) << 24));
245+
yield lo | (hi << 32);
246+
}
247+
default -> throw new IllegalArgumentException("unsupported typeBits: " + typeBits);
248+
};
249+
}
250+
251+
private static void writeWordToBuf(byte[] buf, int off, long value, int typeBits) {
252+
switch (typeBits) {
253+
case 8 -> buf[off] = (byte) value;
254+
case 16 -> {
255+
buf[off] = (byte) value;
256+
buf[off + 1] = (byte) (value >>> 8);
257+
}
258+
case 32 -> {
259+
buf[off] = (byte) value;
260+
buf[off + 1] = (byte) (value >>> 8);
261+
buf[off + 2] = (byte) (value >>> 16);
262+
buf[off + 3] = (byte) (value >>> 24);
182263
}
183-
out[i] = val + frameOfRef;
264+
case 64 -> {
265+
buf[off] = (byte) value;
266+
buf[off + 1] = (byte) (value >>> 8);
267+
buf[off + 2] = (byte) (value >>> 16);
268+
buf[off + 3] = (byte) (value >>> 24);
269+
buf[off + 4] = (byte) (value >>> 32);
270+
buf[off + 5] = (byte) (value >>> 40);
271+
buf[off + 6] = (byte) (value >>> 48);
272+
buf[off + 7] = (byte) (value >>> 56);
273+
}
274+
default -> throw new IllegalArgumentException("unsupported typeBits: " + typeBits);
184275
}
185-
return out;
276+
}
277+
278+
private static long readWordFromSeg(MemorySegment seg, long off, int typeBits) {
279+
return switch (typeBits) {
280+
case 8 -> Byte.toUnsignedLong(seg.get(ValueLayout.JAVA_BYTE, off));
281+
case 16 -> Short.toUnsignedLong(
282+
seg.get(ValueLayout.JAVA_SHORT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN), off));
283+
case 32 -> Integer.toUnsignedLong(
284+
seg.get(ValueLayout.JAVA_INT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN), off));
285+
case 64 -> seg.get(
286+
ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN), off);
287+
default -> throw new IllegalArgumentException("unsupported typeBits: " + typeBits);
288+
};
186289
}
187290

188291
// ── Type conversion ───────────────────────────────────────────────────────
@@ -247,7 +350,11 @@ private static MemorySegment fromLongs(long[] longs, PType ptype) {
247350
return MemorySegment.ofArray(bytes);
248351
}
249352

250-
// ── Stats helpers ─────────────────────────────────────────────────────────
353+
// ── Misc helpers ──────────────────────────────────────────────────────────
354+
355+
private static long typeMask(int typeBits) {
356+
return typeBits == 64 ? -1L : (1L << typeBits) - 1L;
357+
}
251358

252359
private static byte[] statsBytes(PType ptype, long value) {
253360
if (isUnsigned(ptype)) {

0 commit comments

Comments
 (0)