HyukjinKwon commented on code in PR #57559:
URL: https://github.com/apache/spark/pull/57559#discussion_r3708784391
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala:
##########
@@ -339,6 +341,286 @@ case class JsonTuple(children: Seq[Expression])
copy(children = newChildren)
}
+/**
+ * The kind of a single `JSON_TABLE` column.
+ */
+sealed trait JsonTableColumnKind
+object JsonTableColumnKind {
+ /** A `FOR ORDINALITY` column: a 1-based sequential row counter. */
+ case object Ordinality extends JsonTableColumnKind
+ /** A regular value column: extracts the value at `path` and casts it to
`dataType`. */
+ case object Value extends JsonTableColumnKind
+ /** An `EXISTS` column: true when `path` matches, cast to `dataType`. */
+ case object Exists extends JsonTableColumnKind
+}
+
+/**
+ * A single column definition of a `JSON_TABLE` invocation.
+ *
+ * @param name the output column name
+ * @param dataType the declared Spark type of the column (LongType for
ORDINALITY columns)
+ * @param path the SQL/JSON path relative to a row item; None for
ORDINALITY columns
+ * @param kind the column kind (ordinality / value / exists)
+ */
+case class JsonTableColumn(
+ name: String,
+ dataType: DataType,
+ path: Option[String],
+ kind: JsonTableColumnKind)
+
+/**
+ * Behavior when the JSON input is malformed.
+ */
+sealed trait JsonTableErrorMode
+object JsonTableErrorMode {
+ /** Produce no rows on malformed input (the SQL-standard default). */
+ case object NullOnError extends JsonTableErrorMode
+ /** Raise an error on malformed input. */
+ case object ErrorOnError extends JsonTableErrorMode
+}
+
+// scalastyle:off line.size.limit
+/**
+ * The SQL:2016 `JSON_TABLE` table-valued function. Shreds a JSON document
into a relational table:
+ * the `rowPath` selects a sequence of row items and each [[JsonTableColumn]]
projects a value out
+ * of each item. Implemented as a [[Generator]] so it plugs into the existing,
well-tested
+ * [[org.apache.spark.sql.catalyst.plans.logical.Generate]] operator; no new
execution operator is
+ * introduced.
+ *
+ * Only the flat (non-`NESTED PATH`) subset of the standard is supported.
Row-source and value
+ * extraction use the token-aware [[JsonTableEvaluator]], which (unlike
`get_json_object`)
+ * distinguishes a missing path from a JSON `null` value; type coercion reuses
[[Cast]].
+ *
+ * {{{
+ * SELECT t.* FROM json_table(
+ * '{"items":[{"id":1,"n":"a"},{"id":2,"n":"b"}]}',
+ * '$.items[*]'
+ * COLUMNS (seq FOR ORDINALITY, id INT PATH '$.id', name STRING PATH '$.n')
+ * ) AS t;
+ * }}}
+ */
+// scalastyle:on line.size.limit
+case class JsonTable(
+ child: Expression,
+ rowPath: String,
+ columns: Seq[JsonTableColumn],
+ errorMode: JsonTableErrorMode,
+ timeZoneId: Option[String] = None,
+ // Captured at plan-construction time so column casts do not change
behavior if the session's
+ // ANSI mode is flipped between building the plan and executing it
(matching `Cast`, which
+ // fixes its eval mode when the expression is constructed).
+ ansiEnabled: Boolean = SQLConf.get.ansiEnabled)
+ extends UnaryExpression
+ with Generator
+ with TimeZoneAwareExpression
+ with CodegenFallback
+ with ImplicitCastInputTypes
+ with QueryErrorsBase {
+
+ // Declared via ImplicitCastInputTypes so the analyzer coerces the JSON
input to STRING. In
+ // particular an untyped SQL NULL (NullType) is cast to STRING rather than
rejected, so
+ // `JSON_TABLE(NULL, ...)` reaches the runtime and applies the NULL ON ERROR
behavior.
+ override def inputTypes: Seq[AbstractDataType] =
+ Seq(StringTypeWithCollation(supportsTrimCollation = true))
+
+ // ORDINALITY columns always hold a non-null counter; value/EXISTS columns
may be null.
+ override def elementSchema: StructType =
+ StructType(columns.map { c =>
+ val nullable = c.kind != JsonTableColumnKind.Ordinality
+ StructField(c.name, c.dataType, nullable = nullable)
+ })
+
+ override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression =
+ copy(timeZoneId = Option(timeZoneId))
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ // First the standard input-type check (STRING for the JSON input, with
NULL coerced).
+ val inputCheck = super.checkInputDataTypes()
+ if (inputCheck.isFailure) {
+ inputCheck
+ } else {
+ // Validate the row path and every column path. A path is valid here iff
it parses and is
+ // free of wildcards -- except the row path may end in a single `[*]`,
which is stripped into
+ // `containerInstructions`, so the row path is checked on that
already-stripped list.
+ val rowPathValid = JsonPathParser.parse(rowPath).isDefined &&
+ !containerInstructions.contains(PathInstruction.Wildcard)
+ val invalid: Option[(String, String)] = if (!rowPathValid) {
+ Some(("row path", rowPath))
+ } else {
+ columns.iterator.collect { case c if c.path.isDefined => (c.name,
c.path.get) }
+ .collectFirst {
+ // Valid column path: parses and is wildcard-free, i.e.
hasWildcard == Some(false).
+ case (name, path) if
!JsonPathParser.hasWildcard(path).contains(false) =>
+ (s"column '$name'", path)
+ }
+ }
+ invalid match {
+ case Some((location, path)) =>
+ DataTypeMismatch(
+ errorSubClass = "INVALID_JSON_TABLE_PATH",
+ messageParameters = Map("location" -> location, "path" ->
toSQLValue(path)))
+ case None =>
+ // Every value/EXISTS column is produced by casting from a source
type (StringType for
+ // value columns, BooleanType for EXISTS columns) to the declared
column type. Reject a
+ // non-castable declared type (e.g. a value column declared
STRUCT/ARRAY/MAP) here rather
+ // than failing at runtime. Ordinality columns are always LongType
and need no check.
+ // The castability rules differ between ANSI and non-ANSI mode (e.g.
BOOLEAN -> TIMESTAMP
+ // is allowed by non-ANSI casts but not ANSI casts), so this must
use the same eval mode
+ // as the actual per-column `Cast` built in `columnCasts`.
+ def sourceType(c: JsonTableColumn): Option[DataType] = c.kind match {
+ case JsonTableColumnKind.Value => Some(StringType)
+ case JsonTableColumnKind.Exists => Some(BooleanType)
+ case JsonTableColumnKind.Ordinality => None
+ }
+ def castable(src: DataType, target: DataType): Boolean =
+ if (ansiEnabled) Cast.canAnsiCast(src, target) else
Cast.canCast(src, target)
+ columns.iterator.flatMap(c => sourceType(c).map((c, _)))
+ .collectFirst { case (c, src) if !castable(src, c.dataType) => (c,
src) } match {
+ case Some((c, srcType)) =>
+ DataTypeMismatch(
+ errorSubClass = "CAST_WITHOUT_SUGGESTION",
+ messageParameters = Map(
+ "srcType" -> toSQLType(srcType),
+ "targetType" -> toSQLType(c.dataType)))
+ case None =>
+ TypeCheckResult.TypeCheckSuccess
+ }
+ }
+ }
+ }
+
+ // The row path is `containerRowPath` plus an optional trailing `[*]`.
Splitting on the parsed
+ // instruction list (rather than the raw string) is whitespace-insensitive
and unambiguous.
+ // `checkInputDataTypes` guarantees the path parses and is wildcard-free at
this point.
+ @transient private lazy val (containerInstructions, explodeRoot)
+ : (Seq[PathInstruction], Boolean) = {
+ val parsed = JsonPathParser.parse(rowPath).getOrElse(Nil)
+ parsed match {
+ case init :+ PathInstruction.Subscript :+ PathInstruction.Wildcard =>
+ (init, true)
+ case other =>
+ (other, false)
+ }
+ }
+
+ @transient private lazy val rowEvaluator: JsonTableEvaluator =
+ JsonTableEvaluator(containerInstructions, explodeRoot)
+
+ // Parsed instruction list per column (empty for ordinality columns, which
have no path).
+ @transient private lazy val columnPaths: Array[Seq[PathInstruction]] =
+ columns.map(c =>
c.path.flatMap(JsonPathParser.parse).getOrElse(Nil)).toArray
+
+ // Prefix trie over the column paths, built once so every row's value/EXISTS
columns are resolved
+ // in a single traversal of the item rather than one re-parse per column.
Ordinality columns have
+ // no path and are excluded (their empty path must not be confused with a
root path `$`, which is
+ // an included column reading the whole item).
+ @transient private lazy val columnTrie: JsonTablePathTrie = {
+ val include = columns.map(_.kind != JsonTableColumnKind.Ordinality).toArray
+ rowEvaluator.buildPathTrie(columnPaths, include)
+ }
+
+ // One reusable Cast per non-ordinality column, evaluated against a
single-slot mutable input
+ // row. Building the Cast once (over a BoundReference) avoids allocating an
expression tree per
+ // row/column on the hot path. The source type is BooleanType for EXISTS,
StringType otherwise.
+ @transient private lazy val columnCasts: Array[Expression] = {
+ val evalMode = EvalMode.fromBoolean(ansiEnabled)
+ columns.map { c =>
+ c.kind match {
+ case JsonTableColumnKind.Ordinality => null
+ case JsonTableColumnKind.Exists =>
+ Cast(BoundReference(0, BooleanType, nullable = false), c.dataType,
timeZoneId, evalMode)
+ case JsonTableColumnKind.Value =>
+ Cast(BoundReference(0, StringType, nullable = true), c.dataType,
timeZoneId, evalMode)
+ }
+ }.toArray
+ }
+
+ // Reusable single-slot input row for the per-column casts above.
+ @transient private lazy val castInput: GenericInternalRow = new
GenericInternalRow(1)
+
+ private def castColumn(i: Int, value: Any): Any = {
+ castInput.update(0, value)
+ columnCasts(i).eval(castInput)
+ }
+
+ private def projectRow(item: UTF8String, ordinal: Long): InternalRow = {
+ // Resolve every value/EXISTS column in a single traversal of the item;
ordinality slots are
+ // not in the trie and come back as Missing (filled below).
+ val resolved = rowEvaluator.navigateColumns(item, columnTrie,
columns.length)
+ val values = new Array[Any](columns.length)
+ var i = 0
+ while (i < columns.length) {
+ values(i) = columns(i).kind match {
Review Comment:
`columns` is a `List` here (`AstBuilder` builds it with `.map(...).toSeq`),
so `columns(i)` is O(i) and this per-row loop is O(n^2) in the column count --
on every emitted row. The parallel per-column data is already arrayed
(`columnPaths`, `columnCasts`, the trie `include`); snapshot the column kinds
the same way and hoist `columns.length` into a local so `projectRow` is O(n)
per row. Wide projections over large arrays are exactly the JSON_TABLE use
case, so it is worth closing before it compounds.
--
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]
