You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ma...@apache.org on 2024/10/21 18:10:55 UTC
(spark) branch master updated: [SPARK-49967][SQL] Codegen Support for `StructsToJson`(`to_json`)
This is an automated email from the ASF dual-hosted git repository.
maxgekk pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new b4eb03469a90 [SPARK-49967][SQL] Codegen Support for `StructsToJson`(`to_json`)
b4eb03469a90 is described below
commit b4eb03469a90ec5bf65fd51d2bcb41602acad51b
Author: panbingkun <pa...@baidu.com>
AuthorDate: Mon Oct 21 20:10:41 2024 +0200
[SPARK-49967][SQL] Codegen Support for `StructsToJson`(`to_json`)
### What changes were proposed in this pull request?
The pr aims to add `Codegen` Support for `StructsToJson`(`to_json`).
### Why are the changes needed?
- improve codegen coverage.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Pass GA & Existed UT (eg: JsonFunctionsSuite#`*to_json*`)
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #48467 from panbingkun/SPARK-49967.
Authored-by: panbingkun <pa...@baidu.com>
Signed-off-by: Max Gekk <ma...@gmail.com>
---
.../expressions/json/JsonExpressionEvalUtils.scala | 56 ++++++++++++++++++-
.../sql/catalyst/expressions/jsonExpressions.scala | 64 ++++++----------------
.../expressions/ExpressionEvalHelper.scala | 2 +-
.../expressions/JsonExpressionsSuite.scala | 5 +-
.../explain-results/function_to_json.explain | 2 +-
.../query-tests/explain-results/toJSON.explain | 2 +-
.../execution/BaseScriptTransformationExec.scala | 4 +-
7 files changed, 78 insertions(+), 57 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/json/JsonExpressionEvalUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/json/JsonExpressionEvalUtils.scala
index efa5c930b73d..dd7d318f430b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/json/JsonExpressionEvalUtils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/json/JsonExpressionEvalUtils.scala
@@ -16,17 +16,19 @@
*/
package org.apache.spark.sql.catalyst.expressions.json
+import java.io.CharArrayWriter
+
import com.fasterxml.jackson.core.JsonFactory
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.ExprUtils
import org.apache.spark.sql.catalyst.expressions.variant.VariantExpressionEvalUtils
-import org.apache.spark.sql.catalyst.json.{CreateJacksonParser, JacksonParser, JsonInferSchema, JSONOptions}
-import org.apache.spark.sql.catalyst.util.{FailFastMode, FailureSafeParser, PermissiveMode}
+import org.apache.spark.sql.catalyst.json.{CreateJacksonParser, JacksonGenerator, JacksonParser, JsonInferSchema, JSONOptions}
+import org.apache.spark.sql.catalyst.util.{ArrayData, FailFastMode, FailureSafeParser, MapData, PermissiveMode}
import org.apache.spark.sql.errors.QueryCompilationErrors
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructField, StructType, VariantType}
-import org.apache.spark.unsafe.types.UTF8String
+import org.apache.spark.unsafe.types.{UTF8String, VariantVal}
import org.apache.spark.util.Utils
object JsonExpressionEvalUtils {
@@ -111,3 +113,51 @@ case class JsonToStructsEvaluator(
}
}
}
+
+case class StructsToJsonEvaluator(
+ options: Map[String, String],
+ inputSchema: DataType,
+ timeZoneId: Option[String]) extends Serializable {
+
+ @transient
+ private lazy val writer = new CharArrayWriter()
+
+ @transient
+ private lazy val gen = new JacksonGenerator(
+ inputSchema, writer, new JSONOptions(options, timeZoneId.get))
+
+ // This converts rows to the JSON output according to the given schema.
+ @transient
+ private lazy val converter: Any => UTF8String = {
+ def getAndReset(): UTF8String = {
+ gen.flush()
+ val json = writer.toString
+ writer.reset()
+ UTF8String.fromString(json)
+ }
+
+ inputSchema match {
+ case _: StructType =>
+ (row: Any) =>
+ gen.write(row.asInstanceOf[InternalRow])
+ getAndReset()
+ case _: ArrayType =>
+ (arr: Any) =>
+ gen.write(arr.asInstanceOf[ArrayData])
+ getAndReset()
+ case _: MapType =>
+ (map: Any) =>
+ gen.write(map.asInstanceOf[MapData])
+ getAndReset()
+ case _: VariantType =>
+ (v: Any) =>
+ gen.write(v.asInstanceOf[VariantVal])
+ getAndReset()
+ }
+ }
+
+ final def evaluate(value: Any): Any = {
+ if (value == null) return null
+ converter(value)
+ }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala
index d884e76f5256..ac6c233f7d2e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala
@@ -30,16 +30,15 @@ import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch
import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGenerator, CodegenFallback, ExprCode}
import org.apache.spark.sql.catalyst.expressions.codegen.Block.BlockHelper
-import org.apache.spark.sql.catalyst.expressions.json.{JsonExpressionEvalUtils, JsonExpressionUtils, JsonToStructsEvaluator}
+import org.apache.spark.sql.catalyst.expressions.json.{JsonExpressionEvalUtils, JsonExpressionUtils, JsonToStructsEvaluator, StructsToJsonEvaluator}
import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke}
import org.apache.spark.sql.catalyst.json._
import org.apache.spark.sql.catalyst.trees.TreePattern.{JSON_TO_STRUCT, RUNTIME_REPLACEABLE, TreePattern}
-import org.apache.spark.sql.catalyst.util._
import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.internal.types.StringTypeWithCollation
import org.apache.spark.sql.types._
-import org.apache.spark.unsafe.types.{UTF8String, VariantVal}
+import org.apache.spark.unsafe.types.UTF8String
import org.apache.spark.util.Utils
private[this] sealed trait PathInstruction
@@ -748,14 +747,15 @@ case class StructsToJson(
child: Expression,
timeZoneId: Option[String] = None)
extends UnaryExpression
- with TimeZoneAwareExpression
- with CodegenFallback
+ with RuntimeReplaceable
with ExpectsInputTypes
- with NullIntolerant
+ with TimeZoneAwareExpression
with QueryErrorsBase {
override def nullable: Boolean = true
+ override def nodePatternsInternal(): Seq[TreePattern] = Seq(RUNTIME_REPLACEABLE)
+
def this(options: Map[String, String], child: Expression) = this(options, child, None)
// Used in `FunctionRegistry`
@@ -767,44 +767,7 @@ case class StructsToJson(
timeZoneId = None)
@transient
- lazy val writer = new CharArrayWriter()
-
- @transient
- lazy val gen = new JacksonGenerator(
- inputSchema, writer, new JSONOptions(options, timeZoneId.get))
-
- @transient
- lazy val inputSchema = child.dataType
-
- // This converts rows to the JSON output according to the given schema.
- @transient
- lazy val converter: Any => UTF8String = {
- def getAndReset(): UTF8String = {
- gen.flush()
- val json = writer.toString
- writer.reset()
- UTF8String.fromString(json)
- }
-
- inputSchema match {
- case _: StructType =>
- (row: Any) =>
- gen.write(row.asInstanceOf[InternalRow])
- getAndReset()
- case _: ArrayType =>
- (arr: Any) =>
- gen.write(arr.asInstanceOf[ArrayData])
- getAndReset()
- case _: MapType =>
- (map: Any) =>
- gen.write(map.asInstanceOf[MapData])
- getAndReset()
- case _: VariantType =>
- (v: Any) =>
- gen.write(v.asInstanceOf[VariantVal])
- getAndReset()
- }
- }
+ private lazy val inputSchema = child.dataType
override def dataType: DataType = SQLConf.get.defaultStringType
@@ -820,14 +783,23 @@ case class StructsToJson(
override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression =
copy(timeZoneId = Option(timeZoneId))
- override def nullSafeEval(value: Any): Any = converter(value)
-
override def inputTypes: Seq[AbstractDataType] = TypeCollection(ArrayType, StructType) :: Nil
override def prettyName: String = "to_json"
override protected def withNewChildInternal(newChild: Expression): StructsToJson =
copy(child = newChild)
+
+ @transient
+ private lazy val evaluator = StructsToJsonEvaluator(options, inputSchema, timeZoneId)
+
+ override def replacement: Expression = Invoke(
+ Literal.create(evaluator, ObjectType(classOf[StructsToJsonEvaluator])),
+ "evaluate",
+ dataType,
+ Seq(child),
+ Seq(child.dataType)
+ )
}
/**
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionEvalHelper.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionEvalHelper.scala
index 02c7ed727a53..184f5a2a9485 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionEvalHelper.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionEvalHelper.scala
@@ -79,7 +79,7 @@ trait ExpressionEvalHelper extends ScalaCheckDrivenPropertyChecks with PlanTestB
private def prepareEvaluation(expression: Expression): Expression = {
val serializer = new JavaSerializer(new SparkConf()).newInstance()
val resolver = ResolveTimeZone
- val expr = resolver.resolveTimeZones(replace(expression))
+ val expr = replace(resolver.resolveTimeZones(expression))
assert(expr.resolved)
serializer.deserialize(serializer.serialize(expr))
}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/JsonExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/JsonExpressionsSuite.scala
index edb7b93ecdf6..3a58cb92cecf 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/JsonExpressionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/JsonExpressionsSuite.scala
@@ -582,7 +582,7 @@ class JsonExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper with
val schema = StructType(StructField("\"quote", IntegerType) :: Nil)
val struct = Literal.create(create_row(1), schema)
GenerateUnsafeProjection.generate(
- StructsToJson(Map.empty, struct, UTC_OPT) :: Nil)
+ StructsToJson(Map.empty, struct, UTC_OPT).replacement :: Nil)
}
test("to_json - struct") {
@@ -729,8 +729,7 @@ class JsonExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper with
test("from/to json - interval support") {
val schema = StructType(StructField("i", CalendarIntervalType) :: Nil)
checkEvaluation(
- JsonToStructs(schema, Map.empty, Literal.create("""{"i":"1 year 1 day"}""", StringType),
- UTC_OPT),
+ JsonToStructs(schema, Map.empty, Literal.create("""{"i":"1 year 1 day"}""", StringType)),
InternalRow(new CalendarInterval(12, 1, 0)))
Seq(MapType(CalendarIntervalType, IntegerType), MapType(IntegerType, CalendarIntervalType))
diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_to_json.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_to_json.explain
index cd72b12ee19b..da90d9c4c6e1 100644
--- a/sql/connect/common/src/test/resources/query-tests/explain-results/function_to_json.explain
+++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_to_json.explain
@@ -1,2 +1,2 @@
-Project [to_json((timestampFormat,dd/MM/yyyy), d#0, Some(America/Los_Angeles)) AS to_json(d)#0]
+Project [invoke(StructsToJsonEvaluator(Map(timestampFormat -> dd/MM/yyyy),StructType(StructField(id,LongType,true),StructField(a,IntegerType,true),StructField(b,DoubleType,true)),Some(America/Los_Angeles)).evaluate(d#0)) AS to_json(d)#0]
+- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0]
diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/toJSON.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/toJSON.explain
index 1698c562732e..fcb3e173ecaa 100644
--- a/sql/connect/common/src/test/resources/query-tests/explain-results/toJSON.explain
+++ b/sql/connect/common/src/test/resources/query-tests/explain-results/toJSON.explain
@@ -1,2 +1,2 @@
-Project [to_json(struct(id, id#0L, a, a#0, b, b#0, d, d#0, e, e#0, f, f#0, g, g#0), Some(America/Los_Angeles)) AS to_json(struct(id, a, b, d, e, f, g))#0]
+Project [invoke(StructsToJsonEvaluator(Map(),StructType(StructField(id,LongType,true),StructField(a,IntegerType,true),StructField(b,DoubleType,true),StructField(d,StructType(StructField(id,LongType,true),StructField(a,IntegerType,true),StructField(b,DoubleType,true)),true),StructField(e,ArrayType(IntegerType,true),true),StructField(f,MapType(StringType,StructType(StructField(id,LongType,true),StructField(a,IntegerType,true),StructField(b,DoubleType,true)),true),true),StructField(g,String [...]
+- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0]
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala
index 2a1554d287a8..64d2633c3107 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala
@@ -51,8 +51,8 @@ trait BaseScriptTransformationExec extends UnaryExecNode {
child.output.map { in =>
in.dataType match {
case _: ArrayType | _: MapType | _: StructType =>
- new StructsToJson(ioschema.inputSerdeProps.toMap, in)
- .withTimeZone(conf.sessionLocalTimeZone)
+ StructsToJson(ioschema.inputSerdeProps.toMap, in,
+ Some(conf.sessionLocalTimeZone)).replacement
case _ => Cast(in, StringType).withTimeZone(conf.sessionLocalTimeZone)
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org