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 759fbd9ea34c [SPARK-58221][SQL] Add `try_variant_set` expression
759fbd9ea34c is described below
commit 759fbd9ea34c2356e4d042b27d14557908d8c7d3
Author: bojana-db <[email protected]>
AuthorDate: Fri Jul 24 11:25:58 2026 +0000
[SPARK-58221][SQL] Add `try_variant_set` expression
### What changes were proposed in this pull request?
Adds the SQL function `try_variant_set(v, path, val[, create_if_missing])`,
the error-tolerant counterpart of `variant_set` (SPARK-57804). It sets (inserts
or replaces) a value in a Variant value at a single JSONPath location,
returning NULL instead of failing the query when the operation hits a
recoverable error.
Details:
- On a valid set, behaves exactly like `variant_set`: an object path (e.g.
`$.a`) replaces the field if it exists and creates it when `create_if_missing`
is true (the default); an array path (e.g. `$[N]`) replaces the element at
index `N`, and when `N` is at or past the end and `create_if_missing` is true,
pads the array with variant nulls up to `N`; missing intermediate keys/indices
along the path are created when `create_if_missing` is true;
- When `create_if_missing` is false, a missing leaf, a missing
intermediate, or an out-of-range array index leaves `v` unchanged;
- Recoverable errors are caught and return NULL instead of throwing: any
path type mismatch (`VARIANT_PATH_TYPE_MISMATCH`, e.g. an object key applied to
an array);
- Non-recoverable errors still throw, matching `variant_set`: a malformed
path, including the root path `$`, is rejected with `INVALID_VARIANT_PATH` (the
path is parsed before the set), and a result exceeding the size limit raises
`VARIANT_SIZE_LIMIT`;
- Any NULL argument returns NULL;
- The value may be any expression castable to variant.
### Why are the changes needed?
Error-tolerant counterpart of `variant_set`, matching `try_variant_get`,
`try_variant_insert`, and other `try_*` functions.
### Does this PR introduce _any_ user-facing change?
Yes, a new SQL function `try_variant_set` (and the corresponding
Scala/Python `functions` API).
### How was this patch tested?
Unit tests.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Cursor with Claude Opus 4.8
Closes #57373 from bojana-db/try-variant-set.
Authored-by: bojana-db <[email protected]>
Signed-off-by: Uros Bojanic <[email protected]>
---
.../source/reference/pyspark.sql/functions.rst | 1 +
python/pyspark/sql/connect/functions/builtin.py | 15 ++
python/pyspark/sql/functions/__init__.py | 1 +
python/pyspark/sql/functions/builtin.py | 65 ++++++++
python/pyspark/sql/tests/test_functions.py | 13 ++
.../scala/org/apache/spark/sql/functions.scala | 86 +++++++++++
.../sql/catalyst/analysis/FunctionRegistry.scala | 1 +
.../variant/VariantExpressionEvalUtils.scala | 14 +-
.../expressions/variant/variantExpressions.scala | 120 +++++++++++----
.../variant/VariantExpressionSuite.scala | 168 ++++++++++++++-------
.../apache/spark/sql/PlanGenerationTestSuite.scala | 8 +
.../function_try_variant_set.explain | 2 +
..._try_variant_set_with_create_if_missing.explain | 2 +
.../explain-results/function_variant_set.explain | 2 +-
...tion_variant_set_with_create_if_missing.explain | 2 +-
.../queries/function_try_variant_set.json | 125 +++++++++++++++
.../queries/function_try_variant_set.proto.bin | Bin 0 -> 1142 bytes
...ion_try_variant_set_with_create_if_missing.json | 146 ++++++++++++++++++
...ry_variant_set_with_create_if_missing.proto.bin | Bin 0 -> 1338 bytes
.../sql-functions/sql-expression-schema.md | 1 +
.../scala/org/apache/spark/sql/VariantSuite.scala | 49 ++++++
21 files changed, 728 insertions(+), 93 deletions(-)
diff --git a/python/docs/source/reference/pyspark.sql/functions.rst
b/python/docs/source/reference/pyspark.sql/functions.rst
index 7ce413fde0bb..91b722d50efd 100644
--- a/python/docs/source/reference/pyspark.sql/functions.rst
+++ b/python/docs/source/reference/pyspark.sql/functions.rst
@@ -609,6 +609,7 @@ VARIANT Functions
variant_insert
try_variant_insert
variant_set
+ try_variant_set
try_parse_json
to_variant_object
diff --git a/python/pyspark/sql/connect/functions/builtin.py
b/python/pyspark/sql/connect/functions/builtin.py
index acf1ee6b069f..f3e1a0c058f2 100644
--- a/python/pyspark/sql/connect/functions/builtin.py
+++ b/python/pyspark/sql/connect/functions/builtin.py
@@ -2252,6 +2252,21 @@ def variant_set(
variant_set.__doc__ = pysparkfuncs.variant_set.__doc__
+def try_variant_set(
+ v: "ColumnOrName",
+ path: Union[Column, str],
+ value: "ColumnOrName",
+ create_if_missing: bool = True,
+) -> Column:
+ path_col = path if isinstance(path, Column) else lit(path)
+ return _invoke_function(
+ "try_variant_set", _to_col(v), path_col, _to_col(value),
lit(create_if_missing)
+ )
+
+
+try_variant_set.__doc__ = pysparkfuncs.try_variant_set.__doc__
+
+
def variant_array_append(
v: "ColumnOrName", path: Union[Column, str], value: "ColumnOrName"
) -> Column:
diff --git a/python/pyspark/sql/functions/__init__.py
b/python/pyspark/sql/functions/__init__.py
index 7780bae38ee9..50f0fd440dd9 100644
--- a/python/pyspark/sql/functions/__init__.py
+++ b/python/pyspark/sql/functions/__init__.py
@@ -487,6 +487,7 @@ __all__ = [ # noqa: F405
"variant_insert",
"try_variant_insert",
"variant_set",
+ "try_variant_set",
"try_parse_json",
"to_variant_object",
# XML Functions
diff --git a/python/pyspark/sql/functions/builtin.py
b/python/pyspark/sql/functions/builtin.py
index 3cd42dc2e0e2..a9f3769b0831 100644
--- a/python/pyspark/sql/functions/builtin.py
+++ b/python/pyspark/sql/functions/builtin.py
@@ -22853,6 +22853,71 @@ def variant_set(
)
+@_try_remote_functions
+def try_variant_set(
+ v: "ColumnOrName",
+ path: Union[Column, str],
+ value: "ColumnOrName",
+ create_if_missing: bool = True,
+) -> Column:
+ """
+ Sets or upserts a value in a variant at the given JSONPath location. An
existing object field
+ or array element at the target is replaced. A missing field, array index,
or intermediate path
+ is created, unless `create_if_missing` is false, in which case the variant
is left unchanged.
+ Returns NULL if a path segment hits a value of an incompatible type, or if
any argument is NULL.
+
+ .. versionadded:: 4.3.0
+
+ Parameters
+ ----------
+ v : :class:`~pyspark.sql.Column` or str
+ a variant column or column name
+ path : :class:`~pyspark.sql.Column` or str
+ the JSONPath set target. A `str` is a literal path; a
:class:`~pyspark.sql.Column` supplies
+ the path at runtime. A valid path should start with `$` and is
followed by one or more
+ segments like `[123]`, `.name`, `['name']`, or `["name"]`. The root
path `$` is not allowed.
+ value : :class:`~pyspark.sql.Column` or str
+ the value to set. Any expression castable to variant.
+ create_if_missing : bool, optional
+ whether to create missing keys or out-of-range array indices (default
True).
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ a variant column with `value` set at `path`
+
+ Examples
+ --------
+ >>> from pyspark.sql.functions import lit, parse_json, to_json,
try_variant_set
+ >>> df = spark.createDataFrame([{'json': '''{"a": 1, "arr": [1, 2,
3]}'''}])
+ >>> v = parse_json(df.json)
+ >>> df.select(to_json(try_variant_set(v, "$.a",
lit(9))).alias("r")).collect()
+ [Row(r='{"a":9,"arr":[1,2,3]}')]
+ >>> df.select(to_json(try_variant_set(v, "$.b",
lit(2))).alias("r")).collect()
+ [Row(r='{"a":1,"arr":[1,2,3],"b":2}')]
+ >>> df.select(to_json(try_variant_set(v, "$.arr[1]",
lit(9))).alias("r")).collect()
+ [Row(r='{"a":1,"arr":[1,9,3]}')]
+ >>> df.select(to_json(try_variant_set(v, "$.arr[5]",
lit(9))).alias("r")).collect()
+ [Row(r='{"a":1,"arr":[1,2,3,null,null,9]}')]
+ >>> df.select(to_json(try_variant_set(v, "$.b", lit(2),
False)).alias("r")).collect()
+ [Row(r='{"a":1,"arr":[1,2,3]}')]
+ >>> df.select(to_json(try_variant_set(v, "$.a.b",
lit(9))).alias("r")).collect()
+ [Row(r=None)]
+ >>> df.select(to_json(try_variant_set(v, "$.a",
lit(None))).alias("r")).collect()
+ [Row(r=None)]
+ """
+ from pyspark.sql.classic.column import _to_java_column
+
+ path_col = path if isinstance(path, Column) else lit(path)
+ return _invoke_function(
+ "try_variant_set",
+ _to_java_column(v),
+ _to_java_column(path_col),
+ _to_java_column(value),
+ _enum_to_value(create_if_missing),
+ )
+
+
@_try_remote_functions
def variant_array_append(
v: "ColumnOrName", path: Union[Column, str], value: "ColumnOrName"
diff --git a/python/pyspark/sql/tests/test_functions.py
b/python/pyspark/sql/tests/test_functions.py
index 2594c8c60b76..d3a4072c707e 100644
--- a/python/pyspark/sql/tests/test_functions.py
+++ b/python/pyspark/sql/tests/test_functions.py
@@ -3559,6 +3559,19 @@ class FunctionsTestsMixin:
df.select(F.to_json(F.variant_set(v, df.newpath, F.lit(9)))),
['{"a":1,"z":9}', '{"b":2,"z":9}'],
)
+ check(
+ df.select(F.to_json(F.try_variant_set(v, "$.z", F.lit(9)))),
+ ['{"a":1,"z":9}', '{"b":2,"z":9}'],
+ )
+ check(
+ df.select(F.to_json(F.try_variant_set(v, "$.z", F.lit(9), False))),
+ ['{"a":1}', '{"b":2}'],
+ )
+ check(df.select(F.to_json(F.try_variant_set(v, "$[0]", F.lit(9)))),
[None, None])
+ check(
+ df.select(F.to_json(F.try_variant_set(v, df.newpath, F.lit(9)))),
+ ['{"a":1,"z":9}', '{"b":2,"z":9}'],
+ )
arr = F.parse_json(df.arr)
check(
df.select(F.to_json(F.variant_array_append(arr, "$", F.lit(9)))),
diff --git a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala
b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala
index 8b8cd98960e0..3961a2d9de1d 100644
--- a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala
@@ -14324,6 +14324,92 @@ object functions {
def variant_set(v: Column, path: String, value: Column, createIfMissing:
Boolean): Column =
Column.fn("variant_set", v, lit(path), value, lit(createIfMissing))
+ /**
+ * Sets or upserts a value in a variant at the given JSONPath location. An
existing object field
+ * or array element at the target is replaced. A missing field, array index,
or intermediate
+ * path is created. Returns NULL if a path segment hits a value of an
incompatible type, or if
+ * any argument is NULL.
+ *
+ * @param v
+ * a variant column.
+ * @param path
+ * the column containing the JSONPath string identifying the set target. A
valid path should
+ * start with `$` and is followed by one or more segments like `[123]`,
`.name`, `['name']`,
+ * or `["name"]`. The root path `$` is not allowed.
+ * @param value
+ * the value to set. Any expression castable to variant.
+ * @group variant_funcs
+ * @since 4.3.0
+ */
+ def try_variant_set(v: Column, path: Column, value: Column): Column =
+ Column.fn("try_variant_set", v, path, value)
+
+ /**
+ * Sets or upserts a value in a variant at the given JSONPath location. An
existing object field
+ * or array element at the target is replaced. A missing field, array index,
or intermediate
+ * path is created. Returns NULL if a path segment hits a value of an
incompatible type, or if
+ * any argument is NULL.
+ *
+ * @param v
+ * a variant column.
+ * @param path
+ * the JSONPath identifying the set target. A valid path should start with
`$` and is followed
+ * by one or more segments like `[123]`, `.name`, `['name']`, or
`["name"]`. The root path `$`
+ * is not allowed.
+ * @param value
+ * the value to set. Any expression castable to variant.
+ * @group variant_funcs
+ * @since 4.3.0
+ */
+ def try_variant_set(v: Column, path: String, value: Column): Column =
+ Column.fn("try_variant_set", v, lit(path), value)
+
+ /**
+ * Sets or upserts a value in a variant at the given JSONPath location. An
existing object field
+ * or array element at the target is replaced. A missing field, array index,
or intermediate
+ * path is created, unless `createIfMissing` is false, in which case the
variant is left
+ * unchanged. Returns NULL if a path segment hits a value of an incompatible
type, or if any
+ * argument is NULL.
+ *
+ * @param v
+ * a variant column.
+ * @param path
+ * the column containing the JSONPath string identifying the set target. A
valid path should
+ * start with `$` and is followed by one or more segments like `[123]`,
`.name`, `['name']`,
+ * or `["name"]`. The root path `$` is not allowed.
+ * @param value
+ * the value to set. Any expression castable to variant.
+ * @param createIfMissing
+ * whether to create missing keys or out-of-range array indices.
+ * @group variant_funcs
+ * @since 4.3.0
+ */
+ def try_variant_set(v: Column, path: Column, value: Column, createIfMissing:
Boolean): Column =
+ Column.fn("try_variant_set", v, path, value, lit(createIfMissing))
+
+ /**
+ * Sets or upserts a value in a variant at the given JSONPath location. An
existing object field
+ * or array element at the target is replaced. A missing field, array index,
or intermediate
+ * path is created, unless `createIfMissing` is false, in which case the
variant is left
+ * unchanged. Returns NULL if a path segment hits a value of an incompatible
type, or if any
+ * argument is NULL.
+ *
+ * @param v
+ * a variant column.
+ * @param path
+ * the JSONPath identifying the set target. A valid path should start with
`$` and is followed
+ * by one or more segments like `[123]`, `.name`, `['name']`, or
`["name"]`. The root path `$`
+ * is not allowed.
+ * @param value
+ * the value to set. Any expression castable to variant.
+ * @param createIfMissing
+ * whether to create missing keys or out-of-range array indices.
+ * @group variant_funcs
+ * @since 4.3.0
+ */
+ def try_variant_set(v: Column, path: String, value: Column, createIfMissing:
Boolean): Column =
+ Column.fn("try_variant_set", v, lit(path), value, lit(createIfMissing))
+
/**
* Appends a value to the array in a variant at the given JSONPath location.
Returns the variant
* unchanged if a path key or index is absent. Throws an error if a path
segment hits a value of
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
index 5b39aa06f1e4..965330db5828 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
@@ -990,6 +990,7 @@ object FunctionRegistry {
expressionBuilder("variant_insert", VariantInsertExpressionBuilder),
expressionBuilder("try_variant_insert", TryVariantInsertExpressionBuilder),
expressionBuilder("variant_set", VariantSetExpressionBuilder),
+ expressionBuilder("try_variant_set", TryVariantSetExpressionBuilder),
expressionBuilder("variant_array_append",
VariantArrayAppendExpressionBuilder),
expressionBuilder("try_variant_array_append",
TryVariantArrayAppendExpressionBuilder),
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala
index 0872cda99801..da1bfbcea829 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala
@@ -183,7 +183,9 @@ object VariantExpressionEvalUtils {
* Set `input` at `javaSegments` to `value`. `path` is the source string
used in error messages.
* The cast and set share one try, so any size overflow maps to
`VARIANT_SIZE_LIMIT` and a type
* mismatch maps to `VARIANT_PATH_TYPE_MISMATCH`. When `createIfMissing` is
false, a missing
- * key/index leaves the variant unchanged.
+ * key/index leaves the variant unchanged. When `failOnError` is false (the
`try_variant_set`
+ * mode), a path type mismatch returns null instead of throwing; a size
overflow (and a malformed
+ * path, rejected earlier during parsing) is still raised.
*/
def setAtPath(
input: VariantVal,
@@ -192,7 +194,8 @@ object VariantExpressionEvalUtils {
value: Any,
valueDataType: DataType,
createIfMissing: Boolean,
- functionName: String): VariantVal = {
+ functionName: String,
+ failOnError: Boolean): VariantVal = {
val v = new Variant(input.getValue, input.getMetadata)
try {
val valVal = castToVariant(value, valueDataType)
@@ -200,6 +203,7 @@ object VariantExpressionEvalUtils {
val out = VariantBuilder.setAtPath(v, javaSegments, valVariant,
createIfMissing)
new VariantVal(out.getValue, out.getMetadata)
} catch {
+ case _: VariantPathTypeMismatchException if !failOnError => null
case e: VariantPathTypeMismatchException =>
throw QueryExecutionErrors.variantPathTypeMismatch(
path, renderVariantPath(javaSegments.take(e.depth)), functionName)
@@ -214,10 +218,12 @@ object VariantExpressionEvalUtils {
value: Any,
valueDataType: DataType,
createIfMissing: Boolean,
- functionName: String): VariantVal = {
+ functionName: String,
+ failOnError: Boolean): VariantVal = {
val pathStr = path.toString
val javaSegments = toJavaSegments(parseVariantPath(pathStr, functionName))
- setAtPath(input, javaSegments, pathStr, value, valueDataType,
createIfMissing, functionName)
+ setAtPath(input, javaSegments, pathStr, value, valueDataType,
createIfMissing, functionName,
+ failOnError)
}
/**
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 fa6a955097d9..e9ce291b655c 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
@@ -1025,7 +1025,8 @@ case class VariantSet(
input: Expression,
path: Expression,
value: Expression,
- createIfMissing: Expression)
+ createIfMissing: Expression,
+ failOnError: Boolean = true)
extends QuaternaryExpression
with ExpectsInputTypes
with QueryErrorsBase {
@@ -1036,6 +1037,7 @@ case class VariantSet(
override def fourth: Expression = createIfMissing
override def nullIntolerant: Boolean = true
+ override def nullable: Boolean = !failOnError || children.exists(_.nullable)
override def dataType: DataType = VariantType
override def inputTypes: Seq[AbstractDataType] =
@@ -1080,6 +1082,11 @@ case class VariantSet(
}
}
+ override def eval(input: InternalRow): Any = {
+ val _ = foldablePath
+ super.eval(input)
+ }
+
override protected def nullSafeEval(v: Any, p: Any, valValue: Any, create:
Any): Any = {
val inputVariant = v.asInstanceOf[VariantVal]
val createIfMissingValue = create.asInstanceOf[Boolean]
@@ -1087,36 +1094,48 @@ case class VariantSet(
case Some(parsed) =>
VariantExpressionEvalUtils.setAtPath(
inputVariant, parsed.javaSegments, parsed.pathStr, valValue,
value.dataType,
- createIfMissingValue, prettyName)
+ createIfMissingValue, prettyName, failOnError)
case None =>
VariantExpressionEvalUtils.setAtPath(
inputVariant, p.asInstanceOf[UTF8String], valValue, value.dataType,
- createIfMissingValue, prettyName)
+ createIfMissingValue, prettyName, failOnError)
}
}
override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode = {
val cls = VariantExpressionEvalUtils.getClass.getName.stripSuffix("$")
- nullSafeCodeGen(ctx, ev, (vVal, pVal, valVal, createVal) => {
- val fromArg = ctx.addReferenceObj("from", value.dataType)
+ val fromArg = ctx.addReferenceObj("from", value.dataType)
+ def call(vVal: String, pVal: String, valVal: String, createVal: String):
String =
foldablePath match {
case Some(parsed) =>
val parsedArg = ctx.addReferenceObj("setPath", parsed)
- s"""
- |${ev.value} = $cls.setAtPath(
+ s"""$cls.setAtPath(
| $vVal, $parsedArg.javaSegments(), $parsedArg.pathStr(),
$valVal, $fromArg,
- | $createVal, "$prettyName");
- """.stripMargin
+ | $createVal, "$prettyName", $failOnError)""".stripMargin
case None =>
- s"""
- |${ev.value} = $cls.setAtPath(
- | $vVal, $pVal, $valVal, $fromArg, $createVal, "$prettyName");
- """.stripMargin
+ 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)
+ }
}
- override def prettyName: String = "variant_set"
+ override def prettyName: String = if (failOnError) "variant_set" else
"try_variant_set"
override protected def withNewChildrenInternal(
newFirst: Expression,
@@ -1136,6 +1155,22 @@ object VariantSet {
}
}
+abstract class VariantSetExpressionBuilderBase(failOnError: Boolean) extends
ExpressionBuilder {
+ override def functionSignature: Option[FunctionSignature] = {
+ val vArg = InputParameter("v")
+ val pathArg = InputParameter("path")
+ val valArg = InputParameter("val")
+ val createIfMissingArg =
+ InputParameter("create_if_missing", Some(Literal.create(true,
BooleanType)))
+ Some(FunctionSignature(Seq(vArg, pathArg, valArg, createIfMissingArg)))
+ }
+
+ override def build(funcName: String, expressions: Seq[Expression]):
Expression = {
+ assert(expressions.size == 4)
+ VariantSet(expressions(0), expressions(1), expressions(2), expressions(3),
failOnError)
+ }
+}
+
// scalastyle:off line.size.limit
@ExpressionDescription(
usage = "_FUNC_(v, path, val[, create_if_missing]) - Sets or upserts a value
in a variant at " +
@@ -1176,21 +1211,48 @@ object VariantSet {
group = "variant_funcs"
)
// scalastyle:on line.size.limit
-object VariantSetExpressionBuilder extends ExpressionBuilder {
- override def functionSignature: Option[FunctionSignature] = {
- val vArg = InputParameter("v")
- val pathArg = InputParameter("path")
- val valArg = InputParameter("val")
- val createIfMissingArg =
- InputParameter("create_if_missing", Some(Literal.create(true,
BooleanType)))
- Some(FunctionSignature(Seq(vArg, pathArg, valArg, createIfMissingArg)))
- }
+object VariantSetExpressionBuilder extends
VariantSetExpressionBuilderBase(true)
- override def build(funcName: String, expressions: Seq[Expression]):
Expression = {
- assert(expressions.size == 4)
- VariantSet(expressions(0), expressions(1), expressions(2), expressions(3))
- }
-}
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = "_FUNC_(v, path, val[, create_if_missing]) - Sets or upserts a value
in a variant at " +
+ "the given JSONPath location. An existing object field or array element at
the target is " +
+ "replaced. A missing field, array index, or intermediate path is created,
unless " +
+ "`create_if_missing` is false, in which case the variant is left
unchanged. Returns NULL if " +
+ "a path segment hits a value of an incompatible type, or if any argument
is NULL.",
+ arguments = """
+ Arguments:
+ * v - A variant value to mutate.
+ * path - A string expression evaluating to a JSONPath identifying the
set target. A valid
+ path should start with `$` and is followed by one or more segments
like `[123]`,
+ `.name`, `['name']`, or `["name"]`. The root path `$` is not allowed.
+ * val - Any expression castable to variant.
+ * create_if_missing - An optional boolean (default true).
+ """,
+ examples = """
+ Examples:
+ > SELECT _FUNC_(parse_json('{"a": 1}'), '$.a', 2);
+ {"a":2}
+ > SELECT _FUNC_(parse_json('{"a": 1}'), '$.b', 3);
+ {"a":1,"b":3}
+ > SELECT _FUNC_(parse_json('{"a": {"b": 1}}'), '$.a.c.d', 2);
+ {"a":{"b":1,"c":{"d":2}}}
+ > SELECT _FUNC_(parse_json('["a","b","c"]'), '$[1]', 'z');
+ ["a","z","c"]
+ > SELECT _FUNC_(parse_json('["a","b","c"]'), '$[5]', 'z');
+ ["a","b","c",null,null,"z"]
+ > SELECT _FUNC_(parse_json('{"a": 1}'), '$.b', 2, false);
+ {"a":1}
+ > SELECT _FUNC_(parse_json('{"a": 1}'), '$.a.b', 2);
+ NULL
+ > SELECT _FUNC_(parse_json('{"a": 1}'), '$.a', null);
+ NULL
+ """,
+ since = "4.3.0",
+ group = "variant_funcs"
+)
+// scalastyle:on line.size.limit
+object TryVariantSetExpressionBuilder extends
VariantSetExpressionBuilderBase(false)
case class VariantArrayAppend(
input: 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 92f20bcc3127..eaf9ac267ff0 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
@@ -1503,14 +1503,55 @@ class VariantExpressionSuite extends SparkFunSuite with
ExpressionEvalHelper {
value: Expression,
expected: String,
createIfMissing: Boolean = true): Unit = {
- val expr = VariantSet(
- Literal(parseJson(input)),
- Literal.create(path, StringType),
- value,
- Literal(createIfMissing))
+ // On success `variant_set` and `try_variant_set` behave identically.
+ Seq(true, false).foreach { failOnError =>
+ val expr = VariantSet(
+ Literal(parseJson(input)),
+ Literal.create(path, StringType),
+ value,
+ Literal(createIfMissing),
+ failOnError)
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(expr, StringType)),
+ expected)
+ }
+ }
+
+ // A recoverable error: `variant_set` throws, but `try_variant_set`
returns NULL.
+ def checkSetRecoverableError(
+ input: String,
+ path: String,
+ value: Expression,
+ condition: String,
+ parameters: Map[String, String],
+ createIfMissing: Boolean = true): Unit = {
+ checkErrorInExpression[SparkRuntimeException](
+ VariantSet(Literal(parseJson(input)), Literal.create(path,
StringType), value,
+ Literal(createIfMissing)),
+ condition,
+ parameters)
checkEvaluation(
- ResolveTimeZone.resolveTimeZones(Cast(expr, StringType)),
- expected)
+ VariantSet(Literal(parseJson(input)), Literal.create(path,
StringType), value,
+ Literal(createIfMissing), failOnError = false),
+ null)
+ }
+
+ // An unrecoverable error (a malformed path or a size overflow): both
functions throw, and only
+ // the `functionName` in the error message differs.
+ def checkSetUnrecoverableError(
+ input: String,
+ path: String,
+ value: Expression,
+ condition: String,
+ parameters: String => Map[String, String]): Unit = {
+ Seq("variant_set" -> true, "try_variant_set" -> false).foreach {
+ case (name, failOnError) =>
+ checkErrorInExpression[SparkRuntimeException](
+ VariantSet(Literal(parseJson(input)), Literal.create(path,
StringType), value,
+ Literal(true), failOnError),
+ condition,
+ parameters(name))
+ }
}
// Object sets: replace an existing key, or add a missing one.
@@ -1561,59 +1602,57 @@ class VariantExpressionSuite extends SparkFunSuite with
ExpressionEvalHelper {
checkSet(
"{}", "$.a", Literal.create(Array(1, 2, 3), ArrayType(IntegerType)),
"""{"a":[1,2,3]}""")
- // NULL-intolerant: any NULL argument yields NULL.
- checkEvaluation(
- VariantSet(Literal.create(null, VariantType), Literal("$.a"),
Literal(1), Literal(true)),
- null)
+ Seq(true, false).foreach { failOnError =>
+ checkEvaluation(
+ VariantSet(Literal.create(null, VariantType), Literal("$.a"),
Literal(1), Literal(true),
+ failOnError),
+ null)
+ checkEvaluation(
+ VariantSet(Literal(parseJson("""{"a": 1}""")), Literal("$.a"),
Literal(2),
+ Literal.create(null, BooleanType), failOnError),
+ null)
+ }
checkSet("""{"a": 1}""", null, Literal(1), null)
checkSet("""{"a": 1}""", "$.a", Literal.create(null, VariantType), null)
checkSet("""{"a": 1}""", "$.a", Literal.create(null, NullType), null)
- checkEvaluation(
- VariantSet(Literal(parseJson("""{"a": 1}""")), Literal("$.a"),
Literal(2),
- Literal.create(null, BooleanType)),
- null)
- // Dynamic (non-foldable) path and create_if_missing.
- val dynamic = VariantSet(
- Literal(parseJson("""{"a": 1}""")),
- BoundReference(0, StringType, nullable = true),
- Literal(2),
- BoundReference(1, BooleanType, nullable = true))
- checkEvaluation(
- ResolveTimeZone.resolveTimeZones(Cast(dynamic, StringType)),
- """{"a":1,"b":2}""",
- InternalRow(UTF8String.fromString("$.b"), true))
- checkEvaluation(
- ResolveTimeZone.resolveTimeZones(Cast(dynamic, StringType)),
- """{"a":1}""",
- InternalRow(UTF8String.fromString("$.b"), false))
+ Seq(true, false).foreach { failOnError =>
+ val dynamic = VariantSet(
+ Literal(parseJson("""{"a": 1}""")),
+ BoundReference(0, StringType, nullable = true),
+ Literal(2),
+ BoundReference(1, BooleanType, nullable = true),
+ failOnError)
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(dynamic, StringType)),
+ """{"a":1,"b":2}""",
+ InternalRow(UTF8String.fromString("$.b"), true))
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(dynamic, StringType)),
+ """{"a":1}""",
+ InternalRow(UTF8String.fromString("$.b"), false))
+ }
- // Type mismatch throws regardless of create_if_missing.
+ // Recoverable errors: `variant_set` throws; `try_variant_set` returns
NULL. Every shape of
+ // path type mismatch is recoverable, regardless of create_if_missing.
Seq(true, false).foreach { create =>
- checkErrorInExpression[SparkRuntimeException](
- VariantSet(Literal(parseJson("""{"a": 1}""")), Literal("$.a.b"),
Literal(2),
- Literal(create)),
+ checkSetRecoverableError("""{"a": 1}""", "$.a.b", Literal(2),
"VARIANT_PATH_TYPE_MISMATCH",
- Map("path" -> "$.a.b", "failedAt" -> "$.a", "functionName" ->
"`variant_set`"))
+ Map("path" -> "$.a.b", "failedAt" -> "$.a", "functionName" ->
"`variant_set`"),
+ create)
}
- checkErrorInExpression[SparkRuntimeException](
- VariantSet(Literal(parseJson("5")), Literal("$.a"), Literal(2),
Literal(true)),
+ checkSetRecoverableError("5", "$.a", Literal(2),
"VARIANT_PATH_TYPE_MISMATCH",
Map("path" -> "$.a", "failedAt" -> "$", "functionName" ->
"`variant_set`"))
-
// Segment kind not matching the container: array index on an object,
object key on an array.
- checkErrorInExpression[SparkRuntimeException](
- VariantSet(Literal(parseJson("""{"a": 1}""")), Literal("$[0]"),
Literal(2), Literal(true)),
+ checkSetRecoverableError("""{"a": 1}""", "$[0]", Literal(2),
"VARIANT_PATH_TYPE_MISMATCH",
Map("path" -> "$[0]", "failedAt" -> "$", "functionName" ->
"`variant_set`"))
- checkErrorInExpression[SparkRuntimeException](
- VariantSet(Literal(parseJson("[1, 2]")), Literal("$.a"), Literal(2),
Literal(true)),
+ checkSetRecoverableError("[1, 2]", "$.a", Literal(2),
"VARIANT_PATH_TYPE_MISMATCH",
Map("path" -> "$.a", "failedAt" -> "$", "functionName" ->
"`variant_set`"))
// A key that needs bracket notation is rendered with brackets in
`failedAt`.
- checkErrorInExpression[SparkRuntimeException](
- VariantSet(Literal(parseJson("""{"a.b": 5}""")),
Literal("""$['a.b'].c"""), Literal(2),
- Literal(true)),
+ checkSetRecoverableError("""{"a.b": 5}""", """$['a.b'].c""", Literal(2),
"VARIANT_PATH_TYPE_MISMATCH",
Map("path" -> "$['a.b'].c", "failedAt" -> "$['a.b']", "functionName" ->
"`variant_set`"))
@@ -1622,9 +1661,11 @@ class VariantExpressionSuite extends SparkFunSuite with
ExpressionEvalHelper {
Literal.create(null, MapType(StringType, IntegerType)),
Literal.create(null, StructType(Seq(StructField("x", IntegerType))))
).foreach { v =>
- assert(
- VariantSet(Literal(parseJson("{}")), Literal("$.a"), v, Literal(true))
- .checkInputDataTypes().isFailure)
+ Seq(true, false).foreach { failOnError =>
+ assert(
+ VariantSet(Literal(parseJson("{}")), Literal("$.a"), v,
Literal(true), failOnError)
+ .checkInputDataTypes().isFailure)
+ }
}
val structVal = ToVariantObject(Literal.create(create_row(1, "x"),
@@ -1639,20 +1680,31 @@ class VariantExpressionSuite extends SparkFunSuite with
ExpressionEvalHelper {
checkSet("""{"k": 1}""", "$.k", mapVal, """{"k":{"x":3,"y":2,"z":1}}""")
checkSet("""{"a": {"b": 1}}""", "$.a.b", structVal,
"""{"a":{"b":{"a":1,"b":"x"}}}""")
- checkErrorInExpression[SparkRuntimeException](
- VariantSet(Literal(parseJson("{}")), Literal("$"), Literal(1),
Literal(true)),
+ // Unrecoverable errors: both functions throw.
+ checkSetUnrecoverableError("{}", "$", Literal(1),
"INVALID_VARIANT_PATH",
- Map("path" -> "$", "functionName" -> "`variant_set`"))
-
- checkErrorInExpression[SparkRuntimeException](
- VariantSet(Literal(parseJson("{}")), Literal("$.a[2000000000]"),
Literal(1), Literal(true)),
- "VARIANT_SIZE_LIMIT",
- Map("sizeLimit" -> "16.0 MiB", "functionName" -> "`variant_set`"))
+ name => Map("path" -> "$", "functionName" -> s"`$name`"))
val tooBig = "x".repeat(16 * 1024 * 1024)
- checkErrorInExpression[SparkRuntimeException](
- VariantSet(Literal(parseJson("{}")), Literal("$.a"), Literal(tooBig),
Literal(true)),
+ checkSetUnrecoverableError("{}", "$.a[2000000000]", Literal(1),
+ "VARIANT_SIZE_LIMIT",
+ name => Map("sizeLimit" -> "16.0 MiB", "functionName" -> s"`$name`"))
+ checkSetUnrecoverableError("{}", "$.a", Literal(tooBig),
"VARIANT_SIZE_LIMIT",
- Map("sizeLimit" -> "16.0 MiB", "functionName" -> "`variant_set`"))
+ name => Map("sizeLimit" -> "16.0 MiB", "functionName" -> s"`$name`"))
+
+ Seq("variant_set" -> true, "try_variant_set" -> false).foreach {
+ case (name, failOnError) =>
+ checkErrorInExpression[SparkRuntimeException](
+ VariantSet(
+ BoundReference(0, VariantType, nullable = true),
+ Literal.create("$", StringType),
+ Literal(1),
+ Literal(true),
+ failOnError),
+ InternalRow(null),
+ "INVALID_VARIANT_PATH",
+ Map("path" -> "$", "functionName" -> s"`$name`"))
+ }
}
test("variant_array_append") {
diff --git
a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
index 068460293f15..266de35706da 100644
---
a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
+++
b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
@@ -2779,6 +2779,14 @@ class PlanGenerationTestSuite extends ConnectFunSuite
with Logging {
fn.variant_set(fn.parse_json(fn.col("g")), "$.a", fn.lit(1), false)
}
+ functionTest("try_variant_set") {
+ fn.try_variant_set(fn.parse_json(fn.col("g")), "$.a", fn.lit(1))
+ }
+
+ functionTest("try_variant_set with create_if_missing") {
+ fn.try_variant_set(fn.parse_json(fn.col("g")), "$.a", fn.lit(1), false)
+ }
+
functionTest("variant_array_append") {
fn.variant_array_append(fn.parse_json(fn.col("g")), "$.a", fn.lit(1))
}
diff --git
a/sql/connect/common/src/test/resources/query-tests/explain-results/function_try_variant_set.explain
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_try_variant_set.explain
new file mode 100644
index 000000000000..30503b185e1e
--- /dev/null
+++
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_try_variant_set.explain
@@ -0,0 +1,2 @@
+Project
[try_variant_set(static_invoke(VariantExpressionEvalUtils.parseJson(g#0, false,
true, true)), $.a, 1, true, false) AS try_variant_set(parse_json(g), $.a, 1,
true)#0]
++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0]
diff --git
a/sql/connect/common/src/test/resources/query-tests/explain-results/function_try_variant_set_with_create_if_missing.explain
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_try_variant_set_with_create_if_missing.explain
new file mode 100644
index 000000000000..831e27641964
--- /dev/null
+++
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_try_variant_set_with_create_if_missing.explain
@@ -0,0 +1,2 @@
+Project
[try_variant_set(static_invoke(VariantExpressionEvalUtils.parseJson(g#0, false,
true, true)), $.a, 1, false, false) AS try_variant_set(parse_json(g), $.a, 1,
false)#0]
++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0]
diff --git
a/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_set.explain
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_set.explain
index 3f695273d79b..37dbd79d5071 100644
---
a/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_set.explain
+++
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_set.explain
@@ -1,2 +1,2 @@
-Project [variant_set(static_invoke(VariantExpressionEvalUtils.parseJson(g#0,
false, true, true)), $.a, 1, true) AS variant_set(parse_json(g), $.a, 1,
true)#0]
+Project [variant_set(static_invoke(VariantExpressionEvalUtils.parseJson(g#0,
false, true, true)), $.a, 1, true, true) AS variant_set(parse_json(g), $.a, 1,
true)#0]
+- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0]
diff --git
a/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_set_with_create_if_missing.explain
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_set_with_create_if_missing.explain
index 2ad5e4938e92..c0a38e55e436 100644
---
a/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_set_with_create_if_missing.explain
+++
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_set_with_create_if_missing.explain
@@ -1,2 +1,2 @@
-Project [variant_set(static_invoke(VariantExpressionEvalUtils.parseJson(g#0,
false, true, true)), $.a, 1, false) AS variant_set(parse_json(g), $.a, 1,
false)#0]
+Project [variant_set(static_invoke(VariantExpressionEvalUtils.parseJson(g#0,
false, true, true)), $.a, 1, false, true) AS variant_set(parse_json(g), $.a, 1,
false)#0]
+- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0]
diff --git
a/sql/connect/common/src/test/resources/query-tests/queries/function_try_variant_set.json
b/sql/connect/common/src/test/resources/query-tests/queries/function_try_variant_set.json
new file mode 100644
index 000000000000..5f539f810892
--- /dev/null
+++
b/sql/connect/common/src/test/resources/query-tests/queries/function_try_variant_set.json
@@ -0,0 +1,125 @@
+{
+ "common": {
+ "planId": "1"
+ },
+ "project": {
+ "input": {
+ "common": {
+ "planId": "0"
+ },
+ "localRelation": {
+ "schema":
"struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e"
+ }
+ },
+ "expressions": [{
+ "unresolvedFunction": {
+ "functionName": "try_variant_set",
+ "arguments": [{
+ "unresolvedFunction": {
+ "functionName": "parse_json",
+ "arguments": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "g"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "col",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass":
"org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }],
+ "isInternal": false
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "parse_json",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass":
"org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "literal": {
+ "string": "$.a"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "try_variant_set",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass":
"org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "literal": {
+ "integer": 1
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "lit",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass":
"org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }],
+ "isInternal": false
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "try_variant_set",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }]
+ }
+}
\ No newline at end of file
diff --git
a/sql/connect/common/src/test/resources/query-tests/queries/function_try_variant_set.proto.bin
b/sql/connect/common/src/test/resources/query-tests/queries/function_try_variant_set.proto.bin
new file mode 100644
index 000000000000..074a58129fe5
Binary files /dev/null and
b/sql/connect/common/src/test/resources/query-tests/queries/function_try_variant_set.proto.bin
differ
diff --git
a/sql/connect/common/src/test/resources/query-tests/queries/function_try_variant_set_with_create_if_missing.json
b/sql/connect/common/src/test/resources/query-tests/queries/function_try_variant_set_with_create_if_missing.json
new file mode 100644
index 000000000000..323e483f9713
--- /dev/null
+++
b/sql/connect/common/src/test/resources/query-tests/queries/function_try_variant_set_with_create_if_missing.json
@@ -0,0 +1,146 @@
+{
+ "common": {
+ "planId": "1"
+ },
+ "project": {
+ "input": {
+ "common": {
+ "planId": "0"
+ },
+ "localRelation": {
+ "schema":
"struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e"
+ }
+ },
+ "expressions": [{
+ "unresolvedFunction": {
+ "functionName": "try_variant_set",
+ "arguments": [{
+ "unresolvedFunction": {
+ "functionName": "parse_json",
+ "arguments": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "g"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "col",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass":
"org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }],
+ "isInternal": false
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "parse_json",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass":
"org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "literal": {
+ "string": "$.a"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "try_variant_set",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass":
"org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "literal": {
+ "integer": 1
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "lit",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass":
"org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "literal": {
+ "boolean": false
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "try_variant_set",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass":
"org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }],
+ "isInternal": false
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "try_variant_set",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }]
+ }
+}
\ No newline at end of file
diff --git
a/sql/connect/common/src/test/resources/query-tests/queries/function_try_variant_set_with_create_if_missing.proto.bin
b/sql/connect/common/src/test/resources/query-tests/queries/function_try_variant_set_with_create_if_missing.proto.bin
new file mode 100644
index 000000000000..ccf0814fe095
Binary files /dev/null and
b/sql/connect/common/src/test/resources/query-tests/queries/function_try_variant_set_with_create_if_missing.proto.bin
differ
diff --git a/sql/core/src/test/resources/sql-functions/sql-expression-schema.md
b/sql/core/src/test/resources/sql-functions/sql-expression-schema.md
index 99285d126acd..2637fbc6f077 100644
--- a/sql/core/src/test/resources/sql-functions/sql-expression-schema.md
+++ b/sql/core/src/test/resources/sql-functions/sql-expression-schema.md
@@ -558,6 +558,7 @@
|
org.apache.spark.sql.catalyst.expressions.variant.TryVariantArrayAppendExpressionBuilder
| try_variant_array_append | SELECT try_variant_array_append(parse_json('[1,
2, 3]'), '$', 4) | struct<try_variant_array_append(parse_json([1, 2, 3]), $,
4):variant> |
|
org.apache.spark.sql.catalyst.expressions.variant.TryVariantGetExpressionBuilder
| try_variant_get | SELECT try_variant_get(parse_json('{"a": 1}'), '$.a',
'int') | struct<try_variant_get(parse_json({"a": 1}), $.a):int> |
|
org.apache.spark.sql.catalyst.expressions.variant.TryVariantInsertExpressionBuilder
| try_variant_insert | SELECT try_variant_insert(parse_json('{"a": 1}'),
'$.b', 2) | struct<try_variant_insert(parse_json({"a": 1}), $.b, 2):variant> |
+|
org.apache.spark.sql.catalyst.expressions.variant.TryVariantSetExpressionBuilder
| try_variant_set | SELECT try_variant_set(parse_json('{"a": 1}'), '$.a', 2) |
struct<try_variant_set(parse_json({"a": 1}), $.a, 2, true):variant> |
|
org.apache.spark.sql.catalyst.expressions.variant.VariantArrayAppendExpressionBuilder
| variant_array_append | SELECT variant_array_append(parse_json('[1, 2, 3]'),
'$', 4) | struct<variant_array_append(parse_json([1, 2, 3]), $, 4):variant> |
| org.apache.spark.sql.catalyst.expressions.variant.VariantDelete |
variant_delete | SELECT variant_delete(parse_json('{"a": 1, "b": 2, "c": 3,
"items": [1, 2, 3]}'), NULL, '$.a', '$.c') |
struct<variant_delete(parse_json({"a": 1, "b": 2, "c": 3, "items": [1, 2, 3]}),
NULL, $.a, $.c):variant> |
|
org.apache.spark.sql.catalyst.expressions.variant.VariantGetExpressionBuilder |
variant_get | SELECT variant_get(parse_json('{"a": 1}'), '$.a', 'int') |
struct<variant_get(parse_json({"a": 1}), $.a):int> |
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala
index fe4dde44ec99..864d0c9b631b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala
@@ -521,6 +521,55 @@ class VariantSuite extends SharedSparkSession with
ExpressionEvalHelper {
}
}
+ test("try_variant_set with literal arguments") {
+ def rows(results: Any*): Seq[Row] = results.map(Row(_))
+
+ checkAnswer(
+ sql("SELECT to_json(try_variant_set(parse_json('{\"a\": 1}'), '$.a',
2))"),
+ rows("""{"a":2}"""))
+ checkAnswer(
+ sql("SELECT to_json(try_variant_set(parse_json('{\"a\": 1}'), '$.b', 2,
false))"),
+ rows("""{"a":1}"""))
+
+ checkAnswer(
+ sql("SELECT to_json(try_variant_set(parse_json('{\"a\": 1}'), '$.a.b',
2))"),
+ rows(null))
+
+ checkAnswer(
+ sql("SELECT to_json(try_variant_set(CAST(NULL AS VARIANT), '$.a', 1))"),
+ rows(null))
+
+ // A malformed path is still raised.
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT try_variant_set(parse_json('{}'), '$', 1)").collect()
+ },
+ condition = "INVALID_VARIANT_PATH",
+ parameters = Map("path" -> "$", "functionName" ->
toSQLId("try_variant_set")))
+ }
+
+ test("try_variant_set with dynamic arguments") {
+ def rows(results: Any*): Seq[Row] = results.map(Row(_))
+ Seq("CODEGEN_ONLY", "NO_CODEGEN").foreach { codegenMode =>
+ withSQLConf(SQLConf.CODEGEN_FACTORY_MODE.key -> codegenMode) {
+ val df = Seq(
+ ("""{"a": 1}""", "$.a", 2),
+ ("""{"a": 1}""", "$.a.b", 9),
+ (null, "$.a", 2)
+ ).toDF("json", "path", "val")
+ val v = parse_json(col("json"))
+
+ checkAnswer(
+ df.select(to_json(try_variant_set(v, col("path"),
col("val"))).alias("r")),
+ rows("""{"a":2}""", null, null))
+
+ checkAnswer(
+ df.select(to_json(try_variant_set(v, "$.b", col("val"),
false)).alias("r")),
+ rows("""{"a":1}""", """{"a":1}""", null))
+ }
+ }
+ }
+
test("variant_array_append with literal arguments") {
def rows(results: Any*): Seq[Row] = results.map(Row(_))
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
