如何根据值列表在PySpark中对数据进行子集化

eagi6jfj  于 2022-12-27  发布在  Apache
关注(0)|答案(1)|浏览(126)

假设我有一个值列表:

list_codes = ["code_123", "code_456"]

这是PySpark Dataframe

+----------+-----------+
| code     |   value   |
+----------+-----------+
| code_456 | value_456 |
| code_123 | value_123 |
+----------+-----------+

我需要根据列表值的顺序按以下顺序从数据框中提取数据子集

+----------+-----------+
| code     |   value   |
+----------+-----------+
| code_123 | value_123 |
| code_456 | value_456 |
+----------+-----------+

当我使用这个命令时,它保持元素的原始顺序

df_subset = df.filter(f.col("code").isin(list_codes))

如何解决这个排序问题(实际上,当我们有超过2个值在列表中当然)?
谢谢!

eyh26e7m

eyh26e7m1#

我会通过添加一个rank列将排序值列表转换为dataframe,然后执行内部连接以进行过滤,最后通过rank列进行排序。

list_codes = ["code_123", "code_9", "code_456"]
data_codes = [(v,i) for i,v in enumerate(list_codes)]
df_codes = spark.createDataFrame(data=data_codes, schema = ["code", "rank"])
df_codes.show()

+--------+----+
|    code|rank|
+--------+----+
|code_123|   0|
|  code_9|   1|
|code_456|   2|
+--------+----+

而连接+排序部分将是(当然,你可以在末尾删除rank列):

data = [("code_123","value_123"), 
        ("code_456","value_456"), 
       ("code_789","value_789"), 
        ("code_9","value_9"), 
      ]
colums = ["code", "value"]
df = spark.createDataFrame(data=data, schema = colums).join(df_codes, ["code"], "inner") \
        .orderBy("rank") \
#         .drop("rank")
df.show()

+--------+---------+----+
|    code|    value|rank|
+--------+---------+----+
|code_123|value_123|   0|
|  code_9|  value_9|   1|
|code_456|value_456|   2|
+--------+---------+----+

如果答案符合您的需要,请不要忘记标记:)。

相关问题