4
我想计算每行的列子集的最大值,并将其作为现有Dataframe
的新列添加。PySpark:计算列子集的最大行数并添加到现有数据框
我能够做到这一点非常尴尬的方式:
def add_colmax(df,subset_columns,colnm):
'''
calculate the maximum of the selected "subset_columns" from dataframe df for each row,
new column containing row wise maximum is added to dataframe df.
df: dataframe. It must contain subset_columns as subset of columns
colnm: Name of the new column containing row-wise maximum of subset_columns
subset_columns: the subset of columns from w
'''
from pyspark.sql.functions import monotonicallyIncreasingId
from pyspark.sql import Row
def get_max_row_with_None(row):
return float(np.max(row))
df_subset = df.select(subset_columns)
rdd = df_subset.map(get_max_row_with_None)
df_rowsum = rdd.map(Row(colnm)).toDF()
df_rowsum = df_rowsum.withColumn("id",monotonicallyIncreasingId())
df = df.withColumn("id",monotonicallyIncreasingId())
df = df.join(df_rowsum,df.id == df_rowsum.id).drop(df.id).drop(df_rowsum.id)
return df
此功能的作用:
rdd1 = sc.parallelize([("foo", 1.0,3.0,None),
("bar", 2.0,2.0,-10),
("baz", 3.3,1.2,10.0)])
df1 = sqlContext.createDataFrame(rdd1, ('v1', 'v2','v3','v4'))
df_new = add_colmax(df1,['v2','v3','v4'],"rowsum")
df_new.collect()
回报:
[Row(v1=u'bar', v2=2.0, v3=2.0, v4=-10, rowsum=2.0),
Row(v1=u'baz', v2=3.3, v3=1.2, v4=None, rowsum=3.3),
Row(v1=u'foo', v2=1.0, v3=3.0, v4=None, rowsum=3.0)]
我想,如果我可以使用用户用withColumn
定义函数,这可以做得简单得多。但我无法弄清楚如何去做。 请让我知道,如果你有更简单的方法来实现这一点。 我使用的火花1.6