diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedColumnReader.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedColumnReader.java index 54b979240ae79..b3a261eefd11d 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedColumnReader.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedColumnReader.java @@ -98,6 +98,12 @@ public class VectorizedColumnReader { */ private long pageFirstRowIndex; + /** + * The data page the values are being decoded from, for error messages. Only its description is + * kept: for codecs that decompress lazily, the page holds the decompressor stream. + */ + private String currentPage; + private final PageReader pageReader; private final ColumnDescriptor descriptor; private final LogicalTypeAnnotation logicalTypeAnnotation; @@ -146,23 +152,33 @@ public VectorizedColumnReader( int96RebaseMode, int96RebaseTz); - DictionaryPage dictionaryPage = pageReader.readDictionaryPage(); - if (dictionaryPage != null) { - try { - this.dictionary = dictionaryPage.getEncoding().initDictionary(descriptor, dictionaryPage); - this.isCurrentPageDictionaryEncoded = true; - } catch (IOException e) { - throw new IOException("could not decode the dictionary for " + descriptor, e); - } catch (ParquetDecodingException e) { - throw new ParquetDecodingException( - "could not decode the dictionary for " + descriptor + ": " + e.getMessage(), e); + Dictionary dict = null; + try { + // This decompresses the page eagerly for some codecs, e.g. ZSTD. + DictionaryPage dictionaryPage = pageReader.readDictionaryPage(); + if (dictionaryPage != null) { + dict = dictionaryPage.getEncoding().initDictionary(descriptor, dictionaryPage); } - } else { - this.dictionary = null; - this.isCurrentPageDictionaryEncoded = false; + } catch (IOException e) { + throw new IOException("could not decode the dictionary for " + descriptor, e); + } catch (ParquetDecodingException e) { + throw new ParquetDecodingException( + "could not decode the dictionary for " + descriptor + ": " + e.getMessage(), e); + } catch (IndexOutOfBoundsException | NegativeArraySizeException | + UnsupportedOperationException e) { + // Reading the dictionary page decompresses it, for some codecs only when initDictionary + // reads its bytes, and initDictionary parses it. A corrupted page or page header makes + // them throw these: an IndexOutOfBoundsException when a BINARY dictionary ends within the + // length of an entry, a NegativeArraySizeException for a negative size, and an + // UnsupportedOperationException for an encoding that has no dictionary. An entry that + // claims more bytes than the page has is not detected here (SPARK-60316). + throw new ParquetDecodingException( + "could not decode the dictionary for " + descriptor + ": " + e, e); } + this.dictionary = dict; + this.isCurrentPageDictionaryEncoded = dict != null; if (pageReader.getTotalValueCount() == 0) { - throw new IOException("totalValueCount == 0"); + throw new IOException("totalValueCount == 0 in col " + descriptor); } assert "LEGACY".equals(datetimeRebaseMode) || "EXCEPTION".equals(datetimeRebaseMode) || "CORRECTED".equals(datetimeRebaseMode); @@ -272,79 +288,96 @@ void readBatch( } readState.resetForNewPage(pageValueCount, pageFirstRowIndex); } - PrimitiveType.PrimitiveTypeName typeName = - descriptor.getPrimitiveType().getPrimitiveTypeName(); - if (isCurrentPageDictionaryEncoded) { - // Save starting offset in case we need to decode dictionary IDs. - int startOffset = readState.valueOffset; - // Save starting row index so we can check if we need to eagerly decode dict ids later - long startRowId = readState.rowId; - - // Read and decode dictionary ids. - if (readState.maxRepetitionLevel == 0) { - defColumn.readIntegers(readState, dictionaryIds, column, definitionLevels, - (VectorizedValuesReader) dataColumn); - } else { - repColumn.readIntegersRepeated(readState, repetitionLevels, defColumn, definitionLevels, - dictionaryIds, column, (VectorizedValuesReader) dataColumn); - } + // A corrupted page makes the readers below throw a ParquetDecodingException without the + // column. Only that is wrapped, so that other exceptions, e.g. SparkUpgradeException, + // keep their types. + try { + PrimitiveType.PrimitiveTypeName typeName = + descriptor.getPrimitiveType().getPrimitiveTypeName(); + if (isCurrentPageDictionaryEncoded) { + // Save starting offset in case we need to decode dictionary IDs. + int startOffset = readState.valueOffset; + // Save starting row index so we can check if we need to eagerly decode dict ids later + long startRowId = readState.rowId; + + // Read and decode dictionary ids. + if (readState.maxRepetitionLevel == 0) { + defColumn.readIntegers(readState, dictionaryIds, column, definitionLevels, + (VectorizedValuesReader) dataColumn); + } else { + repColumn.readIntegersRepeated(readState, repetitionLevels, defColumn, definitionLevels, + dictionaryIds, column, (VectorizedValuesReader) dataColumn); + } - // TIMESTAMP_MILLIS encoded as INT64 can't be lazily decoded as we need to post process - // the values to add microseconds precision. - if (column.hasDictionary() || (startRowId == pageFirstRowIndex && - isLazyDecodingSupported(typeName, column.dataType()))) { - // Column vector supports lazy decoding of dictionary values so just set the dictionary. - // We can't do this if startRowId is not the first row index in the page AND the column - // doesn't have a dictionary (i.e. some non-dictionary encoded values have already been - // added). - PrimitiveType primitiveType = descriptor.getPrimitiveType(); - - // We need to make sure that we initialize the right type for the dictionary otherwise - // WritableColumnVector will throw an exception when trying to decode to an Int when the - // dictionary is in fact initialized as Long - LogicalTypeAnnotation typeAnnotation = primitiveType.getLogicalTypeAnnotation(); - boolean castLongToInt = - typeAnnotation instanceof DecimalLogicalTypeAnnotation annotation && - annotation.getPrecision() <= Decimal.MAX_INT_DIGITS() && - primitiveType.getPrimitiveTypeName() == INT64; - - // We require a long value, but we need to use dictionary to decode the original - // signed int first - boolean isUnsignedInt32 = updaterFactory.isUnsignedIntTypeMatched(32); - - // We require a decimal value, but we need to use dictionary to decode the original - // signed long first - boolean isUnsignedInt64 = updaterFactory.isUnsignedIntTypeMatched(64); - - boolean needTransform = castLongToInt || isUnsignedInt32 || isUnsignedInt64; - column.setDictionary(new ParquetDictionary(dictionary, needTransform)); - } else { - updater.decodeDictionaryIds(readState.valueOffset - startOffset, startOffset, column, - dictionaryIds, dictionary); - } - } else { - if (column.hasDictionary() && readState.valueOffset != 0) { - // This batch already has dictionary encoded values but this new page is not. The batch - // does not support a mix of dictionary and not so we will decode the dictionary. - updater.decodeDictionaryIds(readState.valueOffset, 0, column, dictionaryIds, dictionary); - } - column.setDictionary(null); - VectorizedValuesReader valuesReader = (VectorizedValuesReader) dataColumn; - if (readState.maxRepetitionLevel == 0) { - defColumn.readBatch(readState, column, definitionLevels, valuesReader, updater); + // TIMESTAMP_MILLIS encoded as INT64 can't be lazily decoded as we need to post process + // the values to add microseconds precision. + if (column.hasDictionary() || (startRowId == pageFirstRowIndex && + isLazyDecodingSupported(typeName, column.dataType()))) { + // Column vector supports lazy decoding of dictionary values so just set the dictionary. + // We can't do this if startRowId is not the first row index in the page AND the column + // doesn't have a dictionary (i.e. some non-dictionary encoded values have already been + // added). + PrimitiveType primitiveType = descriptor.getPrimitiveType(); + + // We need to make sure that we initialize the right type for the dictionary otherwise + // WritableColumnVector will throw an exception when trying to decode to an Int when the + // dictionary is in fact initialized as Long + LogicalTypeAnnotation typeAnnotation = primitiveType.getLogicalTypeAnnotation(); + boolean castLongToInt = + typeAnnotation instanceof DecimalLogicalTypeAnnotation annotation && + annotation.getPrecision() <= Decimal.MAX_INT_DIGITS() && + primitiveType.getPrimitiveTypeName() == INT64; + + // We require a long value, but we need to use dictionary to decode the original + // signed int first + boolean isUnsignedInt32 = updaterFactory.isUnsignedIntTypeMatched(32); + + // We require a decimal value, but we need to use dictionary to decode the original + // signed long first + boolean isUnsignedInt64 = updaterFactory.isUnsignedIntTypeMatched(64); + + boolean needTransform = castLongToInt || isUnsignedInt32 || isUnsignedInt64; + column.setDictionary(new ParquetDictionary(dictionary, needTransform)); + } else { + updater.decodeDictionaryIds(readState.valueOffset - startOffset, startOffset, column, + dictionaryIds, dictionary); + } } else { - repColumn.readBatchRepeated(readState, repetitionLevels, defColumn, definitionLevels, - column, valuesReader, updater); + if (column.hasDictionary() && readState.valueOffset != 0) { + // This batch already has dictionary encoded values but this new page is not. The batch + // does not support a mix of dictionary and not so we will decode the dictionary. + updater.decodeDictionaryIds(readState.valueOffset, 0, column, dictionaryIds, + dictionary); + } + column.setDictionary(null); + VectorizedValuesReader valuesReader = (VectorizedValuesReader) dataColumn; + if (readState.maxRepetitionLevel == 0) { + defColumn.readBatch(readState, column, definitionLevels, valuesReader, updater); + } else { + repColumn.readBatchRepeated(readState, repetitionLevels, defColumn, definitionLevels, + column, valuesReader, updater); + } } + } catch (ParquetDecodingException e) { + throw new ParquetDecodingException("could not decode page " + currentPage + " in col " + + descriptor + ": " + e.getMessage(), e); } } } private int readPage() { - DataPage page = pageReader.readPage(); + DataPage page; + try { + // This decompresses the page eagerly for some codecs, e.g. ZSTD. + page = pageReader.readPage(); + } catch (ParquetDecodingException e) { + throw new ParquetDecodingException( + "could not read the next page in col " + descriptor + ": " + e.getMessage(), e); + } if (page == null) { return -1; } + this.currentPage = page.toString(); this.pageFirstRowIndex = page.getFirstRowIndex().orElse(0L); return page.accept(new DataPage.Visitor() { diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedDeltaBinaryPackedReader.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedDeltaBinaryPackedReader.java index c04fb49f1addf..9e4f0339a2de9 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedDeltaBinaryPackedReader.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedDeltaBinaryPackedReader.java @@ -26,7 +26,6 @@ import org.apache.parquet.column.values.bitpacking.Packer; import org.apache.parquet.io.ParquetDecodingException; -import org.apache.spark.network.util.JavaUtils; import org.apache.spark.sql.catalyst.util.RebaseDateTime; import org.apache.spark.sql.execution.datasources.DataSourceUtils; import org.apache.spark.sql.execution.vectorized.WritableColumnVector; @@ -89,8 +88,10 @@ public class VectorizedDeltaBinaryPackedReader extends VectorizedReaderBase { @Override public void initFromPage(int valueCount, ByteBufferInputStream in) throws IOException { - JavaUtils.checkArgument(valueCount >= 1, - "Page must have at least one value, but it has " + valueCount); + if (valueCount < 1) { + throw new ParquetDecodingException("Invalid DELTA_BINARY_PACKED page: it must have at " + + "least one value, but it has " + valueCount); + } this.in = in; // Read the header this.blockSizeInValues = BytesUtils.readUnsignedVarInt(in); @@ -100,14 +101,17 @@ public void initFromPage(int valueCount, ByteBufferInputStream in) throws IOExce throw new ParquetDecodingException("Invalid DELTA_BINARY_PACKED block size in values: " + blockSizeInValues + " (must be in [1, " + MAX_BLOCK_SIZE_IN_VALUES + "])"); } - if (miniBlockNumInABlock <= 0 || miniBlockNumInABlock > blockSizeInValues) { + if (miniBlockNumInABlock <= 0) { 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) { + throw new ParquetDecodingException("Invalid DELTA_BINARY_PACKED mini block count: " + + miniBlockNumInABlock + " (block size in values: " + blockSizeInValues + + ", the mini block size must be a multiple of 8)"); + } + this.miniBlockSizeInValues = blockSizeInValues / miniBlockNumInABlock; // True value count. May be less than valueCount because of nulls this.totalValueCount = BytesUtils.readUnsignedVarInt(in); if (totalValueCount < 0) { @@ -472,9 +476,15 @@ private void readBlockHeader() { * see org.apache.parquet.column.values.delta.DeltaBinaryPackingValuesReader#unpackMiniBlock */ private void unpackMiniBlock() throws IOException { + int bitWidth = bitWidths[currentMiniBlock]; + // The bit widths of the unused mini blocks of the last block can be anything, so a width is + // checked only when its mini block is unpacked. + if (bitWidth > 64) { + throw new ParquetDecodingException("Invalid DELTA_BINARY_PACKED mini block bit width: " + + bitWidth + " (must be at most 64)"); + } Arrays.fill(this.unpackedValuesBuffer, 0); - BytePackerForLong packer = Packer.LITTLE_ENDIAN.newBytePackerForLong( - bitWidths[currentMiniBlock]); + BytePackerForLong packer = Packer.LITTLE_ENDIAN.newBytePackerForLong(bitWidth); for (int j = 0; j < miniBlockSizeInValues; j += 8) { ByteBuffer buffer = in.slice(packer.getBitWidth()); if (buffer.hasArray()) { diff --git a/sql/core/src/test/java/test/org/apache/parquet/column/page/TestDataPage.java b/sql/core/src/test/java/test/org/apache/parquet/column/page/TestDataPage.java index da30819195891..792735b85e519 100644 --- a/sql/core/src/test/java/test/org/apache/parquet/column/page/TestDataPage.java +++ b/sql/core/src/test/java/test/org/apache/parquet/column/page/TestDataPage.java @@ -41,4 +41,9 @@ public Optional getIndexRowCount() { public T accept(Visitor visitor) { return wrapped.accept(visitor); } + + @Override + public String toString() { + return wrapped.toString(); + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetDeltaEncodingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetDeltaEncodingSuite.scala index bf2bf5b010fad..eb0574a9bfdcf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetDeltaEncodingSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetDeltaEncodingSuite.scala @@ -204,6 +204,25 @@ abstract class ParquetDeltaEncodingSuite[T] extends ParquetCompatibilityTest reader.initFromPage(100, craftHeader(128, 4, -1)) } assert(e3.getMessage.contains("total value count")) + + // The mini blocks must have a multiple of 8 values: 128 values do not split into 3 mini + // blocks, and split into 32 mini blocks of 4 values. + Seq(3, 32).foreach { miniBlockNum => + reader = new VectorizedDeltaBinaryPackedReader + val e4 = intercept[ParquetDecodingException] { + reader.initFromPage(100, craftHeader(128, miniBlockNum, 1)) + } + assert(e4.getMessage == s"Invalid DELTA_BINARY_PACKED mini block count: $miniBlockNum " + + "(block size in values: 128, the mini block size must be a multiple of 8)") + } + + // A page must have at least one value. + reader = new VectorizedDeltaBinaryPackedReader + val e5 = intercept[ParquetDecodingException] { + reader.initFromPage(0, craftHeader(128, 4, 1)) + } + assert(e5.getMessage == + "Invalid DELTA_BINARY_PACKED page: it must have at least one value, but it has 0") } test("skip()") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala index 7fc1bf186f6af..4f51c87a0f3f8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala @@ -17,6 +17,7 @@ package org.apache.spark.sql.execution.datasources.parquet +import java.io.ByteArrayInputStream import java.math.{BigDecimal => JBigDecimal} import java.nio.file.Files import java.time.{LocalDateTime, LocalTime} @@ -36,6 +37,7 @@ import org.apache.parquet.column.{Encoding, ParquetProperties} import org.apache.parquet.column.ParquetProperties.WriterVersion.PARQUET_1_0 import org.apache.parquet.example.data.Group import org.apache.parquet.example.data.simple.{SimpleGroup, SimpleGroupFactory} +import org.apache.parquet.format.{PageType, Util} import org.apache.parquet.hadoop._ import org.apache.parquet.hadoop.example.ExampleParquetWriter import org.apache.parquet.io.ParquetDecodingException @@ -2488,31 +2490,18 @@ 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. + * Writes 100 rows of `column` with the writer `options`, lets `corrupt` modify the bytes of + * the written file, and returns the error of reading the file with the vectorized reader. */ - private def readCorruptedRleSection( - column: String, - writerVersion: String = "PARQUET_1_0")( - corrupt: (Array[Byte], Int) => Unit): Throwable = { + private def readCorruptedFile(column: String, options: Map[String, String])( + corrupt: Array[Byte] => 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) + spark.range(100).selectExpr(column).coalesce(1).write.options(options).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 offset = bytes.indexOfSlice(section) - assert(offset >= 0 && bytes.indexOfSlice(section, offset + 1) < 0) - corrupt(bytes, offset) + corrupt(bytes) Files.write(file.toPath, bytes) dir.listFiles().filter(_.getName.endsWith(".crc")).foreach(_.delete()) @@ -2528,6 +2517,29 @@ class ParquetIOSuite extends ParquetTest with SharedSparkSession { error } + /** + * 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 = { + val options = Map( + ParquetOutputFormat.ENABLE_DICTIONARY -> "false", + ParquetOutputFormat.WRITER_VERSION -> writerVersion, + "compression" -> "uncompressed") + readCorruptedFile(column, options) { bytes => + val section = Array[Byte](3, 0, 0, 0, 0xC8.toByte, 0x01, 1) + val offset = bytes.indexOfSlice(section) + assert(offset >= 0 && bytes.indexOfSlice(section, offset + 1) < 0) + corrupt(bytes, offset) + } + } + private def causes(e: Throwable): Seq[Throwable] = ExceptionUtils.getThrowableList(e).asScala.toSeq @@ -2569,6 +2581,26 @@ 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(PageType.DATA_PAGE -> false, PageType.DICTIONARY_PAGE -> true).foreach { + case (pageType, dictionary) => + val options = Map( + ParquetOutputFormat.ENABLE_DICTIONARY -> dictionary.toString, + "compression" -> "zstd") + val e = readCorruptedFile("cast(id % 10 as int) as a", options) { bytes => + // The first page of the only column chunk starts right after the magic number. + val start = ParquetFileWriter.MAGIC.length + val in = new ByteArrayInputStream(bytes, start, bytes.length - start) + assert(Util.readPageHeader(in).getType == pageType) + // Invert the first bytes of the page, the ZSTD frame magic number. + val pageStart = bytes.length - in.available() + (pageStart until pageStart + 4).foreach(i => bytes(i) = (~bytes(i)).toByte) + } + assert(causes(e).exists(c => + Option(c.getMessage).exists(_.contains("[a] optional int32 a"))), e) + } + } } class JobCommitFailureParquetOutputCommitter(outputPath: Path, context: TaskAttemptContext) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorizedSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorizedSuite.scala index 8611970baa123..5749b602a635b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorizedSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetVectorizedSuite.scala @@ -17,6 +17,7 @@ package org.apache.spark.sql.execution.datasources.parquet +import java.io.IOException import java.util.{Optional, PrimitiveIterator} import scala.collection.mutable.ArrayBuffer @@ -54,6 +55,11 @@ class ParquetVectorizedSuite extends ParquetTest with SharedSparkSession { private val NUM_VALUES: Int = VALUES.length private val BATCH_SIZE_CONFIGS: Seq[Int] = Seq(1, 3, 5, 7, 10, 20, 40) private val PAGE_SIZE_CONFIGS: Seq[Seq[Int]] = Seq(Seq(6, 6, 7, 7), Seq(4, 9, 4, 9)) + // 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) + // One non-null value: the definition levels (length 2, an RLE run of 1 x 1), then the + // dictionary ids (bit width 2, an RLE run of 1 x 0). + private val oneValueWithDictionaryId0 = Array[Byte](2, 0, 0, 0, 2, 1, 2, 2, 0) implicit def toStrings(ints: Seq[Int]): Seq[String] = ints.map(i => ('a' + i).toChar.toString) @@ -526,29 +532,12 @@ class ParquetVectorizedSuite extends ParquetTest with SharedSparkSession { test("SPARK-59832: a corrupted dictionary page reports the column") { // A PLAIN INT32 dictionary page that declares 3 entries but holds only 2. - val parquetSchema = MessageTypeParser.parseMessageType("message root { optional int32 a; }") - val ty = parquetSchema.asGroupType().getType("a").asPrimitiveType() - val cd = new ColumnDescriptor(Array("a"), ty, 0, 1) - val memPageStore = new MemPageStore(1) - val pageWriter = memPageStore.getPageWriter(cd) - pageWriter.writeDictionaryPage(new DictionaryPage( - BytesInput.from(Array[Byte](1, 0, 0, 0, 2, 0, 0, 0)), 3, Encoding.PLAIN_DICTIONARY)) - // One non-null value: the definition levels (length 2, an RLE run of 1 x 1), then the - // dictionary ids (bit width 2, an RLE run of 1 x 0). - pageWriter.writePage(BytesInput.from(Array[Byte](2, 0, 0, 0, 2, 1, 2, 2, 0)), 1, 1, - Statistics.createStats(ty), Encoding.RLE, Encoding.RLE, Encoding.RLE_DICTIONARY) - val recordReader = new VectorizedParquetRecordReader( - DateTimeUtils.getZoneId("EST"), "CORRECTED", "UTC", "CORRECTED", "UTC", true, 1) - try { - recordReader.initialize(parquetSchema, parquetSchema, - TestParquetRowGroupReader(Seq(TestPageReadStore(memPageStore, Seq(0L)))), 1) - val e = intercept[ParquetDecodingException](recordReader.nextKeyValue()) - assert(e.getMessage.contains("could not decode the dictionary for [a]"), e.getMessage) - assert(e.getMessage.contains(e.getCause.getMessage), e.getMessage) - assert(e.getCause.isInstanceOf[ParquetDecodingException], e) - } finally { - recordReader.close() - } + val e = readCorruptedPage("int32", Some(new DictionaryPage( + BytesInput.from(Array[Byte](1, 0, 0, 0, 2, 0, 0, 0)), 3, Encoding.PLAIN_DICTIONARY)), + oneValueWithDictionaryId0, 1, Encoding.RLE_DICTIONARY) + assert(e.getMessage.contains("could not decode the dictionary for [a]"), e.getMessage) + assert(e.getMessage.contains(e.getCause.getMessage), e.getMessage) + assert(e.getCause.isInstanceOf[ParquetDecodingException], e) } test("SPARK-60262: a column requested twice fails instead of leaving a vector unwritten") { @@ -585,6 +574,78 @@ class ParquetVectorizedSuite extends ParquetTest with SharedSparkSession { } } + 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), + // A DELTA_BINARY_PACKED block of 128 values in 4 mini blocks, with 4 values and the first + // value 0, then the min delta 0 and the mini block bit widths 65, 0, 0 and 0. The second + // value unpacks the first mini block. + ("Invalid DELTA_BINARY_PACKED mini block bit width: 65 (must be at most 64)", None, + Array[Byte](0x80.toByte, 1, 4, 4, 0, 0, 65, 0, 0, 0), Encoding.DELTA_BINARY_PACKED) + ).foreach { case (error, dictionaryPage, values, encoding) => + val e = readCorruptedPage("int32", dictionaryPage, fourNonNullValues ++ values, 4, encoding) + assert(e.getMessage.startsWith("could not decode page Page [bytes.size="), e.getMessage) + assert(e.getMessage.contains(s"in col [a] optional int32 a: $error"), e.getMessage) + assert(e.getCause.isInstanceOf[ParquetDecodingException], e) + assert(e.getCause.getMessage == error, e) + } + } + + test("SPARK-59854: a dictionary that cannot be decoded reports the column") { + Seq( + // A PLAIN BINARY dictionary page that declares 3 entries but holds only 1, which cuts off + // the length of the second entry. + ("binary", new DictionaryPage(BytesInput.from(Array[Byte](1, 0, 0, 0, 'x')), 3, + Encoding.PLAIN_DICTIONARY), classOf[IndexOutOfBoundsException]), + // A dictionary page with a negative dictionary size. + ("int32", new DictionaryPage(BytesInput.from(Array[Byte](1, 0, 0, 0)), -1, + Encoding.PLAIN_DICTIONARY), classOf[NegativeArraySizeException]), + // A dictionary page with an encoding that has no dictionary. + ("int32", new DictionaryPage(BytesInput.from(Array[Byte](1, 0, 0, 0)), 1, + Encoding.DELTA_BINARY_PACKED), classOf[UnsupportedOperationException]) + ).foreach { case (columnType, dictionaryPage, errorClass) => + val e = readCorruptedPage(columnType, Some(dictionaryPage), oneValueWithDictionaryId0, 1, + Encoding.RLE_DICTIONARY) + assert(e.getMessage.startsWith( + s"could not decode the dictionary for [a] optional $columnType a: ${errorClass.getName}"), + e.getMessage) + assert(e.getCause.getClass == errorClass, e) + } + } + + test("SPARK-59854: an invalid DELTA_BINARY_PACKED header reports the column") { + // One non-null value. The header has a block of 128 values in 3 mini blocks. + val e = readCorruptedPage("int32", None, + Array[Byte](2, 0, 0, 0, 2, 1, 0x80.toByte, 1, 3, 1, 0), 1, Encoding.DELTA_BINARY_PACKED) + assert(e.getMessage.contains("in col [a] optional int32 a: Invalid DELTA_BINARY_PACKED " + + "mini block count: 3 (block size in values: 128, the mini block size must be a multiple " + + "of 8)"), e.getMessage) + assert(e.getCause.isInstanceOf[ParquetDecodingException], e) + } + + test("SPARK-59854: a column chunk without values reports the column") { + val parquetSchema = MessageTypeParser.parseMessageType("message root { optional int32 a; }") + val ty = parquetSchema.asGroupType().getType("a").asPrimitiveType() + val memPageStore = new MemPageStore(1) + // The column has a page writer, so the page store has a page reader for it, but no page. + memPageStore.getPageWriter(new ColumnDescriptor(Array("a"), ty, 0, 1)) + val recordReader = new VectorizedParquetRecordReader( + DateTimeUtils.getZoneId("EST"), "CORRECTED", "UTC", "CORRECTED", "UTC", true, 1) + try { + recordReader.initialize(parquetSchema, parquetSchema, + TestParquetRowGroupReader(Seq(TestPageReadStore(memPageStore, Seq(0L)))), 1) + val e = intercept[IOException](recordReader.nextKeyValue()) + assert(e.getMessage == "totalValueCount == 0 in col [a] optional int32 a") + } finally { + recordReader.close() + } + } + truncateTypeTest("primitive type", IntegerType, LongType, IntegerType) truncateTypeTest("basic struct", @@ -887,6 +948,37 @@ class ParquetVectorizedSuite extends ParquetTest with SharedSparkSession { columnWriterStore.flush() } + /** + * Reads a row group of the optional column `a` of type `columnType`, with an optional + * dictionary page and one data page of `valueCount` values, and returns the error of reading + * it with the vectorized reader. + */ + private def readCorruptedPage( + columnType: String, + dictionaryPage: Option[DictionaryPage], + data: Array[Byte], + valueCount: Int, + valuesEncoding: Encoding): ParquetDecodingException = { + val parquetSchema = + MessageTypeParser.parseMessageType(s"message root { optional $columnType a; }") + 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() + } + } + private def checkAnswer( totalRowCount: Int, fileSchema: MessageType,