viirya commented on code in PR #59542:
URL: https://github.com/apache/spark/pull/59542#discussion_r4298490028
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedColumnReader.java:
##########
@@ -264,79 +281,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);
- }
-
- // 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);
+ // 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 {
- 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) {
Review Comment:
Right, an out-of-range id fails with an `ArrayIndexOutOfBoundsException`
from `Dictionary.decodeTo*`, and with lazy decoding only after `readBatch`
returns. Covering both needs a check of every id against `getMaxId()`, which is
per-value work on the hot path (the lazy path does not look at the ids at all
today), so I'd rather leave it to a follow-up. I narrowed the description: only
`ParquetDecodingException`s are covered, and out-of-range ids are listed as not
covered.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedColumnReader.java:
##########
@@ -157,6 +169,11 @@ public VectorizedColumnReader(
} catch (ParquetDecodingException e) {
throw new ParquetDecodingException(
"could not decode the dictionary for " + descriptor + ": " +
e.getMessage(), e);
+ } catch (RuntimeException e) {
Review Comment:
Good point, and the comment was wrong: for SNAPPY and GZIP the dictionary is
decompressed when `initDictionary` reads the bytes. The catch now takes only
what a corrupted page or page header causes there: `IndexOutOfBoundsException`
(a truncated BINARY dictionary), `NegativeArraySizeException` (a negative
dictionary size in the page header) and `UnsupportedOperationException` (an
encoding that has no dictionary). The test covers all three, and the comment
now says that decompression can happen here. Decompression errors there are
`IOException`s, which the existing catch already handles.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedColumnReader.java:
##########
@@ -264,79 +281,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);
- }
-
- // 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);
+ // 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 {
- 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;
Review Comment:
Done, `currentPage` is now set before the null check, so it is cleared at
the end of the column chunk.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala:
##########
@@ -2569,6 +2571,47 @@ 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") {
+ // ZSTD pages are decompressed when they are read from the column chunk,
not when decoded.
+ Seq(
+ (false, "could not read the next page in col", "could not decompress
page"),
+ (true, "could not read the dictionary page for", "Could not decompress
dictionary page")
+ ).foreach { case (dictionary, error, cause) =>
+ withTempPath { dir =>
Review Comment:
Done. `readCorruptedFile` now takes the writer options and the corruption
step; `readCorruptedRleSection` and the ZSTD test both use it, with the
60-second `failAfter` guard.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala:
##########
@@ -2569,6 +2571,47 @@ 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") {
+ // ZSTD pages are decompressed when they are read from the column chunk,
not when decoded.
+ Seq(
+ (false, "could not read the next page in col", "could not decompress
page"),
+ (true, "could not read the dictionary page for", "Could not decompress
dictionary page")
+ ).foreach { case (dictionary, error, cause) =>
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ spark.range(1000).selectExpr("cast(id % 100 as int) as a").coalesce(1)
+ .write
+ .option(ParquetOutputFormat.ENABLE_DICTIONARY, dictionary.toString)
+ .option("compression", "zstd")
+ .parquet(path)
+ val file = dir.listFiles().filter(_.getName.endsWith(".parquet")).head
+ val footer = ParquetFileReader.readFooter(
+ spark.sessionState.newHadoopConf(), new Path(file.getCanonicalPath))
+ val chunk = footer.getBlocks.get(0).getColumns.get(0)
+ assert(chunk.hasDictionaryPage == dictionary)
+ // Invert the first bytes of the first page, the ZSTD frame magic
number.
+ val bytes = Files.readAllBytes(file.toPath)
+ val in = new ByteArrayInputStream(bytes, chunk.getStartingPos.toInt,
bytes.length)
+ Util.readPageHeader(in)
+ val pageStart = bytes.length - in.available()
+ (pageStart until pageStart + 4).foreach(i => bytes(i) =
(~bytes(i)).toByte)
+ Files.write(file.toPath, bytes)
+ dir.listFiles().filter(_.getName.endsWith(".crc")).foreach(_.delete())
+
+ withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "true") {
+ val e = intercept[SparkException](spark.read.parquet(path).collect())
+ val columnErrors =
+ causes(e).filter(c =>
Option(c.getMessage).exists(_.contains("optional int32 a")))
+ assert(columnErrors.size == 1, e)
+ assert(columnErrors.head.getMessage ==
Review Comment:
Done, the test now only checks that the column is reported once, on a
`ParquetDecodingException`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala:
##########
@@ -2569,6 +2571,47 @@ 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") {
+ // ZSTD pages are decompressed when they are read from the column chunk,
not when decoded.
+ Seq(
+ (false, "could not read the next page in col", "could not decompress
page"),
+ (true, "could not read the dictionary page for", "Could not decompress
dictionary page")
+ ).foreach { case (dictionary, error, cause) =>
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ spark.range(1000).selectExpr("cast(id % 100 as int) as a").coalesce(1)
+ .write
+ .option(ParquetOutputFormat.ENABLE_DICTIONARY, dictionary.toString)
+ .option("compression", "zstd")
+ .parquet(path)
+ val file = dir.listFiles().filter(_.getName.endsWith(".parquet")).head
+ val footer = ParquetFileReader.readFooter(
Review Comment:
`ParquetTest.readFooter` reads the `_metadata` summary file of a directory
(`new Path(path, PARQUET_METADATA_FILE)`), which is why the tests at L1353 and
L1527 write summary files first, so it would not find the footer of this data
file. I switched to the non-deprecated
`ParquetFooterReader.readFooter(HadoopInputFile.fromPath(...), NO_FILTER)`
instead.
--
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]
