[GitHub] [spark] zhengruifeng commented on a diff in pull request #40263: [SPARK-42659][ML] Reimplement `FPGrowthModel.transform` with dataframe operations

["zhengruifeng (via GitHub)" <gi...@apache.org>]

zhengruifeng commented on code in PR #40263:
URL: https://github.com/apache/spark/pull/40263#discussion_r1125383324


##########
mllib/src/main/scala/org/apache/spark/ml/fpm/FPGrowth.scala:
##########
@@ -275,29 +274,38 @@ class FPGrowthModel private[ml] (
   @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]) => {
-      if (items != null) {
-        val itemset = items.toSet
-        brRules.value.filter(_._1.forall(itemset.contains))
-          .flatMap(_._2.filter(!itemset.contains(_))).distinct
-      } else {
-        Seq.empty
-      }},
-      dt,
-      Nil
+    val arrayType = associationRules.schema("consequent").dataType
+
+    dataset.crossJoin(
+      broadcast(
+        associationRules
+          .where(not(isnull(col("antecedent"))) &&
+            not(isnull(col("consequent"))))
+          .select(
+            collect_list(
+              struct("antecedent", "consequent")
+            ).as($(predictionCol))
+          )
+      )
+    ).withColumn(
+      $(predictionCol),
+      when(not(isnull(col($(itemsCol)))),
+        array_sort(
+          array_distinct(
+            aggregate(
+              col($(predictionCol)),
+              array().cast(arrayType),
+              (r, s) => when(
+                forall(s.getField("antecedent"),
+                  c => array_contains(col($(itemsCol)), c)),
+                array_union(r,
+                  array_except(s.getField("consequent"), col($(itemsCol))))
+              ).otherwise(r)
+            )
+          )
+        )
+      ).otherwise(array().cast(arrayType))

Review Comment:
   actually the main goals are:
   1, try to avoid eager action in transformation; like what we did in https://github.com/apache/spark/commit/6a0713a141fa98d83029d8388508cbbc40fd554e
   2, try to avoid collect the DF `associationRules` to the driver, I feel such collection will be a potential problem in Spark Connect, which there are many Spark Sessions in the driver. 



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org