计算每行并在DataFrame PySpark中添加新列 - 更好的解决方案?

时间:2017-01-21 14:28:10

标签: apache-spark dataframe pyspark apache-spark-sql rdd

我在PySpark中使用Data Frame 我有以下任务:检查每列的“时间”值是多少>所有列均为2。对于u1,它是0,对于u2 => 2等等

user    a   b   c   d   times
   u1   1   0   1   0   0
   u2   0   1   4   3   2
   u3   2   1   7   0   1

我的解决方案如下。它有效,我不确定它是最好的方式,并没有尝试真正的大数据。我不喜欢转换为rdd并返回数据框。有更好的吗?我在开始时每列都会被UDF搞砸,但是没有找到方法来计算每行的所有结果:

def calculate_times(row):
    times = 0
    for index, item in enumerate(row):
        if not isinstance(item, basestring):
           if item > 2:
             times = times+1

return times    

def add_column(pair):
    return dict(pair[0].asDict().items() + [("is_outlier", pair[1])])   

def calculate_times_for_all(df): 
    rdd_with_times = df.map(lambda row: (calculate_times(row))
    rdd_final = df.rdd.zip(rdd_with_times).map(add_column)

    df_final = sqlContext.createDataFrame(rdd_final)
    return  df_final

对于这个解决方案,我使用了这个主题 How do you add a numpy.array as a new column to a pyspark.SQL DataFrame?

谢谢!

1 个答案:

答案 0 :(得分:4)

这只是一个简单的单行。示例数据:

df = sc.parallelize([
    ("u1", 1, 0, 1, 0), ("u2", 0, 1, 4, 3), ("u3", 2, 1, 7, 0)
]).toDF(["user", "a", "b", "c", "d"])

withColumn

df.withColumn("times", sum((df[c] > 2).cast("int") for c in df.columns[1:]))

结果:

+----+---+---+---+---+-----+
|user|  a|  b|  c|  d|times|
+----+---+---+---+---+-----+
|  u1|  1|  0|  1|  0|    0|
|  u2|  0|  1|  4|  3|    2|
|  u3|  2|  1|  7|  0|    1|
+----+---+---+---+---+-----+

注意:

它的列nullable你应该更正,例如使用coalesce

from pyspark.sql.functions import coalesce

sum(coalesce((df[c] > 2).cast("int"), 0) for c in df.columns[1:])