cloud-fan commented on code in PR #58005: URL: https://github.com/apache/spark/pull/58005#discussion_r3853615178
########## sql/core/src/test/scala/org/apache/spark/sql/JsonArraySuite.scala: ########## @@ -0,0 +1,644 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql + +import org.apache.spark.SparkRuntimeException +import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch +import org.apache.spark.sql.catalyst.expressions.{Cast, Collate, JsonArray, JsonConstructorNullBehavior, JsonQuery, JsonQueryBehavior, JsonQueryQuotes, JsonQueryWrapper, Literal, ResolvedCollation} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{CharType, GeographyType, GeometryType, IntegerType, MapType, StringType, VarcharType} + +/** + * Test suite for the `JSON_ARRAY` ANSI SQL:2016 constructor function. + */ +class JsonArraySuite extends QueryTest with SharedSparkSession { + + import testImplicits._ + + test("JSON_ARRAY with simple scalar values") { + checkAnswer( + sql("SELECT JSON_ARRAY(1, 'x', true)"), + Row("""[1,"x",true]""")) + } + + test("JSON_ARRAY with NULL elements - ABSENT ON NULL (default)") { + checkAnswer( + sql("SELECT JSON_ARRAY(1, NULL, 3)"), + Row("[1,3]")) + } + + test("JSON_ARRAY with NULL elements - NULL ON NULL") { + checkAnswer( + sql("SELECT JSON_ARRAY(1, NULL, 3 NULL ON NULL)"), + Row("[1,null,3]")) + } + + test("JSON_ARRAY with NULL elements - explicit ABSENT ON NULL") { + // Exercise the explicit `ABSENT ON NULL` grammar branch (the default is implicit absent, so + // this spelling is otherwise untested); it drops NULL elements just like the default. + checkAnswer( + sql("SELECT JSON_ARRAY(1, NULL, 3 ABSENT ON NULL)"), + Row("[1,3]")) + checkAnswer( + sql("SELECT JSON_ARRAY(1, NULL, 3 ABSENT ON NULL RETURNING STRING)"), + Row("[1,3]")) + } + + test("JSON_ARRAY with empty list") { + checkAnswer( + sql("SELECT JSON_ARRAY()"), + Row("[]")) + } + + test("JSON_ARRAY with floating point numbers") { + checkAnswer( + sql("SELECT JSON_ARRAY(1.5, 2.7)"), + Row("[1.5,2.7]")) + } + + test("JSON_ARRAY with mixed types") { + checkAnswer( + sql("SELECT JSON_ARRAY(1, 'text', 3.14, true, false)"), + Row("""[1,"text",3.14,true,false]""")) + } + + test("JSON_ARRAY with all NULLs and ABSENT ON NULL") { + checkAnswer( + sql("SELECT JSON_ARRAY(NULL, NULL)"), + Row("[]")) + } + + test("JSON_ARRAY with RETURNING STRING (explicit)") { + checkAnswer( + sql("SELECT JSON_ARRAY(1, 2, 3 RETURNING STRING)"), + Row("[1,2,3]")) + } + + test("JSON_ARRAY with both NULL ON NULL and RETURNING clauses") { + // The grammar allows `... ON NULL` and `RETURNING` together, in that order; exercise both. + checkAnswer( + sql("SELECT JSON_ARRAY(1, NULL, 3 NULL ON NULL RETURNING STRING)"), + Row("[1,null,3]")) + } + + test("JSON_ARRAY over non-foldable columns exercises row-wise eval") { + val df = Seq((1, "a", true), (2, "b", false)).toDF("i", "s", "b") + checkAnswer( + df.selectExpr("JSON_ARRAY(i, s, b)"), + Seq(Row("""[1,"a",true]"""), Row("""[2,"b",false]"""))) + } + + test("JSON_ARRAY renders decimals and dates via Jackson, not toString") { + checkAnswer( + sql("SELECT JSON_ARRAY(CAST(1.50 AS DECIMAL(5,2)), DATE'2020-01-02')"), + Row("""[1.50,"2020-01-02"]""")) + } + + test("JSON_ARRAY renders a TIMESTAMP via to_json's writer in the session time zone") { + // The constructor is TimeZoneAware and shares to_json's writer, so a TIMESTAMP element must + // render identically to to_json of the singleton array, formatted in the session time zone. + // Assert agreement with that writer (rather than pinning a fragile format string), and that the + // rendering tracks the session time zone by differing between two zones. + def render(tz: String): String = withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> tz) { + val out = + sql("SELECT JSON_ARRAY(TIMESTAMP'2020-01-02 03:04:05')").collect().head.getString(0) + val expected = + sql("SELECT to_json(array(TIMESTAMP'2020-01-02 03:04:05'))").collect().head.getString(0) + assert(out == expected, s"for tz=$tz") + out + } + assert(render("UTC") != render("America/Los_Angeles")) + } + + test("JSON_ARRAY renders array and map elements as JSON structures, like to_json") { + // The docs state array/map/struct arguments render via the same writer as to_json (as nested + // JSON structures, not quoted strings). Cover arrays and maps explicitly (structs are covered + // by the ignoreNullFields test); a nested array element serializes to [1,2], a map to {"k":1}. + checkAnswer( + sql("SELECT JSON_ARRAY(array(1, 2), map('k', 1))"), + Row("""[[1,2],{"k":1}]""")) + checkAnswer( + sql("SELECT JSON_ARRAY(array(array(1), array(2, 3)))"), + Row("[[[1],[2,3]]]")) + } + + test("JSON_ARRAY strings are escaped") { + checkAnswer( + sql("""SELECT JSON_ARRAY('a"b', 'c\td')"""), + Row("""["a\"b","c\td"]""")) + } + + test("nested JSON_ARRAY is spliced raw, not re-quoted (implicit FORMAT JSON)") { + checkAnswer( + sql("SELECT JSON_ARRAY(JSON_ARRAY(1, 2), 3)"), + Row("[[1,2],3]")) + checkAnswer( + sql("SELECT JSON_ARRAY(JSON_ARRAY(1))"), + Row("[[1]]")) + } + + test("explicit FORMAT JSON splices a string verbatim; a plain string is quoted") { + // A plain string element is quoted and escaped like any other string value... + checkAnswer(sql("""SELECT JSON_ARRAY('[1,2]')"""), Row("""["[1,2]"]""")) + // ...while FORMAT JSON marks it as already-JSON text, spliced in verbatim. + checkAnswer(sql("""SELECT JSON_ARRAY('[1,2]' FORMAT JSON)"""), Row("[[1,2]]")) + checkAnswer( + sql("""SELECT JSON_ARRAY('{"a":1}' FORMAT JSON, 'x')"""), + Row("""[{"a":1},"x"]""")) + } + + test("splicing is decided from the source, not the optimized plan shape") { + // A JSON_ARRAY result surfaced as a column is a plain STRING and must be quoted -- even though + // CollapseProject may inline the inner JSON_ARRAY into the outer argument position. The FORMAT + // JSON decision is frozen from the lexical argument at parse time, so it does not depend on + // whether that inlining happens: the result is ["[1]"], never [[1]]. + val inlined = sql("SELECT JSON_ARRAY(a) AS r FROM (SELECT JSON_ARRAY(1) AS a) t") + checkAnswer(inlined, Row("""["[1]"]""")) + // Referencing the alias twice blocks CollapseProject from inlining it; the result is identical, Review Comment: **Non-blocking:** Use a non-foldable producer and assert the optimized plans so this case actually exercises the non-inlined shape. `JSON_ARRAY(1)` is foldable, and `Optimizer.isCheap` treats every foldable expression as cheap, so `CollapseProject` may inline the alias even when it is referenced twice; both cases can therefore cover the same shape. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ApplyDefaultCollation.scala: ########## @@ -440,6 +443,18 @@ object ApplyDefaultCollation extends Rule[LogicalPlan] { private def hasDefaultStringCharOrVarcharType(dataType: DataType): Boolean = dataType.existsRecursively(isDefaultStringCharOrVarcharType) + /** + * A [[StringType]] that was written explicitly -- a distinct instance rather than the default Review Comment: **Non-blocking:** Narrow this identity contract to explicitly collated `STRING`, and make the same correction in `DefaultCollationTypeCoercion.scala:52` and `jsonExpressions.scala:830,1168,1587`. `DataTypeAstBuilder` returns the `StringType` companion for plain `RETURNING STRING`, so an explicitly written plain `STRING` is not a distinct instance and cannot be distinguished from omission by reference identity. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala: ########## @@ -1201,7 +1221,434 @@ object JsonQuery { } /** - * Converts an json input string to a [[StructType]], [[ArrayType]] or [[MapType]] + * Behavior of `JSON_ARRAY`'s `ON NULL` clause: what to do with NULL elements in the array. + */ +sealed trait JsonConstructorNullBehavior +object JsonConstructorNullBehavior { + /** Include NULL elements as JSON `null` values. */ + case object Null extends JsonConstructorNullBehavior + /** Omit NULL elements from the array. */ + case object Absent extends JsonConstructorNullBehavior +} + +/** + * Marker for expressions whose result is JSON text and therefore carry an implicit SQL/JSON + * `FORMAT JSON`: when such an expression appears as an argument of a JSON constructor (e.g. + * `JSON_ARRAY`), its value is spliced in verbatim rather than quoted as a JSON string: + * + * {{{ + * JSON_ARRAY(JSON_ARRAY(1)) -- '[[1]]' (spliced; not the quoted string '["[1]"]') + * }}} + * + * Crucially, the constructor freezes this decision from the *lexical* argument at parse time (see + * `AstBuilder.visitJsonArray`) rather than re-deriving it from the child expression during + * evaluation, so a later optimizer rewrite (e.g. `CollapseProject` inlining a `JSON_ARRAY` alias + * into an argument position) cannot change whether a value is spliced or quoted. `JSON_OBJECT` + * should extend this as it is added. + * + * Most implementers always emit JSON text, so `emitsImplicitJsonText` defaults to true. An + * implementer with a mode that instead emits a plain (non-JSON) string overrides it so that mode + * takes the ordinary-string (quoted) path in a JSON constructor rather than being spliced: + * `JSON_QUERY(... OMIT QUOTES)` returns a matched scalar string's decoded content (e.g. `Ada`, not + * `"Ada"`), which is an ordinary string, so `JsonQuery` returns true only under the default + * `KEEP QUOTES`. + */ +trait ImplicitlyFormattedAsJson extends Expression { + /** Whether this specific instance actually emits JSON text (see the trait doc). */ + def emitsImplicitJsonText: Boolean = true +} + +// scalastyle:off line.size.limit +/** + * The SQL:2016 `JSON_ARRAY` constructor function (feature T811) constructs a JSON array from + * a list of values, with optional `(NULL | ABSENT) ON NULL` control and `RETURNING` type clause. + * + * `NULL ON NULL` (non-default) keeps NULL elements as JSON `null` values. + * `ABSENT ON NULL` (default per the standard) omits NULL elements. + * RETURNING defaults to STRING. + * + * `formatJson(i)` marks element `i` as already-JSON text (SQL/JSON `FORMAT JSON`), so it is spliced + * in verbatim instead of quoted as a string. It is set once, at parse time, from the lexical + * argument -- implicitly for a nested JSON constructor and explicitly for a `... FORMAT JSON` + * clause -- and is a plain field (not a child), so optimizer rewrites that swap the child + * expression (e.g. `CollapseProject`) leave it unchanged. This keeps the output independent of plan + * shape: a value is spliced iff it was written as JSON in the source, never because a rewrite + * happened to substitute a JSON constructor into an argument position. + * + * {{{ + * JSON_ARRAY(1, 'x', true) -- '[1,"x",true]' + * JSON_ARRAY(1, NULL, 3) -- '[1,3]' (ABSENT ON NULL default) + * JSON_ARRAY(1, NULL, 3 NULL ON NULL) -- '[1,null,3]' + * JSON_ARRAY() -- '[]' + * JSON_ARRAY(JSON_ARRAY(1)) -- '[[1]]' (nested constructor: implicit FORMAT JSON) + * JSON_ARRAY('[1,2]') -- '["[1,2]"]' (plain string: quoted) + * JSON_ARRAY('[1,2]' FORMAT JSON) -- '[[1,2]]' (explicit FORMAT JSON: spliced raw) + * }}} + */ +// scalastyle:on line.size.limit +case class JsonArray( + values: Seq[Expression], + formatJson: Seq[Boolean], + needsValidation: Seq[Boolean], + nullBehavior: JsonConstructorNullBehavior, + returning: DataType, + timeZoneId: Option[String] = None) + extends Expression + with TimeZoneAwareExpression + with CodegenFallback + with ExpectsInputTypes + with QueryErrorsBase + with DefaultStringProducingExpression + with ImplicitlyFormattedAsJson { + + // `formatJson(i)` freezes whether element `i` is spliced raw (vs quoted); `needsValidation(i)` + // freezes whether its raw text is arbitrary user input that must be JSON-validated at eval (an + // explicit `FORMAT JSON` on a non-constructor). Both are decided from the lexical argument at + // parse time (see `AstBuilder.visitJsonArray`) and never re-derived from the child expression, + // so an analyzer/optimizer rewrite that wraps a child (e.g. a default-collation `Cast` around a + // trusted nested constructor) cannot flip splicing or spuriously mark a trusted value as needing + // validation. A validated element implies a spliced one. + assert(values.length == formatJson.length && values.length == needsValidation.length, + "JsonArray requires one formatJson and one needsValidation flag per value") + assert(needsValidation.lazyZip(formatJson).forall((nv, fj) => !nv || fj), + "JsonArray needsValidation implies formatJson") + + // True iff some element carries an *explicit* `FORMAT JSON` on a non-constructor: such a value is + // arbitrary user text validated at eval, so it can throw and must not be constant-folded. A + // nested (implicit) constructor produces well-formed JSON by construction and never throws here. + // Read straight off the frozen `needsValidation` flags -- never re-derived from the (possibly + // rewritten) child expressions -- so it drives both `throwable` and the `foldable` exclusion. + private lazy val hasExplicitFormatJson: Boolean = needsValidation.contains(true) + + // Throwable only when the expression can actually throw at eval: when an element carries an + // explicit `FORMAT JSON` (validated, may throw on malformed text) or when a child is itself + // throwable. A plain `JSON_ARRAY(...)` with no explicit `FORMAT JSON` and no throwable children + // cannot throw (RETURNING is restricted to string types at analysis, so the result cast is + // STRING -> STRING), so leaving it non-throwable lets `PushPredicateThroughJoin` push safe + // filters like `JSON_ARRAY(k) = '[42]'` below a join. + override lazy val throwable: Boolean = children.exists(_.throwable) || hasExplicitFormatJson + + // A non-throwing JSON array constructor always yields a value: NULL elements are dropped or + // rendered as JSON `null` (never propagated), the empty argument list yields `[]`, and the + // STRING -> STRING RETURNING cast cannot null a non-null input. For throwable shapes, report + // nullable conservatively so `NullPropagation` does not fold `JSON_ARRAY(...) IS [NOT] NULL` and + // accidentally skip evaluation-time validation or child exceptions. + override def nullable: Boolean = throwable + + // The default RETURNING is a plain STRING, so mix in `DefaultStringProducingExpression` (above) + // to let `ApplyDefaultCollation` cast the result to a non-default object/session collation (e.g. + // `CREATE TABLE ... DEFAULT COLLATION UTF8_LCASE AS SELECT JSON_ARRAY(...)`). The `dataType` + // override below stays authoritative when RETURNING is given explicitly. + + // A constant argument list has no per-row state, so let `ConstantFolding` evaluate the whole + // constructor once instead of serializing JSON row by row. But only fold shapes that cannot throw + // at eval: an explicit `FORMAT JSON` value is validated and may throw on malformed text, and + // `ConstantFolding` evaluates foldables outside conditional branches eagerly -- folding such a + // shape would surface the error at optimization even for rows a later filter/join would drop. A + // nested (implicit FORMAT JSON) value is produced by a constructor and is never malformed, so it + // stays foldable, and its rawness round-trips through `.sql` via an explicit `FORMAT JSON`. + override def foldable: Boolean = children.forall(_.foldable) && !hasExplicitFormatJson + + override def children: Seq[Expression] = values + + override def inputTypes: Seq[AbstractDataType] = values.map(_ => AnyDataType) + + override def dataType: DataType = returning + + override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression = + copy(timeZoneId = Option(timeZoneId)) + + override def checkInputDataTypes(): TypeCheckResult = { + // A constructor emits JSON text, so RETURNING is restricted to string types (VARIANT is a + // deferred extension). CHAR/VARCHAR are normalized to STRING by the parser. + if (!JsonArray.isValidReturningType(returning)) { + DataTypeMismatch( + errorSubClass = "INVALID_JSON_RETURNING_TYPE", + messageParameters = Map( + "functionName" -> toSQLId(prettyName), "returningType" -> toSQLType(returning))) + } else { + // Validate each element up front rather than failing at runtime: a FORMAT JSON element must + // be a string carrying JSON text, and every other element must be serializable to JSON. The + // latter mirrors `to_json`'s analysis-time `JacksonUtils.verifyType` check. + var result: TypeCheckResult = TypeCheckResult.TypeCheckSuccess + var i = 0 + while (i < values.length && result == TypeCheckResult.TypeCheckSuccess) { + val dt = values(i).dataType + if (formatJson(i) && !(dt.isInstanceOf[StringType] || dt == NullType)) { + // A FORMAT JSON element must carry JSON text (string), but an untyped NULL literal is + // allowed: `eval` handles nulls (ABSENT/NULL ON NULL) before it would ever splice, so + // `JSON_ARRAY(NULL FORMAT JSON)` behaves like any other NULL element. + result = DataTypeMismatch( + errorSubClass = "INVALID_JSON_FORMAT_JSON_INPUT", + messageParameters = Map( + "functionName" -> toSQLId(prettyName), + "position" -> (i + 1).toString, + "inputType" -> toSQLType(dt))) + } else { + val elemCheck = JacksonUtils.verifyType(prettyName, dt) + // `verifyType` accepts every `AtomicType`, but `JacksonGenerator` (the writer this shares + // with `to_json`) has no serializer for the spatial atomics and would fail at runtime. + // Reject them here so a passing analysis implies a serializable element. The scan mirrors + // `verifyType`'s traversal (struct fields, array elements, map *values* -- map keys are + // written via `toString`, so a spatial key is fine). + if (elemCheck.isFailure) { + result = elemCheck + } else if (JsonArray.containsUnsupportedJsonType(dt)) { + result = DataTypeMismatch( + errorSubClass = "CANNOT_CONVERT_TO_JSON", + messageParameters = Map( + "name" -> toSQLId(prettyName), + "type" -> toSQLType(dt))) + } + } + i += 1 + } + result + } + } + + // Reuses the mutable `castInput` row and per-element JSON writers, so it holds evaluation state + // and must be fresh-copied before interpreted execution (matches the neighboring JSON + // expressions). + override def stateful: Boolean = true + + // A JSON array is heterogeneous, so each element is serialized with its own data type rather + // than a single shared element type. We build one serializer per child, each configured as a + // single-element `ArrayType(child.dataType)`; serializing `[value]` yields the text `[<frag>]`, + // whose outer brackets we strip to recover the element fragment `<frag>`. This reuses the same + // Jackson generation path as `to_json`, so numbers, decimals, datetimes, and nested structures + // are rendered correctly. + @transient private lazy val resolvedZoneId: String = + timeZoneId.getOrElse(SQLConf.get.sessionLocalTimeZone) + + @transient private lazy val elementEvaluators: Array[StructsToJsonEvaluator] = + new Array[StructsToJsonEvaluator](values.length) + + @transient private lazy val cachedValidatedFormatJsonTexts: Array[String] = + new Array[String](values.length) + + @transient private lazy val singleElem: Array[Any] = new Array[Any](1) + + // Wraps the reused `singleElem` array by reference (GenericArrayData does not copy), so a single + // instance is shared across elements and rows: `appendRenderedElement` mutates `singleElem` in + // place and the serializer reads it synchronously. + @transient private lazy val singleElemArray: GenericArrayData = new GenericArrayData(singleElem) + + @transient private lazy val castInput: GenericInternalRow = new GenericInternalRow(1) + + @transient private lazy val returningCast: Expression = + Cast(BoundReference(0, StringType, nullable = true), returning, timeZoneId, EvalMode.ANSI) + + @transient private lazy val formatJsonFactory: JsonFactory = new JsonFactory() + + // A FORMAT JSON element is spliced into the result verbatim, so it must itself be exactly one + // well-formed JSON value. `checkInputDataTypes` only guarantees the argument is string-typed; a + // string carrying `1,2` or `{bad` would otherwise corrupt the surrounding array into invalid or + // unintended JSON (e.g. `JSON_ARRAY('1,2' FORMAT JSON)` -> `[1,2]`). Validate the runtime value + // before appending: parse one value and require that nothing follows it. + private def validateJsonText(idx: Int, text: String): Unit = { + val valid = + try { + Utils.tryWithResource(formatJsonFactory.createParser(text)) { parser => + if (parser.nextToken() == null) { + false // empty / whitespace-only input carries no JSON value + } else { + // For a scalar this is a no-op; for an array/object it advances to the matching close. + parser.skipChildren() + parser.nextToken() == null // reject anything trailing the first value + } + } + } catch { + case _: JsonProcessingException => false + } + if (!valid) { + throw QueryExecutionErrors.invalidJsonFormatJsonValueError(prettyName, idx + 1, text) + } + } + + // Render a single non-null element to its JSON fragment via its own-typed serializer, appending + // it straight into the result builder. + // TODO(SPARK-58730): this serializes each element as a one-element array and strips the brackets, + // so a row with N values does N Jackson flushes plus an intermediate string allocation each. A + // JSON_ARRAY-specific evaluator that opens the top-level array once and writes each element into + // the shared generator (and codegen for the whole path) would avoid this per-element trip. + private def appendRenderedElement(sb: java.lang.StringBuilder, idx: Int, value: Any): Unit = { + singleElem(0) = value + var evaluator = elementEvaluators(idx) + if (evaluator == null) { + evaluator = StructsToJsonEvaluator( + Map.empty, ArrayType(values(idx).dataType), Some(resolvedZoneId)) + elementEvaluators(idx) = evaluator + } + // Take the Java String directly (via `evaluateString`) rather than `evaluate`'s `UTF8String`, + // which we would immediately decode back to a String here. + val arrJson = evaluator.evaluateString(singleElemArray) + // `arrJson` is "[<frag>]" (compact, no spaces); append just the interior fragment so the + // serialized text is copied once, without materializing a per-element substring. + sb.append(arrJson, 1, arrJson.length - 1) + } + + private def formatJsonText(idx: Int, value: Any): String = { + val text = value.asInstanceOf[UTF8String].toString Review Comment: **Non-blocking:** Check `cachedValidatedFormatJsonTexts(idx)` before converting the `UTF8String` to a Java `String`. Explicit `FORMAT JSON` keeps `JsonArray` non-foldable, so this method runs per row; the current ordering decodes the same foldable value on every row even after the validated `String` is cached. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
