[SPARK-34859][SQL] parquet vectorized reader - support column index with rowIndexes

sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedColumnReader.java

@@ -21,7 +21,7 @@
 import java.math.BigInteger;
 import java.time.ZoneId;
 import java.time.ZoneOffset;
-import java.util.Arrays;
+import java.util.*;
 
 import org.apache.parquet.bytes.ByteBufferInputStream;
 import org.apache.parquet.bytes.BytesInput;
@@ -40,6 +40,7 @@
 import org.apache.spark.sql.catalyst.util.RebaseDateTime;
 import org.apache.spark.sql.execution.datasources.DataSourceUtils;
 import org.apache.spark.sql.execution.datasources.SchemaColumnConvertNotSupportedException;
+import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector;
 import org.apache.spark.sql.execution.vectorized.WritableColumnVector;
 import org.apache.spark.sql.types.DataType;
 import org.apache.spark.sql.types.DataTypes;
@@ -111,6 +112,12 @@ public class VectorizedColumnReader {
   private final String datetimeRebaseMode;
   private final String int96RebaseMode;
 
+  // TODO handle and init these filed properly
+  private Optional<PrimitiveIterator.OfLong> rowIndexesIterator;
+  private long[] rowIndexes;
+  private WritableColumnVector tempVector;
+  private Long currentRow;
+
   private boolean isDecimalTypeMatched(DataType dt) {
     DecimalType d = (DecimalType) dt;
     DecimalMetadata dm = descriptor.getPrimitiveType().getDecimalMetadata();
@@ -140,7 +147,10 @@ public VectorizedColumnReader(
       PageReader pageReader,
       ZoneId convertTz,
       String datetimeRebaseMode,
-      String int96RebaseMode) throws IOException {
+      String int96RebaseMode,
+      Optional<PrimitiveIterator.OfLong> rowIndexesIterator
+      ) throws IOException {
+    this.rowIndexesIterator = rowIndexesIterator;
     this.descriptor = descriptor;
     this.pageReader = pageReader;
     this.convertTz = convertTz;
@@ -248,7 +258,31 @@ static long rebaseInt96(long julianMicros, final boolean failIfRebase) {
   /**
    * Reads `total` values from this columnReader into column.
    */
-  void readBatch(int total, WritableColumnVector column) throws IOException {
+  void readBatch(int total, WritableColumnVector _column) throws IOException {
+    WritableColumnVector column;
+    if (rowIndexesIterator.isPresent()) {
+      if (tempVector == null) {
+        switch (descriptor.getPrimitiveType().getPrimitiveTypeName()) {
+          case INT64:
+            tempVector = new OnHeapColumnVector(4096, DataTypes.LongType);
+            break;
+          case BINARY:
+            tempVector = new OnHeapColumnVector(4096, DataTypes.BinaryType);
+            break;
+        }
+      }
+      column = tempVector;
+      column.reset();
+
+      rowIndexes = new long[total];
+      for (int i = 0; i < total; i++) {
+        rowIndexes[i] = rowIndexesIterator.get().next();
+      }
+    } else {
+      column = _column;
+    }
+
+
     int rowId = 0;
     WritableColumnVector dictionaryIds = null;
     if (dictionary != null) {
@@ -257,6 +291,7 @@ void readBatch(int total, WritableColumnVector column) throws IOException {
       // page.
       dictionaryIds = column.reserveDictionaryIds(total);
     }
+
     while (total > 0) {
       // Compute the number of values we want to read in this page.
       int leftInPage = (int) (endOfPageValueCount - valuesRead);
@@ -338,9 +373,52 @@ void readBatch(int total, WritableColumnVector column) throws IOException {
         }
       }
 
+      if (rowIndexesIterator.isPresent()) {
+        boolean continuousRange = (rowIndexes[total - 1] - rowIndexes[0] + 1) == total;
+        if (continuousRange) {
+          // skip to offset pos and dump all remaining values
+          int offset = (int) (rowIndexes[rowId] - currentRow);
+          if (offset < num) {
+            switch (typeName) {
+              case INT64:
+                _column.putLongs(rowId, num, column.getLongs(offset, num - offset), 0);
+                break;
+              case BINARY:
+                for (int i = 0; i < num - offset; i++) {
+                  _column.putByteArray(rowId + i, column.getBinary(i + offset));
+                }
+                break;
+            }
+            currentRow += num;
+            rowId += (num - offset);
+            total -= (num - offset);
+          } else {
+            currentRow += num;
+          }
+        } else {
+          // need to check every row
+          for (int i = 0; i < num; ) {
+            while (currentRow < rowIndexes[rowId]) {
+              i++;
+              currentRow++;
+            }
+            switch (typeName) {
+              case INT64:
+                _column.putLong(rowId, column.getLong(i));
+                break;
+              case BINARY:
+                _column.putByteArray(rowId, column.getBinary(i));
+            }
+            rowId++;
+            total--;
+          }
+        }
+      } else {
+        rowId += num;
+        total -= num;
+      }
+
       valuesRead += num;
-      rowId += num;
-      total -= num;
     }
   }
 
@@ -853,6 +931,7 @@ private void initDataReader(Encoding dataEncoding, ByteBufferInputStream in) thr
   }
 
   private void readPageV1(DataPageV1 page) throws IOException {
+    this.currentRow = page.getFirstRowIndex().orElse(0L);
     this.pageValueCount = page.getValueCount();
     ValuesReader rlReader = page.getRlEncoding().getValuesReader(descriptor, REPETITION_LEVEL);
     ValuesReader dlReader;
@@ -878,6 +957,7 @@ private void readPageV1(DataPageV1 page) throws IOException {
   }
 
   private void readPageV2(DataPageV2 page) throws IOException {
+    this.currentRow = page.getFirstRowIndex().orElse(0L);
     this.pageValueCount = page.getValueCount();
     this.repetitionLevelColumn = createRLEIterator(descriptor.getMaxRepetitionLevel(),
         page.getRepetitionLevels(), descriptor);
@@ -894,4 +974,5 @@ private void readPageV2(DataPageV2 page) throws IOException {
       throw new IOException("could not read page " + page + " in col " + descriptor, e);
     }
   }
+
 }

sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java

@@ -336,7 +336,8 @@ private void checkEndOfRowGroup() throws IOException {
         pages.getPageReader(columns.get(i)),
         convertTz,
         datetimeRebaseMode,
-        int96RebaseMode);
+        int96RebaseMode,
+        pages.getRowIndexes());
     }
     totalCountLoadedSoFar += pages.getRowCount();
   }

sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ParquetColumnIndexBenchmark.scala

@@ -0,0 +1,118 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.benchmark
+
+import java.io.File
+
+import scala.util.Random
+
+import org.apache.parquet.hadoop.ParquetInputFormat
+
+import org.apache.spark.SparkConf
+import org.apache.spark.benchmark.Benchmark
+import org.apache.spark.sql.{DataFrame, SparkSession}
+
+/**
+  * Benchmark to measure read performance with Parquet column index.
+  * To run this benchmark:
+  * {{{
+  *   1. without sbt: bin/spark-submit --class <this class> <spark sql test jar>
+  *   2. build/sbt "sql/test:runMain <this class>"
+  *   3. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "sql/test:runMain <this class>"
+  *      Results will be written to "benchmarks/ParquetFilterPushdownBenchmark-results.txt".
+  * }}}
+  */
+object ParquetColumnIndexBenchmark extends SqlBasedBenchmark {
+
+  override def getSparkSession: SparkSession = {
+    val conf = new SparkConf()
+      .setAppName(this.getClass.getSimpleName)
+      // Since `spark.master` always exists, overrides this value
+      .set("spark.master", "local[1]")
+      .setIfMissing("spark.driver.memory", "3g")
+      .setIfMissing("spark.executor.memory", "3g")
+      .setIfMissing("orc.compression", "snappy")
+      .setIfMissing("spark.sql.parquet.compression.codec", "snappy")
+
+    SparkSession.builder().config(conf).getOrCreate()
+  }
+
+  private val numRows = 1024 * 1024 * 15
+  private val width = 5
+  private val mid = numRows / 2
+
+  def withTempTable(tableNames: String*)(f: => Unit): Unit = {
+    try f finally tableNames.foreach(spark.catalog.dropTempView)
+  }
+
+  private def prepareTable(
+                            dir: File, numRows: Int): Unit = {
+    import spark.implicits._
+
+    val df = spark.range(numRows).map(i => (i, i + ":f" + "o" * Random.nextInt(200))).toDF()
+
+    saveAsTable(df, dir)
+  }
+
+  private def saveAsTable(df: DataFrame, dir: File, useDictionary: Boolean = false): Unit = {
+    val parquetPath = dir.getCanonicalPath + "/parquet"
+    df.write.mode("overwrite").parquet(parquetPath)
+    spark.read.parquet(parquetPath).createOrReplaceTempView("parquetTable")
+  }
+
+  def filterPushDownBenchmark(
+                               values: Int,
+                               title: String,
+                               whereExpr: String,
+                               selectExpr: String = "*"): Unit = {
+    val benchmark = new Benchmark(title, values, minNumIters = 5, output = output)
+
+    Seq(false, true).foreach { columnIndexEnabled =>
+      val name = s"Parquet Vectorized ${if (columnIndexEnabled) s"(columnIndex)" else ""}"
+      benchmark.addCase(name) { _ =>
+        withSQLConf(ParquetInputFormat.COLUMN_INDEX_FILTERING_ENABLED -> s"$columnIndexEnabled") {
+          spark.sql(s"SELECT $selectExpr FROM parquetTable WHERE $whereExpr").noop()
+        }
+      }
+    }
+
+    benchmark.run()
+  }
+
+  override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
+    runBenchmark("Pushdown for single value filter") {
+      withTempPath { dir =>
+        withTempTable("parquetTable") {
+          prepareTable(dir, numRows)
+          filterPushDownBenchmark(numRows, "simple filters", s" _1 = $numRows - 100 ")
+        }
+      }
+    }
+
+    runBenchmark("Pushdown for range filter") {
+      withTempPath { dir =>
+        withTempTable("parquetTable") {
+          prepareTable(dir, numRows)
+          filterPushDownBenchmark(numRows,
+            "range filters", s" _1 > ($numRows - 1000000) and _1 < ($numRows - 1000)")
+        }
+      }
+    }
+
+  }
+}