From b9bc4d1420313da0005cb1da5a23e281b2af5563 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 22 Jul 2026 15:41:39 -0700 Subject: [PATCH 01/10] [SPARK-58272][SQL] Enable runtime Bloom filters for materialized cached inputs --- .../optimizer/InjectRuntimeFilter.scala | 98 +++++-- .../catalyst/plans/logical/LogicalPlan.scala | 12 + .../apache/spark/sql/internal/SQLConf.scala | 20 +- .../execution/columnar/InMemoryRelation.scala | 242 +++++++++++++++--- .../execution/datasources/FileScanRDD.scala | 4 + .../sql/util/PartitionKeyedAccumulator.scala | 55 ++-- .../spark/sql/InjectRuntimeFilterSuite.scala | 239 +++++++++++++++++ ...ntimeBloomFilterCachedInputBenchmark.scala | 139 ++++++++++ .../columnar/InMemoryColumnarQuerySuite.scala | 124 ++++++++- .../util/PartitionKeyedAccumulatorSuite.scala | 22 ++ 10 files changed, 865 insertions(+), 90 deletions(-) create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/RuntimeBloomFilterCachedInputBenchmark.scala diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala index a8f5b3ca67ed0..0e16e9ff1f8ba 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala @@ -36,26 +36,37 @@ import org.apache.spark.sql.internal.SQLConf */ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with JoinSelectionHelper { + private case class FilterCreationSide( + key: Expression, + plan: LogicalPlan, + useMaterializedThreshold: Boolean, + hasSelectivePredicate: Boolean, + materializedRowCount: Option[BigInt] = None) + private def injectFilter( filterApplicationSideKey: Expression, filterApplicationSidePlan: LogicalPlan, - filterCreationSideKey: Expression, - filterCreationSidePlan: LogicalPlan): LogicalPlan = { + filterCreationSide: FilterCreationSide): LogicalPlan = { injectBloomFilter( filterApplicationSideKey, filterApplicationSidePlan, - filterCreationSideKey, - filterCreationSidePlan + filterCreationSide ) } private def injectBloomFilter( filterApplicationSideKey: Expression, filterApplicationSidePlan: LogicalPlan, - filterCreationSideKey: Expression, - filterCreationSidePlan: LogicalPlan): LogicalPlan = { + filterCreationSide: FilterCreationSide): LogicalPlan = { + val filterCreationSideKey = filterCreationSide.key + val filterCreationSidePlan = filterCreationSide.plan + val creationSideThreshold = if (filterCreationSide.useMaterializedThreshold) { + conf.runtimeFilterMaterializedCreationSideThreshold + } else { + conf.runtimeFilterCreationSideThreshold + } // Skip if the filter creation side is too big - if (filterCreationSidePlan.stats.sizeInBytes > conf.runtimeFilterCreationSideThreshold) { + if (filterCreationSidePlan.stats.sizeInBytes > creationSideThreshold) { return filterApplicationSidePlan } val rowCount = filterCreationSidePlan.stats.rowCount @@ -81,23 +92,21 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J } /** - * Extracts a sub-plan which is a simple filter over scan from the input plan. The simple - * filter should be selective and the filter condition (including expressions in the child - * plan referenced by the filter condition) should be a simple expression, so that we do - * not add a subquery that might have an expensive computation. The extracted sub-plan should - * produce a superset of the entire creation side output data, so that it's still correct to - * use the sub-plan to build the runtime filter to prune the application side. + * Extracts either a safely materialized leaf with accurate statistics or a simple selective + * filter over a scan. Filter conditions and the expressions they reference must remain simple, + * so the runtime-filter subquery does not introduce expensive computation. The extracted plan + * must produce a superset of the creation side's join keys. */ private def extractSelectiveFilterOverScan( plan: LogicalPlan, - filterCreationSideKey: Expression): Option[(Expression, LogicalPlan)] = { + filterCreationSideKey: Expression): Option[FilterCreationSide] = { def extract( p: LogicalPlan, predicateReference: AttributeSet, hasHitFilter: Boolean, hasHitSelectiveFilter: Boolean, currentPlan: LogicalPlan, - targetKey: Expression): Option[(Expression, LogicalPlan)] = p match { + targetKey: Expression): Option[FilterCreationSide] = p match { case Project(projectList, child) if hasHitFilter => // We need to make sure all expressions referenced by filter predicates are simple // expressions. @@ -175,8 +184,29 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J } else { None } + case leaf: LeafNodeWithAccurateStats => + val safeLineage = currentPlan.deterministic && + findExpressionAndTrackLineageDown(targetKey, currentPlan).exists { + case (trackedKey, _) => isSimpleExpression(trackedKey) + } + if (leaf.statsAvailable && leaf.isOutputRepeatable && safeLineage) { + leaf.stats.rowCount.map { rowCount => + FilterCreationSide( + targetKey, + currentPlan, + useMaterializedThreshold = true, + hasSelectivePredicate = hasHitSelectiveFilter || leaf.hasSelectivePredicate, + materializedRowCount = Some(rowCount)) + } + } else { + None + } case _: LeafNode if hasHitSelectiveFilter => - Some((targetKey, currentPlan)) + Some(FilterCreationSide( + targetKey, + currentPlan, + useMaterializedThreshold = false, + hasSelectivePredicate = true)) case _ => None } @@ -237,18 +267,37 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J * Extracts the beneficial filter creation plan with check show below: * - The filterApplicationSideKey can be pushed down through joins, aggregates and windows * (ie the expression references originate from a single leaf node) - * - The filter creation side has a selective predicate + * - The filter creation side has a selective predicate, or its exact materialized row count + * is smaller than the application side's distinct join-key count * - The max filterApplicationSide scan size is greater than a configurable threshold */ private def extractBeneficialFilterCreatePlan( filterApplicationSide: LogicalPlan, filterCreationSide: LogicalPlan, filterApplicationSideKey: Expression, - filterCreationSideKey: Expression): Option[(Expression, LogicalPlan)] = { + filterCreationSideKey: Expression): Option[FilterCreationSide] = { if (findExpressionAndTrackLineageDown( filterApplicationSideKey, filterApplicationSide).isDefined && satisfyByteSizeRequirement(filterApplicationSide)) { - extractSelectiveFilterOverScan(filterCreationSide, filterCreationSideKey) + extractSelectiveFilterOverScan(filterCreationSide, filterCreationSideKey).filter { + creationSide => + creationSide.hasSelectivePredicate || { + def distinctCount(key: Expression, plan: LogicalPlan): Option[BigInt] = { + key.references.toSeq match { + case Seq(attribute) => plan.stats.attributeStats.get(attribute) + .flatMap(_.distinctCount) + case _ => None + } + } + val applicationDistinctCount = findExpressionAndTrackLineageDown( + filterApplicationSideKey, filterApplicationSide).flatMap { + case (trackedKey, origin) => distinctCount(trackedKey, origin) + }.orElse(distinctCount(filterApplicationSideKey, filterApplicationSide)) + creationSide.materializedRowCount.exists { rowCount => + applicationDistinctCount.exists(_ > rowCount) + } + } + } } else { None } @@ -307,9 +356,8 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J val hasShuffle = isProbablyShuffleJoin(left, right, hint) if (canPruneLeft(joinType) && (hasShuffle || probablyHasShuffle(left)) && !hasBloomFilter(newLeft, l)) { - extractBeneficialFilterCreatePlan(left, right, l, r).foreach { - case (filterCreationSideKey, filterCreationSidePlan) => - newLeft = injectFilter(l, newLeft, filterCreationSideKey, filterCreationSidePlan) + extractBeneficialFilterCreatePlan(left, right, l, r).foreach { creationSide => + newLeft = injectFilter(l, newLeft, creationSide) } } // Did we actually inject on the left? If not, try on the right @@ -320,10 +368,8 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J // 3. There is no bloom filter on the right key yet if (newLeft.fastEquals(oldLeft) && canPruneRight(joinType) && (hasShuffle || probablyHasShuffle(right)) && !hasBloomFilter(newRight, r)) { - extractBeneficialFilterCreatePlan(right, left, r, l).foreach { - case (filterCreationSideKey, filterCreationSidePlan) => - newRight = injectFilter( - r, newRight, filterCreationSideKey, filterCreationSidePlan) + extractBeneficialFilterCreatePlan(right, left, r, l).foreach { creationSide => + newRight = injectFilter(r, newRight, creationSide) } } if (!newLeft.fastEquals(oldLeft) || !newRight.fastEquals(oldRight)) { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala index d573f48862541..f04297dafb509 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala @@ -242,6 +242,18 @@ trait LeafNode extends LogicalPlan with LeafLike[LogicalPlan] { throw new SparkUnsupportedOperationException("_LEGACY_ERROR_TEMP_3114") } +/** A materialized leaf whose output can safely be scanned again to build a runtime filter. */ +private[sql] trait LeafNodeWithAccurateStats extends LeafNode { + /** Whether the current materialized output has complete, accurate statistics. */ + def statsAvailable: Boolean + + /** Whether scanning the materialized output again returns the same rows. */ + def isOutputRepeatable: Boolean + + /** Whether the original plan contains a predicate that is likely to be selective. */ + def hasSelectivePredicate: Boolean +} + /** * A abstract class for LogicalQueryStage that is visible in logical rewrites. */ diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 45908f299bf6f..e072243293326 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -801,8 +801,9 @@ object SQLConf { val RUNTIME_BLOOM_FILTER_ENABLED = buildConf("spark.sql.optimizer.runtime.bloomFilter.enabled") - .doc("When true and if one side of a shuffle join has a selective predicate, we attempt " + - "to insert a bloom filter in the other side to reduce the amount of shuffle data.") + .doc("When true and if one side of a shuffle join has a selective predicate, or is a " + + "fully materialized, repeatable cache with evidence that pruning is beneficial, we " + + "attempt to insert a bloom filter in the other side to reduce shuffle data.") .version("3.3.0") .booleanConf .createWithDefault(true) @@ -810,11 +811,21 @@ object SQLConf { val RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD = buildConf("spark.sql.optimizer.runtime.bloomFilter.creationSideThreshold") .doc("Size threshold of the bloom filter creation side plan. Estimated size needs to be " + - "under this value to try to inject bloom filter.") + "under this value to try to inject a bloom filter, unless the creation side is fully " + + "materialized and has accurate statistics.") .version("3.3.0") .bytesConf(ByteUnit.BYTE) .createWithDefaultString("10MB") + val RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD = + buildConf("spark.sql.optimizer.runtime.bloomFilter.materializedCreationSideThreshold") + .doc("Size threshold of a fully materialized, repeatable bloom filter creation side with " + + "accurate statistics. This replaces the general creation-side threshold because scanning " + + "materialized output does not recompute its original plan.") + .version("5.0.0") + .bytesConf(ByteUnit.BYTE) + .createWithDefaultString("100MB") + val RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD = buildConf("spark.sql.optimizer.runtime.bloomFilter.applicationSideScanSizeThreshold") .doc("Byte size threshold of the Bloom filter application side plan's aggregated scan " + @@ -8271,6 +8282,9 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def runtimeFilterCreationSideThreshold: Long = getConf(RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD) + def runtimeFilterMaterializedCreationSideThreshold: Long = + getConf(RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD) + def runtimeRowLevelOperationGroupFilterEnabled: Boolean = getConf(RUNTIME_ROW_LEVEL_OPERATION_GROUP_FILTER_ENABLED) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala index e19d562c19b5b..68dda3b532fd6 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala @@ -22,17 +22,25 @@ import com.esotericsoftware.kryo.io.{Input => KryoInput, Output => KryoOutput} import org.apache.spark.{SparkException, TaskContext} import org.apache.spark.network.util.JavaUtils -import org.apache.spark.rdd.RDD -import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.rdd.{DeterministicLevel, RDD} +import org.apache.spark.sql.catalyst.{FileSourceOptions, InternalRow} import org.apache.spark.sql.catalyst.analysis.MultiInstanceRelation import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.plans.{logical, QueryPlan} import org.apache.spark.sql.catalyst.plans.logical.{ColumnStat, LogicalPlan, Statistics} +import org.apache.spark.sql.catalyst.trees.TreePattern.CURRENT_LIKE import org.apache.spark.sql.catalyst.types.DataTypeUtils -import org.apache.spark.sql.catalyst.util.truncatedString +import org.apache.spark.sql.catalyst.util.{truncatedString, CaseInsensitiveMap} import org.apache.spark.sql.columnar.{CachedBatch, CachedBatchSerializer, SimpleMetricsCachedBatch, SimpleMetricsCachedBatchSerializer} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec +import org.apache.spark.sql.execution.datasources.{FileFormat, FileScanRDD, HadoopFsRelation, LogicalRelation} +import org.apache.spark.sql.execution.datasources.binaryfile.BinaryFileFormat +import org.apache.spark.sql.execution.datasources.csv.CSVFileFormat +import org.apache.spark.sql.execution.datasources.json.JsonFileFormat +import org.apache.spark.sql.execution.datasources.orc.OrcFileFormat +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.execution.datasources.text.TextFileFormat import org.apache.spark.sql.execution.vectorized.{OffHeapColumnVector, OnHeapColumnVector, WritableColumnVector} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.types._ @@ -257,10 +265,15 @@ case class CachedRDDBuilder( storageLevel: StorageLevel, @transient cachedPlan: SparkPlan, tableName: Option[String], - @transient logicalPlan: LogicalPlan) { + @transient logicalPlan: LogicalPlan, + isCachedLogicalPlanRepeatable: Boolean = false, + hasSelectivePredicate: Boolean = false, + fileSourceOptions: Seq[Map[String, String]] = Seq.empty) { @transient @volatile private var _cachedColumnBuffers: RDD[CachedBatch] = null - @transient @volatile private var _cachedColumnBuffersAreLoaded: Boolean = false + @volatile private var isCachedRDDRepeatable = false + private var hasStrictFileSourceReads = true + private var _materializedStats: Option[(Long, Long)] = None // The cache's materialization bookkeeping: a partition-keyed accumulator storing // (rowCount, sizeInBytes) per partition. AQE creates a separate cache scan stage per reference to @@ -300,35 +313,40 @@ case class CachedRDDBuilder( if (_cachedColumnBuffers != null) { _cachedColumnBuffers.unpersist(blocking) _cachedColumnBuffers = null - // The buffers no longer back a live RDD. Reset the one-way "loaded" latch and install new - // bookkeeping so a rebuild on this builder does not inherit stale state or late updates from - // tasks that captured the previous generation's accumulator. - _cachedColumnBuffersAreLoaded = false + _materializedStats = None partitionStats = newPartitionStats() } + isCachedRDDRepeatable = false } - def isCachedColumnBuffersLoaded: Boolean = synchronized { - _cachedColumnBuffers != null && isCachedRDDLoaded - } + def isCachedColumnBuffersLoaded: Boolean = loadedMaterializedStats.isDefined + + private[sql] def isCachedPlanRepeatable: Boolean = + isCachedLogicalPlanRepeatable && isCachedRDDRepeatable - private def isCachedRDDLoaded: Boolean = { - _cachedColumnBuffersAreLoaded || { - // We must make sure the statistics of `sizeInBytes` and `rowCount` are accurate if - // `isCachedRDDLoaded` return true. Otherwise, AQE would do a wrong optimization, - // e.g., convert a non-empty plan to empty local relation if `rowCount` is 0. - // Count DISTINCT materialized partitions (the keyed accumulator's key set), so the cache is - // only reported loaded once every partition has been computed -- sound even if a partition is - // computed more than once by concurrent or speculative tasks. - val numMaterialized = partitionStats.accumulatedNumPartitions - val rddLoaded = _cachedColumnBuffers.partitions.length.toLong == numMaterialized - if (rddLoaded) { - _cachedColumnBuffersAreLoaded = rddLoaded + /** Reads completeness and exact statistics from one cache generation atomically. */ + private[sql] def loadedMaterializedStats: Option[(Long, Long)] = synchronized { + if (_cachedColumnBuffers == null) { + None + } else { + _materializedStats.orElse { + partitionStats.foldValuesIfComplete( + _cachedColumnBuffers.partitions.length, + (0L, 0L)) { + case ((rows, bytes), (partitionRows, partitionBytes)) => + (rows + partitionRows, bytes + partitionBytes) + }.map { stats => + _materializedStats = Some(stats) + stats + } } - rddLoaded } } + private[sql] def repeatableMaterializedStats: Option[(Long, Long)] = synchronized { + if (isCachedPlanRepeatable) loadedMaterializedStats else None + } + // Reported row count / size for the cache's statistics: exact and de-duplicated, folded over the // distinct materialized partitions. Synchronized so a fold never races a concurrent `clearCache` // reset. @@ -347,16 +365,51 @@ case class CachedRDDBuilder( } private def buildBuffers(): RDD[CachedBatch] = { + def buildInputRDD[T](input: => RDD[T]): RDD[T] = { + if (fileSourceOptions.isEmpty) { + input + } else { + // File scans initialize their input RDD lazily. Check both configuration domains while + // constructing that RDD so a temporary best-effort setting cannot be mistaken for a + // repeatable strict read. + val materializationConf = SQLConf.get.clone() + val cachedPlanConf = cachedPlan.conf.clone() + + def hasStrictReads(conf: SQLConf): Boolean = SQLConf.withExistingConf(conf) { + fileSourceOptions.forall { options => + val effectiveOptions = new FileSourceOptions(options) + !effectiveOptions.ignoreMissingFiles && !effectiveOptions.ignoreCorruptFiles + } + } + + val (inputRDD, strictPhysicalReads) = SQLConf.withExistingConf(materializationConf) { + val result = input + val fileScans = cachedPlan.collect { case scan: FileSourceScanExec => scan } + val scansAreStrict = fileScans.size == fileSourceOptions.size && fileScans.forall { + scan => scan.inputRDD match { + case fileRDD: FileScanRDD => fileRDD.hasStrictFileReads + case _ => false + } + } + (result, scansAreStrict) + } + hasStrictFileSourceReads = hasStrictFileSourceReads && + hasStrictReads(materializationConf) && hasStrictReads(cachedPlanConf) && + strictPhysicalReads + inputRDD + } + } + val cb = try { if (supportsColumnarInput) { serializer.convertColumnarBatchToCachedBatch( - cachedPlan.executeColumnar(), + buildInputRDD(cachedPlan.executeColumnar()), cachedPlan.output, storageLevel, cachedPlan.conf) } else { serializer.convertInternalRowToCachedBatch( - cachedPlan.execute(), + buildInputRDD(cachedPlan.execute()), cachedPlan.output, storageLevel, cachedPlan.conf) @@ -374,9 +427,8 @@ case class CachedRDDBuilder( // id. Bound to a local so the task closure below captures only the accumulator, not the // enclosing CachedRDDBuilder (whose cachedPlan is not serializable). val accumulator = partitionStats - val cached = cb.mapPartitionsInternal { it => + val cached = cb.mapPartitionsWithIndexInternal { (partitionId, it) => val taskContext = TaskContext.get() - val partitionId = taskContext.partitionId() // This task computes exactly one partition. Tally its totals so the completion listener // records them once, keyed by partition id (covering empty-output partitions, which produce // no batches). @@ -403,11 +455,106 @@ case class CachedRDDBuilder( } }.persist(storageLevel) cached.setName(cachedName) + isCachedRDDRepeatable = hasStrictFileSourceReads && + cached.outputDeterministicLevel != DeterministicLevel.INDETERMINATE && + InMemoryRelation.hasRepeatablePhysicalPlan(cachedPlan) cached } } -object InMemoryRelation { +object InMemoryRelation extends PredicateHelper { + + private val trustedFileFormatClasses: Set[Class[_ <: FileFormat]] = Set( + classOf[BinaryFileFormat], + classOf[CSVFileFormat], + classOf[JsonFileFormat], + classOf[OrcFileFormat], + classOf[ParquetFileFormat], + classOf[TextFileFormat]) + + private val trustedExternalFileFormatNames = Set( + "org.apache.spark.sql.avro.AvroFileFormat", + "org.apache.spark.sql.hive.orc.OrcFileFormat") + + private def hasSafeExpressions(plan: QueryPlan[_]): Boolean = { + plan.expressions.forall { expression => + !expression.exists { + case _: AesEncrypt | _: NonSQLExpression | _: UserDefinedExpression => true + case value => !value.deterministic || value.containsPattern(CURRENT_LIKE) || + !value.getClass.getName.startsWith("org.apache.spark.sql.catalyst.expressions.") + } + } + } + + private def hasRepeatableLogicalPlan(analyzedPlan: LogicalPlan, plan: LogicalPlan): Boolean = { + // Runtime-replaceable expressions such as AES encryption can become deterministic-looking + // StaticInvoke nodes during optimization despite using a fresh random initialization vector. + // Inspect the original analyzed expressions before trusting the optimized execution shape. + analyzedPlan.deterministic && analyzedPlan.collectWithSubqueries { + case node if !hasSafeExpressions(node) => true + }.isEmpty && plan.deterministic && plan.collectWithSubqueries { + case node if !hasSafeExpressions(node) => true + case _: logical.Project | _: logical.Filter | _: logical.SubqueryAlias | + _: logical.Range | _: logical.LocalRelation => false + case relation: LogicalRelation => relation.relation match { + case fileRelation: HadoopFsRelation => + val fileFormatClass = fileRelation.fileFormat.getClass + !(trustedFileFormatClasses.contains(fileFormatClass) || + trustedExternalFileFormatNames.contains(fileFormatClass.getName)) + case _ => true + } + case _ => true + }.forall(!_) + } + + private[columnar] def hasRepeatablePhysicalPlan(plan: SparkPlan): Boolean = { + !plan.exists { node => + val supported = node match { + case _: ColumnarToRowExec | _: FileSourceScanExec | _: FilterExec | + _: InputAdapter | _: LocalTableScanExec | _: ProjectExec | + _: RangeExec | _: WholeStageCodegenExec => true + case _ => false + } + !supported || node.subqueries.nonEmpty || !hasSafeExpressions(node) + } + } + + private def collectFileSourceOptions(plan: LogicalPlan): Seq[Map[String, String]] = { + val relevantOptions = Seq( + FileSourceOptions.IGNORE_MISSING_FILES, + FileSourceOptions.IGNORE_CORRUPT_FILES) + plan.collectWithSubqueries { + case relation: LogicalRelation if relation.relation.isInstanceOf[HadoopFsRelation] => + val options = CaseInsensitiveMap(relation.relation.asInstanceOf[HadoopFsRelation].options) + relevantOptions.flatMap { key => options.get(key).map(key -> _) }.toMap + } + } + + private def hasSelectivePredicate(plan: LogicalPlan): Boolean = { + plan.collectWithSubqueries { + case logical.Filter(condition, _) if condition.deterministic && + isLikelySelective(condition) => true + }.nonEmpty + } + + private def newCacheBuilder( + serializer: CachedBatchSerializer, + storageLevel: StorageLevel, + cachedPlan: SparkPlan, + tableName: Option[String], + logicalPlan: LogicalPlan, + analyzedPlan: LogicalPlan, + optimizedPlan: LogicalPlan): CachedRDDBuilder = { + CachedRDDBuilder( + serializer, + storageLevel, + cachedPlan, + tableName, + logicalPlan, + isCachedLogicalPlanRepeatable = hasRepeatableLogicalPlan(analyzedPlan, optimizedPlan), + hasSelectivePredicate = hasSelectivePredicate(optimizedPlan), + fileSourceOptions = collectFileSourceOptions(optimizedPlan)) + } private[this] var ser: Option[CachedBatchSerializer] = None private[this] def getSerializer(sqlConf: SQLConf): CachedBatchSerializer = synchronized { @@ -434,8 +581,8 @@ object InMemoryRelation { } else { qe.executedPlan } - val cacheBuilder = - CachedRDDBuilder(serializer, storageLevel, child, tableName, qe.logical) + val cacheBuilder = newCacheBuilder( + serializer, storageLevel, child, tableName, qe.logical, qe.analyzed, optimizedPlan) val relation = new InMemoryRelation(child.output, cacheBuilder, optimizedPlan.outputOrdering) relation.statsOfPlanToCache = optimizedPlan.stats relation @@ -450,8 +597,8 @@ object InMemoryRelation { child: SparkPlan, tableName: Option[String], optimizedPlan: LogicalPlan): InMemoryRelation = { - val cacheBuilder = - CachedRDDBuilder(serializer, storageLevel, child, tableName, optimizedPlan) + val cacheBuilder = newCacheBuilder( + serializer, storageLevel, child, tableName, optimizedPlan, optimizedPlan, optimizedPlan) val relation = new InMemoryRelation(child.output, cacheBuilder, optimizedPlan.outputOrdering) relation.statsOfPlanToCache = optimizedPlan.stats relation @@ -465,7 +612,14 @@ object InMemoryRelation { } else { qe.executedPlan } - val newBuilder = cacheBuilder.copy(cachedPlan = newCachedPlan, logicalPlan = qe.logical) + val newBuilder = newCacheBuilder( + serializer, + cacheBuilder.storageLevel, + newCachedPlan, + cacheBuilder.tableName, + qe.logical, + qe.analyzed, + optimizedPlan) val relation = new InMemoryRelation( newBuilder.cachedPlan.output, newBuilder, optimizedPlan.outputOrdering) relation.statsOfPlanToCache = optimizedPlan.stats @@ -487,7 +641,7 @@ case class InMemoryRelation( output: Seq[Attribute], @transient cacheBuilder: CachedRDDBuilder, override val outputOrdering: Seq[SortOrder]) - extends logical.LeafNode with MultiInstanceRelation { + extends logical.LeafNodeWithAccurateStats with MultiInstanceRelation { @volatile var statsOfPlanToCache: Statistics = null @@ -500,6 +654,13 @@ case class InMemoryRelation( def cachedPlan: SparkPlan = cacheBuilder.cachedPlan + override def statsAvailable: Boolean = + cacheBuilder.storageLevel.useDisk && cacheBuilder.repeatableMaterializedStats.isDefined + + override def isOutputRepeatable: Boolean = cacheBuilder.repeatableMaterializedStats.isDefined + + override def hasSelectivePredicate: Boolean = cacheBuilder.hasSelectivePredicate + private[sql] def updateStats( rowCount: Long, newColStats: Map[Attribute, ColumnStat]): Unit = this.synchronized { @@ -511,14 +672,11 @@ case class InMemoryRelation( } override def computeStats(): Statistics = { - if (!cacheBuilder.isCachedColumnBuffersLoaded) { + cacheBuilder.loadedMaterializedStats.map { case (rowCount, sizeInBytes) => + statsOfPlanToCache.copy(sizeInBytes = sizeInBytes, rowCount = Some(rowCount)) + }.getOrElse { // Underlying columnar RDD hasn't been materialized, use the stats from the plan to cache. statsOfPlanToCache - } else { - statsOfPlanToCache.copy( - sizeInBytes = cacheBuilder.materializedSizeInBytes, - rowCount = Some(cacheBuilder.materializedRowCount) - ) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileScanRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileScanRDD.scala index b591573c00afe..ee6bf9c0abe90 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileScanRDD.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileScanRDD.scala @@ -89,6 +89,10 @@ class FileScanRDD( private val ignoreCorruptFiles = options.ignoreCorruptFiles private val ignoreMissingFiles = options.ignoreMissingFiles + + /** Whether this reader fails instead of silently skipping missing or corrupt input files. */ + private[sql] def hasStrictFileReads: Boolean = !ignoreCorruptFiles && !ignoreMissingFiles + // Evaluated on the driver (sparkSession is @transient) and serialized to executors so the // `compute` iterator below can pass it through to ColumnVectorUtils.populate. private val memoryMode: MemoryMode = diff --git a/sql/core/src/main/scala/org/apache/spark/sql/util/PartitionKeyedAccumulator.scala b/sql/core/src/main/scala/org/apache/spark/sql/util/PartitionKeyedAccumulator.scala index bb8f04a8a5565..f56a6f882eda0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/util/PartitionKeyedAccumulator.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/util/PartitionKeyedAccumulator.scala @@ -34,12 +34,9 @@ import org.apache.spark.util.AccumulatorV2 * failed/interrupted tasks are dropped by the accumulator framework (it is not * `countFailedValues`), so only complete per-partition values are ever merged. * - * Backed by a `ConcurrentHashMap`, whose per-entry atomicity is sufficient here: `add` and the - * `putAll` in `merge` are last-write-wins per key, and the reads (`value`, - * `accumulatedNumPartitions`, `foldValues`) only require thread-safety and eventual consistency - * -- they are weakly consistent during concurrent updates but exact once all updates have been - * merged. This avoids any explicit locking (and the nested-lock pattern a two-map `merge` would - * otherwise need). + * Backed by a `ConcurrentHashMap`. Mutations and folds are synchronized so a caller can atomically + * verify that every partition completed and read its statistics from the same stable snapshot. + * Framework-facing reads remain safe, weakly consistent views. * * @tparam T the per-partition value type. Must be non-null (`ConcurrentHashMap` forbids nulls). */ @@ -52,23 +49,29 @@ class PartitionKeyedAccumulator[T] extends AccumulatorV2[(Int, T), java.util.Map override def copyAndReset(): PartitionKeyedAccumulator[T] = new PartitionKeyedAccumulator[T] - override def copy(): PartitionKeyedAccumulator[T] = { + override def copy(): PartitionKeyedAccumulator[T] = synchronized { val newAcc = new PartitionKeyedAccumulator[T] newAcc.byPartition.putAll(byPartition) newAcc } - override def reset(): Unit = byPartition.clear() + override def reset(): Unit = synchronized { + byPartition.clear() + } - override def add(v: (Int, T)): Unit = byPartition.put(v._1, v._2) + override def add(v: (Int, T)): Unit = synchronized { + byPartition.put(v._1, v._2) + } - override def merge(other: AccumulatorV2[(Int, T), java.util.Map[Int, T]]): Unit = other match { - case o: PartitionKeyedAccumulator[T] => - // Last-write-wins per partition id: a partition recorded by more than one task replaces - // rather than accumulates, keeping any caller-derived aggregate exact. - byPartition.putAll(o.byPartition) - case _ => throw new UnsupportedOperationException( - s"Cannot merge ${this.getClass.getName} with ${other.getClass.getName}") + override def merge(other: AccumulatorV2[(Int, T), java.util.Map[Int, T]]): Unit = synchronized { + other match { + case o: PartitionKeyedAccumulator[T] => + // Last-write-wins per partition id: a partition recorded by more than one task replaces + // rather than accumulates, keeping any caller-derived aggregate exact. + byPartition.putAll(o.byPartition) + case _ => throw new UnsupportedOperationException( + s"Cannot merge ${this.getClass.getName} with ${other.getClass.getName}") + } } // A read-only VIEW over the live map -- no copy. Only the accumulator framework calls `value` @@ -80,11 +83,27 @@ class PartitionKeyedAccumulator[T] extends AccumulatorV2[(Int, T), java.util.Map /** Number of distinct partitions that have been recorded. */ def accumulatedNumPartitions: Long = byPartition.size().toLong - /** Folds the per-partition values (each partition counted once) into a single aggregate. */ - def foldValues[A](zero: A)(op: (A, T) => A): A = { + private def foldValuesUnsafe[A](zero: A)(op: (A, T) => A): A = { var result = zero val it = byPartition.values().iterator() while (it.hasNext) result = op(result, it.next()) result } + + /** Folds the per-partition values (each partition counted once) into a single aggregate. */ + def foldValues[A](zero: A)(op: (A, T) => A): A = synchronized { + foldValuesUnsafe(zero)(op) + } + + /** Atomically checks that all partitions completed and folds their statistics. */ + def foldValuesIfComplete[A]( + expectedNumPartitions: Int, + zero: A)( + op: (A, T) => A): Option[A] = synchronized { + if (byPartition.size() == expectedNumPartitions) { + Some(foldValuesUnsafe(zero)(op)) + } else { + None + } + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala index f5db5373299ab..ff8c8fb73dc4d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala @@ -22,12 +22,15 @@ import java.io.File import org.apache.spark.sql.catalyst.expressions.{Alias, BloomFilterMightContain, Literal} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPlan} +import org.apache.spark.sql.columnar.CachedBatch import org.apache.spark.sql.execution.{ReusedSubqueryExec, SubqueryExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEPropagateEmptyRelation} +import org.apache.spark.sql.execution.columnar.InMemoryRelation import org.apache.spark.sql.execution.planmerging.MergeSubplans import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{IntegerType, StructType} +import org.apache.spark.storage.StorageLevel class InjectRuntimeFilterSuite extends SharedSparkSession with AdaptiveSparkPlanHelper { @@ -299,6 +302,242 @@ class InjectRuntimeFilterSuite extends SharedSparkSession checkWithAndWithoutFeatureEnabled(query, shouldReplace = false) } + test("SPARK-58272: safely use fully materialized selectively filtered caches") { + val cacheName = "cached_bloom_filter_keys" + val query = s"SELECT * FROM bf1 JOIN $cacheName ON bf1.c1 = $cacheName.c2" + + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true") { + withTempView(cacheName) { + withCache(cacheName) { + spark.range(0, 40, 1, numPartitions = 4) + .where("id = 8") + .selectExpr("CAST(id AS INT) AS c2") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + + def cachedRelation: InMemoryRelation = { + spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case relation: InMemoryRelation => relation + }.get + } + + val before = cachedRelation + assert(!before.statsAvailable) + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) + + val buffers = before.cacheBuilder.cachedColumnBuffers + assert(buffers.getNumPartitions > 1) + spark.sparkContext.runJob( + buffers, + (iterator: Iterator[CachedBatch]) => iterator.size, + Seq(0)) + assert(!before.statsAvailable) + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) + + assert(spark.table(cacheName).count() == 1) + val materialized = cachedRelation + assert(materialized.statsAvailable) + assert(materialized.isOutputRepeatable) + assert(materialized.hasSelectivePredicate) + val cachedSize = materialized.stats.sizeInBytes + assert(cachedSize > 0) + + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> + (cachedSize - 1).toString) { + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) + } + + var expected: Array[Row] = null + withSQLConf(SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "false") { + expected = sql(query).collect() + } + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> + cachedSize.toString) { + val actual = sql(query) + assert(getNumBloomFilters(actual.queryExecution.optimizedPlan) == 1) + checkAnswer(actual, expected) + + val leftOuter = + s"SELECT * FROM bf1 LEFT OUTER JOIN $cacheName ON bf1.c1 = $cacheName.c2" + assert(getNumBloomFilters(sql(leftOuter).queryExecution.optimizedPlan) == 0) + } + } + } + } + } + + test("SPARK-58272: use materialized caches without predicates only with pruning statistics") { + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true", + SQLConf.CBO_ENABLED.key -> "true") { + val factPlan = spark.table("bf1").queryExecution.optimizedPlan + val factKey = factPlan.output.find(_.name == "c1").get + val applicationDistinctCount = + factPlan.stats.attributeStats.get(factKey).flatMap(_.distinctCount).get + assert(applicationDistinctCount > 5 && applicationDistinctCount < 100) + + Seq(("small_cached_bloom_keys", 8L, 13L, true), + ("large_cached_bloom_keys", 0L, 100L, false)).foreach { + case (cacheName, start, end, shouldInject) => + withTempView(cacheName) { + withCache(cacheName) { + spark.range(start, end, 1, numPartitions = 4) + .selectExpr("CAST(id AS INT) AS c2") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + spark.table(cacheName).count() + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(relation.statsAvailable) + assert(!relation.hasSelectivePredicate) + + val query = s"SELECT * FROM bf1 JOIN $cacheName ON bf1.c1 = $cacheName.c2" + val actual = sql(query) + assert(getNumBloomFilters(actual.queryExecution.optimizedPlan) == + (if (shouldInject) 1 else 0)) + + if (shouldInject) { + val projectedFact = spark.table("bf1").selectExpr("c1 AS fact_key") + val projectedKeys = spark.table(cacheName).selectExpr("c2 AS cached_key") + val projectedJoin = projectedFact.join( + projectedKeys, projectedFact("fact_key") === projectedKeys("cached_key")) + assert(getNumBloomFilters(projectedJoin.queryExecution.optimizedPlan) == 1) + + var expected: Array[Row] = null + withSQLConf(SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "false") { + expected = sql(query).collect() + } + checkAnswer(actual, expected) + + withSQLConf(SQLConf.CBO_ENABLED.key -> "false") { + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) + } + } + } + } + } + } + } + + test("SPARK-58272: reject memory-only and non-repeatable materialized caches") { + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true") { + Seq(("memory_only_bloom_keys", false), + ("nondeterministic_bloom_keys", true)).foreach { + case (cacheName, nondeterministic) => + withTempView(cacheName) { + withCache(cacheName) { + val keys = if (nondeterministic) { + spark.range(0, 40, 1, numPartitions = 4) + .selectExpr("CAST(rand() * 20 AS INT) AS c2") + .where("c2 < 10") + .persist(StorageLevel.MEMORY_AND_DISK) + } else { + spark.range(0, 40, 1, numPartitions = 4) + .where("id = 8") + .selectExpr("CAST(id AS INT) AS c2") + .persist(StorageLevel.MEMORY_ONLY) + } + keys.createOrReplaceTempView(cacheName) + keys.count() + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(!relation.statsAvailable) + + val query = s"SELECT * FROM bf1 JOIN $cacheName ON bf1.c1 = $cacheName.c2" + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) + } + } + } + } + } + + test("SPARK-58272: do not trust cached runtime-replaceable encryption") { + val cacheName = "encrypted_cached_bloom_filter_keys" + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true") { + withTempView(cacheName) { + withCache(cacheName) { + spark.range(0, 40, 1, numPartitions = 4) + .where("id = 8") + .selectExpr( + "CAST(id AS INT) AS c2", + "aes_encrypt(CAST(id AS STRING), '0000111122223333') AS encrypted") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + assert(spark.table(cacheName).count() == 1) + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(relation.hasSelectivePredicate) + assert(!relation.statsAvailable) + + val query = s"SELECT * FROM bf1 JOIN $cacheName ON bf1.c1 = $cacheName.c2" + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) + } + } + } + } + + test("SPARK-58272: use materialized selectively filtered file-backed caches") { + val cacheName = "parquet_cached_bloom_filter_keys" + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true", + SQLConf.IGNORE_MISSING_FILES.key -> "false", + SQLConf.IGNORE_CORRUPT_FILES.key -> "false") { + withTempPath { path => + spark.range(0, 40, 1, numPartitions = 4).write.parquet(path.getCanonicalPath) + withTempView(cacheName) { + withCache(cacheName) { + spark.read.parquet(path.getCanonicalPath) + .where("id = 8") + .selectExpr("CAST(id AS INT) AS c2") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + assert(spark.table(cacheName).count() == 1) + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(relation.statsAvailable) + assert(relation.hasSelectivePredicate) + + val query = s"SELECT * FROM bf1 JOIN $cacheName ON bf1.c1 = $cacheName.c2" + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 1) + } + } + } + } + } + + test("SPARK-58272: materialized creation-side threshold exceeds broadcast threshold") { + val conf = new SQLConf + assert(conf.runtimeFilterMaterializedCreationSideThreshold > conf.autoBroadcastJoinThreshold) + } + test("Runtime bloom filter join: simple") { withSQLConf(SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "3000", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "2000") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/RuntimeBloomFilterCachedInputBenchmark.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/RuntimeBloomFilterCachedInputBenchmark.scala new file mode 100644 index 0000000000000..23b8e508d0686 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/RuntimeBloomFilterCachedInputBenchmark.scala @@ -0,0 +1,139 @@ +/* + * 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.benchmark + +import scala.collection.mutable +import scala.concurrent.duration._ + +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions.BloomFilterMightContain +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.metric.SQLShuffleWriteMetricsReporter +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.storage.StorageLevel + +/** + * Measures runtime Bloom filtering when a materialized cache hides a selective predicate. + * + * To run this benchmark: + * {{{ + * build/sbt "sql/Test/runMain + * org.apache.spark.sql.execution.benchmark.RuntimeBloomFilterCachedInputBenchmark" + * }}} + * + * The additional shuffle metric reports only the wide fact-side exchange, independently of + * exchanges used to aggregate results or construct the Bloom filter. + */ +object RuntimeBloomFilterCachedInputBenchmark extends SqlBasedBenchmark { + + private val factRows = 500000L + private val filteringStride = 100L + private val partitions = 4 + + override def getSparkSession: SparkSession = { + SparkSession.builder() + .master("local[4]") + .appName(this.getClass.getCanonicalName) + .config(SQLConf.SHUFFLE_PARTITIONS.key, 4) + .config("spark.ui.enabled", false) + .getOrCreate() + } + + override def runBenchmarkSuite(mainArgs: Array[String]): Unit = { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> partitions.toString, + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "10MB", + SQLConf.IGNORE_CORRUPT_FILES.key -> "false", + SQLConf.IGNORE_MISSING_FILES.key -> "false") { + val fact = spark.range(0, factRows, 1, partitions).selectExpr( + "id AS fact_key", + "sha2(cast(id AS STRING), 256) AS payload") + val dimension = spark.range(0, factRows, 1, partitions) + .filter(s"id % $filteringStride = 0") + .selectExpr("id AS dimension_key") + .persist(StorageLevel.MEMORY_AND_DISK) + + try { + assert(dimension.count() == factRows / filteringStride) + + runBenchmark("Runtime Bloom filter from a materialized selective cache") { + val factShuffleBytes = mutable.Map.empty[Boolean, Long] + val benchmark = new Benchmark( + "Cached input runtime Bloom filter", + factRows, + minNumIters = 2, + warmupTime = 1.second, + minTime = Duration.Zero, + output = output) + + Seq(false, true).foreach { enabled => + val name = if (enabled) { + "Runtime Bloom filter enabled" + } else { + "Runtime Bloom filter disabled" + } + benchmark.addCase(name, numIters = 2) { _ => + withSQLConf(SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> enabled.toString) { + val result = fact.join(dimension, fact("fact_key") === dimension("dimension_key")) + .selectExpr("sum(length(payload)) AS payload_bytes") + val observed = result.collect().head.getLong(0) + assert(observed == factRows / filteringStride * 64) + + val optimizedPlan = result.queryExecution.optimizedPlan + val hasRuntimeBloomFilter = optimizedPlan.exists { node => + node.expressions.exists(_.exists(_.isInstanceOf[BloomFilterMightContain])) + } + assert(hasRuntimeBloomFilter == enabled, + s"Expected runtime Bloom filter enabled=$enabled:\n$optimizedPlan") + + val factExchanges = result.queryExecution.executedPlan.collect { + case exchange: ShuffleExchangeExec + if exchange.output.exists(_.name == "fact_key") => exchange + } + assert(factExchanges.size == 1, + s"Expected one fact-side shuffle exchange, found ${factExchanges.size}") + factShuffleBytes(enabled) = + factExchanges.head.metrics(SQLShuffleWriteMetricsReporter.SHUFFLE_BYTES_WRITTEN) + .value + } + } + } + + benchmark.run() + + val baseline = factShuffleBytes(false) + val optimized = factShuffleBytes(true) + assert(baseline > 0 && optimized < baseline, + s"Expected the Bloom filter to reduce fact shuffle bytes: $baseline -> $optimized") + val reduction = (baseline - optimized) * 100.0 / baseline + // scalastyle:off println + benchmark.out.println(s"Fact-side shuffle bytes without runtime Bloom filter: $baseline") + benchmark.out.println(s"Fact-side shuffle bytes with runtime Bloom filter: $optimized") + benchmark.out.println(f"Fact-side shuffle reduction: $reduction%.2f%%") + // scalastyle:on println + } + } finally { + dimension.unpersist(blocking = true) + } + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala index fcf7edfcf87b0..a6c54cecb75d8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning import org.apache.spark.sql.classic.DataFrame import org.apache.spark.sql.columnar.CachedBatch -import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} +import org.apache.spark.sql.execution.{FileSourceScanExec, FilterExec, InputAdapter, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf @@ -649,4 +649,126 @@ class InMemoryColumnarQuerySuite extends SharedSparkSession with AdaptiveSparkPl assert(exceptionCnt.get == 0) } + + test("SPARK-58272: only fully materialized repeatable disk-backed caches publish exact stats") { + def checkCache(level: StorageLevel, expected: Boolean): Unit = { + val cached = spark.range(0, 20, 1, numPartitions = 2) + .filter($"id" < 10) + .persist(level) + try { + val relation = cached.queryExecution.withCachedData.collectFirst { + case plan: InMemoryRelation => plan + }.get + + assert(relation.hasSelectivePredicate) + assert(!relation.isOutputRepeatable) + assert(!relation.statsAvailable) + relation.cacheBuilder.cachedColumnBuffers.count() + + assert(relation.isOutputRepeatable) + assert(relation.statsAvailable == expected) + assert(relation.computeStats().rowCount.contains(10L)) + } finally { + cached.unpersist(blocking = true) + } + } + + checkCache(MEMORY_AND_DISK, expected = true) + checkCache(MEMORY_ONLY, expected = false) + } + + test("SPARK-58272: materialized caches require trusted strict file reads") { + withTempPath { path => + spark.range(10).write.parquet(path.getCanonicalPath) + + def checkCache(data: DataFrame, expected: Boolean): Unit = { + val cached = data.filter($"id" < 5).persist(MEMORY_AND_DISK) + try { + val relation = cached.queryExecution.withCachedData.collectFirst { + case plan: InMemoryRelation => plan + }.get + relation.cacheBuilder.cachedColumnBuffers.count() + assert(relation.isOutputRepeatable == expected) + assert(relation.statsAvailable == expected) + assert(relation.hasSelectivePredicate) + } finally { + cached.unpersist(blocking = true) + } + } + + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.IGNORE_MISSING_FILES.key -> "false", + SQLConf.IGNORE_CORRUPT_FILES.key -> "false") { + checkCache(spark.read.parquet(path.getCanonicalPath), expected = true) + + val preinitialized = spark.read.parquet(path.getCanonicalPath) + .filter($"id" < 5) + .persist(MEMORY_AND_DISK) + try { + val relation = preinitialized.queryExecution.withCachedData.collectFirst { + case plan: InMemoryRelation => plan + }.get + val fileScan = relation.cacheBuilder.cachedPlan.collectFirst { + case scan: FileSourceScanExec => scan + }.get + withSQLConf(SQLConf.IGNORE_MISSING_FILES.key -> "true") { + fileScan.inputRDD + } + + relation.cacheBuilder.cachedColumnBuffers.count() + assert(!relation.isOutputRepeatable) + assert(!relation.statsAvailable) + } finally { + preinitialized.unpersist(blocking = true) + } + + checkCache( + spark.read.option("ignoreMissingFiles", "true").parquet(path.getCanonicalPath), + expected = false) + checkCache( + spark.read.option("ignoreCorruptFiles", "true").parquet(path.getCanonicalPath), + expected = false) + } + + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.IGNORE_MISSING_FILES.key -> "true", + SQLConf.IGNORE_CORRUPT_FILES.key -> "false") { + checkCache(spark.read.parquet(path.getCanonicalPath), expected = false) + } + } + } + + test("SPARK-58272: unsafe cached lineages cannot supply runtime-filter statistics") { + val nondeterministic = spark.range(10).filter(rand() > 0.5).persist(MEMORY_AND_DISK) + val isSmall = udf((value: Long) => value < 5) + val userDefined = spark.range(10) + .filter(isSmall($"id")) + .persist(MEMORY_AND_DISK) + val randomizedEncryption = spark.range(10) + .selectExpr( + "id", + "aes_encrypt(CAST(id AS STRING), '0000111122223333') AS encrypted") + .filter($"id" < 5) + .persist(MEMORY_AND_DISK) + val currentTime = spark.range(10) + .selectExpr("id", "current_timestamp() AS observed_at") + .filter($"id" < 5) + .persist(MEMORY_AND_DISK) + + Seq(nondeterministic, userDefined, randomizedEncryption, currentTime).foreach { cached => + try { + val relation = cached.queryExecution.withCachedData.collectFirst { + case plan: InMemoryRelation => plan + }.get + relation.cacheBuilder.cachedColumnBuffers.count() + assert(!relation.isOutputRepeatable) + assert(!relation.statsAvailable) + } finally { + cached.unpersist(blocking = true) + } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/util/PartitionKeyedAccumulatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/util/PartitionKeyedAccumulatorSuite.scala index 19e499942e310..17aa7de7f86d5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/util/PartitionKeyedAccumulatorSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/util/PartitionKeyedAccumulatorSuite.scala @@ -104,4 +104,26 @@ class PartitionKeyedAccumulatorSuite extends SparkFunSuite { assert(acc.accumulatedNumPartitions == 2) assert(acc.foldValues("")((s, v) => s + v).length == 2) // "c" + "b" (each partition once) } + + test("SPARK-58272: fold returns an atomic snapshot only after every partition completes") { + val accumulator = new PartitionKeyedAccumulator[Stats] + accumulator.add((0, (10L, 100L))) + + assert(accumulator.foldValuesIfComplete(2, (0L, 0L)) { + case ((rows, bytes), (partitionRows, partitionBytes)) => + (rows + partitionRows, bytes + partitionBytes) + }.isEmpty) + + accumulator.add((1, (5L, 50L))) + assert(accumulator.foldValuesIfComplete(2, (0L, 0L)) { + case ((rows, bytes), (partitionRows, partitionBytes)) => + (rows + partitionRows, bytes + partitionBytes) + }.contains((15L, 150L))) + + accumulator.add((1, (7L, 70L))) + assert(accumulator.foldValuesIfComplete(2, (0L, 0L)) { + case ((rows, bytes), (partitionRows, partitionBytes)) => + (rows + partitionRows, bytes + partitionBytes) + }.contains((17L, 170L))) + } } From 2bbf825e36cf04599828ad053db6b5f48ca45a50 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 22 Jul 2026 18:05:48 -0700 Subject: [PATCH 02/10] [SPARK-58272][SQL] Preserve runtime Bloom filters above cached leaves --- .../optimizer/InjectRuntimeFilter.scala | 6 +++ .../spark/sql/InjectRuntimeFilterSuite.scala | 45 +++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala index 0e16e9ff1f8ba..506615923ccbd 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala @@ -198,6 +198,12 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J hasSelectivePredicate = hasHitSelectiveFilter || leaf.hasSelectivePredicate, materializedRowCount = Some(rowCount)) } + } else if (hasHitSelectiveFilter) { + Some(FilterCreationSide( + targetKey, + currentPlan, + useMaterializedThreshold = false, + hasSelectivePredicate = true)) } else { None } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala index ff8c8fb73dc4d..e9dccc4fa4e71 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala @@ -372,6 +372,51 @@ class InjectRuntimeFilterSuite extends SharedSparkSession } } + test("SPARK-58272: preserve selective runtime filters over memory-only and cold caches") { + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "1MB", + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true") { + Seq( + ("memory_only_filtered_bloom_keys", StorageLevel.MEMORY_ONLY, true), + ("cold_filtered_bloom_keys", StorageLevel.MEMORY_AND_DISK, false)).foreach { + case (cacheName, storageLevel, materialize) => + withTempView(cacheName) { + withCache(cacheName) { + val cached = spark.range(0, 40, 1, numPartitions = 4) + .selectExpr("CAST(id AS INT) AS c2") + .persist(storageLevel) + cached.createOrReplaceTempView(cacheName) + + if (materialize) { + assert(cached.count() == 40) + } + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cachedRelation: InMemoryRelation => cachedRelation + }.get + assert(!relation.statsAvailable) + assert(relation.isOutputRepeatable == materialize) + + val query = s"SELECT * FROM bf1 JOIN " + + s"(SELECT * FROM $cacheName WHERE c2 = 8) filtered_keys " + + "ON bf1.c1 = filtered_keys.c2" + val actual = sql(query) + assert(getNumBloomFilters(actual.queryExecution.optimizedPlan) == 1) + + var expected: Array[Row] = null + withSQLConf(SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "false") { + expected = sql(query).collect() + } + checkAnswer(actual, expected) + } + } + } + } + } + test("SPARK-58272: use materialized caches without predicates only with pruning statistics") { withSQLConf( SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", From 3897fcc22d737365bc99699e675dba4b0f4288b8 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 22 Jul 2026 20:41:54 -0700 Subject: [PATCH 03/10] [SPARK-58272][SQL] Skip cached Bloom filters for nonbinary collations --- .../optimizer/InjectRuntimeFilter.scala | 12 +++++-- .../spark/sql/InjectRuntimeFilterSuite.scala | 35 +++++++++++++++++++ 2 files changed, 44 insertions(+), 3 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala index 506615923ccbd..abe728048a670 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala @@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreePattern.{INVOKE, JSON_TO_STRUCT, LIKE_FAMLIY, PYTHON_UDF, REGEXP_EXTRACT_FAMILY, REGEXP_REPLACE, SCALA_UDF} +import org.apache.spark.sql.catalyst.util.UnsafeRowUtils import org.apache.spark.sql.internal.SQLConf /** @@ -99,7 +100,8 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J */ private def extractSelectiveFilterOverScan( plan: LogicalPlan, - filterCreationSideKey: Expression): Option[FilterCreationSide] = { + filterCreationSideKey: Expression, + allowMaterializedCache: Boolean): Option[FilterCreationSide] = { def extract( p: LogicalPlan, predicateReference: AttributeSet, @@ -189,7 +191,8 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J findExpressionAndTrackLineageDown(targetKey, currentPlan).exists { case (trackedKey, _) => isSimpleExpression(trackedKey) } - if (leaf.statsAvailable && leaf.isOutputRepeatable && safeLineage) { + if (allowMaterializedCache && leaf.statsAvailable && leaf.isOutputRepeatable && + safeLineage) { leaf.stats.rowCount.map { rowCount => FilterCreationSide( targetKey, @@ -285,7 +288,10 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J if (findExpressionAndTrackLineageDown( filterApplicationSideKey, filterApplicationSide).isDefined && satisfyByteSizeRequirement(filterApplicationSide)) { - extractSelectiveFilterOverScan(filterCreationSide, filterCreationSideKey).filter { + val allowMaterializedCache = UnsafeRowUtils.isBinaryStable(filterCreationSideKey.dataType) && + UnsafeRowUtils.isBinaryStable(filterApplicationSideKey.dataType) + extractSelectiveFilterOverScan( + filterCreationSide, filterCreationSideKey, allowMaterializedCache).filter { creationSide => creationSide.hasSelectivePredicate || { def distinctCount(key: Expression, plan: LogicalPlan): Option[BigInt] = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala index e9dccc4fa4e71..23961143ae372 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala @@ -417,6 +417,41 @@ class InjectRuntimeFilterSuite extends SharedSparkSession } } + test("SPARK-58272: reject materialized runtime filters for nonbinary collated join keys") { + val cacheName = "collated_cached_bloom_filter_keys" + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> "1MB", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true") { + withTempView(cacheName) { + withCache(cacheName) { + spark.range(0, 40, 1, numPartitions = 4) + .where("id = 1") + .selectExpr("IF(id = 1, 'a', 'z') COLLATE UTF8_LCASE AS k") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + assert(spark.table(cacheName).count() == 1) + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(relation.statsAvailable) + assert(relation.hasSelectivePredicate) + + val actual = sql( + s"SELECT fact.k FROM " + + "(SELECT (CASE WHEN a1 = 73 THEN 'A' ELSE 'X' END) " + + "COLLATE UTF8_LCASE AS k FROM bf1) fact " + + s"JOIN $cacheName ON fact.k = $cacheName.k COLLATE UTF8_LCASE") + assert(getNumBloomFilters(actual.queryExecution.optimizedPlan) == 0) + checkAnswer(actual, Row("A")) + } + } + } + } + test("SPARK-58272: use materialized caches without predicates only with pruning statistics") { withSQLConf( SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", From 5a822f72c2b2d1a0c2020b7e8942f3a10712fa79 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 22 Jul 2026 20:55:32 -0700 Subject: [PATCH 04/10] [SPARK-58272][SQL] Preserve existing selective runtime Bloom filters --- .../spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala | 4 +++- .../org/apache/spark/sql/InjectRuntimeFilterSuite.scala | 7 ++++--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala index abe728048a670..d49515da15163 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala @@ -192,7 +192,9 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J case (trackedKey, _) => isSimpleExpression(trackedKey) } if (allowMaterializedCache && leaf.statsAvailable && leaf.isOutputRepeatable && - safeLineage) { + safeLineage && (!hasHitSelectiveFilter || + currentPlan.stats.sizeInBytes <= + conf.runtimeFilterMaterializedCreationSideThreshold)) { leaf.stats.rowCount.map { rowCount => FilterCreationSide( targetKey, diff --git a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala index 23961143ae372..eb3767172b64f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala @@ -372,7 +372,7 @@ class InjectRuntimeFilterSuite extends SharedSparkSession } } - test("SPARK-58272: preserve selective runtime filters over memory-only and cold caches") { + test("SPARK-58272: preserve selective runtime filters over all cache states") { withSQLConf( SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", @@ -381,7 +381,8 @@ class InjectRuntimeFilterSuite extends SharedSparkSession SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true") { Seq( ("memory_only_filtered_bloom_keys", StorageLevel.MEMORY_ONLY, true), - ("cold_filtered_bloom_keys", StorageLevel.MEMORY_AND_DISK, false)).foreach { + ("cold_filtered_bloom_keys", StorageLevel.MEMORY_AND_DISK, false), + ("materialized_filtered_bloom_keys", StorageLevel.MEMORY_AND_DISK, true)).foreach { case (cacheName, storageLevel, materialize) => withTempView(cacheName) { withCache(cacheName) { @@ -397,7 +398,7 @@ class InjectRuntimeFilterSuite extends SharedSparkSession val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { case cachedRelation: InMemoryRelation => cachedRelation }.get - assert(!relation.statsAvailable) + assert(relation.statsAvailable == (storageLevel.useDisk && materialize)) assert(relation.isOutputRepeatable == materialize) val query = s"SELECT * FROM bf1 JOIN " + From fe56a0f2097e55065411a4897407d22c2b2baca5 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 22 Jul 2026 21:13:56 -0700 Subject: [PATCH 05/10] [SPARK-58272][SQL] Require attribute statistics for cached Bloom pruning --- .../spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala | 4 ++-- .../scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala | 4 ++++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala index d49515da15163..4c163687704ab 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala @@ -297,8 +297,8 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J creationSide => creationSide.hasSelectivePredicate || { def distinctCount(key: Expression, plan: LogicalPlan): Option[BigInt] = { - key.references.toSeq match { - case Seq(attribute) => plan.stats.attributeStats.get(attribute) + key match { + case attribute: Attribute => plan.stats.attributeStats.get(attribute) .flatMap(_.distinctCount) case _ => None } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala index eb3767172b64f..c70f7c2ce65c3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala @@ -489,6 +489,10 @@ class InjectRuntimeFilterSuite extends SharedSparkSession (if (shouldInject) 1 else 0)) if (shouldInject) { + val derivedJoin = sql(s"SELECT * FROM bf1 JOIN $cacheName " + + s"ON bf1.c1 % 2 = $cacheName.c2 % 2") + assert(getNumBloomFilters(derivedJoin.queryExecution.optimizedPlan) == 0) + val projectedFact = spark.table("bf1").selectExpr("c1 AS fact_key") val projectedKeys = spark.table(cacheName).selectExpr("c2 AS cached_key") val projectedJoin = projectedFact.join( From 4568523554c490069127287632157e446e9be6af Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 22 Jul 2026 22:02:02 -0700 Subject: [PATCH 06/10] SPARK-58272: Preserve selective runtime filter choices --- .../optimizer/InjectRuntimeFilter.scala | 98 ++++++---- .../spark/sql/InjectRuntimeFilterSuite.scala | 169 +++++++++++++++++- 2 files changed, 233 insertions(+), 34 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala index 4c163687704ab..c51a7850aa08f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala @@ -40,9 +40,7 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J private case class FilterCreationSide( key: Expression, plan: LogicalPlan, - useMaterializedThreshold: Boolean, - hasSelectivePredicate: Boolean, - materializedRowCount: Option[BigInt] = None) + useMaterializedThreshold: Boolean) private def injectFilter( filterApplicationSideKey: Expression, @@ -101,7 +99,8 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J private def extractSelectiveFilterOverScan( plan: LogicalPlan, filterCreationSideKey: Expression, - allowMaterializedCache: Boolean): Option[FilterCreationSide] = { + allowMaterializedCache: Boolean, + applicationDistinctCount: => Option[BigInt]): Option[FilterCreationSide] = { def extract( p: LogicalPlan, predicateReference: AttributeSet, @@ -192,23 +191,22 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J case (trackedKey, _) => isSimpleExpression(trackedKey) } if (allowMaterializedCache && leaf.statsAvailable && leaf.isOutputRepeatable && - safeLineage && (!hasHitSelectiveFilter || - currentPlan.stats.sizeInBytes <= - conf.runtimeFilterMaterializedCreationSideThreshold)) { - leaf.stats.rowCount.map { rowCount => + safeLineage && currentPlan.stats.sizeInBytes <= + conf.runtimeFilterMaterializedCreationSideThreshold) { + leaf.stats.rowCount.filter { rowCount => + hasHitSelectiveFilter || leaf.hasSelectivePredicate || + applicationDistinctCount.exists(_ > rowCount) + }.map { rowCount => FilterCreationSide( targetKey, currentPlan, - useMaterializedThreshold = true, - hasSelectivePredicate = hasHitSelectiveFilter || leaf.hasSelectivePredicate, - materializedRowCount = Some(rowCount)) + useMaterializedThreshold = true) } } else if (hasHitSelectiveFilter) { Some(FilterCreationSide( targetKey, currentPlan, - useMaterializedThreshold = false, - hasSelectivePredicate = true)) + useMaterializedThreshold = false)) } else { None } @@ -216,8 +214,7 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J Some(FilterCreationSide( targetKey, currentPlan, - useMaterializedThreshold = false, - hasSelectivePredicate = true)) + useMaterializedThreshold = false)) case _ => None } @@ -292,25 +289,60 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J satisfyByteSizeRequirement(filterApplicationSide)) { val allowMaterializedCache = UnsafeRowUtils.isBinaryStable(filterCreationSideKey.dataType) && UnsafeRowUtils.isBinaryStable(filterApplicationSideKey.dataType) - extractSelectiveFilterOverScan( - filterCreationSide, filterCreationSideKey, allowMaterializedCache).filter { - creationSide => - creationSide.hasSelectivePredicate || { - def distinctCount(key: Expression, plan: LogicalPlan): Option[BigInt] = { - key match { - case attribute: Attribute => plan.stats.attributeStats.get(attribute) - .flatMap(_.distinctCount) - case _ => None - } - } - val applicationDistinctCount = findExpressionAndTrackLineageDown( - filterApplicationSideKey, filterApplicationSide).flatMap { - case (trackedKey, origin) => distinctCount(trackedKey, origin) - }.orElse(distinctCount(filterApplicationSideKey, filterApplicationSide)) - creationSide.materializedRowCount.exists { rowCount => - applicationDistinctCount.exists(_ > rowCount) - } + def distinctCount(key: Expression, plan: LogicalPlan): Option[BigInt] = key match { + case attribute: Attribute => + plan.stats.attributeStats.get(attribute).flatMap(_.distinctCount) + case _ => None + } + def hasOnlyJoinKeyNullChecksOverScan( + plan: LogicalPlan, + targetKey: Expression): Boolean = plan match { + case project: Project => + hasOnlyJoinKeyNullChecksOverScan( + project.child, replaceAlias(targetKey, getAliasMap(project))) + case Filter(condition, child) => + splitConjunctivePredicates(condition).forall { + case IsNotNull(expression) => expression.semanticEquals(targetKey) + case _ => false + } && hasOnlyJoinKeyNullChecksOverScan(child, targetKey) + case _: LeafNode => true + case _ => false + } + lazy val currentDistinctCount = + distinctCount(filterApplicationSideKey, filterApplicationSide) + lazy val lineageDistinctCount = findExpressionAndTrackLineageDown( + filterApplicationSideKey, filterApplicationSide).flatMap { + case (trackedKey, origin) => distinctCount(trackedKey, origin) + } + lazy val applicationDistinctCount = { + if (hasOnlyJoinKeyNullChecksOverScan( + filterApplicationSide, filterApplicationSideKey)) { + lineageDistinctCount.orElse(currentDistinctCount) + } else { + currentDistinctCount + } + } + if (allowMaterializedCache) { + val selectiveCreationSide = extractSelectiveFilterOverScan( + filterCreationSide, + filterCreationSideKey, + allowMaterializedCache = false, + applicationDistinctCount = None) + selectiveCreationSide + .filter(_.plan.stats.sizeInBytes <= conf.runtimeFilterCreationSideThreshold) + .orElse { + extractSelectiveFilterOverScan( + filterCreationSide, + filterCreationSideKey, + allowMaterializedCache = true, + applicationDistinctCount = applicationDistinctCount) } + } else { + extractSelectiveFilterOverScan( + filterCreationSide, + filterCreationSideKey, + allowMaterializedCache = false, + applicationDistinctCount = None) } } else { None diff --git a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala index c70f7c2ce65c3..2fafc2069edb1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql import java.io.File -import org.apache.spark.sql.catalyst.expressions.{Alias, BloomFilterMightContain, Literal} +import org.apache.spark.sql.catalyst.expressions.{Alias, BloomFilterMightContain, Literal, ScalarSubquery} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPlan} import org.apache.spark.sql.columnar.CachedBatch @@ -515,6 +515,173 @@ class InjectRuntimeFilterSuite extends SharedSparkSession } } + test("SPARK-58272: use filtered application-side distinct counts") { + val cacheName = "filtered_application_bloom_keys" + val nullFilteredCacheName = "null_filtered_application_bloom_keys" + + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "1MB", + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> "1MB", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true", + SQLConf.CBO_ENABLED.key -> "true") { + withTempView(cacheName, nullFilteredCacheName) { + withCache(cacheName, nullFilteredCacheName) { + spark.range(8, 13, 1, numPartitions = 4) + .selectExpr("CAST(id AS INT) AS c2") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + assert(spark.table(cacheName).count() == 5) + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(relation.statsAvailable) + assert(!relation.hasSelectivePredicate) + + val filteredFact = spark.table("bf1").where("a1 = 77") + val factPlan = filteredFact.queryExecution.optimizedPlan + val factKey = factPlan.output.find(_.name == "c1").get + assert(factPlan.stats.attributeStats.get(factKey) + .flatMap(_.distinctCount).contains(BigInt(1))) + + Seq( + s"SELECT * FROM (SELECT * FROM bf1 WHERE a1 = 77) fact " + + s"JOIN $cacheName keys ON fact.c1 = keys.c2", + s"SELECT * FROM $cacheName keys " + + "JOIN (SELECT * FROM bf1 WHERE a1 = 77) fact ON keys.c2 = fact.c1" + ).foreach { query => + var expected: Array[Row] = null + withSQLConf(SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "false") { + expected = sql(query).collect() + } + + val actual = sql(query) + val plan = actual.queryExecution.optimizedPlan + assert(getNumBloomFilters(plan) == 1) + val applicationSides = plan.collect { + case Filter(condition, child) + if condition.exists(_.isInstanceOf[BloomFilterMightContain]) => child + } + assert(applicationSides.size == 1) + assert(applicationSides.head.exists(_.isInstanceOf[InMemoryRelation]), plan.treeString) + checkAnswer(actual, expected) + } + + val windowedQuery = + "SELECT * FROM (" + + "SELECT *, ROW_NUMBER() OVER (PARTITION BY c1 ORDER BY a1) AS rn " + + "FROM bf1 WHERE a1 = 77) fact " + + s"JOIN $cacheName keys ON fact.c1 = keys.c2" + assert(getNumBloomFilters(sql(windowedQuery).queryExecution.optimizedPlan) == 0) + + spark.range(0, 6, 1, numPartitions = 4) + .selectExpr( + "CAST(CASE id WHEN 0 THEN 0 WHEN 1 THEN 8 WHEN 2 THEN 23 " + + "WHEN 3 THEN 58 WHEN 4 THEN 74 ELSE 86 END AS INT) AS c2") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(nullFilteredCacheName) + assert(spark.table(nullFilteredCacheName).count() == 6) + + val nullFilteredFact = spark.table("bf1").where("a1 IS NOT NULL") + val nullFilteredPlan = nullFilteredFact.queryExecution.optimizedPlan + val nullFilteredKey = nullFilteredPlan.output.find(_.name == "c1").get + assert(nullFilteredPlan.stats.attributeStats.get(nullFilteredKey) + .flatMap(_.distinctCount).contains(BigInt(6))) + + val query = s"SELECT * FROM (SELECT * FROM bf1 WHERE a1 IS NOT NULL) fact " + + s"JOIN $nullFilteredCacheName keys ON fact.c1 = keys.c2" + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) + } + } + } + } + + test("SPARK-58272: preserve transitive selective filters past materialized caches") { + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "1MB", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true", + SQLConf.CBO_ENABLED.key -> "true") { + val factPlan = spark.table("bf3").queryExecution.optimizedPlan + val factKey = factPlan.output.find(_.name == "c3").get + val applicationDistinctCount = + factPlan.stats.attributeStats.get(factKey).flatMap(_.distinctCount).get + assert(applicationDistinctCount > 5 && applicationDistinctCount < 100) + + Seq( + ("large_transitive_bloom_keys", 0L, 100L, "1MB", false), + ("oversized_transitive_bloom_keys", 8L, 13L, "0", false), + ("selective_transitive_bloom_keys", 0L, 1000L, "1MB", true) + ).foreach { + case (cacheName, start, end, materializedThreshold, hasHiddenSelectivePredicate) => + withTempView(cacheName) { + withCache(cacheName) { + val keys = spark.range(start, end, 1, numPartitions = 4) + val cachedKeys = if (hasHiddenSelectivePredicate) { + keys.where("id < 100") + } else { + keys + } + cachedKeys + .selectExpr("CAST(id AS INT) AS c2") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + val expectedCount = if (hasHiddenSelectivePredicate) 100L else end - start + assert(spark.table(cacheName).count() == expectedCount) + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(relation.statsAvailable) + assert(relation.hasSelectivePredicate == hasHiddenSelectivePredicate) + + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> + materializedThreshold) { + Seq( + s"SELECT /*+ BROADCAST(selective) */ keys.c2 FROM $cacheName keys " + + "JOIN (SELECT c2 FROM bf2 WHERE a2 = 5) selective " + + "ON keys.c2 = selective.c2", + s"SELECT /*+ BROADCAST(selective) */ keys.c2 " + + "FROM (SELECT c2 FROM bf2 WHERE a2 = 5) selective " + + s"JOIN $cacheName keys ON selective.c2 = keys.c2" + ).foreach { creationSide => + val query = s"SELECT * FROM bf3 fact JOIN ($creationSide) creation " + + "ON fact.c3 = creation.c2" + var expected: Array[Row] = null + withSQLConf(SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "false") { + expected = sql(query).collect() + } + + val actual = sql(query) + val plan = actual.queryExecution.optimizedPlan + assert(getNumBloomFilters(plan) == 1) + val applicationSides = plan.collect { + case Filter(condition, child) + if condition.exists(_.isInstanceOf[BloomFilterMightContain]) => child + } + assert(applicationSides.size == 1) + assert(applicationSides.head.output.exists(_.name == "c3"), plan.treeString) + val creationPlans = plan.collect { + case Filter(condition, _) => condition.collect { + case BloomFilterMightContain(subquery: ScalarSubquery, _) => subquery.plan + } + }.flatten + assert(creationPlans.size == 1) + assert(!creationPlans.head.exists(_.isInstanceOf[InMemoryRelation]), + plan.treeString) + checkAnswer(actual, expected) + } + } + } + } + } + } + } + test("SPARK-58272: reject memory-only and non-repeatable materialized caches") { withSQLConf( SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", From 5fd9c141c8faa48b77e09b11dd8a1174d434c8ae Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Thu, 23 Jul 2026 13:25:08 -0700 Subject: [PATCH 07/10] [SPARK-58272][SQL] Fix materialized runtime Bloom filter correctness --- .../optimizer/InjectRuntimeFilter.scala | 14 +-- .../apache/spark/sql/internal/SQLConf.scala | 3 +- .../execution/columnar/InMemoryRelation.scala | 17 ++-- .../spark/sql/InjectRuntimeFilterSuite.scala | 87 +++++++++++++++++++ .../columnar/InMemoryColumnarQuerySuite.scala | 41 ++++++++- 5 files changed, 143 insertions(+), 19 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala index c51a7850aa08f..aeb53f51f7222 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala @@ -40,7 +40,8 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J private case class FilterCreationSide( key: Expression, plan: LogicalPlan, - useMaterializedThreshold: Boolean) + useMaterializedThreshold: Boolean, + materializedRowCount: Option[BigInt] = None) private def injectFilter( filterApplicationSideKey: Expression, @@ -68,7 +69,8 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J if (filterCreationSidePlan.stats.sizeInBytes > creationSideThreshold) { return filterApplicationSidePlan } - val rowCount = filterCreationSidePlan.stats.rowCount + val rowCount = filterCreationSide.materializedRowCount + .orElse(filterCreationSidePlan.stats.rowCount) val bloomFilterAgg = if (rowCount.isDefined && rowCount.get.longValue > 0L) { new BloomFilterAggregate(new XxHash64(Seq(filterCreationSideKey)), rowCount.get.longValue) @@ -194,13 +196,15 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J safeLineage && currentPlan.stats.sizeInBytes <= conf.runtimeFilterMaterializedCreationSideThreshold) { leaf.stats.rowCount.filter { rowCount => - hasHitSelectiveFilter || leaf.hasSelectivePredicate || - applicationDistinctCount.exists(_ > rowCount) + rowCount <= conf.getConf(SQLConf.RUNTIME_BLOOM_FILTER_MAX_NUM_ITEMS) && + (hasHitSelectiveFilter || leaf.hasSelectivePredicate || + applicationDistinctCount.exists(_ > rowCount)) }.map { rowCount => FilterCreationSide( targetKey, currentPlan, - useMaterializedThreshold = true) + useMaterializedThreshold = true, + materializedRowCount = Some(rowCount)) } } else if (hasHitSelectiveFilter) { Some(FilterCreationSide( diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index e072243293326..ae0a46cc3c0d3 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -822,7 +822,8 @@ object SQLConf { .doc("Size threshold of a fully materialized, repeatable bloom filter creation side with " + "accurate statistics. This replaces the general creation-side threshold because scanning " + "materialized output does not recompute its original plan.") - .version("5.0.0") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) .bytesConf(ByteUnit.BYTE) .createWithDefaultString("100MB") diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala index 68dda3b532fd6..56caff226603b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala @@ -273,7 +273,6 @@ case class CachedRDDBuilder( @transient @volatile private var _cachedColumnBuffers: RDD[CachedBatch] = null @volatile private var isCachedRDDRepeatable = false private var hasStrictFileSourceReads = true - private var _materializedStats: Option[(Long, Long)] = None // The cache's materialization bookkeeping: a partition-keyed accumulator storing // (rowCount, sizeInBytes) per partition. AQE creates a separate cache scan stage per reference to @@ -313,7 +312,6 @@ case class CachedRDDBuilder( if (_cachedColumnBuffers != null) { _cachedColumnBuffers.unpersist(blocking) _cachedColumnBuffers = null - _materializedStats = None partitionStats = newPartitionStats() } isCachedRDDRepeatable = false @@ -329,16 +327,11 @@ case class CachedRDDBuilder( if (_cachedColumnBuffers == null) { None } else { - _materializedStats.orElse { - partitionStats.foldValuesIfComplete( - _cachedColumnBuffers.partitions.length, - (0L, 0L)) { - case ((rows, bytes), (partitionRows, partitionBytes)) => - (rows + partitionRows, bytes + partitionBytes) - }.map { stats => - _materializedStats = Some(stats) - stats - } + partitionStats.foldValuesIfComplete( + _cachedColumnBuffers.partitions.length, + (0L, 0L)) { + case ((rows, bytes), (partitionRows, partitionBytes)) => + (rows + partitionRows, bytes + partitionBytes) } } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala index 2fafc2069edb1..e3be750ef77bc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala @@ -372,6 +372,93 @@ class InjectRuntimeFilterSuite extends SharedSparkSession } } + test("SPARK-58272: bound materialized caches by Bloom filter capacity") { + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> "1MB", + SQLConf.RUNTIME_BLOOM_FILTER_MAX_NUM_ITEMS.key -> "4", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true") { + Seq( + ("at_bloom_capacity", 4L, 1), + ("over_bloom_capacity", 5L, 0) + ).foreach { case (cacheName, rowCount, expectedBloomFilters) => + withTempView(cacheName) { + withCache(cacheName) { + spark.range(0, 8, 1, numPartitions = 4) + .where(s"id < $rowCount") + .selectExpr("CAST(id AS INT) AS c2") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + assert(spark.table(cacheName).count() == rowCount) + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(relation.statsAvailable) + assert(relation.hasSelectivePredicate) + assert(relation.stats.rowCount.contains(rowCount)) + + val query = s"SELECT * FROM bf1 JOIN $cacheName ON bf1.c1 = $cacheName.c2" + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == + expectedBloomFilters) + } + } + } + } + } + + test("SPARK-58272: size projected cached runtime filters using exact materialized rows") { + val cacheName = "projected_cached_bloom_filter_keys" + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> "1MB", + SQLConf.RUNTIME_BLOOM_FILTER_EXPECTED_NUM_ITEMS.key -> "1", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true", + SQLConf.CBO_ENABLED.key -> "false") { + withTempView(cacheName) { + withCache(cacheName) { + spark.range(0, 8, 1, numPartitions = 2) + .where("id < 4") + .selectExpr("CAST(id AS INT) AS c2", "CAST(id * 2 AS INT) AS payload") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + assert(spark.table(cacheName).count() == 4) + + val cachedRelation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case relation: InMemoryRelation => relation + }.get + assert(cachedRelation.stats.rowCount.contains(BigInt(4))) + assert(cachedRelation.hasSelectivePredicate) + + val projectedKeys = spark.table(cacheName).selectExpr("c2 + 1 AS projected_key") + assert(projectedKeys.queryExecution.optimizedPlan.stats.rowCount.isEmpty) + + val fact = spark.table("bf1") + val query = fact.join(projectedKeys, fact("c1") === projectedKeys("projected_key")) + val plan = query.queryExecution.optimizedPlan + assert(getNumBloomFilters(plan) == 1) + + val bloomAggregates = plan.collect { + case Filter(condition, _) => condition.collect { + case subquery: ScalarSubquery => subquery.plan.collect { + case Aggregate(_, aggregateExpressions, _, _) => aggregateExpressions.collect { + case Alias(AggregateExpression(aggregate: BloomFilterAggregate, _, _, _, _), _) => + aggregate + } + }.flatten + }.flatten + }.flatten + assert(bloomAggregates.size == 1) + assert(bloomAggregates.head.estimatedNumItemsExpression.eval() == 4L) + } + } + } + } + test("SPARK-58272: preserve selective runtime filters over all cache states") { withSQLConf( SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala index a6c54cecb75d8..c17142aa1a35a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala @@ -35,7 +35,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.test.SQLTestData._ import org.apache.spark.sql.types._ -import org.apache.spark.storage.StorageLevel +import org.apache.spark.storage.{RDDBlockId, StorageLevel} import org.apache.spark.storage.StorageLevel._ class TestCachedBatchSerializer( @@ -677,6 +677,45 @@ class InMemoryColumnarQuerySuite extends SharedSparkSession with AdaptiveSparkPl checkCache(MEMORY_ONLY, expected = false) } + test("SPARK-58272: materialized cache stats follow partition recomputation") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val property = "spark.sql.test.cache.recomputedPartitionHasRows" + System.clearProperty(property) + val liveGate = udf(() => java.lang.Boolean.getBoolean(property)).asNondeterministic() + val cached = spark.range(0, 8, 1, numPartitions = 1) + .filter(liveGate()) + .persist(MEMORY_ONLY) + + try { + assert(cached.count() == 0) + val relation = cached.queryExecution.withCachedData.collectFirst { + case plan: InMemoryRelation => plan + }.get + val builder = relation.cacheBuilder + val blockId = RDDBlockId(builder.cachedColumnBuffers.id, 0) + val blockManager = spark.sparkContext.env.blockManager + + assert(blockManager.getStatus(blockId).nonEmpty) + assert(builder.loadedMaterializedStats.exists(_._1 == 0L)) + + System.setProperty(property, "true") + blockManager.removeBlock(blockId) + assert(blockManager.getStatus(blockId).isEmpty) + assert(cached.count() == 8) + assert(blockManager.getStatus(blockId).nonEmpty) + assert(builder.materializedRowCount == 8L) + assert(builder.loadedMaterializedStats.exists(_._1 == 8L)) + + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + checkAnswer(cached.groupBy("id").count(), (0L until 8L).map(id => Row(id, 1L))) + } + } finally { + cached.unpersist(blocking = true) + System.clearProperty(property) + } + } + } + test("SPARK-58272: materialized caches require trusted strict file reads") { withTempPath { path => spark.range(10).write.parquet(path.getCanonicalPath) From 49a88c6129a5bd5fe18dd3843e8377d9a6f88759 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Mon, 27 Jul 2026 16:46:30 -0700 Subject: [PATCH 08/10] [SPARK-58272][SQL] Refine materialized runtime filter safety --- .../optimizer/InjectRuntimeFilter.scala | 2 +- .../catalyst/plans/logical/LogicalPlan.scala | 14 +++++++-- .../execution/columnar/InMemoryRelation.scala | 11 +++++-- .../columnar/InMemoryColumnarQuerySuite.scala | 29 +++++++++++++++++++ 4 files changed, 50 insertions(+), 6 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala index aeb53f51f7222..bf065eeecb3a8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala @@ -187,7 +187,7 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J } else { None } - case leaf: LeafNodeWithAccurateStats => + case leaf: MaterializedLeafNode => val safeLineage = currentPlan.deterministic && findExpressionAndTrackLineageDown(targetKey, currentPlan).exists { case (trackedKey, _) => isSimpleExpression(trackedKey) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala index f04297dafb509..905b44c3c4b3a 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala @@ -242,9 +242,17 @@ trait LeafNode extends LogicalPlan with LeafLike[LogicalPlan] { throw new SparkUnsupportedOperationException("_LEGACY_ERROR_TEMP_3114") } -/** A materialized leaf whose output can safely be scanned again to build a runtime filter. */ -private[sql] trait LeafNodeWithAccurateStats extends LeafNode { - /** Whether the current materialized output has complete, accurate statistics. */ +/** + * A leaf node that exposes materialization metadata used by + * [[org.apache.spark.sql.catalyst.optimizer.InjectRuntimeFilter]] to determine whether its output + * can be scanned again safely and profitably to build a runtime filter. + */ +private[sql] trait MaterializedLeafNode extends LeafNode { + /** + * Whether the current materialized output has complete, accurate statistics and durable storage + * for another scan. This excludes memory-only storage levels, whose blocks may be discarded + * under memory pressure and recomputed. + */ def statsAvailable: Boolean /** Whether scanning the materialized output again returns the same rows. */ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala index 56caff226603b..ca8352ebfc4ac 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala @@ -315,6 +315,8 @@ case class CachedRDDBuilder( partitionStats = newPartitionStats() } isCachedRDDRepeatable = false + // Read strictness is derived independently for each cache generation. + hasStrictFileSourceReads = true } def isCachedColumnBuffersLoaded: Boolean = loadedMaterializedStats.isDefined @@ -469,12 +471,17 @@ object InMemoryRelation extends PredicateHelper { "org.apache.spark.sql.avro.AvroFileFormat", "org.apache.spark.sql.hive.orc.OrcFileFormat") + private val catalystExpressionPackage = "org.apache.spark.sql.catalyst.expressions." + private def hasSafeExpressions(plan: QueryPlan[_]): Boolean = { + // Treat the Catalyst namespace as the trust boundary for Expression.deterministic's + // repeatability contract. Expressions outside it fail closed; reject AesEncrypt and + // opaque/user-defined expressions explicitly. plan.expressions.forall { expression => !expression.exists { case _: AesEncrypt | _: NonSQLExpression | _: UserDefinedExpression => true case value => !value.deterministic || value.containsPattern(CURRENT_LIKE) || - !value.getClass.getName.startsWith("org.apache.spark.sql.catalyst.expressions.") + !value.getClass.getName.startsWith(catalystExpressionPackage) } } } @@ -634,7 +641,7 @@ case class InMemoryRelation( output: Seq[Attribute], @transient cacheBuilder: CachedRDDBuilder, override val outputOrdering: Seq[SortOrder]) - extends logical.LeafNodeWithAccurateStats with MultiInstanceRelation { + extends logical.MaterializedLeafNode with MultiInstanceRelation { @volatile var statsOfPlanToCache: Statistics = null diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala index c17142aa1a35a..518d6e5d23803 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala @@ -762,6 +762,35 @@ class InMemoryColumnarQuerySuite extends SharedSparkSession with AdaptiveSparkPl preinitialized.unpersist(blocking = true) } + val rebuildable = spark.read.parquet(path.getCanonicalPath) + .filter($"id" < 5) + .persist(MEMORY_AND_DISK) + try { + val relation = rebuildable.queryExecution.withCachedData.collectFirst { + case plan: InMemoryRelation => plan + }.get + val builder = relation.cacheBuilder + val fileScan = builder.cachedPlan.collectFirst { + case scan: FileSourceScanExec => scan + }.get + + // Keep the physical file reader strict while making only this cache generation observe + // best-effort session settings. + fileScan.inputRDD + withSQLConf(SQLConf.IGNORE_MISSING_FILES.key -> "true") { + builder.cachedColumnBuffers.count() + } + assert(!relation.isOutputRepeatable) + assert(!relation.statsAvailable) + + builder.clearCache(blocking = true) + builder.cachedColumnBuffers.count() + assert(relation.isOutputRepeatable) + assert(relation.statsAvailable) + } finally { + rebuildable.unpersist(blocking = true) + } + checkCache( spark.read.option("ignoreMissingFiles", "true").parquet(path.getCanonicalPath), expected = false) From 9fafde74efddcf85586c987f88314b05c96952cc Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Fri, 31 Jul 2026 15:16:51 -0700 Subject: [PATCH 09/10] [SPARK-58272][SQL] Reuse materialized cache metadata snapshots --- .../optimizer/InjectRuntimeFilter.scala | 56 ++++++++++++------- .../catalyst/plans/logical/LogicalPlan.scala | 22 +++++++- .../execution/columnar/InMemoryRelation.scala | 20 +++++-- .../spark/sql/InjectRuntimeFilterSuite.scala | 15 ++++- .../columnar/InMemoryColumnarQuerySuite.scala | 16 ++++++ 5 files changed, 100 insertions(+), 29 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala index bf065eeecb3a8..ded84587cc4e1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala @@ -41,7 +41,8 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J key: Expression, plan: LogicalPlan, useMaterializedThreshold: Boolean, - materializedRowCount: Option[BigInt] = None) + materializedRowCount: Option[BigInt] = None, + materializedSizeInBytes: Option[BigInt] = None) private def injectFilter( filterApplicationSideKey: Expression, @@ -66,7 +67,8 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J conf.runtimeFilterCreationSideThreshold } // Skip if the filter creation side is too big - if (filterCreationSidePlan.stats.sizeInBytes > creationSideThreshold) { + if (filterCreationSide.materializedSizeInBytes + .getOrElse(filterCreationSidePlan.stats.sizeInBytes) > creationSideThreshold) { return filterApplicationSidePlan } val rowCount = filterCreationSide.materializedRowCount @@ -192,28 +194,42 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J findExpressionAndTrackLineageDown(targetKey, currentPlan).exists { case (trackedKey, _) => isSimpleExpression(trackedKey) } - if (allowMaterializedCache && leaf.statsAvailable && leaf.isOutputRepeatable && - safeLineage && currentPlan.stats.sizeInBytes <= - conf.runtimeFilterMaterializedCreationSideThreshold) { - leaf.stats.rowCount.filter { rowCount => - rowCount <= conf.getConf(SQLConf.RUNTIME_BLOOM_FILTER_MAX_NUM_ITEMS) && - (hasHitSelectiveFilter || leaf.hasSelectivePredicate || - applicationDistinctCount.exists(_ > rowCount)) - }.map { rowCount => - FilterCreationSide( - targetKey, - currentPlan, - useMaterializedThreshold = true, - materializedRowCount = Some(rowCount)) + val materializedMetadata = if (allowMaterializedCache && safeLineage && + leaf.mayHaveUsableMaterializedStats) { + leaf.materializedMetadata.filter(_.statsAvailable).flatMap { metadata => + val creationSize = if (currentPlan eq leaf) { + metadata.sizeInBytes + } else { + currentPlan.stats.sizeInBytes + } + Option.when(creationSize <= conf.runtimeFilterMaterializedCreationSideThreshold) { + metadata -> creationSize + } } - } else if (hasHitSelectiveFilter) { - Some(FilterCreationSide( - targetKey, - currentPlan, - useMaterializedThreshold = false)) } else { None } + materializedMetadata match { + case Some((metadata, creationSize)) => + val rowCount = metadata.rowCount + Option.when( + rowCount <= conf.getConf(SQLConf.RUNTIME_BLOOM_FILTER_MAX_NUM_ITEMS) && + (hasHitSelectiveFilter || leaf.hasSelectivePredicate || + applicationDistinctCount.exists(_ > rowCount))) { + FilterCreationSide( + targetKey, + currentPlan, + useMaterializedThreshold = true, + materializedRowCount = Some(rowCount), + materializedSizeInBytes = Some(creationSize)) + } + case None if hasHitSelectiveFilter => + Some(FilterCreationSide( + targetKey, + currentPlan, + useMaterializedThreshold = false)) + case _ => None + } case _: LeafNode if hasHitSelectiveFilter => Some(FilterCreationSide( targetKey, diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala index 905b44c3c4b3a..757da1dccdd9e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala @@ -242,21 +242,39 @@ trait LeafNode extends LogicalPlan with LeafLike[LogicalPlan] { throw new SparkUnsupportedOperationException("_LEGACY_ERROR_TEMP_3114") } +/** + * A single observation of a fully materialized leaf's current cache generation. + */ +private[sql] case class MaterializedLeafMetadata( + rowCount: BigInt, + sizeInBytes: BigInt, + isOutputRepeatable: Boolean, + isDurable: Boolean) { + def statsAvailable: Boolean = isOutputRepeatable && isDurable +} + /** * A leaf node that exposes materialization metadata used by * [[org.apache.spark.sql.catalyst.optimizer.InjectRuntimeFilter]] to determine whether its output * can be scanned again safely and profitably to build a runtime filter. */ private[sql] trait MaterializedLeafNode extends LeafNode { + /** A complete, generation-consistent snapshot, if the leaf is fully materialized. */ + def materializedMetadata: Option[MaterializedLeafMetadata] + + /** Cheap prerequisite for reading potentially usable materialization metadata. */ + def mayHaveUsableMaterializedStats: Boolean + /** * Whether the current materialized output has complete, accurate statistics and durable storage * for another scan. This excludes memory-only storage levels, whose blocks may be discarded * under memory pressure and recomputed. */ - def statsAvailable: Boolean + def statsAvailable: Boolean = + mayHaveUsableMaterializedStats && materializedMetadata.exists(_.statsAvailable) /** Whether scanning the materialized output again returns the same rows. */ - def isOutputRepeatable: Boolean + def isOutputRepeatable: Boolean = materializedMetadata.exists(_.isOutputRepeatable) /** Whether the original plan contains a predicate that is likely to be selective. */ def hasSelectivePredicate: Boolean diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala index ca8352ebfc4ac..8d93b9ff23f69 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala @@ -338,8 +338,14 @@ case class CachedRDDBuilder( } } - private[sql] def repeatableMaterializedStats: Option[(Long, Long)] = synchronized { - if (isCachedPlanRepeatable) loadedMaterializedStats else None + private[sql] def materializedMetadata: Option[logical.MaterializedLeafMetadata] = synchronized { + loadedMaterializedStats.map { case (rowCount, sizeInBytes) => + logical.MaterializedLeafMetadata( + rowCount = rowCount, + sizeInBytes = sizeInBytes, + isOutputRepeatable = isCachedPlanRepeatable, + isDurable = storageLevel.useDisk) + } } // Reported row count / size for the cache's statistics: exact and de-duplicated, folded over the @@ -654,10 +660,14 @@ case class InMemoryRelation( def cachedPlan: SparkPlan = cacheBuilder.cachedPlan - override def statsAvailable: Boolean = - cacheBuilder.storageLevel.useDisk && cacheBuilder.repeatableMaterializedStats.isDefined + override def mayHaveUsableMaterializedStats: Boolean = + cacheBuilder.storageLevel.useDisk && cacheBuilder.isCachedPlanRepeatable + + override def materializedMetadata: Option[logical.MaterializedLeafMetadata] = + cacheBuilder.materializedMetadata - override def isOutputRepeatable: Boolean = cacheBuilder.repeatableMaterializedStats.isDefined + override def isOutputRepeatable: Boolean = + cacheBuilder.isCachedPlanRepeatable && materializedMetadata.exists(_.isOutputRepeatable) override def hasSelectivePredicate: Boolean = cacheBuilder.hasSelectivePredicate diff --git a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala index e3be750ef77bc..beabadd178295 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala @@ -326,6 +326,7 @@ class InjectRuntimeFilterSuite extends SharedSparkSession } val before = cachedRelation + assert(before.materializedMetadata.isEmpty) assert(!before.statsAvailable) assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) @@ -335,11 +336,15 @@ class InjectRuntimeFilterSuite extends SharedSparkSession buffers, (iterator: Iterator[CachedBatch]) => iterator.size, Seq(0)) + assert(before.materializedMetadata.isEmpty) assert(!before.statsAvailable) assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) assert(spark.table(cacheName).count() == 1) val materialized = cachedRelation + val metadata = materialized.materializedMetadata.get + assert(metadata.rowCount == 1L) + assert(metadata.statsAvailable) assert(materialized.statsAvailable) assert(materialized.isOutputRepeatable) assert(materialized.hasSelectivePredicate) @@ -435,11 +440,17 @@ class InjectRuntimeFilterSuite extends SharedSparkSession assert(cachedRelation.hasSelectivePredicate) val projectedKeys = spark.table(cacheName).selectExpr("c2 + 1 AS projected_key") - assert(projectedKeys.queryExecution.optimizedPlan.stats.rowCount.isEmpty) + val projectedPlan = projectedKeys.queryExecution.optimizedPlan + assert(projectedPlan.stats.rowCount.isEmpty) + assert(projectedPlan.stats.sizeInBytes < cachedRelation.stats.sizeInBytes) val fact = spark.table("bf1") val query = fact.join(projectedKeys, fact("c1") === projectedKeys("projected_key")) - val plan = query.queryExecution.optimizedPlan + val plan = withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> + projectedPlan.stats.sizeInBytes.toString) { + query.queryExecution.optimizedPlan + } assert(getNumBloomFilters(plan) == 1) val bloomAggregates = plan.collect { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala index 518d6e5d23803..91b1b388cb8ca 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala @@ -661,10 +661,19 @@ class InMemoryColumnarQuerySuite extends SharedSparkSession with AdaptiveSparkPl }.get assert(relation.hasSelectivePredicate) + assert(!relation.mayHaveUsableMaterializedStats) + assert(relation.materializedMetadata.isEmpty) assert(!relation.isOutputRepeatable) assert(!relation.statsAvailable) relation.cacheBuilder.cachedColumnBuffers.count() + val metadata = relation.materializedMetadata.get + assert(metadata.rowCount == 10L) + assert(metadata.sizeInBytes == relation.computeStats().sizeInBytes) + assert(metadata.isOutputRepeatable) + assert(metadata.isDurable == expected) + assert(metadata.statsAvailable == expected) + assert(relation.mayHaveUsableMaterializedStats == expected) assert(relation.isOutputRepeatable) assert(relation.statsAvailable == expected) assert(relation.computeStats().rowCount.contains(10L)) @@ -697,6 +706,7 @@ class InMemoryColumnarQuerySuite extends SharedSparkSession with AdaptiveSparkPl assert(blockManager.getStatus(blockId).nonEmpty) assert(builder.loadedMaterializedStats.exists(_._1 == 0L)) + assert(relation.materializedMetadata.exists(_.rowCount == 0L)) System.setProperty(property, "true") blockManager.removeBlock(blockId) @@ -705,6 +715,7 @@ class InMemoryColumnarQuerySuite extends SharedSparkSession with AdaptiveSparkPl assert(blockManager.getStatus(blockId).nonEmpty) assert(builder.materializedRowCount == 8L) assert(builder.loadedMaterializedStats.exists(_._1 == 8L)) + assert(relation.materializedMetadata.exists(_.rowCount == 8L)) withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { checkAnswer(cached.groupBy("id").count(), (0L until 8L).map(id => Row(id, 1L))) @@ -780,11 +791,16 @@ class InMemoryColumnarQuerySuite extends SharedSparkSession with AdaptiveSparkPl withSQLConf(SQLConf.IGNORE_MISSING_FILES.key -> "true") { builder.cachedColumnBuffers.count() } + assert(!relation.materializedMetadata.get.isOutputRepeatable) + assert(!relation.mayHaveUsableMaterializedStats) assert(!relation.isOutputRepeatable) assert(!relation.statsAvailable) builder.clearCache(blocking = true) + assert(relation.materializedMetadata.isEmpty) builder.cachedColumnBuffers.count() + assert(relation.materializedMetadata.get.statsAvailable) + assert(relation.mayHaveUsableMaterializedStats) assert(relation.isOutputRepeatable) assert(relation.statsAvailable) } finally { From 06ca37cfa69f9d1da48c43ab321962d854b53a76 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 1 Aug 2026 14:47:37 -0700 Subject: [PATCH 10/10] [SPARK-58272][SQL] Reduce runtime Bloom filter planning overhead --- .../optimizer/InjectRuntimeFilter.scala | 26 ++++-- .../execution/columnar/InMemoryRelation.scala | 93 +++++++++++-------- 2 files changed, 74 insertions(+), 45 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala index ded84587cc4e1..aef521e8d0a7d 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala @@ -104,7 +104,8 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J plan: LogicalPlan, filterCreationSideKey: Expression, allowMaterializedCache: Boolean, - applicationDistinctCount: => Option[BigInt]): Option[FilterCreationSide] = { + applicationDistinctCount: => Option[BigInt], + onMaterializedLeaf: => Unit = ()): Option[FilterCreationSide] = { def extract( p: LogicalPlan, predicateReference: AttributeSet, @@ -190,12 +191,15 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J None } case leaf: MaterializedLeafNode => + onMaterializedLeaf val safeLineage = currentPlan.deterministic && findExpressionAndTrackLineageDown(targetKey, currentPlan).exists { case (trackedKey, _) => isSimpleExpression(trackedKey) } val materializedMetadata = if (allowMaterializedCache && safeLineage && - leaf.mayHaveUsableMaterializedStats) { + leaf.mayHaveUsableMaterializedStats && + (hasHitSelectiveFilter || leaf.hasSelectivePredicate || + applicationDistinctCount.isDefined)) { leaf.materializedMetadata.filter(_.statsAvailable).flatMap { metadata => val creationSize = if (currentPlan eq leaf) { metadata.sizeInBytes @@ -343,19 +347,25 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J } } if (allowMaterializedCache) { + var sawMaterializedLeaf = false val selectiveCreationSide = extractSelectiveFilterOverScan( filterCreationSide, filterCreationSideKey, allowMaterializedCache = false, - applicationDistinctCount = None) + applicationDistinctCount = None, + onMaterializedLeaf = { sawMaterializedLeaf = true }) selectiveCreationSide .filter(_.plan.stats.sizeInBytes <= conf.runtimeFilterCreationSideThreshold) .orElse { - extractSelectiveFilterOverScan( - filterCreationSide, - filterCreationSideKey, - allowMaterializedCache = true, - applicationDistinctCount = applicationDistinctCount) + if (sawMaterializedLeaf) { + extractSelectiveFilterOverScan( + filterCreationSide, + filterCreationSideKey, + allowMaterializedCache = true, + applicationDistinctCount = applicationDistinctCount) + } else { + selectiveCreationSide + } } } else { extractSelectiveFilterOverScan( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala index 8d93b9ff23f69..44197efd0e1ba 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala @@ -319,7 +319,10 @@ case class CachedRDDBuilder( hasStrictFileSourceReads = true } - def isCachedColumnBuffersLoaded: Boolean = loadedMaterializedStats.isDefined + def isCachedColumnBuffersLoaded: Boolean = synchronized { + _cachedColumnBuffers != null && + partitionStats.accumulatedNumPartitions == _cachedColumnBuffers.partitions.length + } private[sql] def isCachedPlanRepeatable: Boolean = isCachedLogicalPlanRepeatable && isCachedRDDRepeatable @@ -492,25 +495,42 @@ object InMemoryRelation extends PredicateHelper { } } - private def hasRepeatableLogicalPlan(analyzedPlan: LogicalPlan, plan: LogicalPlan): Boolean = { + private def hasRepeatableLogicalPlan( + analyzedPlan: LogicalPlan, + plan: LogicalPlan, + onOptimizedNode: LogicalPlan => Unit): Boolean = { // Runtime-replaceable expressions such as AES encryption can become deterministic-looking // StaticInvoke nodes during optimization despite using a fresh random initialization vector. // Inspect the original analyzed expressions before trusting the optimized execution shape. - analyzedPlan.deterministic && analyzedPlan.collectWithSubqueries { - case node if !hasSafeExpressions(node) => true - }.isEmpty && plan.deterministic && plan.collectWithSubqueries { - case node if !hasSafeExpressions(node) => true - case _: logical.Project | _: logical.Filter | _: logical.SubqueryAlias | - _: logical.Range | _: logical.LocalRelation => false - case relation: LogicalRelation => relation.relation match { - case fileRelation: HadoopFsRelation => - val fileFormatClass = fileRelation.fileFormat.getClass - !(trustedFileFormatClasses.contains(fileFormatClass) || - trustedExternalFileFormatNames.contains(fileFormatClass.getName)) - case _ => true + var repeatable = analyzedPlan.deterministic + if (repeatable) { + analyzedPlan.foreachWithSubqueries { node => + if (repeatable && !hasSafeExpressions(node)) { + repeatable = false + } + } + } + if (repeatable) { + repeatable = plan.deterministic + } + plan.foreachWithSubqueries { node => + onOptimizedNode(node) + if (repeatable) { + repeatable = hasSafeExpressions(node) && (node match { + case _: logical.Project | _: logical.Filter | _: logical.SubqueryAlias | + _: logical.Range | _: logical.LocalRelation => true + case relation: LogicalRelation => relation.relation match { + case fileRelation: HadoopFsRelation => + val fileFormatClass = fileRelation.fileFormat.getClass + trustedFileFormatClasses.contains(fileFormatClass) || + trustedExternalFileFormatNames.contains(fileFormatClass.getName) + case _ => false + } + case _ => false + }) } - case _ => true - }.forall(!_) + } + repeatable } private[columnar] def hasRepeatablePhysicalPlan(plan: SparkPlan): Boolean = { @@ -525,24 +545,6 @@ object InMemoryRelation extends PredicateHelper { } } - private def collectFileSourceOptions(plan: LogicalPlan): Seq[Map[String, String]] = { - val relevantOptions = Seq( - FileSourceOptions.IGNORE_MISSING_FILES, - FileSourceOptions.IGNORE_CORRUPT_FILES) - plan.collectWithSubqueries { - case relation: LogicalRelation if relation.relation.isInstanceOf[HadoopFsRelation] => - val options = CaseInsensitiveMap(relation.relation.asInstanceOf[HadoopFsRelation].options) - relevantOptions.flatMap { key => options.get(key).map(key -> _) }.toMap - } - } - - private def hasSelectivePredicate(plan: LogicalPlan): Boolean = { - plan.collectWithSubqueries { - case logical.Filter(condition, _) if condition.deterministic && - isLikelySelective(condition) => true - }.nonEmpty - } - private def newCacheBuilder( serializer: CachedBatchSerializer, storageLevel: StorageLevel, @@ -551,15 +553,32 @@ object InMemoryRelation extends PredicateHelper { logicalPlan: LogicalPlan, analyzedPlan: LogicalPlan, optimizedPlan: LogicalPlan): CachedRDDBuilder = { + val relevantOptions = Seq( + FileSourceOptions.IGNORE_MISSING_FILES, + FileSourceOptions.IGNORE_CORRUPT_FILES) + val fileSourceOptions = Seq.newBuilder[Map[String, String]] + var hasSelectivePredicate = false + val repeatable = hasRepeatableLogicalPlan(analyzedPlan, optimizedPlan, { + case logical.Filter(condition, _) => + if (!hasSelectivePredicate && condition.deterministic && isLikelySelective(condition)) { + hasSelectivePredicate = true + } + case relation: LogicalRelation if relation.relation.isInstanceOf[HadoopFsRelation] => + val options = CaseInsensitiveMap(relation.relation.asInstanceOf[HadoopFsRelation].options) + fileSourceOptions += relevantOptions.flatMap { key => + options.get(key).map(key -> _) + }.toMap + case _ => + }) CachedRDDBuilder( serializer, storageLevel, cachedPlan, tableName, logicalPlan, - isCachedLogicalPlanRepeatable = hasRepeatableLogicalPlan(analyzedPlan, optimizedPlan), - hasSelectivePredicate = hasSelectivePredicate(optimizedPlan), - fileSourceOptions = collectFileSourceOptions(optimizedPlan)) + isCachedLogicalPlanRepeatable = repeatable, + hasSelectivePredicate = hasSelectivePredicate, + fileSourceOptions = fileSourceOptions.result()) } private[this] var ser: Option[CachedBatchSerializer] = None