Question

我在pyspark中制作了一个数据框-

df = spark.createDataFrame([
    ("S1", "S1_P1", "i1"),
    ("S1", "S1_P2", "i2"),
    ("S1", "S1_P3", "i3"),
    ("S2", "S2_P1", "i4"),
    ("S3", "S3_P1", "i5"),
    ("S3", "S3_P2", "i6"),
    ("S4", "S4_P1", "i7")
],["State", "Person", "Item"])

看起来像这样-

+-----+------+----+
|State|Person|Item|
+-----+------+----+
|   S1| S1_P1|  i1|
|   S1| S1_P2|  i2|
|   S1| S1_P3|  i3|
|   S2| S2_P1|  i4|
|   S3| S3_P1|  i5|
|   S3| S3_P2|  i6|
|   S4| S4_P1|  i7|
+-----+------+----+

现在，我要将其分组为“状态”，将前两个组放入一个新的数据帧（df1）中，将其余所有分组放在另一个新的数据帧（df2）中，这样，

df1（具有前两个组）看起来像-

+-----+------+----+
|State|Person|Item|
+-----+------+----+
|   S1| S1_P1|  i1|
|   S1| S1_P2|  i2|
|   S1| S1_P3|  i3|
|   S2| S2_P1|  i4|
+-----+------+----+

和df2（以及所有其余组）看起来像-

+-----+------+----+
|State|Person|Item|
+-----+------+----+
|   S3| S3_P1|  i5|
|   S3| S3_P2|  i6|
|   S4| S4_P1|  i7|
+-----+------+----+

我该如何实现？在此先感谢：）

Answer 1

除非明确地order，否则没有明确的方式说出火花中的前2个。

first2StateList = df.select("State").distinct().orderBy("State").limit(2).collect() 

//This list will contain [S1,S2] , use it to filter your original df

df1 = df.filter(df.State.isin(first2StateList))
df2 = df.filter(~df.State.isin(first2StateList))

注意：仅在用例列表较小时才使用收集

Answer 2

另一种方法：使用联接

如果您有多个列，或者collect很大时，请首选

val selectedDf = df.select("State").distinct().orderBy("State").limit(2)

df1 = df.join(selectedDf, ['State'],how='inner')
df2 = df.join(selectedDf, ['State'],how='left_anti')

如何根据组将pyspark数据帧分为2个数据帧

2 个答案: