Pyspark: агрегировать данные, проверяя, существует ли значение или нет (не подсчет или сумма)

Question

У меня есть такой набор данных,

test = spark.createDataFrame([
    (0, 1, 5, "2018-06-03", "Region A"),
    (1, 1, 2, "2018-06-04", "Region B"),
    (2, 2, 1, "2018-06-03", "Region B"),
    (4, 1, 1, "2018-06-05", "Region C"),
    (5, 3, 2, "2018-06-03", "Region D"),
    (6, 1, 2, "2018-06-03", "Region A"),
    (7, 4, 4, "2018-06-03", "Region A"),
    (8, 4, 4, "2018-06-03", "Region B"),
    (9, 5, 4, "2018-06-03", "Region A"),
    (10, 5, 4, "2018-06-03", "Region B"),
])\
  .toDF("orderid", "customerid", "price", "transactiondate", "location")
test.show()

И я могу объединить заказы каждого клиента для каждого региона следующим образом:

temp_result = test.groupBy("customerid").pivot("location").agg(count("orderid")).na.fill(0)
temp_result.show()

Теперь, вместо sum или count, я хотел бы просто объединить данные, определив, существует ли значение (то есть 0 или 1), что-то вроде этого

---

Я могу получить вышеуказанный результат по

for field in temp_result.schema.fields:
    if str(field.name) not in ['customerid', "overall_count", "overall_amount"]:
        name = str(field.name)
        temp_result = temp_result.withColumn(name, \
                                             when(col(name) >= 1, 1).otherwise(0))

но есть ли более простой способ получить его?

Answer 1

Вы в основном почти на месте - требуется лишь небольшая настройка, чтобы получить желаемый результат. В вашей агрегации добавьте сравнение счетчиков и преобразуйте логическое значение в целое число (при необходимости вообще):

temp_result = test.groupBy("customerid")\
                  .pivot("location")\
                  .agg((count("orderid")>0).cast("integer"))\
                  .na.fill(0)

temp_result.show()

Результаты в:

+----------+--------+--------+--------+--------+
|customerid|Region A|Region B|Region C|Region D|
+----------+--------+--------+--------+--------+
|         5|       1|       1|       0|       0|
|         1|       1|       1|       1|       0|
|         3|       0|       0|       0|       1|
|         2|       0|       1|       0|       0|
|         4|       1|       1|       0|       0|
+----------+--------+--------+--------+--------+

В случае, если вы получите ошибку искры, вы можете вместо этого использовать это решение, которое выполняет сравнение счетчиков с помощью дополнительного шага:

temp_result = test.groupBy("customerId", "location")\
                  .agg(count("orderid").alias("count"))\
                  .withColumn("count", (col("count")>0).cast("integer"))\
                  .groupby("customerId")\
                  .pivot("location")\
                  .agg(sum("count")).na.fill(0)

temp_result.show()