Skip to content

Commit ed658b7

Browse files
dfa1claude
andcommitted
fix(array): broadcast index for constant-encoded arrays
ConstantEncoding stores one element but reports length=n to avoid OOM from attacker-controlled row_count. All primitive array accessors now use i % physicalLength so constant arrays broadcast correctly without OOBing at i=1. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 1f4a972 commit ed658b7

7 files changed

Lines changed: 37 additions & 24 deletions

File tree

core/src/main/java/io/github/dfa1/vortex/core/array/ByteArray.java

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -59,15 +59,15 @@ public MemorySegment buffer(int i) {
5959
/// @param i zero-based logical index (must be in {@code [0, length)})
6060
/// @return the raw signed byte value at position {@code i}
6161
public byte getByte(long i) {
62-
return buffer.get(ValueLayout.JAVA_BYTE, i);
62+
return buffer.get(ValueLayout.JAVA_BYTE, i % buffer.byteSize());
6363
}
6464

6565
/// Returns the int value at the given logical index, applying unsigned widening for U8 columns.
6666
///
6767
/// @param i zero-based logical index (must be in {@code [0, length)})
6868
/// @return the value at position {@code i} as a signed int (U8 values are zero-extended)
6969
public int getInt(long i) {
70-
byte raw = buffer.get(ValueLayout.JAVA_BYTE, i);
70+
byte raw = buffer.get(ValueLayout.JAVA_BYTE, i % buffer.byteSize());
7171
boolean unsigned = dtype instanceof DType.Primitive p && p.ptype() == PType.U8;
7272
return unsigned ? Byte.toUnsignedInt(raw) : raw;
7373
}
@@ -80,9 +80,10 @@ public int getInt(long i) {
8080
public long fold(long identity, LongBinaryOperator op) {
8181
MemorySegment buf = buffer;
8282
long n = length;
83+
long cap = buf.byteSize();
8384
long result = identity;
8485
for (long i = 0; i < n; i++) {
85-
result = op.applyAsLong(result, buf.get(ValueLayout.JAVA_BYTE, i));
86+
result = op.applyAsLong(result, buf.get(ValueLayout.JAVA_BYTE, i % cap));
8687
}
8788
return result;
8889
}

core/src/main/java/io/github/dfa1/vortex/core/array/DoubleArray.java

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,8 @@ public MemorySegment buffer(int i) {
5959
/// @param i zero-based logical index (must be in {@code [0, length)})
6060
/// @return the double value at position {@code i}
6161
public double getDouble(long i) {
62-
return buffer.getAtIndex(PTypeIO.LE_DOUBLE, i);
62+
long cap = buffer.byteSize() / PTypeIO.LE_DOUBLE.byteSize();
63+
return buffer.getAtIndex(PTypeIO.LE_DOUBLE, i % cap);
6364
}
6465

6566
/// Invokes the consumer for each element in order.
@@ -68,8 +69,9 @@ public double getDouble(long i) {
6869
public void forEachDouble(DoubleConsumer c) {
6970
MemorySegment buf = buffer;
7071
long n = length;
72+
long cap = buf.byteSize() / PTypeIO.LE_DOUBLE.byteSize();
7173
for (long i = 0; i < n; i++) {
72-
c.accept(buf.getAtIndex(PTypeIO.LE_DOUBLE, i));
74+
c.accept(buf.getAtIndex(PTypeIO.LE_DOUBLE, i % cap));
7375
}
7476
}
7577

@@ -79,11 +81,10 @@ public void forEachDouble(DoubleConsumer c) {
7981
/// @param op binary operator applied to accumulator and each element in order
8082
/// @return the final accumulated value
8183
public double fold(double identity, DoubleBinaryOperator op) {
82-
MemorySegment buf = buffer;
83-
long n = length;
84+
long cap = buffer.byteSize() / PTypeIO.LE_DOUBLE.byteSize();
8485
double result = identity;
85-
for (long i = 0; i < n; i++) {
86-
result = op.applyAsDouble(result, buf.getAtIndex(PTypeIO.LE_DOUBLE, i));
86+
for (long i = 0; i < length; i++) {
87+
result = op.applyAsDouble(result, buffer.getAtIndex(PTypeIO.LE_DOUBLE, i % cap));
8788
}
8889
return result;
8990
}

core/src/main/java/io/github/dfa1/vortex/core/array/FloatArray.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,8 @@ public MemorySegment buffer(int i) {
5858
/// @param i zero-based index (must be in {@code [0, length)})
5959
/// @return the float value at position {@code i}
6060
public float getFloat(long i) {
61-
return buffer.getAtIndex(PTypeIO.LE_FLOAT, i);
61+
long cap = buffer.byteSize() / PTypeIO.LE_FLOAT.byteSize();
62+
return buffer.getAtIndex(PTypeIO.LE_FLOAT, i % cap);
6263
}
6364

6465
/// Folds all elements using the given binary operator and identity value.
@@ -69,9 +70,10 @@ public float getFloat(long i) {
6970
public double fold(double identity, DoubleBinaryOperator op) {
7071
MemorySegment buf = buffer;
7172
long n = length;
73+
long cap = buf.byteSize() / PTypeIO.LE_FLOAT.byteSize();
7274
double result = identity;
7375
for (long i = 0; i < n; i++) {
74-
result = op.applyAsDouble(result, buf.getAtIndex(PTypeIO.LE_FLOAT, i));
76+
result = op.applyAsDouble(result, buf.getAtIndex(PTypeIO.LE_FLOAT, i % cap));
7577
}
7678
return result;
7779
}

core/src/main/java/io/github/dfa1/vortex/core/array/IntArray.java

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,8 @@ public MemorySegment buffer(int i) {
5959
/// @param i zero-based index (must be in {@code [0, length)})
6060
/// @return the int value at position {@code i}
6161
public int getInt(long i) {
62-
return buffer.getAtIndex(PTypeIO.LE_INT, i);
62+
long cap = buffer.byteSize() / PTypeIO.LE_INT.byteSize();
63+
return buffer.getAtIndex(PTypeIO.LE_INT, i % cap);
6364
}
6465

6566
/// Passes each element to the given consumer in order.
@@ -68,8 +69,9 @@ public int getInt(long i) {
6869
public void forEachInt(IntConsumer c) {
6970
MemorySegment buf = buffer;
7071
long n = length;
72+
long cap = buf.byteSize() / PTypeIO.LE_INT.byteSize();
7173
for (long i = 0; i < n; i++) {
72-
c.accept(buf.getAtIndex(PTypeIO.LE_INT, i));
74+
c.accept(buf.getAtIndex(PTypeIO.LE_INT, i % cap));
7375
}
7476
}
7577

@@ -79,11 +81,10 @@ public void forEachInt(IntConsumer c) {
7981
/// @param op binary operator applied to the accumulator and each int element
8082
/// @return the final accumulated result
8183
public int fold(int identity, IntBinaryOperator op) {
82-
MemorySegment buf = buffer;
83-
long n = length;
84+
long cap = buffer.byteSize() / PTypeIO.LE_INT.byteSize();
8485
int result = identity;
85-
for (long i = 0; i < n; i++) {
86-
result = op.applyAsInt(result, buf.getAtIndex(PTypeIO.LE_INT, i));
86+
for (long i = 0; i < length; i++) {
87+
result = op.applyAsInt(result, buffer.getAtIndex(PTypeIO.LE_INT, i % cap));
8788
}
8889
return result;
8990
}

core/src/main/java/io/github/dfa1/vortex/core/array/LongArray.java

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,8 @@ public MemorySegment buffer(int i) {
5959
/// @param i zero-based index (must be in {@code [0, length)})
6060
/// @return the long value at position {@code i}
6161
public long getLong(long i) {
62-
return buffer.getAtIndex(PTypeIO.LE_LONG, i);
62+
long cap = buffer.byteSize() / PTypeIO.LE_LONG.byteSize();
63+
return buffer.getAtIndex(PTypeIO.LE_LONG, i % cap);
6364
}
6465

6566
/// Passes each element to the given consumer in order.
@@ -68,8 +69,9 @@ public long getLong(long i) {
6869
public void forEachLong(LongConsumer c) {
6970
MemorySegment buf = buffer;
7071
long n = length;
72+
long cap = buf.byteSize() / PTypeIO.LE_LONG.byteSize();
7173
for (long i = 0; i < n; i++) {
72-
c.accept(buf.getAtIndex(PTypeIO.LE_LONG, i));
74+
c.accept(buf.getAtIndex(PTypeIO.LE_LONG, i % cap));
7375
}
7476
}
7577

@@ -81,9 +83,10 @@ public void forEachLong(LongConsumer c) {
8183
public long fold(long identity, LongBinaryOperator op) {
8284
MemorySegment buf = buffer;
8385
long n = length;
86+
long cap = buf.byteSize() / PTypeIO.LE_LONG.byteSize();
8487
long result = identity;
8588
for (long i = 0; i < n; i++) {
86-
result = op.applyAsLong(result, buf.getAtIndex(PTypeIO.LE_LONG, i));
89+
result = op.applyAsLong(result, buf.getAtIndex(PTypeIO.LE_LONG, i % cap));
8790
}
8891
return result;
8992
}

core/src/main/java/io/github/dfa1/vortex/core/array/ShortArray.java

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -59,15 +59,17 @@ public MemorySegment buffer(int i) {
5959
/// @param i zero-based index (must be in {@code [0, length)})
6060
/// @return the signed short value at position {@code i}
6161
public short getShort(long i) {
62-
return buffer.getAtIndex(PTypeIO.LE_SHORT, i);
62+
long cap = buffer.byteSize() / PTypeIO.LE_SHORT.byteSize();
63+
return buffer.getAtIndex(PTypeIO.LE_SHORT, i % cap);
6364
}
6465

6566
/// Returns the element at the given index as an {@code int}, widening to unsigned if the dtype is U16.
6667
///
6768
/// @param i zero-based index (must be in {@code [0, length)})
6869
/// @return the element at position {@code i} as an int (unsigned-widened for U16)
6970
public int getInt(long i) {
70-
short raw = buffer.getAtIndex(PTypeIO.LE_SHORT, i);
71+
long cap = buffer.byteSize() / PTypeIO.LE_SHORT.byteSize();
72+
short raw = buffer.getAtIndex(PTypeIO.LE_SHORT, i % cap);
7173
boolean unsigned = dtype instanceof DType.Primitive p && p.ptype() == PType.U16;
7274
return unsigned ? Short.toUnsignedInt(raw) : raw;
7375
}
@@ -80,9 +82,10 @@ public int getInt(long i) {
8082
public long fold(long identity, LongBinaryOperator op) {
8183
MemorySegment buf = buffer;
8284
long n = length;
85+
long cap = buf.byteSize() / PTypeIO.LE_SHORT.byteSize();
8386
long result = identity;
8487
for (long i = 0; i < n; i++) {
85-
result = op.applyAsLong(result, buf.getAtIndex(PTypeIO.LE_SHORT, i));
88+
result = op.applyAsLong(result, buf.getAtIndex(PTypeIO.LE_SHORT, i % cap));
8689
}
8790
return result;
8891
}

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import io.github.dfa1.vortex.core.array.ShortArray;
2020
import io.github.dfa1.vortex.core.array.VarBinArray;
2121
import io.github.dfa1.vortex.encoding.EncodingRegistry;
22+
import io.github.dfa1.vortex.io.VortexInspector;
2223
import io.github.dfa1.vortex.io.VortexReader;
2324
import io.github.dfa1.vortex.scan.ScanResult;
2425
import org.apache.arrow.memory.BufferAllocator;
@@ -383,7 +384,8 @@ private static Long stringByteLength(Array arr) {
383384
void rust_vs_javaReader_statsMatch(String fixture, @TempDir Path tmp) throws Exception {
384385
// Given
385386
Path local = download(BASE.resolve(fixture), tmp);
386-
387+
String inspect = VortexInspector.inspect(VortexReader.open(local));
388+
System.out.println(inspect);
387389
// When
388390
Stats rustStats = rustStats(local);
389391
Stats javaStats = javaStats(local);

0 commit comments

Comments
 (0)