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

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

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


##########
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:
   Agreed, a join could be a lot slower



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