Ведение только одной записи из спрей-фрейма. Не нужна внятная запись

Question

У меня есть следующие данные:

id  type  price
1    us    1.99
1    ca    2.99
2    ca    1.99
3    au    2.99 
3    us    3.99
3    ca    2.99

По сути, я хочу получить только одну запись на id, а если их больше одной, то я хочу сохранить одну с типом, us поверх других.

Желаемый вывод:

id  type  price
1    us    1.99
2    ca    1.99
3    us    3.99

Как я могу сделать это, используя PySpark или SparkSQL против Spark DataFrame или таблицы?

Answer 1

Можно использовать оконную функцию с пользовательской сортировкой, в Scala ниже, угадайте, можно перевести на Python:

// data
val df = List(
  (1, "us", 1.99),
  (1, "ca", 2.99),
  (2, "ca", 1.99),
  (3, "au", 2.99),
  (3, "us", 3.99),
  (3, "ca", 2.99))
  .toDF("id", "type", "price")

// action
val sorting = when($"type" === "us", 0).otherwise(1)
val window = Window.partitionBy($"id").orderBy(sorting)
val result = df
  .withColumn("sortIndex", row_number().over(window))
  .where($"sortIndex" === 1)
  .drop("sortIndex")

Результат:

+---+----+-----+
|id |type|price|
+---+----+-----+
|1  |us  |1.99 |
|2  |ca  |1.99 |
|3  |us  |3.99 |
+---+----+-----+