chenhao-db commented on code in PR #45708:
URL: https://github.com/apache/spark/pull/45708#discussion_r1540169393
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala:
##########
@@ -63,3 +70,300 @@ case class ParseJson(child: Expression) extends
UnaryExpression
override protected def withNewChildInternal(newChild: Expression): ParseJson
=
copy(child = newChild)
}
+
+// A path segment in the `VariantGet` expression. It represents either an
object key access (when
+// `key` is not null) or an array index access (when `key` is null).
+case class PathSegment(key: String, index: Int)
+
+object VariantPathParser extends RegexParsers {
+ private def root: Parser[Char] = '$'
+
+ // Parse index segment like `[123]`.
+ private def index: Parser[PathSegment] =
+ for {
+ index <- '[' ~> "\\d+".r <~ ']'
+ } yield {
+ PathSegment(null, index.toInt)
+ }
+
+ // Parse key segment like `.name`, `['name']`, or `["name"]`.
+ private def key: Parser[PathSegment] =
+ for {
+ key <- '.' ~> "[^\\.\\[]+".r | "['" ~> "[^\\'\\?]+".r <~ "']" |
+ "[\"" ~> "[^\\\"\\?]+".r <~ "\"]"
+ } yield {
+ PathSegment(key, 0)
+ }
+
+ private val parser: Parser[List[PathSegment]] = phrase(root ~> rep(key |
index))
+
+ def parse(str: String): Option[Array[PathSegment]] = {
+ this.parseAll(parser, str) match {
+ case Success(result, _) => Some(result.toArray)
+ case _ => None
+ }
+ }
+}
+
+/**
+ * The implementation for `variant_get` and `try_variant_get` expressions.
Extracts a sub-variant
+ * value according to a path and cast it into a concrete data type.
+ * @param child The source variant value to extract from.
+ * @param path A literal path expression. It has the same format as the JSON
path.
+ * @param schema The target data type to cast into.
+ * @param failOnError Controls whether the expression should throw an
exception or return null if
+ * the cast fails.
+ * @param timeZoneId A string identifier of a time zone. It is required by
timestamp-related casts.
+ */
+case class VariantGet(
+ child: Expression,
+ path: Expression,
+ schema: DataType,
+ failOnError: Boolean,
+ timeZoneId: Option[String] = None)
+ extends BinaryExpression
+ with TimeZoneAwareExpression
+ with NullIntolerant
+ with ExpectsInputTypes
+ with CodegenFallback
+ with QueryErrorsBase {
+ override def checkInputDataTypes(): TypeCheckResult = {
+ val check = super.checkInputDataTypes()
+ if (check.isFailure) {
+ check
+ } else if (!path.foldable) {
+ DataTypeMismatch(
+ errorSubClass = "NON_FOLDABLE_INPUT",
+ messageParameters = Map(
+ "inputName" -> toSQLId("path"),
+ "inputType" -> toSQLType(path.dataType),
+ "inputExpr" -> toSQLExpr(path)
+ )
+ )
+ } else if (!VariantGet.checkDataType(schema)) {
+ DataTypeMismatch(
+ errorSubClass = "CAST_WITHOUT_SUGGESTION",
+ messageParameters = Map(
+ "srcType" -> toSQLType(VariantType),
+ "targetType" -> toSQLType(schema)
+ )
+ )
+ } else {
+ TypeCheckResult.TypeCheckSuccess
+ }
+ }
+
+ override lazy val dataType: DataType = schema.asNullable
+
+ @transient private lazy val parsedPath = {
+ val pathValue = path.eval().toString
+ VariantPathParser.parse(pathValue).getOrElse {
+ throw QueryExecutionErrors.invalidVariantGetPath(pathValue, prettyName)
+ }
+ }
+
+ final override def nodePatternsInternal(): Seq[TreePattern] =
Seq(VARIANT_GET)
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(VariantType, StringType)
+
+ override def prettyName: String = if (failOnError) "variant_get" else
"try_variant_get"
+
+ override def nullable: Boolean = true
+
+ protected override def nullSafeEval(input: Any, path: Any): Any = {
Review Comment:
There is an issue with this approach: we cannot know the return type of
`StaticInvoke(...)` in advance. It is possible that there are two variant rows
with different content types and they can both be cast into the same type.
E.g., both `variant_get(parse_json('1'), '$', 'int')` and
`variant_get(parse_json('"1"'), '$', 'int')` should return `1`. If the
`StaticInvoke(...)` must have a return type, it can only be the variant type.
We still need the current `VariantGet.cast` function as the internal
implementation of `Cast(Variant, targetType)`, so it doesn't really simplify
anything.
--
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]
