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

Question

我在pyspark中做了一个dataframe-

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

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

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

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 个数据帧

How to split a pyspark dataframe into 2 dataframe on the basis of groups

dataframe

apache-spark

pyspark

pyspark-dataframes