From b5fa32344b334cb06d98ca4272d692cfaec95884 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bojana=20Ze=C4=8Devi=C4=87?= Date: Wed, 1 Jul 2026 13:24:46 +0000 Subject: [PATCH 1/4] init --- .../spark/types/variant/VariantBuilder.java | 113 ++++++++++++ .../reference/pyspark.sql/functions.rst | 1 + .../pyspark/sql/connect/functions/builtin.py | 15 ++ python/pyspark/sql/functions/__init__.py | 1 + python/pyspark/sql/functions/builtin.py | 63 +++++++ python/pyspark/sql/tests/test_functions.py | 13 ++ .../org/apache/spark/sql/functions.scala | 82 +++++++++ .../catalyst/analysis/FunctionRegistry.scala | 1 + .../variant/VariantExpressionEvalUtils.scala | 41 +++++ .../variant/variantExpressions.scala | 170 ++++++++++++++++++ .../variant/VariantExpressionSuite.scala | 132 ++++++++++++++ .../spark/sql/PlanGenerationTestSuite.scala | 8 + .../function_variant_set.explain | 2 + ...variant_set_with_create_if_missing.explain | 2 + .../queries/function_variant_set.json | 125 +++++++++++++ .../queries/function_variant_set.proto.bin | Bin 0 -> 1130 bytes ...on_variant_set_with_create_if_missing.json | 146 +++++++++++++++ ...riant_set_with_create_if_missing.proto.bin | Bin 0 -> 1322 bytes .../sql-functions/sql-expression-schema.md | 1 + .../org/apache/spark/sql/VariantSuite.scala | 82 +++++++++ 20 files changed, 998 insertions(+) create mode 100644 sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_set.explain create mode 100644 sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_set_with_create_if_missing.explain create mode 100644 sql/connect/common/src/test/resources/query-tests/queries/function_variant_set.json create mode 100644 sql/connect/common/src/test/resources/query-tests/queries/function_variant_set.proto.bin create mode 100644 sql/connect/common/src/test/resources/query-tests/queries/function_variant_set_with_create_if_missing.json create mode 100644 sql/connect/common/src/test/resources/query-tests/queries/function_variant_set_with_create_if_missing.proto.bin 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 9685f0a35f85a..07a527e87a6c3 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 static Variant insertAtPath(Variant v, PathSegment[] segments, Variant va 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 @@ private void appendWithInsertionImpl( } } + 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 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 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 70e5811180d0c..7ab82ced8c321 100644 --- a/python/docs/source/reference/pyspark.sql/functions.rst +++ b/python/docs/source/reference/pyspark.sql/functions.rst @@ -605,6 +605,7 @@ VARIANT Functions variant_delete variant_get 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 8e74aeef2828f..04a5318f6eb0c 100644 --- a/python/pyspark/sql/connect/functions/builtin.py +++ b/python/pyspark/sql/connect/functions/builtin.py @@ -2227,6 +2227,21 @@ def variant_insert(v: "ColumnOrName", path: Union[Column, str], value: "ColumnOr variant_insert.__doc__ = pysparkfuncs.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 ee9d06fcaa2d8..94e677f80b208 100644 --- a/python/pyspark/sql/functions/__init__.py +++ b/python/pyspark/sql/functions/__init__.py @@ -483,6 +483,7 @@ "variant_delete", "variant_get", "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 1e348cfd42874..5128c23d5f243 100644 --- a/python/pyspark/sql/functions/builtin.py +++ b/python/pyspark/sql/functions/builtin.py @@ -21786,6 +21786,69 @@ def variant_insert(v: "ColumnOrName", path: Union[Column, str], value: "ColumnOr ) +@_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. + 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), + _to_java_column(lit(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 2ade286763dea..7dbc3cf49f7a9 100644 --- a/python/pyspark/sql/tests/test_functions.py +++ b/python/pyspark/sql/tests/test_functions.py @@ -3534,6 +3534,19 @@ def check(resultDf, expected): 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.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", "OBJECT"]) check(df.select(F.schema_of_variant_agg(v)), ["OBJECT"]) 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 c8b269a63de5e..5fb15afb52e66 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 @@ -9859,6 +9859,88 @@ object functions { def variant_insert(v: Column, path: String, value: Column): Column = Column.fn("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. 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. 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. 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. 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 193e939889dbf..e6d0252f9a1e7 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 @@ -987,6 +987,7 @@ object FunctionRegistry { expression[IsValidVariant]("is_valid_variant"), expression[VariantDelete]("variant_delete"), expression[VariantInsert]("variant_insert"), + 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 20a55bccc7791..b259270f73099 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 @@ -166,6 +166,47 @@ object VariantExpressionEvalUtils { insertAtPath(input, javaSegments, pathStr, value, valueDataType, functionName) } + /** + * Set `value` to `input` at `javaSegments`. `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 0f8995ce829fe..6ac14b234c52c 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 @@ -942,6 +942,176 @@ object VariantInsert { } } +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. 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 6f1a410c5bb91..29ae17db12801 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 @@ -1465,4 +1465,136 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper { "VARIANT_SIZE_LIMIT", Map("sizeLimit" -> "16.0 MiB", "functionName" -> "`variant_insert`")) } + + 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]}""") + + // create_if_missing = false. + checkSet("""{"a": 1, "b": 2}""", "$.a", Literal(99), """{"a":99,"b":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","b"]""", "$[1]", Literal("z"), """["a","z"]""", 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) + } + + checkErrorInExpression[SparkRuntimeException]( + VariantSet(Literal(parseJson("{}")), Literal("$"), Literal(1), Literal(true)), + "INVALID_VARIANT_PATH", + Map("path" -> "$", "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 34faa4545b98f..04299810ebb7a 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 @@ -2767,6 +2767,14 @@ class PlanGenerationTestSuite extends ConnectFunSuite with Logging { fn.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 0000000000000..3f695273d79b5 --- /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 , [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 0000000000000..2ad5e4938e92c --- /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 , [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 0000000000000..ef8e00fe0acb6 --- /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 0000000000000000000000000000000000000000..899cc1c6b0b5423727f98908c440696cde32308c GIT binary patch literal 1130 zcmdUtKTE?v7{>cH_=j__oLfKy=^!Np=@>$WD7Xqbxs zP*K#`&B2f17x6+tIw&nJ#pQK z?kT3}g=N13AvsFLdP9{7`U=z{kW_F71prtY8+ zk>;6D?T!^?X?OMe!c`Hb7$#%d`H|jFQS33nB-S2>OzDGyYt~sSOOtstVsVKg6UMVF zyX(KG#}92aDtHF+2=pGpof#uf6rW#m`UYfSXtiCe;04GN&<6;2XKXO}n{%vSayPf( RFH4hhJ-^)b()`nre*rxZkVOCh literal 0 HcmV?d00001 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 0000000000000..cb795edc537b4 --- /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 0000000000000000000000000000000000000000..02d73bc213938c0d8edddc179c206aaeb4b09977 GIT binary patch literal 1322 zcmdUtu}i~16vq1+L_7z}xkx|+=^!Np=@>$WD7Xqbxs;dW8n338n7arr8JvV}`4=1n zK}AtF!OhWs$qNPPptQJ@E;k(C_ucP%_eiq^<_#egkY=lKG6H%HQsqUcTpl?AkGW75 zI`tViQCb8Evm$4T!n2rzd5)K^_IXy{8Di#Y4;QgD&FWck9j~RHIa)ElmdWY`<{X9U zNv7zPVSWN38RnAp1~L`&4X8mNuHYWX9nf3wU1DICZEvM{Y-5JuDYKrKOVcSU)_fZiUTHCj&;T{lj^A8ignh@QZo;SEGki?LT8p` zck>VR_@SLf1usCJfIfh~HzV>y@%fR{4ufdBvi literal 0 HcmV?d00001 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 8e49cd70641e9..dd6a2fba03f37 100644 --- a/sql/core/src/test/resources/sql-functions/sql-expression-schema.md +++ b/sql/core/src/test/resources/sql-functions/sql-expression-schema.md @@ -558,6 +558,7 @@ | org.apache.spark.sql.catalyst.expressions.variant.VariantDelete | variant_delete | SELECT variant_delete(parse_json('{"a": 1, "b": 2, "c": 3, "items": [1, 2, 3]}'), NULL, '$.a', '$.c') | struct | | org.apache.spark.sql.catalyst.expressions.variant.VariantGetExpressionBuilder | variant_get | SELECT variant_get(parse_json('{"a": 1}'), '$.a', 'int') | struct | | org.apache.spark.sql.catalyst.expressions.variant.VariantInsert | variant_insert | SELECT variant_insert(parse_json('{"a": 1}'), '$.b', 2) | struct | +| org.apache.spark.sql.catalyst.expressions.variant.VariantSetExpressionBuilder | variant_set | SELECT variant_set(parse_json('{"a": 1}'), '$.a', 2) | struct | | org.apache.spark.sql.catalyst.expressions.xml.XPathBoolean | xpath_boolean | SELECT xpath_boolean('1','a/b') | struct1, a/b):boolean> | | org.apache.spark.sql.catalyst.expressions.xml.XPathDouble | xpath_double | SELECT xpath_double('12', 'sum(a/b)') | struct12, sum(a/b)):double> | | org.apache.spark.sql.catalyst.expressions.xml.XPathDouble | xpath_number | SELECT xpath_number('12', 'sum(a/b)') | struct12, 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 0390507261f55..fb0907da41919 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 @@ -384,6 +384,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) + ).toDF("json", "path", "val") + val v = parse_json(col("json")) + val out = df.select(to_json(variant_set(v, col("path"), col("val"))).alias("r")) + checkAnswer(out, rows("""{"a":2}""", """[9,"b"]""", null)) + + // create_if_missing = false leaves a missing target unchanged (Column-path overload). + val flagDf = Seq("""{"a": 1}""").toDF("json") + val flagV = parse_json(flagDf("json")) + checkAnswer( + flagDf.select(to_json(variant_set(flagV, lit("$.b"), lit(2), false)).alias("r")), + rows("""{"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) From 9e736c5b8618590e7d75863609dbb851fc0b07a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bojana=20Ze=C4=8Devi=C4=87?= Date: Thu, 9 Jul 2026 12:22:08 +0000 Subject: [PATCH 2/4] improve --- .../variant/VariantExpressionEvalUtils.scala | 2 +- .../variant/VariantExpressionSuite.scala | 11 +++++++++++ .../org/apache/spark/sql/VariantSuite.scala | 18 +++++++++--------- 3 files changed, 21 insertions(+), 10 deletions(-) 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 b259270f73099..649cf5130fac2 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 @@ -167,7 +167,7 @@ object VariantExpressionEvalUtils { } /** - * Set `value` to `input` at `javaSegments`. `path` is the source string used in error messages. + * 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. 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 29ae17db12801..489d66814c0ae 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 @@ -1506,6 +1506,13 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper { // 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": 1}""", "$.b", Literal(2), """{"a":1}""", false) @@ -1591,6 +1598,10 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper { "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)), 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 fb0907da41919..f8d8a93efb165 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 @@ -440,18 +440,18 @@ class VariantSuite extends SharedSparkSession with ExpressionEvalHelper { val df = Seq( ("""{"a": 1}""", "$.a", 2), ("""["a","b"]""", "$[0]", 9), - (null, "$.a", 2) + (null, "$.a", 2), + ("""{"a": 1}""", "$.b", 2) ).toDF("json", "path", "val") val v = parse_json(col("json")) - val out = df.select(to_json(variant_set(v, col("path"), col("val"))).alias("r")) - checkAnswer(out, rows("""{"a":2}""", """[9,"b"]""", null)) - - // create_if_missing = false leaves a missing target unchanged (Column-path overload). - val flagDf = Seq("""{"a": 1}""").toDF("json") - val flagV = parse_json(flagDf("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( - flagDf.select(to_json(variant_set(flagV, lit("$.b"), lit(2), false)).alias("r")), - rows("""{"a":1}""")) + 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") From f3bcfb402e6d534e5fdc85f3b1c312c33ccc7197 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bojana=20Ze=C4=8Devi=C4=87?= Date: Thu, 9 Jul 2026 16:43:06 +0000 Subject: [PATCH 3/4] lint --- .../org/apache/spark/types/variant/VariantBuilder.java | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) 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 07a527e87a6c3..fa23d610f2395 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 @@ -145,11 +145,11 @@ public static Variant insertAtPath(Variant v, PathSegment[] segments, Variant va } // 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. + // (`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) { From 92adccec827926a55eba1e2f8c75cbf02068a5a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bojana=20Ze=C4=8Devi=C4=87?= Date: Fri, 10 Jul 2026 15:08:11 +0000 Subject: [PATCH 4/4] edit docs --- python/pyspark/sql/functions/builtin.py | 3 ++- .../main/scala/org/apache/spark/sql/functions.scala | 12 ++++++++---- .../expressions/variant/variantExpressions.scala | 5 +++-- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/python/pyspark/sql/functions/builtin.py b/python/pyspark/sql/functions/builtin.py index 5128c23d5f243..372c4cf14cf98 100644 --- a/python/pyspark/sql/functions/builtin.py +++ b/python/pyspark/sql/functions/builtin.py @@ -21797,7 +21797,8 @@ def variant_set( Sets or upserts a value in a variant at the given JSONPath location. An existing object field or array element at the target is replaced. A missing field, array index, or intermediate path is created, unless `create_if_missing` is false, in which case the variant is left unchanged. - Returns NULL if any argument is NULL. + 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 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 5fb15afb52e66..94eb97d77385d 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 @@ -9862,7 +9862,8 @@ object functions { /** * Sets or upserts a value in a variant at the given JSONPath location. An existing object field * or array element at the target is replaced. A missing field, array index, or intermediate - * path is created. Returns NULL if any argument is NULL. + * 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. @@ -9881,7 +9882,8 @@ object functions { /** * Sets or upserts a value in a variant at the given JSONPath location. An existing object field * or array element at the target is replaced. A missing field, array index, or intermediate - * path is created. Returns NULL if any argument is NULL. + * 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. @@ -9901,7 +9903,8 @@ object functions { * Sets or upserts a value in a variant at the given JSONPath location. An existing object field * or array element at the target is replaced. A missing field, array index, or intermediate * path is created, unless `createIfMissing` is false, in which case the variant is left - * unchanged. Returns NULL if any argument is NULL. + * 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. @@ -9923,7 +9926,8 @@ object functions { * Sets or upserts a value in a variant at the given JSONPath location. An existing object field * or array element at the target is replaced. A missing field, array index, or intermediate * path is created, unless `createIfMissing` is false, in which case the variant is left - * unchanged. Returns NULL if any argument is NULL. + * 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. 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 6ac14b234c52c..1797eab2e0088 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 @@ -1062,8 +1062,9 @@ object VariantSet { usage = "_FUNC_(v, path, val[, create_if_missing]) - Sets or upserts a value in a variant at " + "the given JSONPath location. An existing object field or array element at the target is " + "replaced. A missing field, array index, or intermediate path is created, unless " + - "`create_if_missing` is false, in which case the variant is left unchanged. Returns NULL if " + - "any argument is NULL.", + "`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.