Python：如何从包含列表中的值的dataframe中删除所有列？

Question

include_cols_path = sys.argv[5]
with open(include_cols_path) as f:
include_cols = f.read().splitlines()

include_cols是一个字符串列表

df1 = sqlContext.read.csv(input_path + '/' + lot_number +'.csv', header=True).toPandas()

df1是一个大文件的数据帧。我只想保留名称包含include_cols中任何字符串的列。

Answer 1

df1.loc[:, df1.columns.str.contains('|'.join(include_cols))]

例如：

df1 = pd.DataFrame(data=np.random.random((5, 5)), columns=list('ABCDE'))
include_cols = ['A', 'C', 'Z']
df1.loc[:, df1.columns.str.contains('|'.join(include_cols))]
>>>           A         C
    0  0.247271  0.761153
    1  0.390240  0.050055
    2  0.333401  0.823384
    3  0.821196  0.929520
    4  0.210226  0.406168

“|”.join（include_cols）部分将使用输入列表的所有元素创建一个或条件。在上面的例子中，A | C | Z。如果使用列名上的.contains（）方法在列名中包含一个元素，则此条件将为True。

Answer 2

在pandas中这样做肯定是个骗局。但是，您似乎正在将spark数据帧转换为pandas数据帧。

与其执行（昂贵的）收集操作，然后过滤您想要的列，不如使用选择（）在火花侧进行过滤：

df1 = sqlContext.read.csv(input_path + '/' + lot_number +'.csv', header=True)
pandas_df = df1.select(include_cols).toPandas()

您还应该考虑转换为熊猫DataFrame是否真的是您想要做的。你可以在熊猫中做的任何事情也可以在火花中完成。

编辑

我本来误解了你的问题。根据你的评论，我想这就是你要找的：

selected_columns = [c for c in df1.columns if any([x in c for x in include_cols])]
pandas_df = df1.select(selected_columns).toPandas()

说明：遍历df1中的列，并且只保留列名称中至少包含include_cols中的一个字符串的列。any（）函数返回True，如果条件中至少有一个是True。

Answer 3

final_cols = [col for col in df.columns.values if col in include_cols]
df = df[final_cols]