使用不同类型（array[double]vs double）乘以两个pyspark dataframe列，而不使用breeze

Question

我的问题与被问到的问题相同，但我需要一个在pyspark中没有微风的解决方案

例如，如果我的pyspark数据框如下所示：

user    |  weight  |  vec
"u1"    | 0.1      | [2, 4, 6]
"u1"    | 0.5      | [4, 8, 12]
"u2"    | 0.5      | [20, 40, 60]

当column weight的类型为double，column vec的类型为Array[double]时，我想得到每个用户向量的加权和，这样我得到的数据帧如下所示：

user    |  wsum
"u1"    | [2.2, 4.4, 6.6]
"u2"    | [10, 20, 30]

为此，我尝试了以下方法：

df.groupBy('user').agg((F.sum(df.vec* df.weight)).alias("wsum"))

但它失败了，因为vec列和weight列的类型不同

没有微风，我如何解决这个错误

Answer 1

正在使用Spark 2.4中提供的高阶函数transform：

# get size of vec array
n = df.select(size("vec")).first()[0]

# transform each element of the vec array
transform_expr = "transform(vec, x -> x * weight)"

df.withColumn("weighted_vec", expr(transform_expr)) \
  .groupBy("user").agg(array(*[sum(col("weighted_vec")[i]) for i in range(n)]).alias("wsum"))\
  .show()

给出：

+  +         +
|user|              wsum|
+  +         +
|  u1|   [2.2, 4.4, 6.6]|
|  u2|[10.0, 20.0, 30.0]|
+  +         +

对于Spark<；2.4，使用a表示理解，将每个元素乘以weight列，如下所示：

df.withColumn("weighted_vec", array(*[col("vec")[i] * col("weight") for i in range(n)])) \
  .groupBy("user").agg(array(*[sum(col("weighted_vec")[i]) for i in range(n)]).alias("wsum")) \
  .show()