viirya commented on code in PR #59542:
URL: https://github.com/apache/spark/pull/59542#discussion_r4304156546
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorizedSuite.scala:
##########
@@ -551,6 +551,96 @@ class ParquetVectorizedSuite extends ParquetTest with
SharedSparkSession {
}
}
+ /**
+ * Reads a row group of one page of `valueCount` values of the column `a` in
`schema`, after
+ * an optional dictionary page, and returns the error of reading it.
+ */
+ private def readCorruptedPage(
+ schema: String,
+ dictionaryPage: Option[DictionaryPage],
+ data: Array[Byte],
+ valueCount: Int,
+ valuesEncoding: Encoding): Throwable = {
+ val parquetSchema = MessageTypeParser.parseMessageType(schema)
+ val ty = parquetSchema.asGroupType().getType("a").asPrimitiveType()
+ val cd = new ColumnDescriptor(Array("a"), ty, 0, 1)
+ val memPageStore = new MemPageStore(valueCount)
+ val pageWriter = memPageStore.getPageWriter(cd)
+ dictionaryPage.foreach(pageWriter.writeDictionaryPage)
+ pageWriter.writePage(BytesInput.from(data), valueCount, valueCount,
+ Statistics.createStats(ty), Encoding.RLE, Encoding.RLE, valuesEncoding)
+ val recordReader = new VectorizedParquetRecordReader(
+ DateTimeUtils.getZoneId("EST"), "CORRECTED", "UTC", "CORRECTED", "UTC",
true, valueCount)
+ try {
+ recordReader.initialize(parquetSchema, parquetSchema,
+ TestParquetRowGroupReader(Seq(TestPageReadStore(memPageStore,
Seq(0L)))), valueCount)
+ intercept[ParquetDecodingException](recordReader.nextKeyValue())
+ } finally {
+ recordReader.close()
+ }
+ }
+
+ // The definition levels of 4 non-null values: the length 2, then an RLE run
of 4 x 1.
+ private val fourNonNullValues = Array[Byte](2, 0, 0, 0, 4 << 1, 1)
+
+ test("SPARK-59854: a decoding error in a batch reports the page and the
column") {
+ Seq(
+ // Dictionary ids (bit width 1) with an RLE run of only 2 of the 4
values.
+ ("Corrupted RLE data: reading past the end of the encoded values",
Some(new DictionaryPage(
+ BytesInput.from(Array[Byte](1, 0, 0, 0, 2, 0, 0, 0)), 2,
Encoding.PLAIN_DICTIONARY)),
+ Array[Byte](1, 2 << 1, 0), Encoding.RLE_DICTIONARY),
+ // 2 of the 4 PLAIN values.
+ ("Failed to read 16 bytes", None, Array[Byte](1, 0, 0, 0, 2, 0, 0, 0),
Encoding.PLAIN)
+ ).foreach { case (error, dictionaryPage, values, encoding) =>
+ val e = readCorruptedPage("message root { optional int32 a; }",
dictionaryPage,
+ fourNonNullValues ++ values, 4, encoding)
+ assert(e.getMessage.startsWith("could not decode page "), e.getMessage)
Review Comment:
Done. `TestDataPage.toString` now delegates to the wrapped page, and the
test checks `could not decode page Page [bytes.size=`, so it would fail if the
page were missing.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedDeltaBinaryPackedReader.java:
##########
@@ -104,10 +105,13 @@ public void initFromPage(int valueCount,
ByteBufferInputStream in) throws IOExce
throw new ParquetDecodingException("Invalid DELTA_BINARY_PACKED mini
block count: "
+ miniBlockNumInABlock + " (block size in values: " +
blockSizeInValues + ")");
}
- double miniSize = (double) blockSizeInValues / miniBlockNumInABlock;
- JavaUtils.checkArgument(miniSize % 8 == 0,
- "miniBlockSize must be multiple of 8, but it's " + miniSize);
- this.miniBlockSizeInValues = (int) miniSize;
+ if (blockSizeInValues % miniBlockNumInABlock != 0 ||
+ (blockSizeInValues / miniBlockNumInABlock) % 8 != 0) {
Review Comment:
Done. `ParquetDeltaEncodingSuite` now checks 128 values in 3 and in 32 mini
blocks and a page with no values, and I dropped the direct `initFromPage(0,
...)` check and its import from `ParquetVectorizedSuite`.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedColumnReader.java:
##########
@@ -147,7 +152,14 @@ public VectorizedColumnReader(
int96RebaseMode,
int96RebaseTz);
- DictionaryPage dictionaryPage = pageReader.readDictionaryPage();
+ DictionaryPage dictionaryPage;
+ try {
+ // This decompresses the page eagerly for some codecs, e.g. ZSTD.
+ dictionaryPage = pageReader.readDictionaryPage();
+ } catch (ParquetDecodingException e) {
Review Comment:
Done for the first part: `readDictionaryPage()` is now in the same `try` as
`initDictionary`, so a dictionary page that cannot be read fails with `could
not decode the dictionary for ...` for every codec.
For the helper method, I tried it and measured `git blame` on the decode
block. Plain `git blame` attributes 64 of its lines to this PR in both
versions, `-w` 8 in both, and `-w -M` 7 with the `try` and 6 with the method.
Moving the block into a method body re-indents it as well (to 4 spaces, not
back to the original 6), so it does not keep the history. So I kept the `try`.
Agreed that the wrapper must keep catching only `ParquetDecodingException`, so
that a vector growth failure, which `rethrowIfMustPropagate` checks by its
exact class, is not wrapped.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorizedSuite.scala:
##########
@@ -551,6 +551,96 @@ class ParquetVectorizedSuite extends ParquetTest with
SharedSparkSession {
}
}
+ /**
+ * Reads a row group of one page of `valueCount` values of the column `a` in
`schema`, after
+ * an optional dictionary page, and returns the error of reading it.
+ */
+ private def readCorruptedPage(
Review Comment:
Done. The SPARK-59832 test now uses `readCorruptedPage`, which takes the
column type instead of a schema, so it no longer looks free-form. The helper
sits with the other helpers after `writeDataPage`, and the definition level
vals are at the top of the suite.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedDeltaBinaryPackedReader.java:
##########
@@ -104,10 +105,13 @@ public void initFromPage(int valueCount,
ByteBufferInputStream in) throws IOExce
throw new ParquetDecodingException("Invalid DELTA_BINARY_PACKED mini
block count: "
+ miniBlockNumInABlock + " (block size in values: " +
blockSizeInValues + ")");
}
- double miniSize = (double) blockSizeInValues / miniBlockNumInABlock;
- JavaUtils.checkArgument(miniSize % 8 == 0,
- "miniBlockSize must be multiple of 8, but it's " + miniSize);
- this.miniBlockSizeInValues = (int) miniSize;
+ if (blockSizeInValues % miniBlockNumInABlock != 0 ||
Review Comment:
Done, the `miniBlockNumInABlock > blockSizeInValues` arm is gone and `<= 0`
stays. The existing `craftHeader(128, 256, 1)` case still gets the "mini block
count" message.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala:
##########
@@ -2569,6 +2583,30 @@ class ParquetIOSuite extends ParquetTest with
SharedSparkSession {
val e = readCorruptedRleSection("id >= 0 as b", writerVersion =
"PARQUET_2_0")(setMaxLength)
assertInvalidLengthReportsColumn(e, "b")
}
+
+ test("SPARK-59854: a ZSTD page that cannot be decompressed reports the
column") {
+ Seq(false, true).foreach { dictionary =>
+ val options = Map(
+ ParquetOutputFormat.ENABLE_DICTIONARY -> dictionary.toString,
+ "compression" -> "zstd")
+ val e = readCorruptedFile("cast(id % 10 as int) as a", options) { (file,
bytes) =>
+ val footer = ParquetFooterReader.readFooter(
Review Comment:
Done. The test parses the page header at `ParquetFileWriter.MAGIC.length`
and asserts its type, `DATA_PAGE` or `DICTIONARY_PAGE`, and `readCorruptedFile`
no longer passes the path.
--
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]