cloud-fan commented on code in PR #58895:
URL: https://github.com/apache/spark/pull/58895#discussion_r4109912730
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala:
##########
@@ -771,6 +836,28 @@ case class FileSourceScanExec(
readRDD
}
+ // Materialize scalar subqueries inside storage filters to literals and bind
AttributeReferences
+ // to BoundReferences targeting positions in `requiredSchema`. Subqueries
must have been prepared
+ // by SparkPlan before this is forced (same contract as `pushedDownFilters`).
+ @transient
+ protected lazy val preparedStorageFilters: Seq[Expression] = {
+ if (storageFilters.isEmpty) {
+ Nil
+ } else {
+ // No conf check here: the conf decides at planning time whether a scan
is offered storage
+ // filters at all, and re-reading it now could only make this scan drop
work it already has.
+ //
+ // `output` is `readDataColumns ++ generatedMetadataColumns ++
partitionColumns ++
+ // constantMetadataColumns` and `requiredSchema` is the StructType of
the first two groups, so
+ // the first `requiredSchema.length` attributes line up with its fields.
+ val requestedDataAttrs = output.take(requiredSchema.length)
+ storageFilters.map { expr =>
+ val subqueryReplaced = expr.transform { case s:
execution.ScalarSubquery => s.toLiteral }
+ BindReferences.bindReference(subqueryReplaced, requestedDataAttrs)
Review Comment:
**Non-blocking (P2):** `supportsStorageFilter` sees the original predicate,
but this preparation replaces every `AttributeReference` with a
`BoundReference` before `buildReaderWithStorageFilters` runs. That contradicts
the documented same-expression handoff and means a format that opted in using a
column name or field metadata cannot apply the same decision while constructing
its reader. Could we carry both the original expression and an executable bound
form (or otherwise preserve the advertised information)? Please also add an
end-to-end generic `FileFormat` test that exercises the planner-to-reader
handoff and the builder-returning-`None` fallback to
`buildReaderWithPartitionValues`.
**Recommended change:** Represent the original offered expression and its
executable prepared form explicitly at the FileFormat handoff, so the builder
can observe the promised original information while Parquet still evaluates a
materialized, bound form. Add generic FileFormat integration coverage for
identity/metadata preservation and for None falling back to the ordinary reader.
**Why this works:** Prepare each accepted predicate without discarding its
original representation, pass both representations through a single aligned
handoff object or equivalent paired contract, and make Parquet consume the
bound member. Exercise the full planning and scan-construction path with a test
format that records the original expression and can decline specialized
construction.
**Scope:** Refine the generic V1 storage-filter contract, its
planner-to-scan transport, the Parquet consumer adaptation, and regression
coverage.
**Compatibility:** Storage filtering remains a default-off optional
optimization; unsupported formats and runtime give-up paths continue to return
ordinary scan results.
**Risks:** The two representations could become misaligned if they are
stored as independent sequences rather than one paired value. Changing the
newly added FileFormat method shape must keep default decline behavior and
Parquet serialization intact.
**Constraints:** Scalar subqueries must still be materialized before reader
closure serialization. Parquet evaluation must remain bound to requiredSchema
ordinals. The post-scan Filter remains authoritative and the optimization
remains optional.
**Success:** A format can correlate the exact named/metadata-bearing
expression it accepted with the executable predicate used to construct its
reader. Parquet receives a materialized expression bound to the requested
schema and preserves current results and fallback behavior. Returning None from
the specialized builder invokes the ordinary builder and returns its rows. A
regression that canonicalizes or discards the original expression before reader
construction fails the new integration coverage.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -1947,6 +1947,51 @@ 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, allows the vectorized Parquet reader to evaluate runtime
storage filters " +
+ "(e.g. bloom filters from join runtime filtering) at the scan level
using late " +
+ "materialization: read key columns first, evaluate the filter per row,
then read data " +
+ "columns restricted to surviving rows. This is a planning-time
decision only: when " +
+ "false, no storage filter is attached to a scan in the first place and
the filter is " +
+ "applied as an ordinary post-scan filter alone. A pushed filter stays
in the post-scan " +
+ "filter as well, the way a pushed data filter does, so honoring it is
optional: a reader " +
+ "that meets a row group it cannot prune reads it the way a plain scan
would, plus one " +
+ "more read of the key columns, since the phase that evaluated the
filter already read " +
+ "them. A file written with no Parquet page index is read with the
filter applied only " +
+ "where it empties a whole row group, since narrowing to part of one
needs that index, so " +
+ "every row group of it that keeps a row pays that. Setting " +
+ "parquet.filter.columnindex.enabled to false turns this off entirely,
because reading " +
+ "part of a row group goes through the page index. Note that " +
+ "the surviving key values of a whole row group are buffered before the
" +
+ "first batch of that row group 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("Most memory, in bytes, that the vectorized Parquet reader holds
for one row group " +
Review Comment:
**Nit (P3):** The opening `Most memory ...` is a sentence fragment, and the
following independent clauses are joined with commas. Since this text appears
in the generated configuration reference, could you rewrite it as a complete
definition (for example, `The maximum memory ...`) and split the accounting
rules into complete sentences?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/ParquetReadState.java:
##########
@@ -19,30 +19,29 @@
import org.apache.parquet.column.ColumnDescriptor;
-import java.util.ArrayList;
-import java.util.Iterator;
-import java.util.List;
import java.util.PrimitiveIterator;
/**
* Helper class to store intermediate state while reading a Parquet column
chunk.
*/
final class ParquetReadState {
- /** A special row range used when there is no row indexes (hence all rows
must be included) */
- private static final RowRange MAX_ROW_RANGE = new RowRange(Long.MIN_VALUE,
Long.MAX_VALUE);
+ /** The row indexes to include, only not-null if the column index is
present. */
Review Comment:
**Nit (P3):** Non-null row indexes are not conditioned solely on a column
index. The late-materialization path can pass `finalRanges` to
`readFilteredRowGroup`, which derives row indexes from offset indexes even when
the optional column index is absent. Could this comment describe the actual
condition under which `rowIndexes` is populated?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilter.scala:
##########
@@ -0,0 +1,279 @@
+/*
+ * 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.util.concurrent.ConcurrentHashMap
+
+import org.apache.spark.SparkThrowable
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{And, BasePredicate,
BloomFilterMightContain, BoundReference, Expression, Literal, Predicate,
XxHash64}
+import org.apache.spark.sql.catalyst.trees.TreePattern.PLAN_EXPRESSION
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.types.{BinaryType, BooleanType, ByteType,
DataType, DateType, DayTimeIntervalType, DecimalType, DoubleType, FloatType,
IntegerType, LongType, ShortType, StringType, StructType, TimestampNTZType,
TimestampType, TimeType, YearMonthIntervalType}
+
+/**
+ * Optional SQL metrics the reader updates while applying a
[[ParquetStorageFilter]]. Every
+ * counter is scoped to what the storage filter added on top of a read of the
same projection
+ * without one. All fields are nullable; a null field disables that metric.
+ *
+ * - [[rowGroupsSkipped]] counts row groups whose data columns were never
read.
+ * - [[rowsExcludedByRowGroup]] sums the rows those skips excluded, per
skipped 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 what would have
saved it: a row
+ * inside a kept row group is read as part of its page and dropped during
decode, so page
Review Comment:
**Nit (P3):** `rowsExcludedWithinRowGroup` counts all rows rejected inside a
retained row group, but it does not prove that every rejected row's page was
read: phase 2 can skip a page containing no surviving rows. Could this describe
only what the counter establishes, without the row/page IO claim? The adjacent
`rowsExcludedByRowGroup` bullet also needs a complete sentence.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetStorageFilterSuite.scala:
##########
@@ -0,0 +1,2526 @@
+/*
+ * 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.{ByteArrayOutputStream, File}
+import java.net.URI
+import java.time.LocalTime
+import java.util.concurrent.atomic.AtomicLong
+
+import scala.collection.mutable
+import scala.jdk.CollectionConverters._
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{FileStatus, FSDataInputStream, FSInputStream,
Path, RawLocalFileSystem}
+import org.apache.hadoop.mapreduce.Job
+import org.apache.parquet.column.{Encoding, ParquetProperties}
+import org.apache.parquet.column.impl.ColumnWriteStoreV1
+import org.apache.parquet.column.page.DataPageV1
+import org.apache.parquet.column.page.mem.MemPageStore
+import org.apache.parquet.hadoop.{ParquetFileReader, ParquetFileWriter,
ParquetInputFormat, ParquetOutputFormat}
+import org.apache.parquet.hadoop.metadata.{ColumnChunkMetaData,
CompressionCodecName}
+import org.apache.parquet.hadoop.util.HadoopOutputFile
+import org.apache.parquet.schema.MessageTypeParser
+
+import org.apache.spark.paths.SparkPath
+import org.apache.spark.sql.{sources, DataFrame, QueryTest, Row, SparkSession}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute,
AttributeReference, BloomFilterMightContain, BoundReference, Cast, Coalesce,
EqualTo, Expression, GreaterThanOrEqual, IsNull, LessThanOrEqual, Literal, Or,
Predicate, Rand, Remainder, XxHash64}
+import org.apache.spark.sql.catalyst.plans.logical.{Filter => LogicalFilter}
+import org.apache.spark.sql.execution.{CollapseCodegenStages,
ColumnarToRowExec, FileSourceScanExec, FileSourceScanLike, FilterExec,
LocalLimitExec, SparkPlan, WholeStageCodegenExec}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.{FileFormat,
FileSourceStrategy, OutputWriterFactory, PartitionedFile}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+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, ColumnVector}
+import org.apache.spark.unsafe.types.UTF8String
+import org.apache.spark.util.Utils
+import org.apache.spark.util.sketch.BloomFilter
+
+/**
+ * Tests the late-materialization path of [[VectorizedParquetRecordReader]]
driven by a
+ * [[ParquetStorageFilter]]. Writes small multi-row-group parquet files, wires
a hand-built filter
+ * into the reader, and asserts correctness + the two storage-filter metrics.
+ */
+class ParquetStorageFilterSuite extends QueryTest with SharedSparkSession
+ with AdaptiveSparkPlanHelper {
+ import testImplicits._
+
+ // Writes `df` as one parquet file under a fresh directory and returns its
path. Every write
+ // helper in this suite goes through here.
+ private def writeSingleParquetFile(
+ dir: File,
+ df: DataFrame,
+ rowGroupSize: Long,
+ pageSize: Option[Long] = None,
+ dictionary: Boolean = false): String = {
+ val outDir = new File(dir, s"test-${System.nanoTime()}").getAbsolutePath
+ val writer = df
+ .repartition(1)
+ .write
+ .option(ParquetOutputFormat.BLOCK_SIZE, rowGroupSize)
+ // Dictionary encoding off keeps row-group sizing predictable. The
column index is still
+ // written either way.
+ .option(ParquetOutputFormat.ENABLE_DICTIONARY, dictionary.toString)
+ // A small page size gives each row group several pages per column, which
is what lets
+ // column-index filtering produce a row range narrower than the whole row
group.
+ pageSize.foreach(size => writer.option(ParquetOutputFormat.PAGE_SIZE,
size))
+ writer.parquet(outDir)
+ val files = new File(outDir).listFiles((_, name) =>
name.endsWith(".parquet"))
+ assert(files != null && files.length == 1, s"expected exactly one parquet
file under $outDir")
+ files(0).getAbsolutePath
+ }
+
+ // Writes a parquet file with the given rows and row-group size; returns the
path.
+ private def writeParquetFile(
+ dir: File,
+ rows: Seq[(Long, String)],
+ rowGroupSize: Long = 1024L,
+ pageSize: Option[Long] = None): String =
+ writeSingleParquetFile(dir, rows.toDF("k", "v"), rowGroupSize, pageSize)
+
+ // Collects all `(k, v)` rows from a reader initialized with the given
storage filter.
+ private def readAll(
+ filePath: String,
+ storageFilter: ParquetStorageFilter): (Seq[(Long, String)],
VectorizedParquetRecordReader) =
+ readAllWith(filePath, Seq("k", "v"), storageFilter,
+ (batch, i) => (batch.column(0).getLong(i),
batch.column(1).getUTF8String(i).toString))
+
+ // Builds a `k >= threshold` storage filter bound to position 0.
+ private def keyAtLeastFilter(
+ threshold: Long,
+ metrics: StorageFilterMetrics = StorageFilterMetrics()):
ParquetStorageFilter = {
+ val expr = GreaterThanOrEqual(BoundReference(0, LongType, nullable =
false), Literal(threshold))
+ val requested = StructType(Seq(
+ StructField("k", LongType, nullable = false), StructField("v",
StringType, nullable = false)))
+ ParquetStorageFilter.create(Seq(expr), requested, metrics)
+ }
+
+ // Writes a single-column (just `k`) parquet file for the supplied key type
via Spark's
+ // {@code Encoder}.
+ private def writeKeyOnlyParquetFile[T : org.apache.spark.sql.Encoder](
+ dir: File,
+ keys: Seq[T],
+ rowGroupSize: Long = 1024L): String =
+ writeSingleParquetFile(dir, spark.createDataset(keys).toDF("k"),
rowGroupSize)
+
+ // Reads a key-only file, returning the survivor keys and the reader. The
{@code extract} function
+ // pulls one value at a time from the batch's key column.
+ private def readKeyOnlyAll[T](
+ filePath: String,
+ storageFilter: ParquetStorageFilter,
+ extract: (org.apache.spark.sql.vectorized.ColumnVector, Int) => T,
+ capacity: Int = 4096): (Seq[T], VectorizedParquetRecordReader) =
+ readAllWith(filePath, Seq("k"), storageFilter, (batch, i) =>
extract(batch.column(0), i),
+ capacity)
+
+ // Builds a `k >= threshold` storage filter bound to position 0 against a
key-only schema of the
+ // given key type.
+ private def keyOnlyAtLeastFilter(
+ threshold: Literal,
+ keyType: DataType,
+ metrics: StorageFilterMetrics = StorageFilterMetrics()):
ParquetStorageFilter = {
+ val expr = GreaterThanOrEqual(BoundReference(0, keyType, nullable =
false), threshold)
+ val requested = StructType(Seq(StructField("k", keyType, nullable =
false)))
+ ParquetStorageFilter.create(Seq(expr), requested, metrics)
+ }
+
+ test("rejects entire row group: no data-column IO, row-group-skipped metric
incremented") {
+ withTempDir { dir =>
+ // 40 rows in 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 rgSkipped = SQLMetrics.createMetric(spark.sparkContext,
"rowGroupsSkipped")
+ val rowsExcludedRg = SQLMetrics.createMetric(spark.sparkContext,
"rowsExcludedByRowGroup")
+ val rowsExcludedPf = SQLMetrics.createMetric(spark.sparkContext,
"rowsExcludedWithinRowGroup")
+ val bytesAvoidedRg = SQLMetrics.createSizeMetric(spark.sparkContext,
"bytesAvoidedByRg")
+ val bytesAvoidedPf = SQLMetrics.createSizeMetric(spark.sparkContext,
"bytesAvoidedByPf")
+ val filter = keyAtLeastFilter(1000L, StorageFilterMetrics(
+ rowGroupsSkipped = rgSkipped,
+ rowsExcludedByRowGroup = rowsExcludedRg,
+ rowsExcludedWithinRowGroup = rowsExcludedPf,
+ bytesAvoidedByRowGroup = bytesAvoidedRg,
+ bytesAvoidedByPageFiltering = bytesAvoidedPf))
+ val (result, reader) = readAll(path, filter)
+ try {
+ assert(result.isEmpty, "filter rejects all rows; no rows should be
emitted")
+ assert(rgSkipped.value > 0,
+ s"expected at least one row group skipped; got ${rgSkipped.value}")
+ assert(rowsExcludedRg.value > 0,
+ s"expected rows excluded by whole-rowgroup skip; got
${rowsExcludedRg.value}")
+ assert(rowsExcludedPf.value == 0,
+ s"no partial-row-group filtering expected; got
${rowsExcludedPf.value}")
+ // Schema is (k: Long, v: String). Skipping a row group avoids the
v-column bytes the
+ // no-storage-filter path would have read; phase 1 still pays for k.
So avoided > 0.
+ assert(bytesAvoidedRg.value > 0,
+ s"expected non-key bytes avoided by whole row groups; got
${bytesAvoidedRg.value}")
+ assert(bytesAvoidedPf.value == 0,
+ s"no page-filtering bytes expected when all groups skipped; got
${bytesAvoidedPf.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 rgSkipped = SQLMetrics.createMetric(spark.sparkContext,
"rowGroupsSkipped")
+ val rowsExcludedPf = SQLMetrics.createMetric(spark.sparkContext,
"rowsExcludedWithinRowGroup")
+ val filter = keyAtLeastFilter(0L, StorageFilterMetrics(
+ rowGroupsSkipped = rgSkipped, rowsExcludedWithinRowGroup =
rowsExcludedPf))
+ val (result, reader) = readAll(path, filter)
+ try {
+ assert(result.toSet == rows.toSet, s"all rows should round-trip; got
${result.size} rows")
+ assert(rgSkipped.value == 0, s"nothing should be skipped; got
${rgSkipped.value}")
+ assert(rowsExcludedPf.value == 0,
+ s"nothing should be filtered; got ${rowsExcludedPf.value}")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("mixed: some row groups skipped, others partially kept") {
+ withTempDir { dir =>
+ // Many rows + small row groups => guaranteed multiple row groups.
+ val rows = (1L to 200L).map(i => (i, s"v_$i"))
+ val path = writeParquetFile(dir, rows, rowGroupSize = 256L)
+
+ val rgSkipped = SQLMetrics.createMetric(spark.sparkContext,
"rowGroupsSkipped")
+ val rowsExcludedRg = SQLMetrics.createMetric(spark.sparkContext,
"rowsExcludedByRowGroup")
+ val rowsExcludedPf = SQLMetrics.createMetric(spark.sparkContext,
"rowsExcludedWithinRowGroup")
+ // k >= 195 keeps only the last 6 rows; earlier row groups should be
skipped.
+ val filter = keyAtLeastFilter(195L, StorageFilterMetrics(
+ rowGroupsSkipped = rgSkipped,
+ rowsExcludedByRowGroup = rowsExcludedRg,
+ rowsExcludedWithinRowGroup = rowsExcludedPf))
+ val (result, reader) = readAll(path, filter)
+ try {
+ // Output is exact: VectorizedColumnReader uses
PageReadStore.getRowIndexes (driven by
+ // our finalRanges) to skip rows within partial pages, so emitted rows
== survivors.
+ val expected = rows.filter(_._1 >= 195L).toSet
+ assert(result.toSet == expected,
+ s"expected exact filtering; got ${result.map(_._1).sorted}, " +
+ s"expected ${expected.map(_._1).toSeq.sorted}")
+ assert(rgSkipped.value >= 1, s"expected row groups skipped; got
${rgSkipped.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.
+ assert(rowsExcludedPf.value > 0,
+ s"expected rows excluded inside a kept row group; got
${rowsExcludedPf.value}")
+ assert(result.size + rowsExcludedRg.value + rowsExcludedPf.value ==
rows.size,
+ s"${result.size} emitted plus ${rowsExcludedRg.value} plus
${rowsExcludedPf.value} " +
+ s"should account for all ${rows.size} rows")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("key-only projection: phase 2 is skipped and both byte-avoided metrics
are zero") {
+ // When the projected schema contains only the bloom key, phase 2 is
skipped entirely
+ // (`nonKeyColumns == null` in the reader). All output rows come from the
per-key-column
+ // queues populated in phase 1. Total bytes read match the
no-storage-filter path (phase 1 reads
+ // the key column once instead of phase 2 re-reading it), so both
`avoided` metrics are zero:
+ // there are no non-key bytes to skip. This is the shape the design notes
call the biggest win,
+ // so it must not be the shape that pays for metrics.
+ withTempDir { dir =>
+ val keys = (1L to 200L)
+ val path = writeKeyOnlyParquetFile(dir, keys, rowGroupSize = 256L)
+
+ val rgSkipped = SQLMetrics.createMetric(spark.sparkContext,
"rowGroupsSkipped")
+ val rowsExcludedRg = SQLMetrics.createMetric(spark.sparkContext,
"rowsExcludedByRowGroup")
+ val rowsExcludedPf = SQLMetrics.createMetric(spark.sparkContext,
"rowsExcludedWithinRowGroup")
+ val bytesAvoidedRg = SQLMetrics.createSizeMetric(spark.sparkContext,
"bytesAvoidedByRg")
+ val bytesAvoidedPf = SQLMetrics.createSizeMetric(spark.sparkContext,
"bytesAvoidedByPf")
+ // k >= 195 -> last 6 keys survive; preceding row groups skipped or
page-pruned.
+ val filter = keyOnlyAtLeastFilter(Literal(195L), LongType,
StorageFilterMetrics(
+ rowGroupsSkipped = rgSkipped,
+ rowsExcludedByRowGroup = rowsExcludedRg,
+ rowsExcludedWithinRowGroup = rowsExcludedPf,
+ bytesAvoidedByRowGroup = bytesAvoidedRg,
+ bytesAvoidedByPageFiltering = bytesAvoidedPf))
+ val (result, reader) = readKeyOnlyAll(path, filter, (vec, i) =>
vec.getLong(i))
+ try {
+ val expected = keys.filter(_ >= 195L).toSet
+ assert(result.toSet == expected,
+ s"expected exact survivor keys; got ${result.sorted}, expected
${expected.toSeq.sorted}")
+ assert(rgSkipped.value >= 1, s"expected row groups skipped; got
${rgSkipped.value}")
+ assert(rowsExcludedRg.value > 0,
+ s"a skipped row group must count its rows too; got
${rowsExcludedRg.value}")
+ // The all-keys path takes its kept-row count from `finalRowCount`
rather than from a
+ // phase-2 page store, so this identity is the only thing that checks
that arithmetic.
+ assert(result.size + rowsExcludedRg.value + rowsExcludedPf.value ==
keys.size,
+ s"${result.size} emitted plus ${rowsExcludedRg.value} plus
${rowsExcludedPf.value} " +
+ s"should account for all ${keys.size} rows")
+ // For an all-keys projection, the no-storage-filter path would have
read the same key
+ // column phase 1 reads. There are no non-key bytes to avoid; both
metrics are 0.
+ assert(bytesAvoidedRg.value == 0,
+ s"expected no non-key bytes to avoid on all-keys projection; got
${bytesAvoidedRg.value}")
+ assert(bytesAvoidedPf.value == 0,
+ s"expected no non-key bytes to avoid on all-keys projection; got
${bytesAvoidedPf.value}")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("multi-batch emit: survivor count exceeds capacity") {
+ // Drive the reader at capacity = 16 with a row group of 100 surviving
rows. Exercises:
+ // - The per-key-column queue holding multiple full-capacity vectors
plus a partial tail.
+ // - The published queue head getting closed at the start of every
subsequent emit.
+ // - The batch's key slots being rewritten ceil(100/16) = 7 times.
+ withTempDir { dir =>
+ val keys = (1L to 100L)
+ // Big rowGroupSize so all 100 rows fit in one row group.
+ val path = writeKeyOnlyParquetFile(dir, keys, rowGroupSize = 64 * 1024L)
+ // Filter accepts every row so the queue is fully populated.
+ val filter = keyOnlyAtLeastFilter(Literal(0L), LongType)
+ val (result, reader) =
+ readKeyOnlyAll(path, filter, (vec, i) => vec.getLong(i), capacity = 16)
+ try {
+ assert(result == keys.toSeq,
+ s"expected all keys returned in order across multiple batches; got
${result.size} rows")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("int key column: filter survivors round-trip through phase 1
accumulators") {
+ // Covers the IntegerType branch of ValueCopier.
+ withTempDir { dir =>
+ val keys = (1 to 100)
+ val path = writeKeyOnlyParquetFile(dir, keys, rowGroupSize = 256L)
+ val filter = keyOnlyAtLeastFilter(Literal(90), IntegerType)
+ val (result, reader) = readKeyOnlyAll(path, filter, (vec, i) =>
vec.getInt(i))
+ try {
+ assert(result.toSet == keys.filter(_ >= 90).toSet,
+ s"expected int keys >= 90; got ${result.sorted}")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("string key column: filter survivors round-trip through phase 1
accumulators") {
+ // Covers the StringType branch of ValueCopier (variable-length byte copy
via putByteArray).
+ withTempDir { dir =>
+ val keys = (1 to 20).map(i => f"k$i%03d")
+ val path = writeKeyOnlyParquetFile(dir, keys, rowGroupSize = 256L)
+ val filter = keyOnlyAtLeastFilter(
+ Literal.create("k015", StringType), StringType)
+ val (result, reader) =
+ readKeyOnlyAll(path, filter, (vec, i) => vec.getUTF8String(i).toString)
+ try {
+ assert(result.toSet == keys.filter(_ >= "k015").toSet,
+ s"expected string keys >= 'k015'; got ${result.sorted}")
+ } finally {
+ reader.close()
+ }
+ }
+ }
+
+ test("ParquetStorageFilter.create rejects a filter that violates a planner
precondition") {
+ // These are all planner bugs by construction: storageFiltersFor
pre-checks each one, and
+ // by the time create runs the conjunct is gone from the post-scan Filter,
so a soft rejection
Review Comment:
**Nit (P3):** Two comments in this suite describe a different contract from
the one being tested: the suite exercises five storage-filter metrics rather
than two, and `FileSourceStrategy` deliberately retains this conjunct in the
post-scan `Filter` while also attaching it to the scan. Could you update both
comments so they document the actual coverage and fallback-safety mechanism?
--
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]
