通过基于另一个变量保留顺序来收集列表

问题描述 投票：0回答：10

我正在尝试使用现有列集上的 groupby 聚合在 Pyspark 中创建新的列表列。下面提供了一个示例输入数据框：

------------------------
id | date        | value
------------------------
1  |2014-01-03   | 10 
1  |2014-01-04   | 5
1  |2014-01-05   | 15
1  |2014-01-06   | 20
2  |2014-02-10   | 100   
2  |2014-03-11   | 500
2  |2014-04-15   | 1500

预期输出是：

id | value_list
------------------------
1  | [10, 5, 15, 20]
2  | [100, 500, 1500]

列表中的值按日期排序。

我尝试使用collect_list如下：

from pyspark.sql import functions as F
ordered_df = input_df.orderBy(['id','date'],ascending = True)
grouped_df = ordered_df.groupby("id").agg(F.collect_list("value"))

但是即使我在聚合之前按日期对输入数据帧进行排序，collect_list 也不能保证顺序。

有人可以帮助了解如何通过保留基于第二个（日期）变量的顺序来进行聚合吗？

python apache-spark pyspark
10个回答
139
投票
from pyspark.sql import functions as F
from pyspark.sql import Window

w = Window.partitionBy('id').orderBy('date')

sorted_list_df = input_df.withColumn(
            'sorted_list', F.collect_list('value').over(w)
        )\
        .groupBy('id')\
        .agg(F.max('sorted_list').alias('sorted_list'))


Window
用户提供的示例通常并不能真正解释发生了什么，所以让我为您剖析一下。

如您所知，将 

collect_list
与 
groupBy
一起使用将产生一个 无序 值列表。这是因为，根据数据的分区方式，Spark 一旦在组中找到行，就会将值追加到列表中。然后，顺序取决于 Spark 如何规划执行器上的聚合。

Window
函数允许您控制这种情况，按特定值对行进行分组，以便您可以执行操作
over
每个结果组：

w = Window.partitionBy('id').orderBy('date')

  • partitionBy
    - 您想要具有相同 
    id
    • 的行组/分区 
  • orderBy
    - 您希望组中的每一行按 
    date
    • 排序

一旦定义了窗口的范围 - “具有相同 

id
的行，按 
date
排序”-，您就可以使用它对其执行操作，在本例中为 
collect_list
:

F.collect_list('value').over(w)

此时，您创建了一个新列 

sorted_list
，其中包含按日期排序的有序值列表，但每个 
id
仍然有重复的行。要修剪掉您想要 
groupBy
id
的重复行并保留每组的 
max
值：

.groupBy('id')\
.agg(F.max('sorted_list').alias('sorted_list'))
74
投票

您可以使用

sort_array
功能。如果您将日期和值收集为列表，您可以使用 
sort_array
对结果列进行排序，并仅保留您需要的列。

import pyspark.sql.functions as F

grouped_df = (
    input_df
    .groupby("id")
    .agg(
        F.sort_array(F.collect_list(F.struct("date", "value")))
        .alias("collected_list")
    )
    .withColumn("sorted_list", F.col("collected_list.value"))
    .drop("collected_list")
).show(truncate=False)

+---+----------------+
|id |sorted_list     |
+---+----------------+
|1  |[10, 5, 15, 20] |
|2  |[100, 500, 1500]|
+---+----------------+
51
投票

编辑：PySpark 2.4 中添加了 

pyspark.sql.functions.array_sort
，其运行方式与下面定义的 
sorter
UDF 完全相同，并且通常性能更高。为后代留下旧答案。

对于 PySpark <2.4 only:

如果您将日期和值收集为列表，您可以使用 和 

udf
根据日期对结果列进行排序，然后仅保留结果中的值。

import operator
import pyspark.sql.functions as F

# create list column
grouped_df = input_df.groupby("id") \
               .agg(F.collect_list(F.struct("date", "value")) \
               .alias("list_col"))

# define udf
def sorter(l):
  res = sorted(l, key=operator.itemgetter(0))
  return [item[1] for item in res]

sort_udf = F.udf(sorter)

# test
grouped_df.select("id", sort_udf("list_col") \
  .alias("sorted_list")) \
  .show(truncate = False)
+---+----------------+
|id |sorted_list     |
+---+----------------+
|1  |[10, 5, 15, 20] |
|2  |[100, 500, 1500]|
+---+----------------+
22
投票

这个问题是针对 PySpark 的，但对于 Scala Spark 来说也可能会有所帮助。

让我们准备测试数据框：

import org.apache.spark.sql.functions._
import org.apache.spark.sql.{DataFrame, Row, SparkSession}
import org.apache.spark.sql.expressions.{ Window, UserDefinedFunction}

import java.sql.Date
import java.time.LocalDate

val spark: SparkSession = ...

// Out test data set
val data: Seq[(Int, Date, Int)] = Seq(
  (1, Date.valueOf(LocalDate.parse("2014-01-03")), 10),
  (1, Date.valueOf(LocalDate.parse("2014-01-04")), 5),
  (1, Date.valueOf(LocalDate.parse("2014-01-05")), 15),
  (1, Date.valueOf(LocalDate.parse("2014-01-06")), 20),
  (2, Date.valueOf(LocalDate.parse("2014-02-10")), 100),
  (2, Date.valueOf(LocalDate.parse("2014-02-11")), 500),
  (2, Date.valueOf(LocalDate.parse("2014-02-15")), 1500)
)

// Create dataframe
val df: DataFrame = spark.createDataFrame(data)
  .toDF("id", "date", "value")
df.show()
//+---+----------+-----+
//| id|      date|value|
//+---+----------+-----+
//|  1|2014-01-03|   10|
//|  1|2014-01-04|    5|
//|  1|2014-01-05|   15|
//|  1|2014-01-06|   20|
//|  2|2014-02-10|  100|
//|  2|2014-02-11|  500|
//|  2|2014-02-15| 1500|
//+---+----------+-----+

使用UDF

// Group by id and aggregate date and value to new column date_value
val grouped = df.groupBy(col("id"))
  .agg(collect_list(struct("date", "value")) as "date_value")
grouped.show()
grouped.printSchema()
// +---+--------------------+
// | id|          date_value|
// +---+--------------------+
// |  1|[[2014-01-03,10],...|
// |  2|[[2014-02-10,100]...|
// +---+--------------------+

// udf to extract data from Row, sort by needed column (date) and return value
val sortUdf: UserDefinedFunction = udf((rows: Seq[Row]) => {
  rows.map { case Row(date: Date, value: Int) => (date, value) }
    .sortBy { case (date, value) => date }
    .map { case (date, value) => value }
})

// Select id and value_list
val r1 = grouped.select(col("id"), sortUdf(col("date_value")).alias("value_list"))
r1.show()
// +---+----------------+
// | id|      value_list|
// +---+----------------+
// |  1| [10, 5, 15, 20]|
// |  2|[100, 500, 1500]|
// +---+----------------+

使用窗口

val window = Window.partitionBy(col("id")).orderBy(col("date"))
val sortedDf = df.withColumn("values_sorted_by_date", collect_list("value").over(window))
sortedDf.show()
//+---+----------+-----+---------------------+
//| id|      date|value|values_sorted_by_date|
//+---+----------+-----+---------------------+
//|  1|2014-01-03|   10|                 [10]|
//|  1|2014-01-04|    5|              [10, 5]|
//|  1|2014-01-05|   15|          [10, 5, 15]|
//|  1|2014-01-06|   20|      [10, 5, 15, 20]|
//|  2|2014-02-10|  100|                [100]|
//|  2|2014-02-11|  500|           [100, 500]|
//|  2|2014-02-15| 1500|     [100, 500, 1500]|
//+---+----------+-----+---------------------+

val r2 = sortedDf.groupBy(col("id"))
  .agg(max("values_sorted_by_date").as("value_list")) 
r2.show()
//+---+----------------+
//| id|      value_list|
//+---+----------------+
//|  1| [10, 5, 15, 20]|
//|  2|[100, 500, 1500]|
//+---+----------------+
7
投票

为了确保对每个 id 进行排序，我们可以使用 sortWithinPartitions： 

from pyspark.sql import functions as F
ordered_df = (
    input_df
        .repartition(input_df.id)
        .sortWithinPartitions(['date'])


)
grouped_df = ordered_df.groupby("id").agg(F.collect_list("value"))
5
投票

在 Spark SQL 世界中，答案是：

SELECT 
browser, max(list)
from (
  SELECT
    id,
    COLLECT_LIST(value) OVER (PARTITION BY id ORDER BY date DESC) as list
  FROM browser_count
  GROUP BYid, value, date) 
Group by browser;
2
投票

我尝试了 TMichel 方法，但对我不起作用。当我执行最大聚合时，我没有返回列表的最高值。所以对我有用的是：

def max_n_values(df, key, col_name, number):
    '''
    Returns the max n values of a spark dataframe
    partitioned by the key and ranked by the col_name
    '''
    w2 = Window.partitionBy(key).orderBy(f.col(col_name).desc())
    output = df.select('*',
                       f.row_number().over(w2).alias('rank')).filter(
                           f.col('rank') <= number).drop('rank')
    return output

def col_list(df, key, col_to_collect, name, score):
    w = Window.partitionBy(key).orderBy(f.col(score).desc())

    list_df = df.withColumn(name, f.collect_set(col_to_collect).over(w))
    size_df = list_df.withColumn('size', f.size(name))
    output = max_n_values(df=size_df,
                               key=key,
                               col_name='size',
                               number=1)
    return output
2
投票

从 Spark 2.4 开始，@mtoto 的答案中创建的collect_list(ArrayType) 可以使用 SparkSQL 的内置函数transform和array_sort进行后处理（不需要udf）：

from pyspark.sql.functions import collect_list, expr, struct

df.groupby('id') \
  .agg(collect_list(struct('date','value')).alias('value_list')) \
  .withColumn('value_list', expr('transform(array_sort(value_list), x -> x.value)')) \
  .show()
+---+----------------+
| id|      value_list|
+---+----------------+
|  1| [10, 5, 15, 20]|
|  2|[100, 500, 1500]|
+---+----------------+

注意：如果需要降序排列，请将

array_sort(value_list)
更改为
sort_array(value_list, False)

警告：如果项目（在collect_list中）必须按多个字段（列）以混合顺序排序，则array_sort（）和sort_array（）将不起作用，即

orderBy('col1', desc('col2'))
。

2
投票

如果你想使用 Spark sql，这里是你如何实现这一点。假设表名（或临时视图）是

temp_table
。

select
t1.id,
collect_list(value) as value_list
(Select * from temp_table order by id,date) t1
group by 1
0
投票

补充ShadyStego所说的，我一直在Spark上测试sortWithinPartitions和GroupBy的使用，发现它的性能比Window函数或UDF要好得多。尽管如此，使用此方法时仍存在每个分区一次错误排序的问题，但它可以轻松解决。我在这里展示它Spark (pySpark) groupBy 对collect_list 上的第一个元素进行错误排序。

此方法在大型 DataFrame 上特别有用，但如果驱动程序内存不足，可能需要大量分区。

最新问题
© www.soinside.com 2019 - 2024. All rights reserved.