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

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（前 2 组）看起来像-

+-----+------+----+
|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

另一种方式：使用连接

• 如果您有多个列或当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')

Answer 2

在 spark 中没有明确的方式来表示前 2，除非您明确地 order 它。

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))

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