通过评估表达式连接两个 spark 数据帧

Question

我有两个火花数据框

userMemberShipDF:

用户	会员数组
a1	s1、s2、s3
a2	s4, s6
a3	s5、s4、s3
a4	s1,s3,s4,s5
a5	s2、s4、s6
a6	s3、s7、s1
a7	s1、s4、s6

和 categoryDF 与

类别ID	会员表达式	开始日期	期间
c1	s1 || s2	2022-05-01	30
c2	s4 && s6 && !s2	2022-06-20	50
c3	s3 && s4	2022-06-10	60

结果数据框包含user, category_id, start_date, duration列

我已经编写了一个 function ，它将从第二个数据帧中获取membership_expression 以及第一个dataframe 中的membership_array 并评估为真或假。

例如membership_expression = s1 || s2 将匹配所有用户 a1、a4、a5、a6 和 a7，而表达式s4 && s6 !s2将仅匹配 a2 等。

我想根据此表达式的计算结果为真或假来加入两个数据框。 我查看了火花连接，它只会将列作为连接条件，而不是 boolean 表达式。

所以我尝试了以下方法

val matchedUserSegments = userMemberShipDF
      .map { r =>
        {
        // categoryDF is broadcasted
          val category_items_set = categoryDF.value.flatMap { fl =>
            {
              if (CategoryEvaluator.evaluateMemberShipExpression(fl.membership_expression, r.membership_array)) {
                Some(fl.category_id)
              } else {
                None
              }
            }
          }
          (r.user_id, category_items_set)
        }
      }
      .toDF("user_id", "category_items_set")

然后在 category_items_set 上分解生成的 dataframe ，然后在 categoryDF 上加入，得到所需的 output 表。

我知道我做了两次操作，但找不到更好的方法来通过迭代两个数据帧来计算所有内容。

请提出一种有效的方法来做到这一点。

我有很多数据，而 spark 工作需要 24 小时以上才能完成。 谢谢

Answer 1

PS：为简单起见，我没有包含start_date和duration ，并且还将示例user行限制为a1, a2, a3, a4 。 此处显示的 Output 可能与您预期的 output 不完全匹配； 但如果你使用完整的数据，我相信输出会匹配。

import pyspark.sql.functions as F

userMemberShipDF = spark.createDataFrame([
    ("a1",["s1","s2","s3"]),
    ("a2",["s4","s6"]),
    ("a3",["s5","s4","s3"]),
    ("a4",["s1","s3","s4","s5"]),
], ["user","membership_array"])

如果用户具有该成员资格，则将每个成员资格 s1、s2、s3 等转换为单独的列并标记为true ：

userMemberShipDF = userMemberShipDF.withColumn("membership_individual", F.explode("membership_array"))
+----+----------------+---------------------+
|user|membership_array|membership_individual|
+----+----------------+---------------------+
|  a1|    [s1, s2, s3]|                   s1|
|  a1|    [s1, s2, s3]|                   s2|
|  a1|    [s1, s2, s3]|                   s3|
|  a2|        [s4, s6]|                   s4|
|  a2|        [s4, s6]|                   s6|
|  a3|    [s5, s4, s3]|                   s5|
|  a3|    [s5, s4, s3]|                   s4|
|  a3|    [s5, s4, s3]|                   s3|
|  a4|[s1, s3, s4, s5]|                   s1|
|  a4|[s1, s3, s4, s5]|                   s3|
|  a4|[s1, s3, s4, s5]|                   s4|
|  a4|[s1, s3, s4, s5]|                   s5|
+----+----------------+---------------------+


userMemberShipDF = userMemberShipDF.groupBy("user").pivot("membership_individual").agg(F.count("*").isNotNull()).na.fill(False)
+----+-----+-----+-----+-----+-----+-----+
|user|   s1|   s2|   s3|   s4|   s5|   s6|
+----+-----+-----+-----+-----+-----+-----+
|  a3|false|false| true| true| true|false|
|  a4| true|false| true| true| true|false|
|  a2|false|false|false| true|false| true|
|  a1| true| true| true|false|false|false|
+----+-----+-----+-----+-----+-----+-----+

在类别数据中，替换|| , && , ! 有or , and , not :

categoryDF = spark.createDataFrame([
    ("c1", "s1 || s2"),
    ("c2", "s4 && s6 && !s2"),
    ("c3", "s3 && s4"),
], ["category_id", "membership_expression"])

categoryDF = categoryDF.withColumn("membership_expression", F.regexp_replace("membership_expression", "\|\|", " or "))
categoryDF = categoryDF.withColumn("membership_expression", F.regexp_replace("membership_expression", "\&\&", " and "))
categoryDF = categoryDF.withColumn("membership_expression", F.regexp_replace("membership_expression", "\!", " not "))

+-----------+-------------------------+
|category_id|membership_expression    |
+-----------+-------------------------+
|c1         |s1  or  s2               |
|c2         |s4  and  s6  and   not s2|
|c3         |s3  and  s4              |
+-----------+-------------------------+

交叉连接用户和类别数据以针对每个类别评估每个用户：

resultDF_sp = categoryDF.crossJoin(userMemberShipDF)
+-----------+-------------------------+----+-----+-----+-----+-----+-----+-----+
|category_id|membership_expression    |user|s1   |s2   |s3   |s4   |s5   |s6   |
+-----------+-------------------------+----+-----+-----+-----+-----+-----+-----+
|c1         |s1  or  s2               |a3  |false|false|true |true |true |false|
|c1         |s1  or  s2               |a4  |true |false|true |true |true |false|
|c1         |s1  or  s2               |a2  |false|false|false|true |false|true |
|c1         |s1  or  s2               |a1  |true |true |true |false|false|false|
|c2         |s4  and  s6  and   not s2|a3  |false|false|true |true |true |false|
|c2         |s4  and  s6  and   not s2|a4  |true |false|true |true |true |false|
|c2         |s4  and  s6  and   not s2|a2  |false|false|false|true |false|true |
|c2         |s4  and  s6  and   not s2|a1  |true |true |true |false|false|false|
|c3         |s3  and  s4              |a3  |false|false|true |true |true |false|
|c3         |s3  and  s4              |a4  |true |false|true |true |true |false|
|c3         |s3  and  s4              |a2  |false|false|false|true |false|true |
|c3         |s3  and  s4              |a1  |true |true |true |false|false|false|
+-----------+-------------------------+----+-----+-----+-----+-----+-----+-----+

评估membership_expression

啊！ 这部分不优雅

Spark 提供expr function来使用列值评估 SQL 表达式； 但这仅在表达式是 static 字符串时才有效：

resultDF_sp.select(F.expr("s1 or s2"))

但如果表达式"s1 or s2"是一个列值（如上面的membership_expression列），则无法对其进行评估。 这导致错误Column is not iterable ：

resultDF_sp.select(F.expr(F.col("membership_expression")))

关于这个，stackoverflow 有几个问题； 但他们都建议解析表达式，并编写一个评估器来手动评估解析的表达式：

• 如何评估作为列值的表达式？
• 使用列值作为列名

幸运的是，可以使用其他列的参数值将表达式计算为列值。

所以，我不喜欢的部分； 但别无选择是将 dataframe 转换为 pandas，评估表达式并转换回火花（如果有人可以建议如何在火花中实现这一点，我很乐意将其包含在编辑中）：

resultDF_pd = resultDF_sp.toPandas()

def evaluate_expr(row_series):
    df = row_series.to_frame().transpose().infer_objects()
    return df.eval(df["membership_expression"].values[0]).values[0]

resultDF_pd["is_matching_user"] = resultDF_pd.apply(lambda row: evaluate_expr(row), axis=1)
resultDF_sp = spark.createDataFrame(resultDF_pd[["category_id", "user", "is_matching_user"]])
+-----------+----+----------------+
|category_id|user|is_matching_user|
+-----------+----+----------------+
|         c1|  a3|           false|
|         c1|  a4|            true|
|         c1|  a2|           false|
|         c1|  a1|            true|
|         c2|  a3|           false|
|         c2|  a4|           false|
|         c2|  a2|            true|
|         c2|  a1|           false|
|         c3|  a3|            true|
|         c3|  a4|            true|
|         c3|  a2|           false|
|         c3|  a1|           false|
+-----------+----+----------------+

最后，过滤匹配的用户：

resultDF_sp = resultDF_sp.filter("is_matching_user")
+-----------+----+----------------+
|category_id|user|is_matching_user|
+-----------+----+----------------+
|         c1|  a4|            true|
|         c1|  a1|            true|
|         c2|  a2|            true|
|         c3|  a3|            true|
|         c3|  a4|            true|
+-----------+----+----------------+