过滤以多列为条件的数据框，根据列值具有不同的条件

Posted 2023-04-15

【中文标题】过滤以多列为条件的数据框，根据列值具有不同的条件

【英文标题】：Filtering dataframes conditioned on multiple columns, with different conditions depending on column values

【发布时间】：2019-04-05 14:18:52

【问题描述】：

我正在尝试过滤以多列为条件的数据框。问题是条件随列值而变化。

例如，我有以下原始数据集。

df_assets = pd.DataFrame(
    'asset_id': ['10092', '39883', '82883', '28883', '09093', '27817', '38261', '28372', '38472', '39738'],
     'asset_type' :['laptop', 'mousepad', 'mouse', 'pen', 'pen', 'monitor', 'wire', 'laptop', 'keyboard', 'mouse'],
     'quantity': [3, 12, 12, 10, 2, 0, 21, 0, 10, 16],
     'assigned': ['susan', 'jane', 'amy', 'ben', 'donald', 'evan', 'chris', 'fred', 'henry', 'george']
    )

+--------+----------+--------+--------+
|asset_id|asset_type|assigned|quantity|
+--------+----------+--------+--------+
|   10092|    laptop|   susan|     3.0|
|   39883|  mousepad|    jane|    12.0|
|   82883|     mouse|     amy|    12.0|
|   28883|       pen|     ben|    10.0|
|   09093|       pen|  donald|     2.0|
|   27817|   monitor|    evan|     0.0|
|   38261|      wire|   chris|    21.0|
|   28372|    laptop|    fred|     0.0|
|   38472|  keyboard|   henry|    10.0|
|   39738|     mouse|  george|    16.0|
+--------+----------+--------+--------+

我想应用一个过滤器来返回数量高于以下字典中所述数量的行（即数量 >= 10 的键盘、数量 >= 10 的电线、数量 >= 1 的笔记本电脑等）。我还想返回不包含字典中指定的资产类型的行。

d = 'keyboard': 10, 'wire': 20, 'laptop': 1, 'mouse': 15, 'monitor': 1

输出应该是

+--------+----------+--------+--------+
|asset_id|asset_type|assigned|quantity|
+--------+----------+--------+--------+
|   10092|    laptop|   susan|     3.0|
|   38261|      wire|   chris|    21.0|
|   39738|     mouse|  george|    16.0|
|   38472|  keyboard|   henry|    10.0|
|   39883|  mousepad|    jane|    12.0|
|   28883|       pen|     ben|    10.0|
|   09093|       pen|  donald|     2.0|
+--------+----------+--------+--------+

有没有一种优雅的方法来做到这一点？我通过为字典中的每个项目过滤和创建数据框并附加到最​​终数据框来实现上述功能。

count = 0
for k, v in d.items():
    sdf_filter = sdf_assets.filter( (f.col('asset_type') == k) & (f.col('quantity') >= v) )
    if count != 0:
        sdf_thresholded = sdf_thresholded.union(sdf_filter)
    else:
        sdf_thresholded = sdf_filter
    count += 1

sdf_result_test = sdf_thresholded.union(sdf_assets.filter(~f.col('asset_type').isin(d.keys())))

我在 PySpark 中编写了所有内容，因为完整的数据集有很多条目，但我希望能在 python 中提供任何解决方案！

【参考方案1】：

首先将您的字典 d 转换为 spark DataFrame：

d = 'keyboard': 10, 'wire': 20, 'laptop': 1, 'mouse': 15, 'monitor': 1
d = spark.createDataFrame(d.items(), ["asset_type", "min_quantity"])
d.show()
#+----------+------------+
#|asset_type|min_quantity|
#+----------+------------+
#|    laptop|           1|
#|      wire|          20|
#|     mouse|          15|
#|   monitor|           1|
#|  keyboard|          10|
#+----------+------------+

现在您可以将其与df_assets spark DataFrame 一起加入，并使用where 根据您的条件过滤行。由于d 很小，您可以使用broadcast join 来提高性能。

from pyspark.sql.functions import broadcast 

# first convert from pandas DataFrame to spark DataFrame
df_assets = spark.createDataFrame(df_assets)

# left join with d
df_assets.join(broadcast(d), on="asset_type", how="left")\
    .where("min_quantity IS NULL or quantity >= min_quantity")\
    .select("asset_id", "asset_type", "assigned", "quantity")\
    .show()
#+--------+----------+--------+--------+
#|asset_id|asset_type|assigned|quantity|
#+--------+----------+--------+--------+
#|   10092|    laptop|   susan|       3|
#|   39883|  mousepad|    jane|      12|
#|   28883|       pen|     ben|      10|
#|   09093|       pen|  donald|       2|
#|   38261|      wire|   chris|      21|
#|   38472|  keyboard|   henry|      10|
#|   39738|     mouse|  george|      16|
#+--------+----------+--------+--------+