This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 11378c3c3c49 [SPARK-57804][SQL] Add `variant_set` expression
11378c3c3c49 is described below
commit 11378c3c3c49bfcc39b9130484dc7cf13a28d622
Author: bojana-db <[email protected]>
AuthorDate: Wed Jul 15 23:54:08 2026 +0800
[SPARK-57804][SQL] Add `variant_set` expression
### What changes were proposed in this pull request?
Adds the SQL function `variant_set(v, path, val[, create_if_missing])`,
which sets (inserts or replaces) a value in a Variant value at a single
JSONPath location.
Details:
- Object path (e.g. `$.a`): replaces the field if it exists; creates it
when `create_if_missing` is true (the default);
- Array path (e.g. `$[N]`): replaces the element at index `N`; when `N` is
at or past the end and `create_if_missing` is true, the array is padded 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;
- Any NULL argument returns NULL;
- The value may be any expression castable to variant (primitives, arrays,
or another variant; structs and maps are rejected);
- `VARIANT_PATH_TYPE_MISMATCH` is raised when a path segment is applied to
a value of an incompatible type; the root path `$` is rejected with
`INVALID_VARIANT_PATH`, and results exceeding the size limit raise
`VARIANT_SIZE_LIMIT`.
### Why are the changes needed?
Without `variant_set`, updating a variant means converting it to another
datatype (e.g. map), mutating, and converting back.
### Does this PR introduce _any_ user-facing change?
Yes, a new SQL function (and Scala/Python `functions` API).
### How was this patch tested?
Unit tests.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code with Claude Opus 4.8
Closes #57125 from bojana-db/variant-set.
Authored-by: bojana-db <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
(cherry picked from commit f22a4012440c26e91c12190179ddeb5e0bba9628)
Signed-off-by: Wenchen Fan <[email protected]>
---
.../apache/spark/types/variant/VariantBuilder.java | 113 ++++++++++++++
.../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 | 64 ++++++++
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 | 41 +++++
.../expressions/variant/variantExpressions.scala | 171 +++++++++++++++++++++
.../variant/VariantExpressionSuite.scala | 159 +++++++++++++++++++
.../apache/spark/sql/PlanGenerationTestSuite.scala | 8 +
.../explain-results/function_variant_set.explain | 2 +
...tion_variant_set_with_create_if_missing.explain | 2 +
.../query-tests/queries/function_variant_set.json | 125 +++++++++++++++
.../queries/function_variant_set.proto.bin | Bin 0 -> 1130 bytes
...unction_variant_set_with_create_if_missing.json | 146 ++++++++++++++++++
...on_variant_set_with_create_if_missing.proto.bin | Bin 0 -> 1322 bytes
.../sql-functions/sql-expression-schema.md | 1 +
.../scala/org/apache/spark/sql/VariantSuite.scala | 82 ++++++++++
20 files changed, 1031 insertions(+)
diff --git
a/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java
b/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java
index 9685f0a35f85..fa23d610f239 100644
---
a/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java
+++
b/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java
@@ -144,6 +144,22 @@ public class VariantBuilder {
return builder.result();
}
+ // Return a new variant with the field or array element at `segments` set to
the given value
+ // (`segments` must be non-empty). An object leaf replaces the field if
present, otherwise adds
+ // it; an array leaf replaces the element at the index. When
`createIfMissing` is true, missing
+ // leaves and intermediate keys/indices are created; when false, a missing
key/index leaves the
+ // variant unchanged. A segment that targets an incompatible value throws
+ // VariantPathTypeMismatchException, which the caller maps to
VARIANT_PATH_TYPE_MISMATCH.
+ public static Variant setAtPath(
+ Variant v, PathSegment[] segments, Variant val, boolean createIfMissing)
{
+ if (segments.length == 0) {
+ throw new IllegalArgumentException("Segments must be non-empty");
+ }
+ VariantBuilder builder = new VariantBuilder(false);
+ builder.appendWithSetImpl(v.value, v.metadata, v.pos, segments, 0, val,
createIfMissing);
+ return builder.result();
+ }
+
// Build the variant metadata from `dictionaryKeys` and return the variant
result.
public Variant result() {
int numKeys = dictionaryKeys.size();
@@ -666,6 +682,103 @@ public class VariantBuilder {
}
}
+ private void appendWithSetImpl(
+ byte[] value, byte[] metadata, int pos, PathSegment[] segments, int
depth, Variant val,
+ boolean createIfMissing) {
+ checkIndex(pos, value.length);
+ PathSegment seg = segments[depth];
+ boolean isLast = depth == segments.length - 1;
+ int basicType = value[pos] & BASIC_TYPE_MASK;
+ if (seg instanceof ObjectKeySegment && basicType == OBJECT) {
+ String key = ((ObjectKeySegment) seg).key;
+ handleObject(value, pos, (size, idSize, offsetSize, idStart,
offsetStart, dataStart) -> {
+ ArrayList<FieldEntry> fields = new ArrayList<>(size + 1);
+ int start = writePos;
+ boolean found = false;
+ for (int i = 0; i < size; ++i) {
+ int id = readUnsigned(value, idStart + idSize * i, idSize);
+ int offset = readUnsigned(value, offsetStart + offsetSize * i,
offsetSize);
+ int elementPos = dataStart + offset;
+ String fieldKey = getMetadataKey(metadata, id);
+ boolean isTarget = fieldKey.equals(key);
+ found |= isTarget;
+ int newId = addKey(fieldKey);
+ fields.add(new FieldEntry(fieldKey, newId, writePos - start));
+ if (isTarget && isLast) {
+ // Replace the existing field's value in place.
+ appendVariant(val);
+ } else if (isTarget) {
+ appendWithSetImpl(
+ value, metadata, elementPos, segments, depth + 1, val,
createIfMissing);
+ } else {
+ appendVariantImpl(value, metadata, elementPos);
+ }
+ }
+ if (!found && createIfMissing) {
+ // Target key is missing; create it (and any remaining path). When
`createIfMissing` is
+ // false this is a no-op: the fields copied above already reproduce
the input object.
+ int newId = addKey(key);
+ fields.add(new FieldEntry(key, newId, writePos - start));
+ if (isLast) {
+ appendVariant(val);
+ } else {
+ appendNewPath(segments, depth + 1, val);
+ }
+ }
+ finishWritingObject(start, fields);
+ return null;
+ });
+ } else if (seg instanceof ArrayIndexSegment && basicType == ARRAY) {
+ int index = ((ArrayIndexSegment) seg).index;
+ handleArray(value, pos, (size, offsetSize, offsetStart, dataStart) -> {
+ ArrayList<Integer> offsets = new ArrayList<>(size + 1);
+ int start = writePos;
+ if (index < size) {
+ // Replace the element at `index`, or descend into it for an
intermediate segment.
+ for (int i = 0; i < size; ++i) {
+ int offset = readUnsigned(value, offsetStart + offsetSize * i,
offsetSize);
+ int elementPos = dataStart + offset;
+ offsets.add(writePos - start);
+ if (i != index) {
+ appendVariantImpl(value, metadata, elementPos);
+ } else if (isLast) {
+ appendVariant(val);
+ } else {
+ appendWithSetImpl(
+ value, metadata, elementPos, segments, depth + 1, val,
createIfMissing);
+ }
+ }
+ } else {
+ // Index is past the end. Copy existing elements; when
`createIfMissing` is true, pad with
+ // variant nulls up to `index` and create the leaf value or the rest
of the path. When
+ // false this is a no-op: the copied elements already reproduce the
input array.
+ for (int i = 0; i < size; ++i) {
+ int offset = readUnsigned(value, offsetStart + offsetSize * i,
offsetSize);
+ offsets.add(writePos - start);
+ appendVariantImpl(value, metadata, dataStart + offset);
+ }
+ if (createIfMissing) {
+ for (int i = size; i < index; ++i) {
+ offsets.add(writePos - start);
+ appendNull();
+ }
+ offsets.add(writePos - start);
+ if (isLast) {
+ appendVariant(val);
+ } else {
+ appendNewPath(segments, depth + 1, val);
+ }
+ }
+ }
+ finishWritingArray(start, offsets);
+ return null;
+ });
+ } else {
+ // The segment kind does not match the container at this path prefix.
+ throw new VariantPathTypeMismatchException(depth);
+ }
+ }
+
// Build a fresh chain of containers for `segments[depth..]`, terminating in
`val`. Used to
// materialize missing intermediate path segments during insertion. The kind
of each segment
// decides the container created: an object-key segment creates a
single-field object, while an
diff --git a/python/docs/source/reference/pyspark.sql/functions.rst
b/python/docs/source/reference/pyspark.sql/functions.rst
index 9ae6cd12a13c..0f34809c0067 100644
--- a/python/docs/source/reference/pyspark.sql/functions.rst
+++ b/python/docs/source/reference/pyspark.sql/functions.rst
@@ -606,6 +606,7 @@ VARIANT Functions
variant_get
variant_insert
try_variant_insert
+ 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 8a33d3b2b38e..1b88329583e2 100644
--- a/python/pyspark/sql/connect/functions/builtin.py
+++ b/python/pyspark/sql/connect/functions/builtin.py
@@ -2237,6 +2237,21 @@ def try_variant_insert(
try_variant_insert.__doc__ = pysparkfuncs.try_variant_insert.__doc__
+def 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(
+ "variant_set", _to_col(v), path_col, _to_col(value),
lit(create_if_missing)
+ )
+
+
+variant_set.__doc__ = pysparkfuncs.variant_set.__doc__
+
+
def variant_get(v: "ColumnOrName", path: Union[Column, str], targetType: str)
-> Column:
assert isinstance(path, (Column, str))
if isinstance(path, str):
diff --git a/python/pyspark/sql/functions/__init__.py
b/python/pyspark/sql/functions/__init__.py
index 052ca8a51737..cc1433bbb916 100644
--- a/python/pyspark/sql/functions/__init__.py
+++ b/python/pyspark/sql/functions/__init__.py
@@ -484,6 +484,7 @@ __all__ = [ # noqa: F405
"variant_get",
"variant_insert",
"try_variant_insert",
+ "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 25672e06c97a..e4431aac1a67 100644
--- a/python/pyspark/sql/functions/builtin.py
+++ b/python/pyspark/sql/functions/builtin.py
@@ -21847,6 +21847,70 @@ def try_variant_insert(
)
+@_try_remote_functions
+def 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.
+ Throws an error if a path segment hits a value of an incompatible type.
Returns NULL 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, variant_set
+ >>> df = spark.createDataFrame([{'json': '''{"a": 1, "arr": [1, 2,
3]}'''}])
+ >>> v = parse_json(df.json)
+ >>> df.select(to_json(variant_set(v, "$.a", lit(9))).alias("r")).collect()
+ [Row(r='{"a":9,"arr":[1,2,3]}')]
+ >>> df.select(to_json(variant_set(v, "$.b", lit(2))).alias("r")).collect()
+ [Row(r='{"a":1,"arr":[1,2,3],"b":2}')]
+ >>> df.select(to_json(variant_set(v, "$.arr[1]",
lit(9))).alias("r")).collect()
+ [Row(r='{"a":1,"arr":[1,9,3]}')]
+ >>> df.select(to_json(variant_set(v, "$.b", lit(2),
False)).alias("r")).collect()
+ [Row(r='{"a":1,"arr":[1,2,3]}')]
+ >>> df.select(to_json(variant_set(v, "$.a",
parse_json(lit("null")))).alias("r")).collect()
+ [Row(r='{"a":null,"arr":[1,2,3]}')]
+ >>> df.select(to_json(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(
+ "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_get(v: "ColumnOrName", path: Union[Column, str], targetType: str)
-> Column:
"""
diff --git a/python/pyspark/sql/tests/test_functions.py
b/python/pyspark/sql/tests/test_functions.py
index 3ac206d71f92..f168b30c48ca 100644
--- a/python/pyspark/sql/tests/test_functions.py
+++ b/python/pyspark/sql/tests/test_functions.py
@@ -3540,6 +3540,19 @@ class FunctionsTestsMixin:
)
check(df.select(F.to_json(F.try_variant_insert(v, df.path,
F.lit(9)))), [None, None])
check(df.select(F.to_json(F.try_variant_insert(v, "$.z",
F.lit(None)))), [None, None])
+ check(
+ df.select(F.to_json(F.variant_set(v, "$.z", F.lit(9)))),
+ ['{"a":1,"z":9}', '{"b":2,"z":9}'],
+ )
+ check(
+ df.select(F.to_json(F.variant_set(v, "$.z", F.lit(9), False))),
+ ['{"a":1}', '{"b":2}'],
+ )
+ check(df.select(F.to_json(F.variant_set(v, "$.z", F.lit(None)))),
[None, None])
+ check(
+ 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.schema_of_variant(v)), ["OBJECT<a: BIGINT>",
"OBJECT<b: BIGINT>"])
check(df.select(F.schema_of_variant_agg(v)), ["OBJECT<a: BIGINT, b:
BIGINT>"])
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 cc27a985330a..d15b01ec34c6 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
@@ -9901,6 +9901,92 @@ object functions {
def try_variant_insert(v: Column, path: String, value: Column): Column =
Column.fn("try_variant_insert", 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. Throws an error if a path segment hits a value of an
incompatible type.
+ * Returns NULL 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 variant_set(v: Column, path: Column, value: Column): Column =
+ Column.fn("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. Throws an error if a path segment hits a value of an
incompatible type.
+ * Returns NULL 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 variant_set(v: Column, path: String, value: Column): Column =
+ Column.fn("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. Throws an error if a path segment hits a value of an
incompatible type. Returns
+ * NULL 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 variant_set(v: Column, path: Column, value: Column, createIfMissing:
Boolean): Column =
+ Column.fn("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. Throws an error if a path segment hits a value of an
incompatible type. Returns
+ * NULL 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 variant_set(v: Column, path: String, value: Column, createIfMissing:
Boolean): Column =
+ Column.fn("variant_set", v, lit(path), value, lit(createIfMissing))
+
/**
* Extracts a sub-variant from `v` according to `path` string, and then cast
the sub-variant to
* `targetType`. Returns null if the path does not exist. Throws an
exception if the cast fails.
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 f1b788eb3e81..9f818c109a59 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
@@ -988,6 +988,7 @@ object FunctionRegistry {
expression[VariantDelete]("variant_delete"),
expressionBuilder("variant_insert", VariantInsertExpressionBuilder),
expressionBuilder("try_variant_insert", TryVariantInsertExpressionBuilder),
+ expressionBuilder("variant_set", VariantSetExpressionBuilder),
// Spatial
expression[ST_AsBinary]("st_asbinary"),
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 bf5d22ee7482..47c555c1e691 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
@@ -175,6 +175,47 @@ object VariantExpressionEvalUtils {
insertAtPath(input, javaSegments, pathStr, value, valueDataType,
functionName, failOnError)
}
+ /**
+ * 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.
+ */
+ def setAtPath(
+ input: VariantVal,
+ javaSegments: Array[VariantBuilder.PathSegment],
+ path: String,
+ value: Any,
+ valueDataType: DataType,
+ createIfMissing: Boolean,
+ functionName: String): VariantVal = {
+ val v = new Variant(input.getValue, input.getMetadata)
+ try {
+ val valVal = castToVariant(value, valueDataType)
+ val valVariant = new Variant(valVal.getValue, valVal.getMetadata)
+ val out = VariantBuilder.setAtPath(v, javaSegments, valVariant,
createIfMissing)
+ new VariantVal(out.getValue, out.getMetadata)
+ } catch {
+ case e: VariantPathTypeMismatchException =>
+ throw QueryExecutionErrors.variantPathTypeMismatch(
+ path, renderVariantPath(javaSegments.take(e.depth)), functionName)
+ case _: VariantSizeLimitException =>
+ throw
QueryExecutionErrors.variantSizeLimitError(VariantUtil.SIZE_LIMIT, functionName)
+ }
+ }
+
+ def setAtPath(
+ input: VariantVal,
+ path: UTF8String,
+ value: Any,
+ valueDataType: DataType,
+ createIfMissing: Boolean,
+ functionName: String): VariantVal = {
+ val pathStr = path.toString
+ val javaSegments = toJavaSegments(parseVariantPath(pathStr, functionName))
+ setAtPath(input, javaSegments, pathStr, value, valueDataType,
createIfMissing, functionName)
+ }
+
/** Cast a Spark value from `dataType` into the variant type. */
def castToVariant(input: Any, dataType: DataType): VariantVal = {
// Enforce strict check because it is illegal for input struct/map/variant
to contain duplicate
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 3e404cc9c190..41552cb2d383 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
@@ -1011,6 +1011,177 @@ object VariantInsertExpressionBuilder extends
VariantInsertExpressionBuilderBase
// scalastyle:on line.size.limit
object TryVariantInsertExpressionBuilder extends
VariantInsertExpressionBuilderBase(false)
+case class VariantSet(
+ input: Expression,
+ path: Expression,
+ value: Expression,
+ createIfMissing: Expression)
+ extends QuaternaryExpression
+ with ExpectsInputTypes
+ with QueryErrorsBase {
+
+ override def first: Expression = input
+ override def second: Expression = path
+ override def third: Expression = value
+ override def fourth: Expression = createIfMissing
+
+ override def nullIntolerant: Boolean = true
+
+ override def dataType: DataType = VariantType
+ override def inputTypes: Seq[AbstractDataType] =
+ Seq(
+ VariantType,
+ StringTypeWithCollation(supportsTrimCollation = true),
+ AnyDataType,
+ BooleanType)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ 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
+ }
+ }
+
+ // 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
+ }
+ }
+
+ override protected def nullSafeEval(v: Any, p: Any, valValue: Any, create:
Any): Any = {
+ val inputVariant = v.asInstanceOf[VariantVal]
+ val createIfMissingValue = create.asInstanceOf[Boolean]
+ foldablePath match {
+ case Some(parsed) =>
+ VariantExpressionEvalUtils.setAtPath(
+ inputVariant, parsed.javaSegments, parsed.pathStr, valValue,
value.dataType,
+ createIfMissingValue, prettyName)
+ case None =>
+ VariantExpressionEvalUtils.setAtPath(
+ inputVariant, p.asInstanceOf[UTF8String], valValue, value.dataType,
+ createIfMissingValue, prettyName)
+ }
+ }
+
+ 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)
+ foldablePath match {
+ case Some(parsed) =>
+ val parsedArg = ctx.addReferenceObj("setPath", parsed)
+ s"""
+ |${ev.value} = $cls.setAtPath(
+ | $vVal, $parsedArg.javaSegments(), $parsedArg.pathStr(),
$valVal, $fromArg,
+ | $createVal, "$prettyName");
+ """.stripMargin
+ case None =>
+ s"""
+ |${ev.value} = $cls.setAtPath(
+ | $vVal, $pVal, $valVal, $fromArg, $createVal, "$prettyName");
+ """.stripMargin
+ }
+ })
+ }
+
+ override def prettyName: String = "variant_set"
+
+ override protected def withNewChildrenInternal(
+ newFirst: Expression,
+ newSecond: Expression,
+ newThird: Expression,
+ newFourth: Expression): 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)
+ }
+}
+
+// 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. Throws an error " +
+ "if a path segment hits a value of an incompatible type. Returns NULL 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', parse_json('null'));
+ {"a":null}
+ > SELECT _FUNC_(parse_json('{"a": 1}'), '$.a', null);
+ NULL
+ """,
+ since = "4.3.0",
+ 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)))
+ }
+
+ override def build(funcName: String, expressions: Seq[Expression]):
Expression = {
+ assert(expressions.size == 4)
+ VariantSet(expressions(0), expressions(1), expressions(2), expressions(3))
+ }
+}
+
case class VariantExplode(child: Expression) extends UnaryExpression with
Generator
with ExpectsInputTypes {
override def inputTypes: Seq[AbstractDataType] = Seq(VariantType)
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 ec28ed514c19..1f872d7f17ef 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
@@ -1495,4 +1495,163 @@ class VariantExpressionSuite extends SparkFunSuite with
ExpressionEvalHelper {
"VARIANT_SIZE_LIMIT",
name => Map("sizeLimit" -> "16.0 MiB", "functionName" -> s"`$name`"))
}
+
+ test("variant_set") {
+ def checkSet(
+ input: String,
+ path: String,
+ value: Expression,
+ expected: String,
+ createIfMissing: Boolean = true): Unit = {
+ val expr = VariantSet(
+ Literal(parseJson(input)),
+ Literal.create(path, StringType),
+ value,
+ Literal(createIfMissing))
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(expr, StringType)),
+ expected)
+ }
+
+ // Object sets: replace an existing key, or add a missing one.
+ checkSet("""{"a": 1}""", "$.a", Literal(2), """{"a":2}""")
+ checkSet("""{"a": 1}""", "$.b", Literal(3), """{"a":1,"b":3}""")
+ checkSet("""{"a": {"b": 1}}""", "$.a.b", Literal(9), """{"a":{"b":9}}""")
+
+ // Missing intermediate keys are created when create_if_missing is true.
+ checkSet("{}", "$.a.b", Literal(1), """{"a":{"b":1}}""")
+ checkSet("{}", "$.a[0]", Literal(1), """{"a":[1]}""")
+
+ // Array sets replace the element at the index (no shifting).
+ checkSet("""["a","b","c"]""", "$[1]", Literal("z"), """["a","z","c"]""")
+ // N == length appends; N > length pads with variant nulls.
+ checkSet("""["a","b","c"]""", "$[3]", Literal("z"),
"""["a","b","c","z"]""")
+ checkSet("""["a","b"]""", "$[5]", Literal("z"),
"""["a","b",null,null,null,"z"]""")
+ checkSet("""{"a": [1, 2, 3]}""", "$.a[1]", Literal(9), """{"a":[1,9,3]}""")
+
+ // Array index as an intermediate segment.
+ checkSet("""[{"a": 1}]""", "$[0].b", Literal(2), """[{"a":1,"b":2}]""")
+ checkSet("""[{"a": 1}]""", "$[0].a", Literal(9), """[{"a":9}]""")
+ checkSet("""{"a": [1, 2]}""", "$.a[3].b", Literal(9),
"""{"a":[1,2,null,{"b":9}]}""")
+ // A missing intermediate array is created and padded with variant nulls
up to the index.
+ checkSet("{}", "$.a[2]", Literal(1), """{"a":[null,null,1]}""")
+
+ // scalastyle:off nonascii
+ // Non-ASCII (multi-byte UTF-8).
+ checkSet("""{"你好": 1}""", """$['你好']""", Literal(2), """{"你好":2}""")
+ checkSet("""{"a": 1}""", """$['世界']""", Literal(2), """{"a":1,"世界":2}""")
+ checkSet("{}", """$['café']['über']""", Literal("naïve"),
"""{"café":{"über":"naïve"}}""")
+ // scalastyle:on nonascii
+
+ // create_if_missing = false.
+ checkSet("""{"a": 1, "b": 2}""", "$.a", Literal(99), """{"a":99,"b":2}""",
false)
+ checkSet("""["a","b"]""", "$[1]", Literal("z"), """["a","z"]""", false)
+ checkSet("""{"a": {"b": 1}}""", "$.a.b", Literal(9), """{"a":{"b":9}}""",
false)
+ checkSet("""{"a": [1, 2]}""", "$.a[1]", Literal(9), """{"a":[1,9]}""",
false)
+ checkSet("""[{"a": 1}]""", "$[0].a", Literal(9), """[{"a":9}]""", false)
+ checkSet("""{"a": {"b": [1, 2]}}""", "$.a.b[0]", Literal(9),
"""{"a":{"b":[9,2]}}""", false)
+ checkSet("""{"a": 1}""", "$.b", Literal(2), """{"a":1}""", false)
+ checkSet("""{"a": {"b": 1}}""", "$.a.c", Literal(2), """{"a":{"b":1}}""",
false)
+ checkSet("""{"a": [1, 2]}""", "$.a[5]", Literal(9), """{"a":[1,2]}""",
false)
+
+ // Set a variant null vs. a verbatim string vs. structured JSON.
+ checkSet("""{"a": 1}""", "$.a", Literal(parseJson("null")),
"""{"a":null}""")
+ checkSet("{}", "$.a", Literal("""{"x":1}"""), """{"a":"{\"x\":1}"}""")
+ checkSet("{}", "$.a", Literal(parseJson("""{"x":1}""")),
"""{"a":{"x":1}}""")
+ 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)
+ 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))
+
+ // Type mismatch throws regardless of create_if_missing.
+ Seq(true, false).foreach { create =>
+ checkErrorInExpression[SparkRuntimeException](
+ VariantSet(Literal(parseJson("""{"a": 1}""")), Literal("$.a.b"),
Literal(2),
+ Literal(create)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$.a.b", "failedAt" -> "$.a", "functionName" ->
"`variant_set`"))
+ }
+ checkErrorInExpression[SparkRuntimeException](
+ VariantSet(Literal(parseJson("5")), Literal("$.a"), Literal(2),
Literal(true)),
+ "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)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$[0]", "failedAt" -> "$", "functionName" ->
"`variant_set`"))
+ checkErrorInExpression[SparkRuntimeException](
+ VariantSet(Literal(parseJson("[1, 2]")), Literal("$.a"), Literal(2),
Literal(true)),
+ "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)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$['a.b'].c", "failedAt" -> "$['a.b']", "functionName" ->
"`variant_set`"))
+
+ // Structs and maps are rejected at analysis; arrays of scalars remain
allowed.
+ Seq(
+ 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)
+ }
+
+ val structVal = ToVariantObject(Literal.create(create_row(1, "x"),
+ StructType(Array(StructField("a", IntegerType), StructField("b",
StringType)))))
+ val mapVal = ToVariantObject(Literal.create(Map("z" -> 1, "y" -> 2, "x" ->
3)))
+ Seq(structVal, mapVal).foreach { v =>
+ assert(
+ VariantSet(Literal(parseJson("{}")), Literal("$.k"), v, Literal(true))
+ .checkInputDataTypes().isSuccess)
+ }
+ checkSet("""{"k": 1}""", "$.k", structVal, """{"k":{"a":1,"b":"x"}}""")
+ 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)),
+ "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`"))
+ val tooBig = "x".repeat(16 * 1024 * 1024)
+ checkErrorInExpression[SparkRuntimeException](
+ VariantSet(Literal(parseJson("{}")), Literal("$.a"), Literal(tooBig),
Literal(true)),
+ "VARIANT_SIZE_LIMIT",
+ Map("sizeLimit" -> "16.0 MiB", "functionName" -> "`variant_set`"))
+ }
}
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 0222a40219ae..747e348891bc 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
@@ -2771,6 +2771,14 @@ class PlanGenerationTestSuite extends ConnectFunSuite
with Logging {
fn.try_variant_insert(fn.parse_json(fn.col("g")), "$.a", fn.lit(1))
}
+ functionTest("variant_set") {
+ fn.variant_set(fn.parse_json(fn.col("g")), "$.a", fn.lit(1))
+ }
+
+ functionTest("variant_set with create_if_missing") {
+ fn.variant_set(fn.parse_json(fn.col("g")), "$.a", fn.lit(1), false)
+ }
+
functionTest("variant_get") {
fn.variant_get(fn.parse_json(fn.col("g")), "$", "int")
}
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
new file mode 100644
index 000000000000..3f695273d79b
--- /dev/null
+++
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_set.explain
@@ -0,0 +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]
++- 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
new file mode 100644
index 000000000000..2ad5e4938e92
--- /dev/null
+++
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_set_with_create_if_missing.explain
@@ -0,0 +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]
++- 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_variant_set.json
b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_set.json
new file mode 100644
index 000000000000..ef8e00fe0acb
--- /dev/null
+++
b/sql/connect/common/src/test/resources/query-tests/queries/function_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": "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": "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": "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_variant_set.proto.bin
b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_set.proto.bin
new file mode 100644
index 000000000000..899cc1c6b0b5
Binary files /dev/null and
b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_set.proto.bin
differ
diff --git
a/sql/connect/common/src/test/resources/query-tests/queries/function_variant_set_with_create_if_missing.json
b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_set_with_create_if_missing.json
new file mode 100644
index 000000000000..cb795edc537b
--- /dev/null
+++
b/sql/connect/common/src/test/resources/query-tests/queries/function_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": "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": "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": "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": "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_variant_set_with_create_if_missing.proto.bin
b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_set_with_create_if_missing.proto.bin
new file mode 100644
index 000000000000..02d73bc21393
Binary files /dev/null and
b/sql/connect/common/src/test/resources/query-tests/queries/function_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 ecdb01a2e034..d999d9c96e79 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
@@ -559,6 +559,7 @@
| 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> |
|
org.apache.spark.sql.catalyst.expressions.variant.VariantInsertExpressionBuilder
| variant_insert | SELECT variant_insert(parse_json('{"a": 1}'), '$.b', 2) |
struct<variant_insert(parse_json({"a": 1}), $.b, 2):variant> |
+|
org.apache.spark.sql.catalyst.expressions.variant.VariantSetExpressionBuilder |
variant_set | SELECT variant_set(parse_json('{"a": 1}'), '$.a', 2) |
struct<variant_set(parse_json({"a": 1}), $.a, 2, true):variant> |
| org.apache.spark.sql.catalyst.expressions.xml.XPathBoolean | xpath_boolean |
SELECT xpath_boolean('<a><b>1</b></a>','a/b') |
struct<xpath_boolean(<a><b>1</b></a>, a/b):boolean> |
| org.apache.spark.sql.catalyst.expressions.xml.XPathDouble | xpath_double |
SELECT xpath_double('<a><b>1</b><b>2</b></a>', 'sum(a/b)') |
struct<xpath_double(<a><b>1</b><b>2</b></a>, sum(a/b)):double> |
| org.apache.spark.sql.catalyst.expressions.xml.XPathDouble | xpath_number |
SELECT xpath_number('<a><b>1</b><b>2</b></a>', 'sum(a/b)') |
struct<xpath_number(<a><b>1</b><b>2</b></a>, sum(a/b)):double> |
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 8d25633887f9..300479cf5d49 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
@@ -439,6 +439,88 @@ class VariantSuite extends SharedSparkSession with
ExpressionEvalHelper {
}
}
+ test("variant_set with literal arguments") {
+ def rows(results: Any*): Seq[Row] = results.map(Row(_))
+
+ // A basic invocation, exercising the SQL parse/registration path end to
end.
+ checkAnswer(
+ sql("SELECT to_json(variant_set(parse_json('{\"a\": 1}'), '$.a', 2))"),
+ rows("""{"a":2}"""))
+
+ // The optional create_if_missing arg, positionally and via named-argument
syntax.
+ checkAnswer(
+ sql("SELECT to_json(variant_set(parse_json('{\"a\": 1}'), '$.b', 2,
false))"),
+ rows("""{"a":1}"""))
+ checkAnswer(
+ sql("SELECT to_json(variant_set(parse_json('{\"a\": 1}'), '$.a', 2, " +
+ "create_if_missing => false))"),
+ rows("""{"a":2}"""))
+
+ // NULL-intolerant.
+ checkAnswer(
+ sql("SELECT to_json(variant_set(CAST(NULL AS VARIANT), '$.a', 1))"),
+ rows(null))
+ checkAnswer(
+ sql("SELECT to_json(variant_set(parse_json('{\"a\": 1}'), '$.b',
NULL))"),
+ rows(null))
+ checkAnswer(
+ sql("SELECT to_json(variant_set(parse_json('{\"a\": 1}'), NULL, 1))"),
+ rows(null))
+
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT variant_set(parse_json('{\"a\": 1}'), '$.a.b',
2)").collect()
+ },
+ condition = "VARIANT_PATH_TYPE_MISMATCH",
+ parameters = Map(
+ "path" -> "$.a.b", "failedAt" -> "$.a", "functionName" ->
toSQLId("variant_set")))
+
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT variant_set(parse_json('{}'), '$', 1)").collect()
+ },
+ condition = "INVALID_VARIANT_PATH",
+ parameters = Map("path" -> "$", "functionName" ->
toSQLId("variant_set")))
+
+ // A raw struct/map value is not castable to variant and is rejected at
analysis.
+ assert(intercept[AnalysisException] {
+ sql("SELECT variant_set(parse_json('{}'), '$.a', named_struct('x', 1))")
+ }.getCondition == "DATATYPE_MISMATCH.CAST_WITHOUT_SUGGESTION")
+ }
+
+ test("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","b"]""", "$[0]", 9),
+ (null, "$.a", 2),
+ ("""{"a": 1}""", "$.b", 2)
+ ).toDF("json", "path", "val")
+ val v = parse_json(col("json"))
+ // Default create_if_missing = true: the last row's missing `$.b`
target is created.
+ checkAnswer(
+ df.select(to_json(variant_set(v, col("path"),
col("val"))).alias("r")),
+ rows("""{"a":2}""", """[9,"b"]""", null, """{"a":1,"b":2}"""))
+ // create_if_missing = false (Column-path overload): only the last row
differs.
+ checkAnswer(
+ df.select(to_json(variant_set(v, col("path"), col("val"),
false)).alias("r")),
+ rows("""{"a":2}""", """[9,"b"]""", null, """{"a":1}"""))
+
+ // String-path overloads, with and without the boolean flag.
+ val objDf = Seq(("""{"a": 1}""", 2), (null, 2)).toDF("json", "val")
+ val objV = parse_json(objDf("json"))
+ checkAnswer(
+ objDf.select(to_json(variant_set(objV, "$.a",
col("val"))).alias("r")),
+ rows("""{"a":2}""", null))
+ checkAnswer(
+ objDf.select(to_json(variant_set(objV, "$.b", col("val"),
false)).alias("r")),
+ rows("""{"a":1}""", null))
+ }
+ }
+ }
+
test("round trip tests") {
withSQLConf(SQLConf.VARIANT_INFER_SHREDDING_SCHEMA.key -> "false") {
val rand = new Random(42)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
