peter-toth commented on code in PR #58895:
URL: https://github.com/apache/spark/pull/58895#discussion_r4145455892
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java:
##########
@@ -476,34 +506,80 @@ private boolean containsPath(Type parquetType, String[]
path, int depth) {
return false;
}
- private void checkEndOfRowGroup() throws IOException {
- if (rowsReturned != totalCountLoadedSoFar) return;
+ /**
+ * Loads the row group the read continues with, and returns false if there
is none left. Only a
+ * subclass that picks its own rows can answer false, because this class
reads every row of every
+ * row group, and then {@code totalRowCount} already says when they run out.
+ */
+ protected boolean loadNextRowGroup() throws IOException {
PageReadStore pages = reader.readNextRowGroup();
if (pages == null) {
throw new IOException("expecting more rows but reached last block. Read "
+ rowsReturned + " out of " + totalRowCount);
}
+ // Null ranges, because this path never picks the rows itself, and whether
parquet narrowed the
+ // row group by a pushed filter's column index is known only to the store
it handed back.
+ installRowGroup(pages, null, pages.getRowCount());
+ return true;
+ }
+
+ /**
+ * Points the column readers at one row group's {@code pages}, from which
they are to read the
+ * rows {@code rowRanges} names, or every row of the row group when it is
null. {@code rowCount}
+ * is how many rows that is, which is what the read has left before it needs
another row group.
+ */
+ protected final void installRowGroup(PageReadStore pages, RowRanges
rowRanges, long rowCount)
Review Comment:
Removed. `installRowGroup` reads `pages.getRowCount()` itself, as master
did, and the assertion from your store-count item is what keeps the subclass's
`finalRowCount` equal to it.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/VectorizedParquetRecordReader.java:
##########
@@ -80,7 +85,7 @@ public class VectorizedParquetRecordReader extends
SpecificParquetRecordReaderBa
* Encapsulate writable column vectors with other Parquet related info such
as
* repetition / definition levels.
*/
- private ParquetColumnVector[] columnVectors;
+ protected ParquetColumnVector[] columnVectors;
Review Comment:
Both done, and taken further. `keyVectorsFromBatch` reads
`persistentBatchColumns`, and `columnVectors` is private again. The
`resetBatch` and `readBatchColumns` hooks are gone as well as `finishBatch`, so
the base's `nextBatch` is master's bar the line that calls the load. The
subclass overrides `nextBatch` whole and points the key slots after the base
has read. `closeAllocated` went the same way, so the base's `close` is master's
too. The base class has six protected methods and four protected fields now.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -0,0 +1,2755 @@
+/*
+ * 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.datasources.parquet
+
+import java.io.{File, RandomAccessFile}
+import java.net.URI
+import java.time.LocalTime
+import java.util.concurrent.atomic.AtomicLong
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{FileStatus, FSDataInputStream, FSInputStream,
Path, RawLocalFileSystem}
+import org.apache.hadoop.mapred.FileSplit
+import org.apache.hadoop.mapreduce.{Job, TaskAttemptID}
+import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl
+import org.apache.parquet.column.{Encoding, ParquetProperties}
+import org.apache.parquet.column.impl.ColumnWriteStoreV1
+import org.apache.parquet.column.page.DataPageV1
+import org.apache.parquet.column.page.mem.MemPageStore
+import org.apache.parquet.format.converter.ParquetMetadataConverter
+import org.apache.parquet.hadoop.{ParquetFileWriter, ParquetInputFormat,
ParquetOutputFormat}
+import org.apache.parquet.hadoop.metadata.{ColumnChunkMetaData,
CompressionCodecName, ParquetMetadata}
+import org.apache.parquet.hadoop.util.{HadoopInputFile, HadoopOutputFile}
+import org.apache.parquet.schema.MessageTypeParser
+
+import org.apache.spark.SparkException
+import org.apache.spark.paths.SparkPath
+import org.apache.spark.sql.{sources, DataFrame, QueryTest, Row, SparkSession}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute,
AttributeReference, BloomFilterMightContain, BoundReference, Cast, EqualTo,
Expression, GreaterThanOrEqual, IsNull, LessThanOrEqual, Literal, Or,
Predicate, Rand, Remainder, SecondsToTimestamp, UnaryExpression, XxHash64}
+import org.apache.spark.sql.catalyst.expressions.aggregate.BloomFilterAggregate
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.catalyst.plans.logical.{Filter => LogicalFilter}
+import org.apache.spark.sql.execution.{CollapseCodegenStages,
ColumnarToRowExec, FileSourceScanExec, FilterExec, LocalLimitExec, SparkPlan,
WholeStageCodegenExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.{FileFormat,
FileSourceStrategy, OutputWriterFactory, PartitionedFile}
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.functions.col
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.vectorized.ColumnarBatch
+import org.apache.spark.util.Utils
+import org.apache.spark.util.sketch.BloomFilter
+
+/**
+ * Tests [[LateMaterializationParquetRecordReader]] driven by a
[[ParquetStorageFilter]]. Writes
+ * small multi-row-group parquet files, wires a hand-built filter into the
reader, and asserts
+ * correctness and the five storage-filter metrics.
+ */
+class ParquetStorageFilterSuite extends QueryTest with SharedSparkSession
+ with AdaptiveSparkPlanHelper {
+ import testImplicits._
+
+ // Writes `df` as one parquet file under a fresh directory and returns its
path. Every write
+ // helper in this suite goes through here.
+ private def writeSingleParquetFile(
+ dir: File,
+ df: DataFrame,
+ rowGroupSize: Long,
+ pageSize: Option[Long] = None,
+ dictionary: Boolean = false): String = {
+ val outDir = new File(dir, s"test-${System.nanoTime()}").getAbsolutePath
+ val writer = df
+ .repartition(1)
+ .write
+ .option(ParquetOutputFormat.BLOCK_SIZE, rowGroupSize)
+ // Dictionary encoding off keeps row-group sizing predictable. The
column index is still
+ // written either way.
+ .option(ParquetOutputFormat.ENABLE_DICTIONARY, dictionary.toString)
+ // A small page size gives each row group several pages per column, which
is what lets
+ // column-index filtering produce a row range narrower than the whole row
group.
+ pageSize.foreach(size => writer.option(ParquetOutputFormat.PAGE_SIZE,
size))
+ writer.parquet(outDir)
+ val files = new File(outDir).listFiles((_, name) =>
name.endsWith(".parquet"))
+ assert(files != null && files.length == 1, s"expected exactly one parquet
file under $outDir")
+ files(0).getAbsolutePath
+ }
+
+ // Writes a parquet file with the given rows and row-group size; returns the
path.
+ private def writeParquetFile(
+ dir: File,
+ rows: Seq[(Long, String)],
+ rowGroupSize: Long = 1024L,
+ pageSize: Option[Long] = None): String =
+ writeSingleParquetFile(dir, rows.toDF("k", "v"), rowGroupSize, pageSize)
+
+ // The scan node of a plan, which every test that inspects or rewrites
`storageFilters` needs.
+ private def scanOf(plan: SparkPlan): FileSourceScanExec =
+ plan.collect { case s: FileSourceScanExec => s }.headOption
+ .getOrElse(fail(s"No FileSourceScanExec found in plan: $plan"))
+
+ // A hadoop conf whose reads go through `CountingLocalFileSystem`, so a test
can measure them.
+ private def countingHadoopConf(): Configuration = {
+ val hadoopConf = spark.sessionState.newHadoopConf()
+ hadoopConf.set(s"fs.${CountingLocalFileSystem.scheme}.impl",
+ classOf[CountingLocalFileSystem].getName)
+
hadoopConf.setBoolean(s"fs.${CountingLocalFileSystem.scheme}.impl.disable.cache",
true)
+ hadoopConf
+ }
+
+ // Reads `path` whole through `readerFn`, and returns the rows it emitted
with the bytes the read
+ // transferred. The reader must have been built with `countingHadoopConf`.
+ private def readCounting(
+ path: String, readerFn: PartitionedFile => Iterator[_]): (Int, Long) = {
+ val file = PartitionedFile(
+ InternalRow.empty,
+ SparkPath.fromUrlString(s"${CountingLocalFileSystem.scheme}://$path"),
+ 0,
+ new File(path).length())
+ CountingLocalFileSystem.reset()
+ val emitted = readerFn(file).asInstanceOf[Iterator[Object]].map {
+ case batch: ColumnarBatch => batch.numRows()
+ case _ => 1
+ }.sum
+ (emitted, CountingLocalFileSystem.bytesRead())
+ }
+
+ // The metric map a scan hands the format, for a test that builds a reader
by hand.
+ private def metricMap(): Map[String, SQLMetric] =
+ new ParquetFileFormat().storageFilterMetrics(spark.sparkContext)
+
+ // Every metric the reader updates, created the way the production path
creates them. A test that
+ // asserts one holds the whole set and reads the counter off it.
+ private def allMetrics(): StorageFilterMetrics =
+ StorageFilterMetrics.create(spark.sparkContext)
+
+ // `ParquetStorageFilter.create` takes what the production path always
supplies; most tests care
+ // about neither the metrics nor the cap.
+ private def createFilter(
+ boundExpressions: Seq[Expression],
+ requestedSchema: StructType,
+ metrics: StorageFilterMetrics = allMetrics(),
+ maxSplicedRowGroupBytes: Long = Long.MaxValue): ParquetStorageFilter =
+ ParquetStorageFilter.create(
+ boundExpressions, requestedSchema, metrics, maxSplicedRowGroupBytes)
+
+ // Collects all `(k, v)` rows from a reader initialized with the given
storage filter.
+ private def readAll(
+ filePath: String,
+ storageFilter: ParquetStorageFilter): (Seq[(Long, String)],
VectorizedParquetRecordReader) =
+ readAllWith(filePath, Seq("k", "v"), storageFilter,
+ (batch, i) => (batch.column(0).getLong(i),
batch.column(1).getUTF8String(i).toString))
+
+ // Builds a `k >= threshold` storage filter bound to position 0, over a long
key.
+ private def keyAtLeastFilter(
+ threshold: Long,
+ metrics: StorageFilterMetrics = allMetrics()): ParquetStorageFilter =
+ keyTypeAtLeastFilter(Literal(threshold), LongType, metrics)
+
+ // A key column of the caller's type plus one non-key column, written via
Spark's `Encoder`. Every
+ // key-type fixture carries that second column, because a projection of key
columns alone is one
+ // the reader declines. It would read what a plain scan reads, so there is
nothing for phase 2 to
+ // prune.
+ private def writeKeyParquetFile[T : org.apache.spark.sql.Encoder](
+ dir: File,
+ keys: Seq[T],
+ rowGroupSize: Long = 1024L): String =
+ writeSingleParquetFile(
+ dir,
+ spark.createDataset(keys).toDF("k").selectExpr("k", "CAST(k AS STRING)
AS v"),
+ rowGroupSize)
+
+ // Reads a key-type fixture, returning the survivor keys and the reader. The
`extract` function
+ // pulls one value at a time from the batch's key column.
+ private def readKeysAll[T](
+ filePath: String,
+ storageFilter: ParquetStorageFilter,
+ extract: (org.apache.spark.sql.vectorized.ColumnVector, Int) => T,
+ capacity: Int = 4096): (Seq[T], VectorizedParquetRecordReader) =
+ readAllWith(filePath, Seq("k", "v"), storageFilter, (batch, i) =>
extract(batch.column(0), i),
+ capacity)
+
+ // Builds a `k >= threshold` storage filter bound to position 0 against a
`(k, v)` schema whose
+ // key has the given type.
+ private def keyTypeAtLeastFilter(
+ threshold: Literal,
+ keyType: DataType,
+ metrics: StorageFilterMetrics = allMetrics()): ParquetStorageFilter = {
+ val expr = GreaterThanOrEqual(BoundReference(0, keyType, nullable =
false), threshold)
+ val requested = StructType(Seq(
+ StructField("k", keyType, nullable = false), StructField("v",
StringType, nullable = true)))
+ createFilter(Seq(expr), requested, metrics)
+ }
+
+ test("rejects entire row group: no data-column IO, row-group-skipped metric
incremented") {
+ withTempDir { dir =>
+ // 40 rows come out as one row group. Parquet's first row-group size
check is at record 100,
+ // so `rowGroupSize` cannot split a file this small. The one row group
is the one rejected.
+ val rows = (1L to 40L).map(i => (i, s"v_$i"))
+ val path = writeParquetFile(dir, rows, rowGroupSize = 256L)
+
+ val m = allMetrics()
+ val filter = keyAtLeastFilter(1000L, m)
+ val (result, reader) = readAll(path, filter)
+ try {
+ assert(result.isEmpty, "filter rejects all rows; no rows should be
emitted")
+ assert(m.rowGroupsSkipped.value > 0,
+ s"expected at least one row group skipped; got
${m.rowGroupsSkipped.value}")
+ assert(m.rowsExcludedByRowGroup.value > 0,
+ s"expected rows excluded by whole-rowgroup skip; got
${m.rowsExcludedByRowGroup.value}")
+ assert(m.rowsExcludedWithinRowGroup.value == 0,
+ s"no partial-row-group filtering expected; got
${m.rowsExcludedWithinRowGroup.value}")
+ // Schema is (k: Long, v: String). Skipping a row group avoids the
v-column bytes the
+ // no-storage-filter path would have read; phase 1 still pays for k.
So avoided > 0.
+ assert(m.bytesAvoidedByRowGroup.value > 0,
+ s"expected non-key bytes avoided by whole row groups; got " +
+ s"${m.bytesAvoidedByRowGroup.value}")
+ assert(m.bytesAvoidedByPageFiltering.value == 0,
+ s"no page-filtering bytes expected when all groups skipped; got " +
+ s"${m.bytesAvoidedByPageFiltering.value}")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("all rows survive: no skipping and no filtering") {
+ withTempDir { dir =>
+ val rows = (1L to 40L).map(i => (i, s"v_$i"))
+ val path = writeParquetFile(dir, rows, rowGroupSize = 256L)
+
+ val m = allMetrics()
+ val filter = keyAtLeastFilter(0L, m)
+ val (result, reader) = readAll(path, filter)
+ try {
+ assert(result.toSet == rows.toSet, s"all rows should round-trip; got
${result.size} rows")
+ assert(m.rowGroupsSkipped.value == 0,
+ s"nothing should be skipped; got ${m.rowGroupsSkipped.value}")
+ assert(m.rowsExcludedWithinRowGroup.value == 0,
+ s"nothing should be filtered; got
${m.rowsExcludedWithinRowGroup.value}")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("mixed: some row groups skipped, others partially kept") {
+ withTempDir { dir =>
+ // Many rows + small row groups => guaranteed multiple row groups.
+ val rows = (1L to 200L).map(i => (i, s"v_$i"))
+ val path = writeParquetFile(dir, rows, rowGroupSize = 256L)
+
+ val m = allMetrics()
+ // k >= 195 keeps only the last 6 rows; earlier row groups should be
skipped.
+ val filter = keyAtLeastFilter(195L, m)
+ val (result, reader) = readAll(path, filter)
+ try {
+ // Output is exact. VectorizedColumnReader uses
PageReadStore.getRowIndexes (driven by
+ // our finalRanges) to skip rows within partial pages, so emitted rows
== survivors.
+ val expected = rows.filter(_._1 >= 195L).toSet
+ assert(result.toSet == expected,
+ s"expected exact filtering; got ${result.map(_._1).sorted}, " +
+ s"expected ${expected.map(_._1).toSeq.sorted}")
+ assert(m.rowGroupsSkipped.value >= 1,
+ s"expected row groups skipped; got ${m.rowGroupsSkipped.value}")
+ // "Partially kept" is the `rowsExcludedWithinRowGroup` half of the
accounting, and only
+ // this identity establishes it. A row group that was neither skipped
whole nor emitted has
+ // to have its rows counted there.
+ val excludedRg = m.rowsExcludedByRowGroup.value
+ val excludedPf = m.rowsExcludedWithinRowGroup.value
+ assert(excludedPf > 0, s"expected rows excluded inside a kept row
group; got $excludedPf")
+ assert(result.size + excludedRg + excludedPf == rows.size,
+ s"${result.size} emitted plus $excludedRg plus $excludedPf should
account for all " +
+ s"${rows.size} rows")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("an all-keys projection: the reader declines the filter and reads
plainly") {
+ // A scan whose projected columns are all key columns of the filter reads
exactly what a plain
+ // scan reads, because phase 1 has to read a key column to evaluate the
filter on it, and
+ // phase 2 is then left with nothing to prune. All the filter can add
there is cost, so the
+ // reader declines it rather than carrying a second emit path. The planner
declines that shape
+ // too, which is why this is asserted here, against the reader driven
directly.
+ withTempDir { dir =>
+ val keys = (1L to 200L)
+ val path = writeKeyParquetFile(dir, keys, rowGroupSize = 256L)
+ val requested = StructType(Seq(StructField("k", LongType, nullable =
false)))
+ val expr = GreaterThanOrEqual(BoundReference(0, LongType, nullable =
false), Literal(195L))
+ val metrics = allMetrics()
+ val filter = createFilter(Seq(expr), requested, metrics)
+ val (result, reader) = readAllWith(
+ path, Seq("k"), filter, (batch, i) => batch.column(0).getLong(i))
+ try {
+ assert(result == keys.toSeq,
+ s"every row of the file must come back unfiltered; got
${result.size} of ${keys.size}")
+ val reported = Seq(metrics.rowGroupsSkipped,
metrics.rowsExcludedByRowGroup,
+ metrics.rowsExcludedWithinRowGroup, metrics.bytesAvoidedByRowGroup,
+ metrics.bytesAvoidedByPageFiltering).map(_.value)
+ assert(reported.forall(_ == 0L), s"and the filter must report nothing:
$reported")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("Hadoop's 2-arg initialize initializes the reader once") {
+ // That overload delegates to the 5-arg one, which runs the reader's own
initialization, and it
+ // used to run that a second time on top. A reader that declined the
filter on the first pass,
+ // here because every projected column is a key column, then reached the
filter it had just
+ // released and failed with an NPE instead of reading plainly.
+ withTempDir { dir =>
+ val keys = (1L to 20L)
+ val path = writeKeyParquetFile(dir, keys)
+ val requested = StructType(Seq(StructField("k", LongType, nullable =
false)))
+ val filter = createFilter(
+ Seq(GreaterThanOrEqual(BoundReference(0, LongType, nullable = false),
Literal(15L))),
+ requested)
+ // What `ParquetFileFormat` sets up for the read, down to what the
schema converter needs.
+ val conf = spark.sessionState.newHadoopConf()
+ conf.set(ParquetInputFormat.READ_SUPPORT_CLASS,
classOf[ParquetReadSupport].getName)
+ conf.set(ParquetReadSupport.SPARK_ROW_REQUESTED_SCHEMA, requested.json)
+ Seq(SQLConf.CASE_SENSITIVE, SQLConf.PARQUET_BINARY_AS_STRING,
+ SQLConf.PARQUET_INT96_AS_TIMESTAMP,
SQLConf.PARQUET_INFER_TIMESTAMP_NTZ_ENABLED,
+ SQLConf.LEGACY_PARQUET_NANOS_AS_LONG).foreach { entry =>
+ conf.set(entry.key, spark.sessionState.conf.getConfString(entry.key))
+ }
+ val split = new FileSplit(new Path(path), 0, new File(path).length(),
Array.empty[String])
+ val context = new TaskAttemptContextImpl(conf, new TaskAttemptID())
+ Utils.tryWithResource(new LateMaterializationParquetRecordReader(false,
4096, filter)) {
+ reader =>
+ reader.initialize(split, context)
+ reader.initBatch(new StructType(), null)
+ val read = mutable.ArrayBuffer[Long]()
+ while (reader.nextBatch()) {
+ val batch = reader.resultBatch()
+ (0 until batch.numRows()).foreach(i => read +=
batch.column(0).getLong(i))
+ }
+ assert(read == keys, s"every row must come back unfiltered; got
${read.size}")
+ }
+ }
+ }
+
+ test("multi-batch emit: survivor count exceeds capacity") {
+ // Drive the reader at capacity = 16 with a row group of 71 survivors out
of 100. Exercises:
+ // - The per-key-column queue holding multiple full-capacity vectors
plus a partial tail.
+ // - The published queue head getting closed at the start of every
subsequent emit.
+ // - The batch's key slots being rewritten ceil(71/16) = 5 times.
+ // The filter keeps most of the row group rather than all of it, so a
reader that declined the
+ // filter altogether would return more rows than this asserts rather than
passing.
+ withTempDir { dir =>
+ val keys = (1L to 100L)
+ // Big rowGroupSize so all 100 rows fit in one row group.
+ val path = writeKeyParquetFile(dir, keys, rowGroupSize = 64 * 1024L)
+ val filter = keyTypeAtLeastFilter(Literal(30L), LongType)
+ val (result, reader) =
+ readKeysAll(path, filter, (vec, i) => vec.getLong(i), capacity = 16)
+ try {
+ assert(result == keys.filter(_ >= 30L).toSeq,
+ s"expected the surviving keys in order across multiple batches; got
${result.size} rows")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("int key column: filter survivors round-trip through phase 1
accumulators") {
+ // Covers the IntegerType branch of ValueCopier.
+ withTempDir { dir =>
+ val keys = (1 to 100)
+ val path = writeKeyParquetFile(dir, keys, rowGroupSize = 256L)
+ val filter = keyTypeAtLeastFilter(Literal(90), IntegerType)
+ val (result, reader) = readKeysAll(path, filter, (vec, i) =>
vec.getInt(i))
+ try {
+ assert(result.toSet == keys.filter(_ >= 90).toSet,
+ s"expected int keys >= 90; got ${result.sorted}")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("string key column: filter survivors round-trip through phase 1
accumulators") {
+ // Covers the StringType branch of ValueCopier (variable-length byte copy
via putByteArray).
+ withTempDir { dir =>
+ val keys = (1 to 20).map(i => f"k$i%03d")
+ val path = writeKeyParquetFile(dir, keys, rowGroupSize = 256L)
+ val filter = keyTypeAtLeastFilter(
+ Literal.create("k015", StringType), StringType)
+ val (result, reader) =
+ readKeysAll(path, filter, (vec, i) => vec.getUTF8String(i).toString)
+ try {
+ assert(result.toSet == keys.filter(_ >= "k015").toSet,
+ s"expected string keys >= 'k015'; got ${result.sorted}")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("ParquetStorageFilter.create rejects a filter that violates a planner
precondition") {
+ // These are all planner bugs by construction, because storageFiltersFor
pre-checks each one.
+ // The rows are safe either way, since the conjunct stays in the post-scan
Filter, but a soft
+ // rejection would hide the bug behind a silently slower read. create
fails instead.
+ val requested = StructType(Seq(
+ StructField("k", LongType, nullable = false),
+ StructField("v", StringType, nullable = false)))
+
+ // Nothing to push. The caller is supposed to check this before calling.
+ val empty = intercept[IllegalArgumentException] {
+ createFilter(Seq.empty, requested)
+ }
+ assert(empty.getMessage.contains("must be non-empty"), empty.getMessage)
+
+ // Ordinal 5 is out of range for a two-field requested schema.
+ val outOfRange = intercept[IllegalArgumentException] {
+ createFilter(
+ Seq(GreaterThanOrEqual(BoundReference(5, LongType, nullable = false),
Literal(0L))),
+ requested)
+ }
+ assert(outOfRange.getMessage.contains("outside the 2 fields"),
outOfRange.getMessage)
+
+ // No bound reference at all, so there is no key column to read in phase 1.
+ val noRefs = intercept[IllegalArgumentException] {
+ createFilter(Seq(GreaterThanOrEqual(Literal(1L), Literal(0L))),
requested)
+ }
+ assert(noRefs.getMessage.contains("no bound reference"), noRefs.getMessage)
+
+ // A key type the reader has no value copier for.
+ val variantSchema = StructType(Seq(StructField("k", VariantType, nullable
= true)))
+ val badType = intercept[IllegalArgumentException] {
+ createFilter(
+ Seq(IsNull(BoundReference(0, VariantType, nullable = true))),
variantSchema)
+ }
+ assert(badType.getMessage.contains("isSupportedKeyType"),
badType.getMessage)
+
+ // A key column named like the synthetic row-index metadata column. This
reader finds that
+ // column by name and writes row indexes over whatever the file holds, so
what a key column of
+ // that name reads back would depend on how its row group was read. Both
gates reject it. The
+ // planner's means the filter is never offered, and this one means a
caller past the planner
+ // fails rather than reading a key the reader itself overwrote.
+ val rowIndexName = ParquetFileFormat.ROW_INDEX_TEMPORARY_COLUMN_NAME
+ val rowIndexSchema = StructType(Seq(
+ StructField(rowIndexName, LongType, nullable = false),
+ StructField("v", StringType, nullable = false)))
+ val rowIndexKey = intercept[IllegalArgumentException] {
+ createFilter(
+ Seq(GreaterThanOrEqual(BoundReference(0, LongType, nullable = false),
Literal(0L))),
+ rowIndexSchema)
+ }
+ assert(rowIndexKey.getMessage.contains(rowIndexName),
rowIndexKey.getMessage)
+ val rowIndexAttr = AttributeReference(rowIndexName, LongType, nullable =
false)()
+ assert(!ParquetStorageFilter.isSupportedStorageFilter(
+ BloomFilterMightContain(
+ Literal(bloomBytes(BloomFilter.create(10))),
+ new XxHash64(Seq(rowIndexAttr)))),
+ "a bloom over a column named like the row-index column must not be
offered")
+ }
+
+ // Serializes a [[BloomFilter]] to bytes suitable for a [[Literal]].
+ // The same serialization a runtime bloom's Literal carries, so these tests
build theirs the way
+ // the planner does.
+ private def bloomBytes(bf: BloomFilter): Array[Byte] =
BloomFilterAggregate.serialize(bf)
+
+ // Computes `XxHash64(v)` using the same seed Spark uses for runtime bloom
filters.
+ private def xxHash64(v: Any, dt: DataType): Long = {
+ new XxHash64(Seq(Literal(v,
dt))).eval(InternalRow.empty).asInstanceOf[Long]
+ }
+
+ test("a plain JDK exception from the key expression gives the row group up")
{
+ // The reader's fail-open is not narrowed to Spark's own errors, and this
is why. A built-in
+ // expression that `InjectRuntimeFilter` is happy to build a bloom over
can raise a plain JDK
+ // one. Rethrowing it would fail a query a plain read answers, and under
`ignoreCorruptFiles`
+ // an exception from a reader is read as a corrupt file, which drops the
rest of it silently.
+ //
+ // The filter is selective, so returning every row is what proves the
give-up. A filter that
+ // ran to the end would have dropped the rows below 10 seconds.
+ withTempDir { dir =>
+ val df = spark.range(0, 20).selectExpr(
+ "CAST(CASE WHEN id = 7 THEN 1.0000001 ELSE id END AS DECIMAL(20,7)) AS
k",
+ "CAST(id AS STRING) AS v")
+ val path = writeSingleParquetFile(dir, df, rowGroupSize = 64 * 1024L)
+ val decimalType = DecimalType(20, 7)
+ val requested = StructType(Seq(
+ StructField("k", decimalType, nullable = true),
+ StructField("v", StringType, nullable = true)))
+ // Row 7 throws, since 1.0000001 seconds is not a whole number of
microseconds.
+ val secondsKey = GreaterThanOrEqual(
+ SecondsToTimestamp(BoundReference(0, decimalType, nullable = true)),
+ Literal.create(10L * 1000000L, TimestampType))
+ val filter = createFilter(Seq(secondsKey), requested)
+ val (result, reader) = readAllWith(
+ path, Seq("k", "v"), filter, (b, i) =>
b.column(1).getUTF8String(i).toString)
+ try {
+ assert(result == (0 until 20).map(_.toString),
+ s"the row group must come back whole; got ${result.size} rows:
$result")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("a checked exception from the key expression gives the row group up") {
+ // Nor is the fail-open narrowed to unchecked exceptions. A Hive UDF in
the key wraps what it
+ // throws in a checked SparkException (FAILED_EXECUTE_UDF), which no
signature declares, so a
+ // catch of RuntimeException alone would fail a query a plain read
answers. The filter is
+ // selective, so returning every row is what proves the give-up.
+ withTempDir { dir =>
+ val rows = (0L until 20L).map(i => (i, i.toString))
+ val path = writeParquetFile(dir, rows, rowGroupSize = 64 * 1024L)
+ val requested = StructType(Seq(
+ StructField("k", LongType, nullable = false),
+ StructField("v", StringType, nullable = true)))
+ val key = ThrowsCheckedOn(BoundReference(0, LongType, nullable = false),
bad = 7L)
+ val filter = createFilter(Seq(GreaterThanOrEqual(key, Literal(10L))),
requested)
+ val (result, reader) = readAll(path, filter)
+ try {
+ assert(result == rows, s"the row group must come back whole; got
${result.size} rows")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ // Returns the FileSourceScanExec for a parquet read, with `storageFilters`
attached. Goes through
+ // Spark's normal planning/execution machinery (not the test-only reader
init), so it exercises
+ // preparedStorageFilters' subquery materialization + bind, the SQL conf
check,
+ // ParquetFileFormat.buildReaderWithStorageFilters, and metric propagation.
+ private def scanWithStorageFilter(
+ path: String,
+ keyName: String,
+ threshold: Long): FileSourceScanExec = {
+ val df = spark.read.parquet(path).select("k", "v")
+ val plan = df.queryExecution.executedPlan
+ val scan = scanOf(plan)
+ val keyAttr = scan.output.find(_.name == keyName).getOrElse(fail(s"No
$keyName in scan output"))
+ val expr = GreaterThanOrEqual(keyAttr, Literal(threshold))
+ scan.copy(storageFilters = Seq(expr))
+ }
+
+ // Executes a SparkPlan that may produce columnar batches. When the plan
supports columnar output
+ // (typical for parquet scans with WSCG enabled), Spark's planner normally
inserts a
+ // ColumnarToRowExec; since these tests bypass the planner, we wrap manually.
+ private def executePlanCollect(plan: SparkPlan): Array[(Long, String)] = {
+ val rowPlan = if (plan.supportsColumnar) ColumnarToRowExec(plan) else plan
+ rowPlan.executeCollect().map(r => (r.getLong(0), r.getString(1)))
+ }
+
+ test("end-to-end via FileSourceScanExec: the filter is applied and the
metrics populated") {
+ withTempDir { dir =>
+ val rows = (1L to 200L).map(i => (i, s"v_$i"))
+ val path = writeParquetFile(dir, rows, rowGroupSize = 256L)
+
+ withSQLConf(SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED.key ->
"true") {
+ val scan = scanWithStorageFilter(path, "k", threshold = 195L)
+ val collected = executePlanCollect(scan).toSet
+ val expected = rows.filter(_._1 >= 195L).toSet
+ assert(collected == expected, s"got ${collected.toSeq.sortBy(_._1)};
expected $expected")
+
+ val rgSkipped = scan.metrics(StorageFilterMetrics.ROW_GROUPS_SKIPPED)
+ assert(rgSkipped.value >= 1,
+ s"expected at least one row group skipped via storage filter; got
${rgSkipped.value}")
+ }
+ }
+ }
+
+ test("pushed data filter on a non-key column + storage filter on key:
cross-propagation works") {
+ // A pushed data filter on a non-key column and a storage filter on the
key column have to
+ // compose. Phase 0 derives its row ranges from the pushed filter, phase 1
narrows them with the
+ // storage filter, and phase 2 reads the non-key columns under the
intersection.
+ //
+ // Note the predicate deliberately uses `>` rather than `!=`.
ColumnIndexFilter substitutes
+ // `rangesForMissingColumns` for a predicate over a column outside its
path set, and that is
+ // EMPTY for Gt/GtEq/Lt/LtEq/Eq but allRows for NotEq, so a `!=` predicate
here would be
+ // satisfied by a phase 0 that saw the wrong schema, and would prove
nothing.
+ withTempDir { dir =>
+ val rows = (1L to 200L).map(i => (i, f"v_$i%03d"))
+ val path = writeParquetFile(dir, rows, rowGroupSize = 256L)
+
+ withSQLConf(SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED.key ->
"true") {
+ val df = spark.read.parquet(path).select("k", "v").filter("v >
'v_000'")
+ val plan = df.queryExecution.executedPlan
+ val scan = scanOf(plan)
+ assert(scan.simpleString(200).contains("GreaterThan(v,"),
+ s"the data filter must actually be pushed for this test to mean
anything: " +
+ scan.simpleString(200))
+ // Storage filter on `k` (key column).
+ val keyAttr = scan.output.find(_.name == "k").get
+ val withSF = scan.copy(
+ storageFilters = Seq(GreaterThanOrEqual(keyAttr, Literal(100L))))
+
+ val collected = executePlanCollect(withSF).toSet
+ // Every row satisfies v > 'v_000', so the storage filter alone
decides the result.
+ val expected = rows.filter(_._1 >= 100L).toSet
+ assert(collected == expected,
+ s"got ${collected.toSeq.sortBy(_._1)}; expected
${expected.toSeq.sortBy(_._1)}")
+ }
+ }
+ }
+
+ // ----- FileSourceStrategy bloom-filter extraction -----
+
+ // Counts BloomFilterMightContain expressions inside FilterExec nodes of a
physical plan.
+ // `collect` from AdaptiveSparkPlanHelper rather than SparkPlan's, which
stops at an
+ // `AdaptiveSparkPlanExec` and would report zero for every plan AQE wrapped.
+ private def countBloomFiltersInPostScanFilters(plan: SparkPlan): Int = {
+ collect(plan) {
+ case f: FilterExec =>
+ f.condition.collect { case _: BloomFilterMightContain => 1 }.sum
+ }.sum
+ }
+
+ // Counts BloomFilterMightContain expressions inside
FileSourceScanExec.storageFilters.
+ private def countBloomFiltersInStorageFilters(plan: SparkPlan): Int = {
+ collect(plan) {
+ case s: FileSourceScanExec =>
+ s.storageFilters.map(_.collect { case _: BloomFilterMightContain => 1
}.sum).sum
+ }.sum
+ }
+
+ // Sets up two parquet tables and runs a join that triggers
`InjectRuntimeFilter` for the
+ // application-side scan. Returns the executed plan and the query result for
inspection. Tables
+ // are cleaned up automatically by the calling test (table names are passed
through).
+ private def runBloomFilterJoin(): (SparkPlan, Array[Row]) = {
+ val query =
+ """SELECT bf1.k, bf1.v
+ |FROM bf1 JOIN bf2 ON bf1.k = bf2.k
+ |WHERE bf2.v = 5
+ |""".stripMargin
+ val df = spark.sql(query)
+ (df.queryExecution.executedPlan, df.collect())
+ }
+
+ // Creates two parquet tables for join tests, a large `bf1` and a small
`bf2` with a selective
+ // filter, and sets what it takes for the optimizer to build a bloom over
them at all. The
+ // application side must look big enough to be worth one, and the join must
not become a
+ // broadcast. Every test below needs both, so they are here rather than in
each, and a test states
+ // only what it varies.
+ private def withBloomFilterTables(body: => Unit): Unit = {
+ withTable("bf1", "bf2") {
+ // bf1 is the large application side, 600 rows.
+ spark.range(600).selectExpr("id AS k", "id AS
v").write.format("parquet").saveAsTable("bf1")
+ // bf2 is the small creation side, 30 rows, with a selective filter (v =
5 keeps 1 row).
+ spark.range(30).selectExpr("id AS k", "id AS
v").write.format("parquet").saveAsTable("bf2")
+ withSQLConf(
+
SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "1000",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "200") {
+ body
+ }
+ }
+ }
+
+ test("FileSourceStrategy offers the bloom to the scan and keeps it post-scan
too") {
+ // The scan gets the conjunct to prune with, and the post-scan Filter
keeps it, so the answer
+ // never depends on what the reader managed to do with it.
+ withBloomFilterTables {
+ withSQLConf(
+ SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val (plan, _) = runBloomFilterJoin()
+ val storageBlooms = countBloomFiltersInStorageFilters(plan)
+ val postScanBlooms = countBloomFiltersInPostScanFilters(plan)
+ assert(storageBlooms >= 1,
+ s"expected >= 1 bloom filter on scan.storageFilters; got
$storageBlooms.\n" +
+ s"Plan:\n$plan")
+ assert(postScanBlooms >= 1,
+ s"expected the bloom to stay in a post-scan FilterExec; got
$postScanBlooms.\n" +
+ s"Plan:\n$plan")
+ }
+ }
+ }
+
+ test("FileSourceStrategy leaves the bloom behind when every projected column
is a key column") {
+ // Nothing is left for phase 2 to prune. The reader would read the same
column for the same rows
+ // as a plain scan, since it has to read a key column to evaluate the
filter on it, and would
+ // add only the cost of evaluating the predicate outside the generated
code.
+ withBloomFilterTables {
+ withSQLConf(
+ SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = spark.sql("SELECT bf1.k FROM bf1 JOIN bf2 ON bf1.k = bf2.k
WHERE bf2.v = 5")
+ val plan = df.queryExecution.executedPlan
+ val storageBlooms = countBloomFiltersInStorageFilters(plan)
+ val postScanBlooms = countBloomFiltersInPostScanFilters(plan)
+ assert(storageBlooms == 0,
+ s"expected no bloom on scan.storageFilters for an all-keys
projection; got " +
+ s"$storageBlooms.\nPlan:\n$plan")
+ assert(postScanBlooms >= 1,
+ s"expected the bloom to stay in a post-scan FilterExec; got
$postScanBlooms.\n" +
+ s"Plan:\n$plan")
+ assert(df.collect().map(_.getLong(0)).toSet == Set(5L),
+ "and the query must still return the joined key")
+ }
+ }
+ }
+
+ test("FileSourceStrategy leaves bloom filter as post-scan FilterExec when
conf is off") {
+ withBloomFilterTables {
+ withSQLConf(
+ SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "false",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val (plan, _) = runBloomFilterJoin()
+ val storageBlooms = countBloomFiltersInStorageFilters(plan)
+ val postScanBlooms = countBloomFiltersInPostScanFilters(plan)
+ assert(storageBlooms == 0,
+ s"expected no bloom on scan.storageFilters when conf is off; got
$storageBlooms")
+ assert(postScanBlooms >= 1,
+ s"expected bloom in post-scan FilterExec when conf is off; got
$postScanBlooms.\n" +
+ s"Plan:\n$plan")
+ }
+ }
+ }
+
+ test("the feature still engages when ignoreCorruptFiles is on") {
+ withBloomFilterTables {
+ withSQLConf(
+ SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true",
+ SQLConf.IGNORE_CORRUPT_FILES.key -> "true",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val (plan, rows) = runBloomFilterJoin()
+ val storageBlooms = countBloomFiltersInStorageFilters(plan)
+ assert(storageBlooms == 1,
+ s"expected the bloom on scan.storageFilters; got
$storageBlooms.\nPlan:\n$plan")
+ assert(countBloomFiltersInPostScanFilters(plan) >= 1,
+ s"and the post-scan Filter keeps it.\nPlan:\n$plan")
+ assert(rows.map(r => (r.getLong(0), r.getLong(1))).toSet == Set((5L,
5L)),
+ s"and the query must still return the joined row; got
${rows.mkString(", ")}")
+ // The three assertions above hold whether or not the reader honored
the filter, since the
+ // post-scan Filter answers the query either way. The metrics are what
say it did.
+ val scan = plan.collect { case s: FileSourceScanExec => s }
+ .find(_.storageFilters.nonEmpty).getOrElse(fail(s"no scan with
storage filters:\n$plan"))
+ val excluded =
+ scan.metrics(StorageFilterMetrics.ROWS_EXCLUDED_BY_ROW_GROUP).value +
+
scan.metrics(StorageFilterMetrics.ROWS_EXCLUDED_WITHIN_ROW_GROUP).value
+ assert(excluded > 0, s"the reader must have excluded rows; the metrics
say $excluded")
+ }
+ }
+ }
+
+ test("a row group over the splice cap is read the plain way, same rows, more
bytes") {
+ // Past the cap a row group is read the plain way. Phase 2 takes every
projected column, key
+ // columns included, so nothing is buffered. Two things have to hold. The
rows must not
+ // change, since a fallback that quietly dropped the predicate would
return extra rows. And
+ // the fallback must actually have happened, which is observable. It reads
the key column a
+ // second time, so it transfers strictly more bytes. Without that second
assertion the test
+ // would pass even if the cap never reached the reader.
+ //
+ // The batch size is not what makes the cap reachable, since the count is
examined after
+ // every surviving row. It is here because a row group read the plain way
past the cap still
+ // has to emit in batches, and 51 survivors over a capacity of 16 span
several of them.
+ withTempDir { dir =>
+ val rows = (1L to 400L).map(i => (i, f"v_$i%04d"))
+ val path = writeParquetFile(dir, rows, rowGroupSize = 64 * 1024L,
pageSize = Some(512L))
+ val fileSchema = StructType(Seq(
+ StructField("k", LongType, nullable = true),
+ StructField("v", StringType, nullable = true)))
+ val storageFilters =
+ Seq(GreaterThanOrEqual(BoundReference(0, LongType, nullable = true),
Literal(350L)))
+
+ def run(maxSplicedBytes: String): (Int, Long) = {
+ withSQLConf(
+
SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_MAX_SPLICED_ROW_GROUP_BYTES.key ->
+ maxSplicedBytes,
+ SQLConf.PARQUET_VECTORIZED_READER_BATCH_SIZE.key -> "16") {
+ val readerFn = new ParquetFileFormat().buildReaderWithStorageFilters(
+ spark, fileSchema, new StructType(), fileSchema, Nil,
storageFilters,
+ Map(FileFormat.OPTION_RETURNING_BATCH -> "true"),
countingHadoopConf(), metricMap())
+ .getOrElse(fail("ParquetFileFormat must answer with a reader"))
+ readCounting(path, readerFn)
+ }
+ }
+
+ val (splicedRows, splicedBytes) = run("64MB")
+ // 100 bytes is past what 16 survivors of a long key buffer, and still
leaves room for the
+ // row ranges. `k >= 350` keeps a contiguous run, so there is one range
to hold.
+ val (plainRows, plainBytes) = run("100b")
+ assert(splicedRows == 51, s"the filter keeps keys 350..400; got
$splicedRows")
+ assert(plainRows == splicedRows,
+ s"rows differ past the cap: plain=$plainRows spliced=$splicedRows")
+ assert(plainBytes > splicedBytes,
+ s"past the cap the key column is read twice, so the read must be
larger; " +
+ s"plain=$plainBytes spliced=$splicedBytes")
+ }
+ }
+
+ test("row groups that splice and row groups over the cap, in both orders") {
+ // The emitted batch is one object for the whole read, so its key slots
have to follow the path
+ // each row group took. This file makes all four cases occur in order, at
100 rows per row group
+ // and a batch capacity of 16:
+ // - rows 1-100: 8 survivors, which buffer 72 bytes and fall in one range
of 40, inside the
+ // 150-byte cap, so the row group splices;
+ // - rows 101-200: 51 survivors, so the buffer passes the cap and phase 2
reads every projected
+ // column. The key slots must go back to the persistent vectors here,
or the batch reads keys
+ // out of a vector the previous row group's last batch already released;
+ // - rows 201-300: 6 survivors, splicing again, which needs the
accumulators that giving
+ // splicing up released to be allocated afresh;
+ // - rows 301-400: no survivor at all, so the row group is skipped whole.
+ //
+ // Hence the assertion on the values rather than on the row count alone,
and the skipped row
+ // group metric, which pins that the file really did split into several
row groups.
+ withTempDir { dir =>
+ val rows = (1L to 400L).map(i => (i, f"v_$i%04d"))
+ val path = writeParquetFile(dir, rows, rowGroupSize = 256L)
+ val m = allMetrics()
+ val k = BoundReference(0, LongType, nullable = true)
+ def between(lo: Long, hi: Long): Expression =
+ And(GreaterThanOrEqual(k, Literal(lo)), LessThanOrEqual(k,
Literal(hi)))
+ val filter = createFilter(
+ Seq(Or(Or(LessThanOrEqual(k, Literal(8L)), between(150L, 200L)),
between(250L, 255L))),
+ StructType(Seq(
+ StructField("k", LongType, nullable = true),
+ StructField("v", StringType, nullable = true))),
+ m,
+ maxSplicedRowGroupBytes = 150L)
+ val (result, reader) = readAllWith(path, Seq("k", "v"), filter,
+ (batch, i) => (batch.column(0).getLong(i),
batch.column(1).getUTF8String(i).toString),
+ capacity = 16)
+ try {
+ val expected = rows.filter { case (key, _) =>
+ key <= 8L || (key >= 150L && key <= 200L) || (key >= 250L && key <=
255L)
+ }
+ assert(result == expected,
+ s"expected each surviving key with its own value; got
${result.take(20)}")
+ assert(m.rowGroupsSkipped.value >= 1,
+ s"expected a row group with no survivor at all; got
${m.rowGroupsSkipped.value}")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("a row group past the cap charges its second key read against the byte
metric") {
+ // Giving splicing up means phase 2 reads the key columns a second time,
while the baseline
+ // counts them once, in phase 1. That extra read is a cost against the
saving, so the same file
+ // and filter must report less avoided than they do while splicing.
Without that term the metric
+ // would credit the feature with bytes it did transfer.
+ withTempDir { dir =>
+ val rows = (1L to 400L).map(i => (i, f"v_$i%04d"))
+ val path = writeParquetFile(dir, rows, rowGroupSize = 64 * 1024L,
pageSize = Some(512L))
+ val fileSchema = StructType(Seq(
+ StructField("k", LongType, nullable = true),
+ StructField("v", StringType, nullable = true)))
+
+ def avoidedBytes(cap: Long, threshold: Long): Long = {
+ val m = allMetrics()
+ val filter = createFilter(
+ Seq(GreaterThanOrEqual(BoundReference(0, LongType, nullable = true),
Literal(threshold))),
+ fileSchema, m, cap)
+ val (_, reader) = readAllWith(path, Seq("k", "v"), filter,
+ (b, i) => b.column(0).getLong(i), capacity = 16)
+ try m.bytesAvoidedByPageFiltering.value finally reader.close()
+ }
+
+ val spliced = avoidedBytes(64L * 1024 * 1024, 350L)
+ val plain = avoidedBytes(100L, 350L)
+ assert(spliced > 0, s"page filtering has to avoid something here; got
$spliced")
+ assert(plain < spliced,
+ s"the second key read must count against the saving; plain=$plain
spliced=$spliced")
+
+ // The two halves of the budget are weighed together, and this is the
case that says so. Eight
+ // survivors of a long key buffer 8 * (1 + 8) = 72 bytes and fall in one
range of 40, so
+ // neither half reaches the 100-byte cap on its own while the sum passes
it at the seventh.
+ // Weighing them separately would keep splicing here, and report the
larger saving for it.
+ val eightSpliced = avoidedBytes(64L * 1024 * 1024, 393L)
+ val eightPlain = avoidedBytes(100L, 393L)
+ assert(eightSpliced > 0, s"and avoid something with eight survivors; got
$eightSpliced")
+ assert(eightPlain < eightSpliced,
+ s"the sum of the two halves must cross the cap; plain=$eightPlain
spliced=$eightSpliced")
+ }
+ }
+
+ // Writes a parquet file the low-level way, with no offset index for any
column.
+ private def writeParquetFileWithoutOffsetIndex(
+ dir: File,
+ blocks: Seq[Seq[(Long, Long)]]): String =
+ writeParquetFileByHand(dir, blocks, (_, _) => false)
+
+ // Writes a parquet file the low-level way, with an offset index for the
chunks `hasOffsetIndex`
+ // names by block and column. The `ParquetFileWriter.writeDataPage`
overloads that take no row
+ // count use parquet's no-op offset index builder, which is how a writer
other than parquet-mr's
+ // own produces a chunk that cannot be read in part. Every column is a
required int64, so one page
+ // writer serves all.
+ private def writeParquetFileByHand(
+ dir: File,
+ blocks: Seq[Seq[(Long, Long)]],
+ hasOffsetIndex: (Int, Int) => Boolean): String = {
+ val schema = MessageTypeParser.parseMessageType(
+ "message spark_schema { required int64 k; required int64 v; }")
+ val file = new File(dir, s"no-offset-index-${System.nanoTime()}.parquet")
+ val hadoopPath = new Path(file.getAbsolutePath)
+ val writer = new ParquetFileWriter(
+ HadoopOutputFile.fromPath(hadoopPath,
spark.sessionState.newHadoopConf()),
+ schema, ParquetFileWriter.Mode.CREATE, 128L * 1024 * 1024, 8)
+ writer.start()
+ blocks.zipWithIndex.foreach { case (block, blockIdx) =>
+ writer.startBlock(block.size)
+ schema.getColumns.asScala.zipWithIndex.foreach { case (cd, colIdx) =>
+ val pageStore = new MemPageStore(block.size)
+ val writeStore = new ColumnWriteStoreV1(pageStore,
+ // The row-count check is what decides when a page is flushed, so
without lowering it a
+ // block of a few hundred rows comes out as one page and no column
index can narrow it.
+
ParquetProperties.builder().withPageSize(256).withMinRowCountForPageSizeCheck(1)
+ .withDictionaryEncoding(false).build())
+ val columnWriter = writeStore.getColumnWriter(cd)
+ block.foreach { row =>
+ columnWriter.write(if (colIdx == 0) row._1 else row._2, 0, 0)
+ writeStore.endRecord()
+ }
+ writeStore.flush()
+ writer.startColumn(cd, block.size, CompressionCodecName.UNCOMPRESSED)
+ val pageReader = pageStore.getPageReader(cd)
+ var written = 0L
+ while (written < block.size) {
+ val page = pageReader.readPage().asInstanceOf[DataPageV1]
+ if (hasOffsetIndex(blockIdx, colIdx)) {
+ writer.writeDataPage(page.getValueCount, page.getUncompressedSize,
page.getBytes,
+ page.getStatistics, page.getValueCount.toLong,
page.getRlEncoding,
+ page.getDlEncoding, page.getValueEncoding)
+ } else {
+ writer.writeDataPage(page.getValueCount, page.getUncompressedSize,
page.getBytes,
+ page.getStatistics, page.getRlEncoding, page.getDlEncoding,
page.getValueEncoding)
+ }
+ written += page.getValueCount
+ }
+ writer.endColumn()
+ }
+ writer.endBlock()
+ }
+ writer.end(new java.util.HashMap[String, String]())
+ file.getAbsolutePath
+ }
+
+ test("a file written with no offset index is read with the filter given up")
{
+ // Reading part of a row group needs an offset index, so a file without
one cannot be filtered
+ // at page level. What it can still do is skip a row group the filter
empties, which needs no
+ // index at all, and that is the split this test pins. The three row
groups are ordered so that
+ // the file's first partial read is what discovers the missing index, and
the two row groups
+ // after that discovery still get phase 1. The middle one is emptied by
the filter and skipped
+ // whole, the last keeps rows and is read whole. The missing index is
learned from the read that
+ // fails, so the first row group is the one that discovers it. Spark's own
writer always writes
+ // the page index, hence the hand-built file.
+ withTempDir { dir =>
+ val blocks = Seq(
+ (1L to 100L).map(i => (i, i * 10)),
+ (101L to 200L).map(i => (i, i * 10)),
+ (201L to 300L).map(i => (i, i * 10)))
+ val path = writeParquetFileWithoutOffsetIndex(dir, blocks)
+ // The fixture's whole point, asserted rather than assumed.
+ assert(footerChunks(path).forall(_.getOffsetIndexReference == null),
+ "no chunk of the hand-built file may have an offset index")
+ val m = allMetrics()
+ val requested = StructType(Seq(
+ StructField("k", LongType, nullable = false),
+ StructField("v", LongType, nullable = false)))
+ // Keeps part of the first row group, none of the second, part of the
third.
+ val k = BoundReference(0, LongType, nullable = false)
+ val filter = createFilter(
+ Seq(Or(
+ And(GreaterThanOrEqual(k, Literal(50L)), LessThanOrEqual(k,
Literal(60L))),
+ And(GreaterThanOrEqual(k, Literal(250L)), LessThanOrEqual(k,
Literal(260L))))),
+ requested,
+ m)
+ // `withLogAppender` returns Unit, so the read's result is captured out
of the block.
+ val logAppender = new LogAppender("no offset index")
+ var read: (Seq[(Long, Long)], VectorizedParquetRecordReader) = null
+ withLogAppender(logAppender) {
+ read = readAllWith(path, Seq("k", "v"), filter,
+ (b, i) => (b.column(0).getLong(i), b.column(1).getLong(i)))
+ }
+ val (result, reader) = read
+ try {
+ assert(result == blocks(0) ++ blocks(2),
+ s"the emptied row group is skipped and the other two read whole; got
${result.size} rows")
+ // The footer says which columns have an offset index, so the reader
knows before phase 1
+ // buffers anything, and the phase-2 read never asks parquet for a
narrowed one. Without
+ // that, the first partially surviving row group copies its survivors
and the exception path
+ // throws them away, which is what this WARN reports.
+ assert(!logAppender.loggingEvents.map(_.getMessage.getFormattedMessage)
+ .exists(_.contains("without page-level storage filtering")),
+ "the footer check must get there before the exception path")
+ assert(m.rowGroupsSkipped.value == 1 && m.rowsExcludedByRowGroup.value
== 100,
+ s"one row group skipped whole, with its rows counted; got " +
+ s"${m.rowGroupsSkipped.value} and
${m.rowsExcludedByRowGroup.value}")
+ assert(m.rowsExcludedWithinRowGroup.value == 0,
+ s"and nothing excluded inside a row group, which needs the index;
got " +
+ s"${m.rowsExcludedWithinRowGroup.value}")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("a key column with no offset index does not cost the file its page
filtering") {
+ // Reading part of a row group goes through the block's column index
store, and parquet builds
+ // that store over every column of the requested schema, emptying it
altogether if one of them
+ // has no offset index. So a row group whose key column has none cannot be
narrowed either, and
+ // the footer check has to ask about the whole projection. Asking only
about the columns phase 2
+ // reads would let such a row group buffer its survivors and then have
phase 2 throw, and the
+ // reader answers that by not trusting the page index for the rest of the
file. Every row group
+ // after it would then be read whole, which is more than a plain read of
the same query, since
+ // that keeps the pushed data filter's page pruning.
+ //
+ // In the file, the first row group's key column has no offset index while
its value column
+ // does, and the other two have one for both. The pushed data filter
matches one value per row
+ // group, so parquet's statistics filter keeps all three while its column
index narrows the two
+ // it can to the page that value is in.
+ withTempDir { dir =>
+ val blocks = (0 until 3).map(b => ((b * 100 + 1L) to (b * 100 +
100L)).map(i => (i, i)))
+ val path = writeParquetFileByHand(dir, blocks,
+ (blockIdx, colIdx) => blockIdx > 0 || colIdx == 1)
+ // The fixture's whole point, asserted rather than assumed.
+ val indexed = footerChunks(path)
+ .map(c => (c.getPath.toDotString, c.getOffsetIndexReference != null))
+ assert(indexed == Seq(("k", false), ("v", true), ("k", true), ("v",
true),
+ ("k", true), ("v", true)),
+ s"the first row group's key column must be the only chunk without one;
got $indexed")
+ val schema = StructType(Seq(
+ StructField("k", LongType, nullable = false),
+ StructField("v", LongType, nullable = false)))
+ val pushed = Seq(sources.In("v", Array[Any](1L, 101L, 201L)))
+ // Keeps part of each row group, so a row group that can be narrowed is
narrowed.
+ val storageFilters =
+ Seq(GreaterThanOrEqual(BoundReference(0, LongType, nullable = false),
Literal(50L)))
+ val readerFn = new ParquetFileFormat().buildReaderWithStorageFilters(
+ spark, schema, new StructType(), schema, pushed, storageFilters,
+ Map(FileFormat.OPTION_RETURNING_BATCH -> "true"),
+ spark.sessionState.newHadoopConf(), metricMap())
+ .getOrElse(fail("ParquetFileFormat must answer with a reader"))
+ val file = PartitionedFile(
+ InternalRow.empty, SparkPath.fromPathString(path), 0, new
File(path).length())
+ val logAppender = new LogAppender("no offset index")
+ var values: Seq[Long] = null
+ withLogAppender(logAppender) {
+ values = readerFn(file).asInstanceOf[Iterator[Object]].flatMap {
Review Comment:
Fixed. A `drain` helper consumes the iterator and closes it in a `finally`,
and all three sites use it.
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java:
##########
@@ -0,0 +1,1200 @@
+/*
+ * 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.datasources.parquet;
+
+import java.io.IOException;
+import java.time.ZoneId;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Objects;
+import java.util.PrimitiveIterator;
+import java.util.Set;
+
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.page.PageReadStore;
+import org.apache.parquet.filter2.columnindex.RowRanges;
+import org.apache.parquet.hadoop.ParquetFileReader;
+import org.apache.parquet.hadoop.ParquetInputFormat;
+import org.apache.parquet.hadoop.metadata.BlockMetaData;
+import org.apache.parquet.hadoop.metadata.ColumnPath;
+import
org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore.MissingOffsetIndexException;
+import org.apache.parquet.schema.Type;
+import org.apache.parquet.schema.Types;
+
+import org.apache.spark.internal.LogKeys;
+import org.apache.spark.internal.MDC;
+import org.apache.spark.internal.SparkLogger;
+import org.apache.spark.internal.SparkLoggerFactory;
+import org.apache.spark.memory.MemoryMode;
+import org.apache.spark.sql.catalyst.expressions.BasePredicate;
+import org.apache.spark.sql.catalyst.InternalRow;
+import org.apache.spark.sql.execution.vectorized.OffHeapColumnVector;
+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.StructField;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.vectorized.ColumnVector;
+import org.apache.spark.sql.vectorized.ColumnarBatch;
+
+/**
+ * A {@link VectorizedParquetRecordReader} that reads a row group in phases
under a
+ * {@code ParquetStorageFilter}, so the projected columns of the rows the
filter rejects are never
+ * read. The filter is a deterministic conjunct of the post-scan {@code
Filter}, which still holds
+ * it, so any phase below may stop short and read the file the way a plain
scan would.
+ *
+ * <p>Per row group, all three phases driven by the one {@link #fileReader}:
+ * <ul>
+ * <li>phase 0 asks the column index which rows the pushed data filter
allows, metadata only;</li>
+ * <li>phase 1 reads the key columns the filter references over those rows,
evaluates the filter
+ * per row, and buffers the surviving key values;</li>
+ * <li>phase 2 reads the remaining projected columns over the surviving row
ranges.</li>
+ * </ul>
+ *
+ * <p>The emit path then splices the buffered key vectors into the batch, so
phase 2 never reads a
+ * key column again. A row group whose buffer would pass
+ * {@code spark.sql.parquet.storageFilterPushdown.maxSplicedRowGroupBytes}
gives splicing up, and
+ * past that the filter itself, and is read the plain way.
+ *
+ * <p>Splicing does not engage at all in one real case, when every key column
is missing from this
+ * physical file under schema evolution. The predicate is then constant over
the file and is
+ * evaluated once, against the constants the reader materializes for those
columns. True keeps the
+ * file unfiltered, false skips it.
+ *
+ * <p>Everything else the filter needs is guaranteed by {@code
FileSourceStrategy.storageFiltersFor}
+ * and {@code ParquetStorageFilter.create}, so a violation of it here is a
planner bug and is
+ * asserted rather than handled. A file the reader simply cannot prune is a
different matter, and it
+ * reads that one the way a plain scan would.
+ */
+public class LateMaterializationParquetRecordReader extends
VectorizedParquetRecordReader {
+
+ private static final SparkLogger LOG =
+
SparkLoggerFactory.getLogger(LateMaterializationParquetRecordReader.class);
+
+ /**
+ * What one buffered variable-length value costs the vector that holds it,
as a multiple of the
+ * value's own bytes. Those bytes go in the byte child, with one null byte
per element of that
+ * child, and up to as much again because {@link
WritableColumnVector#reserve} doubles the child
+ * when it grows. A factor rather than an exact figure, so this half of the
charge over-counts
+ * rather than under-counts, since nothing above this buffer can spill it. A
fixed-width key needs
+ * no factor, since its accumulator is allocated at capacity and never grows.
+ */
+ private static final long VARIABLE_LENGTH_BYTES_FACTOR = 4L;
+
+ /**
+ * The storage filter this reader applies. It reads key-column pages first,
evaluates the filter
+ * per row to build {@link RowRanges}, then reads data-column pages
restricted to those ranges.
+ * Nulled once the reader gives it up for the whole file, after which every
path here behaves as
+ * the plain reader does.
+ */
+ private ParquetStorageFilter storageFilter;
+
+ /**
+ * Set to true once this reader has no rows left to return, which is when
the row groups run out
+ * or when the filter rejects the whole file.
+ */
+ private boolean hitEndOfData = false;
+
+ /**
+ * Late-materialization state, populated by {@link
#initializeLateMaterialization()}.
+ *
+ * <p>The base class's {@link #fileReader} drives all three phases. Its
requested
+ * schema is mutated per phase via {@link
ParquetFileReader#setRequestedSchema}: all projected
+ * columns for phase 0 ({@code getRowRanges}), the key columns for phase 1,
and for phase 2 the
+ * non-key columns while splicing or the whole projection once a row group
gives it up.
+ *
+ * <p>The three sets are held as leaf-column lists rather than as {@code
MessageType}s because
+ * that is what both the reader and the byte metrics consume, and because
+ * {@code MessageType.getColumns()} rebuilds the list on every call.
+ *
+ * <p>Each list's leaf {@link ColumnPath}s are resolved alongside it,
because that is all the
+ * footer walks below need and {@link ColumnPath#get} is not free, since it
allocates and interns
+ * through parquet's canonicalizer, which those walks would otherwise repeat
per column on every
+ * one of their several calls per row group. As sets, because a walk goes
over the block's column
+ * chunks and asks which of them it is about, and a requested schema's leaf
paths are unique
+ * anyway.
+ */
+ private List<ColumnDescriptor> requestedColumns;
+ private List<ColumnDescriptor> keyOnlyColumns;
+ private List<ColumnDescriptor> nonKeyColumns;
+ private Set<ColumnPath> requestedPaths;
+ private Set<ColumnPath> keyOnlyPaths;
+ private Set<ColumnPath> nonKeyPaths;
+ /**
+ * Whether a row group can be read in part, and what reading part of one
transfers. Both are read
+ * off the footer {@link #fileReader} already holds, so the byte metrics
cost no IO.
+ */
+ private ParquetBlockAccounting blockAccounting;
+ /**
+ * What an accumulator holds per row for each key column, not counting the
value bytes of a
+ * variable-length type. Together with those value bytes it is what the
per-row-group buffer is
+ * measured against. See
+ * {@code spark.sql.parquet.storageFilterPushdown.maxSplicedRowGroupBytes}.
+ */
+ private int keyFixedBytesPerRow;
+ /**
+ * The pages phase 2 read for the row group being emitted. Held because the
column readers draw
+ * from it for the whole row group, and closed when the next one is loaded,
since parquet closes
+ * neither the stores `readFilteredRowGroup` hands back nor the ones from
`readNextRowGroup`.
+ */
+ private PageReadStore dataPages;
+
+ /** Whether the filter has already been reported as failing to evaluate on
this file. */
+ private boolean loggedFilterEvaluationError;
+
+ /**
+ * Whether {@code parquet.filter.columnindex.enabled} says this file's page
index is not to be
+ * used. That covers phase 0 as well as the narrowing below, because the
conf is the escape hatch
+ * for a file whose index is wrong, and phase 2 reads part of a row group
through the offset index
+ * whatever the conf says. A wrong index there pairs a row's key with
another row's values, which
+ * the post-scan Filter cannot catch, since the key it sees is the right one.
+ */
+ private boolean pageIndexDisabled;
+
+ /**
+ * Whether a read of part of a row group has already failed for want of an
offset index, which
+ * says this file is not to be read in part again. Only the narrowing goes.
The rows the pushed
+ * data filter's own page pruning excludes are still excluded, because
parquet built those ranges
+ * itself, and a block it has no index for it hands back whole.
+ */
+ private boolean partialReadUnavailable;
+
+ /**
+ * Whether the row group currently loaded is spliced. It starts true unless
phase 2 will have to
+ * read every projected column anyway, and turns false in phase 1 once the
survivors buffered pass
+ * the cap. False means phase 2 read every projected column, key columns
included, so the emit
+ * path takes them straight from the persistent batch.
+ */
+ private boolean spliceCurrentRowGroup;
+ private int nextBlockIndex;
+ /** This file's row groups, as the footer lists them after parquet's own
row-group filtering. */
+ private List<BlockMetaData> blocks;
+ /**
+ * The key columns this file has, in the filter's key order, with everything
init resolves about
+ * each. Phase 1 reads them, and the survivor queues and accumulators are
indexed the same way.
+ */
+ private KeyColumn[] keyColumns;
+ /**
+ * Whether every key column is missing from this file, which makes the
predicate constant for it.
+ * Decided in {@link #initBatch}, since the constants only exist once its
vectors are built.
+ */
+ private boolean allKeysMissing;
+ private WritableColumnVector[] keyScratchVectors;
+ private ColumnarBatch keyScratchBatch;
+
+ /**
+ * Splicing state. Phase 1 keeps the surviving key values it has already
decoded, and the emit
+ * path splices them back into the output batch, so phase 2 never reads the
key columns.
+ *
+ * <p>The alternative is for phase 2 to read the key columns again under the
surviving row ranges,
+ * which needs none of this state but pays a second read of them for every
row group. Nothing
+ * absorbs that read on object storage, where it is a new GET rather than a
page-cache hit.
+ *
+ * <p>{@link #isKeyTopLevel} marks the top-level slots that emit takes from
the queues rather
+ * than from a phase-2 read. {@link #keyVectorQueues} holds one queue per
present key column of
+ * capacity-sized survivor vectors, and {@link #currentKeyAccumulators} the
vectors still filling.
+ * A queue keeps owning its head while the batch is built on it, until the
next emit closes it, so
+ * survivor memory drains as the row group is emitted. {@link
#persistentBatchColumns} is
+ * {@link #initBatch}'s vector array and {@link #spliceBatchColumns} the
array the emitted batch
+ * is built over, which is why the two are closed separately.
+ *
+ * <p>Phase 1 evaluates a whole row group before its first batch is emitted,
so the queues hold
+ * every surviving key value of one row group at once. That is bounded per
row group by
+ * {@code spark.sql.parquet.storageFilterPushdown.maxSplicedRowGroupBytes},
past which the row
+ * group is read the plain way and nothing is buffered.
+ */
+ private boolean[] isKeyTopLevel;
+ private ArrayDeque<WritableColumnVector>[] keyVectorQueues;
+ private WritableColumnVector[] currentKeyAccumulators;
+ /** Row count of {@link #currentKeyAccumulators}; all key columns advance in
lockstep. */
+ private int currentKeyAccumulatorRowCount;
+ /** Whether each queue's head is the vector the current batch is built on. */
+ private boolean keyVectorsPublished;
+ private ColumnVector[] persistentBatchColumns;
+ private ColumnVector[] spliceBatchColumns;
+
+ /**
+ * The filter is a constructor argument rather than a setter, because
+ * {@link #initializeLateMaterialization()} inspects the per-file schema
against it and runs from
+ * {@code initialize}: there is no point in the reader's life at which it is
useful without one.
+ */
+ public LateMaterializationParquetRecordReader(
+ ZoneId convertTz,
+ String datetimeRebaseMode,
+ String datetimeRebaseTz,
+ String int96RebaseMode,
+ String int96RebaseTz,
+ boolean useOffHeap,
+ int capacity,
+ ParquetStorageFilter storageFilter) {
+ super(convertTz, datetimeRebaseMode, datetimeRebaseTz, int96RebaseMode,
int96RebaseTz,
+ useOffHeap, capacity);
+ this.storageFilter = Objects.requireNonNull(storageFilter,
"storageFilter");
+ }
+
+ // For test only.
+ public LateMaterializationParquetRecordReader(
+ boolean useOffHeap, int capacity, ParquetStorageFilter storageFilter) {
+ super(useOffHeap, capacity);
+ this.storageFilter = Objects.requireNonNull(storageFilter,
"storageFilter");
+ }
+
+ @Override
+ protected void closeAllocated() {
+ // Each release below is independent, and the base class's close() owns
the file handle and
+ // input stream, so they are chained through finally blocks, and one
failing vector close must
+ // not leak the rest.
+ try {
+ // The batch's vectors are closed through `persistentBatchColumns`
rather than through
+ // `columnarBatch.close()`, which lets both paths share one body. While
splicing the emitted
+ // batch is a view whose slots alias these vectors (non-key and
partition) and the head of
+ // each survivor queue (key slots), so closing it would re-close a
shared vector and free the
+ // same buffer twice. Without splicing its array is this one.
+ try {
+ if (persistentBatchColumns != null) {
+ closeAll(persistentBatchColumns);
+ } else if (columnarBatch != null) {
+ // `takeOverBatch` has not run, or threw part way, and the batch
super built is then the
+ // only thing holding its vectors.
+ columnarBatch.close();
+ }
+ persistentBatchColumns = null;
+ spliceBatchColumns = null;
+ columnarBatch = null;
+ } finally {
+ closeSplicingState();
+ }
+ } finally {
+ try {
+ // Through the array, not through `keyScratchBatch`, which is assigned
only after the
+ // allocation loop finishes, so a partial failure leaves vectors only
the array can reach.
+ closeAll(keyScratchVectors);
+ keyScratchVectors = null;
+ keyScratchBatch = null;
+ } finally {
+ closeDataPages();
+ }
+ }
+ }
+
+ @Override
+ public float getProgress() {
+ // Under a storage filter, rowsReturned counts survivors while
totalRowCount is the count before
+ // the filter narrowed anything, so the ratio neither reaches 1 on its own
nor stays below it. A
+ // row group read whole after the page index is given up can return more
rows than that count.
+ // `hitEndOfData` is the real terminator here, and the ratio is clamped to
what it means.
+ if (hitEndOfData) return 1.0f;
+ return Math.min(1.0f, super.getProgress());
+ }
+
+ /**
+ * Builds the batch the way the plain reader does, takes it over for
splicing, and then answers
+ * the filter for a file that has none of its key columns.
+ *
+ * <p>Both after super's body, because the batch it hands out is what the
splicing one is built
+ * over, and the all-keys-missing decision evaluates the predicate against
the constant vectors
+ * this reader returns for those columns, which exist only once super has
wrapped the allocated
+ * vectors in {@link ParquetColumnVector}s.
+ */
+ @Override
+ protected void initBatch(
+ MemoryMode memMode,
+ StructType partitionColumns,
+ InternalRow partitionValues) {
+ super.initBatch(memMode, partitionColumns, partitionValues);
+ takeOverBatch();
+ if (allKeysMissing) {
+ decideAllKeysMissingFile();
+ }
+ }
+
+ /**
+ * Takes the vectors of the batch super built, and while splicing hands out
a batch of this
+ * reader's own over a copy of that array, whose key slots the emit path
rewrites in place.
+ * {@link ColumnarBatch} holds the array by reference, its staging row
included, so rewriting a
+ * slot is what publishes it, and the object being reused is what makes a
consumer see it.
+ *
+ * <p>Every slot got a vector, key columns included. While splicing a key
slot's vector is
+ * unused, since the emitted batch takes that slot from the survivor queues.
A row group read the
+ * plain way past the buffer cap does read into it, and one capacity-sized
vector per key column
+ * is cheap next to the buffer the cap is there to bound.
+ */
+ private void takeOverBatch() {
+ persistentBatchColumns = new ColumnVector[columnarBatch.numCols()];
+ for (int i = 0; i < persistentBatchColumns.length; i++) {
+ persistentBatchColumns[i] = columnarBatch.column(i);
+ }
+ if (keyColumns == null) return;
+ spliceBatchColumns = persistentBatchColumns.clone();
+ columnarBatch = new ColumnarBatch(spliceBatchColumns);
+ }
+
+ /**
+ * Answers the storage filter for a file that has none of its key columns,
which makes the
+ * predicate constant over the file. It is evaluated against the vectors
just built for those
+ * columns, so the value it sees is the one the scan returns for them.
+ *
+ * <p>False skips the file. True, or an error, reads it the way a plain scan
would, under the same
+ * fail-open rule as a row's value, since the constant can be one the
predicate throws on.
+ */
+ private void decideAllKeysMissingFile() {
+ allKeysMissing = false;
+ // Not closed, and must not be, since the vectors in it belong to the
batch this reader emits.
+ ColumnarBatch keyRow = new ColumnarBatch(keyVectorsFromBatch(), 1);
+ boolean keep;
+ try {
+ keep = storageFilter.preparedPredicate().eval(keyRow.getRow(0));
+ } catch (RuntimeException e) {
+ logFilterGivenUpOnError(e);
+ keep = true;
+ }
+ if (!keep) {
+ recordFileSkipped();
+ hitEndOfData = true;
+ }
+ storageFilter = null;
+ }
+
+ /**
+ * Whether this read has returned every row it is going to.
+ *
+ * <p>{@code totalRowCount} is how many rows the file yields before this
filter narrows anything,
+ * so it is a terminator only while there is no filter. Under one it is an
upper bound the read
+ * never reaches, and a row group read whole after the page index was given
up can pass it. What
+ * ends such a read is {@link #hitEndOfData}, set when the row groups run
out or when the filter
+ * rejects the whole file.
+ */
+ @Override
+ protected boolean noMoreRows() {
+ return hitEndOfData || (storageFilter == null && super.noMoreRows());
+ }
+
+ /**
+ * Releases the key vectors the previous batch was built on, before the
batch is emptied, so a
+ * terminal call cannot leave a spliced batch pointing at vectors just
released, which off heap is
+ * freed memory.
+ */
+ @Override
+ protected void resetBatch() {
+ releasePublishedKeyVectors();
+ super.resetBatch();
+ }
+
+ /**
+ * Points the emitted batch's key slots at what this batch takes them from,
then reads the rest
+ * the way the plain path does. Three cases reach the base's loop. A spliced
row group, whose key
+ * slots have no phase-2 reader because the emitted batch takes them from
the survivor queues. A
+ * row group read the plain way past the buffer cap, which has a reader for
every slot. And a file
+ * the filter was given up for, which is the plain path exactly.
+ *
+ * <p>The row-index column the base fills there is identified by name
+ * (ROW_INDEX_TEMPORARY_COLUMN_NAME), a synthetic metadata column no storage
filter references, so
+ * its slot is always a non-key one and its ParquetColumnVector is always
the persistent one.
+ */
+ @Override
+ protected void readBatchColumns(int num) throws IOException {
+ // Without a storage filter the emitted batch is `initBatch`'s array and
this does not apply.
+ if (spliceBatchColumns != null) {
+ pointKeySlotsAtThisBatch(num);
+ }
+ super.readBatchColumns(num);
+ }
+
+ /**
+ * Points the emitted batch's key slots at the vectors this batch takes them
from. They are the
+ * only slots that ever change, since a non-key slot holds the vector {@link
#initBatch} built for
+ * it for the reader's life. A spliced row group takes the head of each
survivor queue, which is
+ * what publishes the buffered keys, and any other row group takes the
persistent vector, which
+ * is what a row group following a spliced one needs.
+ *
+ * <p>The same {@link ColumnarBatch} is handed out every time, over an array
it holds by
+ * reference. The queues keep owning the vectors they publish until the next
batch releases them,
+ * so a phase-2 read that throws afterwards leaves them reachable for {@link
#close()}.
+ */
+ private void pointKeySlotsAtThisBatch(int num) {
+ keyVectorsPublished = spliceCurrentRowGroup;
+ // Queue k holds the survivors of key column k, and its `batchSlot` is
where that key column
+ // sits in the batch, so the pairing is read off rather than re-derived
from slot order.
+ for (int k = 0; k < keyColumns.length; k++) {
+ int slot = keyColumns[k].batchSlot();
+ spliceBatchColumns[slot] =
+ spliceCurrentRowGroup ? nextSurvivorVector(k, num) :
persistentBatchColumns[slot];
+ }
+ }
+
+ /** The head of survivor queue {@code k}, which is what publishes the keys
it buffered. */
+ private WritableColumnVector nextSurvivorVector(int k, int num) {
+ WritableColumnVector vector = keyVectorQueues[k].peekFirst();
+ if (vector == null) {
+ // Unreachable, since the queues hold exactly the survivors phase 1
accumulated and the emit
+ // loop is driven by that same count. Named here rather than left to a
null batch slot and an
+ // NPE somewhere in the read that follows.
+ throw new IllegalStateException(String.format(
+ "Storage-filter survivor queue %d of row group %d in %s ran out with
%d rows still "
+ + "to emit", k, nextBlockIndex - 1, fileReader.getFile(), num));
+ }
+ return vector;
+ }
+
+ /**
+ * Closes the key vectors the previous batch was built on. The queues own
them until here, so
+ * survivor memory drains as a row group is emitted rather than all at its
end.
+ */
+ private void releasePublishedKeyVectors() {
+ if (!keyVectorsPublished) return;
+ keyVectorsPublished = false;
+ for (ArrayDeque<WritableColumnVector> queue : keyVectorQueues) {
+ queue.removeFirst().close();
+ }
+ }
+
+ @Override
+ protected void initializeInternal() throws IOException,
UnsupportedOperationException {
+ super.initializeInternal();
+ initializeLateMaterialization();
+ }
+
+ private void initializeLateMaterialization() throws IOException {
+ if (fileReader == null) {
+ // The three phases drive a ParquetFileReader directly, and this reader
was handed a row-group
+ // reader with none behind it, so the filter is simply not applied and
the plain read path
+ // runs. The conjunct is in the post-scan filter as well, so the rows it
would have dropped
+ // are dropped above the scan.
+ storageFilter = null;
+ return;
+ }
+ blockAccounting = new ParquetBlockAccounting(fileReader);
+ blocks = fileReader.getRowGroups();
+ // Both of the flags the narrowing hangs on start here. The conf one
closes phase 0 too, and the
+ // other is what the failed read below latches. Neither costs the filter
the row groups it
+ // empties, which need no index at all.
+ pageIndexDisabled =
+
!configuration.getBoolean(ParquetInputFormat.COLUMN_INDEX_FILTERING_ENABLED,
true);
+ // Resolve each key column's top-level ParquetColumn, keeping the ones
this file has. A missing
+ // one can happen under schema evolution, where a column is in the
requested schema but not in
+ // the file. For any non-primitive key we still bail, since phase 1 reads
only primitive leaves.
+ int[] keyIndices = storageFilter.keyColumnIndices();
Review Comment:
A correction to my reply above. Deleting the override tripped MiMa, as a
`DirectAbstractMethodProblem` reported against
`org.apache.hadoop.mapreduce.RecordReader.initialize`, which the
`sql.execution.*` exclude does not cover. So the override is back and only
calls `super.initialize`. The double initialization is gone either way, and the
test still covers it.
--
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]
