在 PySpark 中提取特定行

Question

我有一个这样的数据框

data = [(("ID1", "A", 1)), (("ID1", "B", 5)), (("ID2", "A", 12)), 
       (("ID3", "A", 3)), (("ID3", "B", 3)), (("ID3", "C", 5)), (("ID4", "A", 10))]
df = spark.createDataFrame(data, ["ID", "Type", "Value"])
df.show()

+---+----+-----+
| ID|Type|Value|
+---+----+-----+
|ID1|   A|    1|
|ID1|   B|    5|
|ID2|   A|   12|
|ID3|   A|    3|
|ID3|   B|    3|
|ID3|   C|    5|
|ID4|   A|   10|
+---+----+-----+

我只想提取那些只包含一个特定类型 - “A”的行（或 ID）

因此我的预期输出将包含以下行

+---+----+-----+
| ID|Type|Value|
+---+----+-----+
|ID2|   A|    1|
|ID4|   A|   10|
+---+----+-----+

对于每个 ID 可以包含任何类型 - A、B、C 等。我想提取那些包含一个且仅一个类型的 ID - 'A'

如何在 PySpark 中实现这一点

Answer 1

您可以对其应用过滤器。

import pyspark.sql.functions as f

data = [(("ID1", "A", 1)), (("ID1", "B", 5)), (("ID2", "A", 12)), 
       (("ID3", "A", 3)), (("ID3", "B", 3)), (("ID3", "C", 5)), (("ID4", "A", 10))]
df = spark.createDataFrame(data, ["ID", "Type", "Value"])
df.show()

+---+----+-----+
| ID|Type|Value|
+---+----+-----+
|ID1|   A|    1|
|ID1|   B|    5|
|ID2|   A|   12|
|ID3|   A|    3|
|ID3|   B|    3|
|ID3|   C|    5|
|ID4|   A|   10|
+---+----+-----+

x= df.filter(f.col('Type')=='A')

x.show()

如果我们需要过滤所有只有一条记录且类型为“A”的 ID，那么下面的代码可能是解决方案

df.registerTempTable('table1')


sqlContext.sql('select a.ID, a.Type,a.Value from table1 as a, (select ID, count(*) as cnt_val from table1 group by ID) b where a.ID = b.ID and (a.Type=="A" and b.cnt_val ==1)').show()


+---+----+-----+
| ID|Type|Value|
+---+----+-----+
|ID2|   A|   12|
|ID4|   A|   10|
+---+----+-----+

会有更好的替代方法来找到相同的。

Answer 2

根据 OP 的要求，我正在记下我在 cmets 下写的答案。

当前问题的目的是过滤掉DataFrame，其中每个特定的ID 都只有Type A 的一个元素，而没有其他元素。

# Loading the requisite packages
from pyspark.sql.functions import col, collect_set, array_contains, size, first

我们的想法是先将aggregate() DataFrame 与ID 组合在一起，由此我们将Type 的所有unique 元素使用collect_set() 分组到一个数组中。拥有unique 元素很重要，因为对于特定的ID 可能有两行，这两行的Type 都具有A。这就是我们应该使用collect_set() 而不是collect_list() 的原因，因为后者不会返回唯一元素，而是返回所有元素。

那么我们应该使用first() 来获取Type 和Value 在一个组中的第一个值。如果A 是特定ID 唯一可能的unique Type，则first() 将返回A 的唯一值，以防A 出现一次，如果有则返回最大值A 的重复。

df = df = df.groupby(['ID']).agg(first(col('Type')).alias('Type'),
                                 first(col('Value')).alias('Value'),
                                 collect_set('Type').alias('Type_Arr'))
df.show()
+---+----+-----+---------+
| ID|Type|Value| Type_Arr|
+---+----+-----+---------+
|ID2|   A|   12|      [A]|
|ID3|   A|    3|[A, B, C]|
|ID1|   A|    1|   [A, B]|
|ID4|   A|   10|      [A]|
+---+----+-----+---------+

最后，我们将同时设置两个条件来过滤出所需的数据集。

条件 1： 它使用array_contains() 检查Type 数组中是否存在A。

条件2：检查数组的size。如果大小大于1，则应该有多个Types。

df = df.where(array_contains(col('Type_Arr'),'A') & (size(col('Type_Arr'))==1)).drop('Type_Arr')
df.show()
+---+----+-----+
| ID|Type|Value|
+---+----+-----+
|ID2|   A|   12|
|ID4|   A|   10|
+---+----+-----+

Answer 3

我不熟悉 Python，这里有一个 Scala 可能的解决方案：

df.groupBy("ID").agg(collect_set("Type").as("Types"))
  .select("ID").where((size($"Types")===1).and(array_contains($"Types", "A"))).show()
+---+
| ID|
+---+
|ID2|
|ID4|
+---+

这个想法是按ID 分组并仅过滤包含A 值的大小为1 的Types。