apache-spark - 如何迭代一个组并使用 Pyspark 创建一个数组列？

Question

我有一个包含组和百分比的数据框

| Group | A % | B % | Target % |
| ----- | --- | --- | -------- |
| A     | .05 | .85 | 1.0      |
| A     | .07 | .75 | 1.0      |
| A     | .08 | .95 | 1.0      |
| B     | .03 | .80 | 1.0      |
| B     | .05 | .83 | 1.0      |
| B     | .04 | .85 | 1.0      |

我希望能够A %逐列迭代Group并从 column 中找到一个值数组B %，当与 column 中的每个值相加时A%小于或等于 column Target %。

| Group | A % | B % | Target % | SumArray     |
| ----- | --- | --- | -------- | ------------ |
| A     | .05 | .85 | 1.0      | [.85,.75,.95]|
| A     | .07 | .75 | 1.0      | [.85,.75]    |
| A     | .08 | .95 | 1.0      | [.85,.75]   |
| B     | .03 | .80 | 1.0      | [.80,.83,.85]|
| B     | .05 | .83 | 1.0      | [.80,.83,.85]|
| B     | .04 | .85 | 1.0      | [.80,.83,.85]|

我希望能够使用 PySpark 来解决这个问题。任何想法如何解决这个问题？

标签： apache-sparkpysparkgroup-byapache-spark-sql

Answer 1

您可以使用collect_list函数获取B %按列分组的列值数组，Group然后filter使用您的条件生成结果数组A + B <= Target：

from pyspark.sql import Window
import pyspark.sql.functions as F

df2 = df.withColumn(
    "SumArray",
    F.collect_list(F.col("B")).over(Window.partitionBy("Group"))
).withColumn(
    "SumArray",
    F.expr("filter(SumArray, x -> x + A <= Target)")
)
df2.show()

# +-----+----+----+------+------------------+
# |Group|   A|   B|Target|          SumArray|
# +-----+----+----+------+------------------+
# |    B|0.03| 0.8|   1.0| [0.8, 0.83, 0.85]|
# |    B|0.05|0.83|   1.0| [0.8, 0.83, 0.85]|
# |    B|0.04|0.85|   1.0| [0.8, 0.83, 0.85]|
# |    A|0.05|0.85|   1.0|[0.85, 0.75, 0.95]|
# |    A|0.07|0.75|   1.0|      [0.85, 0.75]|
# |    A|0.08|0.95|   1.0|      [0.85, 0.75]|
# +-----+----+----+------+------------------+