You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by su...@apache.org on 2022/05/04 16:38:22 UTC

[spark] branch branch-3.3 updated: [SPARK-38891][SQL] Skipping allocating vector for repetition & definition levels when possible

This is an automated email from the ASF dual-hosted git repository.

sunchao pushed a commit to branch branch-3.3
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-3.3 by this push:
     new 94d3d6b5fce [SPARK-38891][SQL] Skipping allocating vector for repetition & definition levels when possible
94d3d6b5fce is described below

commit 94d3d6b5fcea7866ba0a88f7ae5699f96857a7de
Author: Chao Sun <su...@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 <su...@apple.com>
    Signed-off-by: Chao Sun <su...@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-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org