viirya commented on code in PR #59117:
URL: https://github.com/apache/spark/pull/59117#discussion_r4232802819
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala:
##########
@@ -2480,6 +2484,93 @@ class ParquetIOSuite extends ParquetTest with
SharedSparkSession {
}
}
}
+
+ /**
+ * Writes 100 rows of `column` to a single uncompressed data page, lets
`corrupt` modify the
+ * RLE section that starts with `03 00 00 00 C8 01 01` (the 4-byte length 3,
then one RLE run
+ * with header 100 << 1 and value 1), and returns the error of reading the
file with the
+ * vectorized reader. For a nullable INT32 column that is the definition
levels; for a
+ * BOOLEAN column written as a V2 page it is the RLE-encoded values.
+ */
+ private def readCorruptedRleSection(
+ column: String,
+ writerVersion: String = "PARQUET_1_0")(
+ corrupt: (Array[Byte], Int) => Unit): Throwable = {
+ var error: Throwable = null
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ spark.range(100).selectExpr(column).coalesce(1)
+ .write
+ .option(ParquetOutputFormat.ENABLE_DICTIONARY, "false")
+ .option(ParquetOutputFormat.WRITER_VERSION, writerVersion)
+ .option("compression", "uncompressed")
+ .parquet(path)
+ val file = dir.listFiles().filter(_.getName.endsWith(".parquet")).head
+ val bytes = Files.readAllBytes(file.toPath)
+ val section = Array[Byte](3, 0, 0, 0, 0xC8.toByte, 0x01, 1)
+ val offsets = bytes.indices.filter { i =>
+ bytes.slice(i, i + section.length).sameElements(section)
+ }
+ assert(offsets.size == 1)
+ corrupt(bytes, offsets.head)
+ Files.write(file.toPath, bytes)
+ dir.listFiles().filter(_.getName.endsWith(".crc")).foreach(_.delete())
+
+ withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "true") {
+ // Before SPARK-59832 a truncated read never returned, so run it with
a timeout.
+ // scalastyle:off awaitresult
+ error = intercept[SparkException] {
+ Await.result(
+
Future(spark.read.parquet(path).collect())(ExecutionContext.global), 1.minute)
+ }
+ // scalastyle:on awaitresult
+ }
+ }
+ error
+ }
+
+ private def causes(e: Throwable): Seq[Throwable] =
+ Iterator.iterate(e)(_.getCause).takeWhile(_ != null).toSeq
+
+ /** Sets the 4-byte length at `offset` to Int.MaxValue, far more than the
rest of the page. */
+ private def setMaxLength(bytes: Array[Byte], offset: Int): Unit =
+ Array(0xFF, 0xFF, 0xFF, 0x7F).zipWithIndex.foreach { case (b, i) =>
+ bytes(offset + i) = b.toByte
+ }
+
+ /**
+ * Checks that the page error is reported once with the column, directly on
top of the
+ * `ParquetDecodingException` about the invalid length.
+ */
+ private def assertInvalidLengthReportsColumn(e: Throwable, column: String):
Unit = {
+ val pageErrors = causes(e).filter(_.getMessage.startsWith("could not read
page"))
Review Comment:
Fixed with `Option(c.getMessage)`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala:
##########
@@ -2480,6 +2484,93 @@ class ParquetIOSuite extends ParquetTest with
SharedSparkSession {
}
}
}
+
+ /**
+ * Writes 100 rows of `column` to a single uncompressed data page, lets
`corrupt` modify the
+ * RLE section that starts with `03 00 00 00 C8 01 01` (the 4-byte length 3,
then one RLE run
+ * with header 100 << 1 and value 1), and returns the error of reading the
file with the
+ * vectorized reader. For a nullable INT32 column that is the definition
levels; for a
+ * BOOLEAN column written as a V2 page it is the RLE-encoded values.
+ */
+ private def readCorruptedRleSection(
+ column: String,
+ writerVersion: String = "PARQUET_1_0")(
+ corrupt: (Array[Byte], Int) => Unit): Throwable = {
+ var error: Throwable = null
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ spark.range(100).selectExpr(column).coalesce(1)
+ .write
+ .option(ParquetOutputFormat.ENABLE_DICTIONARY, "false")
+ .option(ParquetOutputFormat.WRITER_VERSION, writerVersion)
+ .option("compression", "uncompressed")
+ .parquet(path)
+ val file = dir.listFiles().filter(_.getName.endsWith(".parquet")).head
+ val bytes = Files.readAllBytes(file.toPath)
+ val section = Array[Byte](3, 0, 0, 0, 0xC8.toByte, 0x01, 1)
+ val offsets = bytes.indices.filter { i =>
+ bytes.slice(i, i + section.length).sameElements(section)
+ }
+ assert(offsets.size == 1)
+ corrupt(bytes, offsets.head)
+ Files.write(file.toPath, bytes)
+ dir.listFiles().filter(_.getName.endsWith(".crc")).foreach(_.delete())
+
+ withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "true") {
+ // Before SPARK-59832 a truncated read never returned, so run it with
a timeout.
+ // scalastyle:off awaitresult
Review Comment:
Done, `failAfter` with a `ThreadSignaler`, and the `scala.concurrent`
imports are gone.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReaderSuite.scala:
##########
@@ -19,25 +19,66 @@ package org.apache.spark.sql.execution.datasources.parquet
import java.nio.ByteBuffer
import java.util.PrimitiveIterator
+import java.util.concurrent.{Callable, ExecutionException, TimeUnit}
import scala.jdk.CollectionConverters._
-import org.apache.parquet.bytes.ByteBufferInputStream
+import org.apache.parquet.bytes.{ByteBufferInputStream, BytesInput, BytesUtils}
+import org.apache.parquet.column.ColumnDescriptor
+import org.apache.parquet.io.ParquetDecodingException
+import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName
+import org.apache.parquet.schema.Type.Repetition
+import org.apache.parquet.schema.Types
import org.apache.spark.SparkFunSuite
import
org.apache.spark.sql.execution.datasources.parquet.VectorizedRleValuesReaderTestUtils._
import org.apache.spark.sql.execution.vectorized.{OnHeapColumnVector,
WritableColumnVector}
-import org.apache.spark.sql.types.IntegerType
+import org.apache.spark.sql.types.{BooleanType, IntegerType}
+import org.apache.spark.util.ThreadUtils
/**
* Focused correctness tests for `VectorizedRleValuesReader.readBatch`
PACKED-mode decoding,
* covering patterns the P0 optimization cares about: run boundaries, batch
boundaries, and
* nested def-level grouping. Uses the same reflection bridge as the benchmark.
+ *
+ * The SPARK-59832 tests check that corrupted or truncated input fails with a
+ * `ParquetDecodingException` instead of returning wrong values, looping
forever or allocating
+ * a huge buffer, and that reads ending exactly at the end of the encoded
values still work.
*/
class VectorizedRleValuesReaderSuite extends SparkFunSuite {
import VectorizedRleValuesReaderSuite._
+ // Runs the reads that may never return before SPARK-59832. A thread per
read, so that a read
+ // stuck in a loop does not block the following ones.
+ private lazy val executor =
ThreadUtils.newDaemonCachedThreadPool("rle-reader-suite")
+
+ override def afterAll(): Unit = {
+ try {
+ executor.shutdownNow()
+ } finally {
+ super.afterAll()
+ }
+ }
+
+ /**
+ * Runs `f` in another thread, so that a reader that never returns fails the
test, and checks
Review Comment:
Added `assertCorrupted(expected)(f)`, which `interceptCorrupted` uses around
`future.get`. The three `initFromPage` tests use it, so the prefix is only
there and in its doc.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReaderSuite.scala:
##########
@@ -185,10 +226,266 @@ 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(PastEnd)(reader.readIntegers(20, c, 0))
+ }
+ streams(dictIdPage(Array.fill(10)(3), bitWidth = 4)).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ reader.initFromPage(20, in())
+ interceptCorrupted(PastEnd)(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(PastEnd)(reader.readInteger())
+ }
+ }
+
+ test("SPARK-59832: truncated booleans fail instead of looping forever") {
+ val encoded = encodeRle(Array.fill(10)(1), 1)
+ streams(BytesUtils.intToBytes(encoded.length) ++ encoded).foreach { in =>
+ val reader = new VectorizedRleValuesReader(1)
+ reader.initFromPage(20, in())
+ val c = new OnHeapColumnVector(20, BooleanType)
+ interceptCorrupted(PastEnd)(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(PastEnd) {
+ 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") {
+ // Every row is a top-level row (repetition level 0), but only 64 of 100
are encoded.
+ readRepeated(repLevels = Array.fill(64)(0), defLevels =
Array.fill(100)(1), n = 100)
+ }
+
+ test("SPARK-59832: truncated definition levels fail in readBatchRepeated") {
+ // The repetition levels are complete, but only 64 of 100 definition
levels are encoded.
+ // Before the fix the missing definition levels were silently left as 0.
+ readRepeated(repLevels = Array.fill(100)(0), defLevels =
Array.fill(64)(1), n = 100)
+ // Skips the truncated levels to reach rows 80 to 90.
+ readRepeated(repLevels = Array.fill(100)(0), defLevels =
Array.fill(64)(1), n = 100,
+ rowIndexes = longIterator((80 to 90).toArray))
+ }
+
+ private def readRepeated(
+ repLevels: Array[Int],
+ defLevels: Array[Int],
+ n: Int,
+ rowIndexes: PrimitiveIterator.OfLong = null): Unit = {
+ 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(repLevels, 1))))
+ val defReader = new VectorizedRleValuesReader(1, false)
+ defReader.initFromPage(
+ n, ByteBufferInputStream.wrap(ByteBuffer.wrap(encodeRle(defLevels, 1))))
+ val valueReader = new VectorizedPlainValuesReader
+ valueReader.initFromPage(
+ n,
ByteBufferInputStream.wrap(ByteBuffer.wrap(plainIntBytes(n)(valueAt))))
+ val state = ParquetTestAccess.newState(descriptor, false, rowIndexes)
+ ParquetTestAccess.resetForNewPage(state, n, 0L)
+ ParquetTestAccess.resetForNewBatch(state, n)
+ val repLevelsVec = new OnHeapColumnVector(n, IntegerType)
+ val defLevelsVec = new OnHeapColumnVector(n, IntegerType)
+ val values = new OnHeapColumnVector(n, IntegerType)
+ interceptCorrupted(PastEnd) {
+ ParquetTestAccess.readBatchRepeated(repReader, state, repLevelsVec,
defReader,
+ defLevelsVec, 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 = Array[Byte](3, 1, 2, 3)
+ val reader = new VectorizedRleValuesReader(0)
+ reader.initFromPage(5,
ByteBufferInputStream.wrap(ByteBuffer.wrap(trailing)))
+ val c = new OnHeapColumnVector(6, IntegerType)
+ c.putInts(0, 6, -1)
+ reader.readIntegers(5, c, 0)
+ assert((0 until 5).forall(c.getInt(_) == 0))
+ interceptCorrupted(PastEnd)(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(BytesUtils.intToBytes(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: invalid bit-packed run is rejected before allocating its
buffer") {
+ def check(bitWidth: Int, numGroups: Long, data: Array[Byte])(expected:
String): 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(expected)(reader.readIntegers(10, c, 0))
+ }
+ }
+ // Truncated last group. This was already rejected before (by `in.slice`),
even when the
+ // bytes cover every value read. parquet-java and Arrow accept it, for
compatibility with
+ // writers that do not pad the last group.
+ check(bitWidth = 4, numGroups = 2, data = Array.fill[Byte](5)(0x11))(
+ "bit-packed run of 2 groups needs 8 bytes, but only 5 bytes are left")
+ // numGroups * 8 and numGroups * bitWidth overflow to 0.
+ check(bitWidth = 8, numGroups = 1L << 29, data =
encodeRle(Array.fill(10)(3), 8))(
+ s"bit-packed run of ${1 << 29} groups is too long")
+ // numGroups * 8 and numGroups * bitWidth overflow to negative values.
+ check(bitWidth = 2, numGroups = Int.MaxValue, data =
Array.fill[Byte](16)(0))(
+ s"bit-packed run of ${Int.MaxValue} groups is too long")
+ // 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))(
+ s"bit-packed run of ${1 << 27} groups needs ${1 << 27} bytes, but only 3
bytes are left")
+ }
+
+ test("SPARK-59832: data ending inside a run header or an RLE value fails") {
+ // The continuation byte of the run header is missing, and an RLE run of
10 without its
+ // value.
+ Seq(Array[Byte](4, 0x80.toByte), Array[Byte](4, 0x14)).foreach { page =>
+ streams(page).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ reader.initFromPage(10, in())
+ val c = new OnHeapColumnVector(10, IntegerType)
+ interceptCorrupted("failed to read from input
stream")(reader.readIntegers(10, c, 0))
+ }
+ }
+ }
+
+ test("SPARK-59832: cut-off level length is rejected") {
+ streams(Array[Byte](3, 0)).foreach { in =>
+ val reader = new VectorizedRleValuesReader(1)
+ val e = intercept[ParquetDecodingException](reader.initFromPage(8, in()))
+ assert(e.getMessage.contains(
+ "Corrupted RLE data: the 4-byte length is cut off, only 2 bytes are
left in the page"),
+ e.getMessage)
+ }
+ }
+
+ test("SPARK-59832: invalid dictionary id bit width is rejected") {
+ Seq(33, 255).foreach { bitWidth =>
+ streams(Array[Byte](bitWidth.toByte, 0x14, 0)).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ val e = intercept[ParquetDecodingException](reader.initFromPage(10,
in()))
+ assert(e.getMessage.contains(s"Corrupted RLE data: invalid bit width
$bitWidth"),
+ e.getMessage)
+ }
+ }
+ }
+
+ 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(PastEnd)(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(BytesUtils.intToBytes(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 1-byte buffers. The split
keeps any part of the
+ * page longer than 1 byte, including a slice of it, on a
`MultiBufferInputStream`.
+ */
+ private def streams(bytes: Array[Byte]): Seq[() => ByteBufferInputStream] = {
Review Comment:
Fixed in the description: 8 of 14, and only the tests that read values run
on a separate thread.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReaderSuite.scala:
##########
@@ -185,10 +226,266 @@ 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 =>
Review Comment:
Added `assertReadIntegersFails(page, n, expected)`. The three blocks and
`check` use it, and the test builds the page once.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala:
##########
@@ -2480,6 +2484,93 @@ class ParquetIOSuite extends ParquetTest with
SharedSparkSession {
}
}
}
+
+ /**
+ * Writes 100 rows of `column` to a single uncompressed data page, lets
`corrupt` modify the
+ * RLE section that starts with `03 00 00 00 C8 01 01` (the 4-byte length 3,
then one RLE run
+ * with header 100 << 1 and value 1), and returns the error of reading the
file with the
+ * vectorized reader. For a nullable INT32 column that is the definition
levels; for a
+ * BOOLEAN column written as a V2 page it is the RLE-encoded values.
+ */
+ private def readCorruptedRleSection(
+ column: String,
+ writerVersion: String = "PARQUET_1_0")(
+ corrupt: (Array[Byte], Int) => Unit): Throwable = {
+ var error: Throwable = null
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ spark.range(100).selectExpr(column).coalesce(1)
+ .write
+ .option(ParquetOutputFormat.ENABLE_DICTIONARY, "false")
+ .option(ParquetOutputFormat.WRITER_VERSION, writerVersion)
+ .option("compression", "uncompressed")
+ .parquet(path)
+ val file = dir.listFiles().filter(_.getName.endsWith(".parquet")).head
+ val bytes = Files.readAllBytes(file.toPath)
+ val section = Array[Byte](3, 0, 0, 0, 0xC8.toByte, 0x01, 1)
+ val offsets = bytes.indices.filter { i =>
+ bytes.slice(i, i + section.length).sameElements(section)
+ }
+ assert(offsets.size == 1)
+ corrupt(bytes, offsets.head)
+ Files.write(file.toPath, bytes)
+ dir.listFiles().filter(_.getName.endsWith(".crc")).foreach(_.delete())
+
+ withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "true") {
+ // Before SPARK-59832 a truncated read never returned, so run it with
a timeout.
+ // scalastyle:off awaitresult
+ error = intercept[SparkException] {
+ Await.result(
+
Future(spark.read.parquet(path).collect())(ExecutionContext.global), 1.minute)
+ }
+ // scalastyle:on awaitresult
+ }
+ }
+ error
+ }
+
+ private def causes(e: Throwable): Seq[Throwable] =
+ Iterator.iterate(e)(_.getCause).takeWhile(_ != null).toSeq
+
+ /** Sets the 4-byte length at `offset` to Int.MaxValue, far more than the
rest of the page. */
+ private def setMaxLength(bytes: Array[Byte], offset: Int): Unit =
Review Comment:
Done.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReaderSuite.scala:
##########
@@ -19,25 +19,66 @@ package org.apache.spark.sql.execution.datasources.parquet
import java.nio.ByteBuffer
import java.util.PrimitiveIterator
+import java.util.concurrent.{Callable, ExecutionException, TimeUnit}
import scala.jdk.CollectionConverters._
-import org.apache.parquet.bytes.ByteBufferInputStream
+import org.apache.parquet.bytes.{ByteBufferInputStream, BytesInput, BytesUtils}
+import org.apache.parquet.column.ColumnDescriptor
+import org.apache.parquet.io.ParquetDecodingException
+import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName
+import org.apache.parquet.schema.Type.Repetition
+import org.apache.parquet.schema.Types
import org.apache.spark.SparkFunSuite
import
org.apache.spark.sql.execution.datasources.parquet.VectorizedRleValuesReaderTestUtils._
import org.apache.spark.sql.execution.vectorized.{OnHeapColumnVector,
WritableColumnVector}
-import org.apache.spark.sql.types.IntegerType
+import org.apache.spark.sql.types.{BooleanType, IntegerType}
+import org.apache.spark.util.ThreadUtils
/**
* Focused correctness tests for `VectorizedRleValuesReader.readBatch`
PACKED-mode decoding,
* covering patterns the P0 optimization cares about: run boundaries, batch
boundaries, and
* nested def-level grouping. Uses the same reflection bridge as the benchmark.
+ *
+ * The SPARK-59832 tests check that corrupted or truncated input fails with a
+ * `ParquetDecodingException` instead of returning wrong values, looping
forever or allocating
+ * a huge buffer, and that reads ending exactly at the end of the encoded
values still work.
*/
class VectorizedRleValuesReaderSuite extends SparkFunSuite {
import VectorizedRleValuesReaderSuite._
+ // Runs the reads that may never return before SPARK-59832. A thread per
read, so that a read
+ // stuck in a loop does not block the following ones.
+ private lazy val executor =
ThreadUtils.newDaemonCachedThreadPool("rle-reader-suite")
+
+ override def afterAll(): Unit = {
+ try {
+ executor.shutdownNow()
Review Comment:
Added `awaitTermination(10, SECONDS)` after `shutdownNow()`.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReader.java:
##########
@@ -99,13 +99,30 @@ public void initFromPage(int valueCount,
ByteBufferInputStream in) throws IOExce
if (fixedWidth) {
// Initialize for repetition and definition levels
if (readLength) {
+ if (in.available() < 4) {
+ throw new ParquetDecodingException("Corrupted RLE data: the 4-byte
length is cut off, " +
+ "only " + in.available() + " bytes are left in the page");
+ }
int length = readIntLittleEndian();
+ if (length < 0 || length > in.available()) {
+ throw new ParquetDecodingException("Corrupted RLE data: invalid
length " + length +
+ ", " + in.available() + " bytes are left in the page");
+ }
this.in = in.sliceStream(length);
}
} else {
// Initialize for values
if (in.available() > 0) {
- init(in.read());
+ int width = in.read();
+ if (width > 32) {
+ throw new ParquetDecodingException("Corrupted RLE data: invalid bit
width " + width);
+ }
+ init(width);
+ } else {
+ // No encoded values (e.g. an all-null page). Do not treat this as a
bit width of 0,
Review Comment:
Done in 7cdea62. `initFromPage` now remembers that the section is empty, and
the first read fails with "Corrupted RLE data: the page has no data, not even
the bit width".
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReaderSuite.scala:
##########
@@ -196,10 +237,266 @@ 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(PastEnd)(reader.readIntegers(20, c, 0))
+ }
+ streams(dictIdPage(Array.fill(10)(3), bitWidth = 4)).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ reader.initFromPage(20, in())
+ interceptCorrupted(PastEnd)(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(PastEnd)(reader.readInteger())
+ }
+ }
+
+ test("SPARK-59832: truncated booleans fail instead of looping forever") {
+ val encoded = encodeRle(Array.fill(10)(1), 1)
+ streams(BytesUtils.intToBytes(encoded.length) ++ encoded).foreach { in =>
+ val reader = new VectorizedRleValuesReader(1)
+ reader.initFromPage(20, in())
+ val c = new OnHeapColumnVector(20, BooleanType)
+ interceptCorrupted(PastEnd)(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(PastEnd) {
+ 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") {
+ // Every row is a top-level row (repetition level 0), but only 64 of 100
are encoded.
+ readRepeated(repLevels = Array.fill(64)(0), defLevels =
Array.fill(100)(1), n = 100)
+ }
+
+ test("SPARK-59832: truncated definition levels fail in readBatchRepeated") {
+ // The repetition levels are complete, but only 64 of 100 definition
levels are encoded.
+ // Before the fix the missing definition levels were silently left as 0.
+ readRepeated(repLevels = Array.fill(100)(0), defLevels =
Array.fill(64)(1), n = 100)
+ // Skips the truncated levels to reach rows 80 to 90.
+ readRepeated(repLevels = Array.fill(100)(0), defLevels =
Array.fill(64)(1), n = 100,
+ rowIndexes = longIterator((80 to 90).toArray))
+ }
+
+ private def readRepeated(
+ repLevels: Array[Int],
+ defLevels: Array[Int],
+ n: Int,
+ rowIndexes: PrimitiveIterator.OfLong = null): Unit = {
+ 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(repLevels, 1))))
+ val defReader = new VectorizedRleValuesReader(1, false)
+ defReader.initFromPage(
+ n, ByteBufferInputStream.wrap(ByteBuffer.wrap(encodeRle(defLevels, 1))))
+ val valueReader = new VectorizedPlainValuesReader
+ valueReader.initFromPage(
+ n,
ByteBufferInputStream.wrap(ByteBuffer.wrap(plainIntBytes(n)(valueAt))))
+ val state = ParquetTestAccess.newState(descriptor, false, rowIndexes)
+ ParquetTestAccess.resetForNewPage(state, n, 0L)
+ ParquetTestAccess.resetForNewBatch(state, n)
+ val repLevelsVec = new OnHeapColumnVector(n, IntegerType)
+ val defLevelsVec = new OnHeapColumnVector(n, IntegerType)
+ val values = new OnHeapColumnVector(n, IntegerType)
+ interceptCorrupted(PastEnd) {
+ ParquetTestAccess.readBatchRepeated(repReader, state, repLevelsVec,
defReader,
+ defLevelsVec, 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 = Array[Byte](3, 1, 2, 3)
+ val reader = new VectorizedRleValuesReader(0)
+ reader.initFromPage(5,
ByteBufferInputStream.wrap(ByteBuffer.wrap(trailing)))
+ val c = new OnHeapColumnVector(6, IntegerType)
+ c.putInts(0, 6, -1)
+ reader.readIntegers(5, c, 0)
+ assert((0 until 5).forall(c.getInt(_) == 0))
+ interceptCorrupted(PastEnd)(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(BytesUtils.intToBytes(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: invalid bit-packed run is rejected before allocating its
buffer") {
+ def check(bitWidth: Int, numGroups: Long, data: Array[Byte])(expected:
String): 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(expected)(reader.readIntegers(10, c, 0))
+ }
+ }
+ // Truncated last group. This was already rejected before (by `in.slice`),
even when the
+ // bytes cover every value read. parquet-java and Arrow accept it, for
compatibility with
+ // writers that do not pad the last group.
+ check(bitWidth = 4, numGroups = 2, data = Array.fill[Byte](5)(0x11))(
+ "bit-packed run of 2 groups needs 8 bytes, but only 5 bytes are left")
+ // numGroups * 8 and numGroups * bitWidth overflow to 0.
+ check(bitWidth = 8, numGroups = 1L << 29, data =
encodeRle(Array.fill(10)(3), 8))(
+ s"bit-packed run of ${1 << 29} groups is too long")
+ // numGroups * 8 and numGroups * bitWidth overflow to negative values.
+ check(bitWidth = 2, numGroups = Int.MaxValue, data =
Array.fill[Byte](16)(0))(
+ s"bit-packed run of ${Int.MaxValue} groups is too long")
+ // 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))(
+ s"bit-packed run of ${1 << 27} groups needs ${1 << 27} bytes, but only 3
bytes are left")
+ }
+
+ test("SPARK-59832: data ending inside a run header or an RLE value fails") {
+ // The continuation byte of the run header is missing, and an RLE run of
10 without its
+ // value.
+ Seq(Array[Byte](4, 0x80.toByte), Array[Byte](4, 0x14)).foreach { page =>
+ streams(page).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ reader.initFromPage(10, in())
+ val c = new OnHeapColumnVector(10, IntegerType)
+ interceptCorrupted("failed to read from input
stream")(reader.readIntegers(10, c, 0))
+ }
+ }
+ }
+
+ test("SPARK-59832: cut-off level length is rejected") {
+ streams(Array[Byte](3, 0)).foreach { in =>
+ val reader = new VectorizedRleValuesReader(1)
+ val e = intercept[ParquetDecodingException](reader.initFromPage(8, in()))
+ assert(e.getMessage.contains(
+ "Corrupted RLE data: the 4-byte length is cut off, only 2 bytes are
left in the page"),
+ e.getMessage)
+ }
+ }
+
+ test("SPARK-59832: invalid dictionary id bit width is rejected") {
+ Seq(33, 255).foreach { bitWidth =>
+ streams(Array[Byte](bitWidth.toByte, 0x14, 0)).foreach { in =>
+ val reader = new VectorizedRleValuesReader()
+ val e = intercept[ParquetDecodingException](reader.initFromPage(10,
in()))
+ assert(e.getMessage.contains(s"Corrupted RLE data: invalid bit width
$bitWidth"),
+ e.getMessage)
+ }
+ }
+ }
+
+ 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(PastEnd)(reader.readIntegers(5, c, 0))
+ }
+
+ test("SPARK-59832: runs of length 0 are skipped") {
Review Comment:
Done, both reads in that test now go through `streams(page)`.
--
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]
