Spark 之 内置函数 JSON

转自：http://blog.csdn.net/yizishou/article/details/52398665

需求

 

将DataFrame中的StructType类型字段下的所有内容转换为Json字符串。

Spark版本: 1.6.1

 

 

思路

 

• DataFrame有toJSON方法，可将每个Row都转为一个Json字符串，并返回RDD[String]
• DataFrame.write.json方法，可将数据写为Json格式文件

跟踪上述两处代码，发现最终都会调用Spark源码中的org.apache.spark.sql.execution.datasources.json.JacksonGenerator类，使用Jackson，根据传入的StructType、JsonGenerator和InternalRow，生成Json字符串。

 

 

开发

 

我们的函数只需传入一个参数，就是需要转换的列，因此需要实现org.apache.spark.sql.catalyst.expressions包下的UnaryExpression。

后续对功能进行了扩展，不是StructType类型的输入也可以转换。

 

[plain]  view plain  copy

1. package org.apache.spark.sql.catalyst.expressions  
2.   
3.   
4. import java.io.CharArrayWriter  
5.   
6.   
7. import org.apache.spark.sql.catalyst.InternalRow  
8. import org.apache.spark.sql.catalyst.analysis.TypeCheckResult  
9. import org.apache.spark.sql.catalyst.expressions.codegen.CodeGenContext  
10. import org.apache.spark.sql.catalyst.expressions.codegen.GeneratedExpressionCode  
11. import org.apache.spark.sql.execution.datasources.json.JacksonGenerator  
12. import org.apache.spark.sql.types.DataType  
13. import org.apache.spark.sql.types.Metadata  
14. import org.apache.spark.sql.types.StringType  
15. import org.apache.spark.sql.types.StructField  
16. import org.apache.spark.sql.types.StructType  
17.   
18.   
19. import com.fasterxml.jackson.core.JsonFactory  
20. import org.apache.spark.unsafe.types.UTF8String  
21.   
22.   
23. /**  
24.  * 将StructType类型的字段转换为Json String  
25.  * @author yizhu.sun 2016年8月30日  
26.  */  
27. case class Json(child: Expression) extends UnaryExpression with ImplicitCastInputTypes {  
28.   
29.   
30.   override def dataType: DataType = StringType  
31.   override def inputTypes: Seq[DataType] = Seq(child.dataType)  
32.   
33.   
34.   val inputStructType: StructType = child.dataType match {  
35.     case st: StructType => st  
36.     case _ => StructType(Seq(StructField("col", child.dataType, child.nullable, Metadata.empty)))  
37.   }  
38.   
39.   
40.   override def checkInputDataTypes(): TypeCheckResult = TypeCheckResult.TypeCheckSuccess  
41.   
42.   
43.   // 参照 org.apache.spark.sql.DataFrame.toJSON  
44.   // 参照 org.apache.spark.sql.execution.datasources.json.JsonOutputWriter.writeInternal  
45.   protected override def nullSafeEval(data: Any): UTF8String = {  
46.     val writer = new CharArrayWriter  
47.     val gen = new JsonFactory().createGenerator(writer).setRootValueSeparator(null)  
48.     val internalRow = child.dataType match {  
49.       case _: StructType => data.asInstanceOf[InternalRow]  
50.       case _ => InternalRow(data)  
51.     }  
52.     JacksonGenerator(inputStructType, gen)(internalRow)  
53.     gen.flush  
54.     gen.close  
55.     val json = writer.toString  
56.     UTF8String.fromString(  
57.       child.dataType match {  
58.         case _: StructType => json  
59.         case _ => json.substring(json.indexOf(":") + 1, json.lastIndexOf("}"))  
60.       })  
61.   }  
62.   
63.   
64.   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {  
65.     val writer = ctx.freshName("writer")  
66.     val gen = ctx.freshName("gen")  
67.     val st = ctx.freshName("st")  
68.     val json = ctx.freshName("json")  
69.     val typeJson = inputStructType.json  
70.     def getDataExp(data: Any) =  
71.       child.dataType match {  
72.         case _: StructType => s"$data"  
73.         case _ => s"new org.apache.spark.sql.catalyst.expressions.GenericInternalRow(new Object[]{$data})"  
74.       }  
75.     def formatJson(json: String) =  
76.       child.dataType match {  
77.         case _: StructType => s"$json"  
78.         case _ => s"""$json.substring($json.indexOf(":") + 1, $json.lastIndexOf("}"))"""  
79.       }  
80.     nullSafeCodeGen(ctx, ev, (row) => {  
81.       s"""  
82.         | com.fasterxml.jackson.core.JsonGenerator $gen = null;  
83.         | try {  
84.         |   org.apache.spark.sql.types.StructType $st = ${classOf[Json].getName}.getStructType("${typeJson.replace("\"", "\\\"")}");  
85.         |   java.io.CharArrayWriter $writer = new java.io.CharArrayWriter();  
86.         |   $gen = new com.fasterxml.jackson.core.JsonFactory().createGenerator($writer).setRootValueSeparator(null);  
87.         |   org.apache.spark.sql.execution.datasources.json.JacksonGenerator.apply($st, $gen, ${getDataExp(row)});  
88.         |   $gen.flush();  
89.         |   String $json = $writer.toString();  
90.         |   ${ev.value} = UTF8String.fromString(${formatJson(json)});  
91.         | } catch (Exception e) {  
92.         |   ${ev.isNull} = true;  
93.         | } finally {  
94.         |   if ($gen != null) $gen.close();  
95.         | }  
96.        """.stripMargin  
97.     })  
98.   }  
99.   
100.   
101. }  
102.   
103.   
104. object Json {  
105.   
106.   
107.   val structTypeCache = collection.mutable.Map[String, StructType]() // [json, type]  
108.   
109.   
110.   def getStructType(json: String): StructType = {  
111.     structTypeCache.getOrElseUpdate(json, {  
112.       println(">>>>> get StructType from json:")  
113.       println(json)  
114.       DataType.fromJson(json).asInstanceOf[StructType]  
115.     })  
116.   }  
117.   
118.   
119. }  

 

 

注册

 

注意，SQLContext.functionRegistry的可见性为protected[sql]

 

 

[plain]  view plain  copy

1. val (name, (info, builder)) = FunctionRegistry.expression[Json]("json")  
2.   
3. sqlContext.functionRegistry.registerFunction(name, info, builder)  

 

 

 

测试

 

 

[plain]  view plain  copy

1. val subSchema = StructType(Array(  
2.   StructField("a", StringType, true),  
3.   StructField("b", StringType, true),  
4.   StructField("c", IntegerType, true)))  
5.   
6. val schema = StructType(Array(  
7.   StructField("x", subSchema, true)))  
8.   
9. val rdd = sc.makeRDD(Seq(Row(Row("12", null, 123)), Row(Row(null, "2222", null))))  
10.   
11. val df = sqlContext.createDataFrame(rdd, schema)  
12.   
13. df.registerTempTable("df")  
14.   
15. import sqlContext.sql  
16.   
17. sql("select x, x.a from df").show  
18. sql("select x, x.a from df").printSchema  
19. sql("select json(x), json(x.a) from df").show  
20. sql("select json(x), json(x.a) from df").printSchema  

 

结果

[plain]  view plain  copy

1. +----------------+----+  
2. |x               |a   |  
3. +----------------+----+  
4. |[12,null,123]   |12  |  
5. |[null,2222,null]|null|  
6. +----------------+----+  
7.   
8. root  
9.  |-- x: struct (nullable = true)  
10.  |    |-- a: string (nullable = true)  
11.  |    |-- b: string (nullable = true)  
12.  |    |-- c: integer (nullable = true)  
13.  |-- a: string (nullable = true)  
14.   
15. >>>>> get StructType from json:  
16. {"type":"struct","fields":[{"name":"a","type":"string","nullable":true,"metadata":{}},{"name":"b","type":"string","nullable":true,"metadata":{}},{"name":"c","type":"integer","nullable":true,"metadata":{}}]}  
17. >>>>> get StructType from json:  
18. {"type":"struct","fields":[{"name":"col","type":"string","nullable":true,"metadata":{}}]}  
19.   
20. +------------------+----+  
21. |_c0               |_c1 |  
22. +------------------+----+  
23. |{"a":"12","c":123}|"12"|  
24. |{"b":"2222"}      |null|  
25. +------------------+----+  
26.   
27. root  
28.  |-- _c0: string (nullable = true)  
29.  |-- _c1: string (nullable = true)  

 

 

需要注意的点

 

1. 使用SparkSQL自定义函数一般有两种方法，一种是使用开放的api注册简单函数，即调用sqlContext.udf.register方法。另一种就是使用SparkSQL内置函数的注册方法（本例就是使用的这种方法）。前者优势是开发简单，但是实现不了较为复杂的功能，例如本例中需要获取传入的InternalRow的StructType，或者需要实现类似 def fun(arg: Seq[T]): T 这种泛型相关的功能（sqlContext.udf.register的注册方式无法注册返回值为Any的函数）。
2. 本例中实现genCode函数时遇到了困难，即需要在生成的Java代码中构建StructType对象。这个最终通过序列化的思路解决，先使用StructType.json方法将StructType对象序列化为String，然后在Java代码中调用DataType.fromJson反序列化为StructType对象。