如何切片和求和数组列的元素

我想使用SparkSQL对数组列进行sum(或执行其他聚合函数)。

我有一个表

+-------+-------+---------------------------------+
|dept_id|dept_nm|                      emp_details|
+-------+-------+---------------------------------+
|     10|Finance|        [100, 200, 300, 400, 500]|
|     20|     IT|                [10, 20, 50, 100]|
+-------+-------+---------------------------------+

我想对emp_details列的值求和。

预期查询:

sqlContext.sql("select sum(emp_details) from mytable").show

预期的结果

1500
180

我也应该能够对范围元素求和，像:

sqlContext.sql("select sum(slice(emp_details,0,3)) from mytable").show

结果

600
80

当按预期对数组类型执行sum时，它显示sum期望参数为数字类型而不是数组类型。

我认为我们需要为此创建UDF。但是怎么做呢?

使用udf会对性能造成影响吗?

Spark 2.4.0

从Spark 2.4开始，Spark SQL支持高阶函数，用于操作复杂的数据结构，包括数组。

"现代"的解决方案如下:

scala> input.show(false)
+-------+-------+-------------------------+
|dept_id|dept_nm|emp_details              |
+-------+-------+-------------------------+
|10     |Finance|[100, 200, 300, 400, 500]|
|20     |IT     |[10, 20, 50, 100]        |
+-------+-------+-------------------------+
input.createOrReplaceTempView("mytable")
val sqlText = "select dept_id, dept_nm, aggregate(emp_details, 0, (acc, value) -> acc + value) as sum from mytable"
scala> sql(sqlText).show
+-------+-------+----+
|dept_id|dept_nm| sum|
+-------+-------+----+
|     10|Finance|1500|
|     20|     IT| 180|
+-------+-------+----+

你可以在以下文章和视频中找到关于高阶函数的很好的阅读材料:

为Apache Spark 2.4中的复杂数据类型引入新的内置高阶函数
在Databricks上使用SQL中的高阶函数处理嵌套数据
Spark SQL高阶函数介绍与Herman van Hovell (Databricks)

Spark 2.3.2及之前版本

免责声明我不推荐这种方法(尽管它得到了最多的好评)，因为Spark SQL执行Dataset.map时做了反序列化。该查询强制Spark对数据进行反序列化并将其加载到JVM(从由JVM外部的Spark管理的内存区域)。这将不可避免地导致更频繁的gc，从而使性能更差。

一个解决方案是使用Dataset解决方案，其中Spark SQL和Scala的组合可以显示其功能。

scala> val inventory = Seq(
     |   (10, "Finance", Seq(100, 200, 300, 400, 500)),
     |   (20, "IT", Seq(10, 20, 50, 100))).toDF("dept_id", "dept_nm", "emp_details")
inventory: org.apache.spark.sql.DataFrame = [dept_id: int, dept_nm: string ... 1 more field]
// I'm too lazy today for a case class
scala> inventory.as[(Long, String, Seq[Int])].
  map { case (deptId, deptName, details) => (deptId, deptName, details.sum) }.
  toDF("dept_id", "dept_nm", "sum").
  show
+-------+-------+----+
|dept_id|dept_nm| sum|
+-------+-------+----+
|     10|Finance|1500|
|     20|     IT| 180|
+-------+-------+----+

我将切片部分作为练习，因为它同样简单。

从Spark 2.4开始，您可以使用slice函数进行切片:

import org.apache.spark.sql.functions.slice
val df = Seq(
  (10, "Finance", Seq(100, 200, 300, 400, 500)),
  (20, "IT", Seq(10, 20, 50, 100))
).toDF("dept_id", "dept_nm", "emp_details")
val dfSliced = df.withColumn(
   "emp_details_sliced",
   slice($"emp_details", 1, 3)
)
dfSliced.show(false)

+-------+-------+-------------------------+------------------+
|dept_id|dept_nm|emp_details              |emp_details_sliced|
+-------+-------+-------------------------+------------------+
|10     |Finance|[100, 200, 300, 400, 500]|[100, 200, 300]   |
|20     |IT     |[10, 20, 50, 100]        |[10, 20, 50]      |
+-------+-------+-------------------------+------------------+

和aggregate:

的求和数组

dfSliced.selectExpr(
  "*", 
  "aggregate(emp_details, 0, (x, y) -> x + y) as details_sum",  
  "aggregate(emp_details_sliced, 0, (x, y) -> x + y) as details_sliced_sum"
).show

+-------+-------+--------------------+------------------+-----------+------------------+
|dept_id|dept_nm|         emp_details|emp_details_sliced|details_sum|details_sliced_sum|
+-------+-------+--------------------+------------------+-----------+------------------+
|     10|Finance|[100, 200, 300, 4...|   [100, 200, 300]|       1500|               600|
|     20|     IT|   [10, 20, 50, 100]|      [10, 20, 50]|        180|                80|
+-------+-------+--------------------+------------------+-----------+------------------+

一种可能的方法是在Array列上使用explode()，从而按唯一键聚合输出。例如:

import sqlContext.implicits._
import org.apache.spark.sql.functions._
(mytable
  .withColumn("emp_sum",
    explode($"emp_details"))
  .groupBy("dept_nm")
  .agg(sum("emp_sum")).show)
+-------+------------+
|dept_nm|sum(emp_sum)|
+-------+------------+
|Finance|        1500|
|     IT|         180|
+-------+------------+

要在数组中只选择特定的值，我们可以使用链接问题的答案，并通过稍微修改应用它:

val slice = udf((array : Seq[Int], from : Int, to : Int) => array.slice(from,to))
(mytable
  .withColumn("slice", 
    slice($"emp_details", 
      lit(0), 
      lit(3)))
  .withColumn("emp_sum",
    explode($"slice"))
  .groupBy("dept_nm")
  .agg(sum("emp_sum")).show)
+-------+------------+
|dept_nm|sum(emp_sum)|
+-------+------------+
|Finance|         600|
|     IT|          80|
+-------+------------+

val data = Seq((10, "Finance", Array(100,200,300,400,500)),
               (20, "IT", Array(10,20,50,100)))
val mytable = sc.parallelize(data).toDF("dept_id", "dept_nm","emp_details")

这是一个不使用groupBy的mtoto答案的替代方案(我真的不知道哪个是最快的:UDF, mtoto解决方案还是我的，欢迎评论)

一般来说，使用UDF会对性能产生影响。有一个答案，您可能想要阅读，这个资源是一个很好的阅读UDF。

对于您的问题，您可以避免使用UDF。我将使用的是用Scala逻辑生成的Column表达式。

数据:

val df = Seq((10, "Finance", Array(100,200,300,400,500)),
                  (20, "IT", Array(10,  20, 50,100)))
          .toDF("dept_id", "dept_nm","emp_details")

您需要一些技巧才能遍历ArrayType，您可以稍微使用解决方案来发现各种问题(参见底部的slice部分编辑)。这是我的建议，但你可能会发现更好的。首先取最大长度

val maxLength = df.select(size('emp_details).as("l")).groupBy().max("l").first.getInt(0)

然后使用它，测试是否有更短的数组

val sumArray = (1 until maxLength)
      .map(i => when(size('emp_details) > i,'emp_details(i)).otherwise(lit(0)))
      .reduce(_ + _)
      .as("sumArray")
val res = df
  .select('dept_id,'dept_nm,'emp_details,sumArray)

结果:

+-------+-------+--------------------+--------+
|dept_id|dept_nm|         emp_details|sumArray|
+-------+-------+--------------------+--------+
|     10|Finance|[100, 200, 300, 4...|    1500|
|     20|     IT|   [10, 20, 50, 100]|     180|
+-------+-------+--------------------+--------+

我建议你看看sumArray，了解它在做什么。

编辑:当然我只看了问题的一半…但是，如果你想改变求和的项目，你可以看到这个解决方案变得很明显(即你不需要切片函数)，只需改变(0 until maxLength)与你需要的索引范围:

def sumArray(from: Int, max: Int) = (from until max)
      .map(i => when(size('emp_details) > i,'emp_details(i)).otherwise(lit(0)))
      .reduce(_ + _)
      .as("sumArray")

基于zero323的出色答案;如果你有一个长整数数组，即BIGINT，你需要将初始值从0更改为BIGINT(0)，如第一段所述这里有

dfSliced.selectExpr(
  "*", 
  "aggregate(emp_details, BIGINT(0), (x, y) -> x + y) as details_sum",  
  "aggregate(emp_details_sliced, BIGINT(0), (x, y) -> x + y) as details_sliced_sum"
).show

缺少rdd方式，所以让我添加它。

val df = Seq((10, "Finance", Array(100,200,300,400,500)),(20, "IT", Array(10,20,50,100))).toDF("dept_id", "dept_nm","emp_details")
import scala.collection.mutable._
val rdd1 = df.rdd.map( x=> {val p = x.getAs[mutable.WrappedArray[Int]]("emp_details").toArray; Row.merge(x,Row(p.sum,p.slice(0,2).sum)) })
spark.createDataFrame(rdd1,df.schema.add(StructField("sumArray",IntegerType)).add(StructField("sliceArray",IntegerType))).show(false)

输出:

+-------+-------+-------------------------+--------+----------+
|dept_id|dept_nm|emp_details              |sumArray|sliceArray|
+-------+-------+-------------------------+--------+----------+
|10     |Finance|[100, 200, 300, 400, 500]|1500    |300       |
|20     |IT     |[10, 20, 50, 100]        |180     |30        |
+-------+-------+-------------------------+--------+----------+

Spark 2.4.0

Spark 2.3.2及之前版本

相关内容

最新更新

热门标签：