peter-toth commented on code in PR #58895: URL: https://github.com/apache/spark/pull/58895#discussion_r4204523275
########## sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala: ########## @@ -0,0 +1,3708 @@ +/* + * 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.{ByteArrayInputStream, ByteArrayOutputStream, File, RandomAccessFile} +import java.net.URI +import java.time.LocalTime +import java.util.Locale +import java.util.concurrent.atomic.AtomicLong + +import scala.collection.mutable +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +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.{ColumnDescriptor, Encoding, ParquetProperties} +import org.apache.parquet.column.impl.ColumnWriteStoreV1 +import org.apache.parquet.column.page.{DataPageV1, PageReadStore} +import org.apache.parquet.column.page.mem.MemPageStore +import org.apache.parquet.filter2.columnindex.RowRanges +import org.apache.parquet.format.Util +import org.apache.parquet.format.converter.ParquetMetadataConverter +import org.apache.parquet.hadoop.{ParquetFileReader, ParquetFileWriter, ParquetInputFormat, ParquetOutputFormat, ParquetWriter} +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.mockito.ArgumentMatchers.{any, anyInt} +import org.mockito.Mockito.{doAnswer, doThrow, spy} +import org.mockito.invocation.InvocationOnMock + +import org.apache.spark.{SparkEnv, SparkException, TaskContext, TaskKilledException} +import org.apache.spark.internal.config.KILL_ON_FATAL_ERROR_DEPTH +import org.apache.spark.paths.SparkPath +import org.apache.spark.sql.{sources, DataFrame, QueryTest, Row, SparkSession} +import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow} +import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeReference, BloomFilterMightContain, BoundReference, Cast, Coalesce, EqualTo, Expression, GreaterThanOrEqual, In, IsNull, LessThan, LessThanOrEqual, Literal, Or, Predicate, Rand, Remainder, SecondsToTimestamp, Substring, 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.{DataSourceUtils, FileFormat, FileSourceStrategy, OutputWriterFactory, PartitionedFile} +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.execution.vectorized.{OffHeapColumnVector, OnHeapColumnVector, WritableColumnVector} +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, + options: Map[String, String] = Map.empty): 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)) + // Last, so a caller's options override the ones above. + writer.options(options).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) + + // 400 rows whose values sort with their keys, the fixture of the page-level tests. + private val rows400 = (1L to 400L).map(i => (i, f"v_$i%04d")) + + // `rows400` in one row group of 100-row pages. + private def writePagedFile(dir: File): String = + writeParquetFile(dir, rows400, rowGroupSize = 64 * 1024L, pageSize = Some(512L)) + + // 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 = drain(readerFn(file)) { it => + it.map { + case batch: ColumnarBatch => batch.numRows() + case _ => 1 + }.sum + } + (emitted, CountingLocalFileSystem.bytesRead()) + } + + // Consumes what a reader builder returned for one file through `consume`, and closes the reader + // behind it whatever happens. Exhausting the iterator closes it too, but a failure part way + // through, which is what these tests look for, would otherwise leave it and its stream open. + private def drain[T](iter: Iterator[_])(consume: Iterator[Object] => T): T = { + try { + consume(iter.asInstanceOf[Iterator[Object]]) + } finally { + iter match { + case closeable: java.io.Closeable => closeable.close() + case _ => + } + } + } + + // A batch reader for a file of `schema`, built through `ParquetFileFormat` the way a scan with + // storage filters builds one. + private def formatReader( + schema: StructType, + pushed: Seq[sources.Filter], + storageFilters: Seq[Expression], + hadoopConf: Configuration = spark.sessionState.newHadoopConf(), + metrics: Map[String, SQLMetric] = metricMap()): PartitionedFile => Iterator[InternalRow] = + new ParquetFileFormat().buildReaderWithStorageFilters( + spark, schema, new StructType(), schema, pushed, storageFilters, + Map(FileFormat.OPTION_RETURNING_BATCH -> "true"), hadoopConf, metrics) + .getOrElse(fail("ParquetFileFormat must answer with a reader")) + + // Asserts that a row group costs exactly `charge`: a cap of it keeps splicing, which reads the + // key column once, and a cap one byte under it gives splicing up, which reads it again. `read` + // reads the file under a cap and returns the rows it emitted with the bytes it transferred. + private def assertSpliceCharge(charge: Long, rows: Int)(read: String => (Int, Long)): Unit = { + val (uncappedRows, uncapped) = read("64MB") + val (atChargeRows, atCharge) = read(s"${charge}b") + val (underChargeRows, underCharge) = read(s"${charge - 1}b") + assert(Seq(uncappedRows, atChargeRows, underChargeRows) == Seq(rows, rows, rows), + s"every cap must keep the same $rows rows; got $uncappedRows, $atChargeRows and " + + underChargeRows) + assert(atCharge == uncapped, + s"a cap of exactly the charge, $charge, must keep splicing; got $atCharge against $uncapped") + assert(underCharge > uncapped, + s"a cap one byte under it must give splicing up; got $underCharge against $uncapped") + } + + // Reads `path` whole through the format under a splice cap and a batch size, and returns the rows + // it emitted with the bytes it transferred. + private def readUnderCap( + path: String, + schema: StructType, + storageFilters: Seq[Expression], + maxSplicedBytes: String, + batchSize: Int, + metrics: Map[String, SQLMetric] = metricMap()): (Int, Long) = withSQLConf( + SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_MAX_SPLICED_ROW_GROUP_BYTES.key -> maxSplicedBytes, + SQLConf.PARQUET_VECTORIZED_READER_BATCH_SIZE.key -> batchSize.toString) { + readCounting(path, formatReader(schema, Nil, storageFilters, countingHadoopConf(), metrics)) + } + + // Reads `path` whole through `readerFn`, and returns the long column at `ordinal` of every row. + private def readLongs( + readerFn: PartitionedFile => Iterator[InternalRow], path: String, ordinal: Int): Seq[Long] = { + val file = PartitionedFile( + InternalRow.empty, SparkPath.fromPathString(path), 0, new File(path).length()) + drain(readerFn(file)) { it => + it.flatMap { + case batch: ColumnarBatch => (0 until batch.numRows()).map(batch.column(ordinal).getLong) + case row: InternalRow => Seq(row.getLong(ordinal)) + }.toSeq + } + } + + // The native address of an off-heap vector's values, 0 once it is freed. Only for a fixed-width + // vector. A string or binary one keeps its bytes in a child and reports 0 even while live. + private def address(vector: Any): Long = + vector.asInstanceOf[OffHeapColumnVector].valuesNativeAddress() + + // Runs `body` as a task under `context`. The block manager has to know the task for a format + // reader to read its broadcast hadoop conf there. + private def asTask[T](context: TaskContext)(body: => T): T = { + val blockManager = SparkEnv.get.blockManager + blockManager.registerTask(context.taskAttemptId()) + try { + TaskContext.withTaskContext(context)(body) + } finally { + blockManager.releaseAllLocksForTask(context.taskAttemptId()) + } + } + + // Runs `body`, and returns its result with how many log lines containing `text` it logged. + private def withLogged[T](text: String)(body: => T): (T, Int) = { + val logAppender = new LogAppender(text) + var result: Option[T] = None + withLogAppender(logAppender) { + result = Some(body) + } + val logged = logAppender.loggingEvents.map(_.getMessage.getFormattedMessage) + .count(_.contains(text)) + (result.get, logged) + } + + // What the reader logs when it reads a row group without page-level storage filtering. + private val partialReadWarning = "without page-level storage filtering" + + // What the reader logs when filling the survivor buffer fails and it gives splicing up. + private val splicingGivenUpOnError = "buffering a surviving key value raised an error" + + // What the reader logs when applying the filter to a row, or decoding the key pages, raises an + // error and it gives the filter up for that row group. + private val evaluationGivenUp = "applying it to a row raised an error" + private val decodingGivenUp = "decoding its key columns raised an error" + + + // 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.fromMap(metricMap()) + + // `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) + + // The `(k, v)` schema most fixtures read. `create` reads only the types and names, and the reader + // takes each column's repetition from the file, so `nullable` matters only to a format read. + private def kvSchema( + keyType: DataType = LongType, + valueType: DataType = StringType, + nullable: Boolean = true): StructType = + StructType(Seq(StructField("k", keyType, nullable), StructField("v", valueType, nullable))) + + // Collects all `(k, v)` rows from a reader initialized with the given storage filter. + private def readAll( + filePath: String, + storageFilter: ParquetStorageFilter, + capacity: Int = 4096, + useOffHeap: Boolean = false, + onKeyReader: KeyReaderHook.OnKeyReader = null) + : (Seq[(Long, String)], VectorizedParquetRecordReader) = + readAllWith(filePath, Seq("k", "v"), storageFilter, + (batch, i) => (batch.column(0).getLong(i), batch.column(1).getUTF8String(i).toString), + capacity, useOffHeap, onKeyReader = onKeyReader) + + // `k >= threshold` bound to position 0, over a long key. + private def keyAtLeast(threshold: Long): Expression = + GreaterThanOrEqual(BoundReference(0, LongType, nullable = true), Literal(threshold)) + + // Builds a `k >= threshold` storage filter bound to position 0, over a long key. + private def keyAtLeastFilter( + threshold: Long, + metrics: StorageFilterMetrics = allMetrics()): ParquetStorageFilter = { + val expr = keyAtLeast(threshold) + val requested = kvSchema() + createFilter(Seq(expr), requested, 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) + + 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 == 1, + s"expected the one row group skipped; got ${m.rowGroupsSkipped.value}") + assert(m.rowsExcludedByRowGroup.value == 40, + s"expected its 40 rows excluded; 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 the row group avoids exactly the v chunk a + // plain read of this projection reads, and phase 1 still pays for k. + val vChunkBytes = footerOf(path).getBlocks.get(0).getColumns.get(1).getTotalSize + assert(m.bytesAvoidedByRowGroup.value == vChunkBytes, + s"expected the v chunk's $vChunkBytes bytes avoided; 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 == rows, 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() + } + } + } + + Seq(false, true).foreach { useOffHeap => + val mode = if (useOffHeap) "off-heap" else "on-heap" + test(s"mixed, $mode vectors: 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, useOffHeap = useOffHeap) + try { + // Output is exact. Phase 2's column readers are handed `finalRanges` and skip the rows + // outside them within a partly kept page, so emitted rows == survivors. + val expected = rows.filter(_._1 >= 195L) + assert(result == expected, + s"expected exact filtering; got ${result.map(_._1)}, expected ${expected.map(_._1)}") + 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") + // A mixed projection has non-key bytes to avoid, in both the skipped row groups and + // the partially kept one. Asserting they are non-negative would prove nothing, since + // `SQLMetric.add` drops a negative and the value cannot go below zero. + assert(m.bytesAvoidedByRowGroup.value + m.bytesAvoidedByPageFiltering.value > 0, + "a mixed projection with skipped row groups should avoid some non-key bytes") + } 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 = keyAtLeast(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 `initializeInternal`, and it used to run + // it a second time on top. + withTempDir { dir => + val path = writeKeyParquetFile(dir, 1L to 20L) + val requested = kvSchema() + // 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()) + var initializations = 0 + val reader = new VectorizedParquetRecordReader(false, 4096) { + override protected def initializeInternal(): Unit = { + initializations += 1 + super.initializeInternal() + } + } + Utils.tryWithResource(reader) { reader => + reader.initialize(split, context) + assert(initializations == 1, s"initialized $initializations times") + } + } + } + + 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 = kvSchema() + + // 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(keyAtLeast(0L)), + rowIndexSchema) + } + assert(rowIndexKey.getMessage.contains(rowIndexName), rowIndexKey.getMessage) + val rowIndexAttr = AttributeReference(rowIndexName, LongType, nullable = false)() + assert(!ParquetStorageFilter.isSupportedStorageFilter( + BloomFilterMightContain( + bloomLiteralOf(42L), + new XxHash64(Seq(rowIndexAttr)))), + "a bloom over a column named like the row-index column must not be offered") + } + + // A runtime bloom's `Literal` holding one long key, built the way the planner builds one. It is + // serialized as `BloomFilterAggregate` serializes it, and the key is hashed with the same seed. + private def bloomLiteralOf(key: Long): Literal = { + val bf = BloomFilter.create(10, 128) + bf.putLong(new XxHash64(Seq(Literal(key))).eval(InternalRow.empty).asInstanceOf[Long]) + Literal(BloomFilterAggregate.serialize(bf), BinaryType) + } + + 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 = kvSchema(decimalType) + // 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 = kvSchema() + val key = OnKey(BoundReference(0, LongType, nullable = false), at = 7L, + error = new SparkException(_)) + 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() + } + } + } + + test("a task kill or a fatal error behind the key expression's error is not absorbed") { + // The fail-open must not swallow what the executor has to see. A kill interrupts the task's + // thread, and a UDF that wraps the interrupt raises its own error, so the kill is read off the + // task context rather than off the error. It goes out wrapped in a checked exception, since + // `FileScanRDD` reads a RuntimeException from a reader as a corrupt file under + // `ignoreCorruptFiles` and would skip the rest of it. A fatal error an expression wrapped is + // found in the cause chain, which is where the executor looks for one, and goes out as itself. + withTempDir { dir => + val rows = (0L until 20L).map(i => (i, i.toString)) + val path = writeParquetFile(dir, rows, rowGroupSize = 64 * 1024L) + val requested = kvSchema() + val bound = BoundReference(0, LongType, nullable = false) + def filterOn(key: Expression): ParquetStorageFilter = + createFilter(Seq(GreaterThanOrEqual(key, Literal(10L))), requested) + def read(key: Expression): Unit = readAll(path, filterOn(key))._2.close() + def assertPropagates(thrown: SparkException, original: Class[_]): Unit = { + assert(!DataSourceUtils.shouldIgnoreCorruptFileException(thrown), + s"ignoreCorruptFiles must not be able to read this as a corrupt file; got $thrown") + assert(original.isInstance(thrown.getCause) && + thrown.getCause.getMessage.contains("refuses 7"), + s"the key expression's own error must reach the task behind it; got ${thrown.getCause}") + } + + def killedWhile(key: Expression): SparkException = + TaskContext.withTaskContext(TaskContext.empty())(intercept[SparkException](read(key))) + // The reader does not give the filter up on a kill, so it does not report doing so. + val (killed, warnings) = withLogged(evaluationGivenUp) { + killedWhile(OnKey(bound, at = 7L, killsTask = true, error = new SparkException(_))) + } + assertPropagates(killed, classOf[SparkException]) + assert(warnings == 0, s"a kill must not be reported as the filter given up; got $warnings") + assertPropagates( + killedWhile(OnKey(bound, at = 7L, killsTask = true, error = new IllegalStateException(_))), + classOf[IllegalStateException]) + + // An unchecked error around an OutOfMemoryError, the way a UDF's catch-all wraps one. + val fatal = intercept[OutOfMemoryError] { + read(OnKey(bound, at = 7L, + error = new IllegalStateException(_, new OutOfMemoryError("wrapped")))) + } + assert(fatal.getMessage == "wrapped", s"the wrapped error itself must go out; got $fatal") + + // An `InternalError` goes out in an `UnknownError`, since `ignoreCorruptFiles` would take + // the bare one for a corrupt file, where a plain plan raises it from the post-scan Filter. + val internal = intercept[UnknownError] { + read(OnKey(bound, at = 7L, error = new InternalError(_))) + } + assert(!DataSourceUtils.shouldIgnoreCorruptFileException(internal) && + internal.getCause.isInstanceOf[InternalError], s"got $internal") + + // An AssertionError is not fatal, so the reader falls back on it as on any other error. + val (result, reader) = + readAll(path, filterOn(OnKey(bound, at = 7L, error = new AssertionError(_)))) + try { + assert(result == rows, s"the row group must come back whole; got ${result.size} rows") + } finally { + reader.close() + } + } + } + + test("a fatal error goes out as itself, found as deep as the executor would find it") { + // A plain plan raises a filter's error from the post-scan Filter as it is, and a decoding error + // from the reader, which `FileScanRDD` wraps in FAILED_READ_FILE. The executor looks + // `spark.executor.killOnFatalError.depth` levels down from what reaches it, so the fatal cause + // of a `NonFatal` decoding error counts down to one level less. An `InternalError` from + // decoding is corrupt data, which the plain read meets again. + val depth = SparkEnv.get.conf.get(KILL_ON_FATAL_ERROR_DEPTH) + def under(levels: Int, fatal: Throwable): Throwable = + (0 until levels).foldLeft(fatal)((cause, _) => new IllegalStateException("wraps", cause)) + def rethrow(e: Throwable, fromRead: Boolean): Unit = + ParquetStorageFilter.rethrowIfMustPropagate(e, fromRead) + def fallsBack(e: Throwable, fromRead: Boolean): Unit = + try rethrow(e, fromRead) catch { + // ScalaTest would abort the whole suite on these rather than fail this test. + case t @ (_: VirtualMachineError | _: LinkageError) => + fail(s"the reader must fall back on $e; got $t", t) + } + val oom = new OutOfMemoryError("fatal") + + assert(intercept[OutOfMemoryError](rethrow(under(depth - 1, oom), fromRead = false)) eq oom, + "the deepest level the executor sees counts") + fallsBack(under(depth, oom), fromRead = false) + + assert(intercept[OutOfMemoryError](rethrow(under(depth - 2, oom), fromRead = true)) eq oom, + "a decoding error's chain counts one level less deep") + fallsBack(under(depth - 1, oom), fromRead = true) + + // The depth is the executor's conf, not its default. + withSparkEnvConfs(KILL_ON_FATAL_ERROR_DEPTH.key -> "2") { + intercept[OutOfMemoryError](rethrow(under(1, oom), fromRead = false)) + fallsBack(under(2, oom), fromRead = false) + } + // `FileScanRDD` wraps only a `NonFatal` error, so a decoding `LinkageError` is searched as + // deep as the executor searches it. The OOM sits three levels down here. + withSparkEnvConfs(KILL_ON_FATAL_ERROR_DEPTH.key -> "3") { + val linkage = new ExceptionInInitializerError(new RuntimeException("static init", oom)) + assert(intercept[OutOfMemoryError](rethrow(linkage, fromRead = true)) eq oom, + "a decoding error that is not NonFatal reaches the executor unwrapped") + // One level further down is past what the executor looks at. + fallsBack(new ExceptionInInitializerError( + new RuntimeException("static init", new RuntimeException("wraps", oom))), fromRead = true) + } + + // A codec's `InternalError` while decoding is corrupt data, at the top or in the chain. + fallsBack(new InternalError("codec"), fromRead = true) + fallsBack(under(1, new InternalError("codec")), fromRead = true) + // One from an expression goes out in an `UnknownError`, raw or wrapped. `ignoreCorruptFiles` + // would take the bare `InternalError` for a corrupt file. + Seq(new InternalError("expression"), under(1, new InternalError("expression"))).foreach { e => + val thrown = intercept[UnknownError](rethrow(e, fromRead = false)) + assert(thrown.getCause.isInstanceOf[InternalError], s"got $thrown") + assert(!DataSourceUtils.shouldIgnoreCorruptFileException(thrown), s"got $thrown") + } + fallsBack(new AssertionError("not fatal"), fromRead = false) + + // A vector only the reader allocates that cannot grow, the way `reserve` reports it, is the + // reader's own extra work while decoding. Raised from an expression it is a wrapped OOM. The + // reader matches the exact class `reserve` throws, which a real one pins. + val vector = new OnHeapColumnVector(1, LongType) + val reserved = try intercept[RuntimeException](vector.reserve(-1)) finally vector.close() + assert(reserved.getClass == classOf[RuntimeException], + s"`reserve` must still throw the class the reader matches; got ${reserved.getClass}") + val growth = new RuntimeException("Cannot reserve", new OutOfMemoryError("reserve")) + fallsBack(growth, fromRead = true) + assert(intercept[OutOfMemoryError](rethrow(growth, fromRead = false)).getMessage == "reserve") + // Only that exact shape. A plain exception with a fatal error deeper down is not one, and + // neither is a subclass with an OutOfMemoryError cause. + assert(intercept[OutOfMemoryError] { + rethrow(new RuntimeException("wraps", under(1, oom)), fromRead = true) + } eq oom) + assert(intercept[OutOfMemoryError] { + rethrow(new IllegalStateException("wraps", oom), fromRead = true) + } eq oom) + + // A fatal error at the top goes out as it is even with a kill pending, as the executor checks + // it first. One deeper in the chain goes out as the kill, which the executor reports. So does a + // vector that cannot grow. An error the executor's kill case does not take goes out as it is. + val context = TaskContext.empty() + context.markInterrupted("killed by the test") + TaskContext.withTaskContext(context) { + assert(intercept[OutOfMemoryError](rethrow(oom, fromRead = false)) eq oom) + // A `NonFatal` wrapper would be reported as killed, and the executor would run on. + assert(intercept[UnknownError] { + rethrow(new InternalError("expression"), fromRead = false) + }.getCause.isInstanceOf[InternalError]) + val killed = intercept[SparkException](rethrow(under(1, oom), fromRead = false)) + assert(killed.getCause.isInstanceOf[IllegalStateException], s"got $killed") + assert(intercept[SparkException](rethrow(growth, fromRead = true)).getCause eq growth) + val codec = new InternalError("codec") + assert(intercept[InternalError](rethrow(codec, fromRead = true)) eq codec, + "a plain read raises the codec's error as it is, which the executor does not report as " + + "killed unless ignoreCorruptFiles skips it") + val initError = new ExceptionInInitializerError(new IllegalStateException("static init")) + assert(intercept[ExceptionInInitializerError](rethrow(initError, fromRead = false)) eq + initError, "a LinkageError is neither NonFatal nor fatal, and goes out as it is") + } + } + + test("all keys missing: a kill or a fatal error on the constant is not absorbed") { + // A file with none of the key columns evaluates the filter once, on the value the scan returns + // for them. That catch gives the filter up on the filter's own errors like the per-row one, + // and has to let a kill and a fatal error out like it too. + withSQLConf(SQLConf.ENABLE_DEFAULT_COLUMNS.key -> "true") { + withTempDir { dir => + val path = + writeSingleParquetFile(dir, spark.range(0, 10).selectExpr("id AS v"), 1024L) + val schema = StructType(Seq( + StructField("k", LongType).withExistenceDefaultValue("7").withCurrentDefaultValue("7"), + StructField("v", LongType))) + val bound = BoundReference(0, LongType, nullable = true) + def read(key: Expression): Seq[Long] = readLongs( + formatReader(schema, Nil, Seq(GreaterThanOrEqual(key, Literal(0L)))), path, ordinal = 1) + + assert(read(bound) == (0L until 10L), "the default passes the filter, so the file is read") + val fatal = intercept[OutOfMemoryError] { + read(OnKey(bound, at = 7L, + error = new IllegalStateException(_, new OutOfMemoryError("w")))) + } + assert(fatal.getMessage == "w", s"the wrapped error itself must go out; got $fatal") + // As for a row, an `InternalError` here is the filter's own, not corrupt data. + val internal = intercept[UnknownError] { + read(OnKey(bound, at = 7L, error = new InternalError(_))) + } + assert(internal.getCause.isInstanceOf[InternalError], s"got $internal") + val killed = asTask(TaskContext.empty()) { + intercept[SparkException] { + read(OnKey(bound, at = 7L, killsTask = true, error = new SparkException(_))) + } + } + assert(killed.getMessage.contains("task kill") && + !DataSourceUtils.shouldIgnoreCorruptFileException(killed), s"got $killed") + } + } + } + + test("a give-up WARN is reported once per split, and a give-up stays in its row group") { + // The first and the last of three row groups give the filter up on an evaluation error, and the + // reader reports that once. The middle one raises nothing, so the filter must still drop its + // keys 151 to 200, which a reader that gave the filter up for the rest of the split would not. + withTempDir { dir => + val rows = (1L to 300L).map(i => (i, s"v_$i")) + val path = writeSingleParquetFile(dir, rows.toDF("k", "v"), + rowGroupSize = ParquetWriter.DEFAULT_BLOCK_SIZE, + options = Map(ParquetOutputFormat.BLOCK_ROW_COUNT_LIMIT -> "100")) + assert(footerOf(path).getBlocks.asScala.map(_.getRowCount) == Seq(100L, 100L, 100L), + "the fixture must be three row groups of 100 rows") + val bound = BoundReference(0, LongType, nullable = true) + def failingAt(key: Long): Expression = GreaterThanOrEqual( + OnKey(bound, at = key, error = new IllegalStateException(_)), Literal(0L)) + val dropsMiddle = Or(LessThanOrEqual(bound, Literal(150L)), LessThan(Literal(200L), bound)) + val filter = createFilter(Seq(And(And(failingAt(5L), failingAt(250L)), dropsMiddle)), + kvSchema()) + val ((result, reader), warnings) = withLogged(evaluationGivenUp)(readAll(path, filter)) + try { + assert(result == rows.filter { case (k, _) => k <= 150L || k > 200L }, + s"the middle row group must be filtered and the other two come back whole; got " + + s"${result.size} rows") + assert(warnings == 1, s"the evaluation WARN must be reported once; got $warnings") + } finally { + reader.close() + } + } + } + + test("a task kill is seen while the reader works through a row group") { + // A plain read meets a kill in `FileScanRDD` after every batch. This reader decodes and + // evaluates a whole row group before its first batch, and can pass row groups the filter + // empties without emitting one, so it checks for a kill itself, per chunk and per row group. + withTempDir { dir => + val rows = (0L until 100L).map(i => (i, i.toString)) + val path = writeParquetFile(dir, rows, rowGroupSize = 64 * 1024L) + assert(footerOf(path).getBlocks.size == 1, "the fixture must be one row group") + val bound = BoundReference(0, LongType, nullable = false) + // Marks the task killed at row 5, inside phase 1's first chunk of 16 rows. With one row + // group, only the check at the next chunk can see it before the read ends. + val filter = createFilter( + Seq(GreaterThanOrEqual(OnKey(bound, at = 5L, killsTask = true), Literal(50L))), kvSchema()) + val thrown = TaskContext.withTaskContext(TaskContext.empty()) { + intercept[SparkException] { + readAllWith(path, Seq("k", "v"), filter, (b, i) => b.column(0).getLong(i), + capacity = 16)._2.close() + } + } + assert(thrown.getCause.isInstanceOf[TaskKilledException], + s"the kill must reach the task; got $thrown") + assert(!DataSourceUtils.shouldIgnoreCorruptFileException(thrown), + s"and ignoreCorruptFiles must not be able to swallow it; got $thrown") + } + } + + // Returns the FileSourceScanExec `df` plans to, with the storage filters `build` makes attached. + // `build` resolves a column by name against the scan's own output. Goes through Spark's normal + // planning/execution machinery (not the test-only reader init), so it exercises + // preparedStorageFilters' subquery materialization + bind, + // ParquetFileFormat.buildReaderWithStorageFilters, and metric propagation. + private def withStorageFilters(df: DataFrame)( + build: (String => Attribute) => Seq[Expression]): FileSourceScanExec = { + val scan = scanOf(df.queryExecution.executedPlan) + def attr(name: String): Attribute = + scan.output.find(_.name == name).getOrElse(fail(s"No $name in scan output")) + scan.copy(storageFilters = build(attr)) + } + + // The scan of `df` with a `k >= threshold` storage filter attached. + private def scanWithStorageFilter(df: DataFrame, threshold: Long): FileSourceScanExec = + withStorageFilters(df)(attr => Seq(GreaterThanOrEqual(attr("k"), Literal(threshold)))) + + // The same over `SELECT k, v` of the file at `path`. + private def scanWithStorageFilter(path: String, threshold: Long): FileSourceScanExec = + scanWithStorageFilter(spark.read.parquet(path).select("k", "v"), threshold) + + // 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 collectRows(plan: SparkPlan): Array[InternalRow] = { + val rowPlan = if (plan.supportsColumnar) ColumnarToRowExec(plan) else plan + rowPlan.executeCollect() + } + + // `collectRows` over a `(k, v)` output. + private def executePlanCollect(plan: SparkPlan): Array[(Long, String)] = + collectRows(plan).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, threshold = 195L) + val collected = executePlanCollect(scan).toSeq.sorted + val expected = rows.filter(_._1 >= 195L) + assert(collected == expected, s"got $collected; 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. The pushed filter drops pages, phase 1 reads the key over the pages it kept and + // narrows them with the storage filter, and phase 2 reads the non-key column under the + // intersection. Each filter alone keeps rows the other drops, so the scan returns exactly the + // rows both keep. + withTempDir { dir => + val rows = rows400 + val path = writePagedFile(dir) + assert(footerOf(path).getBlocks.size == 1, "the fixture must be one row group") + + withSQLConf(SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true") { + // Pages hold 100 rows, so this drops the first page and keeps the other three. + val df = spark.read.parquet(path).select("k", "v").filter("v > 'v_0100'") + // Keeps the first page, which the pushed filter drops, and drops the last 150 rows. + val withSF = + withStorageFilters(df)(attr => Seq(LessThanOrEqual(attr("k"), Literal(250L)))) + val pushedOnly = executePlanCollect(withSF.copy(storageFilters = Nil)).toSeq.sorted + assert(pushedOnly == rows.drop(100), + s"the pushed filter must drop the first page; got ${pushedOnly.size} rows") + + val collected = executePlanCollect(withSF).toSeq.sorted + assert(collected == rows.slice(100, 250), + s"expected keys 101..250, the rows both filters keep; got ${collected.size} rows") + val excludedWithin = + withSF.metrics(StorageFilterMetrics.ROWS_EXCLUDED_WITHIN_ROW_GROUP).value + assert(excludedWithin == 150, + s"the storage filter excludes 150 of the 300 rows the pushed filter kept; got " + + excludedWithin) + } + } + } + + // ----- 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 + } + + // The rows the scans with storage filters excluded, by both row metrics. Through `collect(plan)`, + // which finds those scans in a plan AQE wrapped too. + private def rowsExcludedByStorageFilters(plan: SparkPlan): Long = + collect(plan) { + case s: FileSourceScanExec if s.storageFilters.nonEmpty => + s.metrics(StorageFilterMetrics.ROWS_EXCLUDED_BY_ROW_GROUP).value + + s.metrics(StorageFilterMetrics.ROWS_EXCLUDED_WITHIN_ROW_GROUP).value + }.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 + } + } + } + + Seq(false, true).foreach { ignoreCorruptFiles => + test("FileSourceStrategy offers the bloom to the scan and keeps it post-scan too " + + s"(ignoreCorruptFiles = $ignoreCorruptFiles)") { + // 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. ignoreCorruptFiles does not turn + // the feature off. + withBloomFilterTables { + withSQLConf( + SQLConf.PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED.key -> "true", + SQLConf.IGNORE_CORRUPT_FILES.key -> ignoreCorruptFiles.toString, + 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))).toSeq == Seq((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 excluded = rowsExcludedByStorageFilters(plan) + assert(excluded > 0, s"the reader must have excluded rows; the metrics say $excluded") + } + } + } + } + + 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)).toSeq == Seq(5L), + "and the query must still return the joined key") + } + } + } + + test("a row group over the splice cap reads its key columns again: same rows, more bytes, " + + "less avoided") { + // Past the cap a row group gives splicing up. Phase 2 takes every projected column, key + // columns included, so nothing is buffered. Three things have to hold. The rows must not + // change, since a fallback that quietly dropped the predicate would return extra rows. The + // fallback must actually have happened, which the second read of the key column shows as + // strictly more bytes transferred. And that second read is a cost against the saving, so the + // byte metric must report less avoided than while splicing. + // + // 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 that gave splicing up still has to emit + // in batches, and 51 survivors over a capacity of 16 span several of them. + withTempDir { dir => + val path = writePagedFile(dir) + + def run(maxSplicedBytes: String, threshold: Long): (Int, Long, Long) = { + val metrics = metricMap() + val (rows, bytes) = readUnderCap(path, kvSchema(), Seq(keyAtLeast(threshold)), + maxSplicedBytes, batchSize = 16, metrics) + (rows, bytes, metrics(StorageFilterMetrics.BYTES_AVOIDED_BY_PAGE_FILTERING).value) + } + + val (splicedRows, splicedBytes, splicedAvoided) = run("64MB", 350L) + // `k >= 350` keeps one contiguous range of 40 bytes, and a long key buffers 9 bytes per + // survivor, so a 100-byte cap is passed at the seventh. + val (plainRows, plainBytes, plainAvoided) = run("100b", 350L) + 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") + assert(splicedAvoided > 0, s"page filtering has to avoid something here; got $splicedAvoided") + assert(plainAvoided < splicedAvoided, + s"the second key read must count against the saving; plain=$plainAvoided " + + s"spliced=$splicedAvoided") + + // The two halves of the budget are weighed together, and this is the case that says so. + // Eight survivors buffer 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) = run("64MB", 393L) + val (_, _, eightPlain) = run("100b", 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") + + // The exact charge. Seven survivors in one range cost 7 * 9 + 40 = 103 bytes, which a cap of + // 103 holds and one of 102 does not. A charge one byte off either way flips one of the two. + assertSpliceCharge(103L, rows = 7) { cap => + val (rows, bytes, _) = run(cap, 394L) + (rows, bytes) + } + } + } + + 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. Each row + // group's path is asserted too, as whether its batches hold the persistent key vector, which + // also pins the charge: at twice the long key's 9 bytes the first row group would pass the cap. + withTempDir { dir => + val rows = rows400 + val path = writeParquetFile(dir, rows, rowGroupSize = 256L) + assert(footerOf(path).getBlocks.size == 4, "the fixture must be four row groups") + 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))), + kvSchema(), + m, + maxSplicedRowGroupBytes = 150L) + val reader = new LateMaterializationParquetRecordReader(false, 16, filter) + try { + reader.initialize(path, Seq("k", "v").asJava) + reader.initBatch(new StructType(), null) + // The batch holds the persistent vectors until the first batch splices. + val persistentKeys = reader.resultBatch().column(0) + val result = mutable.ArrayBuffer[(Long, String)]() + val splicedByRowGroup = mutable.LinkedHashMap[Long, Set[Boolean]]() + while (reader.nextBatch()) { + val batch = reader.resultBatch() + val spliced = batch.column(0) ne persistentKeys + (0 until batch.numRows()).foreach { i => + val key = batch.column(0).getLong(i) + result += ((key, batch.column(1).getUTF8String(i).toString)) + val rowGroup = (key - 1) / 100 + splicedByRowGroup(rowGroup) = splicedByRowGroup.getOrElse(rowGroup, Set()) + spliced + } + } + 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(splicedByRowGroup == Map(0L -> Set(true), 1L -> Set(false), 2L -> Set(true)), + s"rows 1-100 and 201-300 must splice and rows 101-200 must not; got $splicedByRowGroup") + assert(m.rowGroupsSkipped.value == 1, + s"expected the row group with no survivor skipped; got ${m.rowGroupsSkipped.value}") + } finally { + reader.close() + } + } + } + + // The later-set tests' file: twenty string keys, `0001` to `0020` padded with `x` to 1000 bytes, + // the last two to `lastTwoBytes`, in one row group. Returns it with the rows `laterSetFilter` + // keeps. + private def writeLaterSetFile(dir: File, lastTwoBytes: Int): (String, Seq[(String, String)]) = { + val df = spark.range(1, 21).selectExpr( + "CONCAT(LPAD(CAST(id AS STRING), 4, '0'), " + + s"REPEAT('x', IF(id >= 19, ${lastTwoBytes - 4}, 996))) AS k", + "CAST(id AS STRING) AS v") + val path = writeSingleParquetFile(dir, df, rowGroupSize = 64 * 1024L) + assert(footerOf(path).getBlocks.size == 1, "the fixture must be one row group") + val kept = (15 to 20).map { id => + (f"$id%04d" + "x" * ((if (id >= 19) lastTwoBytes else 1000) - 4), id.toString) + } + (path, kept) + } + + private val laterSetSchema: StructType = kvSchema(keyType = StringType) + + // `k >= "0015"`, which keeps the last six keys. Over a capacity of 4 that is one set of four, + // then two more. + private val laterSetFilter: Seq[Expression] = Seq(GreaterThanOrEqual( + BoundReference(0, StringType, nullable = true), Literal.create("0015", StringType))) + + // Reads a later-set file over a capacity of 4 under `cap`, and returns the rows. + private def readLaterSet(path: String, cap: Long): Seq[(String, String)] = { + val (rows, reader) = readAllWith( + path, Seq("k", "v"), createFilter(laterSetFilter, laterSetSchema, allMetrics(), cap), + (b, i) => (b.column(0).getUTF8String(i).toString, b.column(1).getUTF8String(i).toString), + capacity = 4) + reader.close() + rows + } + + test("a later accumulator set of a variable-length key is charged per value") { + // Six survivors over a capacity of 4 fill one set of four 1000-byte keys and a second of two + // 2000-byte ones. Every set grows as its values arrive, the way the first one does. Each value + // is charged four times its bytes plus its row's fixed part of 17, and the one range adds 40. + // So the row group is charged 4 * (4000 + 17) + 2 * (8000 + 17) + 40 = 32142 bytes. A cap of + // 32142 keeps splicing and one of 32141 gives it up, which a missing or a double charge would + // each get wrong. And the keys must come back right whichever path the row group took. + withTempDir { dir => + val (path, expected) = writeLaterSetFile(dir, lastTwoBytes = 2000) + assertSpliceCharge(32142L, rows = 6) { cap => + readUnderCap(path, laterSetSchema, laterSetFilter, cap, batchSize = 4) + } + Seq(Long.MaxValue, 32142L, 32141L).foreach { cap => + val rows = readLaterSet(path, cap) + assert(rows == expected, + s"every survivor must come back with its own key under cap $cap; got ${rows.map(_._2)}") + } + } + } + + test("what a row group of scattered survivors is charged does not depend on the cap") { + // Keys 9 to 12 fill the first set over a capacity of 4, and keys 14 and 17 come scattered, each + // in a range of its own. Every set grows as values arrive and is charged per value, so the row + // group's charge is one number: the first set's 16068, three ranges of 40 and the second + // set's 8034, 24222 in all. A cap of it splices and one byte under it gives splicing up. + // Larger caps splice too, which a few of them check, among them 32216 to 32221. Those gave + // splicing up when a later set was sized from the one before it with no room for its ranges. + // The second set is not sized at all, so its byte child grows to 2000 bytes for its two + // 1000-byte values, where sizing it from the first set would have reserved 8000. + withTempDir { dir => + val (path, _) = writeLaterSetFile(dir, lastTwoBytes = 1000) + val prefix = Substring(BoundReference(0, StringType, nullable = true), Literal(1), Literal(4)) + val keys = Seq("0009", "0010", "0011", "0012", "0014", "0017").map(Literal(_)) + def readUnder(cap: String): (Int, Long) = + readUnderCap(path, laterSetSchema, Seq(In(prefix, keys)), cap, batchSize = 4) + assertSpliceCharge(24222L, rows = 6)(readUnder) + val (_, uncapped) = readUnder("64MB") + Seq(32215L, 32216L, 32221L, 32222L).foreach { cap => + val (rows, bytes) = readUnder(s"${cap}b") + assert(rows == 6 && bytes == uncapped, + s"a cap of $cap must keep the six rows and splicing; got $rows rows") + } + val (children, reader) = readAllWith(path, Seq("k", "v"), + createFilter(Seq(In(prefix, keys)), laterSetSchema), + (b, _) => ParquetTestAccess.byteChildCapacity( + b.column(0).asInstanceOf[WritableColumnVector]), + capacity = 4) + reader.close() + assert(children.drop(4).distinct == Seq(2000), + s"the second set's byte child must grow as values arrive; got $children") + } + } + + test("two key columns are both charged") { + // Seven survivors in one range cost 7 * (9 + 9) + 40 = 166 bytes with two long keys. A cap of + // 166 keeps splicing and one of 165 gives it up, so charging one key alone would show. + withTempDir { dir => + val df = spark.range(1, 401).selectExpr("id AS k", "id AS k2", "CAST(id AS STRING) AS v") + val path = writeSingleParquetFile(dir, df, rowGroupSize = 64 * 1024L, pageSize = Some(512L)) + val schema = StructType(Seq(StructField("k", LongType), StructField("k2", LongType), + StructField("v", StringType))) + val storageFilters = Seq(And( + GreaterThanOrEqual(BoundReference(0, LongType, nullable = true), Literal(394L)), + GreaterThanOrEqual(BoundReference(1, LongType, nullable = true), Literal(0L)))) + assertSpliceCharge(166L, rows = 7) { cap => + readUnderCap(path, schema, storageFilters, cap, batchSize = 16) + } + } + } + + // 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 footer check sees the missing index + // for every row group before anything is read, so no read fails here. The middle row group is + // emptied by the filter and skipped whole, and the other two keep rows and are read whole. + // Spark's own writer always writes the page index, hence the hand-built file. The retry after a + // failing read is covered by the test on a file whose index parquet cannot parse. + 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 = writeParquetFileByHand(dir, blocks, hasOffsetIndex = (_, _) => false) + // 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 = kvSchema(valueType = LongType) + // 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) + val ((result, reader), warnings) = withLogged(partialReadWarning) { + readAllWith(path, Seq("k", "v"), filter, + (b, i) => (b.column(0).getLong(i), b.column(1).getLong(i))) + } + 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(warnings == 0, + s"the footer check must get there before the exception path; got $warnings") + 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 is declined by the footer check, not by the read") { + // 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, which + // wastes the buffer and reads the row group again whole. The rows would not show that, since + // the row group comes back whole either way. The WARN the exception path logs is what does. + // + // 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 = kvSchema(valueType = 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(keyAtLeast(50L)) + val (values, warnings) = withLogged(partialReadWarning) { + readLongs(formatReader(schema, pushed, storageFilters), path, ordinal = 1) + } + // The footer check gets there before the phase-2 read does, which is what keeps the rest of + // the file's page filtering. This WARN is what the exception path would have reported. + assert(warnings == 0, "the footer check must decline the row group rather than the read " + + s"discovering it; got $warnings") + // The row group that cannot be narrowed comes back whole, since its filter is given up and + // the empty store widens its ranges to the block. + assert((1L to 100L).forall(values.contains), + s"the first row group must come back whole; got ${values.count(_ <= 100L)} of its rows") + // The other two keep their page pruning. Read whole they would contribute 200 rows rather + // than one page each. + assert(values.size < 300, + s"the row groups that can be narrowed must still be; got all ${values.size} rows") + // And every row the pushed filter actually matches survives the narrowing. + assert(Seq(1L, 101L, 201L).forall(values.contains), + s"the matching rows must be in the output; got ${values.size} rows") + } + } + + // Overwrites the bytes of one chunk's offset index while leaving the footer's reference to it. + // Parquet then reads the reference, fails to parse what it points at, and reports the whole block + // as having no column index, which is the one file shape that reaches phase 2's retry. + private def breakOffsetIndex(filePath: String, blockIdx: Int, columnIdx: Int): Unit = { + val chunk = footerOf(filePath).getBlocks.get(blockIdx).getColumns.get(columnIdx) + val ref = chunk.getOffsetIndexReference + assert(ref != null, s"block $blockIdx column $columnIdx has no offset index to damage") + overwriteBytes(filePath, ref.getOffset, ref.getLength) + } + + // Overwrites the back half of one data page of `column` in the first row group, which leaves the + // page header, and so everything parquet reads eagerly, intact and the page body undecodable. + private def corruptDataPage(filePath: String, column: String, pageIdx: Int): Unit = { + val (offset, size) = Utils.tryWithResource(ParquetFileReader.open( + HadoopInputFile.fromPath(new Path(filePath), spark.sessionState.newHadoopConf()))) { + reader => + val chunk = reader.getRowGroups.get(0).getColumns.asScala + .find(_.getPath.toDotString == column).getOrElse(fail(s"no $column chunk")) + val index = reader.readOffsetIndex(chunk) + (index.getOffset(pageIdx), index.getCompressedPageSize(pageIdx)) + } + overwriteBytes(filePath, offset + size / 2, size - size / 2) + } + + // Rewrites the dictionary page header of `column` in the first row group to hold only `entries` + // values, so the ids past them point at nothing. The page itself is untouched, and a column that + // decodes lazily meets the damage only when one of those values is read. + private def shrinkDictionary(filePath: String, column: String, entries: Int): Unit = { + val chunk = footerOf(filePath).getBlocks.get(0).getColumns.asScala + .find(_.getPath.toDotString == column).getOrElse(fail(s"no $column chunk")) + assert(chunk.hasDictionaryPage, s"$column has no dictionary page to damage") + val bytes = new Array[Byte](256) + val raf = new RandomAccessFile(filePath, "r") + try { + raf.seek(chunk.getDictionaryPageOffset) + raf.readFully(bytes) + } finally { + raf.close() + } + val in = new ByteArrayInputStream(bytes) + val header = Util.readPageHeader(in) + val headerLength = bytes.length - in.available() + assert(header.getDictionary_page_header.getNum_values > entries, + s"$column's dictionary must hold more than $entries values") + header.getDictionary_page_header.setNum_values(entries) + val rewritten = new ByteArrayOutputStream() + Util.writePageHeader(header, rewritten) + assert(rewritten.size() == headerLength, "the rewritten header must keep its length") + writeBytes(filePath, chunk.getDictionaryPageOffset, rewritten.toByteArray) + } + + // Overwrites `length` bytes at `offset` with 0xFF. + private def overwriteBytes(filePath: String, offset: Long, length: Int): Unit = + writeBytes(filePath, offset, Array.fill(length)(0xFF.toByte)) + + // Writes `bytes` over the file at `offset`. + private def writeBytes(filePath: String, offset: Long, bytes: Array[Byte]): Unit = { + val file = new File(filePath) + val raf = new RandomAccessFile(file, "rw") + try { + raf.seek(offset) + raf.write(bytes) + } finally { + raf.close() + } + // Hadoop's local filesystem checks a side-car checksum that the write above invalidates, and a + // mismatch there would fail the read for the wrong reason. + val crc = new File(file.getParentFile, s".${file.getName}.crc") + if (crc.exists()) assert(crc.delete(), s"could not remove $crc") + } + + test("a corrupt key page gives the row group up and fails where a plain read fails") { + // Parquet reads a chunk's bytes up front but decodes a page only when it is read, so a corrupt + // key page shows in phase 1. When its decoding raises an exception, as it does here, giving the + // row group up hands it to phase 2 whole, which reads the same pages in the same batches as a + // plain read, so the page fails at the same batch, after the same rows. That is what makes + // `ignoreCorruptFiles` keep exactly what it keeps of a plain read, and it is asserted here on + // the readers themselves, as the rows each hands back before it throws. The filter is + // selective, so a reader that kept applying it would return fewer. A codec that reports + // corruption as an `InternalError`, as hadoop-lzo does, takes the same path. + withTempDir { dir => + val rows = (0L until 1000L).map(i => (i, s"v_$i")) + // One row group of 100-row pages, so the fifth key page covers rows 400 to 499. + val path = writeSingleParquetFile(dir, rows.toDF("k", "v"), rowGroupSize = 64L * 1024 * 1024, + options = Map(ParquetOutputFormat.PAGE_ROW_COUNT_LIMIT -> "100")) + corruptDataPage(path, column = "k", pageIdx = 4) + + // The rows a reader hands back before it throws, what it throws, and, for the filtering + // reader, whether its scratch vectors were replaced on the way and the old ones freed. Both + // read off heap, where an old scratch vector left open would leak native memory. + def readUntilFailure(storageFilter: ParquetStorageFilter): (Seq[Long], Throwable, Boolean) = { + val reader = if (storageFilter == null) { + new VectorizedParquetRecordReader(true, 16) + } else { + new LateMaterializationParquetRecordReader(true, 16, storageFilter) + } + def scratch: Seq[WritableColumnVector] = reader match { + case filtering: LateMaterializationParquetRecordReader => + ParquetTestAccess.keyScratchVectors(filtering) + case _ => Nil + } + val read = mutable.ArrayBuffer[Long]() + try { + reader.initialize(path, Seq("k", "v").asJava) + reader.initBatch(new StructType(), null) + val scratchBefore = scratch + val error = try { + while (reader.nextBatch()) { + val batch = reader.resultBatch() + (0 until batch.numRows()).foreach(i => read += batch.column(0).getLong(i)) + } + null + } catch { + case NonFatal(e) => e + } + val replacedAndFreed = scratch.zip(scratchBefore).exists { case (now, before) => + (now ne before) && address(before) == 0 + } + (read.toSeq, error, replacedAndFreed) + } finally { + reader.close() + } + } + + val (plainRows, plainError, _) = readUntilFailure(null) + assert(plainError != null, "the damaged page must fail a plain read, or this proves nothing") + assert(plainRows == (0L until 400L), s"a plain read returns the rows before that page's " + + s"batch; got ${plainRows.size}") + val (filteredRows, filteredError, scratchReplaced) = + readUntilFailure(keyAtLeastFilter(500L)) + assert(filteredError != null && filteredError.getClass == plainError.getClass, + s"the filtering read must fail the same way; got $filteredError, not $plainError") + assert(filteredRows == plainRows, + s"and after the same rows; got ${filteredRows.size}, not ${plainRows.size}") + // A failed decode can leave the scratch vectors grown. Phase 2 reads the same pages into + // vectors of its own, so the reader replaces them with capacity-sized ones rather than keep + // the grown ones for the rest of the split. + assert(scratchReplaced, + "the decoding fallback must replace the scratch vectors and free the old ones") + + // A kill that shows only at the decoding catch, after the per-chunk check has passed, is + // seen there. Without that catch's own check the row group would fall back and phase 2 + // would fail on the page instead. + val killedAtTheCatch = spy(TaskContext.empty()) + doAnswer((_: InvocationOnMock) => Boolean.box(Thread.currentThread().getStackTrace + .exists(_.getMethodName.startsWith("rethrowIfMustPropagate")))) + .when(killedAtTheCatch).isInterrupted() + val (_, killed, _) = TaskContext.withTaskContext(killedAtTheCatch) { + readUntilFailure(keyAtLeastFilter(500L)) + } + assert(killed.isInstanceOf[SparkException] && killed.getMessage.contains("task kill") && + killed.getCause != null && killed.getCause.getClass == plainError.getClass, + s"the kill must go out from the decoding catch, around the decoding error; got $killed") + assert(!DataSourceUtils.shouldIgnoreCorruptFileException(killed), s"got $killed") + } + } + + test("a key vector that cannot grow, or a codec's InternalError, gives the row group up") { + // Both are fatal-looking errors that decoding the key pages can raise, and the reader falls + // back on both. A scratch vector that cannot grow is phase 1's own, and phase 2 then reads + // what a plain read reads. An `InternalError` from decoding is how a codec like hadoop-lzo + // reports a corrupt page, which the plain read of the same pages meets again. Raised from the + // filter instead, the first goes out as its OutOfMemoryError and the second in an + // `UnknownError`. Off heap, the scratch vectors the failed decode used must be freed, and the + // second row group's phase 1 runs over the ones that replace them. + withTempDir { dir => + val rows = (0L until 40L).map(i => (i, i.toString)) + val path = writeSingleParquetFile(dir, rows.toDF("k", "v"), + rowGroupSize = ParquetWriter.DEFAULT_BLOCK_SIZE, + options = Map(ParquetOutputFormat.BLOCK_ROW_COUNT_LIMIT -> "20")) + assert(footerOf(path).getBlocks.asScala.map(_.getRowCount) == Seq(20L, 20L), + "the fixture must be two row groups of 20 rows") + Seq( + new RuntimeException("Cannot reserve", new OutOfMemoryError("reserve")), + new InternalError("codec")).foreach { error => + var failedScratch: Seq[WritableColumnVector] = Nil + // Only the first row group's key reader fails. + val failing: KeyReaderHook.OnKeyReader = (reader, keyReader) => { + val scratch = ParquetTestAccess.keyScratchVectors( + reader.asInstanceOf[LateMaterializationParquetRecordReader]) + if (failedScratch.nonEmpty) { + // Checked before phase 1 writes into them, since a freed one would crash the JVM. + assert(scratch.forall(v => !failedScratch.exists(_ eq v) && address(v) != 0), + "the second row group must decode into new scratch vectors") + keyReader + } else { + failedScratch = scratch + val failingReader = spy(keyReader) + doThrow(error).when(failingReader).readBatch(anyInt(), any(), any(), any()) + failingReader + } + } + val ((result, reader), warnings) = try { + withLogged(decodingGivenUp) { + readAll(path, keyAtLeastFilter(30L), useOffHeap = true, onKeyReader = failing) + } + } catch { + // ScalaTest would abort the whole suite on these rather than fail this test. + case e: VirtualMachineError => fail(s"the reader must fall back on $error; got $e", e) + } + try { + assert(result == rows.take(20) ++ rows.drop(30), + s"the first row group must come back whole and the second filtered on $error; got " + + s"${result.map(_._1)}") + assert(warnings == 1, s"the decoding fallback must report $error; got $warnings") + assert(failedScratch.forall(v => address(v) == 0), + s"the scratch vectors of the failed decode must be freed on $error") + } finally { + reader.close() + } + } + } + } + + test("a key value the filter does not read and cannot be decoded gives splicing up") { + // The filter need not read every key on every row, as `coalesce(a, b)` reads `b` only where + // `a` is null. The copy into the survivor buffer reads every key, so it can meet a value no + // plain read decodes. A dictionary-encoded column decodes a value only when it is read, so a + // dictionary that holds fewer values than its ids point at is such a value. Copying it gives + // splicing up for the row group, and phase 2 reads the key columns the way a plain read does. + // So the read answers where a plain one does, and the filter still applies. + withTempDir { dir => + // Two copies of each `b`, so parquet keeps its dictionary. Cut to 64 values, it leaves `b` + // undecodable from row 128 on. `a` is null only before that, so the filter never reads a + // damaged `b`. + val rows = (0 until 200).map { i => + (if (i < 100 && i % 5 == 0) null else s"a_$i", s"b_${i / 2}", i) + } + val path = writeSingleParquetFile(dir, rows.toDF("a", "b", "v"), + rowGroupSize = ParquetWriter.DEFAULT_BLOCK_SIZE, dictionary = true) + assert(encodingsOf(path, "b").exists(_.usesDictionary), + s"the fixture's whole point; got ${encodingsOf(path, "b")}") + shrinkDictionary(path, column = "b", entries = 64) + val columns = Seq("a", "b", "v") + intercept[ArrayIndexOutOfBoundsException] { + readAllWith(path, columns, null, (b, i) => b.column(1).getUTF8String(i).toString)._2 + .close() + } + + val schema = StructType(Seq( + StructField("a", StringType), StructField("b", StringType), StructField("v", IntegerType))) + val coalesced = Coalesce(Seq(BoundReference(0, StringType, nullable = true), + BoundReference(1, StringType, nullable = true))) + val filter = createFilter( + Seq(GreaterThanOrEqual(coalesced, Literal.create("a_15", StringType))), schema) + val ((result, reader), givenUp) = withLogged(splicingGivenUpOnError) { + readAllWith(path, columns, filter, (b, i) => b.column(2).getInt(i), capacity = 16) + } + try { + // Keeps rows 150-199 among others, whose `b` is damaged. + val expected = rows.filter { case (a, b, _) => Option(a).getOrElse(b) >= "a_15" }.map(_._3) + assert(expected.contains(150), "the filter must keep a row whose `b` is damaged") + assert(result == expected, s"got ${result.size} rows, expected ${expected.size}") + assert(givenUp == 1, + "the copy must have failed and given splicing up once, or this proves nothing; " + + s"logged $givenUp") + } finally { + reader.close() + } + } + } + + test("a row group whose offset index parquet cannot parse leaves the rest of the file its " + + "page pruning") { + // The retry reads the damaged row group whole over the pushed data filter's ranges, and only + // that one, since parquet keeps a column index store per block. What it must not give up is + // the pushed data filter's own page pruning, which a plain read of the same query keeps, nor + // the other row groups' narrowing. Dropping either would read more than it needs to. + // + // One file shape reaches the retry, a footer that references an offset index parquet cannot + // then parse, so the fixture is hand-built and then damaged. + 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, (_, _) => true) + breakOffsetIndex(path, blockIdx = 1, columnIdx = 0) + // The fixture's whole point. The footer still offers the index the read cannot use, which is + // what gets past the footer check and into the retry. + assert(footerOf(path).getBlocks.get(1).getColumns.get(0).getOffsetIndexReference != null, + "the damaged index must still be referenced, or the footer check declines the row group") + val schema = kvSchema(valueType = LongType, nullable = false) + // One value per row group, so parquet's statistics keep every one of them and its column + // index narrows each to the page holding that value. All three are even, so all three also + // survive the storage filter below. + val pushed = Seq(sources.In("v", Array[Any](2L, 102L, 202L))) + val evenKeys = Seq(EqualTo( + Remainder(BoundReference(0, LongType, nullable = false), Literal(2L)), Literal(0L))) + def read(storageFilters: Seq[Expression]): Seq[Long] = + readLongs(formatReader(schema, pushed, storageFilters), path, ordinal = 0) + val (filtered, warnings) = withLogged(partialReadWarning)(read(evenKeys)) + assert(warnings == 1, + "the read must have discovered the damaged index, or this test proves nothing; got " + + warnings) + // A plain read of the same query is the bar. The filtering one must never hand the post-scan + // Filter more rows than that. + val plain = read(Nil) + assert(filtered.size <= plain.size, + s"the filtering read must not read more than a plain one; got ${filtered.size} rows " + + s"against ${plain.size}") + // The sharp form of the same thing. The last row group's index is intact, so its pages are + // still pruned after the retry. + assert(filtered.count(_ > 200L) < 100, + s"the last row group must keep its pruning; got ${filtered.count(_ > 200L)} of its 100") + // The damaged row group comes back whole, which is the cost that is inherent to it. + assert((101L to 200L).forall(filtered.contains), + s"the damaged row group must come back whole; got " + + s"${filtered.count(k => k > 100L && k <= 200L)} of its 100") + // And every row the pushed filter matches is in the output, all three being even. + assert(Seq(2L, 102L, 202L).forall(filtered.contains), Review Comment: Thanks, done. The test compares the exact ordered output against that, which also replaces the size and membership checks it covers. ########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java: ########## @@ -0,0 +1,1032 @@ +/* + * 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.metadata.BlockMetaData; +import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore.MissingOffsetIndexException; + +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.DecimalType; +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 has passed + * {@code spark.sql.parquet.storageFilterPushdown.maxSplicedRowGroupBytes} gives splicing up, and + * phase 2 reads its key columns again with the rest. If the survivors' row ranges alone pass it, + * the filter is given up for that row group too. + * + * <p>A whole file is read without splicing in these cases: + * <ul> + * <li>every key column is missing from it 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;</li> + * <li>the file has none of the projected non-key columns, or none is projected. Phase 2 would + * have nothing to read, so the reader declines the filter and reads the file plainly;</li> + * <li>{@code parquet.filter.columnindex.enabled} is false, so its page index is not used. The + * filter then only skips row groups it rejects whole.</li> + * </ul> + * + * <p>A row group whose columns do not all carry an offset index gets the same as the last case, + * and so does one whose offset index parquet fails to read. + * + * <p>{@link #resultBatch()} is one batch object for the whole read, as for the plain reader, but a + * spliced key slot holds a different vector in each batch, and {@link #nextBatch()} frees the + * previous one. A caller must fetch {@code batch.column(i)} again after every {@code nextBatch()}. + * The slots are swapped rather than copied into the persistent vectors on purpose. A copy would + * pay a second value copy per surviving row, a byte copy for a variable-length key, where the + * swap is one array write per key column per batch. + * + * <p>What {@code FileSourceStrategy.storageFiltersFor} and {@code ParquetStorageFilter.create} + * guarantee is asserted here, since a violation is a planner bug. + */ +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, by default. A factor rather than an exact figure, so with that default this + * half of the charge over-counts rather than under-counts, since nothing above this buffer can + * spill it. A {@code hugeVectorReserveRatio} above 2 makes it under-count for a child past + * {@code hugeVectorThreshold}. That threshold is off by default. 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; + + /** + * What one surviving row range costs to hold, which is the half of the budget a row group + * buffering nothing still pays. Scattered survivors make one range each, and phase 2 needs the + * whole set to select its pages. One set is held, which every column reader of the row group + * walks with a cursor of its own. The figure is parquet's {@code RowRanges.Range}, two longs, + * plus its object header and a list slot. That is what is held. At the end of phase 1 the + * builder's list can have up to half as many slots again, and {@code build()} copies it, so the + * peak is somewhat higher. + */ + private static final long ROW_RANGE_BYTES = 40L; + + /** + * The byte child a variable-length vector is allocated with, per row, which is + * `WritableColumnVector.DEFAULT_ARRAY_LENGTH`. It is part of `keyFixedBytesPerRow`, counted + * twice, for its data byte and its null byte. + */ + private static final int DEFAULT_CHILD_BYTES_PER_ROW = 4; + + /** + * The filter this reader applies. {@link #keyColumns} is null for a file it is not applied to row + * by row, and such a file is read the way the plain reader reads it. + */ + private final ParquetStorageFilter storageFilter; + + /** + * Late-materialization state, set by {@link #initializeLateMaterialization()}. + * + * <p>{@link #fileReader} drives all three phases, and its requested schema is switched per phase. + * Phase 0 reads the whole projection, phase 1 the key columns, and phase 2 the non-key columns + * while splicing or the whole projection otherwise. + * + * <p>Held as leaf-column lists, since that is what the reader and the byte metrics take and + * {@code MessageType.getColumns()} rebuilds its list per call. The leaf paths of the whole + * projection and of the non-key columns are resolved once, for the footer walks that run per row + * group. The key columns need none, since phase 1 always reads them and no walk measures them + * apart. + */ + private List<ColumnDescriptor> requestedColumns; + private List<ColumnDescriptor> keyOnlyColumns; + private List<ColumnDescriptor> nonKeyColumns; + private Set<ColumnPath> requestedPaths; + private Set<ColumnPath> nonKeyPaths; + /** Answers from the footer whether a row group can be read in part, and what that transfers. */ + 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 does not + * close the stores {@code readFilteredRowGroup} hands out. + */ + private PageReadStore dataPages; + + /** The messages this reader has reported, each once per split. */ + private final Set<String> reportedMessages = new HashSet<>(); + + /** + * Whether {@code parquet.filter.columnindex.enabled} is false, so this file's page index must not + * be used. That turns off phase 0 and every partial read. A wrong index would pair a row's key + * with another row's values, which the post-scan Filter cannot catch. Parquet's partial read does + * not check the conf itself. Read off {@link #readOptions}, so it agrees with the row count + * {@link #fileReader} reported. + */ + private boolean pageIndexDisabled; + + /** + * Whether the row group currently loaded is spliced. It starts as whether the row group can be + * read in part, and {@link #abandonSplicing} clears it. False means phase 2 read every projected + * column, so the emit path takes the key columns from the persistent batch. + */ + private boolean spliceCurrentRowGroup; + private int nextBlockIndex; + /** + * The key columns this file has, in the filter's key order, with what init resolves about each. + * Null when the filter is not applied to this file row by row. + */ + private KeyColumn[] keyColumns; + /** Whether this file has none of the key columns, decided in {@link #initBatch}. */ + private boolean allKeysMissing; + private WritableColumnVector[] keyScratchVectors; + private ColumnarBatch keyScratchBatch; + + /** + * Splicing state. Phase 1 keeps the surviving key values it decoded, and the emit path splices + * them into the output batch. That saves phase 2 a second read of the key columns, which parquet + * does not cache, so it goes back to the filesystem. + * + * <p>{@link #survivorBatches} holds one array of survivor vectors per batch to emit, indexed like + * {@link #keyColumns}, and {@link #currentKeyAccumulators} the array still filling. Together they + * hold every survivor of one row group at once, outside any MemoryConsumer. + */ + private final ArrayDeque<WritableColumnVector[]> survivorBatches = new ArrayDeque<>(); + /** The survivor vectors still filling, or null when there are none. */ + private WritableColumnVector[] currentKeyAccumulators; + /** Row count of {@link #currentKeyAccumulators}. */ + private int currentKeyAccumulatorRowCount; + /** + * The survivor vectors the current batch is built on. This reader owns them until the next batch + * releases them, so a consumer that stops between batches, such as a limit or the end of the + * task, leaves them for {@link #close()} to free. + */ + private WritableColumnVector[] publishedKeyVectors; + /** {@link #initBatch}'s vectors, which this reader closes. */ + private ColumnVector[] persistentBatchColumns; + /** The array the emitted batch is built over, whose key slots each batch rewrites. */ + private ColumnVector[] spliceBatchColumns; + + 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"); + } + + /** + * Releases what this reader allocated, then what the base owns. Chained through finally blocks, + * so one failing close does not leak the rest. + */ + @Override + public void close() throws IOException { + try { + try { + if (persistentBatchColumns != null) { + // Through the persistent vectors rather than the emitted batch, whose key slots may hold + // survivor vectors the splicing state closes instead. Super then has no batch to close. + columnarBatch = null; + closeAll(persistentBatchColumns); + persistentBatchColumns = null; + spliceBatchColumns = 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 { + try { + closeDataPages(); + } finally { + super.close(); + } + } + } + } + + /** + * Builds the batch as the plain reader does and takes it over for splicing. A file with none of + * the key columns is decided here, because the constant vectors the predicate reads for those + * columns exist only once super has built them. + */ + @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. For a file the filter is applied to row by row, it + * then hands out a batch of its own over a copy of that array, whose key slots the emit path + * rewrites in place, and allocates phase 1's key vectors. {@link ColumnarBatch} holds the array + * by reference, its staging row included, so rewriting a slot is what publishes it. + * + * <p>Every slot got a vector, key columns included. A spliced row group leaves a key slot's + * vector unused, and a row group that gave splicing up reads into it. + */ + 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); + allocateKeyScratch(); + } + + /** + * Answers the filter for a file with none of its key columns, which makes the predicate constant + * over the file. + * + * <p>False skips the file, which then yields no rows. 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() { + // Not closed, and must not be, since the vectors in it belong to the batch this reader emits. + ColumnarBatch keyRow = new ColumnarBatch(keyVectorsFromBatch(), 1); + // Resolved outside the `try`, so a broken invariant behind it fails rather than falls back. + BasePredicate predicate = storageFilter.preparedPredicate(); + boolean keep; + try { + keep = predicate.eval(keyRow.getRow(0)); + } catch (Throwable t) { + giveUpOnError(t, FILE_GIVEN_UP_ON_EVALUATION, false); + keep = true; + } + if (!keep) { + recordFileSkipped(); + totalRowCount = 0; + } + } + + /** + * Releases the survivor vectors the previous batch was built on, reads the next batch the way the + * plain reader does, and then points the emitted batch's key slots at this batch's vectors. A + * spliced row group takes them from the head of the survivor queue, and any other row group from + * the persistent vectors, which is what a row group following a spliced one needs. Non-key slots + * keep {@link #initBatch}'s vectors for the reader's life. + * + * <p>The row-index column the base fills is never a key column, since + * {@code ParquetStorageFilter.create} rejects its name. So it always lands in a persistent + * vector. + */ + @Override + public boolean nextBatch() throws IOException { + closeAll(publishedKeyVectors); + publishedKeyVectors = null; + if (!super.nextBatch()) return false; + if (keyColumns != null) { + if (spliceCurrentRowGroup) { + publishedKeyVectors = survivorBatches.pollFirst(); + if (publishedKeyVectors == null) { + // Unreachable, since the queue holds exactly the survivors phase 1 accumulated and the + // emit loop is driven by that same count. Without this, the key slots would point at the + // persistent vectors, which a spliced row group resets but never reads into, so the batch + // would pair stale key values with this row group's rows as if they were real ones. + throw ParquetStorageFilter.internalError(String.format( + "Storage-filter survivor queue of row group %d in %s ran out with %d rows still " + + "to emit", nextBlockIndex - 1, fileReader.getFile(), columnarBatch.numRows())); + } + } + for (int k = 0; k < keyColumns.length; k++) { + int slot = keyColumns[k].batchSlot(); + spliceBatchColumns[slot] = + publishedKeyVectors != null ? publishedKeyVectors[k] : persistentBatchColumns[slot]; + } + } + return true; + } + + @Override + protected void initializeInternal() throws IOException, UnsupportedOperationException { + super.initializeInternal(); + initializeLateMaterialization(); + } + + private void initializeLateMaterialization() throws IOException { + if (fileReader == null) { + // A row-group reader was handed in with no file reader behind it, so the phases cannot run. + // The file is read the plain way. + return; + } + blockAccounting = new ParquetBlockAccounting(fileReader); + pageIndexDisabled = !readOptions.useColumnIndexFilter(); + // Resolve each key column's top-level ParquetColumn and keep the ones this file has. A key + // column can be missing under schema evolution. + int[] keyIndices = storageFilter.keyColumnIndices(); + List<Integer> presentKeyPositions = new ArrayList<>(keyIndices.length); + Set<String> keyTopLevelNames = new HashSet<>(); + for (int i = 0; i < keyIndices.length; i++) { + int idx = keyIndices[i]; + if (idx < 0 || idx >= parquetColumn.children().size()) { + // Unreachable: ParquetStorageFilter.create rejects out-of-range ordinals. + throw ParquetStorageFilter.internalError(String.format( + "Storage-filter key ordinal %d is out of range for a %d-column requested schema", + idx, parquetColumn.children().size())); + } + ParquetColumn column = parquetColumn.children().apply(idx); + if (!column.isPrimitive()) { + // Unreachable: ParquetStorageFilter.isSupportedKeyType admits only types with a primitive + // Parquet leaf, and it gates both planning and ParquetStorageFilter.create. + throw ParquetStorageFilter.internalError( + "Storage-filter key column is not a primitive Parquet column: " + column.path()); + } + if (!missingColumns.contains(column)) { + presentKeyPositions.add(i); + keyTopLevelNames.add(column.descriptor().get().getPath()[0]); + } + } + requestedColumns = requestedSchema.getColumns(); + requestedPaths = ParquetBlockAccounting.pathsOf(requestedColumns); + if (presentKeyPositions.isEmpty()) { + // Nothing for phase 1 to read, so the predicate is constant for this file. `initBatch` + // decides it, once the vectors it evaluates against exist. + allKeysMissing = true; + return; + } + + // The non-key columns phase 2 reads. If this file has none of them, the filter saves nothing, + // since phase 1 already reads every column the file has for every row. The reader then + // declines it. The planner already declines a projection of key columns alone, so this catches + // a file missing every projected non-key leaf, which also covers the row-index column. A struct + // the file has, but none of whose requested fields it has, counts only under the legacy + // `returnNullStructIfAllFieldsMissing`, since otherwise the clipped schema reads one of its + // other fields to tell a null struct from one whose requested fields are all null. + List<ColumnDescriptor> nonKey = requestedColumns.stream() + .filter(column -> !keyTopLevelNames.contains(column.getPath()[0])) + .toList(); + if (nonKey.stream().noneMatch(column -> fileSchema.containsPath(column.getPath()))) return; + nonKeyColumns = nonKey; + nonKeyPaths = ParquetBlockAccounting.pathsOf(nonKeyColumns); + + // Resolved here, past the two returns above, because `ValueCopier.forType` rejects a type the + // reader cannot copy, and a file the reader has already declined must not be failed over one. + KeyColumn[] keys = new KeyColumn[presentKeyPositions.size()]; + StructField[] fields = sparkRequestedSchema.fields(); + int fixedBytesPerRow = 0; + for (int i = 0; i < keys.length; i++) { + int rowPosition = presentKeyPositions.get(i); + int batchSlot = keyIndices[rowPosition]; + ParquetColumn column = parquetColumn.children().apply(batchSlot); + DataType type = fields[batchSlot].dataType(); + boolean variableLength = ValueCopier.isVariableLength(type); + keys[i] = new KeyColumn(batchSlot, rowPosition, column.descriptor().get(), + column.required(), type, ValueCopier.forType(type), variableLength); + // One null byte per row either way. A fixed-width value adds its own width. A variable-length + // one adds the int offset and int length that point at the byte child, plus the + // `DEFAULT_CHILD_BYTES_PER_ROW` the child is allocated per row, and a null byte for each. + // A decimal of up to 9 digits is held as an int, though its default size is a long's. + int valueWidth = type instanceof DecimalType && DecimalType.is32BitDecimalType(type) + ? 4 : type.defaultSize(); + fixedBytesPerRow += + variableLength ? 1 + 4 + 4 + 2 * DEFAULT_CHILD_BYTES_PER_ROW : 1 + valueWidth; + } + keyFixedBytesPerRow = fixedBytesPerRow; + keyOnlyColumns = Arrays.stream(keys).map(KeyColumn::descriptor).toList(); + // Last, since a non-null `keyColumns` is what says the filter is applied to this file. + keyColumns = keys; + } + + /** + * Loads the next row group with rows left to emit, through the three phases of the class javadoc. + * Phase 2 reads the non-key columns over the survivors while splicing, the whole projection over + * the survivors once splicing is given up, and the whole projection over + * {@code pushedFilterRanges} once the filter is. Returns false when the row groups run out, which + * ends a read the filter narrowed. + */ + @Override + protected boolean loadNextRowGroup() throws IOException { + if (keyColumns == null) return super.loadNextRowGroup(); + // The previous row group is fully emitted by the time this is called, so its pages are done + // with. Released here rather than at the next assignment, so a row group the filter empties, + // or the end of the file, does not keep the previous one's pages alive. Its column readers go + // too. They hold the row ranges it was read over, which the budget no longer counts, and the + // decoders of their last page. + closeDataPages(); + releaseRowGroupReaders(); + StorageFilterMetrics m = storageFilter.metrics(); + List<BlockMetaData> blocks = fileReader.getRowGroups(); + while (nextBlockIndex < blocks.size()) { + // A plain read goes back to `FileScanRDD`, which checks for a task kill, after every batch. + // This loop can pass any number of row groups the filter empties without emitting one. + ParquetStorageFilter.throwIfKilled(); + int blockIdx = nextBlockIndex++; + long blockRowCount = blocks.get(blockIdx).getRowCount(); + + // Phase 0: rows the pushed data filter allows, by column index. The full projection goes back + // on first, because `getRowRanges` builds an unseen block's store over the current paths. + fileReader.setRequestedSchema(requestedColumns); + RowRanges pushedFilterRanges = pushedFilterRangesFor(blockIdx, blockRowCount); + // RowRanges.rowCount() walks every range, so resolve each range set's count once. + long baselineRows = pushedFilterRanges.rowCount(); + if (baselineRows == 0) { + // An empty block, or one the pushed data filter rejects whole. Not this filter's skip, so + // no storage-filter metric counts it. + continue; + } + + // Whether phase 2 can read part of this row group. That needs an offset index for every + // column of the block's column index store, and parquet empties the store if one column + // lacks it. With a data filter pushed, the store is built at initialize over the whole + // projection, so every projected column is checked. Files written before parquet-mr 1.11, + // or by pyarrow's `write_table` by default, have no offset index at all. With no data filter + // pushed the check is stricter than needed, since the store would cover only the non-key + // columns. False means nothing is buffered and the filter can only skip the row group whole. + // Phase 2 then reads the whole projection over `pushedFilterRanges`. Parquet narrows those + // only by an index it could read, so that read cannot throw. + boolean canNarrowRowGroup = + !pageIndexDisabled && blockAccounting.hasOffsetIndexes(blockIdx, requestedPaths); + spliceCurrentRowGroup = canNarrowRowGroup; + + // Phase 1 reads the key columns over `pushedFilterRanges` and evaluates the filter per row. + fileReader.setRequestedSchema(keyOnlyColumns); + RowRanges survivors; + // Closed at the end of the phase that reads it, since parquet closes no store it hands out. + try (PageReadStore keyPages = + fileReader.readFilteredRowGroup(blockIdx, pushedFilterRanges)) { + requireRowCount(keyPages, baselineRows, blockIdx); + survivors = evaluateStorageFilter( + keyPages, pushedFilterRanges, baselineRows, blockRowCount, canNarrowRowGroup); + } + // Null means phase 1 gave the filter up for this row group, and phase 2 then reads every row + // the pushed filter kept, as a plain read would. Phase 2 can also give it up, in the catch + // below. + boolean filterGivenUp = survivors == null; + if (filterGivenUp) abandonSplicing(); + RowRanges finalRanges = filterGivenUp ? pushedFilterRanges : survivors; + long finalRowCount = filterGivenUp ? baselineRows : survivors.rowCount(); + // The byte metrics' baseline, the non-key bytes a plain read asks for over the rows the + // pushed filter kept. Not walked for a row group whose filter was given up. + long nonKeyBaselineBytes = filterGivenUp ? 0L : blockAccounting.compressedBytesForRowRanges( + blockIdx, nonKeyPaths, pushedFilterRanges, baselineRows); + if (finalRowCount == 0) { + // The filter rejected every row, so the row group is skipped whole. + m.recordRowGroupSkipped(baselineRows, nonKeyBaselineBytes); + continue; + } + + fileReader.setRequestedSchema(spliceCurrentRowGroup ? nonKeyColumns : requestedColumns); Review Comment: Thanks. I'd keep this one for a follow-up, as you suggest for further optimizations. One batch of the longest keys bounds it, so it needs a threshold to be worth the reallocation. It is in the description's follow-up list now. -- 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]
