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


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala:
##########
@@ -979,6 +979,222 @@ case class JsonExists(
     copy(child = newChild)
 }
 
+/**
+ * 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. `raw` is 
its verbatim JSON text,
+  // `unquoted` is the OMIT QUOTES form (a string's decoded content; `raw` 
otherwise, so OMIT QUOTES
+  // is a no-op for objects, arrays, and non-string scalars), and `structural` 
is true for an object
+  // or array match (rather than a scalar, incl. JSON null).
+  private def wrapAndQuote(raw: UTF8String, unquoted: UTF8String, structural: 
Boolean): UTF8String =
+    wrapper match {
+      case JsonQueryWrapper.Without =>
+        // OMIT QUOTES reuses the string decoded during the lookup rather than 
re-parsing the
+        // serialized fragment; OMIT QUOTES combined with a wrapper is 
rejected at parse time.

Review Comment:
   Done.



##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/json/JsonExpressionEvalUtils.scala:
##########
@@ -550,6 +575,61 @@ case class JsonTableEvaluator(containerPath: 
Seq[PathInstruction], explodeRoot:
     }
   }
 
+  /**
+   * Resolves `containerPath` against a single JSON value for `JSON_QUERY`, 
serializing the matched
+   * value as verbatim JSON text. Returns:
+   *
+   *   - `None` if the input is not a single well-formed JSON value (malformed 
/ trailing garbage /
+   *     empty), which the caller maps to ON ERROR;
+   *   - `Some(Missing)` if the path matches nothing (ON EMPTY);
+   *   - `Some(Found(raw, structural, unquoted))` if the path matches, where 
`raw` is the value's
+   *     verbatim JSON text, `structural` is true for an object or array (as 
opposed to a scalar,
+   *     including a JSON `null`, whose text is `null`), and `unquoted` is the 
`OMIT QUOTES` form
+   *     (a matched JSON string's decoded content; `raw` for every other 
value).
+   *
+   * A `null` input is the caller's responsibility. Like [[lookup]] this 
navigates and validates
+   * with a single parser: after the matched value is serialized (which 
consumes it),
+   * [[drainToRootEnd]]
+   * walks out of the enclosing containers and rejects any trailing content, 
so a valid prefix
+   * followed by garbage is rejected exactly as a fully malformed document is.
+   */
+  final def queryLookup(json: UTF8String): Option[JsonQueryLookup] = {
+    Utils.tryWithResource(CreateJacksonParser.utf8String(jsonFactory, json)) { 
parser =>
+      try {
+        if (parser.nextToken() == null) {
+          None // empty or whitespace-only
+        } else {
+          val result = positionAt(parser, containerPath) match {
+            case PositionResult.Missing => JsonQueryLookup.Missing
+            // A JSON `null` literal is a scalar value for JSON_QUERY: 
serialize it to the text
+            // `null` rather than reporting it specially. The parser is 
positioned on the token.
+            case PositionResult.NullValue =>
+              val raw = serializeCurrentValue(parser)
+              JsonQueryLookup.Found(raw, structural = false, unquoted = raw)
+            case PositionResult.AtValue =>
+              parser.currentToken match {
+                case JsonToken.START_OBJECT | JsonToken.START_ARRAY =>
+                  val raw = serializeCurrentValue(parser)
+                  JsonQueryLookup.Found(raw, structural = true, unquoted = raw)
+                case JsonToken.VALUE_STRING =>
+                  // Capture the decoded string straight from the parser so 
`OMIT QUOTES` need not
+                  // re-parse the serialized (re-quoted) form.
+                  val unquoted = UTF8String.fromString(parser.getText)

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