Лучшие n элементов по столбцу

Question

Допустим, у меня есть этот фрейм данных Spark:

col1 | col2 | col3 | col4
   a |    g |    h |    p
   r |    i |    h |    l
   f |    j |    z |    d
   a |    j |    m |    l
   f |    g |    h |    q
   f |    z |    z |    a
 ...

Я хочу отключить столбцы и получить массив из первых n элементов по происшествию.Например, при n = 3:

columnName |   content
      col1 | [f, a, r]
      col2 | [g, j, i]
      col3 | [h, z, m]
      col4 | [l, a, d]

Мне удалось объединить имена столбцов в один столбец с помощью этого кода:

columnNames = output_df.columns
output_df = output_df.withColumn("columns", F.array([F.lit(x) for x in columnNames]))

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

Любое предложение?

Спасибо

Answer 1

Я не вижу ничего, кроме как вручную вычисляя все случаи, это не очень эффективно, и я был бы рад услышать о других методах.

Но если вы не боитесь проблем с производительностью, это поможет!

Помните, что я пишу это в scala, я постараюсь перевести это как pyspark, но, как я никогда раньше этого не делал, это будет сложно.

// Let's create a dataframe for reproductibility
val data = Seq(("a", "g", "h", "p"),
("r", "i", "h", "l"),
("f", "j", "z", "d"),
("a", "j", "m", "l"),
("f", "g", "h", "q"),
("f", "z", "z", "a"))

val df = data.toDF("col1", "col2", "col3", "col4")

// Let's add a constant 1, with the groupby sum that will give us the occurencies !
val dfWithFuturOccurences = df.withColumn("futur_occurences", F.lit(1))

// Your n value
val n = 3

// Here goes the magic
df.columns // For each column
    .map(x => 
        (x, dfWithFuturOccurences
            .groupBy(x)
            .agg(sum("futur_occurences").alias("occurences")) // Count occurences here
            .orderBy(desc("occurences"))
            .select(x)
            .limit(n) // Select the top n elements
            .rdd.map(r => r(0).toString).collect().toSeq) //  Collect them and store them as a Seq of string
        )
    .toSeq
    .toDF("col", "top_elements")

В PySpark это, вероятно, что-то вроде этого:

import pyspark.sql.functions as F

data = list(map(lambda x: 
            (x,
            [str(row[x]) for row in 
             dfWithFuturOccurences
            .groupBy(x)
            .agg(F.sum("futur_occurences").alias("occurences"))
            .orderBy(desc("occurences"))
            .select(x)
            .limit(n)
            .collect()]
            )
        , df.columns))

А затем преобразовать ваши данные в информационный фрейм, и все готово!

Выход:

+----+------------+
| col|top_elements|
+----+------------+
|col1|   [f, a, r]|
|col2|   [g, j, z]|
|col3|   [h, z, m]|
|col4|   [l, p, d]|
+----+------------+