peter-toth commented on code in PR #58895: URL: https://github.com/apache/spark/pull/58895#discussion_r4136317873
########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilter.scala: ########## @@ -0,0 +1,241 @@ +/* + * 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 org.apache.spark.SparkContext +import org.apache.spark.sql.catalyst.expressions.{And, BasePredicate, BloomFilterMightContain, BoundReference, Expression, Predicate} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.types.{DataType, StructType} + +/** + * The SQL metrics the reader updates while applying a [[ParquetStorageFilter]], created by + * `ParquetFileFormat.storageFilterMetrics`. Every counter is scoped to what the storage filter + * added on top of a read of the same projection without one. + * + * - [[rowGroupsSkipped]] counts row groups whose data columns were never read. + * - [[rowsExcludedByRowGroup]] sums those skipped row groups' rows, counting per block the rows + * that survived the pushed data filter. + * - [[rowsExcludedWithinRowGroup]] sums rows excluded inside row groups that were kept. + * - [[bytesAvoidedByRowGroup]] sums, per skipped row group, the non-key bytes a plain read of this + * projection would have transferred for the rows that survived the pushed data filter. Phase 1 + * reads the key columns of every block, so key bytes are never part of it, and it is zero on an + * all-keys projection, which can avoid nothing. + * - [[bytesAvoidedByPageFiltering]] sums, per kept row group, that same non-key baseline minus the + * bytes phase 2 read, which is what `finalRanges` page selection pruned. + * + * The row counters' suffix says where a row was excluded, not whether reading it was avoided. A row + * excluded inside a kept row group may have been read as part of a page that held a survivor, or + * not read at all because phase 2 skipped its page. An all-keys projection has no page filtering at + * all, and [[rowsExcludedWithinRowGroup]] still counts every row the filter dropped. + */ +case class StorageFilterMetrics( + rowGroupsSkipped: SQLMetric, + rowsExcludedByRowGroup: SQLMetric, + rowsExcludedWithinRowGroup: SQLMetric, + bytesAvoidedByRowGroup: SQLMetric, + bytesAvoidedByPageFiltering: SQLMetric) { + + /** Counts a row group whose data columns the filter kept the reader from touching at all. */ + def recordRowGroupSkipped(excludedRows: Long, avoidedBytes: Long): Unit = { + rowGroupsSkipped.add(1L) + rowsExcludedByRowGroup.add(excludedRows) + bytesAvoidedByRowGroup.add(avoidedBytes) + } + + /** These counters keyed for the scan, which carries them without naming any of them. */ + def toMap: Map[String, SQLMetric] = Map( + StorageFilterMetrics.ROW_GROUPS_SKIPPED -> rowGroupsSkipped, + StorageFilterMetrics.ROWS_EXCLUDED_BY_ROW_GROUP -> rowsExcludedByRowGroup, + StorageFilterMetrics.ROWS_EXCLUDED_WITHIN_ROW_GROUP -> rowsExcludedWithinRowGroup, + StorageFilterMetrics.BYTES_AVOIDED_BY_ROW_GROUP -> bytesAvoidedByRowGroup, + StorageFilterMetrics.BYTES_AVOIDED_BY_PAGE_FILTERING -> bytesAvoidedByPageFiltering) +} + +object StorageFilterMetrics { + + // The keys these counters are exposed under, next to the fields they name. What they count is + // this reader's vocabulary, which is why the format creates them rather than the scan. + val ROW_GROUPS_SKIPPED = "storageFilterRowGroupsSkipped" + val ROWS_EXCLUDED_BY_ROW_GROUP = "storageFilterRowsExcludedByRowGroup" + val ROWS_EXCLUDED_WITHIN_ROW_GROUP = "storageFilterRowsExcludedWithinRowGroup" + val BYTES_AVOIDED_BY_ROW_GROUP = "storageFilterBytesAvoidedByRowGroup" + val BYTES_AVOIDED_BY_PAGE_FILTERING = "storageFilterBytesAvoidedByPageFiltering" + + /** A fresh set of counters, with the labels the SQL UI shows. */ + def create(sparkContext: SparkContext): StorageFilterMetrics = StorageFilterMetrics( + rowGroupsSkipped = + SQLMetrics.createMetric(sparkContext, "row groups skipped by storage filter"), + rowsExcludedByRowGroup = + SQLMetrics.createMetric(sparkContext, "rows excluded by storage filter (whole row group)"), + rowsExcludedWithinRowGroup = + SQLMetrics.createMetric(sparkContext, "rows excluded by storage filter (within row group)"), + bytesAvoidedByRowGroup = SQLMetrics.createSizeMetric( + sparkContext, "bytes avoided by storage filter (whole row group)"), + bytesAvoidedByPageFiltering = SQLMetrics.createSizeMetric( + sparkContext, "bytes avoided by storage filter (page filtering)")) + + /** + * The counters the scan carries back, which are the ones [[create]] made: the map is this + * format's own transport, and a scan with storage filters on a Parquet relation has no other. + */ + def fromMap(metrics: Map[String, SQLMetric]): StorageFilterMetrics = StorageFilterMetrics( + rowGroupsSkipped = metrics(ROW_GROUPS_SKIPPED), + rowsExcludedByRowGroup = metrics(ROWS_EXCLUDED_BY_ROW_GROUP), + rowsExcludedWithinRowGroup = metrics(ROWS_EXCLUDED_WITHIN_ROW_GROUP), + bytesAvoidedByRowGroup = metrics(BYTES_AVOIDED_BY_ROW_GROUP), + bytesAvoidedByPageFiltering = metrics(BYTES_AVOIDED_BY_PAGE_FILTERING)) +} + +/** + * A runtime filter that the vectorized Parquet reader uses to drive late materialization: read + * key-column pages first, evaluate this filter per row to decide which rows survive, and skip + * data-column pages that do not overlap any surviving row range. + * + * [[keyColumnIndices]] are indices into the scan's requested data schema identifying the leaf + * columns referenced by the filter. [[preparedPredicate]] expects a row whose fields are those key + * columns in that order. The predicate's references are rewritten to [[BoundReference]]s pointing + * at positions 0..(keyColumnIndices.length - 1) when the filter is built. + */ +class ParquetStorageFilter private ( + val keyColumnIndices: Array[Int], + private var boundExpression: Expression, + val metrics: StorageFilterMetrics, + val maxSplicedRowGroupBytes: Long) extends Serializable { + + // Codegen-produced predicates can be awkward to serialize from driver to executor, so we defer + // construction to first use on the executor. + @transient private lazy val predicate: BasePredicate = { + // The field below is nulled the first time this runs, and the copy that reaches every task is + // the driver's. So a driver-side evaluation would leave each executor building a predicate from + // nothing, which is named here rather than left to an NPE from `Predicate.create`. + require(boundExpression != null, + "a storage filter must only be evaluated on an executor, and this copy has released the " + + "expression it would build its predicate from") + val created = Predicate.create(boundExpression) + // A prepared bloom reaches the executor as a binary Literal inside this expression, up to + // `spark.sql.optimizer.runtime.bloomFilter.maxNumBits` of it, while `created` holds the + // deserialized filter instead. Dropping the expression lets the task reclaim those bytes for Review Comment: Qualified. The comment says the reclaim holds for a columnar scan's task, and that a row-based one keeps the bytes through the scan node its closure captures. I added it to the broadcast item in the description as well, since a broadcast handle in the plan instead of the bytes would fix both. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala: ########## @@ -1947,6 +1947,48 @@ object SQLConf { .checkValue(threshold => threshold >= 0, "The threshold must not be negative.") .createWithDefault(10) + val PARQUET_STORAGE_FILTER_PUSHDOWN_ENABLED = + buildConf("spark.sql.parquet.storageFilterPushdown.enabled") + .doc("If true, the vectorized Parquet reader may apply a runtime storage filter, such as a " + + "bloom filter from join runtime filtering, while it reads. It reads the columns the " + + "filter needs first, evaluates the filter per row, and then reads the remaining columns " + + "only for the rows that survived. This is a planning-time decision. With it false, no " + + "storage filter is attached to a scan at all. A filter that is attached also stays in " + + "the post-scan filter, the way a pushed data filter does, so the reader is free to stop " + + "applying it wherever doing so would cost more than it saves, and the answer does not " + + "change. Narrowing a row group to part of its rows trusts the Parquet page index, which " + + "a read with no pushed data filter never consults, so on a file whose page index is " + + "wrong this can pair a row's key with another row's values. Setting " + + "parquet.filter.columnindex.enabled to false is the escape hatch, and it leaves the " + + "filter with the row groups it empties, which need no index at all. Note that the " + + "surviving key values of a whole row group are buffered before that row group's first " + + "batch is produced, so a task holds up to one extra copy of the key columns for one row " + + "group.") + .version("5.0.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(false) + + val PARQUET_STORAGE_FILTER_PUSHDOWN_MAX_SPLICED_ROW_GROUP_BYTES = + buildConf("spark.sql.parquet.storageFilterPushdown.maxSplicedRowGroupBytes") + .internal() + .doc("The maximum memory, in bytes, that the vectorized Parquet reader holds for one row " + + "group while it applies a storage filter. Two things count against it, and both grow " + + "with the number of surviving rows. One is the key values buffered to splice into the " + + "output batches, which follow the reader's memory mode and so can be off heap. The other " + + "is the row ranges those rows fall into, always on heap, which the second phase needs to " + + "select its pages. The count is examined after every surviving row. Past the limit the " + + "reader releases the buffer and reads every projected column of the surviving rows " + + "instead, which costs one extra read of the key columns. Past it again it reads the row " + + "group with no filter applied at all, which is correct but as slow as not pushing the " + Review Comment: Fixed in all three places. The budget doc says the give-up pays that same extra key read and so is slower than not pushing the filter. The `enabled` doc says a row group where the reader stops reads its key columns twice. It also says the columnindex escape hatch keeps the row groups the filter empties while every other row group reads its key columns twice, and it names turning the conf off as the other way out. The reader comment says a given-up row group read what a plain scan reads plus the key columns phase 1 had already read. The description's list of give-up paths says the same. ########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java: ########## @@ -0,0 +1,1200 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.datasources.parquet; + +import java.io.IOException; +import java.time.ZoneId; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashSet; +import java.util.List; +import java.util.Objects; +import java.util.PrimitiveIterator; +import java.util.Set; + +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.column.page.PageReadStore; +import org.apache.parquet.filter2.columnindex.RowRanges; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.ParquetInputFormat; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore.MissingOffsetIndexException; +import org.apache.parquet.schema.Type; +import org.apache.parquet.schema.Types; + +import org.apache.spark.internal.LogKeys; +import org.apache.spark.internal.MDC; +import org.apache.spark.internal.SparkLogger; +import org.apache.spark.internal.SparkLoggerFactory; +import org.apache.spark.memory.MemoryMode; +import org.apache.spark.sql.catalyst.expressions.BasePredicate; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.execution.vectorized.OffHeapColumnVector; +import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector; +import org.apache.spark.sql.execution.vectorized.WritableColumnVector; +import org.apache.spark.sql.types.DataType; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.vectorized.ColumnVector; +import org.apache.spark.sql.vectorized.ColumnarBatch; + +/** + * A {@link VectorizedParquetRecordReader} that reads a row group in phases under a + * {@code ParquetStorageFilter}, so the projected columns of the rows the filter rejects are never + * read. The filter is a deterministic conjunct of the post-scan {@code Filter}, which still holds + * it, so any phase below may stop short and read the file the way a plain scan would. + * + * <p>Per row group, all three phases driven by the one {@link #fileReader}: + * <ul> + * <li>phase 0 asks the column index which rows the pushed data filter allows, metadata only;</li> + * <li>phase 1 reads the key columns the filter references over those rows, evaluates the filter + * per row, and buffers the surviving key values;</li> + * <li>phase 2 reads the remaining projected columns over the surviving row ranges.</li> + * </ul> + * + * <p>The emit path then splices the buffered key vectors into the batch, so phase 2 never reads a + * key column again. A row group whose buffer would pass + * {@code spark.sql.parquet.storageFilterPushdown.maxSplicedRowGroupBytes} gives splicing up, and + * past that the filter itself, and is read the plain way. + * + * <p>Splicing does not engage at all in one real case, when every key column is missing from this + * physical file under schema evolution. The predicate is then constant over the file and is + * evaluated once, against the constants the reader materializes for those columns. True keeps the + * file unfiltered, false skips it. + * + * <p>Everything else the filter needs is guaranteed by {@code FileSourceStrategy.storageFiltersFor} + * and {@code ParquetStorageFilter.create}, so a violation of it here is a planner bug and is + * asserted rather than handled. A file the reader simply cannot prune is a different matter, and it + * reads that one the way a plain scan would. + */ +public class LateMaterializationParquetRecordReader extends VectorizedParquetRecordReader { + + private static final SparkLogger LOG = + SparkLoggerFactory.getLogger(LateMaterializationParquetRecordReader.class); + + /** + * What one buffered variable-length value costs the vector that holds it, as a multiple of the + * value's own bytes. Those bytes go in the byte child, with one null byte per element of that + * child, and up to as much again because {@link WritableColumnVector#reserve} doubles the child + * when it grows. A factor rather than an exact figure, so this half of the charge over-counts + * rather than under-counts, since nothing above this buffer can spill it. A fixed-width key needs + * no factor, since its accumulator is allocated at capacity and never grows. + */ + private static final long VARIABLE_LENGTH_BYTES_FACTOR = 4L; + + /** + * The storage filter this reader applies. It reads key-column pages first, evaluates the filter + * per row to build {@link RowRanges}, then reads data-column pages restricted to those ranges. + * Nulled once the reader gives it up for the whole file, after which every path here behaves as + * the plain reader does. + */ + private ParquetStorageFilter storageFilter; + + /** + * Set to true once this reader has no rows left to return, which is when the row groups run out + * or when the filter rejects the whole file. + */ + private boolean hitEndOfData = false; + + /** + * Late-materialization state, populated by {@link #initializeLateMaterialization()}. + * + * <p>The base class's {@link #fileReader} drives all three phases. Its requested + * schema is mutated per phase via {@link ParquetFileReader#setRequestedSchema}: all projected + * columns for phase 0 ({@code getRowRanges}), the key columns for phase 1, and for phase 2 the + * non-key columns while splicing or the whole projection once a row group gives it up. + * + * <p>The three sets are held as leaf-column lists rather than as {@code MessageType}s because + * that is what both the reader and the byte metrics consume, and because + * {@code MessageType.getColumns()} rebuilds the list on every call. + * + * <p>Each list's leaf {@link ColumnPath}s are resolved alongside it, because that is all the + * footer walks below need and {@link ColumnPath#get} is not free, since it allocates and interns + * through parquet's canonicalizer, which those walks would otherwise repeat per column on every + * one of their several calls per row group. As sets, because a walk goes over the block's column + * chunks and asks which of them it is about, and a requested schema's leaf paths are unique + * anyway. + */ + private List<ColumnDescriptor> requestedColumns; + private List<ColumnDescriptor> keyOnlyColumns; + private List<ColumnDescriptor> nonKeyColumns; + private Set<ColumnPath> requestedPaths; + private Set<ColumnPath> keyOnlyPaths; + private Set<ColumnPath> nonKeyPaths; + /** + * Whether a row group can be read in part, and what reading part of one transfers. Both are read + * off the footer {@link #fileReader} already holds, so the byte metrics cost no IO. + */ + private ParquetBlockAccounting blockAccounting; + /** + * What an accumulator holds per row for each key column, not counting the value bytes of a + * variable-length type. Together with those value bytes it is what the per-row-group buffer is + * measured against. See + * {@code spark.sql.parquet.storageFilterPushdown.maxSplicedRowGroupBytes}. + */ + private int keyFixedBytesPerRow; + /** + * The pages phase 2 read for the row group being emitted. Held because the column readers draw + * from it for the whole row group, and closed when the next one is loaded, since parquet closes + * neither the stores `readFilteredRowGroup` hands back nor the ones from `readNextRowGroup`. + */ + private PageReadStore dataPages; + + /** Whether the filter has already been reported as failing to evaluate on this file. */ + private boolean loggedFilterEvaluationError; + + /** + * Whether {@code parquet.filter.columnindex.enabled} says this file's page index is not to be + * used. That covers phase 0 as well as the narrowing below, because the conf is the escape hatch + * for a file whose index is wrong, and phase 2 reads part of a row group through the offset index + * whatever the conf says. A wrong index there pairs a row's key with another row's values, which + * the post-scan Filter cannot catch, since the key it sees is the right one. + */ + private boolean pageIndexDisabled; + + /** + * Whether a read of part of a row group has already failed for want of an offset index, which + * says this file is not to be read in part again. Only the narrowing goes. The rows the pushed + * data filter's own page pruning excludes are still excluded, because parquet built those ranges + * itself, and a block it has no index for it hands back whole. + */ + private boolean partialReadUnavailable; + + /** + * Whether the row group currently loaded is spliced. It starts true unless phase 2 will have to + * read every projected column anyway, and turns false in phase 1 once the survivors buffered pass + * the cap. False means phase 2 read every projected column, key columns included, so the emit + * path takes them straight from the persistent batch. + */ + private boolean spliceCurrentRowGroup; + private int nextBlockIndex; + /** This file's row groups, as the footer lists them after parquet's own row-group filtering. */ + private List<BlockMetaData> blocks; + /** + * The key columns this file has, in the filter's key order, with everything init resolves about + * each. Phase 1 reads them, and the survivor queues and accumulators are indexed the same way. + */ + private KeyColumn[] keyColumns; + /** + * Whether every key column is missing from this file, which makes the predicate constant for it. + * Decided in {@link #initBatch}, since the constants only exist once its vectors are built. + */ + private boolean allKeysMissing; + private WritableColumnVector[] keyScratchVectors; + private ColumnarBatch keyScratchBatch; + + /** + * Splicing state. Phase 1 keeps the surviving key values it has already decoded, and the emit + * path splices them back into the output batch, so phase 2 never reads the key columns. + * + * <p>The alternative is for phase 2 to read the key columns again under the surviving row ranges, + * which needs none of this state but pays a second read of them for every row group. Nothing + * absorbs that read on object storage, where it is a new GET rather than a page-cache hit. + * + * <p>{@link #isKeyTopLevel} marks the top-level slots that emit takes from the queues rather + * than from a phase-2 read. {@link #keyVectorQueues} holds one queue per present key column of + * capacity-sized survivor vectors, and {@link #currentKeyAccumulators} the vectors still filling. + * A queue keeps owning its head while the batch is built on it, until the next emit closes it, so + * survivor memory drains as the row group is emitted. {@link #persistentBatchColumns} is + * {@link #initBatch}'s vector array and {@link #spliceBatchColumns} the array the emitted batch + * is built over, which is why the two are closed separately. + * + * <p>Phase 1 evaluates a whole row group before its first batch is emitted, so the queues hold + * every surviving key value of one row group at once. That is bounded per row group by + * {@code spark.sql.parquet.storageFilterPushdown.maxSplicedRowGroupBytes}, past which the row + * group is read the plain way and nothing is buffered. + */ + private boolean[] isKeyTopLevel; + private ArrayDeque<WritableColumnVector>[] keyVectorQueues; + private WritableColumnVector[] currentKeyAccumulators; + /** Row count of {@link #currentKeyAccumulators}; all key columns advance in lockstep. */ + private int currentKeyAccumulatorRowCount; + /** Whether each queue's head is the vector the current batch is built on. */ + private boolean keyVectorsPublished; + private ColumnVector[] persistentBatchColumns; + private ColumnVector[] spliceBatchColumns; + + /** + * The filter is a constructor argument rather than a setter, because + * {@link #initializeLateMaterialization()} inspects the per-file schema against it and runs from + * {@code initialize}: there is no point in the reader's life at which it is useful without one. + */ + public LateMaterializationParquetRecordReader( + ZoneId convertTz, + String datetimeRebaseMode, + String datetimeRebaseTz, + String int96RebaseMode, + String int96RebaseTz, + boolean useOffHeap, + int capacity, + ParquetStorageFilter storageFilter) { + super(convertTz, datetimeRebaseMode, datetimeRebaseTz, int96RebaseMode, int96RebaseTz, + useOffHeap, capacity); + this.storageFilter = Objects.requireNonNull(storageFilter, "storageFilter"); + } + + // For test only. + public LateMaterializationParquetRecordReader( + boolean useOffHeap, int capacity, ParquetStorageFilter storageFilter) { + super(useOffHeap, capacity); + this.storageFilter = Objects.requireNonNull(storageFilter, "storageFilter"); + } + + @Override + protected void closeAllocated() { + // Each release below is independent, and the base class's close() owns the file handle and + // input stream, so they are chained through finally blocks, and one failing vector close must + // not leak the rest. + try { + // The batch's vectors are closed through `persistentBatchColumns` rather than through + // `columnarBatch.close()`, which lets both paths share one body. While splicing the emitted + // batch is a view whose slots alias these vectors (non-key and partition) and the head of + // each survivor queue (key slots), so closing it would re-close a shared vector and free the Review Comment: Rewritten, in the places you name and one more. `closeAllocated`'s comment says the persistent key vectors are what closing the spliced batch would leave unreached. The base class's `closeAllocated` javadoc made the same double-free claim, and it now just says a subclass closes what its batch does not hold. The suite comment says a read after close is the off-heap hazard and that a second close is harmless. ########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/LateMaterializationParquetRecordReader.java: ########## @@ -0,0 +1,1200 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.datasources.parquet; + +import java.io.IOException; +import java.time.ZoneId; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashSet; +import java.util.List; +import java.util.Objects; +import java.util.PrimitiveIterator; +import java.util.Set; + +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.column.page.PageReadStore; +import org.apache.parquet.filter2.columnindex.RowRanges; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.ParquetInputFormat; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore.MissingOffsetIndexException; +import org.apache.parquet.schema.Type; +import org.apache.parquet.schema.Types; + +import org.apache.spark.internal.LogKeys; +import org.apache.spark.internal.MDC; +import org.apache.spark.internal.SparkLogger; +import org.apache.spark.internal.SparkLoggerFactory; +import org.apache.spark.memory.MemoryMode; +import org.apache.spark.sql.catalyst.expressions.BasePredicate; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.execution.vectorized.OffHeapColumnVector; +import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector; +import org.apache.spark.sql.execution.vectorized.WritableColumnVector; +import org.apache.spark.sql.types.DataType; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.vectorized.ColumnVector; +import org.apache.spark.sql.vectorized.ColumnarBatch; + +/** + * A {@link VectorizedParquetRecordReader} that reads a row group in phases under a + * {@code ParquetStorageFilter}, so the projected columns of the rows the filter rejects are never + * read. The filter is a deterministic conjunct of the post-scan {@code Filter}, which still holds + * it, so any phase below may stop short and read the file the way a plain scan would. + * + * <p>Per row group, all three phases driven by the one {@link #fileReader}: + * <ul> + * <li>phase 0 asks the column index which rows the pushed data filter allows, metadata only;</li> + * <li>phase 1 reads the key columns the filter references over those rows, evaluates the filter + * per row, and buffers the surviving key values;</li> + * <li>phase 2 reads the remaining projected columns over the surviving row ranges.</li> + * </ul> + * + * <p>The emit path then splices the buffered key vectors into the batch, so phase 2 never reads a + * key column again. A row group whose buffer would pass + * {@code spark.sql.parquet.storageFilterPushdown.maxSplicedRowGroupBytes} gives splicing up, and + * past that the filter itself, and is read the plain way. + * + * <p>Splicing does not engage at all in one real case, when every key column is missing from this + * physical file under schema evolution. The predicate is then constant over the file and is + * evaluated once, against the constants the reader materializes for those columns. True keeps the + * file unfiltered, false skips it. + * + * <p>Everything else the filter needs is guaranteed by {@code FileSourceStrategy.storageFiltersFor} + * and {@code ParquetStorageFilter.create}, so a violation of it here is a planner bug and is + * asserted rather than handled. A file the reader simply cannot prune is a different matter, and it + * reads that one the way a plain scan would. + */ +public class LateMaterializationParquetRecordReader extends VectorizedParquetRecordReader { + + private static final SparkLogger LOG = + SparkLoggerFactory.getLogger(LateMaterializationParquetRecordReader.class); + + /** + * What one buffered variable-length value costs the vector that holds it, as a multiple of the + * value's own bytes. Those bytes go in the byte child, with one null byte per element of that + * child, and up to as much again because {@link WritableColumnVector#reserve} doubles the child + * when it grows. A factor rather than an exact figure, so this half of the charge over-counts + * rather than under-counts, since nothing above this buffer can spill it. A fixed-width key needs + * no factor, since its accumulator is allocated at capacity and never grows. + */ + private static final long VARIABLE_LENGTH_BYTES_FACTOR = 4L; + + /** + * The storage filter this reader applies. It reads key-column pages first, evaluates the filter + * per row to build {@link RowRanges}, then reads data-column pages restricted to those ranges. + * Nulled once the reader gives it up for the whole file, after which every path here behaves as + * the plain reader does. + */ + private ParquetStorageFilter storageFilter; + + /** + * Set to true once this reader has no rows left to return, which is when the row groups run out + * or when the filter rejects the whole file. + */ + private boolean hitEndOfData = false; + + /** + * Late-materialization state, populated by {@link #initializeLateMaterialization()}. + * + * <p>The base class's {@link #fileReader} drives all three phases. Its requested + * schema is mutated per phase via {@link ParquetFileReader#setRequestedSchema}: all projected + * columns for phase 0 ({@code getRowRanges}), the key columns for phase 1, and for phase 2 the + * non-key columns while splicing or the whole projection once a row group gives it up. + * + * <p>The three sets are held as leaf-column lists rather than as {@code MessageType}s because + * that is what both the reader and the byte metrics consume, and because + * {@code MessageType.getColumns()} rebuilds the list on every call. + * + * <p>Each list's leaf {@link ColumnPath}s are resolved alongside it, because that is all the + * footer walks below need and {@link ColumnPath#get} is not free, since it allocates and interns + * through parquet's canonicalizer, which those walks would otherwise repeat per column on every + * one of their several calls per row group. As sets, because a walk goes over the block's column + * chunks and asks which of them it is about, and a requested schema's leaf paths are unique + * anyway. + */ + private List<ColumnDescriptor> requestedColumns; + private List<ColumnDescriptor> keyOnlyColumns; + private List<ColumnDescriptor> nonKeyColumns; + private Set<ColumnPath> requestedPaths; + private Set<ColumnPath> keyOnlyPaths; + private Set<ColumnPath> nonKeyPaths; + /** + * Whether a row group can be read in part, and what reading part of one transfers. Both are read + * off the footer {@link #fileReader} already holds, so the byte metrics cost no IO. + */ + private ParquetBlockAccounting blockAccounting; + /** + * What an accumulator holds per row for each key column, not counting the value bytes of a + * variable-length type. Together with those value bytes it is what the per-row-group buffer is + * measured against. See + * {@code spark.sql.parquet.storageFilterPushdown.maxSplicedRowGroupBytes}. + */ + private int keyFixedBytesPerRow; + /** + * The pages phase 2 read for the row group being emitted. Held because the column readers draw + * from it for the whole row group, and closed when the next one is loaded, since parquet closes + * neither the stores `readFilteredRowGroup` hands back nor the ones from `readNextRowGroup`. + */ + private PageReadStore dataPages; + + /** Whether the filter has already been reported as failing to evaluate on this file. */ + private boolean loggedFilterEvaluationError; + + /** + * Whether {@code parquet.filter.columnindex.enabled} says this file's page index is not to be + * used. That covers phase 0 as well as the narrowing below, because the conf is the escape hatch + * for a file whose index is wrong, and phase 2 reads part of a row group through the offset index + * whatever the conf says. A wrong index there pairs a row's key with another row's values, which + * the post-scan Filter cannot catch, since the key it sees is the right one. + */ + private boolean pageIndexDisabled; + + /** + * Whether a read of part of a row group has already failed for want of an offset index, which + * says this file is not to be read in part again. Only the narrowing goes. The rows the pushed + * data filter's own page pruning excludes are still excluded, because parquet built those ranges + * itself, and a block it has no index for it hands back whole. + */ + private boolean partialReadUnavailable; + + /** + * Whether the row group currently loaded is spliced. It starts true unless phase 2 will have to + * read every projected column anyway, and turns false in phase 1 once the survivors buffered pass + * the cap. False means phase 2 read every projected column, key columns included, so the emit + * path takes them straight from the persistent batch. + */ + private boolean spliceCurrentRowGroup; + private int nextBlockIndex; + /** This file's row groups, as the footer lists them after parquet's own row-group filtering. */ + private List<BlockMetaData> blocks; + /** + * The key columns this file has, in the filter's key order, with everything init resolves about + * each. Phase 1 reads them, and the survivor queues and accumulators are indexed the same way. + */ + private KeyColumn[] keyColumns; + /** + * Whether every key column is missing from this file, which makes the predicate constant for it. + * Decided in {@link #initBatch}, since the constants only exist once its vectors are built. + */ + private boolean allKeysMissing; + private WritableColumnVector[] keyScratchVectors; + private ColumnarBatch keyScratchBatch; + + /** + * Splicing state. Phase 1 keeps the surviving key values it has already decoded, and the emit + * path splices them back into the output batch, so phase 2 never reads the key columns. + * + * <p>The alternative is for phase 2 to read the key columns again under the surviving row ranges, + * which needs none of this state but pays a second read of them for every row group. Nothing + * absorbs that read on object storage, where it is a new GET rather than a page-cache hit. + * + * <p>{@link #isKeyTopLevel} marks the top-level slots that emit takes from the queues rather + * than from a phase-2 read. {@link #keyVectorQueues} holds one queue per present key column of + * capacity-sized survivor vectors, and {@link #currentKeyAccumulators} the vectors still filling. + * A queue keeps owning its head while the batch is built on it, until the next emit closes it, so + * survivor memory drains as the row group is emitted. {@link #persistentBatchColumns} is + * {@link #initBatch}'s vector array and {@link #spliceBatchColumns} the array the emitted batch + * is built over, which is why the two are closed separately. + * + * <p>Phase 1 evaluates a whole row group before its first batch is emitted, so the queues hold + * every surviving key value of one row group at once. That is bounded per row group by + * {@code spark.sql.parquet.storageFilterPushdown.maxSplicedRowGroupBytes}, past which the row + * group is read the plain way and nothing is buffered. + */ + private boolean[] isKeyTopLevel; + private ArrayDeque<WritableColumnVector>[] keyVectorQueues; + private WritableColumnVector[] currentKeyAccumulators; + /** Row count of {@link #currentKeyAccumulators}; all key columns advance in lockstep. */ + private int currentKeyAccumulatorRowCount; + /** Whether each queue's head is the vector the current batch is built on. */ + private boolean keyVectorsPublished; + private ColumnVector[] persistentBatchColumns; + private ColumnVector[] spliceBatchColumns; + + /** + * The filter is a constructor argument rather than a setter, because + * {@link #initializeLateMaterialization()} inspects the per-file schema against it and runs from + * {@code initialize}: there is no point in the reader's life at which it is useful without one. + */ + public LateMaterializationParquetRecordReader( + ZoneId convertTz, + String datetimeRebaseMode, + String datetimeRebaseTz, + String int96RebaseMode, + String int96RebaseTz, + boolean useOffHeap, + int capacity, + ParquetStorageFilter storageFilter) { + super(convertTz, datetimeRebaseMode, datetimeRebaseTz, int96RebaseMode, int96RebaseTz, + useOffHeap, capacity); + this.storageFilter = Objects.requireNonNull(storageFilter, "storageFilter"); + } + + // For test only. + public LateMaterializationParquetRecordReader( + boolean useOffHeap, int capacity, ParquetStorageFilter storageFilter) { + super(useOffHeap, capacity); + this.storageFilter = Objects.requireNonNull(storageFilter, "storageFilter"); + } + + @Override + protected void closeAllocated() { + // Each release below is independent, and the base class's close() owns the file handle and + // input stream, so they are chained through finally blocks, and one failing vector close must + // not leak the rest. + try { + // The batch's vectors are closed through `persistentBatchColumns` rather than through + // `columnarBatch.close()`, which lets both paths share one body. While splicing the emitted + // batch is a view whose slots alias these vectors (non-key and partition) and the head of + // each survivor queue (key slots), so closing it would re-close a shared vector and free the + // same buffer twice. Without splicing its array is this one. + try { + if (persistentBatchColumns != null) { + closeAll(persistentBatchColumns); + } else if (columnarBatch != null) { + // `takeOverBatch` has not run, or threw part way, and the batch super built is then the + // only thing holding its vectors. + columnarBatch.close(); + } + persistentBatchColumns = null; + spliceBatchColumns = null; + columnarBatch = null; + } finally { + closeSplicingState(); + } + } finally { + try { + // Through the array, not through `keyScratchBatch`, which is assigned only after the + // allocation loop finishes, so a partial failure leaves vectors only the array can reach. + closeAll(keyScratchVectors); + keyScratchVectors = null; + keyScratchBatch = null; + } finally { + closeDataPages(); + } + } + } + + @Override + public float getProgress() { + // Under a storage filter, rowsReturned counts survivors while totalRowCount is the count before + // the filter narrowed anything, so the ratio neither reaches 1 on its own nor stays below it. A + // row group read whole after the page index is given up can return more rows than that count. + // `hitEndOfData` is the real terminator here, and the ratio is clamped to what it means. + if (hitEndOfData) return 1.0f; + return Math.min(1.0f, super.getProgress()); + } + + /** + * Builds the batch the way the plain reader does, takes it over for splicing, and then answers + * the filter for a file that has none of its key columns. + * + * <p>Both after super's body, because the batch it hands out is what the splicing one is built + * over, and the all-keys-missing decision evaluates the predicate against the constant vectors + * this reader returns for those columns, which exist only once super has wrapped the allocated + * vectors in {@link ParquetColumnVector}s. + */ + @Override + protected void initBatch( + MemoryMode memMode, + StructType partitionColumns, + InternalRow partitionValues) { + super.initBatch(memMode, partitionColumns, partitionValues); + takeOverBatch(); + if (allKeysMissing) { + decideAllKeysMissingFile(); + } + } + + /** + * Takes the vectors of the batch super built, and while splicing hands out a batch of this + * reader's own over a copy of that array, whose key slots the emit path rewrites in place. + * {@link ColumnarBatch} holds the array by reference, its staging row included, so rewriting a + * slot is what publishes it, and the object being reused is what makes a consumer see it. + * + * <p>Every slot got a vector, key columns included. While splicing a key slot's vector is + * unused, since the emitted batch takes that slot from the survivor queues. A row group read the + * plain way past the buffer cap does read into it, and one capacity-sized vector per key column + * is cheap next to the buffer the cap is there to bound. + */ + private void takeOverBatch() { + persistentBatchColumns = new ColumnVector[columnarBatch.numCols()]; + for (int i = 0; i < persistentBatchColumns.length; i++) { + persistentBatchColumns[i] = columnarBatch.column(i); + } + if (keyColumns == null) return; + spliceBatchColumns = persistentBatchColumns.clone(); + columnarBatch = new ColumnarBatch(spliceBatchColumns); + } + + /** + * Answers the storage filter for a file that has none of its key columns, which makes the + * predicate constant over the file. It is evaluated against the vectors just built for those + * columns, so the value it sees is the one the scan returns for them. + * + * <p>False skips the file. True, or an error, reads it the way a plain scan would, under the same + * fail-open rule as a row's value, since the constant can be one the predicate throws on. + */ + private void decideAllKeysMissingFile() { + allKeysMissing = false; + // Not closed, and must not be, since the vectors in it belong to the batch this reader emits. + ColumnarBatch keyRow = new ColumnarBatch(keyVectorsFromBatch(), 1); + boolean keep; + try { + keep = storageFilter.preparedPredicate().eval(keyRow.getRow(0)); + } catch (RuntimeException e) { + logFilterGivenUpOnError(e); + keep = true; + } + if (!keep) { + recordFileSkipped(); + hitEndOfData = true; + } + storageFilter = null; + } + + /** + * Whether this read has returned every row it is going to. + * + * <p>{@code totalRowCount} is how many rows the file yields before this filter narrows anything, + * so it is a terminator only while there is no filter. Under one it is an upper bound the read + * never reaches, and a row group read whole after the page index was given up can pass it. What + * ends such a read is {@link #hitEndOfData}, set when the row groups run out or when the filter + * rejects the whole file. + */ + @Override + protected boolean noMoreRows() { + return hitEndOfData || (storageFilter == null && super.noMoreRows()); + } + + /** + * Releases the key vectors the previous batch was built on, before the batch is emptied, so a + * terminal call cannot leave a spliced batch pointing at vectors just released, which off heap is + * freed memory. + */ + @Override + protected void resetBatch() { + releasePublishedKeyVectors(); + super.resetBatch(); + } + + /** + * Points the emitted batch's key slots at what this batch takes them from, then reads the rest + * the way the plain path does. Three cases reach the base's loop. A spliced row group, whose key + * slots have no phase-2 reader because the emitted batch takes them from the survivor queues. A + * row group read the plain way past the buffer cap, which has a reader for every slot. And a file + * the filter was given up for, which is the plain path exactly. + * + * <p>The row-index column the base fills there is identified by name + * (ROW_INDEX_TEMPORARY_COLUMN_NAME), a synthetic metadata column no storage filter references, so + * its slot is always a non-key one and its ParquetColumnVector is always the persistent one. + */ + @Override + protected void readBatchColumns(int num) throws IOException { + // Without a storage filter the emitted batch is `initBatch`'s array and this does not apply. + if (spliceBatchColumns != null) { + pointKeySlotsAtThisBatch(num); + } + super.readBatchColumns(num); + } + + /** + * Points the emitted batch's key slots at the vectors this batch takes them from. They are the + * only slots that ever change, since a non-key slot holds the vector {@link #initBatch} built for + * it for the reader's life. A spliced row group takes the head of each survivor queue, which is + * what publishes the buffered keys, and any other row group takes the persistent vector, which + * is what a row group following a spliced one needs. + * + * <p>The same {@link ColumnarBatch} is handed out every time, over an array it holds by + * reference. The queues keep owning the vectors they publish until the next batch releases them, + * so a phase-2 read that throws afterwards leaves them reachable for {@link #close()}. + */ + private void pointKeySlotsAtThisBatch(int num) { + keyVectorsPublished = spliceCurrentRowGroup; + // Queue k holds the survivors of key column k, and its `batchSlot` is where that key column + // sits in the batch, so the pairing is read off rather than re-derived from slot order. + for (int k = 0; k < keyColumns.length; k++) { + int slot = keyColumns[k].batchSlot(); + spliceBatchColumns[slot] = + spliceCurrentRowGroup ? nextSurvivorVector(k, num) : persistentBatchColumns[slot]; + } + } + + /** The head of survivor queue {@code k}, which is what publishes the keys it buffered. */ + private WritableColumnVector nextSurvivorVector(int k, int num) { + WritableColumnVector vector = keyVectorQueues[k].peekFirst(); + if (vector == null) { + // Unreachable, since the queues hold exactly the survivors phase 1 accumulated and the emit + // loop is driven by that same count. Named here rather than left to a null batch slot and an + // NPE somewhere in the read that follows. + throw new IllegalStateException(String.format( + "Storage-filter survivor queue %d of row group %d in %s ran out with %d rows still " + + "to emit", k, nextBlockIndex - 1, fileReader.getFile(), num)); + } + return vector; + } + + /** + * Closes the key vectors the previous batch was built on. The queues own them until here, so + * survivor memory drains as a row group is emitted rather than all at its end. + */ + private void releasePublishedKeyVectors() { + if (!keyVectorsPublished) return; + keyVectorsPublished = false; + for (ArrayDeque<WritableColumnVector> queue : keyVectorQueues) { + queue.removeFirst().close(); + } + } + + @Override + protected void initializeInternal() throws IOException, UnsupportedOperationException { + super.initializeInternal(); + initializeLateMaterialization(); + } + + private void initializeLateMaterialization() throws IOException { + if (fileReader == null) { + // The three phases drive a ParquetFileReader directly, and this reader was handed a row-group + // reader with none behind it, so the filter is simply not applied and the plain read path + // runs. The conjunct is in the post-scan filter as well, so the rows it would have dropped + // are dropped above the scan. + storageFilter = null; + return; + } + blockAccounting = new ParquetBlockAccounting(fileReader); + blocks = fileReader.getRowGroups(); + // Both of the flags the narrowing hangs on start here. The conf one closes phase 0 too, and the + // other is what the failed read below latches. Neither costs the filter the row groups it + // empties, which need no index at all. + pageIndexDisabled = + !configuration.getBoolean(ParquetInputFormat.COLUMN_INDEX_FILTERING_ENABLED, true); + // Resolve each key column's top-level ParquetColumn, keeping the ones this file has. A missing + // one can happen under schema evolution, where a column is in the requested schema but not in + // the file. For any non-primitive key we still bail, since phase 1 reads only primitive leaves. + int[] keyIndices = storageFilter.keyColumnIndices(); + List<ParquetColumn> presentKeyColumns = new ArrayList<>(keyIndices.length); + List<Integer> presentKeyPositions = new ArrayList<>(keyIndices.length); + 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 new IllegalStateException(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 new IllegalStateException( + "Storage-filter key column is not a primitive Parquet column: " + column.path()); + } + if (!missingColumns.contains(column)) { + presentKeyPositions.add(i); + presentKeyColumns.add(column); + } + } + + // A key column this file does not have is not substituted into the predicate. Its slot in the + // row phase 1 evaluates points at the vector `initBatch` builds for it, which holds exactly + // what the scan will return for that column, its existence DEFAULT or null where it has none. + // So there is one rule for what a missing column reads as, ParquetColumnVector's, and the + // filter cannot decide on a value the rows never show. + requestedColumns = requestedSchema.getColumns(); + requestedPaths = ParquetBlockAccounting.pathsOf(requestedColumns); + if (presentKeyColumns.isEmpty()) { + // Nothing for phase 1 to read, so the predicate is constant for this file. Deciding needs the + // vectors `initBatch` has not built yet, so it happens there, in `decideAllKeysMissingFile`. + allKeysMissing = true; + return; + } + + int numKeys = presentKeyColumns.size(); + ColumnDescriptor[] keyDescriptors = new ColumnDescriptor[numKeys]; + Types.MessageTypeBuilder keySchemaBuilder = Types.buildMessage(); + Set<String> keyTopLevelNames = new HashSet<>(); + for (int i = 0; i < numKeys; i++) { + keyDescriptors[i] = presentKeyColumns.get(i).descriptor().get(); + // Preserve the field name/type as it appears at the top of requestedSchema. + String topLevelName = keyDescriptors[i].getPath()[0]; + keySchemaBuilder.addField(requestedSchema.getType(topLevelName)); + keyTopLevelNames.add(topLevelName); + } + keyOnlyColumns = keySchemaBuilder.named(requestedSchema.getName()).getColumns(); + keyOnlyPaths = ParquetBlockAccounting.pathsOf(keyOnlyColumns); + + // Build the non-key (complement) schema, which phase 2 reads under finalRanges. With nothing in + // it there is nothing to prune, since such a scan reads the same columns for the same rows as a + // plain one, phase 1 having to read a key column to evaluate the filter on it. So the reader + // declines rather than carrying a second emit path for it. The planner does not push this shape + // either, which makes the two consistent without depending on each other. + Types.MessageTypeBuilder nonKeyBuilder = Types.buildMessage(); + int nonKeyFieldCount = 0; + for (Type field : requestedSchema.getFields()) { + if (!keyTopLevelNames.contains(field.getName())) { + nonKeyBuilder.addField(field); + nonKeyFieldCount++; + } + } + if (nonKeyFieldCount == 0) { + storageFilter = null; + return; + } + nonKeyColumns = nonKeyBuilder.named(requestedSchema.getName()).getColumns(); + 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. + keyColumns = new KeyColumn[numKeys]; + keyFixedBytesPerRow = 0; + StructField[] fields = sparkRequestedSchema.fields(); + for (int i = 0; i < numKeys; i++) { + int rowPosition = presentKeyPositions.get(i); + int batchSlot = keyIndices[rowPosition]; + DataType type = fields[batchSlot].dataType(); + boolean variableLength = ValueCopier.isVariableLength(type); + keyColumns[i] = new KeyColumn(batchSlot, rowPosition, keyDescriptors[i], + presentKeyColumns.get(i).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 what the child + // itself is allocated per row before a value is written, which is four bytes of data, one per + // WritableColumnVector.DEFAULT_ARRAY_LENGTH, and a null byte for each of them. + keyFixedBytesPerRow += variableLength ? 1 + 4 + 4 + 4 + 4 : 1 + type.defaultSize(); + } + initializeSplicingState(); + } + + /** + * Populates the splicing bookkeeping. {@code keyColumnIndices} already index the top-level + * slots of {@link #sparkSchema}, the same indexing as {@link #columnVectors}, so they map + * directly onto {@link #isKeyTopLevel}, which says which batch slots the emit path may take from + * the survivor queues. + */ + @SuppressWarnings("unchecked") + private void initializeSplicingState() { + isKeyTopLevel = new boolean[sparkSchema.fields().length]; + // Only the key columns this file has, since a missing one keeps the constant vector `initBatch` + // built for its slot, there being nothing to buffer and nothing to splice back. + for (KeyColumn key : keyColumns) { + isKeyTopLevel[key.batchSlot()] = true; + } + // Both arrays before either is filled, so `closeSplicingState` finding one of them finds both, + // whatever failed part way through here. + keyVectorQueues = new ArrayDeque[keyColumns.length]; + currentKeyAccumulators = new WritableColumnVector[keyColumns.length]; + for (int i = 0; i < keyColumns.length; i++) { + keyVectorQueues[i] = new ArrayDeque<>(); + } + currentKeyAccumulatorRowCount = 0; + } + + /** + * Loads the next row group using the three-phase late-materialization pattern, all driven by the + * one {@link #fileReader} with its requested schema mutated per phase: + * - Phase 0, over the full schema, computes {@code pushedFilterRanges} from the pushed data + * filter via the column index, metadata only, using + * {@link ParquetFileReader#getRowRanges}. + * - Phase 1, over the key-only schema, reads key-column pages restricted to + * {@code pushedFilterRanges}, evaluates the storage filter per row and builds + * {@code finalRanges}. + * - Phase 2: read the non-key columns restricted to {@code finalRanges}. A row group that gave + * splicing up reads the whole projection instead, still under {@code finalRanges}, and one + * that gave the filter up reads it under {@code pushedFilterRanges}, which is what a plain + * scan reads. + * + * Row groups for which {@code finalRanges} is empty are skipped entirely (no phase-2 IO). False + * says the row groups have run out, which is what ends a read the filter narrowed. + */ + @Override + protected boolean loadNextRowGroup() throws IOException { + // Nulled once the reader gave the filter up for the whole file, and then the plain path is + // exactly what this read wants. + if (storageFilter == 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 an all-keys row group, which reads + // no data pages at all, does not keep the one before it alive. + closeDataPages(); + while (nextBlockIndex < blocks.size()) { + int blockIdx = nextBlockIndex++; + long blockRowCount = blocks.get(blockIdx).getRowCount(); + if (blockRowCount == 0) { + // parquet-mr never writes these, but an empty block makes parquet's own getRowRanges build + // Range(0, -1) and trip its `from <= to` assertion. The plain read path skips them too. + continue; + } + // Splicing buffers one key value per surviving row of the whole row group before it can emit + // the first batch, and that buffer is outside any MemoryConsumer, so phase 1 counts what it + // holds against `maxSplicedRowGroupBytes` together with the row ranges phase 2 will hold. + // Past that it gives splicing up, and past it again the filter itself. + + // Phase 0: rows allowed by the pushed data filter, at column-index granularity. The full + // requestedSchema goes back on first, because phases 1 and 2 narrow it and + // ParquetFileReader.getRowRanges computes ranges against the reader's 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) { + // Pushed data filter rejects this block entirely via column index. Not a storage-filter + // skip, so we don't increment storage-filter metrics. + continue; + } + + StorageFilterMetrics m = storageFilter.metrics(); + + // Whether phase 2 can read part of this row group at all, which is what everything below + // buffers for. Reading part of one goes through the block's column index store, and parquet + // empties that store altogether if one column it was built over has no offset index. Which + // columns those are depends on when it is first built. With a pushed data filter it is phase + // 0's `getRowRanges` that builds it, over the whole projection, and a key column without an + // offset index is then enough to make phase 2 throw. So the question is asked of the whole + // projection, which also makes the answer the same for every row group of a file whatever + // else is pushed. It is stricter than the read needs with no data filter pushed, where the + // store would be built over the non-key columns alone, and such a row group loses its page + // filtering for nothing. The shape is a file whose columns do not all carry an offset index, + // which no writer produces by itself. + // False says the filter is left with emptying this row group whole, and then nothing is + // buffered either, and phase 2 reads the key columns again along with everything else, over + // ranges an empty store has already widened to the whole block, which is why it cannot throw. + // Parquet's own exception stays as the backstop for the one shape the footer cannot see, an + // offset index it claims but cannot produce. + boolean canNarrowRowGroup = !pageIndexDisabled && !partialReadUnavailable + && blockAccounting.hasOffsetIndexes(blockIdx, requestedPaths); + spliceCurrentRowGroup = canNarrowRowGroup; + + // Phase 1: switch to key-only schema, read key columns under pushedFilterRanges, evaluate the + // storage filter per row. The defaults below are what a row group whose filter is given up + // emits, which is every row of `pushedFilterRanges`, exactly what a plain read would. + RowRanges finalRanges = pushedFilterRanges; + long finalRowCount = baselineRows; + fileReader.setRequestedSchema(keyOnlyColumns); + // Closed at the end of the phase that reads it. `readFilteredRowGroup` hands out a store + // the file reader does not track, unlike `readNextRowGroup`, so nothing else would. + RowRanges survivors; + try (PageReadStore keyPages = + fileReader.readFilteredRowGroup(blockIdx, pushedFilterRanges)) { + if (keyPages == null) { + // Unreachable, since readFilteredRowGroup returns null only for an empty block or empty + // ranges, and this block has rows the pushed filter kept. Should parquet ever say + // otherwise, the filter is given up for this row group, which reads what a plain read + // reads. The alternative is to fail a query a plain read answers. + abandonSplicing(); + survivors = null; + } else { + survivors = evaluateStorageFilter( + keyPages, pushedFilterRanges, baselineRows, blockRowCount, canNarrowRowGroup); + } + } + // Null says phase 1 gave the filter up for this row group, and then the ranges it built are + // incomplete, and what gets read is every row the pushed filter kept, exactly what a plain + // read would read. Phase 2 can give it up too, in the catch below. + boolean filterGivenUp = survivors == null; + // What this feature can avoid reading is the non-key columns of the rows the storage filter + // rejects, so that is the baseline the byte metrics are measured against, the non-key bytes a + // plain read of this projection would transfer for every row the pushed filter kept. Nothing + // is measured for a row group whose filter was given up, so nothing is walked for it either. + // compressedBytesForRowRanges never does IO of its own. + long nonKeyBaselineBytes = filterGivenUp ? 0L : blockAccounting.compressedBytesForRowRanges( + blockIdx, nonKeyPaths, pushedFilterRanges, baselineRows); + if (survivors != null) { + finalRanges = survivors; + finalRowCount = survivors.rowCount(); + if (finalRowCount == 0) { + // Every surviving row was rejected by the storage filter, so the block is skipped + // entirely, which avoids the whole non-key baseline. Phase 1 still paid to read the key + // columns, and that cost is not part of the baseline, so nothing is subtracted from it + // here. + m.recordRowGroupSkipped(baselineRows, nonKeyBaselineBytes); + continue; + } + } + + // Phase 2 reads the non-key columns under the surviving rows, or the whole projection under + // `pushedFilterRanges` for a row group whose filter was given up. + fileReader.setRequestedSchema(spliceCurrentRowGroup ? nonKeyColumns : requestedColumns); + // Reading a strict subset of a block's rows needs a Parquet offset index, and parquet + // enforces that itself, resolving every requested column's offset index before it reads + // anything, and a column without one makes its column index store throw + // MissingOffsetIndexException. Files written before parquet-mr 1.11, or by a writer that + // omits the page index (pyarrow's `write_table` defaults to `write_page_index=False`), have + // none. The filter is then given up for this row group and the read retried over + // `pushedFilterRanges`, which is what a plain scan reads. That retry cannot hit the same + // wall, because a store missing one column's offset index reports no column index either, so + // `getRowRanges` could not have narrowed anything and the ranges cover the whole block. + // + // The footer check above answers every shape a writer produces, so what is left here is an + // index the footer claims and parquet cannot then produce. This is also where such a throw + // costs least, landing before any data page is read. + try { + dataPages = fileReader.readFilteredRowGroup(blockIdx, finalRanges); + } catch (MissingOffsetIndexException e) { + // The exception's own message rather than the exception, since this is a file shape the + // feature expects and degrades on, so a stack trace per file of a table written without a + // page index is noise. Parquet logs the same condition on its own path without one either. + LOG.warn("Reading {} without page-level storage filtering. Reading part of a row group " + + "needs a Parquet offset index, and this file has none for at least one column the " + + "read needs ({}). Row groups the filter empties are still skipped whole", + MDC.of(LogKeys.PATH, fileReader.getFile()), + MDC.of(LogKeys.REASON, e.getMessage())); + // Only the narrowing is given up for the rest of the file. The pushed data filter's own + // page pruning stays, because parquet built those ranges itself and a block it has no index + // for it hands back whole, and dropping it here would read more than a plain scan does. + partialReadUnavailable = true; + filterGivenUp = true; + abandonSplicing(); + finalRanges = pushedFilterRanges; + finalRowCount = baselineRows; + if (baselineRows != blockRowCount) { Review Comment: Done as you suggest. The widening branch is gone, and so are the `noMoreRows()` and `getProgress()` overrides and the `hitEndOfData` flag. `rowsReturned <= totalRowCount` holds on every path now, so the base terminator stays and `nextBatch` is back to master's check. A narrowed read ends when `loadNextRowGroup` runs out of row groups, and a file the filter rejects whole sets `totalRowCount` to zero. The retry now depends on the parquet behavior you traced, that a store missing an offset index leaves the ranges whole. If a later parquet broke that, the retry would throw rather than read more rows. `storageFilter`'s javadoc now says it is nulled only before the first row group is loaded, and that nulling it later would restart the read at the first row group. ########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ValueCopier.java: ########## @@ -0,0 +1,130 @@ +/* + * 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 org.apache.spark.sql.execution.vectorized.WritableColumnVector; +import org.apache.spark.sql.types.*; + +/** + * Copies one key value between column vectors. Picked per key column at init time by + * {@link #forType(DataType)} and called per surviving row; the caller handles null sources. + * + * <p>Nothing here is about Parquet. A copier is chosen from the Spark type alone, and both sides of + * the copy are {@link WritableColumnVector}s. + * + * <p>The types this handles are also the types a storage filter's key columns may have: + * {@code ParquetStorageFilter.isSupportedKeyType} answers from {@link #supports}, so the planner's + * gate, the filter's own checks and the copy itself read one list. That list is narrower than Review Comment: Named in the javadoc as a third exclusion. It gives the two-child layout as the reason and says what you describe, that such a filter is not offered and the scan reads as it would without the feature. I left the copiers for later, since the types sit behind `spark.sql.timestampNanosTypes.enabled` and each needs its own byte charge and tests. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategy.scala: ########## @@ -151,6 +151,66 @@ object FileSourceStrategy extends Strategy with PredicateHelper with Logging { } } + /** + * The conjuncts of `afterScanFilters` the file format can evaluate at the storage layer for late + * materialization, to prune the IO of the columns the filter does not reference. + * + * They stay in the post-scan `Filter` as well, the way a pushed data filter does. The reader is + * offered them, not obliged to honor them, so the plan keeps the exact check. What it costs is + * evaluating the conjunct a second time for the rows the scan emits. What it buys is a reader + * free to give up on a file it cannot read in part, or on a stretch of rows whose survivors cost + * too much to keep track of, without the answer depending on it. + * + * Two of the conditions are per scan, and failing either offers nothing: + * - [[FileFormat.supportsStorageFilterPushdown]] holds. That is where a format reads the conf + * that enables this, so a format's own conf never decides for another format, and asking it + * first keeps everything below off the path of a scan that will not use it. + * - [[FileFormat.supportBatch]] holds for the schema the reader will see, + * `partitionSchema ++ outputDataSchema`, which is the schema a format's reader builder derives + * its own vectorized-read decision from. Late materialization needs a batch read, so this asks + * about batch support rather than naming a format. + * + * The rest are per conjunct: + * - It is deterministic. The reader drops the rows the conjunct rejects and the post-scan + * `Filter` evaluates it again on the rows it keeps, so the two evaluations have to agree. Such + * a conjunct would also fail at task time, since a reader evaluates the predicate without Review Comment: Removed that sentence. The deterministic check keeps the reason in its first sentence, and the suite comment gives that same reason now instead of naming `ParquetStorageFilter.test`. -- 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]
