根据条件 pyspark 删除列

Question

能否根据 Pyspark 中的条件删除列

我要删除列的条件：

df_train.groupby().sum() == 0

这是 pandas 中的一个简单示例：

import pandas as pd
#create dataframe
df = pd.DataFrame(np.array([[0,2,1],[0,2,8],[0,6,2]]), columns=['a','b', 'c']) 

#remove columns with only zero value
df.loc[:,df.sum(axis=0) != 0 ]

如果有多种方式，你会首选哪一种？

Answer 1

如果我没理解错的话，您想删除所有列的总和等于 0 的列。

您可以先计算每列的总和，然后过滤 sum = 0 所在的列列表并将该列表传递给 df.drop() 方法：

from pyspark.sql import functions as F


df = spark.createDataFrame([(0, 1, 2), (-1, 3, -6), (1, 4, 0)], ["col1", "col2", "col3"])

sums = df.select(*[F.sum(c).alias(c) for c in df.columns]).first()

cols_to_dop = [c for c in sums.asDict() if sums[c] == 0]

df = df.drop(*cols_to_dop)

df.show()
#+----+----+
#|col2|col3|
#+----+----+
#|   1|   2|
#|   3|  -6|
#|   4|   0|
#+----+----+