I made a dataframe in 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"])

```

and it looks this-

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

```

Now I want to group it on "State" put the first two groups into one new dataframe (df1) and all the rest in another new dataframe (df2) such that,

df1 (with first 2 groups) would look like-

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

```

and df2 (with all the remaining groups) would look like-

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

```

How do I achieve this?
Thanks in advance :)