sql - 使用复杂类型查询 Spark SQL DataFrame

Question

如何查询具有复杂类型(例如映射/数组)的 RDD？例如，当我编写这个测试代码时:

case class Test(name: String, map: Map[String, String])
val map = Map("hello" -> "world", "hey" -> "there")
val map2 = Map("hello" -> "people", "hey" -> "you")
val rdd = sc.parallelize(Array(Test("first", map), Test("second", map2)))

我认为语法会是这样的:

sqlContext.sql("SELECT * FROM rdd WHERE map.hello = world")

或

sqlContext.sql("SELECT * FROM rdd WHERE map[hello] = world")

但我明白了

Can't access nested field in type MapType(StringType,StringType,true)

和

org.apache.spark.sql.catalyst.errors.package$TreeNodeException: Unresolved attributes

分别。

Answer 1

这取决于列的类型。让我们从一些虚拟数据开始:

import org.apache.spark.sql.functions.{udf, lit}
import scala.util.Try

case class SubRecord(x: Int)
case class ArrayElement(foo: String, bar: Int, vals: Array[Double])
case class Record(
  an_array: Array[Int], a_map: Map[String, String], 
  a_struct: SubRecord, an_array_of_structs: Array[ArrayElement])


val df = sc.parallelize(Seq(
  Record(Array(1, 2, 3), Map("foo" -> "bar"), SubRecord(1),
         Array(
           ArrayElement("foo", 1, Array(1.0, 2.0, 2.0)),
           ArrayElement("bar", 2, Array(3.0, 4.0, 5.0)))),
  Record(Array(4, 5, 6), Map("foz" -> "baz"), SubRecord(2),
         Array(ArrayElement("foz", 3, Array(5.0, 6.0)), 
               ArrayElement("baz", 4, Array(7.0, 8.0))))
)).toDF

df.registerTempTable("df")
df.printSchema

// root
// |-- an_array: array (nullable = true)
// |    |-- element: integer (containsNull = false)
// |-- a_map: map (nullable = true)
// |    |-- key: string
// |    |-- value: string (valueContainsNull = true)
// |-- a_struct: struct (nullable = true)
// |    |-- x: integer (nullable = false)
// |-- an_array_of_structs: array (nullable = true)
// |    |-- element: struct (containsNull = true)
// |    |    |-- foo: string (nullable = true)
// |    |    |-- bar: integer (nullable = false)
// |    |    |-- vals: array (nullable = true)
// |    |    |    |-- element: double (containsNull = false)

• 数组(ArrayType)列:

  • Column.getItem方法

    df.select($"an_array".getItem(1)).show
    
    // +-----------+
    // |an_array[1]|
    // +-----------+
    // |          2|
    // |          5|
    // +-----------+
    

  • Hive 括号语法:

    sqlContext.sql("SELECT an_array[1] FROM df").show
    
    // +---+
    // |_c0|
    // +---+
    // |  2|
    // |  5|
    // +---+
    

  • UDF

    val get_ith = udf((xs: Seq[Int], i: Int) => Try(xs(i)).toOption)
    
    df.select(get_ith($"an_array", lit(1))).show
    
    // +---------------+
    // |UDF(an_array,1)|
    // +---------------+
    // |              2|
    // |              5|
    // +---------------+
    

  • 除了上面列出的方法之外，Spark 支持越来越多的对复杂类型进行操作的内置函数。值得注意的示例包括高阶函数，例如 transform(SQL 2.4+、Scala 3.0+、PySpark/SparkR 3.1+):

    df.selectExpr("transform(an_array, x -> x + 1) an_array_inc").show
    // +------------+
    // |an_array_inc|
    // +------------+
    // |   [2, 3, 4]|
    // |   [5, 6, 7]|
    // +------------+
    
    import org.apache.spark.sql.functions.transform
    
    df.select(transform($"an_array", x => x + 1) as "an_array_inc").show
    // +------------+
    // |an_array_inc|
    // +------------+
    // |   [2, 3, 4]|
    // |   [5, 6, 7]|
    // +------------+
    

  • 过滤器(SQL 2.4+、Scala 3.0+、Python/SparkR 3.1+)

    df.selectExpr("filter(an_array, x -> x % 2 == 0) an_array_even").show
    // +-------------+
    // |an_array_even|
    // +-------------+
    // |          [2]|
    // |       [4, 6]|
    // +-------------+
    
    import org.apache.spark.sql.functions.filter
    
    df.select(filter($"an_array", x => x % 2 === 0) as "an_array_even").show
    // +-------------+
    // |an_array_even|
    // +-------------+
    // |          [2]|
    // |       [4, 6]|
    // +-------------+
    

  • 聚合(SQL 2.4+、Scala 3.0+、PySpark/SparkR 3.1+):

    df.selectExpr("aggregate(an_array, 0, (acc, x) -> acc + x, acc -> acc) an_array_sum").show
    // +------------+
    // |an_array_sum|
    // +------------+
    // |           6|
    // |          15|
    // +------------+
    
    import org.apache.spark.sql.functions.aggregate
    
    df.select(aggregate($"an_array", lit(0), (x, y) => x + y) as "an_array_sum").show
    // +------------+                                                                  
    // |an_array_sum|
    // +------------+
    // |           6|
    // |          15|
    // +------------+
    

  • 数组处理函数 (array_*)，例如 array_distinct (2.4+):

    import org.apache.spark.sql.functions.array_distinct
    
    df.select(array_distinct($"an_array_of_structs.vals"(0))).show
    // +-------------------------------------------+
    // |array_distinct(an_array_of_structs.vals[0])|
    // +-------------------------------------------+
    // |                                 [1.0, 2.0]|
    // |                                 [5.0, 6.0]|
    // +-------------------------------------------+
    

  • array_max(array_min，2.4+):

    import org.apache.spark.sql.functions.array_max
    
    df.select(array_max($"an_array")).show
    // +-------------------+
    // |array_max(an_array)|
    // +-------------------+
    // |                  3|
    // |                  6|
    // +-------------------+
    

  • 展平 (2.4+)

    import org.apache.spark.sql.functions.flatten
    
    df.select(flatten($"an_array_of_structs.vals")).show
    // +---------------------------------+
    // |flatten(an_array_of_structs.vals)|
    // +---------------------------------+
    // |             [1.0, 2.0, 2.0, 3...|
    // |             [5.0, 6.0, 7.0, 8.0]|
    // +---------------------------------+
    

  • arrays_zip(2.4+):

    import org.apache.spark.sql.functions.arrays_zip
    
    df.select(arrays_zip($"an_array_of_structs.vals"(0), $"an_array_of_structs.vals"(1))).show(false)
    // +--------------------------------------------------------------------+
    // |arrays_zip(an_array_of_structs.vals[0], an_array_of_structs.vals[1])|
    // +--------------------------------------------------------------------+
    // |[[1.0, 3.0], [2.0, 4.0], [2.0, 5.0]]                                |
    // |[[5.0, 7.0], [6.0, 8.0]]                                            |
    // +--------------------------------------------------------------------+
    

  • array_union(2.4+):

    import org.apache.spark.sql.functions.array_union
    
    df.select(array_union($"an_array_of_structs.vals"(0), $"an_array_of_structs.vals"(1))).show
    // +---------------------------------------------------------------------+
    // |array_union(an_array_of_structs.vals[0], an_array_of_structs.vals[1])|
    // +---------------------------------------------------------------------+
    // |                                                 [1.0, 2.0, 3.0, 4...|
    // |                                                 [5.0, 6.0, 7.0, 8.0]|
    // +---------------------------------------------------------------------+
    

  • 切片(2.4+):

    import org.apache.spark.sql.functions.slice
    
    df.select(slice($"an_array", 2, 2)).show
    // +---------------------+
    // |slice(an_array, 2, 2)|
    // +---------------------+
    // |               [2, 3]|
    // |               [5, 6]|
    // +---------------------+
    

• map (MapType)列

  • 使用Column.getField方法:

    df.select($"a_map".getField("foo")).show
    
    // +----------+
    // |a_map[foo]|
    // +----------+
    // |       bar|
    // |      null|
    // +----------+
    

  • 使用 Hive 括号语法:

    sqlContext.sql("SELECT a_map['foz'] FROM df").show
    
    // +----+
    // | _c0|
    // +----+
    // |null|
    // | baz|
    // +----+
    

  • 使用带点语法的完整路径:

    df.select($"a_map.foo").show
    
    // +----+
    // | foo|
    // +----+
    // | bar|
    // |null|
    // +----+
    

  • 使用 UDF

    val get_field = udf((kvs: Map[String, String], k: String) => kvs.get(k))
    
    df.select(get_field($"a_map", lit("foo"))).show
    
    // +--------------+
    // |UDF(a_map,foo)|
    // +--------------+
    // |           bar|
    // |          null|
    // +--------------+
    

  • map_* 函数的数量不断增加，例如 map_keys (2.3+)

    import org.apache.spark.sql.functions.map_keys
    
    df.select(map_keys($"a_map")).show
    // +---------------+
    // |map_keys(a_map)|
    // +---------------+
    // |          [foo]|
    // |          [foz]|
    // +---------------+
    

  • 或map_values (2.3+)

    import org.apache.spark.sql.functions.map_values
    
    df.select(map_values($"a_map")).show
    // +-----------------+
    // |map_values(a_map)|
    // +-----------------+
    // |            [bar]|
    // |            [baz]|
    // +-----------------+
    

  请查看SPARK-23899获取详细列表。

• 使用带点语法的完整路径的结构 (StructType) 列:

  • 使用 DataFrame API

    df.select($"a_struct.x").show
    
    // +---+
    // |  x|
    // +---+
    // |  1|
    // |  2|
    // +---+
    

  • 使用原始 SQL

    sqlContext.sql("SELECT a_struct.x FROM df").show
    
    // +---+
    // |  x|
    // +---+
    // |  1|
    // |  2|
    // +---+
    

• 可以使用点语法、名称和标准Column方法访问结构数组内的字段:

  df.select($"an_array_of_structs.foo").show
  
  // +----------+
  // |       foo|
  // +----------+
  // |[foo, bar]|
  // |[foz, baz]|
  // +----------+
  
  sqlContext.sql("SELECT an_array_of_structs[0].foo FROM df").show
  
  // +---+
  // |_c0|
  // +---+
  // |foo|
  // |foz|
  // +---+
  
  df.select($"an_array_of_structs.vals".getItem(1).getItem(1)).show
  
  // +------------------------------+
  // |an_array_of_structs.vals[1][1]|
  // +------------------------------+
  // |                           4.0|
  // |                           8.0|
  // +------------------------------+
  

• 可以使用 UDF 访问用户定义类型 (UDT) 字段。请参阅Spark SQL referencing attributes of UDT了解详情。

注释:

• 根据 Spark 版本，其中一些方法只能在 HiveContext 中使用。 UDF 应独立于标准 SQLContext 和 HiveContext 的版本而工作。

• 一般来说，嵌套值是二等公民。嵌套字段并不支持所有典型操作。根据上下文，最好扁平化架构和/或分解集合

  df.select(explode($"an_array_of_structs")).show
  
  // +--------------------+
  // |                 col|
  // +--------------------+
  // |[foo,1,WrappedArr...|
  // |[bar,2,WrappedArr...|
  // |[foz,3,WrappedArr...|
  // |[baz,4,WrappedArr...|
  // +--------------------+
  

• 点语法可以与通配符 (*) 结合使用来选择(可能是多个)字段，而无需显式指定名称:

  df.select($"a_struct.*").show
  // +---+
  // |  x|
  // +---+
  // |  1|
  // |  2|
  // +---+
  

• 可以使用 get_json_object 和 from_json 函数查询 JSON 列。请参阅How to query JSON data column using Spark DataFrames?了解详情。