This is an automated email from the ASF dual-hosted git repository.
uros-b pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 409b37efc3ab [SPARK-59823][SQL] Deduplicate variant manipulation
expressions
409b37efc3ab is described below
commit 409b37efc3abaf390a772f225d6442240892233e
Author: bojana-db <[email protected]>
AuthorDate: Tue Sep 29 20:20:44 2026 +0000
[SPARK-59823][SQL] Deduplicate variant manipulation expressions
### What changes were proposed in this pull request?
This PR consolidates parsed-path caching, value-type validation, and
nullable code-generation handling shared by variant manipulation expressions.
### Why are the changes needed?
The refactor removes duplicated implementations while preserving existing
behavior and making the manipulation expressions easier to maintain.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
`VariantExpressionSuite`: five focused manipulation-expression tests passed
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (Astra)
Closes #59073 from bojana-db/refactor-variant-manipulation.
Authored-by: bojana-db <[email protected]>
Signed-off-by: Uros Bojanic <[email protected]>
---
.../expressions/variant/variantExpressions.scala | 340 +++++++--------------
.../variant/VariantExpressionSuite.scala | 2 +-
2 files changed, 118 insertions(+), 224 deletions(-)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala
index 0affceb09f1e..da0c45a981c8 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala
@@ -391,6 +391,89 @@ object VariantPathParser extends RegexParsers {
}
}
+sealed trait VariantPathArg
+
+case class FoldableVariantPath(segments: Array[VariantPathSegment], pathStr:
String)
+ extends VariantPathArg {
+ // Retain the original path for insert/set/append error messages.
`VariantBuilder.PathSegment`
+ // is not `Serializable`, so the Java form is `@transient` and re-derived
once per executor task
+ // after deserialization.
+ @transient lazy val javaSegments: Array[VariantBuilder.PathSegment] =
+ VariantExpressionEvalUtils.toJavaSegments(segments)
+}
+
+case class DynamicVariantPath(expr: Expression) extends VariantPathArg
+
+object VariantPathArg {
+ private def parseFoldableExpression(
+ child: Expression,
+ functionName: String,
+ allowRoot: Boolean): Option[FoldableVariantPath] = {
+ val value = child.eval()
+ if (value == null) {
+ None
+ } else {
+ val pathStr = value.asInstanceOf[UTF8String].toString
+ Some(FoldableVariantPath(
+ VariantExpressionEvalUtils.parseVariantPath(pathStr, functionName,
allowRoot), pathStr))
+ }
+ }
+
+ def fromExpression(
+ child: Expression,
+ functionName: String,
+ allowRoot: Boolean = false): Option[VariantPathArg] = {
+ if (child.foldable) {
+ parseFoldableExpression(child, functionName, allowRoot)
+ } else {
+ Some(DynamicVariantPath(child))
+ }
+ }
+
+ def foldableFromExpression(
+ child: Expression,
+ functionName: String,
+ allowRoot: Boolean = false): Option[FoldableVariantPath] = {
+ if (child.foldable) {
+ parseFoldableExpression(child, functionName, allowRoot)
+ } else {
+ None
+ }
+ }
+}
+
+object VariantManipulationExpressionUtils extends QueryErrorsBase {
+ // Callers cache this result and force parsing before their NULL
short-circuit.
+ // A dynamic path or constant NULL has no cached parsed path.
+ def foldablePath(
+ path: Expression,
+ functionName: String,
+ allowRoot: Boolean = false): Option[FoldableVariantPath] =
+ VariantPathArg.foldableFromExpression(path, functionName, allowRoot)
+
+ def checkValueType(dataType: DataType): TypeCheckResult = {
+ if (dataType == NullType || VariantGet.checkDataType(dataType,
allowStructsAndMaps = false)) {
+ TypeCheckResult.TypeCheckSuccess
+ } else {
+ DataTypeMismatch(
+ errorSubClass = "CAST_WITHOUT_SUGGESTION",
+ messageParameters =
+ Map("srcType" -> toSQLType(dataType), "targetType" ->
toSQLType(VariantType)))
+ }
+ }
+
+ def assignVariantResult(ev: ExprCode, call: String, nullable: Boolean):
String = {
+ if (nullable) {
+ s"""
+ |${ev.value} = $call;
+ |${ev.isNull} = ${ev.value} == null;
+ """.stripMargin
+ } else {
+ s"${ev.value} = $call;"
+ }
+ }
+}
+
/**
* 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.
@@ -912,8 +995,8 @@ case class VariantDelete(children: Seq[Expression])
private def variantChild: Expression = children.head
private def pathChildren: Seq[Expression] = children.tail
- @transient private lazy val pathArgs: Seq[VariantDelete.DeletePathArg] =
- pathChildren.flatMap(VariantDelete.toPathArg)
+ @transient private lazy val pathArgs: Seq[VariantPathArg] =
+ pathChildren.flatMap(VariantPathArg.fromExpression(_, prettyName))
override def eval(input: InternalRow): Any = {
val inputVariant = variantChild.eval(input).asInstanceOf[VariantVal]
@@ -923,9 +1006,9 @@ case class VariantDelete(children: Seq[Expression])
var i = 0
while (i < args.length) {
args(i) match {
- case parsed: VariantDelete.ParsedDeletePath =>
+ case parsed: FoldableVariantPath =>
current = VariantExpressionEvalUtils.deleteAtPath(current,
parsed.javaSegments)
- case VariantDelete.DynamicDeletePath(expr) =>
+ case DynamicVariantPath(expr) =>
val pathVal = expr.eval(input).asInstanceOf[UTF8String]
if (pathVal != null) {
current = VariantExpressionEvalUtils.deleteAtPath(current, pathVal)
@@ -943,10 +1026,10 @@ case class VariantDelete(children: Seq[Expression])
val current = ctx.freshName("vdCurrent")
val perPath = pathArgs.map {
- case parsed: VariantDelete.ParsedDeletePath =>
+ case parsed: FoldableVariantPath =>
val parsedArg = ctx.addReferenceObj("vdParsed", parsed)
s"$current = $cls.deleteAtPath($current, $parsedArg.javaSegments());"
- case VariantDelete.DynamicDeletePath(expr) =>
+ case DynamicVariantPath(expr) =>
val pCode = expr.genCode(ctx)
s"""
|${pCode.code}
@@ -975,31 +1058,6 @@ case class VariantDelete(children: Seq[Expression])
newChildren: IndexedSeq[Expression]): VariantDelete = copy(children =
newChildren)
}
-object VariantDelete {
- sealed trait DeletePathArg
- case class ParsedDeletePath(segments: Array[VariantPathSegment]) extends
DeletePathArg {
- // `VariantBuilder.PathSegment` is not `Serializable`, so the cached Java
form is
- // `@transient` and re-initialized once per executor task after
deserialization.
- @transient lazy val javaSegments: Array[VariantBuilder.PathSegment] =
- VariantExpressionEvalUtils.toJavaSegments(segments)
- }
- case class DynamicDeletePath(expr: Expression) extends DeletePathArg
-
- private[variant] def toPathArg(child: Expression): Option[DeletePathArg] = {
- if (child.foldable) {
- val v = child.eval()
- if (v == null) {
- None
- } else {
- Some(ParsedDeletePath(VariantExpressionEvalUtils.parseVariantPath(
- v.asInstanceOf[UTF8String].toString, "variant_delete")))
- }
- } else {
- Some(DynamicDeletePath(child))
- }
- }
-}
-
// scalastyle:off line.size.limit
@ExpressionDescription(
usage = "_FUNC_(v, path1[, path2, ...]) - Keeps only the fields or array
elements of a variant " +
@@ -1060,16 +1118,16 @@ case class VariantPick(children: Seq[Expression])
private def variantChild: Expression = children.head
private def pathChildren: Seq[Expression] = children.tail
- @transient private lazy val pathArgs: Seq[VariantPick.PickPathArg] =
- pathChildren.flatMap(VariantPick.toPathArg)
+ @transient private lazy val pathArgs: Seq[VariantPathArg] =
+ pathChildren.flatMap(VariantPathArg.fromExpression(_, prettyName,
allowRoot = true))
// When every path is constant, the keep-tree is the same for all rows, so
build it once here and
// reuse it rather than rebuilding it per row. `None` means at least one
path is dynamic.
@transient private lazy val foldableTree: Option[VariantBuilder.PickNode] = {
- if (pathArgs.forall(_.isInstanceOf[VariantPick.ParsedPickPath])) {
+ if (pathArgs.forall(_.isInstanceOf[FoldableVariantPath])) {
val paths = new
java.util.ArrayList[Array[VariantBuilder.PathSegment]](pathArgs.length)
pathArgs.foreach {
- case parsed: VariantPick.ParsedPickPath =>
paths.add(parsed.javaSegments)
+ case parsed: FoldableVariantPath => paths.add(parsed.javaSegments)
case _ =>
}
Some(VariantBuilder.buildPickTree(paths))
@@ -1085,9 +1143,9 @@ case class VariantPick(children: Seq[Expression])
var i = 0
while (i < pathArgs.length) {
pathArgs(i) match {
- case parsed: VariantPick.ParsedPickPath =>
+ case parsed: FoldableVariantPath =>
paths.add(parsed.javaSegments)
- case VariantPick.DynamicPickPath(expr) =>
+ case DynamicVariantPath(expr) =>
val pathVal = expr.eval(input).asInstanceOf[UTF8String]
if (pathVal != null) {
paths.add(VariantExpressionEvalUtils.parsePickPath(pathVal))
@@ -1125,10 +1183,10 @@ case class VariantPick(children: Seq[Expression])
// At least one dynamic path: gather the paths per row, then project.
val paths = ctx.freshName("vpPaths")
val addPaths = pathArgs.map {
- case parsed: VariantPick.ParsedPickPath =>
+ case parsed: FoldableVariantPath =>
val parsedArg = ctx.addReferenceObj("vpParsed", parsed)
s"$paths.add($parsedArg.javaSegments());"
- case VariantPick.DynamicPickPath(expr) =>
+ case DynamicVariantPath(expr) =>
val pCode = expr.genCode(ctx)
s"""
|${pCode.code}
@@ -1161,39 +1219,13 @@ case class VariantPick(children: Seq[Expression])
newChildren: IndexedSeq[Expression]): VariantPick = copy(children =
newChildren)
}
-object VariantPick {
- sealed trait PickPathArg
- case class ParsedPickPath(segments: Array[VariantPathSegment]) extends
PickPathArg {
- // `VariantBuilder.PathSegment` is not `Serializable`, so the cached Java
form is
- // `@transient` and re-initialized once per executor task after
deserialization.
- @transient lazy val javaSegments: Array[VariantBuilder.PathSegment] =
- VariantExpressionEvalUtils.toJavaSegments(segments)
- }
- case class DynamicPickPath(expr: Expression) extends PickPathArg
-
- private[variant] def toPathArg(child: Expression): Option[PickPathArg] = {
- if (child.foldable) {
- val v = child.eval()
- if (v == null) {
- None
- } else {
- Some(ParsedPickPath(VariantExpressionEvalUtils.parseVariantPath(
- v.asInstanceOf[UTF8String].toString, "variant_pick", allowRoot =
true)))
- }
- } else {
- Some(DynamicPickPath(child))
- }
- }
-}
-
case class VariantInsert(
input: Expression,
path: Expression,
value: Expression,
failOnError: Boolean = true)
extends TernaryExpression
- with ExpectsInputTypes
- with QueryErrorsBase {
+ with ExpectsInputTypes {
override def first: Expression = input
override def second: Expression = path
@@ -1210,36 +1242,13 @@ case class VariantInsert(
val result = super.checkInputDataTypes()
if (result.isFailure) {
result
- } else if (value.dataType == NullType) {
- TypeCheckResult.TypeCheckSuccess
- } else if (!VariantGet.checkDataType(value.dataType, allowStructsAndMaps =
false)) {
- DataTypeMismatch(
- errorSubClass = "CAST_WITHOUT_SUGGESTION",
- messageParameters =
- Map("srcType" -> toSQLType(value.dataType), "targetType" ->
toSQLType(VariantType)))
} else {
- TypeCheckResult.TypeCheckSuccess
+ VariantManipulationExpressionUtils.checkValueType(value.dataType)
}
}
- // When the path is a foldable expression, parse it once at planning time
and cache it. The Java
- // segments are derived once per task (see `ParsedInsertPath`), avoiding a
per-row conversion.
- // `None` means the path is dynamic (or a foldable NULL, which makes the
whole expression NULL and
- // is never evaluated).
- @transient private lazy val foldablePath:
Option[VariantInsert.ParsedInsertPath] = {
- if (path.foldable) {
- val p = path.eval()
- if (p == null) {
- None
- } else {
- val s = p.asInstanceOf[UTF8String].toString
- Some(VariantInsert.ParsedInsertPath(
- VariantExpressionEvalUtils.parseVariantPath(s, prettyName), s))
- }
- } else {
- None
- }
- }
+ @transient private lazy val foldablePath: Option[FoldableVariantPath] =
+ VariantManipulationExpressionUtils.foldablePath(path, prettyName)
override def eval(input: InternalRow): Any = {
val _ = foldablePath
@@ -1272,22 +1281,9 @@ case class VariantInsert(
case None =>
s"""$cls.insertAtPath($vVal, $pVal, $valVal, $fromArg, "$prettyName",
$failOnError)"""
}
- if (failOnError) {
- nullSafeCodeGen(ctx, ev,
- (vVal, pVal, valVal) => s"${ev.value} = ${call(vVal, pVal, valVal)};")
- } else {
- val resultType = CodeGenerator.javaType(VariantType)
- val tmp = ctx.freshName("insertResult")
- nullSafeCodeGen(ctx, ev, (vVal, pVal, valVal) =>
- s"""
- |$resultType $tmp = ${call(vVal, pVal, valVal)};
- |if ($tmp == null) {
- | ${ev.isNull} = true;
- |} else {
- | ${ev.value} = $tmp;
- |}
- """.stripMargin)
- }
+ nullSafeCodeGen(ctx, ev, (vVal, pVal, valVal) =>
+ VariantManipulationExpressionUtils.assignVariantResult(
+ ev, call(vVal, pVal, valVal), nullable))
}
override def prettyName: String = if (failOnError) "variant_insert" else
"try_variant_insert"
@@ -1297,16 +1293,6 @@ case class VariantInsert(
copy(input = newFirst, path = newSecond, value = newThird)
}
-object VariantInsert {
- // Caches a foldable path. `VariantBuilder.PathSegment` is not
`Serializable`, so the Java form is
- // `@transient` and re-derived once per executor task after deserialization.
`pathStr` is the
- // source string, retained for error messages.
- case class ParsedInsertPath(segments: Array[VariantPathSegment], pathStr:
String) {
- @transient lazy val javaSegments: Array[VariantBuilder.PathSegment] =
- VariantExpressionEvalUtils.toJavaSegments(segments)
- }
-}
-
abstract class VariantInsertExpressionBuilderBase(failOnError: Boolean)
extends ExpressionBuilder {
override def build(funcName: String, expressions: Seq[Expression]):
Expression = {
val numArgs = expressions.length
@@ -1431,36 +1417,13 @@ case class VariantSet(
"inputName" -> toSQLId("create_if_missing"),
"inputType" -> toSQLType(createIfMissing.dataType),
"inputExpr" -> toSQLExpr(createIfMissing)))
- } else if (value.dataType == NullType) {
- TypeCheckResult.TypeCheckSuccess
- } else if (!VariantGet.checkDataType(value.dataType, allowStructsAndMaps =
false)) {
- DataTypeMismatch(
- errorSubClass = "CAST_WITHOUT_SUGGESTION",
- messageParameters =
- Map("srcType" -> toSQLType(value.dataType), "targetType" ->
toSQLType(VariantType)))
} else {
- TypeCheckResult.TypeCheckSuccess
+ VariantManipulationExpressionUtils.checkValueType(value.dataType)
}
}
- // When the path is a foldable expression, parse it once at planning time
and cache it. The Java
- // segments are derived once per task (see `ParsedSetPath`), avoiding a
per-row conversion. `None`
- // means the path is dynamic (or a foldable NULL, which makes the whole
expression NULL and is
- // never evaluated).
- @transient private lazy val foldablePath: Option[VariantSet.ParsedSetPath] =
{
- if (path.foldable) {
- val p = path.eval()
- if (p == null) {
- None
- } else {
- val s = p.asInstanceOf[UTF8String].toString
- Some(VariantSet.ParsedSetPath(
- VariantExpressionEvalUtils.parseVariantPath(s, prettyName), s))
- }
- } else {
- None
- }
- }
+ @transient private lazy val foldablePath: Option[FoldableVariantPath] =
+ VariantManipulationExpressionUtils.foldablePath(path, prettyName)
override def eval(input: InternalRow): Any = {
val _ = foldablePath
@@ -1496,23 +1459,10 @@ case class VariantSet(
s"""$cls.setAtPath(
| $vVal, $pVal, $valVal, $fromArg, $createVal, "$prettyName",
$failOnError)"""
.stripMargin
- }
- if (failOnError) {
- nullSafeCodeGen(ctx, ev,
- (vVal, pVal, valVal, createVal) => s"${ev.value} = ${call(vVal, pVal,
valVal, createVal)};")
- } else {
- val resultType = CodeGenerator.javaType(VariantType)
- val tmp = ctx.freshName("setResult")
- nullSafeCodeGen(ctx, ev, (vVal, pVal, valVal, createVal) =>
- s"""
- |$resultType $tmp = ${call(vVal, pVal, valVal, createVal)};
- |if ($tmp == null) {
- | ${ev.isNull} = true;
- |} else {
- | ${ev.value} = $tmp;
- |}
- """.stripMargin)
}
+ nullSafeCodeGen(ctx, ev, (vVal, pVal, valVal, createVal) =>
+ VariantManipulationExpressionUtils.assignVariantResult(
+ ev, call(vVal, pVal, valVal, createVal), nullable))
}
override def prettyName: String = if (failOnError) "variant_set" else
"try_variant_set"
@@ -1525,16 +1475,6 @@ case class VariantSet(
copy(input = newFirst, path = newSecond, value = newThird, createIfMissing
= newFourth)
}
-object VariantSet {
- // Caches a foldable path. `VariantBuilder.PathSegment` is not
`Serializable`, so the Java form is
- // `@transient` and re-derived once per executor task after deserialization.
`pathStr` is the
- // source string, retained for error messages.
- case class ParsedSetPath(segments: Array[VariantPathSegment], pathStr:
String) {
- @transient lazy val javaSegments: Array[VariantBuilder.PathSegment] =
- VariantExpressionEvalUtils.toJavaSegments(segments)
- }
-}
-
abstract class VariantSetExpressionBuilderBase(failOnError: Boolean) extends
ExpressionBuilder {
override def functionSignature: Option[FunctionSignature] = {
val vArg = InputParameter("v")
@@ -1640,8 +1580,7 @@ case class VariantArrayAppend(
value: Expression,
failOnError: Boolean = true)
extends TernaryExpression
- with ExpectsInputTypes
- with QueryErrorsBase {
+ with ExpectsInputTypes {
override def first: Expression = input
override def second: Expression = path
@@ -1658,35 +1597,13 @@ case class VariantArrayAppend(
val result = super.checkInputDataTypes()
if (result.isFailure) {
result
- } else if (value.dataType == NullType) {
- TypeCheckResult.TypeCheckSuccess
- } else if (!VariantGet.checkDataType(value.dataType, allowStructsAndMaps =
false)) {
- DataTypeMismatch(
- errorSubClass = "CAST_WITHOUT_SUGGESTION",
- messageParameters =
- Map("srcType" -> toSQLType(value.dataType), "targetType" ->
toSQLType(VariantType)))
} else {
- TypeCheckResult.TypeCheckSuccess
+ VariantManipulationExpressionUtils.checkValueType(value.dataType)
}
}
- // When the path is a foldable expression, parse it once at planning time
and cache it. `None`
- // means the path is dynamic (or a foldable NULL, which makes the whole
expression NULL
- // and is never evaluated).
- @transient private lazy val foldablePath:
Option[VariantArrayAppend.ParsedAppendPath] = {
- if (path.foldable) {
- val p = path.eval()
- if (p == null) {
- None
- } else {
- val s = p.asInstanceOf[UTF8String].toString
- Some(VariantArrayAppend.ParsedAppendPath(
- VariantExpressionEvalUtils.parseVariantPath(s, prettyName, allowRoot
= true), s))
- }
- } else {
- None
- }
- }
+ @transient private lazy val foldablePath: Option[FoldableVariantPath] =
+ VariantManipulationExpressionUtils.foldablePath(path, prettyName,
allowRoot = true)
override def eval(input: InternalRow): Any = {
val _ = foldablePath
@@ -1719,22 +1636,9 @@ case class VariantArrayAppend(
case None =>
s"""$cls.arrayAppendAtPath($vVal, $pVal, $valVal, $fromArg,
"$prettyName", $failOnError)"""
}
- if (failOnError) {
- nullSafeCodeGen(ctx, ev,
- (vVal, pVal, valVal) => s"${ev.value} = ${call(vVal, pVal, valVal)};")
- } else {
- val resultType = CodeGenerator.javaType(VariantType)
- val tmp = ctx.freshName("appendResult")
- nullSafeCodeGen(ctx, ev, (vVal, pVal, valVal) =>
- s"""
- |$resultType $tmp = ${call(vVal, pVal, valVal)};
- |if ($tmp == null) {
- | ${ev.isNull} = true;
- |} else {
- | ${ev.value} = $tmp;
- |}
- """.stripMargin)
- }
+ nullSafeCodeGen(ctx, ev, (vVal, pVal, valVal) =>
+ VariantManipulationExpressionUtils.assignVariantResult(
+ ev, call(vVal, pVal, valVal), nullable))
}
override def prettyName: String =
@@ -1745,16 +1649,6 @@ case class VariantArrayAppend(
copy(input = newFirst, path = newSecond, value = newThird)
}
-object VariantArrayAppend {
- // Caches a foldable path. `VariantBuilder.PathSegment` is not
`Serializable`, so the Java form is
- // `@transient` and re-derived once per executor task after deserialization.
`pathStr` is the
- // source string, retained for error messages.
- case class ParsedAppendPath(segments: Array[VariantPathSegment], pathStr:
String) {
- @transient lazy val javaSegments: Array[VariantBuilder.PathSegment] =
- VariantExpressionEvalUtils.toJavaSegments(segments)
- }
-}
-
abstract class VariantArrayAppendExpressionBuilderBase(failOnError: Boolean)
extends ExpressionBuilder {
override def build(funcName: String, expressions: Seq[Expression]):
Expression = {
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala
index 75018a572b8b..634100197987 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala
@@ -1459,7 +1459,7 @@ class VariantExpressionSuite extends SparkFunSuite with
ExpressionEvalHelper {
checkDelete("""{"a": null, "b": 2}""", Seq("$.a"), """{"b":2}""")
checkDelete("[null, 1, null]", Seq("$[0]"), "[1,null]")
- // Mixed literal + dynamic path exercises both `ParsedDeletePath` and
`DynamicDeletePath`
+ // Mixed literal + dynamic path exercises both `FoldableVariantPath` and
`DynamicVariantPath`
// arms of `eval` in a single call.
val mixedLitDyn = VariantDelete(Seq(
Literal(parseJson("""{"a": 1, "b": 2, "c": 3}""")),
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
