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

时间:2019-04-05 14:18:52

标签: python apache-spark pyspark apache-spark-sql

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

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

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 个答案:

答案 0 :(得分:0)

首先将字典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|
#+--------+----------+--------+--------+