viirya commented on code in PR #59117:
URL: https://github.com/apache/spark/pull/59117#discussion_r4130646619
##########
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") {
Review Comment:
Fixed, thanks.
##########
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:
Good catch, the description was wrong there. I updated it: this PR keeps
rejecting a short last group as before (`in.slice` already threw), and notes
that parquet-java and Arrow C++ accept it. The test case now has a comment
saying the same. I filed SPARK-59853 to accept it the way Arrow does.
##########
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:
Added "truncated definition levels fail in readBatchRepeated", with and
without row indexes. They reach the throws in `readValues` and in
`skipValues(int, ParquetReadState, ...)` respectively. On master this case
returns without an error, so I also added it to the description.
##########
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:
Agreed. `readPageV1` and `initDataReader` now also catch
`ParquetDecodingException` and wrap it with the page and column like an
`IOException`, and a new end-to-end test checks that an oversized definition
level length reports the column. Decoding errors thrown later from `readBatch`
never had the column, including the ones from SPARK-59812, so I filed
SPARK-59854 for that.
##########
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:
Right, thanks. `streams` now splits into 1-byte buffers, so any part of the
page longer than 1 byte, including a slice of it, stays on a
`MultiBufferInputStream`. I didn't reuse
`VectorizedPlainValuesReaderSuite.toStream` because its 3-byte split has the
same problem after a slice.
##########
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:
Done. The existing body is now `readGroup()` and `readNextGroup()` loops
over it. The `int totalBytes = numGroups * bitWidth; in.slice(totalBytes)`
lines are unchanged, and the long multiplication is only in the new 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]
