Spark使用同一DataFrame执行cogroup报错的解决方案问询

问题场景 #

• 运行环境：Spark 3.3.2
• 操作步骤：用同一个DataFrame两次调用groupByKey生成两个KeyedDataset，执行cogroup操作后，再调用groupByKey和mapGroups
• 报错信息：抛出java.lang.IllegalStateException，提示找不到name字段
• 临时规避方案：复制一份相同的DataFrame分别使用，代码可正常运行

错误复现代码 #

import org.apache.spark.sql.{SparkSession, Dataset}
import org.apache.spark.sql.functions._

case class Person(name: String, age: Int, city: String)

object CogroupBug {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName("CogroupBug")
      .master("local[*]")
      .getOrCreate()
    import spark.implicits._

    val df = Seq(
      Person("Alice", 30, "NY"),
      Person("Bob", 25, "LA"),
      Person("Alice", 35, "NY"),
      Person("Bob", 30, "LA")
    ).toDF()

    val keyed1 = df.groupByKey(_.getAs[String]("name"))
    val keyed2 = df.groupByKey(_.getAs[String]("name"))

    val cogrouped = keyed1.cogroup(keyed2) {
      (key, iter1, iter2) =>
        iter1 ++ iter2
    }

    val result = cogrouped.groupByKey(_.getAs[String]("name"))
      .mapGroups { (key, iter) =>
        (key, iter.map(_.getAs[Int]("age")).sum)
      }

    result.show()
    spark.stop()
  }
}

核心错误栈信息 #

java.lang.IllegalStateException: Could not find field name in Row (Alice,30,NY)
    at org.apache.spark.sql.Row$class.getAs(Row.scala:342)
    at org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema.getAs(rows.scala:278)
    at CogroupBug$$anonfun$2.apply(CogroupBug.scala:23)
    at CogroupBug$$anonfun$2.apply(CogroupBug.scala:23)
    ...

正常运行的复制DataFrame版本代码 #

import org.apache.spark.sql.{SparkSession, Dataset}
import org.apache.spark.sql.functions._

case class Person(name: String, age: Int, city: String)

object CogroupFixed {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName("CogroupFixed")
      .master("local[*]")
      .getOrCreate()
    import spark.implicits._

    val df = Seq(
      Person("Alice", 30, "NY"),
      Person("Bob", 25, "LA"),
      Person("Alice", 35, "NY"),
      Person("Bob", 30, "LA")
    ).toDF()
    // 复制一份DataFrame
    val df2 = df.copy()

    val keyed1 = df.groupByKey(_.getAs[String]("name"))
    val keyed2 = df2.groupByKey(_.getAs[String]("name"))

    val cogrouped = keyed1.cogroup(keyed2) {
      (key, iter1, iter2) =>
        iter1 ++ iter2
    }

    val result = cogrouped.groupByKey(_.getAs[String]("name"))
      .mapGroups { (key, iter) =>
        (key, iter.map(_.getAs[Int]("age")).sum)
      }

    result.show()
    spark.stop()
  }
}

build.sbt配置 #

name := "SparkCogroupBug"
version := "0.1"
scalaVersion := "2.12.15"

libraryDependencies += "org.apache.spark" %% "spark-sql" % "3.3.2" % Provided

无需复制DataFrame的解决方案 #

这个问题是Spark 3.3.2中Dataset API的已知bug，根源是同一个DataFrame生成的两个KeyedDataset在cogroup后，元数据出现冲突，导致后续操作无法正确解析字段。以下是几种无需复制DataFrame的解决方法：

方法1：缓存原始DataFrame #

在生成KeyedDataset前，对原始DataFrame执行cache()，强制Spark生成独立的执行计划，避免元数据冲突：

val df = Seq(
  Person("Alice", 30, "NY"),
  Person("Bob", 25, "LA"),
  Person("Alice", 35, "NY"),
  Person("Bob", 30, "LA")
).toDF().cache() // 添加缓存操作

val keyed1 = df.groupByKey(_.getAs[String]("name"))
val keyed2 = df.groupByKey(_.getAs[String]("name"))

方法2：逻辑复制DataFrame（无额外开销） #

通过select("*")生成一个逻辑上独立的DataFrame，不会触发物理数据复制，只是生成新的逻辑执行计划：

val df = Seq(
  Person("Alice", 30, "NY"),
  Person("Bob", 25, "LA"),
  Person("Alice", 35, "NY"),
  Person("Bob", 30, "LA")
).toDF()

val df2 = df.select("*") // 仅逻辑复制，无数据复制开销

val keyed1 = df.groupByKey(_.getAs[String]("name"))
val keyed2 = df2.groupByKey(_.getAs[String]("name"))

方法3：升级Spark版本 #

该bug在Spark 3.4.0及以上版本已被官方修复，升级到更高版本可以彻底解决问题。

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