This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new aa4df79f810c [SPARK-58606][ML] Move association-rule collection and
preparation in FPGrowthModel transform from driver to executors
aa4df79f810c is described below
commit aa4df79f810c025a2cc1da9ccf37c586c7f441fc
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Thu Aug 6 15:22:54 2026 +0800
[SPARK-58606][ML] Move association-rule collection and preparation in
FPGrowthModel transform from driver to executors
### What changes were proposed in this pull request?
This PR replaces the manually collected association-rule array and
user-defined function in `FPGrowthModel.transform` with DataFrame operations.
It joins input transactions with matching rules, aggregates consequents for
each input row, and uses `array_except` to exclude items already present in the
transaction.
The join has no broadcast hint, so Catalyst may use a broadcast nested-loop
join only when the plan statistics consider a side small enough.
### Why are the changes needed?
Collecting every association rule before a transform creates avoidable
driver-memory pressure, particularly for Spark Connect server workloads with
many rules.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- `build/sbt -java-home /usr/lib/jvm/java-17-openjdk-amd64 mllib/compile`
- The focused `FPGrowthSuite` has not been run locally yet.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57806 from zhengruifeng/ml_fpgrowth_transform_join.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 21a1ae1ef937633ee3135a0963922bb0df39ddf9)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../scala/org/apache/spark/ml/fpm/FPGrowth.scala | 34 +++++++++-------------
1 file changed, 14 insertions(+), 20 deletions(-)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/fpm/FPGrowth.scala
b/mllib/src/main/scala/org/apache/spark/ml/fpm/FPGrowth.scala
index 0a5213228a8b..d92a8f4c27a5 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/fpm/FPGrowth.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/fpm/FPGrowth.scala
@@ -273,35 +273,29 @@ class FPGrowthModel private[ml] (
* from all the applicable rules as prediction. The prediction column has
the same data type as
* the input column(Array[T]) and will not contain existing items in the
input column. The null
* values in the itemsCol columns are treated as empty sets.
- * WARNING: internally it collects association rules to the driver and uses
broadcast for
- * efficiency. This may bring pressure to driver memory for large set of
association rules.
+ * Internally, transform aggregates association rules into a single-row
DataFrame and joins it
+ * with the input dataset.
*/
@Since("2.2.0")
override def transform(dataset: Dataset[_]): DataFrame = {
transformSchema(dataset.schema, logging = true)
- genericTransform(dataset)
- }
-
- private def genericTransform(dataset: Dataset[_]): DataFrame = {
- val rules: Array[(Seq[Any], Seq[Any])] =
associationRules.select("antecedent", "consequent")
- .rdd.map(r => (r.getSeq(0), r.getSeq(1)))
- .collect().asInstanceOf[Array[(Seq[Any], Seq[Any])]]
- val brRules = dataset.sparkSession.sparkContext.broadcast(rules)
-
val dt = dataset.schema($(itemsCol)).dataType
- // For each rule, examine the input items and summarize the consequents
- val predictUDF = SparkUserDefinedFunction((items: Seq[Any]) => {
+ // For each rule, examine the input items and summarize the consequents.
+ val predictFunc = (items: Seq[Any], rules: Seq[Row]) => {
if (items != null) {
val itemset = items.toSet
- brRules.value.filter(_._1.forall(itemset.contains))
- .flatMap(_._2.filter(!itemset.contains(_))).distinct
+ rules.filter(_.getSeq[Any](0).forall(itemset.contains))
+ .flatMap(_.getSeq[Any](1).filter(!itemset.contains(_))).distinct
} else {
Seq.empty
- }},
- dt,
- Nil
- )
- dataset.withColumn($(predictionCol), predictUDF(col($(itemsCol))))
+ }
+ }
+ val predictUDF = SparkUserDefinedFunction(predictFunc, dt, Nil)
+ dataset.join(
+ associationRules.select("antecedent", "consequent")
+ .agg(collect_set(struct("antecedent", "consequent")).as("rules")))
+ .withColumn($(predictionCol), predictUDF(col($(itemsCol)), col("rules")))
+ .drop("rules")
}
@Since("2.2.0")
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
