ganeshashree commented on code in PR #57957:
URL: https://github.com/apache/spark/pull/57957#discussion_r3776717833


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala:
##########
@@ -859,6 +859,222 @@ object JsonValue {
   }
 }
 
+/**
+ * Behavior of `JSON_QUERY`'s `ON EMPTY` / `ON ERROR` clause: what to produce 
when the path matches
+ * nothing (`ON EMPTY`) or the input is not valid JSON (`ON ERROR`).
+ */
+sealed trait JsonQueryBehavior
+object JsonQueryBehavior {
+  /** Produce SQL NULL (the SQL-standard default for both clauses). */
+  case object Null extends JsonQueryBehavior
+  /** Raise an error. */
+  case object Error extends JsonQueryBehavior
+  /** Produce an empty JSON array `[]`. */
+  case object EmptyArray extends JsonQueryBehavior
+  /** Produce an empty JSON object `{}`. */
+  case object EmptyObject extends JsonQueryBehavior
+}
+
+/**
+ * The array-wrapper behavior of `JSON_QUERY` (SQL:2016 `... ARRAY WRAPPER`). 
This implementation
+ * resolves a single value per path (wildcard-free paths only), so a wrapper 
wraps that value in a
+ * one-element array:
+ *   - `Without` (default): return the value unwrapped;
+ *   - `Unconditional` (`WITH [UNCONDITIONAL] ARRAY WRAPPER`): always wrap;
+ *   - `Conditional` (`WITH CONDITIONAL ARRAY WRAPPER`): wrap only a scalar; 
leave an object or
+ *     array as is.
+ */
+sealed trait JsonQueryWrapper
+object JsonQueryWrapper {
+  case object Without extends JsonQueryWrapper
+  case object Conditional extends JsonQueryWrapper
+  case object Unconditional extends JsonQueryWrapper
+}
+
+/** The quotes behavior of `JSON_QUERY`: `KEEP QUOTES` (default) or `OMIT 
QUOTES`. */
+sealed trait JsonQueryQuotes
+object JsonQueryQuotes {
+  case object Keep extends JsonQueryQuotes
+  case object Omit extends JsonQueryQuotes
+}
+
+// scalastyle:off line.size.limit
+/**
+ * The SQL:2016 `JSON_QUERY` function (feature T828): extracts the JSON value 
located by `path` from
+ * a JSON input and returns it as JSON text (STRING):
+ *
+ *   - missing path                        -> ON EMPTY behavior
+ *   - malformed / non-single-value input  -> ON ERROR behavior
+ *   - matched object / array / scalar     -> its verbatim JSON text, after 
applying the array
+ *                                            wrapper and quotes clauses
+ *
+ * A matched scalar (including a JSON `null`) is not an error under the 
default `WITHOUT ARRAY
+ * WRAPPER`; it is emitted as JSON text (`JSON_QUERY('{"id":7}', '$.id')` -> 
`7`). `OMIT QUOTES`
+ * strips the surrounding quotes from a scalar string result (and cannot be 
combined with a wrapper).
+ * Both `ON EMPTY` and `ON ERROR` default to NULL per the standard, and a 
`null` JSON input yields
+ * SQL NULL directly. `RETURNING` is restricted to string types here (VARIANT 
is deferred); the
+ * result is always JSON text.
+ *
+ * {{{
+ *   JSON_QUERY('{"a":{"x":1}}', '$.a')                          -- '{"x":1}'
+ *   JSON_QUERY('{"t":["x","y"]}', '$.t')                        -- '["x","y"]'
+ *   JSON_QUERY('{"t":["x","y"]}', '$.t[0]' WITH ARRAY WRAPPER)  -- '["x"]'
+ *   JSON_QUERY('{"n":"Ada"}', '$.n' OMIT QUOTES)                -- 'Ada'
+ * }}}
+ */
+// scalastyle:on line.size.limit
+case class JsonQuery(
+    child: Expression,
+    path: String,
+    returning: DataType,
+    wrapper: JsonQueryWrapper,
+    quotes: JsonQueryQuotes,
+    onEmpty: JsonQueryBehavior,
+    onError: JsonQueryBehavior)
+  extends UnaryExpression
+  with CodegenFallback
+  with ExpectsInputTypes
+  with QueryErrorsBase {
+
+  override def nullable: Boolean = true
+
+  // The JSON input must be a STRING; the result is JSON text.
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(StringTypeWithCollation(supportsTrimCollation = true))
+
+  override def dataType: DataType = returning
+
+  override def checkInputDataTypes(): TypeCheckResult = {
+    val inputCheck = super.checkInputDataTypes()
+    if (inputCheck.isFailure) {
+      inputCheck
+    } else if (!JsonPathParser.hasWildcard(path).contains(false)) {
+      // The path must parse and be wildcard-free (a single value is resolved).
+      DataTypeMismatch(
+        errorSubClass = "INVALID_JSON_PATH",
+        messageParameters = Map(
+          "functionName" -> toSQLId(prettyName), "path" -> toSQLValue(path)))
+    } else if (!JsonQuery.isValidReturningType(returning)) {
+      // RETURNING is restricted to string types (the result is JSON text; 
VARIANT is deferred).
+      DataTypeMismatch(
+        errorSubClass = "INVALID_JSON_QUERY_RETURNING_TYPE",
+        messageParameters = Map(
+          "functionName" -> toSQLId(prettyName), "returningType" -> 
toSQLType(returning)))
+    } else if (quotes == JsonQueryQuotes.Omit && wrapper != 
JsonQueryWrapper.Without) {
+      // OMIT QUOTES applies only to an unwrapped scalar; the SQL standard 
forbids pairing it with
+      // an array wrapper. Enforced here (not only in the parser) so a 
directly-constructed
+      // expression cannot silently ignore the quotes clause.
+      DataTypeMismatch(
+        errorSubClass = "INVALID_JSON_QUERY_WRAPPER_AND_QUOTES",
+        messageParameters = Map("functionName" -> toSQLId(prettyName)))
+    } else {
+      TypeCheckResult.TypeCheckSuccess
+    }
+  }
+
+  // Path parsed once (the grammar makes it a string literal). 
`checkInputDataTypes` guarantees it
+  // parses and is wildcard-free, so the evaluator is only built for a valid 
path.
+  @transient private lazy val evaluator: JsonTableEvaluator =
+    JsonTableEvaluator(JsonPathParser.parse(path).getOrElse(Nil), explodeRoot 
= false)
+
+  // Handle the ON EMPTY / ON ERROR case per the configured behavior.
+  private def onEmptyResult(): Any = behaviorResult(onEmpty, isEmpty = true)
+  private def onErrorResult(): Any = behaviorResult(onError, isEmpty = false)
+
+  private def behaviorResult(behavior: JsonQueryBehavior, isEmpty: Boolean): 
Any = behavior match {
+    case JsonQueryBehavior.Null => null
+    case JsonQueryBehavior.EmptyArray => JsonQuery.EmptyArrayText
+    case JsonQueryBehavior.EmptyObject => JsonQuery.EmptyObjectText
+    case JsonQueryBehavior.Error =>
+      if (isEmpty) throw 
QueryExecutionErrors.jsonQueryOnEmptyError(prettyName, path, cause = null)
+      else throw QueryExecutionErrors.jsonQueryOnErrorError(prettyName, path, 
cause = null)
+  }
+
+  // Apply the array-wrapper and quotes clauses to a matched value's verbatim 
JSON text.
+  // `structural` is true for an object or array match (rather than a scalar, 
incl. JSON null).
+  private def wrapAndQuote(raw: UTF8String, structural: Boolean): UTF8String = 
wrapper match {
+    case JsonQueryWrapper.Without =>
+      // OMIT QUOTES unquotes a scalar string result; a no-op for objects, 
arrays, and non-string
+      // scalars. OMIT QUOTES combined with a wrapper is rejected at parse 
time. `unquotedString`
+      // re-parses the serialized scalar rather than reading the parser's 
`getText` during the
+      // lookup: this keeps the evaluator agnostic to the quotes clause and 
reuses a tested helper,
+      // at the cost of one extra parse of the (small) scalar on the opt-in 
OMIT QUOTES path only --
+      // the common KEEP QUOTES path needs the serialized form and pays 
nothing extra.
+      if (quotes == JsonQueryQuotes.Omit) evaluator.unquotedString(raw) else 
raw

Review Comment:
   Done.



-- 
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]

Reply via email to