zhengruifeng commented on code in PR #40263:
URL: https://github.com/apache/spark/pull/40263#discussion_r1125384107
##########
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:
> It looks more complex than the spark UDF code before
I think I can simplify it a bit.
> a join could be a lot slower
Since it is a BroadcastHashJoin, so I think will not be very bad
--
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-unsubscr...@spark.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
