peter-toth commented on code in PR #58895:
URL: https://github.com/apache/spark/pull/58895#discussion_r4133528530
##########
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:
The binding is documented now rather than carried twice.
`buildReaderWithStorageFilters`'s scaladoc says the expressions arrive bound to
`requiredSchema`, as `BoundReference`s over the row a reader assembles from it,
which is the only form a reader can evaluate per row without resolving names
again. `supportsStorageFilter` is asked about the unbound expression, where the
names and the metadata are still there, so that is where a format decides. The
two are index-aligned, one prepared filter per offered conjunct in the same
order.
I did not add a second sequence or a pair type. It doubles a public
`FileFormat` method's contract for a format that would then have to correlate
the two itself, and the drift you name is the reason not to. If a format ever
needs the original at build time, the aligned pair is the change to make then.
The fallback is covered end to end now. A scan whose relation carries a
`ParquetFileFormat` subclass and a storage filter returns every row of the file
through `buildReaderWithPartitionValues`, and the test fails if the subclass
guard is removed.
##########
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:
Rewritten. It opens with `The maximum memory, in bytes, that ...` and the
accounting rules are separate sentences. The other conf doc was shortened in
the same pass.
##########
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:
The comment is gone. `ParquetReadState` is a small class per way of saying
which rows to include, and one factory picks between them. Its javadoc states
the actual condition: ranges when the caller of the read holds them, which is
this feature's phase 2, the page store's row indexes when only the store can
describe them, which is parquet's own filtering, and neither, which means every
row of the chunk.
##########
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:
Rewritten. The doc says the counters' suffix is where a row was excluded,
not whether reading it was avoided, and it names both cases, a row read as part
of a page that held a survivor and a row whose page phase 2 skipped. What a
page-level saving was is what the byte counters say, and they are measured
against the non-key bytes a plain read of the same projection would have
transferred.
##########
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:
Both rewritten. The metric test asserts all five counters, and the `create`
test says the rows are safe either way because the conjunct stays in the
post-scan `Filter`, which is what makes a hard rejection the right answer to a
planner bug.
--
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]
