[SPARK-57804][SQL] Add variant_set expression

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<FieldEntry> fields = new ArrayList<>(size + 1);
+        int start = writePos;
+        boolean found = false;
+        for (int i = 0; i < size; ++i) {
+          int id = readUnsigned(value, idStart + idSize * i, idSize);
+          int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize);
+          int elementPos = dataStart + offset;
+          String fieldKey = getMetadataKey(metadata, id);
+          boolean isTarget = fieldKey.equals(key);
+          found |= isTarget;
+          int newId = addKey(fieldKey);
+          fields.add(new FieldEntry(fieldKey, newId, writePos - start));
+          if (isTarget && isLast) {
+            // Replace the existing field's value in place.
+            appendVariant(val);
+          } else if (isTarget) {
+            appendWithSetImpl(
+                value, metadata, elementPos, segments, depth + 1, val, createIfMissing);
+          } else {
+            appendVariantImpl(value, metadata, elementPos);
+          }
+        }
+        if (!found && createIfMissing) {
+          // Target key is missing; create it (and any remaining path). When `createIfMissing` is
+          // false this is a no-op: the fields copied above already reproduce the input object.
+          int newId = addKey(key);
+          fields.add(new FieldEntry(key, newId, writePos - start));
+          if (isLast) {
+            appendVariant(val);
+          } else {
+            appendNewPath(segments, depth + 1, val);
+          }
+        }
+        finishWritingObject(start, fields);
+        return null;
+      });
+    } else if (seg instanceof ArrayIndexSegment && basicType == ARRAY) {
+      int index = ((ArrayIndexSegment) seg).index;
+      handleArray(value, pos, (size, offsetSize, offsetStart, dataStart) -> {
+        ArrayList<Integer> offsets = new ArrayList<>(size + 1);
+        int start = writePos;
+        if (index < size) {
+          // Replace the element at `index`, or descend into it for an intermediate segment.
+          for (int i = 0; i < size; ++i) {
+            int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize);
+            int elementPos = dataStart + offset;
+            offsets.add(writePos - start);
+            if (i != index) {
+              appendVariantImpl(value, metadata, elementPos);
+            } else if (isLast) {
+              appendVariant(val);
+            } else {
+              appendWithSetImpl(
+                  value, metadata, elementPos, segments, depth + 1, val, createIfMissing);
+            }
+          }
+        } else {
+          // Index is past the end. Copy existing elements; when `createIfMissing` is true, pad with
+          // variant nulls up to `index` and create the leaf value or the rest of the path. When
+          // false this is a no-op: the copied elements already reproduce the input array.
+          for (int i = 0; i < size; ++i) {
+            int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize);
+            offsets.add(writePos - start);
+            appendVariantImpl(value, metadata, dataStart + offset);
+          }
+          if (createIfMissing) {
+            for (int i = size; i < index; ++i) {
+              offsets.add(writePos - start);
+              appendNull();
+            }
+            offsets.add(writePos - start);
+            if (isLast) {
+              appendVariant(val);
+            } else {
+              appendNewPath(segments, depth + 1, val);
+            }
+          }
+        }
+        finishWritingArray(start, offsets);
+        return null;
+      });
+    } else {
+      // The segment kind does not match the container at this path prefix.
+      throw new VariantPathTypeMismatchException(depth);
+    }
+  }
+
   // Build a fresh chain of containers for `segments[depth..]`, terminating in `val`. Used to
   // materialize missing intermediate path segments during insertion. The kind of each segment
   // decides the container created: an object-key segment creates a single-field object, while an

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
 

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):

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

python/pyspark/sql/functions/builtin.py

@@ -21786,6 +21786,70 @@ 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.
+    Throws an error if a path segment hits a value of an incompatible type. Returns NULL if any
+    argument is NULL.
+
+    .. versionadded:: 4.3.0
+
+    Parameters
+    ----------
+    v : :class:`~pyspark.sql.Column` or str
+        a variant column or column name
+    path : :class:`~pyspark.sql.Column` or str
+        the JSONPath set target. A `str` is a literal path; a :class:`~pyspark.sql.Column` supplies
+        the path at runtime. A valid path should start with `$` and is followed by one or more
+        segments like `[123]`, `.name`, `['name']`, or `["name"]`. The root path `$` is not allowed.
+    value : :class:`~pyspark.sql.Column` or str
+        the value to set. Any expression castable to variant.
+    create_if_missing : bool, optional
+        whether to create missing keys or out-of-range array indices (default True).
+
+    Returns
+    -------
+    :class:`~pyspark.sql.Column`
+        a variant column with `value` set at `path`
+
+    Examples
+    --------
+    >>> from pyspark.sql.functions import lit, parse_json, to_json, variant_set
+    >>> df = spark.createDataFrame([{'json': '''{"a": 1, "arr": [1, 2, 3]}'''}])
+    >>> v = parse_json(df.json)
+    >>> df.select(to_json(variant_set(v, "$.a", lit(9))).alias("r")).collect()
+    [Row(r='{"a":9,"arr":[1,2,3]}')]
+    >>> df.select(to_json(variant_set(v, "$.b", lit(2))).alias("r")).collect()
+    [Row(r='{"a":1,"arr":[1,2,3],"b":2}')]
+    >>> df.select(to_json(variant_set(v, "$.arr[1]", lit(9))).alias("r")).collect()
+    [Row(r='{"a":1,"arr":[1,9,3]}')]
+    >>> df.select(to_json(variant_set(v, "$.b", lit(2), False)).alias("r")).collect()
+    [Row(r='{"a":1,"arr":[1,2,3]}')]
+    >>> df.select(to_json(variant_set(v, "$.a", parse_json(lit("null")))).alias("r")).collect()
+    [Row(r='{"a":null,"arr":[1,2,3]}')]
+    >>> df.select(to_json(variant_set(v, "$.a", lit(None))).alias("r")).collect()
+    [Row(r=None)]
+    """
+    from pyspark.sql.classic.column import _to_java_column
+
+    path_col = path if isinstance(path, Column) else lit(path)
+    return _invoke_function(
+        "variant_set",
+        _to_java_column(v),
+        _to_java_column(path_col),
+        _to_java_column(value),
+        _to_java_column(lit(create_if_missing)),
+    )
+
+
 @_try_remote_functions
 def variant_get(v: "ColumnOrName", path: Union[Column, str], targetType: str) -> Column:
     """

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<a: BIGINT>", "OBJECT<b: BIGINT>"])
         check(df.select(F.schema_of_variant_agg(v)), ["OBJECT<a: BIGINT, b: BIGINT>"])
 

sql/api/src/main/scala/org/apache/spark/sql/functions.scala

@@ -9859,6 +9859,92 @@ 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. Throws an error if a path segment hits a value of an incompatible type.
+   * Returns NULL if any argument is NULL.
+   *
+   * @param v
+   *   a variant column.
+   * @param path
+   *   the column containing the JSONPath string identifying the set target. A valid path should
+   *   start with `$` and is followed by one or more segments like `[123]`, `.name`, `['name']`,
+   *   or `["name"]`. The root path `$` is not allowed.
+   * @param value
+   *   the value to set. Any expression castable to variant.
+   * @group variant_funcs
+   * @since 4.3.0
+   */
+  def variant_set(v: Column, path: Column, value: Column): Column =
+    Column.fn("variant_set", v, path, value)
+
+  /**
+   * Sets or upserts a value in a variant at the given JSONPath location. An existing object field
+   * or array element at the target is replaced. A missing field, array index, or intermediate
+   * path is created. Throws an error if a path segment hits a value of an incompatible type.
+   * Returns NULL if any argument is NULL.
+   *
+   * @param v
+   *   a variant column.
+   * @param path
+   *   the JSONPath identifying the set target. A valid path should start with `$` and is followed
+   *   by one or more segments like `[123]`, `.name`, `['name']`, or `["name"]`. The root path `$`
+   *   is not allowed.
+   * @param value
+   *   the value to set. Any expression castable to variant.
+   * @group variant_funcs
+   * @since 4.3.0
+   */
+  def variant_set(v: Column, path: String, value: Column): Column =
+    Column.fn("variant_set", v, lit(path), value)
+
+  /**
+   * Sets or upserts a value in a variant at the given JSONPath location. An existing object field
+   * or array element at the target is replaced. A missing field, array index, or intermediate
+   * path is created, unless `createIfMissing` is false, in which case the variant is left
+   * unchanged. Throws an error if a path segment hits a value of an incompatible type. Returns
+   * NULL if any argument is NULL.
+   *
+   * @param v
+   *   a variant column.
+   * @param path
+   *   the column containing the JSONPath string identifying the set target. A valid path should
+   *   start with `$` and is followed by one or more segments like `[123]`, `.name`, `['name']`,
+   *   or `["name"]`. The root path `$` is not allowed.
+   * @param value
+   *   the value to set. Any expression castable to variant.
+   * @param createIfMissing
+   *   whether to create missing keys or out-of-range array indices.
+   * @group variant_funcs
+   * @since 4.3.0
+   */
+  def variant_set(v: Column, path: Column, value: Column, createIfMissing: Boolean): Column =
+    Column.fn("variant_set", v, path, value, lit(createIfMissing))
+
+  /**
+   * Sets or upserts a value in a variant at the given JSONPath location. An existing object field
+   * or array element at the target is replaced. A missing field, array index, or intermediate
+   * path is created, unless `createIfMissing` is false, in which case the variant is left
+   * unchanged. Throws an error if a path segment hits a value of an incompatible type. Returns
+   * NULL if any argument is NULL.
+   *
+   * @param v
+   *   a variant column.
+   * @param path
+   *   the JSONPath identifying the set target. A valid path should start with `$` and is followed
+   *   by one or more segments like `[123]`, `.name`, `['name']`, or `["name"]`. The root path `$`
+   *   is not allowed.
+   * @param value
+   *   the value to set. Any expression castable to variant.
+   * @param createIfMissing
+   *   whether to create missing keys or out-of-range array indices.
+   * @group variant_funcs
+   * @since 4.3.0
+   */
+  def variant_set(v: Column, path: String, value: Column, createIfMissing: Boolean): Column =
+    Column.fn("variant_set", v, lit(path), value, lit(createIfMissing))
+
   /**
    * Extracts a sub-variant from `v` according to `path` string, and then cast the sub-variant to
    * `targetType`. Returns null if the path does not exist. Throws an exception if the cast fails.

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"),