Spark DataFrame数组字符串前缀匹配分组统计（无需UDF）

针对包含a、b（字符串数组）、c列的DataFrame，要按a列及数组b中是否存在以xyz开头的字符串分组统计数量，可通过Spark内置函数高效实现，完全无需自定义UDF：

方案一：Spark 3.0+ 推荐用法（array_exists） #

利用Spark 3.0新增的array_exists函数直接判断数组中是否存在符合条件的元素，逻辑简洁且性能最优：

Python 代码示例 #

from pyspark.sql import functions as F

# 新增标识列，标记数组b是否包含以'xyz'开头的元素
df_with_flag = df.withColumn(
    "has_xyz_prefix",
    F.array_exists(F.col("b"), lambda x: F.startswith(x, "xyz"))
)

# 按a列和标识列分组，统计每组数量
result_df = df_with_flag.groupBy("a", "has_xyz_prefix").count()

Scala 代码示例 #

import org.apache.spark.sql.functions._

val dfWithFlag = df.withColumn(
    "has_xyz_prefix",
    array_exists(col("b"), x => startsWith(x, "xyz"))
)

val resultDF = dfWithFlag.groupBy("a", "has_xyz_prefix").count()

方案二：Spark 2.4+ 兼容方案（filter + size） #

如果你的Spark版本低于3.0，可通过filter函数过滤数组中符合条件的元素，再结合size判断是否存在匹配项：

Python 代码示例 #

from pyspark.sql import functions as F

# 过滤数组中以'xyz'开头的元素，判断剩余数组长度是否大于0
df_with_flag = df.withColumn(
    "has_xyz_prefix",
    F.greater(F.size(F.filter(F.col("b"), lambda x: F.startswith(x, "xyz"))), 0)
)

# 分组统计
result_df = df_with_flag.groupBy("a", "has_xyz_prefix").count()

Scala 代码示例 #

import org.apache.spark.sql.functions._

val dfWithFlag = df.withColumn(
    "has_xyz_prefix",
    greater(size(filter(col("b"), x => startsWith(x, "xyz"))), 0)
)

val resultDF = dfWithFlag.groupBy("a", "has_xyz_prefix").count()

性能说明 #

上述方案均使用Spark原生内置函数，这些函数会被Catalyst优化器解析并生成高效的执行计划，避免了UDF带来的序列化/反序列化开销及优化限制，在大数据量场景下性能远优于UDF实现。

内容的提问来源于stack exchange，提问作者Lorenzo Cazador