dongjoon-hyun commented on code in PR #59117:
URL: https://github.com/apache/spark/pull/59117#discussion_r4128797803
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReaderSuite.scala:
##########
@@ -185,10 +203,210 @@ class VectorizedRleValuesReaderSuite extends
SparkFunSuite {
produced += toRead
}
}
+ test("SPARK-59832: truncated dictionary ids fail instead of being partially
read") {
+ // A dictionary id page with 10 ids, read as if it had 20.
+ streams(dictIdPage(Array.fill(10)(3), bitWidth = 4)).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ reader.initFromPage(20, in())
+ val c = new OnHeapColumnVector(20, IntegerType)
+ interceptCorrupted(reader.readIntegers(20, c, 0))
+ }
+ streams(dictIdPage(Array.fill(10)(3), bitWidth = 4)).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ reader.initFromPage(20, in())
+ interceptCorrupted(reader.skipIntegers(20))
+ }
+ streams(dictIdPage(Array.fill(10)(3), bitWidth = 4)).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ reader.initFromPage(20, in())
+ (0 until 10).foreach(_ => assert(reader.readInteger() == 3))
+ interceptCorrupted(reader.readInteger())
+ }
+ }
+
+ test("SPARK-59832: truncated booleans fail instead of looping forever") {
+ val encoded = encodeRle(Array.fill(10)(1), 1)
+ streams(intLE(encoded.length) ++ encoded).foreach { in =>
+ val reader = new VectorizedRleValuesReader(1)
+ reader.initFromPage(20, in())
+ val c = new OnHeapColumnVector(20, BooleanType)
+ interceptCorrupted(reader.readBooleans(20, c, 0))
+ }
+ }
+
+ test("SPARK-59832: truncated definition levels fail in readBatch") {
+ // The page declares 100 values but its definition levels only encode 64.
Before the fix
+ // readBatch returned without progress, so
VectorizedColumnReader.readBatch spun forever.
+ val n = 100
+ val encoded = encodeRle(Array.fill(64)(1), 1)
+ val plainBytes = plainIntBytes(n)(valueAt)
+ def run(withDefLevels: Boolean, rowIndexes: PrimitiveIterator.OfLong):
Unit = {
+ val reader = new VectorizedRleValuesReader(1, false)
+ reader.initFromPage(n,
ByteBufferInputStream.wrap(ByteBuffer.wrap(encoded)))
+ val valueReader = new VectorizedPlainValuesReader
+ valueReader.initFromPage(n,
ByteBufferInputStream.wrap(ByteBuffer.wrap(plainBytes)))
+ val state = ParquetTestAccess.newState(intColumnDescriptor(1), false,
rowIndexes)
+ ParquetTestAccess.resetForNewPage(state, n, 0L)
+ ParquetTestAccess.resetForNewBatch(state, n)
+ val values = new OnHeapColumnVector(n, IntegerType)
+ val defLevels = if (withDefLevels) new OnHeapColumnVector(n,
IntegerType) else null
+ interceptCorrupted {
+ ParquetTestAccess.readBatch(reader, state, values, defLevels,
valueReader, integerUpdater)
+ }
+ }
+ run(withDefLevels = false, rowIndexes = null)
+ run(withDefLevels = true, rowIndexes = null)
+ // Skips the truncated levels to reach rows 80 to 90.
+ run(withDefLevels = false, rowIndexes = longIterator((80 to 90).toArray))
+ }
+
+ test("SPARK-59832: truncated repetition levels fail in readBatchRepeated") {
Review Comment:
None of the new tests reach the new throw sites in `readValues` (L649) or
`skipValues(int, ParquetReadState, ...)` (L743). This test truncates only the
repetition levels, and in the flat row-range case `readBatchInternal` never
calls `skipValues` with more than `currentCount`, so it throws at L220 instead.
On master, truncated *definition* levels in the repeated path don't hang; they
silently produce wrong nested values because the trailing def levels stay 0.
Could you add a `readBatchRepeated` case that truncates only the def-level
stream, with and without row indexes?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReader.java:
##########
@@ -100,12 +100,21 @@ public void initFromPage(int valueCount,
ByteBufferInputStream in) throws IOExce
// Initialize for repetition and definition levels
if (readLength) {
int length = readIntLittleEndian();
+ if (length < 0 || length > in.available()) {
Review Comment:
Before this change, a positive length larger than the rest of the page made
`sliceStream` throw an `EOFException`, and `readPageV1` (or `initDataReader`
for RLE booleans) wrapped it as "could not read page ... in col ...". The new
`ParquetDecodingException` is unchecked, so it skips those wrappers and the
error no longer says which column is corrupt. It matches SPARK-59812, but for
wide tables it may be worth also catching `ParquetDecodingException` there (or
in `VectorizedColumnReader.readBatch`) and adding the column descriptor.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReader.java:
##########
@@ -976,48 +985,58 @@ private int readIntLittleEndianPaddedOnBitWidth() throws
IOException {
}
/**
- * Reads the next group. Returns false if no more group available.
+ * Reads the next non-empty group, so that `currentCount` is positive on
return.
+ * <p>
+ * Callers never read more values than the page declares, so running out of
encoded data here
+ * means the page is corrupted. With a bit width of 0, `initFromPage`
already provides all the
+ * values of the page, so reaching here also means reading past the end.
*/
- private boolean readNextGroup() {
- if (in.available() <= 0) {
- currentCount = 0;
- return false;
- }
-
+ private void readNextGroup() {
try {
- int header = readUnsignedVarInt();
- this.mode = (header & 1) == 0 ? MODE.RLE : MODE.PACKED;
- switch (mode) {
- case RLE -> {
- this.currentCount = header >>> 1;
- this.currentValue = readIntLittleEndianPaddedOnBitWidth();
+ do {
+ if (bitWidth == 0 || in.available() <= 0) {
+ throw new ParquetDecodingException(
+ "Corrupted RLE data: reading past the end of the encoded values");
}
- case PACKED -> {
- int numGroups = header >>> 1;
- this.currentCount = numGroups * 8;
-
- if (this.currentBuffer.length < this.currentCount) {
- this.currentBuffer = new int[this.currentCount];
+ int header = readUnsignedVarInt();
+ this.mode = (header & 1) == 0 ? MODE.RLE : MODE.PACKED;
+ switch (mode) {
+ case RLE -> {
+ this.currentCount = header >>> 1;
+ this.currentValue = readIntLittleEndianPaddedOnBitWidth();
}
- currentBufferIdx = 0;
- // Slice all packed bytes in one call (numGroups groups x bitWidth
bytes each)
- // instead of one slice per group, avoiding per-group ByteBuffer
allocation.
- int totalBytes = numGroups * bitWidth;
- ByteBuffer packed = in.slice(totalBytes);
- int pos = packed.position();
- int valueIndex = 0;
- while (valueIndex < this.currentCount) {
- this.packer.unpack8Values(packed, pos, this.currentBuffer,
valueIndex);
- pos += bitWidth;
- valueIndex += 8;
+ case PACKED -> {
+ int numGroups = header >>> 1;
+ // Validate before allocating the buffer, so a corrupted header
can neither overflow
+ // `currentCount` nor make us allocate more than the remaining
bytes can fill.
+ long totalBytes = (long) numGroups * bitWidth;
+ if (numGroups > Integer.MAX_VALUE / 8 || totalBytes >
in.available()) {
Review Comment:
This check also rejects a final bit-packed run whose trailing zero padding
is missing, even when every value the page needs is present. parquet-java's
`RunLengthBitPackingHybridDecoder` tolerates this (`bytesToRead =
Math.min(bytesToRead, in.available())`), and Arrow C++ accepts it since
apache/arrow#47992. Such files exist in the wild: Polars before
pola-rs/polars#13883 (see pola-rs/polars#13818) and Impala 1.1.1.
It's not a regression, because the previous `in.slice(totalBytes)` threw an
`EOFException` too. But the description says "The non-vectorized reader already
fails on these pages", which isn't true here, and the "Truncated last group"
test now pins the difference: `[4][0x05][0x11 x 5]` read as 10 values decodes
fine with `spark.sql.parquet.enableVectorizedReader=false`. Could we fix the
description, and maybe file a follow-up to accept a short last group the way
Arrow does (only values fully covered by the remaining bytes, only within the
last group)?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReaderSuite.scala:
##########
@@ -185,10 +203,210 @@ class VectorizedRleValuesReaderSuite extends
SparkFunSuite {
produced += toRead
}
}
+ test("SPARK-59832: truncated dictionary ids fail instead of being partially
read") {
+ // A dictionary id page with 10 ids, read as if it had 20.
+ streams(dictIdPage(Array.fill(10)(3), bitWidth = 4)).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ reader.initFromPage(20, in())
+ val c = new OnHeapColumnVector(20, IntegerType)
+ interceptCorrupted(reader.readIntegers(20, c, 0))
+ }
+ streams(dictIdPage(Array.fill(10)(3), bitWidth = 4)).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ reader.initFromPage(20, in())
+ interceptCorrupted(reader.skipIntegers(20))
+ }
+ streams(dictIdPage(Array.fill(10)(3), bitWidth = 4)).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ reader.initFromPage(20, in())
+ (0 until 10).foreach(_ => assert(reader.readInteger() == 3))
+ interceptCorrupted(reader.readInteger())
+ }
+ }
+
+ test("SPARK-59832: truncated booleans fail instead of looping forever") {
+ val encoded = encodeRle(Array.fill(10)(1), 1)
+ streams(intLE(encoded.length) ++ encoded).foreach { in =>
+ val reader = new VectorizedRleValuesReader(1)
+ reader.initFromPage(20, in())
+ val c = new OnHeapColumnVector(20, BooleanType)
+ interceptCorrupted(reader.readBooleans(20, c, 0))
+ }
+ }
+
+ test("SPARK-59832: truncated definition levels fail in readBatch") {
+ // The page declares 100 values but its definition levels only encode 64.
Before the fix
+ // readBatch returned without progress, so
VectorizedColumnReader.readBatch spun forever.
+ val n = 100
+ val encoded = encodeRle(Array.fill(64)(1), 1)
+ val plainBytes = plainIntBytes(n)(valueAt)
+ def run(withDefLevels: Boolean, rowIndexes: PrimitiveIterator.OfLong):
Unit = {
+ val reader = new VectorizedRleValuesReader(1, false)
+ reader.initFromPage(n,
ByteBufferInputStream.wrap(ByteBuffer.wrap(encoded)))
+ val valueReader = new VectorizedPlainValuesReader
+ valueReader.initFromPage(n,
ByteBufferInputStream.wrap(ByteBuffer.wrap(plainBytes)))
+ val state = ParquetTestAccess.newState(intColumnDescriptor(1), false,
rowIndexes)
+ ParquetTestAccess.resetForNewPage(state, n, 0L)
+ ParquetTestAccess.resetForNewBatch(state, n)
+ val values = new OnHeapColumnVector(n, IntegerType)
+ val defLevels = if (withDefLevels) new OnHeapColumnVector(n,
IntegerType) else null
+ interceptCorrupted {
+ ParquetTestAccess.readBatch(reader, state, values, defLevels,
valueReader, integerUpdater)
+ }
+ }
+ run(withDefLevels = false, rowIndexes = null)
+ run(withDefLevels = true, rowIndexes = null)
+ // Skips the truncated levels to reach rows 80 to 90.
+ run(withDefLevels = false, rowIndexes = longIterator((80 to 90).toArray))
+ }
+
+ test("SPARK-59832: truncated repetition levels fail in readBatchRepeated") {
+ val n = 100
+ val prim = Types.primitive(PrimitiveTypeName.INT32,
Repetition.REPEATED).named("col")
+ val descriptor = new ColumnDescriptor(Array("col"), prim, 1, 1)
+ val repReader = new VectorizedRleValuesReader(1, false)
+ repReader.initFromPage(
+ n,
ByteBufferInputStream.wrap(ByteBuffer.wrap(encodeRle(Array.fill(64)(0), 1))))
+ val defReader = new VectorizedRleValuesReader(1, false)
+ defReader.initFromPage(
+ n,
ByteBufferInputStream.wrap(ByteBuffer.wrap(encodeRle(Array.fill(n)(1), 1))))
+ val valueReader = new VectorizedPlainValuesReader
+ valueReader.initFromPage(
+ n,
ByteBufferInputStream.wrap(ByteBuffer.wrap(plainIntBytes(n)(valueAt))))
+ val state = ParquetTestAccess.newState(descriptor, false)
+ ParquetTestAccess.resetForNewPage(state, n, 0L)
+ ParquetTestAccess.resetForNewBatch(state, n)
+ val repLevels = new OnHeapColumnVector(n, IntegerType)
+ val defLevels = new OnHeapColumnVector(n, IntegerType)
+ val values = new OnHeapColumnVector(n, IntegerType)
+ interceptCorrupted {
+ repReader.readBatchRepeated(state.asInstanceOf[ParquetReadState],
repLevels, defReader,
+ defLevels, values, valueReader, integerUpdater)
+ }
+ }
+
+ test("SPARK-59832: reading past a page with bit width 0 fails") {
+ // With bit width 0 the page is one implicit run of `valueCount` zeros;
the bytes that
+ // follow belong to the next section of the page and must not be decoded
as levels.
+ val trailing = varint(((1L << 27) << 1) | 1) ++ Array[Byte](1, 2, 3)
+ val reader = new VectorizedRleValuesReader(0)
+ reader.initFromPage(5,
ByteBufferInputStream.wrap(ByteBuffer.wrap(trailing)))
+ val c = new OnHeapColumnVector(6, IntegerType)
+ reader.readIntegers(5, c, 0)
+ assert((0 until 5).forall(c.getInt(_) == 0))
+ interceptCorrupted(reader.readIntegers(1, c, 5))
+ }
+
+ test("SPARK-59832: invalid level length is rejected") {
+ val payload = Array[Byte](10, 20, 30, 40)
+ Seq(-1, -4, -5, payload.length + 1, Int.MaxValue).foreach { length =>
+ streams(intLE(length) ++ payload).foreach { in =>
+ val reader = new VectorizedRleValuesReader(1)
+ val e = intercept[ParquetDecodingException](reader.initFromPage(8,
in()))
+ assert(e.getMessage.contains(s"Corrupted RLE data: invalid length
$length"),
+ e.getMessage)
+ }
+ }
+ }
+
+ test("SPARK-59832: bit-packed run longer than the remaining bytes is
rejected") {
+ def check(bitWidth: Int, numGroups: Long, data: Array[Byte]): Unit = {
+ val page = Array(bitWidth.toByte) ++ varint((numGroups << 1) | 1) ++ data
+ streams(page).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ reader.initFromPage(10, in())
+ val c = new OnHeapColumnVector(10, IntegerType)
+ interceptCorrupted(reader.readIntegers(10, c, 0))
+ }
+ }
+ // Truncated last group.
+ check(bitWidth = 4, numGroups = 2, data = Array.fill[Byte](5)(0x11))
+ // numGroups * 8 and numGroups * bitWidth overflow to 0.
+ check(bitWidth = 8, numGroups = 1L << 29, data =
encodeRle(Array.fill(10)(3), 8))
+ // numGroups * 8 and numGroups * bitWidth overflow to negative values.
+ check(bitWidth = 2, numGroups = Int.MaxValue, data =
Array.fill[Byte](16)(0))
+ // Would allocate an int[2^30] buffer before noticing that only 3 bytes
are left.
+ check(bitWidth = 1, numGroups = 1L << 27, data = Array[Byte](1, 2, 3))
+ }
+
+ test("SPARK-59832: dictionary ids of an empty page are not decoded as 0") {
+ val reader = new VectorizedRleValuesReader()
+ reader.initFromPage(5,
ByteBufferInputStream.wrap(ByteBuffer.wrap(Array.emptyByteArray)))
+ // An all-null page reads no ids.
+ val c = new OnHeapColumnVector(5, IntegerType)
+ reader.readIntegers(0, c, 0)
+ reader.skipIntegers(0)
+ interceptCorrupted(reader.readIntegers(5, c, 0))
+ }
+
+ test("SPARK-59832: runs of length 0 are skipped") {
+ // [bit width 4][RLE run of 0 x 7][bit-packed run of 0 groups][RLE run of
5 x 3]
+ val page = Array[Byte](4, 0, 7, 1, 10, 3)
+ val r1 = new VectorizedRleValuesReader()
+ r1.initFromPage(5, ByteBufferInputStream.wrap(ByteBuffer.wrap(page)))
+ assert((0 until 5).map(_ => r1.readInteger()) == Seq.fill(5)(3))
+ val r2 = new VectorizedRleValuesReader()
+ r2.initFromPage(5, ByteBufferInputStream.wrap(ByteBuffer.wrap(page)))
+ val c = new OnHeapColumnVector(5, IntegerType)
+ r2.readIntegers(5, c, 0)
+ assert((0 until 5).map(c.getInt) == Seq.fill(5)(3))
+ }
+
+ test("SPARK-59832: reads that end exactly at the end of the encoded values
succeed") {
+ // A run of 20 identical ids (RLE) followed by 17 mixed ids (bit-packed,
padded to 24).
+ val ids = Array.tabulate(37)(i => if (i < 20) 5 else i % 7)
+ streams(dictIdPage(ids, bitWidth = 3)).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ reader.initFromPage(ids.length, in())
+ val c = new OnHeapColumnVector(ids.length, IntegerType)
+ reader.readIntegers(ids.length, c, 0)
+ assert((0 until ids.length).map(c.getInt) == ids.toSeq)
+ }
+ streams(dictIdPage(ids, bitWidth = 3)).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ reader.initFromPage(ids.length, in())
+ reader.skipIntegers(30)
+ assert((30 until ids.length).map(_ => reader.readInteger()) ==
ids.drop(30).toSeq)
+ }
+
+ val booleans = Array.tabulate(37)(i => if (i < 20) 1 else i % 2)
+ val encoded = encodeRle(booleans, 1)
+ // The level length covers exactly the rest of the page.
+ streams(intLE(encoded.length) ++ encoded).foreach { in =>
+ val reader = new VectorizedRleValuesReader(1)
+ reader.initFromPage(booleans.length, in())
+ val c = new OnHeapColumnVector(booleans.length, BooleanType)
+ reader.readBooleans(booleans.length, c, 0)
+ assert((0 until booleans.length).map(c.getBoolean) == booleans.map(_ ==
1).toSeq)
+ }
+ }
}
private object VectorizedRleValuesReaderSuite {
+ /** The page as a single buffer and split into 3-byte buffers. */
+ private def streams(bytes: Array[Byte]): Seq[() => ByteBufferInputStream] =
Seq(
Review Comment:
When a page is 3 bytes or less, `bytes.grouped(3)` yields one chunk, and
`ByteBufferInputStream.wrap(List)` then returns a `SingleBufferInputStream`. So
the "multiple buffers" variant is single-buffer for the truncated dictionary id
page (`[4, 20, 3]`). In the boolean test, the RLE bytes also end up in a single
buffer after `sliceStream(2)`. As a result, no test hits the end-of-data check
on a `MultiBufferInputStream`. Could we reuse
`VectorizedPlainValuesReaderSuite.toStream(bytes, split = true)`, which asserts
that the page spans more than one buffer, or split into 1-byte buffers?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReader.java:
##########
@@ -976,48 +985,58 @@ private int readIntLittleEndianPaddedOnBitWidth() throws
IOException {
}
/**
- * Reads the next group. Returns false if no more group available.
+ * Reads the next non-empty group, so that `currentCount` is positive on
return.
+ * <p>
+ * Callers never read more values than the page declares, so running out of
encoded data here
+ * means the page is corrupted. With a bit width of 0, `initFromPage`
already provides all the
+ * values of the page, so reaching here also means reading past the end.
*/
- private boolean readNextGroup() {
- if (in.available() <= 0) {
- currentCount = 0;
- return false;
- }
-
+ private void readNextGroup() {
try {
- int header = readUnsignedVarInt();
- this.mode = (header & 1) == 0 ? MODE.RLE : MODE.PACKED;
- switch (mode) {
- case RLE -> {
- this.currentCount = header >>> 1;
- this.currentValue = readIntLittleEndianPaddedOnBitWidth();
+ do {
Review Comment:
For backports: wrapping the whole body in `do { ... } while (currentCount ==
0)` re-indents about 20 unchanged lines. branch-4.2/4.1/4.0 still slice per
group (they don't have SPARK-56895), so the cherry-pick conflicts in the PACKED
loop. If the current body becomes a `readGroup()` helper called from `do {
readGroup(); } while (currentCount == 0);`, and the new checks are added
without touching the existing `int totalBytes = numGroups * bitWidth;
in.slice(totalBytes)` lines (use the long multiplication only inside the
check), it applies cleanly to every maintenance branch in my simulation. Either
way, the backports need to keep the `numGroups > Integer.MAX_VALUE / 8 ||
totalBytes > in.available()` pre-check before `currentBuffer` is allocated.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
