如何在Spark DataFrame的不同长度数组列中随机选取元素？

我来帮你搞定这两个问题，先从你遇到的UDF错误说起，再给你更简洁的解决方案：

一、修复UDF的类型转换错误 #

你遇到的ClassCastException是因为Spark的数组列在传递给Scala UDF时，并不是直接的Array[Int]，而是被包装成了scala.collection.mutable.WrappedArray类型，直接用Array[Int]作为参数会导致类型转换失败。

你只需要把UDF的参数类型改成WrappedArray[Int]或者更通用的Seq[Int]（因为WrappedArray实现了Seq接口），就能解决这个问题：

方案1：使用WrappedArray作为参数 #

import scala.util.Random
import org.apache.spark.sql.functions.udf
import scala.collection.mutable.WrappedArray

// 修改函数参数类型为WrappedArray[Int]
def getRandomElement(arr: WrappedArray[Int]): Int = {
  arr(Random.nextInt(arr.size))
}

val getRandomElementUdf = udf(getRandomElement _)
sampleDf.withColumn("randomItem", getRandomElementUdf('arrays)).show

方案2：使用更通用的Seq[Int] #

如果不想依赖具体的WrappedArray类型，用Seq更灵活：

import scala.util.Random
import org.apache.spark.sql.functions.udf

def getRandomElement(arr: Seq[Int]): Int = {
  arr(Random.nextInt(arr.size))
}

val getRandomElementUdf = udf(getRandomElement _)
sampleDf.withColumn("randomItem", getRandomElementUdf('arrays)).show

二、更高效的非UDF方案（推荐） #

Spark 2.4及以上版本提供了element_at内置函数，它支持用列作为索引来获取数组元素（注意Spark数组是1-based索引，和Scala的0-based不同）。结合你之前生成随机索引的思路，我们可以直接用内置函数实现，不需要UDF，性能更好：

完整代码 #

import org.apache.spark.sql.functions.{size, rand, floor, element_at}

// 直接生成随机元素，不需要中间列
sampleDf.withColumn("chosen_item", 
  element_at('arrays, floor(rand() * size('arrays)) + 1) // 把0-based索引转成1-based
).show

分步解释 #

1. size('arrays)：获取每个数组的长度
2. rand() * size('arrays)：生成0到数组长度之间的随机浮点数
3. floor(...)：把浮点数转成0-based的整数索引
4. +1：转成Spark数组需要的1-based索引
5. element_at('arrays, ...)：用计算出的索引取数组元素

如果你需要保留中间的数组长度和索引列，也可以修改你之前的choice函数：

def choice(df: DataFrame, colName: String): DataFrame = {
  df.withColumn("array_size", size(col(colName)))
    .withColumn("random_idx_0based", floor(rand * 'array_size))
    .withColumn("chosen_item", element_at(col(colName), 'random_idx_0based + 1))
}

choice(sampleDf, "arrays").show

为什么之前的getItem不行？因为getItem是Column类的方法，它只能接受字面量整数作为索引，不能接受列对象，而element_at是Spark提供的高阶函数，支持列作为参数，完美适配你的需求。

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