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..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 @@ -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 /** @@ -36,29 +37,42 @@ 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, + materializedRowCount: Option[BigInt] = None, + materializedSizeInBytes: 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 (filterCreationSide.materializedSizeInBytes + .getOrElse(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) @@ -81,23 +95,24 @@ 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, + allowMaterializedCache: Boolean, + applicationDistinctCount: => Option[BigInt], + onMaterializedLeaf: => Unit = ()): 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 +190,55 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J } else { None } + case leaf: MaterializedLeafNode => + onMaterializedLeaf + val safeLineage = currentPlan.deterministic && + findExpressionAndTrackLineageDown(targetKey, currentPlan).exists { + case (trackedKey, _) => isSimpleExpression(trackedKey) + } + val materializedMetadata = if (allowMaterializedCache && safeLineage && + leaf.mayHaveUsableMaterializedStats && + (hasHitSelectiveFilter || leaf.hasSelectivePredicate || + applicationDistinctCount.isDefined)) { + 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 { + 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((targetKey, currentPlan)) + Some(FilterCreationSide( + targetKey, + currentPlan, + useMaterializedThreshold = false)) case _ => None } @@ -237,18 +299,81 @@ 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) + val allowMaterializedCache = UnsafeRowUtils.isBinaryStable(filterCreationSideKey.dataType) && + UnsafeRowUtils.isBinaryStable(filterApplicationSideKey.dataType) + 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) { + var sawMaterializedLeaf = false + val selectiveCreationSide = extractSelectiveFilterOverScan( + filterCreationSide, + filterCreationSideKey, + allowMaterializedCache = false, + applicationDistinctCount = None, + onMaterializedLeaf = { sawMaterializedLeaf = true }) + selectiveCreationSide + .filter(_.plan.stats.sizeInBytes <= conf.runtimeFilterCreationSideThreshold) + .orElse { + if (sawMaterializedLeaf) { + extractSelectiveFilterOverScan( + filterCreationSide, + filterCreationSideKey, + allowMaterializedCache = true, + applicationDistinctCount = applicationDistinctCount) + } else { + selectiveCreationSide + } + } + } else { + extractSelectiveFilterOverScan( + filterCreationSide, + filterCreationSideKey, + allowMaterializedCache = false, + applicationDistinctCount = None) + } } else { None } @@ -307,9 +432,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 +444,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..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,6 +242,44 @@ 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 = + mayHaveUsableMaterializedStats && materializedMetadata.exists(_.statsAvailable) + + /** Whether scanning the materialized output again returns the same rows. */ + def isOutputRepeatable: Boolean = materializedMetadata.exists(_.isOutputRepeatable) + + /** 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..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 @@ -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,22 @@ 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("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .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 +8283,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..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 @@ -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,14 @@ 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 // 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,32 +312,42 @@ 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 partitionStats = newPartitionStats() } + isCachedRDDRepeatable = false + // Read strictness is derived independently for each cache generation. + hasStrictFileSourceReads = true } def isCachedColumnBuffersLoaded: Boolean = synchronized { - _cachedColumnBuffers != null && isCachedRDDLoaded - } - - 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 + _cachedColumnBuffers != null && + partitionStats.accumulatedNumPartitions == _cachedColumnBuffers.partitions.length + } + + private[sql] def isCachedPlanRepeatable: Boolean = + isCachedLogicalPlanRepeatable && isCachedRDDRepeatable + + /** Reads completeness and exact statistics from one cache generation atomically. */ + private[sql] def loadedMaterializedStats: Option[(Long, Long)] = synchronized { + if (_cachedColumnBuffers == null) { + None + } else { + partitionStats.foldValuesIfComplete( + _cachedColumnBuffers.partitions.length, + (0L, 0L)) { + case ((rows, bytes), (partitionRows, partitionBytes)) => + (rows + partitionRows, bytes + partitionBytes) } - rddLoaded + } + } + + private[sql] def materializedMetadata: Option[logical.MaterializedLeafMetadata] = synchronized { + loadedMaterializedStats.map { case (rowCount, sizeInBytes) => + logical.MaterializedLeafMetadata( + rowCount = rowCount, + sizeInBytes = sizeInBytes, + isOutputRepeatable = isCachedPlanRepeatable, + isDurable = storageLevel.useDisk) } } @@ -347,16 +369,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 +431,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 +459,127 @@ 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 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(catalystExpressionPackage) + } + } + } + + 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. + 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 + }) + } + } + repeatable + } + + 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 newCacheBuilder( + serializer: CachedBatchSerializer, + storageLevel: StorageLevel, + cachedPlan: SparkPlan, + tableName: Option[String], + 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 = repeatable, + hasSelectivePredicate = hasSelectivePredicate, + fileSourceOptions = fileSourceOptions.result()) + } private[this] var ser: Option[CachedBatchSerializer] = None private[this] def getSerializer(sqlConf: SQLConf): CachedBatchSerializer = synchronized { @@ -434,8 +606,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 +622,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 +637,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 +666,7 @@ case class InMemoryRelation( output: Seq[Attribute], @transient cacheBuilder: CachedRDDBuilder, override val outputOrdering: Seq[SortOrder]) - extends logical.LeafNode with MultiInstanceRelation { + extends logical.MaterializedLeafNode with MultiInstanceRelation { @volatile var statsOfPlanToCache: Statistics = null @@ -500,6 +679,17 @@ case class InMemoryRelation( def cachedPlan: SparkPlan = cacheBuilder.cachedPlan + override def mayHaveUsableMaterializedStats: Boolean = + cacheBuilder.storageLevel.useDisk && cacheBuilder.isCachedPlanRepeatable + + override def materializedMetadata: Option[logical.MaterializedLeafMetadata] = + cacheBuilder.materializedMetadata + + override def isOutputRepeatable: Boolean = + cacheBuilder.isCachedPlanRepeatable && materializedMetadata.exists(_.isOutputRepeatable) + + override def hasSelectivePredicate: Boolean = cacheBuilder.hasSelectivePredicate + private[sql] def updateStats( rowCount: Long, newColStats: Map[Attribute, ColumnStat]): Unit = this.synchronized { @@ -511,14 +701,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..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 @@ -19,15 +19,18 @@ 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 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,592 @@ 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.materializedMetadata.isEmpty) + 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.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) + 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: 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") + 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 = 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 { + 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", + 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), + ("materialized_filtered_bloom_keys", StorageLevel.MEMORY_AND_DISK, true)).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 == (storageLevel.useDisk && materialize)) + 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: 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", + 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 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( + 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: 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", + 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..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 @@ -28,14 +28,14 @@ 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 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( @@ -649,4 +649,210 @@ 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.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)) + } finally { + cached.unpersist(blocking = true) + } + } + + checkCache(MEMORY_AND_DISK, expected = true) + 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)) + assert(relation.materializedMetadata.exists(_.rowCount == 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)) + 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))) + } + } 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) + + 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) + } + + 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.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 { + rebuildable.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))) + } }