Skip to content

Commit ede7f84

Browse files
dfa1claude
dfa1
and
claude
committed
fix(encoding): broadcast 1-element constant child segments on bulk read
Commit 10a7776 made ConstantEncoding allocate a single element regardless of declared rowCount as a zip-bomb defense. Consumers that bulk-read a child segment via raw `MemorySegment.copy(seg, i * elemBytes, ...)` or `seg.get(LAYOUT, i * sz)` then ran off the end of the 1-element segment as soon as i > 0 — surfacing as IndexOutOfBoundsException in RustVsJavaReadBenchmark.javaReadClose (ALP F64 column with bitpacked patches whose values child was constant-encoded). LongArray.getLong and sibling typed accessors already wrap with `i % cap`; this commit teaches the segment-level consumers the same broadcast arithmetic via a new SegmentBroadcast helper exposing elementOffset, capacity, and broadcastCopy. Patched call sites (every decodeChildSegment consumer that reads more than one element by offset): - BitpackedEncoding.applyPatches (idx + val) - SparseEncoding.applyPatches, decodeBool, decodeVarBin - AlpEncoding.decodeF64/decodeF32 source loop and applyPatches - AlpRdEncoding.decodeF64/decodeF32 left+right loops and patches - PatchedEncoding.decode bulk inner copy and applyPatches - ZigZagEncoding.decode per-ptype loops - DeltaEncoding.readLongs - RleEncoding.readLongs, readIndices, readUnsignedLongs - RunEndEncoding.expandByte/Short/Int/Long/Bool/Strings - DictEncoding.readCode and expandU8/U16/U32 - FsstEncoding decode (uncompressed lens + codes offsets) - ChunkedEncoding.readOffsets - VarBinEncoding.decode (materialize 1-element offsets before VarBinArray) Added PatchesBroadcastRegressionTest with a hand-built ArrayNode tree that puts ConstantEncoding under bitpacked patches, plus direct coverage for SegmentBroadcast.elementOffset/capacity/broadcastCopy. Bench (RustVsJavaReadBenchmark): javaReadClose now passes at 60.7 ops/s vs JNI 49.1 ops/s (Java 1.24× faster). Volume, Symbol, Cascading, and JNI methods remain within run-to-run noise of the prior numbers — modulo arithmetic is hoisted by the JIT as a loop invariant. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1 parent e696f07 commit ede7f84

15 files changed

Lines changed: 409 additions & 113 deletions

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

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -423,11 +423,13 @@ private static Array decodeF64(DecodeContext ctx, EncodingProtos.ALPMetadata met
423423

424424
MemorySegment src = ctx.decodeChildSegment(0, I64_DTYPE, n);
425425
// In-place when the child returned a writable arena buffer (e.g. BitpackedEncoding, DeltaEncoding).
426-
// Fall back to a new allocation when the source is a read-only mmap slice (PrimitiveEncoding).
426+
// Fall back to a new allocation when the source is a read-only mmap slice (PrimitiveEncoding)
427+
// or a broadcasted single element (ConstantEncoding).
427428
MemorySegment buf = src.isReadOnly() ? ctx.arena().allocate(n * 8, 8) : src;
428429
if (src.isReadOnly()) {
430+
long srcCap = SegmentBroadcast.capacity(src, 8);
429431
for (long i = 0; i < n; i++) {
430-
buf.setAtIndex(PTypeIO.LE_DOUBLE, i, (double) src.getAtIndex(PTypeIO.LE_LONG, i) * df * de);
432+
buf.setAtIndex(PTypeIO.LE_DOUBLE, i, (double) src.getAtIndex(PTypeIO.LE_LONG, i % srcCap) * df * de);
431433
}
432434
} else {
433435
for (long i = 0; i < n; i++) {
@@ -450,8 +452,9 @@ private static Array decodeF32(DecodeContext ctx, EncodingProtos.ALPMetadata met
450452
MemorySegment src32 = ctx.decodeChildSegment(0, I32_DTYPE, n);
451453
MemorySegment buf32 = src32.isReadOnly() ? ctx.arena().allocate(n * 4, 4) : src32;
452454
if (src32.isReadOnly()) {
455+
long srcCap = SegmentBroadcast.capacity(src32, 4);
453456
for (long i = 0; i < n; i++) {
454-
buf32.setAtIndex(PTypeIO.LE_FLOAT, i, (float) src32.getAtIndex(PTypeIO.LE_INT, i) * df * de);
457+
buf32.setAtIndex(PTypeIO.LE_FLOAT, i, (float) src32.getAtIndex(PTypeIO.LE_INT, i % srcCap) * df * de);
455458
}
456459
} else {
457460
for (long i = 0; i < n; i++) {
@@ -471,22 +474,24 @@ private static void applyPatches(DecodeContext ctx, EncodingProtos.PatchesMetada
471474
long numPatches = pm.getLen();
472475
long offset = pm.getOffset();
473476
PType idxPtype = ptypeFromProto(pm.getIndicesPtype());
477+
int idxBytes = idxPtype.byteSize();
474478

475479
MemorySegment idxSeg = ctx.decodeChildSegment(1, new DType.Primitive(idxPtype, false), numPatches);
476480
MemorySegment valSeg = ctx.decodeChildSegment(2, ctx.dtype(), numPatches);
477481

478482
for (long i = 0; i < numPatches; i++) {
479-
long absIdx = readUnsigned(idxSeg, i, idxPtype) - offset;
480-
MemorySegment.copy(valSeg, i * elemBytes, out, absIdx * elemBytes, elemBytes);
483+
long absIdx = readUnsigned(idxSeg, SegmentBroadcast.elementOffset(idxSeg, i, idxBytes), idxPtype) - offset;
484+
MemorySegment.copy(valSeg, SegmentBroadcast.elementOffset(valSeg, i, elemBytes),
485+
out, absIdx * elemBytes, elemBytes);
481486
}
482487
}
483488

484-
private static long readUnsigned(MemorySegment seg, long i, PType ptype) {
489+
private static long readUnsigned(MemorySegment seg, long off, PType ptype) {
485490
return switch (ptype) {
486-
case U8 -> Byte.toUnsignedLong(seg.get(ValueLayout.JAVA_BYTE, i));
487-
case U16 -> Short.toUnsignedLong(seg.get(PTypeIO.LE_SHORT, i * 2));
488-
case U32 -> Integer.toUnsignedLong(seg.get(PTypeIO.LE_INT, i * 4));
489-
case U64 -> seg.get(PTypeIO.LE_LONG, i * 8);
491+
case U8 -> Byte.toUnsignedLong(seg.get(ValueLayout.JAVA_BYTE, off));
492+
case U16 -> Short.toUnsignedLong(seg.get(PTypeIO.LE_SHORT, off));
493+
case U32 -> Integer.toUnsignedLong(seg.get(PTypeIO.LE_INT, off));
494+
case U64 -> seg.get(PTypeIO.LE_LONG, off);
490495
default -> throw new VortexException(EncodingId.VORTEX_ALP, "non-unsigned patch index ptype " + ptype);
491496
};
492497
}

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

Lines changed: 28 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -399,17 +399,19 @@ private static Array decodeF64(DecodeContext ctx,
399399
short[] dict, int rightBitWidth, long n) {
400400
MemorySegment leftSeg = ctx.decodeChildSegment(0, U16_DTYPE, n);
401401
MemorySegment rightSeg = ctx.decodeChildSegment(1, U64_DTYPE, n);
402+
long leftCap = SegmentBroadcast.capacity(leftSeg, 2);
403+
long rightCap = SegmentBroadcast.capacity(rightSeg, 8);
402404
MemorySegment out = ctx.arena().allocate(n * Long.BYTES, Long.BYTES);
403405

404406
for (long i = 0; i < n; i++) {
405-
int code = Short.toUnsignedInt(leftSeg.getAtIndex(PTypeIO.LE_SHORT, i));
407+
int code = Short.toUnsignedInt(leftSeg.getAtIndex(PTypeIO.LE_SHORT, i % leftCap));
406408
long leftBits = (long) (dict[code] & 0xFFFF) << rightBitWidth;
407-
long rightBits = rightSeg.getAtIndex(PTypeIO.LE_LONG, i);
409+
long rightBits = rightSeg.getAtIndex(PTypeIO.LE_LONG, i % rightCap);
408410
out.setAtIndex(PTypeIO.LE_LONG, i, leftBits | rightBits);
409411
}
410412

411413
if (meta.hasPatches()) {
412-
applyPatchesF64(ctx, meta.getPatches(), out, rightSeg, rightBitWidth);
414+
applyPatchesF64(ctx, meta.getPatches(), out, rightSeg, rightCap, rightBitWidth);
413415
}
414416

415417
return new DoubleArray(ctx.dtype(), n, out.asReadOnly());
@@ -420,64 +422,71 @@ private static Array decodeF32(DecodeContext ctx,
420422
short[] dict, int rightBitWidth, long n) {
421423
MemorySegment leftSeg = ctx.decodeChildSegment(0, U16_DTYPE, n);
422424
MemorySegment rightSeg = ctx.decodeChildSegment(1, U32_DTYPE, n);
425+
long leftCap = SegmentBroadcast.capacity(leftSeg, 2);
426+
long rightCap = SegmentBroadcast.capacity(rightSeg, 4);
423427
MemorySegment out = ctx.arena().allocate(n * Integer.BYTES, Integer.BYTES);
424428

425429
for (long i = 0; i < n; i++) {
426-
int code = Short.toUnsignedInt(leftSeg.getAtIndex(PTypeIO.LE_SHORT, i));
430+
int code = Short.toUnsignedInt(leftSeg.getAtIndex(PTypeIO.LE_SHORT, i % leftCap));
427431
int leftBits = (dict[code] & 0xFFFF) << rightBitWidth;
428-
int rightBits = rightSeg.getAtIndex(PTypeIO.LE_INT, i);
432+
int rightBits = rightSeg.getAtIndex(PTypeIO.LE_INT, i % rightCap);
429433
out.setAtIndex(PTypeIO.LE_INT, i, leftBits | rightBits);
430434
}
431435

432436
if (meta.hasPatches()) {
433-
applyPatchesF32(ctx, meta.getPatches(), out, rightSeg, rightBitWidth);
437+
applyPatchesF32(ctx, meta.getPatches(), out, rightSeg, rightCap, rightBitWidth);
434438
}
435439

436440
return new FloatArray(ctx.dtype(), n, out.asReadOnly());
437441
}
438442

439443
private static void applyPatchesF64(DecodeContext ctx,
440444
EncodingProtos.PatchesMetadata pm,
441-
MemorySegment out, MemorySegment rightSeg, int rightBitWidth) {
445+
MemorySegment out, MemorySegment rightSeg, long rightCap, int rightBitWidth) {
442446
long numPatches = pm.getLen();
443447
long offset = pm.getOffset();
444448
PType idxPtype = PType.fromOrdinal(pm.getIndicesPtype().getNumber());
445449

446450
MemorySegment idxSeg = ctx.decodeChildSegment(2, new DType.Primitive(idxPtype, false), numPatches);
447451
MemorySegment valSeg = ctx.decodeChildSegment(3, U16_DTYPE, numPatches);
452+
int idxBytes = idxPtype.byteSize();
453+
long valCap = SegmentBroadcast.capacity(valSeg, 2);
448454

449455
for (long j = 0; j < numPatches; j++) {
450-
long absIdx = readUnsigned(idxSeg, j, idxPtype) - offset;
451-
short actualLeftU16 = valSeg.getAtIndex(PTypeIO.LE_SHORT, j);
456+
long absIdx = readUnsigned(idxSeg, SegmentBroadcast.elementOffset(idxSeg, j, idxBytes), idxPtype) - offset;
457+
short actualLeftU16 = valSeg.getAtIndex(PTypeIO.LE_SHORT, j % valCap);
452458
long leftBits = (long) (actualLeftU16 & 0xFFFF) << rightBitWidth;
453-
long rightBits = rightSeg.getAtIndex(PTypeIO.LE_LONG, absIdx);
459+
long rightBits = rightSeg.getAtIndex(PTypeIO.LE_LONG, absIdx % rightCap);
454460
out.setAtIndex(PTypeIO.LE_LONG, absIdx, leftBits | rightBits);
455461
}
456462
}
457463

458-
private static void applyPatchesF32(DecodeContext ctx, EncodingProtos.PatchesMetadata pm, MemorySegment out, MemorySegment rightSeg, int rightBitWidth) {
464+
private static void applyPatchesF32(DecodeContext ctx, EncodingProtos.PatchesMetadata pm,
465+
MemorySegment out, MemorySegment rightSeg, long rightCap, int rightBitWidth) {
459466
long numPatches = pm.getLen();
460467
long offset = pm.getOffset();
461468
PType idxPtype = PType.fromOrdinal(pm.getIndicesPtype().getNumber());
462469

463470
MemorySegment idxSeg = ctx.decodeChildSegment(2, new DType.Primitive(idxPtype, false), numPatches);
464471
MemorySegment valSeg = ctx.decodeChildSegment(3, U16_DTYPE, numPatches);
472+
int idxBytes = idxPtype.byteSize();
473+
long valCap = SegmentBroadcast.capacity(valSeg, 2);
465474

466475
for (long j = 0; j < numPatches; j++) {
467-
long absIdx = readUnsigned(idxSeg, j, idxPtype) - offset;
468-
short actualLeftU16 = valSeg.getAtIndex(PTypeIO.LE_SHORT, j);
476+
long absIdx = readUnsigned(idxSeg, SegmentBroadcast.elementOffset(idxSeg, j, idxBytes), idxPtype) - offset;
477+
short actualLeftU16 = valSeg.getAtIndex(PTypeIO.LE_SHORT, j % valCap);
469478
int leftBits = (actualLeftU16 & 0xFFFF) << rightBitWidth;
470-
int rightBits = rightSeg.getAtIndex(PTypeIO.LE_INT, absIdx);
479+
int rightBits = rightSeg.getAtIndex(PTypeIO.LE_INT, absIdx % rightCap);
471480
out.setAtIndex(PTypeIO.LE_INT, (int) absIdx, leftBits | rightBits);
472481
}
473482
}
474483

475-
private static long readUnsigned(MemorySegment seg, long i, PType ptype) {
484+
private static long readUnsigned(MemorySegment seg, long off, PType ptype) {
476485
return switch (ptype) {
477-
case U8 -> Byte.toUnsignedLong(seg.get(java.lang.foreign.ValueLayout.JAVA_BYTE, i));
478-
case U16 -> Short.toUnsignedLong(seg.get(PTypeIO.LE_SHORT, i * 2));
479-
case U32 -> Integer.toUnsignedLong(seg.get(PTypeIO.LE_INT, i * 4));
480-
case U64 -> seg.get(PTypeIO.LE_LONG, i * 8);
486+
case U8 -> Byte.toUnsignedLong(seg.get(java.lang.foreign.ValueLayout.JAVA_BYTE, off));
487+
case U16 -> Short.toUnsignedLong(seg.get(PTypeIO.LE_SHORT, off));
488+
case U32 -> Integer.toUnsignedLong(seg.get(PTypeIO.LE_INT, off));
489+
case U64 -> seg.get(PTypeIO.LE_LONG, off);
481490
default -> throw new VortexException(EncodingId.VORTEX_ALPRD,
482491
"non-unsigned patch index ptype " + ptype);
483492
};

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

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -619,23 +619,25 @@ private static void applyPatches(DecodeContext ctx, EncodingProtos.PatchesMetada
619619
MemorySegment idxSeg = ctx.decodeChildSegment(0, new DType.Primitive(idxPtype, false), numPatches);
620620
MemorySegment valSeg = ctx.decodeChildSegment(1, ctx.dtype(), numPatches);
621621

622+
int idxBytes = idxPtype.byteSize();
622623
long n = ctx.rowCount();
623624
for (long i = 0; i < numPatches; i++) {
624-
long absIdx = readUnsignedIdx(idxSeg, i, idxPtype) - offset;
625+
long absIdx = readUnsignedIdx(idxSeg, SegmentBroadcast.elementOffset(idxSeg, i, idxBytes), idxPtype) - offset;
625626
if (absIdx < 0 || absIdx >= n) {
626627
throw new VortexException(EncodingId.FASTLANES_BITPACKED,
627628
"patch index " + absIdx + " out of range [0," + n + ")");
628629
}
629-
MemorySegment.copy(valSeg, i * elemBytes, out, absIdx * elemBytes, elemBytes);
630+
MemorySegment.copy(valSeg, SegmentBroadcast.elementOffset(valSeg, i, elemBytes),
631+
out, absIdx * elemBytes, elemBytes);
630632
}
631633
}
632634

633-
private static long readUnsignedIdx(MemorySegment seg, long i, PType ptype) {
635+
private static long readUnsignedIdx(MemorySegment seg, long off, PType ptype) {
634636
return switch (ptype) {
635-
case U8 -> Byte.toUnsignedLong(seg.get(ValueLayout.JAVA_BYTE, i));
636-
case U16 -> Short.toUnsignedLong(seg.get(PTypeIO.LE_SHORT, i * 2));
637-
case U32 -> Integer.toUnsignedLong(seg.get(PTypeIO.LE_INT, i * 4));
638-
case U64 -> seg.get(PTypeIO.LE_LONG, i * 8);
637+
case U8 -> Byte.toUnsignedLong(seg.get(ValueLayout.JAVA_BYTE, off));
638+
case U16 -> Short.toUnsignedLong(seg.get(PTypeIO.LE_SHORT, off));
639+
case U32 -> Integer.toUnsignedLong(seg.get(PTypeIO.LE_INT, off));
640+
case U64 -> seg.get(PTypeIO.LE_LONG, off);
639641
default -> throw new VortexException(EncodingId.FASTLANES_BITPACKED,
640642
"non-unsigned patch index ptype " + ptype);
641643
};

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,9 +139,10 @@ static Array decode(DecodeContext ctx) {
139139
private static long[] readOffsets(DecodeContext ctx, int nchunks) {
140140
DType u64 = new DType.Primitive(PType.U64, false);
141141
MemorySegment offsetsBuf = ctx.decodeChildSegment(0, u64, nchunks + 1L);
142+
long cap = SegmentBroadcast.capacity(offsetsBuf, 8);
142143
long[] offsets = new long[nchunks + 1];
143144
for (int i = 0; i <= nchunks; i++) {
144-
offsets[i] = offsetsBuf.get(LE_LONG, (long) i * 8);
145+
offsets[i] = offsetsBuf.get(LE_LONG, (i % cap) * 8);
145146
}
146147
return offsets;
147148
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -357,8 +357,9 @@ private static void undeltaChunk(long[] deltas, long[] bases, int lanes, int typ
357357
private static long[] readLongs(MemorySegment buf, int count, PType ptype) {
358358
long[] out = new long[count];
359359
int elemSize = ptype.byteSize();
360+
long cap = SegmentBroadcast.capacity(buf, elemSize);
360361
for (int i = 0; i < count; i++) {
361-
long off = (long) i * elemSize;
362+
long off = (i % cap) * elemSize;
362363
out[i] = switch (ptype) {
363364
case I8 -> buf.get(ValueLayout.JAVA_BYTE, off);
364365
case U8 -> Byte.toUnsignedLong(buf.get(ValueLayout.JAVA_BYTE, off));

0 commit comments

Comments
 (0)