[SPARK-58035] Add try_variant_insert expression

python/docs/source/reference/pyspark.sql/functions.rst

@@ -605,6 +605,7 @@ VARIANT Functions
     variant_delete
     variant_get
     variant_insert
+    try_variant_insert
     try_parse_json
     to_variant_object
 

python/pyspark/sql/connect/functions/builtin.py

@@ -2227,6 +2227,16 @@ def variant_insert(v: "ColumnOrName", path: Union[Column, str], value: "ColumnOr
 variant_insert.__doc__ = pysparkfuncs.variant_insert.__doc__
 
 
+def try_variant_insert(
+    v: "ColumnOrName", path: Union[Column, str], value: "ColumnOrName"
+) -> Column:
+    path_col = path if isinstance(path, Column) else lit(path)
+    return _invoke_function("try_variant_insert", _to_col(v), path_col, _to_col(value))
+
+
+try_variant_insert.__doc__ = pysparkfuncs.try_variant_insert.__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",
+    "try_variant_insert",
     "try_parse_json",
     "to_variant_object",
     # XML Functions

python/pyspark/sql/functions/builtin.py

@@ -21731,7 +21731,8 @@ def variant_insert(v: "ColumnOrName", path: Union[Column, str], value: "ColumnOr
     """
     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.
+    right. Missing intermediate keys are created. 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
 
@@ -21786,6 +21787,66 @@ def variant_insert(v: "ColumnOrName", path: Union[Column, str], value: "ColumnOr
     )
 
 
+@_try_remote_functions
+def try_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;
+    an array path inserts at the index, shifting later elements right. Missing intermediate keys
+    are created. Returns NULL if the field already exists or a path segment hits a value of an
+    incompatible type, or if any argument is NULL.
+
+    .. versionadded:: 4.3.0
+
+    Parameters
+    ----------
+    v : :class:`~pyspark.sql.Column` or str
+        a variant column or column name
+    path : :class:`~pyspark.sql.Column` or str
+        the JSONPath 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`, or NULL if the insertion fails
+
+    Examples
+    --------
+    >>> from pyspark.sql.functions import lit, parse_json, to_json, try_variant_insert
+    >>> df = spark.createDataFrame([{'json': '''{ "a": 1, "arr": ["x", "y"] }'''}])
+    >>> v = parse_json(df.json)
+    >>> df.select(to_json(try_variant_insert(v, "$.b", lit(2))).alias("r")).collect()
+    [Row(r='{"a":1,"arr":["x","y"],"b":2}')]
+    >>> df.select(to_json(try_variant_insert(v, "$.c.d", lit(3))).alias("r")).collect()
+    [Row(r='{"a":1,"arr":["x","y"],"c":{"d":3}}')]
+    >>> df.select(to_json(try_variant_insert(v, "$.arr[1]", lit("z"))).alias("r")).collect()
+    [Row(r='{"a":1,"arr":["x","z","y"]}')]
+    >>> df.select(to_json(try_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(try_variant_insert(v, "$.a", lit(2))).alias("r")).collect()
+    [Row(r=None)]
+    >>> df.select(to_json(try_variant_insert(v, "$.a.b", lit(2))).alias("r")).collect()
+    [Row(r=None)]
+    >>> df.select(to_json(try_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(
+        "try_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:
     """

python/pyspark/sql/tests/test_functions.py

@@ -3534,6 +3534,12 @@ 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.try_variant_insert(v, "$.z", F.lit(9)))),
+            ['{"a":1,"z":9}', '{"b":2,"z":9}'],
+        )
+        check(df.select(F.to_json(F.try_variant_insert(v, df.path, F.lit(9)))), [None, None])
+        check(df.select(F.to_json(F.try_variant_insert(v, "$.z", F.lit(None)))), [None, None])
         check(df.select(F.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

@@ -9824,7 +9824,8 @@ object functions {
   /**
    * 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.
+   * elements right. Missing intermediate keys are 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.
@@ -9843,7 +9844,8 @@ object functions {
   /**
    * 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.
+   * elements right. Missing intermediate keys are 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.
@@ -9859,6 +9861,46 @@ object functions {
   def variant_insert(v: Column, path: String, value: Column): Column =
     Column.fn("variant_insert", v, lit(path), value)
 
+  /**
+   * Inserts a value into a variant at the given JSONPath location. An object path adds a new
+   * field; an array path inserts at the index, shifting later elements right. Missing
+   * intermediate keys are created. Returns NULL if the field already exists or a path segment
+   * hits a value of an incompatible type, or if any argument is NULL.
+   *
+   * @param v
+   *   a variant column.
+   * @param path
+   *   the column containing the JSONPath string identifying the 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 try_variant_insert(v: Column, path: Column, value: Column): Column =
+    Column.fn("try_variant_insert", v, path, value)
+
+  /**
+   * Inserts a value into a variant at the given JSONPath location. An object path adds a new
+   * field; an array path inserts at the index, shifting later elements right. Missing
+   * intermediate keys are created. Returns NULL if the field already exists or a path segment
+   * hits a value of an incompatible type, or if any argument is NULL.
+   *
+   * @param v
+   *   a variant column.
+   * @param path
+   *   the JSONPath identifying the 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 try_variant_insert(v: Column, path: String, value: Column): Column =
+    Column.fn("try_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.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -986,7 +986,8 @@ object FunctionRegistry {
     expression[ToVariantObject]("to_variant_object"),
     expression[IsValidVariant]("is_valid_variant"),
     expression[VariantDelete]("variant_delete"),
-    expression[VariantInsert]("variant_insert"),
+    expressionBuilder("variant_insert", VariantInsertExpressionBuilder),
+    expressionBuilder("try_variant_insert", TryVariantInsertExpressionBuilder),
 
     // Spatial
     expression[ST_AsBinary]("st_asbinary"),

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala

@@ -19,6 +19,7 @@ package org.apache.spark.sql.catalyst.expressions.variant
 
 import scala.util.control.NonFatal
 
+import org.apache.spark.SparkRuntimeException
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.util.{ArrayData, MapData}
 import org.apache.spark.sql.errors.QueryExecutionErrors
@@ -131,22 +132,29 @@ object VariantExpressionEvalUtils {
   /**
    * 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`.
+   * and a type mismatch maps to `VARIANT_PATH_TYPE_MISMATCH`. When `failOnError` is false (the
+   * `try_variant_insert` mode), a duplicate key or path type mismatch returns null instead of
+   * throwing; a size overflow (and a malformed path, rejected earlier during parsing) is still
+   * raised.
    */
   def insertAtPath(
       input: VariantVal,
       javaSegments: Array[VariantBuilder.PathSegment],
       path: String,
       value: Any,
       valueDataType: DataType,
-      functionName: String): VariantVal = {
+      functionName: String,
+      failOnError: Boolean): 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 _: VariantPathTypeMismatchException if !failOnError => null
+      case e: SparkRuntimeException if !failOnError && e.getCondition == "VARIANT_DUPLICATE_KEY" =>
+        null
       case e: VariantPathTypeMismatchException =>
         throw QueryExecutionErrors.variantPathTypeMismatch(
           path, renderVariantPath(javaSegments.take(e.depth)), functionName)
@@ -160,10 +168,11 @@ object VariantExpressionEvalUtils {
       path: UTF8String,
       value: Any,
       valueDataType: DataType,
-      functionName: String): VariantVal = {
+      functionName: String,
+      failOnError: Boolean): VariantVal = {
     val pathStr = path.toString
     val javaSegments = toJavaSegments(parseVariantPath(pathStr, functionName))
-    insertAtPath(input, javaSegments, pathStr, value, valueDataType, functionName)
+    insertAtPath(input, javaSegments, pathStr, value, valueDataType, functionName, failOnError)
   }
 
   /** Cast a Spark value from `dataType` into the variant type. */