Pyspark Dataframe Pivot и счетчик групп

Question

Я работаю над фреймом данных pyspark, который выглядит следующим образом:

id Категория

1 A

1 A

1 B

2 B

2 A

3 B

3 B

3 B

Я хочу снять столбец категории ипосчитать их появления. Итак, желаемый результат показан ниже

id AB

1 2 1

2 1 1

3 Null 3

Я попытался найти в Интернете что-то, что могло бы помочь мне, но я не смог найти ничего, что могло бы дать мне этот конкретный результат.

Answer 1

Попробуйте это - (Не уверен, что это оптимизировано)

df = spark.createDataFrame([(1,'A'),(1,'A'),(1,'B'),(2,'B'),(2,'A'),(3,'B'),(3,'B'),(3,'B')],['id','category'])
df = df.groupBy('id','category').count()
df.groupBy('id').pivot('category').sum('count').show()

Answer 2

Я не уверен, что результат правильный. Если это так, возможно, вы могли бы описать причину, по которой 2 вместо 3 на B значении

я бы написал так: избегая подробностей в контексте Spark

    import pandas as pd
    from pyspark import SparkContext
    from pyspark.sql import SparkSession
    from pyspark.sql.types import *

    schema = StructType([
        StructField("id", IntegerType()),
        StructField("category", StringType())
    ])

    sample = {'id': [1, 1, 1, 2, 2, 3, 3, 3], 'category': ['A', 'A', 'B', 'B', 'A', 'B', 'B', 'B']}
    df1 = pd.DataFrame(sample)

    sc = SparkContext()
    spark = SparkSession(sparkContext=sc)
    df1_spark = spark.createDataFrame(data=df1, schema=schema)

    # original dataframe
    df1_spark.show()

    df2 = df1_spark.groupBy('id').pivot('category').count().drop('category').sort('id').distinct()

    # resulting dataframe
    df2.show()

вывод 1

    +---+--------+
    | id|category|
    +---+--------+
    |  1|       A|
    |  1|       A|
    |  1|       B|
    |  2|       B|
    |  2|       A|
    |  3|       B|
    |  3|       B|
    |  3|       B|
    +---+--------+

вывод 2

    +---+----+---+
    | id|   A|  B|
    +---+----+---+
    |  1|   2|  1|
    |  2|   1|  1|
    |  3|null|  3|
    +---+----+---+

Надеюсь, это поможет (: