PySpark:一次性在多个列上运行相同的操作

时间:2016-03-18 15:15:02

标签: python sql select dataframe pyspark

我的DataFrame table包含

等行
['row1', 'col_1', 'col_2', 'col_3', ..., 'col_N', 'alpha']

N(除了第一个和最后一个之外的列数)相对较大。

现在,我需要通过将名为col_i的每个列乘以列alpha来创建另一个DataFrame。是否有一种更聪明的方法,而不是每个列都进行手动乘法,如:

sc = SparkContext()
sqlc = SQLContext(sc)

sqlc.sql('SELECT col_1 * alpha, col_2 * alpha, ..., col_N * alpha FROM table')

所以我想知道是否可以对每一列进行相同的操作,而不是为每一列专门编写。

1 个答案:

答案 0 :(得分:2)

不确定这是多么有效,但我可能会这样做:

col_names = df.columns
# start from one to exclude row, end at -1 to exclude alpha column
for x in range(1, len(col_names) - 1): 
    new_column_name = col_names[x] + "_x_alpha" # get good new column names
    df = df.withColumn(new_column_name, (getattr(df, col_names[x]) * getattr(df, col_names[3])))

这将生成您最初使用的相同数据框,但新列会将每个column_*条目乘以alpha中的条目。

当我在我的小例子上运行df.show()时,我得到了这个输出:

row col_1 col_2 alpha col_1_x_alpha col_2_x_alpha
1   2     3     4     8             12           
2   3     4     5     15            20           
3   4     5     6     24            30  

然后,您可以运行SQL查询以仅获取column_*_x_alpha类型的列。