[SPARK-58272][SQL] Enable runtime Bloom filters for materialized cached inputs

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)) {

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.
  */

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -801,20 +801,32 @@ 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)
 
   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)