使用列表中的列对PySpark中的多个列进行分区

时间:2018-03-12 17:32:29

标签: apache-spark pyspark window-functions

我的问题类似于这个帖子: Partitioning by multiple columns in Spark SQL

但是我在Pyspark而不是Scala工作,我希望将列表列表作为列表传递。我想做这样的事情:

column_list = ["col1","col2"]
win_spec = Window.partitionBy(column_list)

我可以使用以下方法:

win_spec = Window.partitionBy(col("col1"))

这也有效:

col_name = "col1"
win_spec = Window.partitionBy(col(col_name))

这也有效:

win_spec = Window.partitionBy([col("col1"), col("col2")])

3 个答案:

答案 0 :(得分:6)

使用列表推导[col(x) for x in column_list]将列名转换为列表达式:

from pyspark.sql.functions import col
column_list = ["col1","col2"]
win_spec = Window.partitionBy([col(x) for x in column_list])

答案 1 :(得分:1)

你的第一次尝试应该有效。

考虑以下示例:

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

df = sqlCtx.createDataFrame(
    [
        ("a", "apple", 1),
        ("a", "orange", 2),
        ("a", "orange", 3),
        ("b", "orange", 3),
        ("b", "orange", 5)
    ],
    ["name", "fruit","value"]
)
df.show()
#+----+------+-----+
#|name| fruit|value|
#+----+------+-----+
#|   a| apple|    1|
#|   a|orange|    2|
#|   a|orange|    3|
#|   b|orange|    3|
#|   b|orange|    5|
#+----+------+-----+

假设您想计算每行总和的一小部分,按前两列进行分组:

cols = ["name", "fruit"]
w = Window.partitionBy(cols)
df.select(cols + [(f.col('value') / f.sum('value').over(w)).alias('fraction')]).show()

#+----+------+--------+
#|name| fruit|fraction|
#+----+------+--------+
#|   a| apple|     1.0|
#|   b|orange|   0.375|
#|   b|orange|   0.625|
#|   a|orange|     0.6|
#|   a|orange|     0.4|
#+----+------+--------+

答案 2 :(得分:0)

PySpark >= 2.4,这也有效 =>

column_list = ["col1","col2"]

win_spec = Window.partitionBy(*column_list)