Spark滑动窗口中如何将整行/对象传入自定义UDF？

你遇到的问题是Spark窗口函数无法直接将整行作为参数传给UDF——因为窗口函数的over子句只能配合聚合函数使用，不能直接传递原始行。不过我们可以通过将整行打包为结构体，再用collect_list聚合窗口内的结构体序列来解决这个问题，完美适配你“列增减无需修改代码”的需求。

核心修改步骤 #

1. 打包整行为结构体：用struct($"*")把所有列打包成一个结构体，这样不管后续列怎么增减，这个操作都会自动包含全部字段。
2. 聚合窗口内的结构体序列：用collect_list函数在滑动窗口内收集所有结构体，得到窗口内的行序列。
3. 调整UDF适配输入类型：因为collect_list(struct($"*"))返回的是Seq[Row]，所以UDF需要接收这个类型，并在内部将Row转换为你的Foo对象（或者直接提取字段）。

修改后的完整代码 #

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions.desc
import org.apache.spark.sql.Dataset
import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.{col, udf, struct, collect_list}
import org.apache.spark.sql.Row

case class Foo(a:Int, b:String )
case class FooResult(a1:Int, b1:String, a2:Int, b2:String, a3:Int, b3:String, computation1:Int, computation2:String )

object Main extends App {
  val ss = SparkSession.builder()
    .appName("DataSet Test")
    .master("local[*]").getOrCreate()
  
  import ss.implicits._
  
  val foos = Seq(Foo(-18, "Z"), Foo(-11, "G"), Foo(-8, "A"), Foo(-4, "C"), Foo(-1,"F")).toDS()
  
  // 滑动窗口定义保持不变
  val sliding_window_spec = Window.orderBy(desc("a")).rowsBetween( -2, 0)
  
  // 调整UDF：接收Row序列，转换为Foo后处理
  def slidingUdf = udf((rows: Seq[Row]) => {
    if(rows.size < 3) {
      null
    } else {
      // 将Row转换为Foo对象
      val foo0 = Foo(rows(0).getInt(0), rows(0).getString(1))
      val foo1 = Foo(rows(1).getInt(0), rows(1).getString(1))
      val foo2 = Foo(rows(2).getInt(0), rows(2).getString(1))
      
      FooResult(
        foo0.a, foo0.b, foo1.a, foo1.b, foo2.a, foo2.b,
        foo0.a + foo1.a, foo0.b + foo1.b + foo2.b
      )
    }
  })
  
  // 关键修改：先collect_list整行结构体，再传给UDF
  foos.select(
    slidingUdf(collect_list(struct($"*")).over(sliding_window_spec)).as("test")
  )
  .filter(col("test").isNotNull)
  .select(col("test.*"))
  .show(false)
  
  ss.stop()
}

为什么这个方法有效？ #

• struct($"*")会自动捕获所有列，后续列增减时无需修改这部分代码，完全符合你简化维护的需求。
• collect_list是Spark的内置聚合函数，可以在窗口内收集所有行的结构体，生成一个包含窗口内所有行的序列。
• UDF接收这个序列后，将Row转换为Foo对象（或者直接通过Row.getXXX提取字段），就能实现你原本的业务逻辑。

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