This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new c20d5195c34 [SPARK-38891][SQL] Skipping allocating vector for
repetition & definition levels when possible
c20d5195c34 is described below
commit c20d5195c343b0a64a588ca56baacb52000e1c90
Author: Chao Sun <sunc...@apple.com>
AuthorDate: Wed May 4 09:37:06 2022 -0700
[SPARK-38891][SQL] Skipping allocating vector for repetition & definition
levels when possible
### What changes were proposed in this pull request?
This PR adds two optimization on the vectorized Parquet reader for complex
types:
- avoid allocating vectors for repetition and definition levels whenever
can be applied
- avoid reading definition levels whenever can be applied
### Why are the changes needed?
At the moment, Spark will allocate vectors for repetition and definition
levels, and also read definition levels even if it's not necessary, for
instance, when reading primitive types. This may add extra memory footprint
especially when reading wide tables. Therefore, we should avoid them if
possible.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
Existing tests.
Closes #36202 from sunchao/SPARK-38891.
Authored-by: Chao Sun <sunc...@apple.com>
Signed-off-by: Chao Sun <sunc...@apple.com>
---
.../datasources/parquet/ParquetColumnVector.java | 39 +++++++---
.../parquet/VectorizedRleValuesReader.java | 88 +++++++++++++++++++++-
2 files changed, 115 insertions(+), 12 deletions(-)
diff --git
a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetColumnVector.java
b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetColumnVector.java
index 4b29520d30f..c8399d9137f 100644
---
a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetColumnVector.java
+++
b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetColumnVector.java
@@ -61,7 +61,16 @@ final class ParquetColumnVector {
int capacity,
MemoryMode memoryMode,
Set<ParquetColumn> missingColumns) {
+ this(column, vector, capacity, memoryMode, missingColumns, true);
+ }
+ ParquetColumnVector(
+ ParquetColumn column,
+ WritableColumnVector vector,
+ int capacity,
+ MemoryMode memoryMode,
+ Set<ParquetColumn> missingColumns,
+ boolean isTopLevel) {
DataType sparkType = column.sparkType();
if (!sparkType.sameType(vector.dataType())) {
throw new IllegalArgumentException("Spark type: " + sparkType +
@@ -79,20 +88,27 @@ final class ParquetColumnVector {
}
if (isPrimitive) {
- // TODO: avoid allocating these if not necessary, for instance, the node
is of top-level
- // and is not repeated, or the node is not top-level but its max
repetition level is 0.
- repetitionLevels = allocateLevelsVector(capacity, memoryMode);
- definitionLevels = allocateLevelsVector(capacity, memoryMode);
+ if (column.repetitionLevel() > 0) {
+ repetitionLevels = allocateLevelsVector(capacity, memoryMode);
+ }
+ // We don't need to create and store definition levels if the column is
top-level.
+ if (!isTopLevel) {
+ definitionLevels = allocateLevelsVector(capacity, memoryMode);
+ }
} else {
Preconditions.checkArgument(column.children().size() ==
vector.getNumChildren());
+ boolean allChildrenAreMissing = true;
+
for (int i = 0; i < column.children().size(); i++) {
ParquetColumnVector childCv = new
ParquetColumnVector(column.children().apply(i),
- vector.getChild(i), capacity, memoryMode, missingColumns);
+ vector.getChild(i), capacity, memoryMode, missingColumns, false);
children.add(childCv);
+
// Only use levels from non-missing child, this can happen if only
some but not all
// fields of a struct are missing.
if (!childCv.vector.isAllNull()) {
+ allChildrenAreMissing = false;
this.repetitionLevels = childCv.repetitionLevels;
this.definitionLevels = childCv.definitionLevels;
}
@@ -100,7 +116,7 @@ final class ParquetColumnVector {
// This can happen if all the fields of a struct are missing, in which
case we should mark
// the struct itself as a missing column
- if (repetitionLevels == null) {
+ if (allChildrenAreMissing) {
vector.setAllNull();
}
}
@@ -163,8 +179,12 @@ final class ParquetColumnVector {
if (vector.isAllNull()) return;
vector.reset();
- repetitionLevels.reset();
- definitionLevels.reset();
+ if (repetitionLevels != null) {
+ repetitionLevels.reset();
+ }
+ if (definitionLevels != null) {
+ definitionLevels.reset();
+ }
for (ParquetColumnVector child : children) {
child.reset();
}
@@ -289,7 +309,8 @@ final class ParquetColumnVector {
vector.reserve(definitionLevels.getElementsAppended());
int rowId = 0;
- boolean hasRepetitionLevels = repetitionLevels.getElementsAppended() > 0;
+ boolean hasRepetitionLevels =
+ repetitionLevels != null && repetitionLevels.getElementsAppended() > 0;
for (int i = 0; i < definitionLevels.getElementsAppended(); i++) {
// If repetition level > maxRepetitionLevel, the value is a nested
element (e.g., an array
// element in struct<array<int>>), and we should skip the definition
level since it doesn't
diff --git
a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReader.java
b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReader.java
index 2cc763a5b72..12fe5697954 100644
---
a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReader.java
+++
b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedRleValuesReader.java
@@ -172,7 +172,11 @@ public final class VectorizedRleValuesReader extends
ValuesReader
WritableColumnVector defLevels,
VectorizedValuesReader valueReader,
ParquetVectorUpdater updater) {
- readBatchInternal(state, values, values, defLevels, valueReader, updater);
+ if (defLevels == null) {
+ readBatchInternal(state, values, values, valueReader, updater);
+ } else {
+ readBatchInternalWithDefLevels(state, values, values, defLevels,
valueReader, updater);
+ }
}
/**
@@ -185,11 +189,89 @@ public final class VectorizedRleValuesReader extends
ValuesReader
WritableColumnVector nulls,
WritableColumnVector defLevels,
VectorizedValuesReader valueReader) {
- readBatchInternal(state, values, nulls, defLevels, valueReader,
- new ParquetVectorUpdaterFactory.IntegerUpdater());
+ if (defLevels == null) {
+ readBatchInternal(state, values, nulls, valueReader,
+ new ParquetVectorUpdaterFactory.IntegerUpdater());
+ } else {
+ readBatchInternalWithDefLevels(state, values, nulls, defLevels,
valueReader,
+ new ParquetVectorUpdaterFactory.IntegerUpdater());
+ }
}
private void readBatchInternal(
+ ParquetReadState state,
+ WritableColumnVector values,
+ WritableColumnVector nulls,
+ VectorizedValuesReader valueReader,
+ ParquetVectorUpdater updater) {
+
+ long rowId = state.rowId;
+ int leftInBatch = state.rowsToReadInBatch;
+ int leftInPage = state.valuesToReadInPage;
+
+ while (leftInBatch > 0 && leftInPage > 0) {
+ if (currentCount == 0 && !readNextGroup()) break;
+ int n = Math.min(leftInBatch, Math.min(leftInPage, this.currentCount));
+
+ long rangeStart = state.currentRangeStart();
+ long rangeEnd = state.currentRangeEnd();
+
+ if (rowId + n < rangeStart) {
+ skipValues(n, state, valueReader, updater);
+ rowId += n;
+ leftInPage -= n;
+ } else if (rowId > rangeEnd) {
+ state.nextRange();
+ } else {
+ // The range [rowId, rowId + n) overlaps with the current row range in
state
+ long start = Math.max(rangeStart, rowId);
+ long end = Math.min(rangeEnd, rowId + n - 1);
+
+ // Skip the part [rowId, start)
+ int toSkip = (int) (start - rowId);
+ if (toSkip > 0) {
+ skipValues(toSkip, state, valueReader, updater);
+ rowId += toSkip;
+ leftInPage -= toSkip;
+ }
+
+ // Read the part [start, end]
+ n = (int) (end - start + 1);
+
+ switch (mode) {
+ case RLE:
+ if (currentValue == state.maxDefinitionLevel) {
+ updater.readValues(n, state.valueOffset, values, valueReader);
+ } else {
+ nulls.putNulls(state.valueOffset, n);
+ }
+ state.valueOffset += n;
+ break;
+ case PACKED:
+ for (int i = 0; i < n; ++i) {
+ int currentValue = currentBuffer[currentBufferIdx++];
+ if (currentValue == state.maxDefinitionLevel) {
+ updater.readValue(state.valueOffset++, values, valueReader);
+ } else {
+ nulls.putNull(state.valueOffset++);
+ }
+ }
+ break;
+ }
+ state.levelOffset += n;
+ leftInBatch -= n;
+ rowId += n;
+ leftInPage -= n;
+ currentCount -= n;
+ }
+ }
+
+ state.rowsToReadInBatch = leftInBatch;
+ state.valuesToReadInPage = leftInPage;
+ state.rowId = rowId;
+ }
+
+ private void readBatchInternalWithDefLevels(
ParquetReadState state,
WritableColumnVector values,
WritableColumnVector nulls,
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@spark.apache.org
For additional commands, e-mail: commits-h...@spark.apache.org
