This is an automated email from the ASF dual-hosted git repository.
cloud-fan 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 b830367a0d5f [SPARK-57641][SQL] Add `variant_insert` expression
b830367a0d5f is described below
commit b830367a0d5fba3d07452fc9a73e47a3293f51e1
Author: bojana-db <[email protected]>
AuthorDate: Wed Jul 8 00:49:12 2026 +0800
[SPARK-57641][SQL] Add `variant_insert` expression
### What changes were proposed in this pull request?
Adds the SQL function `variant_insert(v, path, val)`, which inserts a value
into a Variant value at a single JSONPath location.
Details:
- Object path (e.g. `$.b`): adds a new field; errors with
`VARIANT_DUPLICATE_KEY`
if the key already exists;
- Array path (e.g. `$[N]`): inserts at index `N`, shifting later elements
right; an
index past the end pads the gap with variant nulls;
- Missing intermediate keys/indices along the path are created;
- Any NULL argument returns NULL;
- The value may be any expression castable to variant;
- A new error condition `VARIANT_PATH_TYPE_MISMATCH` is raised when a path
segment
is applied to an incompatible value (e.g. an object key on an array); the
root
path `$` is rejected with `INVALID_VARIANT_PATH`, and results exceeding
the size
limit raise `VARIANT_SIZE_LIMIT`.
### Why are the changes needed?
There is no way to add a field or element to a variant; today the only
option is to convert it to another datatype (e.g. map), mutate, and convert
back.
### Does this PR introduce _any_ user-facing change?
Yes, a new SQL function (and Scala/Python `functions` API) plus a new error
condition `VARIANT_PATH_TYPE_MISMATCH`.
### 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 #56703 from bojana-db/variant-insert.
Authored-by: bojana-db <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
.../src/main/resources/error/error-conditions.json | 6 +
.../apache/spark/types/variant/VariantBuilder.java | 154 +++++++++++++++++++++
.../variant/VariantPathTypeMismatchException.java | 31 +++++
.../source/reference/pyspark.sql/functions.rst | 1 +
python/pyspark/sql/connect/functions/builtin.py | 8 ++
python/pyspark/sql/functions/__init__.py | 1 +
python/pyspark/sql/functions/builtin.py | 60 ++++++++
python/pyspark/sql/tests/test_functions.py | 14 +-
.../scala/org/apache/spark/sql/functions.scala | 38 +++++
.../sql/catalyst/analysis/FunctionRegistry.scala | 1 +
.../variant/VariantExpressionEvalUtils.scala | 73 +++++++++-
.../expressions/variant/variantExpressions.scala | 138 +++++++++++++++++-
.../spark/sql/errors/QueryExecutionErrors.scala | 7 +
.../variant/VariantExpressionSuite.scala | 139 +++++++++++++++++++
.../apache/spark/sql/PlanGenerationTestSuite.scala | 4 +
.../function_variant_insert.explain | 2 +
.../queries/function_variant_insert.json | 125 +++++++++++++++++
.../queries/function_variant_insert.proto.bin | Bin 0 -> 1139 bytes
.../sql-functions/sql-expression-schema.md | 1 +
.../scala/org/apache/spark/sql/VariantSuite.scala | 89 ++++++++++++
20 files changed, 883 insertions(+), 9 deletions(-)
diff --git a/common/utils/src/main/resources/error/error-conditions.json
b/common/utils/src/main/resources/error/error-conditions.json
index cc2d223a74b0..dd3220d10b28 100644
--- a/common/utils/src/main/resources/error/error-conditions.json
+++ b/common/utils/src/main/resources/error/error-conditions.json
@@ -9160,6 +9160,12 @@
],
"sqlState" : "22023"
},
+ "VARIANT_PATH_TYPE_MISMATCH" : {
+ "message" : [
+ "The path '<path>' is incompatible with the variant value at
'<failedAt>' in <functionName>."
+ ],
+ "sqlState" : "22023"
+ },
"VARIANT_SIZE_LIMIT" : {
"message" : [
"Cannot build variant bigger than <sizeLimit> in <functionName>.",
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 79f869ba6be7..9685f0a35f85 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
@@ -130,6 +130,20 @@ public class VariantBuilder {
return builder.result();
}
+ // Return a new variant equal to `v` with `val` inserted at `segments`
(which must be non-empty).
+ // An object leaf adds a field (or throws VARIANT_DUPLICATE_KEY if the key
exists); an array leaf
+ // inserts at the index, shifting elements right and padding with nulls past
the end. Missing
+ // intermediate keys/indices are created. A segment applied to an
incompatible value throws
+ // VariantPathTypeMismatchException, which the caller maps to
VARIANT_PATH_TYPE_MISMATCH.
+ public static Variant insertAtPath(Variant v, PathSegment[] segments,
Variant val) {
+ if (segments.length == 0) {
+ throw new IllegalArgumentException("Segments must be non-empty");
+ }
+ VariantBuilder builder = new VariantBuilder(false);
+ builder.appendWithInsertionImpl(v.value, v.metadata, v.pos, segments, 0,
val);
+ return builder.result();
+ }
+
// Build the variant metadata from `dictionaryKeys` and return the variant
result.
public Variant result() {
int numKeys = dictionaryKeys.size();
@@ -549,6 +563,146 @@ public class VariantBuilder {
}
}
+ private void appendWithInsertionImpl(
+ byte[] value, byte[] metadata, int pos, PathSegment[] segments, int
depth, Variant val) {
+ 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) {
+ appendWithInsertionImpl(value, metadata, elementPos, segments,
depth + 1, val);
+ } else {
+ appendVariantImpl(value, metadata, elementPos);
+ }
+ }
+ if (isLast) {
+ // Append the new field unconditionally; if `key` already exists,
the duplicate is
+ // detected by `finishWritingObject`, which raises
VARIANT_DUPLICATE_KEY.
+ int newId = addKey(key);
+ fields.add(new FieldEntry(key, newId, writePos - start));
+ appendVariant(val);
+ } else if (!found) {
+ // Intermediate object key is missing; create the rest of the path
under it.
+ int newId = addKey(key);
+ fields.add(new FieldEntry(key, newId, writePos - start));
+ 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 (isLast) {
+ // Insert `val` at `index`, shifting existing elements right.
+ for (int i = 0; i < size; ++i) {
+ if (i == index) {
+ offsets.add(writePos - start);
+ appendVariant(val);
+ }
+ int offset = readUnsigned(value, offsetStart + offsetSize * i,
offsetSize);
+ offsets.add(writePos - start);
+ appendVariantImpl(value, metadata, dataStart + offset);
+ }
+ if (index >= size) {
+ // Pad with variant nulls up to `index`, then append `val`.
+ for (int i = size; i < index; ++i) {
+ offsets.add(writePos - start);
+ appendNull();
+ }
+ offsets.add(writePos - start);
+ appendVariant(val);
+ }
+ } else if (index < size) {
+ // Descend into the existing element at `index`.
+ 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) {
+ appendWithInsertionImpl(value, metadata, elementPos, segments,
depth + 1, val);
+ } else {
+ appendVariantImpl(value, metadata, elementPos);
+ }
+ }
+ } else {
+ // Intermediate index is past the end; copy existing elements, pad
with variant nulls up
+ // to `index`, then create the rest of the path at `index`.
+ for (int i = 0; i < size; ++i) {
+ int offset = readUnsigned(value, offsetStart + offsetSize * i,
offsetSize);
+ offsets.add(writePos - start);
+ appendVariantImpl(value, metadata, dataStart + offset);
+ }
+ for (int i = size; i < index; ++i) {
+ offsets.add(writePos - start);
+ appendNull();
+ }
+ offsets.add(writePos - start);
+ 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
+ // array-index segment creates an array padded with variant nulls up to the
index.
+ private void appendNewPath(PathSegment[] segments, int depth, Variant val) {
+ PathSegment seg = segments[depth];
+ boolean isLast = depth == segments.length - 1;
+ if (seg instanceof ObjectKeySegment) {
+ String key = ((ObjectKeySegment) seg).key;
+ ArrayList<FieldEntry> fields = new ArrayList<>(1);
+ int start = writePos;
+ int id = addKey(key);
+ fields.add(new FieldEntry(key, id, writePos - start));
+ if (isLast) {
+ appendVariant(val);
+ } else {
+ appendNewPath(segments, depth + 1, val);
+ }
+ finishWritingObject(start, fields);
+ } else {
+ int index = ((ArrayIndexSegment) seg).index;
+ ArrayList<Integer> offsets = new ArrayList<>();
+ int start = writePos;
+ for (int i = 0; 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);
+ }
+ }
+
// Append the variant value without rewriting or creating any metadata. This
is used when
// building an object during shredding, where there is a fixed pre-existing
metadata that
// all shredded values will refer to.
diff --git
a/common/variant/src/main/java/org/apache/spark/types/variant/VariantPathTypeMismatchException.java
b/common/variant/src/main/java/org/apache/spark/types/variant/VariantPathTypeMismatchException.java
new file mode 100644
index 000000000000..83db5e5b83e1
--- /dev/null
+++
b/common/variant/src/main/java/org/apache/spark/types/variant/VariantPathTypeMismatchException.java
@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.types.variant;
+
+/**
+ * An exception indicating that an insertion path segment is incompatible with
the value it targets
+ * (e.g. an object key on a scalar, or an array index on an object). {@code
depth} is the index of
+ * the failing segment.
+ */
+public class VariantPathTypeMismatchException extends RuntimeException {
+ public final int depth;
+
+ public VariantPathTypeMismatchException(int depth) {
+ this.depth = depth;
+ }
+}
diff --git a/python/docs/source/reference/pyspark.sql/functions.rst
b/python/docs/source/reference/pyspark.sql/functions.rst
index 45e0eaff0296..70e5811180d0 100644
--- a/python/docs/source/reference/pyspark.sql/functions.rst
+++ b/python/docs/source/reference/pyspark.sql/functions.rst
@@ -604,6 +604,7 @@ VARIANT Functions
try_variant_get
variant_delete
variant_get
+ variant_insert
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 7a4d2b0e1b6c..8e74aeef2828 100644
--- a/python/pyspark/sql/connect/functions/builtin.py
+++ b/python/pyspark/sql/connect/functions/builtin.py
@@ -2219,6 +2219,14 @@ def variant_delete(v: "ColumnOrName", *paths:
Union[Column, str]) -> Column:
variant_delete.__doc__ = pysparkfuncs.variant_delete.__doc__
+def variant_insert(v: "ColumnOrName", path: Union[Column, str], value:
"ColumnOrName") -> Column:
+ path_col = path if isinstance(path, Column) else lit(path)
+ return _invoke_function("variant_insert", _to_col(v), path_col,
_to_col(value))
+
+
+variant_insert.__doc__ = pysparkfuncs.variant_insert.__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 5083d5c0db16..ee9d06fcaa2d 100644
--- a/python/pyspark/sql/functions/__init__.py
+++ b/python/pyspark/sql/functions/__init__.py
@@ -482,6 +482,7 @@ __all__ = [ # noqa: F405
"try_variant_get",
"variant_delete",
"variant_get",
+ "variant_insert",
"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 44e483ecfb6b..1e348cfd4287 100644
--- a/python/pyspark/sql/functions/builtin.py
+++ b/python/pyspark/sql/functions/builtin.py
@@ -21726,6 +21726,66 @@ def variant_delete(v: "ColumnOrName", *paths:
Union[Column, str]) -> Column:
)
+@_try_remote_functions
+def variant_insert(v: "ColumnOrName", path: Union[Column, str], value:
"ColumnOrName") -> Column:
+ """
+ Inserts a value into a variant at the given JSONPath location. An object
path adds a new field
+ (error if it already exists); an array path inserts at the index, shifting
later elements
+ right. Missing intermediate keys are created. 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 insertion 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 insert. Any expression castable to variant.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ a variant column with `value` inserted at `path`
+
+ Examples
+ --------
+ >>> from pyspark.sql.functions import lit, parse_json, to_json,
variant_insert
+ >>> df = spark.createDataFrame([{
+ ... 'json': '''{ "a": 1, "arr": ["x", "y"] }''',
+ ... 'path': '$.d'
+ ... }])
+ >>> v = parse_json(df.json)
+ >>> df.select(to_json(variant_insert(v, "$.b",
lit(2))).alias("r")).collect()
+ [Row(r='{"a":1,"arr":["x","y"],"b":2}')]
+ >>> df.select(to_json(variant_insert(v, "$.c.d",
lit(3))).alias("r")).collect()
+ [Row(r='{"a":1,"arr":["x","y"],"c":{"d":3}}')]
+ >>> df.select(to_json(variant_insert(v, "$.arr[1]",
lit("z"))).alias("r")).collect()
+ [Row(r='{"a":1,"arr":["x","z","y"]}')]
+ >>> df.select(to_json(variant_insert(v, "$.arr[5]",
lit("z"))).alias("r")).collect()
+ [Row(r='{"a":1,"arr":["x","y",null,null,null,"z"]}')]
+ >>> df.select(to_json(variant_insert(v, df.path,
lit(9))).alias("r")).collect()
+ [Row(r='{"a":1,"arr":["x","y"],"d":9}')]
+ >>> df.select(to_json(variant_insert(v, "$.b",
parse_json(lit('null')))).alias("r")).collect()
+ [Row(r='{"a":1,"arr":["x","y"],"b":null}')]
+ >>> df.select(variant_insert(v, "$.b", 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_insert",
+ _to_java_column(v),
+ _to_java_column(path_col),
+ _to_java_column(value),
+ )
+
+
@_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 8599d0dd46e1..2ade286763de 100644
--- a/python/pyspark/sql/tests/test_functions.py
+++ b/python/pyspark/sql/tests/test_functions.py
@@ -3507,7 +3507,10 @@ class FunctionsTestsMixin:
def test_variant_expressions(self):
df = self.spark.createDataFrame(
- [Row(json="""{ "a" : 1 }""", path="$.a"), Row(json="""{ "b" : 2
}""", path="$.b")]
+ [
+ Row(json="""{ "a" : 1 }""", path="$.a", newpath="$.z"),
+ Row(json="""{ "b" : 2 }""", path="$.b", newpath="$.z"),
+ ]
)
v = F.parse_json(df.json)
@@ -3522,6 +3525,15 @@ class FunctionsTestsMixin:
df.select(F.to_json(F.variant_delete(v, F.lit(None)))),
['{"a":1}', '{"b":2}'],
)
+ check(
+ df.select(F.to_json(F.variant_insert(v, "$.z", F.lit(9)))),
+ ['{"a":1,"z":9}', '{"b":2,"z":9}'],
+ )
+ check(df.select(F.to_json(F.variant_insert(v, "$.z", F.lit(None)))),
[None, None])
+ check(
+ df.select(F.to_json(F.variant_insert(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 0dc667a00f6d..c8b269a63de5 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
@@ -9821,6 +9821,44 @@ object functions {
def variant_delete(v: Column, path: String, paths: String*): Column =
Column.fn("variant_delete", (v +: lit(path) +: paths.map(lit)): _*)
+ /**
+ * Inserts a value into a variant at the given JSONPath location. An object
path adds a new
+ * field (error if it already exists); an array path inserts at the index,
shifting later
+ * elements right. Missing intermediate keys are created. Returns NULL if
any argument is NULL.
+ *
+ * @param v
+ * a variant column.
+ * @param path
+ * the column containing the JSONPath string identifying the insertion
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 insert. Any expression castable to variant.
+ * @group variant_funcs
+ * @since 4.3.0
+ */
+ def variant_insert(v: Column, path: Column, value: Column): Column =
+ Column.fn("variant_insert", v, path, value)
+
+ /**
+ * Inserts a value into a variant at the given JSONPath location. An object
path adds a new
+ * field (error if it already exists); an array path inserts at the index,
shifting later
+ * elements right. Missing intermediate keys are created. Returns NULL if
any argument is NULL.
+ *
+ * @param v
+ * a variant column.
+ * @param path
+ * the JSONPath identifying the insertion 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 insert. Any expression castable to variant.
+ * @group variant_funcs
+ * @since 4.3.0
+ */
+ def variant_insert(v: Column, path: String, value: Column): Column =
+ Column.fn("variant_insert", v, lit(path), value)
+
/**
* 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 a8654491a269..193e939889db 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
@@ -986,6 +986,7 @@ object FunctionRegistry {
expression[ToVariantObject]("to_variant_object"),
expression[IsValidVariant]("is_valid_variant"),
expression[VariantDelete]("variant_delete"),
+ expression[VariantInsert]("variant_insert"),
// 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 47afbcd78837..20a55bccc779 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
@@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.util.{ArrayData, MapData}
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.types._
-import org.apache.spark.types.variant.{Variant, VariantBuilder,
VariantSizeLimitException, VariantUtil}
+import org.apache.spark.types.variant.{Variant, VariantBuilder,
VariantPathTypeMismatchException, VariantSizeLimitException, VariantUtil}
import org.apache.spark.unsafe.types.{UTF8String, VariantVal}
/**
@@ -75,17 +75,40 @@ object VariantExpressionEvalUtils {
def isValidVariant(input: VariantVal): Boolean =
VariantUtil.isValidVariant(input.getValue, input.getMetadata)
- /** Throws `INVALID_VARIANT_PATH` on a malformed path or on the empty (root
`$`) path. */
- def parseVariantDeletePath(pathValue: String): Array[VariantPathSegment] = {
+ /**
+ * Parse a JSONPath for a variant manipulation function. Throws
`INVALID_VARIANT_PATH` on a
+ * malformed path or on the empty (root `$`) path.
+ */
+ def parseVariantPath(pathValue: String, functionName: String):
Array[VariantPathSegment] = {
val parsed = VariantPathParser.parse(pathValue).getOrElse {
- throw QueryExecutionErrors.invalidVariantPath(pathValue,
"variant_delete")
+ throw QueryExecutionErrors.invalidVariantPath(pathValue, functionName)
}
if (parsed.isEmpty) {
- throw QueryExecutionErrors.invalidVariantPath(pathValue,
"variant_delete")
+ throw QueryExecutionErrors.invalidVariantPath(pathValue, functionName)
}
parsed
}
+ /** Render a parsed path prefix back to a JSONPath string for error
messages. */
+ private def renderVariantPath(segments: Array[VariantBuilder.PathSegment]):
String = {
+ val sb = new StringBuilder("$")
+ segments.foreach {
+ case o: VariantBuilder.ObjectKeySegment =>
+ val key = o.key
+ // Dot notation only parses keys with no `.` or `[` (and at least one
char); anything else
+ // must use bracket notation so the rendered path round-trips to the
same segments.
+ if (key.nonEmpty && !key.contains('.') && !key.contains('[')) {
+ sb.append('.').append(key)
+ } else if (!key.contains('\'')) {
+ sb.append("['").append(key).append("']")
+ } else {
+ sb.append("[\"").append(key).append("\"]")
+ }
+ case a: VariantBuilder.ArrayIndexSegment =>
sb.append('[').append(a.index).append(']')
+ }
+ sb.toString
+ }
+
def toJavaSegments(
segments: Array[VariantPathSegment]): Array[VariantBuilder.PathSegment]
= {
segments.map {
@@ -103,7 +126,45 @@ object VariantExpressionEvalUtils {
}
def deleteAtPath(input: VariantVal, path: UTF8String): VariantVal =
- deleteAtPath(input, toJavaSegments(parseVariantDeletePath(path.toString)))
+ deleteAtPath(input, toJavaSegments(parseVariantPath(path.toString,
"variant_delete")))
+
+ /**
+ * Insert `value` into `input` at `javaSegments`. `path` is the source
string used in error
+ * messages. The cast and insert share one try, so any size overflow maps to
`VARIANT_SIZE_LIMIT`
+ * and a type mismatch maps to `VARIANT_PATH_TYPE_MISMATCH`.
+ */
+ def insertAtPath(
+ input: VariantVal,
+ javaSegments: Array[VariantBuilder.PathSegment],
+ path: String,
+ value: Any,
+ valueDataType: DataType,
+ 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.insertAtPath(v, javaSegments, valVariant)
+ 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 insertAtPath(
+ input: VariantVal,
+ path: UTF8String,
+ value: Any,
+ valueDataType: DataType,
+ functionName: String): VariantVal = {
+ val pathStr = path.toString
+ val javaSegments = toJavaSegments(parseVariantPath(pathStr, functionName))
+ insertAtPath(input, javaSegments, pathStr, value, valueDataType,
functionName)
+ }
/** Cast a Spark value from `dataType` into the variant type. */
def castToVariant(input: Any, dataType: DataType): VariantVal = {
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 e0ac976e5100..0f8995ce829f 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
@@ -799,8 +799,8 @@ object VariantDelete {
if (v == null) {
None
} else {
- Some(ParsedDeletePath(
-
VariantExpressionEvalUtils.parseVariantDeletePath(v.asInstanceOf[UTF8String].toString)))
+ Some(ParsedDeletePath(VariantExpressionEvalUtils.parseVariantPath(
+ v.asInstanceOf[UTF8String].toString, "variant_delete")))
}
} else {
Some(DynamicDeletePath(child))
@@ -808,6 +808,140 @@ object VariantDelete {
}
}
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = "_FUNC_(v, path, val) - Inserts a value into a variant at the given
JSONPath " +
+ "location. An object path adds a new field (error if it already exists);
an array path " +
+ "inserts at the index, shifting later elements right. Missing intermediate
keys are " +
+ "created. 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
insertion 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.
+ """,
+ examples = """
+ Examples:
+ > SELECT _FUNC_(parse_json('{"a": 1}'), '$.b', 2);
+ {"a":1,"b":2}
+ > SELECT _FUNC_(parse_json('{}'), '$.a.b', 1);
+ {"a":{"b":1}}
+ > SELECT _FUNC_(parse_json('["a","b","c"]'), '$[1]', 'z');
+ ["a","z","b","c"]
+ > SELECT _FUNC_(parse_json('["a","b"]'), '$[5]', 'z');
+ ["a","b",null,null,null,"z"]
+ > SELECT _FUNC_(parse_json('{}'), '$.a', parse_json('{"x":1}'));
+ {"a":{"x":1}}
+ > SELECT _FUNC_(parse_json('{}'), '$.a', parse_json('null'));
+ {"a":null}
+ > SELECT _FUNC_(parse_json('{}'), '$.a', null);
+ NULL
+ """,
+ since = "4.3.0",
+ group = "variant_funcs"
+)
+// scalastyle:on line.size.limit
+case class VariantInsert(input: Expression, path: Expression, value:
Expression)
+ extends TernaryExpression
+ with ExpectsInputTypes
+ with QueryErrorsBase {
+
+ override def first: Expression = input
+ override def second: Expression = path
+ override def third: Expression = value
+
+ override def nullIntolerant: Boolean = true
+
+ override def dataType: DataType = VariantType
+ override def inputTypes: Seq[AbstractDataType] =
+ Seq(VariantType, StringTypeWithCollation(supportsTrimCollation = true),
AnyDataType)
+
+ 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 `ParsedInsertPath`), avoiding a
per-row conversion.
+ // `None` means the path is dynamic (or a foldable NULL, which makes the
whole expression NULL and
+ // is never evaluated).
+ @transient private lazy val foldablePath:
Option[VariantInsert.ParsedInsertPath] = {
+ if (path.foldable) {
+ val p = path.eval()
+ if (p == null) {
+ None
+ } else {
+ val s = p.asInstanceOf[UTF8String].toString
+ Some(VariantInsert.ParsedInsertPath(
+ VariantExpressionEvalUtils.parseVariantPath(s, prettyName), s))
+ }
+ } else {
+ None
+ }
+ }
+
+ override protected def nullSafeEval(v: Any, p: Any, valValue: Any): Any = {
+ val inputVariant = v.asInstanceOf[VariantVal]
+ foldablePath match {
+ case Some(parsed) =>
+ VariantExpressionEvalUtils.insertAtPath(
+ inputVariant, parsed.javaSegments, parsed.pathStr, valValue,
value.dataType, prettyName)
+ case None =>
+ VariantExpressionEvalUtils.insertAtPath(
+ inputVariant, p.asInstanceOf[UTF8String], valValue, value.dataType,
prettyName)
+ }
+ }
+
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode = {
+ val cls = VariantExpressionEvalUtils.getClass.getName.stripSuffix("$")
+ nullSafeCodeGen(ctx, ev, (vVal, pVal, valVal) => {
+ val fromArg = ctx.addReferenceObj("from", value.dataType)
+ foldablePath match {
+ case Some(parsed) =>
+ val parsedArg = ctx.addReferenceObj("insertPath", parsed)
+ s"""
+ |${ev.value} = $cls.insertAtPath(
+ | $vVal, $parsedArg.javaSegments(), $parsedArg.pathStr(),
$valVal, $fromArg,
+ | "$prettyName");
+ """.stripMargin
+ case None =>
+ s"""
+ |${ev.value} = $cls.insertAtPath($vVal, $pVal, $valVal, $fromArg,
"$prettyName");
+ """.stripMargin
+ }
+ })
+ }
+
+ override def prettyName: String = "variant_insert"
+
+ override protected def withNewChildrenInternal(
+ newFirst: Expression, newSecond: Expression, newThird: Expression):
VariantInsert =
+ copy(input = newFirst, path = newSecond, value = newThird)
+}
+
+object VariantInsert {
+ // Caches a foldable path. `VariantBuilder.PathSegment` is not
`Serializable`, so the Java form is
+ // `@transient` and re-derived once per executor task after deserialization.
`pathStr` is the
+ // source string, retained for error messages.
+ case class ParsedInsertPath(segments: Array[VariantPathSegment], pathStr:
String) {
+ @transient lazy val javaSegments: Array[VariantBuilder.PathSegment] =
+ VariantExpressionEvalUtils.toJavaSegments(segments)
+ }
+}
+
case class VariantExplode(child: Expression) extends UnaryExpression with
Generator
with ExpectsInputTypes {
override def inputTypes: Seq[AbstractDataType] = Seq(VariantType)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
index 71ae98253853..1ffe59d6ab2a 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
@@ -3096,6 +3096,13 @@ private[sql] object QueryExecutionErrors extends
QueryErrorsBase with ExecutionE
messageParameters = Map("path" -> path, "functionName" ->
toSQLId(functionName)))
}
+ def variantPathTypeMismatch(path: String, failedAt: String, functionName:
String): Throwable = {
+ new SparkRuntimeException(
+ errorClass = "VARIANT_PATH_TYPE_MISMATCH",
+ messageParameters =
+ Map("path" -> path, "failedAt" -> failedAt, "functionName" ->
toSQLId(functionName)))
+ }
+
def malformedVariant(): Throwable = new SparkRuntimeException(
"MALFORMED_VARIANT",
Map.empty
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 73fdc48fa585..6f1a410c5bb9 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
@@ -1326,4 +1326,143 @@ class VariantExpressionSuite extends SparkFunSuite with
ExpressionEvalHelper {
noPaths.checkInputDataTypes()
}
}
+
+ test("variant_insert") {
+ def checkInsert(input: String, path: String, value: Expression, expected:
String): Unit = {
+ val expr = VariantInsert(
+ Literal(parseJson(input)), Literal.create(path, StringType), value)
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(expr, StringType)),
+ expected)
+ }
+
+ // Object inserts.
+ checkInsert("""{"a": 1}""", "$.b", Literal(2), """{"a":1,"b":2}""")
+ checkInsert("{}", "$.a", Literal(1), """{"a":1}""")
+ checkInsert("""{"a": {"b": 1}}""", "$.a.c", Literal(3),
"""{"a":{"b":1,"c":3}}""")
+
+ // Missing intermediate keys are created; the next segment's kind picks
the container.
+ checkInsert("{}", "$.a.b.c", Literal(99), """{"a":{"b":{"c":99}}}""")
+ checkInsert("{}", "$.a[0]", Literal(1), """{"a":[1]}""")
+
+ // Array inserts shift existing elements right.
+ checkInsert("""["a","b","c"]""", "$[1]", Literal("z"),
"""["a","z","b","c"]""")
+ checkInsert("""["a","b","c"]""", "$[0]", Literal("z"),
"""["z","a","b","c"]""")
+ checkInsert("""["a","b","c"]""", "$[3]", Literal("z"),
"""["a","b","c","z"]""")
+ checkInsert("""["a","b"]""", "$[5]", Literal("z"),
"""["a","b",null,null,null,"z"]""")
+
+ // A missing intermediate created as an array is also padded with nulls up
to the index.
+ checkInsert("""{"a": []}""", "$.a[2].b", Literal(1),
"""{"a":[null,null,{"b":1}]}""")
+ checkInsert("""{"a": [1, 2]}""", "$.a[3].b", Literal(2),
"""{"a":[1,2,null,{"b":2}]}""")
+ checkInsert("{}", "$.a[2]", Literal(1), """{"a":[null,null,1]}""")
+ checkInsert("{}", "$.a[2].b", Literal(1), """{"a":[null,null,{"b":1}]}""")
+
+ // Descending through an existing container, then inserting into it.
+ checkInsert("""[{"a": 1}]""", "$[0].b", Literal(2), """[{"a":1,"b":2}]""")
+ checkInsert("""[{"x": 0}, {"a": 1}]""", "$[1].b", Literal(2),
"""[{"x":0},{"a":1,"b":2}]""")
+ checkInsert("""{"a": [1, 2, 3]}""", "$.a[1]", Literal(9),
"""{"a":[1,9,2,3]}""")
+
+ // Bracket and dot notations are interchangeable: `['k']`, `["k"]`, and
`.k` in one path.
+ checkInsert(
+ """{"a": {"b": {}}}""", """$['a'].b["c"]""", Literal(1),
"""{"a":{"b":{"c":1}}}""")
+
+ // Special / edge-case keys.
+ checkInsert("""{"a": 1}""", "$['']", Literal(2), """{"":2,"a":1}""")
+ checkInsert("""{"a": 1}""", "$['?']", Literal(2), """{"?":2,"a":1}""")
+ checkInsert(
+ """{"a": 1}""", "$['key with spaces']", Literal(2), """{"a":1,"key with
spaces":2}""")
+ checkInsert("""{"a": 1}""", "$.fb:testid", Literal(2),
"""{"a":1,"fb:testid":2}""")
+
+ // Insert a variant null (the JSON null literal) vs. a verbatim string vs.
structured JSON.
+ checkInsert("""{"a": 1}""", "$.b", Literal(parseJson("null")),
"""{"a":1,"b":null}""")
+ checkInsert("{}", "$.a", Literal("""{"x":1}"""), """{"a":"{\"x\":1}"}""")
+ checkInsert("{}", "$.a", Literal(parseJson("""{"x":1}""")),
"""{"a":{"x":1}}""")
+
+ // Values of various castable types exercise the corresponding
`castToVariant` branches.
+ checkInsert("{}", "$.a", Literal(true), """{"a":true}""")
+ checkInsert("{}", "$.a", Literal(7L), """{"a":7}""")
+ checkInsert("{}", "$.a", Literal(2.5), """{"a":2.5}""")
+ checkInsert(
+ "{}", "$.a", Literal.create(Array(1, 2, 3), ArrayType(IntegerType)),
"""{"a":[1,2,3]}""")
+
+ // NULL-intolerant: any NULL argument yields NULL.
+ checkEvaluation(
+ VariantInsert(Literal.create(null, VariantType), Literal("$.a"),
Literal(1)),
+ null)
+ checkInsert("""{"a": 1}""", null, Literal(1), null)
+ checkInsert("""{"a": 1}""", "$.b", Literal.create(null, VariantType), null)
+ checkInsert("""{"a": 1}""", "$.b", Literal.create(null, NullType), null)
+
+ val dynamic = VariantInsert(
+ Literal(parseJson("""{"a": 1}""")),
+ BoundReference(0, StringType, nullable = true),
+ Literal(2))
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(dynamic, StringType)),
+ """{"a":1,"b":2}""",
+ InternalRow(UTF8String.fromString("$.b")))
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("""{"a": 1}""")), Literal("$.a"),
Literal(2)),
+ "VARIANT_DUPLICATE_KEY",
+ Map("key" -> "a"))
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("""{"a": 1}""")), Literal("$.a.b"),
Literal(2)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$.a.b", "failedAt" -> "$.a", "functionName" ->
"`variant_insert`"))
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("""{"a": 1}""")), Literal("$[0]"),
Literal(2)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$[0]", "failedAt" -> "$", "functionName" ->
"`variant_insert`"))
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("""{"a": [1, 2]}""")), Literal("$.a.b"),
Literal(2)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$.a.b", "failedAt" -> "$.a", "functionName" ->
"`variant_insert`"))
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("""{"a": 5}""")), Literal("$.a[0]"),
Literal(2)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$.a[0]", "failedAt" -> "$.a", "functionName" ->
"`variant_insert`"))
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("5")), Literal("$.a"), Literal(2)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$.a", "failedAt" -> "$", "functionName" ->
"`variant_insert`"))
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("""{"a.b": 5}""")),
Literal("""$['a.b'].c"""), Literal(2)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$['a.b'].c", "failedAt" -> "$['a.b']", "functionName" ->
"`variant_insert`"))
+
+ // Structs and maps are rejected at analysis.
+ Seq(
+ Literal.create(null, MapType(StringType, IntegerType)),
+ Literal.create(null, StructType(Seq(StructField("x", IntegerType))))
+ ).foreach { v =>
+ assert(
+ VariantInsert(Literal(parseJson("{}")), Literal("$.a"),
v).checkInputDataTypes().isFailure)
+ }
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("{}")), Literal("$"), Literal(1)),
+ "INVALID_VARIANT_PATH",
+ Map("path" -> "$", "functionName" -> "`variant_insert`"))
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("{}")), Literal("abc"), Literal(1)),
+ "INVALID_VARIANT_PATH",
+ Map("path" -> "abc", "functionName" -> "`variant_insert`"))
+
+ val tooBig = "x".repeat(16 * 1024 * 1024)
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("{}")), Literal("$.a[2000000000]"),
Literal(1)),
+ "VARIANT_SIZE_LIMIT",
+ Map("sizeLimit" -> "16.0 MiB", "functionName" -> "`variant_insert`"))
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("{}")), Literal("$.a"), Literal(tooBig)),
+ "VARIANT_SIZE_LIMIT",
+ Map("sizeLimit" -> "16.0 MiB", "functionName" -> "`variant_insert`"))
+ }
}
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 36b2cab4d586..34faa4545b98 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
@@ -2763,6 +2763,10 @@ class PlanGenerationTestSuite extends ConnectFunSuite
with Logging {
fn.variant_delete(fn.parse_json(fn.col("g")), "$.a", "$.b")
}
+ functionTest("variant_insert") {
+ fn.variant_insert(fn.parse_json(fn.col("g")), "$.a", fn.lit(1))
+ }
+
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_insert.explain
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_insert.explain
new file mode 100644
index 000000000000..2b4cdcebf273
--- /dev/null
+++
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_insert.explain
@@ -0,0 +1,2 @@
+Project
[variant_insert(static_invoke(VariantExpressionEvalUtils.parseJson(g#0, false,
true, true)), $.a, 1) AS variant_insert(parse_json(g), $.a, 1)#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_insert.json
b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_insert.json
new file mode 100644
index 000000000000..5f1dad489e3d
--- /dev/null
+++
b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_insert.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_insert",
+ "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_insert",
+ "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_insert",
+ "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_insert.proto.bin
b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_insert.proto.bin
new file mode 100644
index 000000000000..d677b5fe6f4d
Binary files /dev/null and
b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_insert.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 5134975689e3..8e49cd70641e 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
@@ -557,6 +557,7 @@
|
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.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.VariantInsert |
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.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 787292f55b5a..0390507261f5 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
@@ -295,6 +295,95 @@ class VariantSuite extends SharedSparkSession with
ExpressionEvalHelper {
}
}
+ test("variant_insert with literal arguments") {
+ def rows(results: Any*): Seq[Row] = results.map(Row(_))
+
+ checkAnswer(
+ sql("SELECT to_json(variant_insert(parse_json('{\"a\": 1}'), '$.b',
2))"),
+ rows("""{"a":1,"b":2}"""))
+
+ checkAnswer(
+ sql("SELECT to_json(variant_insert(parse_json('[\"a\",\"b\",\"c\"]'),
'$[1]', 'z'))"),
+ rows("""["a","z","b","c"]"""))
+
+ checkAnswer(
+ sql("SELECT to_json(variant_insert(parse_json('{}'), '$.a',
parse_json('{\"x\":1}')))"),
+ rows("""{"a":{"x":1}}"""))
+
+ checkAnswer(
+ sql("SELECT to_json(variant_insert(parse_json('{}'), '$.a',
parse_json('null')))"),
+ rows("""{"a":null}"""))
+
+ checkAnswer(
+ sql("SELECT to_json(variant_insert(CAST(NULL AS VARIANT), '$.a', 1))"),
+ rows(null))
+
+ checkAnswer(
+ sql("SELECT to_json(variant_insert(parse_json('{\"a\": 1}'), '$.b',
NULL))"),
+ rows(null))
+
+ checkAnswer(
+ sql("SELECT to_json(variant_insert(parse_json('{\"a\": 1}'), NULL, 1))"),
+ rows(null))
+
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT variant_insert(parse_json('{\"a\": 1}'), '$.a',
2)").collect()
+ },
+ condition = "VARIANT_DUPLICATE_KEY",
+ parameters = Map("key" -> "a"))
+
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT variant_insert(parse_json('{\"a\": 1}'), '$.a.b',
2)").collect()
+ },
+ condition = "VARIANT_PATH_TYPE_MISMATCH",
+ parameters = Map(
+ "path" -> "$.a.b", "failedAt" -> "$.a", "functionName" ->
toSQLId("variant_insert")))
+
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT variant_insert(parse_json('{}'), '$', 1)").collect()
+ },
+ condition = "INVALID_VARIANT_PATH",
+ parameters = Map("path" -> "$", "functionName" ->
toSQLId("variant_insert")))
+
+ // A raw struct/map value is not castable to variant and is rejected at
analysis.
+ assert(intercept[AnalysisException] {
+ sql("SELECT variant_insert(parse_json('{}'), '$.a', named_struct('x',
1))")
+ }.getCondition == "DATATYPE_MISMATCH.CAST_WITHOUT_SUGGESTION")
+ }
+
+ test("variant_insert 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}""", "$.b", 2),
+ ("""["a","b"]""", "$[0]", 9),
+ (null, "$.b", 2)
+ ).toDF("json", "path", "val")
+ val v = parse_json(col("json"))
+ val out = df.select(to_json(variant_insert(v, col("path"),
col("val"))).alias("r"))
+ checkAnswer(out, rows("""{"a":1,"b":2}""", """[9,"a","b"]""", null))
+
+ // A foldable literal path combined with a dynamic value.
+ val objDf = Seq(
+ ("""{"a": 1}""", 2),
+ ("""{"x": 5}""", 9),
+ (null, 2)
+ ).toDF("json", "val")
+ val objV = parse_json(objDf("json"))
+ val out2 = objDf.select(to_json(variant_insert(objV, lit("$.z"),
col("val"))).alias("r"))
+ checkAnswer(out2, rows("""{"a":1,"z":2}""", """{"x":5,"z":9}""", null))
+
+ // String-path overload of the DataFrame API.
+ val out3 = objDf.select(to_json(variant_insert(objV, "$.z",
col("val"))).alias("r"))
+ checkAnswer(out3, rows("""{"a":1,"z":2}""", """{"x":5,"z":9}""", 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]
