PySpark: отфильтровать строки, в которых значение столбца появляется несколько раз в кадре данных

Question

Учитывая фрейм данных PySpark с двумя столбцами, я хочу разделить набор данных на два фрейма данных: один, где комбинация ColA и ColB уникальна, и другой, где он не уникален.

пример:

ColA | ColB | ColCDoesntMatter
1    | 2    | 4
1    | 3    | 3
2    | 2    | 2
1    | 2    | 6
2    | 2    | 8
2    | 3    | 9

Уникальные пары A / B:

1-3
2-3

Неуникальные пары A / B:

1-2
2-2

Таким образом, один кадр данных будет иметьвсе строки с уникальными значениями AB, а другие будут содержать неуникальные значения AB.ColC не имеет значения для фильтра, но должен быть сохранен.Как мне запустить этот фильтр?

Редактировать: обратите внимание, что я не могу использовать dropDuplicates, потому что на самом деле комбинация A и B должна быть уникальной, а не просто A или B.

Answer 1

# count rows per A-B
dfWithRowCount = df.selectExpr("*", "count('*') over (partition by (ColA, ColB)) as nrows")

# filter
uniqueAB = dfWithRowCount.filter(dfWithRowCount.nrows == 1).drop('nrows')
nonUniqueAB = dfWithRowCount.filter(dfWithRowCount.nrows > 1).drop('nrows')

nonUniqueAB.show()
+----+----+----+
|ColA|ColB|ColC|
+----+----+----+
|   2|   2|   2|
|   2|   2|   8|
|   1|   2|   4|
|   1|   2|   6|
+----+----+----+

uniqueAB.show()
+----+----+----+
|ColA|ColB|ColC|
+----+----+----+
|   1|   3|   3|
|   2|   3|   9|
+----+----+----+