web-dev-qa-db-ja.com

Spark StructType / RowのUDF

sparkサブフィールドとして配列と文字列を持つデータフレームに「StructType」列があります。配列を変更し、同じタイプの新しい列を返します。 UDFで処理しますか、それとも代替手段は何ですか?

import org.Apache.spark.sql.types._
import org.Apache.spark.sql.Row
val sub_schema = StructType(StructField("col1",ArrayType(IntegerType,false),true) :: StructField("col2",StringType,true)::Nil)
val schema = StructType(StructField("subtable", sub_schema,true) :: Nil)
val data = Seq(Row(Row(Array(1,2),"eb")),  Row(Row(Array(3,2,1), "dsf")) )
val rd = sc.parallelize(data)
val df = spark.createDataFrame(rd, schema)
df.printSchema

root
 |-- subtable: struct (nullable = true)
 |    |-- col1: array (nullable = true)
 |    |    |-- element: integer (containsNull = false)
 |    |-- col2: string (nullable = true)

Row型のUDFが必要なようです。

val u =  udf((x:Row) => x)
       >> Schema for type org.Apache.spark.sql.Row is not supported

Sparkは戻り値の型のスキーマを認識していないため、残念ながら、udf.registerも失敗します。

spark.udf.register("foo", (x:Row)=> Row, sub_schema)
     <console>:30: error: overloaded method value register with alternatives: ...
16
Danil Kirsanov

結果スキーマを2番目のUDFパラメーターとして渡すことができます。

val u =  udf((x:Row) => x, sub_schema)
15
Danil Kirsanov

はい、UDFでこれを行うことができます。簡単にするために、ケースクラスの例を取り上げ、すべての値に2を追加して配列を変更しました。

case class Root(subtable: Subtable)
case class Subtable(col1: Seq[Int], col2: String)

val df = spark.createDataFrame(Seq(
  Root(Subtable(Seq(1, 2, 3), "toto")),
  Root(Subtable(Seq(10, 20, 30), "tata"))
))

val myUdf = udf((subtable: Row) =>
  Subtable(subtable.getSeq[Int](0).map(_ + 2), subtable.getString(1))
)
val result = df.withColumn("subtable_new", myUdf(df("subtable")))
result.printSchema()
result.show(false)

印刷されます:

root
 |-- subtable: struct (nullable = true)
 |    |-- col1: array (nullable = true)
 |    |    |-- element: integer (containsNull = false)
 |    |-- col2: string (nullable = true)
 |-- subtable_new: struct (nullable = true)
 |    |-- col1: array (nullable = true)
 |    |    |-- element: integer (containsNull = false)
 |    |-- col2: string (nullable = true)

+-------------------------------+-------------------------------+
|subtable                       |subtable_new                   |
+-------------------------------+-------------------------------+
|[WrappedArray(1, 2, 3),toto]   |[WrappedArray(3, 4, 5),toto]   |
|[WrappedArray(10, 20, 30),tata]|[WrappedArray(12, 22, 32),tata]|
+-------------------------------+-------------------------------+
5
L. CWI

あなたは正しい軌道に乗っています。このシナリオでは、UDFはあなたの人生を楽にします。すでに遭遇したように、UDFはsparkが知らない型を返すことができません。したがって、基本的にはsparkが簡単にシリアル化できるものを返す必要があります。 case classであるか、(Seq[Int], String)のようなタプルを返すことができるため、コードの修正バージョンを以下に示します。

def main(args: Array[String]): Unit = {
  import org.Apache.spark.sql.Row
  import org.Apache.spark.sql.functions._
  import org.Apache.spark.sql.types._
  val sub_schema = StructType(StructField("col1", ArrayType(IntegerType, false), true) :: StructField("col2", StringType, true) :: Nil)
  val schema = StructType(StructField("subtable", sub_schema, true) :: Nil)
  val data = Seq(Row(Row(Array(1, 2), "eb")), Row(Row(Array(3, 2, 1), "dsf")))
  val rd = spark.sparkContext.parallelize(data)
  val df = spark.createDataFrame(rd, schema)

  df.printSchema()
  df.show(false)

  val mapArray = (subRows: Row) => {
    // I prefer reading values from row by specifying column names, you may use index also
    val col1 = subRows.getAs[Seq[Int]]("col1")
    val mappedCol1 = col1.map(x => x * x) // Use map based on your requirements
    (mappedCol1, subRows.getAs[String]("col2")) // now mapping is done for col2
  }
  val mapUdf = udf(mapArray)

  val newDf = df.withColumn("col1_mapped", mapUdf(df("subtable")))
  newDf.show(false)
  newDf.printSchema()
}

これらのリンクをご覧ください。より多くの洞察が得られる場合があります。

  1. 複雑なスキーマの操作に関する最も包括的な答え: https://stackoverflow.com/a/33850490/4046067
  2. Sparkがサポートするデータ型: https://spark.Apache.org/docs/latest/sql-programming-guide.html#data-types
4
Tawkir