pan3793 commented on code in PR #58635:
URL: https://github.com/apache/spark/pull/58635#discussion_r4260741518


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala:
##########
@@ -644,3 +683,214 @@ private[sql] object NullAwareAntiJoinPlanning extends 
JoinSelectionHelper {
     }
   }
 }
+
+/**
+ * Decides whether a plan can serve as the source of a runtime filter on a 
join key. A runtime
+ * filter evaluates its source separately from the join, so the key values the 
source produces
+ * must be the same in both evaluations, or the filter could prune rows the 
join itself matches.
+ * `deterministic` is not enough for that: Spark flags order-dependent 
computations such as
+ * first, last, row_number or an unordered LIMIT as deterministic.
+ *
+ * The plan is walked bottom-up, tracking the output attributes whose values 
are unstable: they
+ * come from a non-deterministic expression, an order-dependent aggregate or 
window function, or
+ * an expression over such an attribute. The plan is rejected outright when 
its row set is
+ * unstable: a filter, join condition or grouping consumes an unstable 
attribute, an inner
+ * generate uses an unstable generator, a sample is unseeded or over anything 
but a scan, or a
+ * limit is over anything but a total order; and, with its own reason, when an 
operator's effect
+ * on the rows is not analyzed. The source qualifies when the key references 
no unstable
+ * attribute. Values that are unstable but only carried to the output (a 
`first(name)` next to a
+ * `GROUP BY id`, a row number next to the key) do not disqualify it.
+ */
+private[optimizer] object RuntimeFilterSourceAnalysis extends AliasHelper {
+
+  /**
+   * @param unstable output attributes whose values depend on evaluation order 
or on chance.
+   * @param totallyOrdered whether the rows are in a total order on stable 
keys, so that a limit
+   *                       over them keeps the same rows every time.
+   */
+  private case class Taint(unstable: AttributeSet, totallyOrdered: Boolean = 
false)
+
+  private val NotRepeatable =
+    "the hinted side may produce different rows or join keys when evaluated 
again"
+
+  /** Why `plan` is not a repeatable source of `key`, or None when it is. */
+  def rejection(plan: LogicalPlan, key: Expression): Option[String] = {
+    if (plan.isStreaming) {
+      Some("the hinted side is a stream")
+    } else if (!key.deterministic) {
+      Some(NotRepeatable)
+    } else {
+      analyze(plan) match {
+        case Left(reason) => Some(reason)
+        case Right(t) if key.references.intersect(t.unstable).nonEmpty => 
Some(NotRepeatable)
+        case _ => None
+      }
+    }
+  }
+
+  /**
+   * Whether `e` yields the same value on every evaluation. A subquery counts 
as deterministic
+   * when its plan is, which is the very check this analysis replaces, so its 
plan is analyzed
+   * too.
+   */
+  private def isStable(e: Expression, unstable: AttributeSet): Boolean = {
+    e.deterministic && e.references.intersect(unstable).isEmpty && !e.exists {
+      case s: SubqueryExpression => s.plan.isStreaming ||
+        analyze(s.plan).forall(t => 
s.plan.outputSet.intersect(t.unstable).nonEmpty)
+      case _ => false
+    }
+  }
+
+  /** Returns the taint of `plan`'s output, or the reason its row set is not 
repeatable. */
+  private def analyze(plan: LogicalPlan): Either[String, Taint] = plan match {
+    case _: LeafNode => Right(Taint(AttributeSet.empty))
+
+    case p: Project => analyze(p.child).map { t =>
+      Taint(
+        AttributeSet(p.projectList.filterNot(isStable(_, 
t.unstable)).map(_.toAttribute)),
+        t.totallyOrdered)
+    }
+
+    case f: Filter => analyze(f.child).flatMap { t =>
+      if (isStable(f.condition, t.unstable)) Right(t) else Left(NotRepeatable)
+    }
+
+    case j: Join => analyze(j.left).flatMap { l =>
+      analyze(j.right).flatMap { r =>
+        val unstable = l.unstable ++ r.unstable
+        if (j.condition.forall(isStable(_, unstable))) {
+          Right(Taint(unstable))
+        } else {
+          Left(NotRepeatable)
+        }
+      }
+    }
+
+    case a: Aggregate => analyze(a.child).map { t =>
+      // Grouping on an unstable value changes which rows form a group, so 
every aggregate result
+      // then depends on it; a grouping expression's own value is as stable as 
its input.
+      val stableGroups = a.groupingExpressions.forall(isStable(_, t.unstable))
+      val unstable = a.aggregateExpressions.filter { e =>
+        !isStable(e, t.unstable) ||
+          (e.exists(_.isInstanceOf[AggregateExpression]) &&
+            (!stableGroups || !isOrderIrrelevantAggregate(e)))
+      }
+      Taint(AttributeSet(unstable.map(_.toAttribute)))
+    }
+
+    case w: Window => analyze(w.child).map { t =>
+      val stablePartitions = w.partitionSpec.forall(isStable(_, t.unstable))
+      val unstable = w.windowExpressions.filter { e =>
+        !stablePartitions || !isStable(e, t.unstable) || 
!isOrderIrrelevantWindow(e)
+      }
+      Taint(t.unstable ++ AttributeSet(unstable.map(_.toAttribute)))
+    }
+
+    case u: Union =>
+      val taints = u.children.map(analyze)
+      taints.collectFirst { case Left(reason) => Left(reason) }.getOrElse {
+        val unstable = u.output.zipWithIndex.collect {
+          case (attr, i) if u.children.zip(taints).exists {
+            case (child, Right(taint)) => 
taint.unstable.contains(child.output(i))
+            case _ => false
+          } => attr
+        }
+        Right(Taint(AttributeSet(unstable)))
+      }
+
+    // An inner generate drops the rows for which the generator yields 
nothing, so an unstable
+    // generator changes the row set; an outer generate keeps them.
+    case g: Generate => analyze(g.child).flatMap { t =>
+      if (isStable(g.generator, t.unstable)) {
+        Right(t)
+      } else if (g.outer) {

Review Comment:
   Fixed in cbef594a838. Grouping on an unstable value now marks the 
multiplicity unstable, as does `Distinct` over an unstable column; tests cover 
the `rnd3` shape and its `DISTINCT` variant.



##########
sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala:
##########
@@ -285,6 +285,10 @@ case class DataSourceV2ScanRelation(
     s"RelationV2$outputString $nameWithTimeTravelSpec"
   }
 
+  // The scan is planned once, and each of its executions reads what that plan 
describes, as the
+  // two sides of a self-join do.
+  override def isOutputRepeatable: Boolean = true

Review Comment:
   Fixed in cbef594a838. `DataSourceV2ScanRelation` has a `repeatableScan` 
field (default false), which `V2ScanRelationPushDown` sets unless the scan 
absorbed a limit, offset, top-N, sample or aggregate. A JDBC test covers the 
Bloom filter and DPP paths.



##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala:
##########
@@ -578,3 +623,286 @@ trait JoinSelectionHelper extends Logging {
       conf.getConfString("spark.sql.join.forceApplyShuffledHashJoin", "false") 
== "true"
   }
 }
+
+/**
+ * Decides whether a plan can serve as the source of a runtime filter on a 
join key. A runtime
+ * filter evaluates its source separately from the join, so the rows the 
source produces and the
+ * key values they carry must be the same in both evaluations, or the filter 
could prune rows the
+ * join itself matches. `deterministic` is not enough for that: Spark flags 
order-dependent
+ * computations such as first, last, row_number or an unordered LIMIT as 
deterministic, and a leaf
+ * such as an RDD that is not checkpointed may return different rows on each 
computation.
+ *
+ * The plan is walked bottom-up, tracking:
+ * - output attributes whose values are unstable: they come from a 
non-deterministic expression, an
+ *   order-dependent aggregate or window function, a subquery whose plan fails 
this same analysis,
+ *   or an expression over such an attribute;
+ * - whether the row multiplicity is unstable, i.e. every row is still present 
but may be repeated
+ *   a different number of times, which an outer generate with an unstable 
generator causes and
+ *   which makes every aggregate over the rows unstable.
+ * The plan is rejected outright when its row set is unstable (a filter, join 
condition or
+ * grouping consumes an unstable attribute, an inner generate uses an unstable 
generator, a sample
+ * is unseeded or over anything but a repeatable leaf, a limit is over 
anything but a total order),
+ * when a leaf cannot promise the same rows again, and, each with its own 
reason, when an operator
+ * would duplicate an observation or is not analyzed. The source qualifies 
when the key references
+ * no unstable attribute. Values that are unstable but only carried to the 
output (a `first(name)`
+ * next to a `GROUP BY id`, a row number next to the key) do not disqualify it.
+ */
+private[optimizer] object RuntimeFilterSourceAnalysis extends AliasHelper {
+
+  /**
+   * @param unstable output attributes whose values depend on evaluation order 
or on chance.
+   * @param unstableMultiplicity whether rows may be repeated a different 
number of times per
+   *                             evaluation, while each row stays present.
+   * @param totallyOrdered whether the rows are in a total order on stable 
keys, so that a limit
+   *                       over them keeps the same rows every time.
+   */
+  private case class Taint(
+      unstable: AttributeSet,
+      unstableMultiplicity: Boolean = false,
+      totallyOrdered: Boolean = false)
+
+  private val NotRepeatable =
+    "the hinted side may produce different rows or join keys when evaluated 
again"
+
+  /** Why `plan` is not a repeatable source of `key`, or None when it is. */
+  def rejection(plan: LogicalPlan, key: Expression): Option[String] = {
+    if (plan.isStreaming) {
+      Some("the hinted side is a stream")
+    } else {
+      analyze(plan) match {
+        case Left(reason) => Some(reason)
+        case Right(t) if !isStable(key, t.unstable) => Some(NotRepeatable)
+        case _ => None
+      }
+    }
+  }
+
+  def isRepeatableExpression(e: Expression): Boolean = isStable(e, 
AttributeSet.empty)
+
+  /**
+   * Whether `e` yields the same value on every evaluation. A subquery counts 
as deterministic
+   * when its plan is, which is the very check this analysis replaces, so its 
plan is analyzed
+   * too.
+   */
+  private def isStable(e: Expression, unstable: AttributeSet): Boolean = {
+    e.deterministic && e.references.intersect(unstable).isEmpty && !e.exists {
+      case s: SubqueryExpression => !isRepeatableSubquery(s)
+      case _ => false
+    }
+  }
+
+  /**
+   * A subquery is repeatable when its row set is, and so is whatever of its 
output the subquery
+   * exposes: the build keys of a DPP filter, nothing but existence for 
EXISTS, the output columns
+   * otherwise.
+   */
+  private def isRepeatableSubquery(s: SubqueryExpression): Boolean = {
+    !s.plan.isStreaming && analyze(s.plan).exists { t =>
+      s match {
+        case d: DynamicPruningSubquery =>
+          d.broadcastKeyIndices.forall(i => isStable(d.buildKeys(i), 
t.unstable))
+        case _: Exists => true
+        case _ => s.plan.outputSet.intersect(t.unstable).isEmpty
+      }
+    }
+  }
+
+  /** Returns the taint of `plan`'s output, or the reason its row set is not 
repeatable. */
+  private def analyze(plan: LogicalPlan): Either[String, Taint] = plan match {
+    case l: LeafNode =>
+      if (l.isOutputRepeatable) {
+        Right(Taint(AttributeSet.empty))
+      } else {
+        Left(s"the hinted side reads ${l.nodeName}, which may return different 
rows when " +
+          "evaluated again")
+      }
+
+    case p: Project => analyze(p.child).map { t =>
+      t.copy(unstable =
+        AttributeSet(p.projectList.filterNot(isStable(_, 
t.unstable)).map(_.toAttribute)))
+    }
+
+    case f: Filter => analyze(f.child).flatMap { t =>
+      if (isStable(f.condition, t.unstable)) Right(t) else Left(NotRepeatable)
+    }
+
+    case j: Join => analyze(j.left).flatMap { l =>
+      analyze(j.right).flatMap { r =>
+        val unstable = l.unstable ++ r.unstable
+        if (j.condition.forall(isStable(_, unstable))) {
+          Right(Taint(unstable, l.unstableMultiplicity || 
r.unstableMultiplicity))
+        } else {
+          Left(NotRepeatable)
+        }
+      }
+    }
+
+    case a: Aggregate => analyze(a.child).map { t =>
+      // Grouping on an unstable value changes which rows form a group, and 
unstable multiplicity
+      // changes how many rows each group has, so every aggregate result then 
depends on them; a
+      // grouping expression's own value is as stable as its input. The groups 
themselves occur
+      // once each, whatever the input multiplicity.
+      val stableGroups = a.groupingExpressions.forall(isStable(_, t.unstable))
+      val unstable = a.aggregateExpressions.filter { e =>
+        !isStable(e, t.unstable) ||
+          (e.exists(_.isInstanceOf[AggregateExpression]) &&
+            (!stableGroups || t.unstableMultiplicity || !isOrderIrrelevant(e)))
+      }
+      Taint(AttributeSet(unstable.map(_.toAttribute)))
+    }
+
+    case w: Window => analyze(w.child).map { t =>
+      val stablePartitions = w.partitionSpec.forall(isStable(_, t.unstable))
+      val unstable = w.windowExpressions.filter { e =>
+        !stablePartitions || t.unstableMultiplicity || !isStable(e, 
t.unstable) ||
+          !isOrderIrrelevantWindow(e)
+      }
+      t.copy(unstable = t.unstable ++ 
AttributeSet(unstable.map(_.toAttribute)),
+        totallyOrdered = false)
+    }
+
+    case u: Union =>
+      val taints = u.children.map(analyze)
+      taints.collectFirst { case Left(reason) => Left(reason) }.getOrElse {
+        val unstable = u.output.zipWithIndex.collect {
+          case (attr, i) if u.children.zip(taints).exists {
+            case (child, Right(taint)) => 
taint.unstable.contains(child.output(i))
+            case _ => false
+          } => attr
+        }
+        Right(Taint(AttributeSet(unstable),
+          taints.exists(_.exists(_.unstableMultiplicity))))
+      }
+
+    // An inner generate drops the rows for which the generator yields 
nothing, so an unstable
+    // generator changes the row set; an outer generate keeps every row but 
repeats it as many
+    // times as the generator yields values, so an unstable one changes the 
multiplicity.
+    case g: Generate => analyze(g.child).flatMap { t =>
+      if (isStable(g.generator, t.unstable)) {
+        Right(t.copy(totallyOrdered = false))
+      } else if (g.outer) {
+        Right(Taint(t.unstable ++ AttributeSet(g.generatorOutput), 
unstableMultiplicity = true))
+      } else {
+        Left(NotRepeatable)
+      }
+    }
+
+    case e: Expand => analyze(e.child).map { t =>
+      val unstable = e.output.zipWithIndex.collect {
+        case (attr, i) if e.projections.exists(p => !isStable(p(i), 
t.unstable)) => attr
+      }
+      Taint(AttributeSet(unstable), t.unstableMultiplicity)
+    }
+
+    case s: Sort => analyze(s.child).map { t =>
+      val stableOrder = s.order.forall(o => isStable(o.child, t.unstable))
+      t.copy(totallyOrdered = s.global && stableOrder && 
!t.unstableMultiplicity &&
+        sortedOnUniqueKey(s.child, s.order.map(_.child)))
+    }
+
+    // A limit keeps whichever rows arrive first unless the order is total.
+    case l @ (_: GlobalLimit | _: LocalLimit | _: Offset | _: Tail) =>
+      analyze(l.children.head).flatMap { t =>
+        if (t.totallyOrdered) Right(t) else Left(NotRepeatable)
+      }
+
+    // A sample draws a fresh seed per evaluation unless one is given, and 
depends on the input
+    // row order and partitioning even then: only a seeded sample over a 
repeatable leaf, through
+    // projections and filters that keep both, is repeatable. A predicate with 
a subquery is not
+    // such a filter, as its plan may differ between the two evaluations.
+    case s: Sample =>
+      @tailrec
+      def isOrderPreservingScan(p: LogicalPlan): Boolean = p match {
+        case Project(projectList, child) if 
projectList.forall(_.deterministic) =>
+          isOrderPreservingScan(child)
+        case Filter(condition, child)
+            if condition.deterministic && 
!SubqueryExpression.hasSubquery(condition) =>
+          isOrderPreservingScan(child)
+        case l: LeafNode => l.isOutputRepeatable
+        case _ => false
+      }
+      if (s.seed.isDefined && isOrderPreservingScan(s.child)) {

Review Comment:
   Fixed in cbef594a838. Any `Sample` in a hinted side is rejected; the 
seeded-sample tests now expect rejection.



##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala:
##########
@@ -306,10 +316,24 @@ object PartitionPruning extends Rule[LogicalPlan] with 
PredicateHelper with Join
    * meet the following requirements:
    *   (1) it can not be a stream
    *   (2) it needs to contain a selective predicate or a cheaply-recomputable 
materialized input
+   *
+   * (2) is evidence that pruning pays off, which a [[RuntimeFilterHint]] on 
the filtering side
+   * (`hinted`) supplies directly. A hinted side only has to be a repeatable 
source of the
+   * filtering key, see `JoinSelectionHelper.isRepeatableRuntimeFilterSource`, 
since DPP
+   * re-evaluates it. It must not contain a Python UDF either: the filtering 
side is planned as a
+   * standalone subquery without the Python UDF extraction the main plan gets, 
so a Python UDF
+   * could not be evaluated there.
    */
-  private def hasPartitionPruningFilter(plan: LogicalPlan): Boolean = {
-    !plan.isStreaming &&
-      (hasSelectivePredicate(plan) || 
isCheaplyRecomputableMaterializedPlan(plan))
+  private def hasPartitionPruningFilter(
+      plan: LogicalPlan,
+      hinted: Boolean,
+      filteringKey: Expression): Boolean = {
+    if (hinted) {
+      isRepeatableRuntimeFilterSource(plan, filteringKey) && 
!plan.containsPattern(PYTHON_UDF)

Review Comment:
   Fixed in cbef594a838. One check in the shared predicate rejects Python UDFs, 
`Generate(PythonUDTF)` and external UDFs on both paths; tests use a 
deterministic UDTF.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to