Spark Sql 复杂类型高阶函数

spark sql 2.4 新增了高阶函数功能，允许在数组类型中像 scala/python 一样使用高阶函数

背景

复杂类型的数据和真实数据模型相像，但是使用sql操作较为困难，一般需要借助于 explod/collect_list 等方法，或者使用 scala / python 编写UDF，但是对每个方法都要定义并且注册，较为繁琐，其中 python udf 的性能由于需要在 JVM 和 Python 进程中进行序列化，效率更低。

例如现在有这样一种需求，对 t1 表中某个 array 字段 values 的每个元素加1

1. 只使用 sql 实现

此类方法会带来 shuffle 的开销，collect_list 也不能保证数据的顺序，同时要保证 group 字段全局唯一，否则结果会出错。

SELECT key,
	values,
	collect_list(value + 1) AS values_plus_one
FROM t1
	LATERAL VIEW explode(values) T AS value
GROUP BY key,values

2. 使用 udf 的方式

使用 scala 定义 udf

def addOneToElements(elements: Seq[Int]) = elements.map(element => element + 1)
spark.udf.register("plusOneInt", addOneToElements(_:Seq[Int]):Seq[Int])

或者使用 python 定义 udf

from pyspark.sql.types import IntegerType
from pyspark.sql.types import ArrayType

def add_one_to_els(elements):
		return [el + 1 for el in elements]
spark.udf.register("plusOneInt", add_one_to_els, ArrayType(IntegerType()))

在 sql 中使用 udf

SELECT key,
	values,
	plusOneInt(values) AS values_plus_one,
FROM t1

3. 使用高阶函数的方式

SELECT key,
	values,
	TRANSFORM(values, value -> value + 1) AS values_plus_one
FROM t1

三种方式的性能对比图:

使用

Array 高阶函数

目前支持 transform / filter / exists / agregate / zip_with 方法

id	arr_values	nested_values
1	[1,2,3]	[[1,2],[3,4]]

1. transform

对一个数组应用 function 产生另一个数组

如果 lambda function 中有两个参数，第一个参数为数组中的元素，第二个参数代表该元素的索引(从0开始)

-- array 中每个元素加1
SELECT id,transform(arr_values, value -> value + 1) as arr_values;
> [2,3,4]
-- array 中每个元素加上索引大小
SELECT id,transform(arr_values, (value,i) -> value + i) as arr_values;
> [1,3,5]
-- 处理嵌套类型的数组，对每个数组元素中的值加一
select transform(nested_values,value -> transform(value ,e -> e + 1)) as nested_values;
> [[2, 3], [4, 5]]
-- 使用其他列的值，每个元素加上 id 值
select transform(arr_values,arrs -> transform(arrs ,e -> e + id)) as arr_values;
> [2,3,4]

2. filter

过滤出数组中符合条件的元素

SELECT filter(arr_values, e -> e % 2 == 1) as arr_values;
> [1,3]

3. exists

数组中的一个或多个元素是否满足条件

SELECT exists(arr_values, e ->  e % 2 == 0) as arr_values;
> true

4. aggregate

给定初始值，并对数组中所有的元素都应用 function ，如果需要的话还可以加上 finish function

-- aggregate(expr,start,merge) ，将所有的元素相加
SELECT aggregate(arr_values, 0, (acc, e) -> acc + e) as summed_values;
> 6
-- aggregate(expr,start,merge,finish)，先将所有的元素相加，最后 * 10
SELECT aggregate(arr_values, 0, (acc, e) -> acc + e, acc -> acc * 10) as aggregated_values;
> 60

5. zip_with

将两个数组根据 function 合并为一个数组，较短的那个数组会以填充 null 的方式匹配较长的数组

-- 每个元素进行换位
SELECT zip_with(array(1, 2, 3), array('a', 'b', 'c'), (x, y) -> (y, x))
> [[a, 1], [b, 2], [c, 3]]
-- 两个数值型数组相加
SELECT zip_with(array(1, 2), array(3, 4), (x, y) -> x + y);
> [4,6]
-- 两个字符串数组相加
SELECT zip_with(array('a', 'b', 'c'), array('d', 'e', 'f'), (x, y) -> concat(x, y));
> ["ad","be","cf"]

复杂类型内置函数

spark 2.4 增加了大量的内置函数

• array: array_distinct / array_intersect / array_union / array_except 等
• map: map_form_arrays / map_from_entries / map_concat
• array & map : element_at / cardinality

具体可参考 spark sql api

总结

1. spark sql 高阶函数可以避免用户维护大量的 udf ，且提高了性能，增强了复杂类型的处理能力。
2. collect_list / collect_set 返回的结构为 array ，可以直接使用高阶函数进行操作。
3. spark 2.4 高阶函数目前只支持 Array，在后续的 spark 3.0 中会增加 Map 等类型的支持 Built-in SQL Function Improvement

参考

spark-2.4-functions

Apache Spark 2.4 新增内置函数和高阶函数使用介绍

An Introduction to Higher Order Functions in Spark SQL with Herman van Hovell