[SPARK-58061][SQL] Add variant_array_append expression

common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java

@@ -144,6 +144,17 @@ public static Variant insertAtPath(Variant v, PathSegment[] segments, Variant va
     return builder.result();
   }
 
+  // Return a new variant with `val` appended to the array at `segments`, where an empty `segments`
+  // refers to the whole variant. If the path does not match (missing key or out-of-range index), a
+  // semantically equivalent variant is returned. A target that is not an array, or a segment
+  // applied to an incompatible container type, throws VariantPathTypeMismatchException, which the
+  // caller maps to VARIANT_PATH_TYPE_MISMATCH.
+  public static Variant arrayAppendAtPath(Variant v, PathSegment[] segments, Variant val) {
+    VariantBuilder builder = new VariantBuilder(false);
+    builder.appendWithArrayAppendImpl(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();
@@ -666,6 +677,78 @@ private void appendWithInsertionImpl(
     }
   }
 
+  private void appendWithArrayAppendImpl(
+      byte[] value, byte[] metadata, int pos, PathSegment[] segments, int depth, Variant val) {
+    checkIndex(pos, value.length);
+    int basicType = value[pos] & BASIC_TYPE_MASK;
+    if (depth == segments.length) {
+      // Reached the target: it must be an array, to which `val` is appended.
+      if (basicType != ARRAY) {
+        throw new VariantPathTypeMismatchException(depth);
+      }
+      handleArray(value, pos, (size, offsetSize, offsetStart, dataStart) -> {
+        ArrayList<Integer> offsets = new ArrayList<>(size + 1);
+        int start = writePos;
+        for (int i = 0; i < size; ++i) {
+          int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize);
+          offsets.add(writePos - start);
+          appendVariantImpl(value, metadata, dataStart + offset);
+        }
+        offsets.add(writePos - start);
+        appendVariant(val);
+        finishWritingArray(start, offsets);
+        return null;
+      });
+      return;
+    }
+    PathSegment seg = segments[depth];
+    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);
+        int start = writePos;
+        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);
+          int newId = addKey(fieldKey);
+          fields.add(new FieldEntry(fieldKey, newId, writePos - start));
+          if (fieldKey.equals(key)) {
+            appendWithArrayAppendImpl(value, metadata, elementPos, segments, depth + 1, val);
+          } else {
+            appendVariantImpl(value, metadata, elementPos);
+          }
+        }
+        // A missing key is a no-op: the fields copied above already reproduce the input object.
+        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);
+        int start = writePos;
+        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) {
+            appendWithArrayAppendImpl(value, metadata, elementPos, segments, depth + 1, val);
+          } else {
+            appendVariantImpl(value, metadata, elementPos);
+          }
+        }
+        // An out-of-range index is a no-op: the elements copied above reproduce the input array.
+        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

@@ -602,6 +602,7 @@ VARIANT Functions
     schema_of_variant
     schema_of_variant_agg
     try_variant_get
+    variant_array_append
     variant_delete
     variant_get
     variant_insert

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 variant_array_append(
+    v: "ColumnOrName", path: Union[Column, str], value: "ColumnOrName"
+) -> Column:
+    path_col = path if isinstance(path, Column) else lit(path)
+    return _invoke_function("variant_array_append", _to_col(v), path_col, _to_col(value))
+
+
+variant_array_append.__doc__ = pysparkfuncs.variant_array_append.__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

@@ -480,6 +480,7 @@
     "schema_of_variant",
     "schema_of_variant_agg",
     "try_variant_get",
+    "variant_array_append",
     "variant_delete",
     "variant_get",
     "variant_insert",

python/pyspark/sql/functions/builtin.py

@@ -21786,6 +21786,64 @@ def variant_insert(v: "ColumnOrName", path: Union[Column, str], value: "ColumnOr
     )
 
 
+@_try_remote_functions
+def variant_array_append(
+    v: "ColumnOrName", path: Union[Column, str], value: "ColumnOrName"
+) -> Column:
+    """
+    Appends a value to the array in a variant at the given JSONPath location. Returns the variant
+    unchanged if a path key or index is absent, throws an error if a path segment hits a value of
+    an incompatible type or the target is not an array, and 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 target array. 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 zero or
+        more segments like `[123]`, `.name`, `['name']`, or `["name"]`.
+    value : :class:`~pyspark.sql.Column` or str
+        the value to append. Any expression castable to variant.
+
+    Returns
+    -------
+    :class:`~pyspark.sql.Column`
+        a variant column with `value` appended to the array at `path`
+
+    Examples
+    --------
+    >>> from pyspark.sql.functions import lit, parse_json, to_json, variant_array_append
+    >>> df = spark.createDataFrame([{
+    ...     'json': '''[[1, 2], 5]''',
+    ...     'path': '$[0]'
+    ... }])
+    >>> v = parse_json(df.json)
+    >>> df.select(to_json(variant_array_append(v, "$", lit(3))).alias("r")).collect()
+    [Row(r='[[1,2],5,3]')]
+    >>> df.select(to_json(variant_array_append(v, "$[5]", lit(3))).alias("r")).collect()
+    [Row(r='[[1,2],5]')]
+    >>> df.select(to_json(variant_array_append(v, df.path, lit(9))).alias("r")).collect()
+    [Row(r='[[1,2,9],5]')]
+    >>> nested = variant_array_append(v, "$", parse_json(lit('[4, 5]')))
+    >>> df.select(to_json(nested).alias("r")).collect()
+    [Row(r='[[1,2],5,[4,5]]')]
+    >>> df.select(variant_array_append(v, "$", 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_array_append",
+        _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

@@ -3508,8 +3508,14 @@ def test_parse_json(self):
     def test_variant_expressions(self):
         df = self.spark.createDataFrame(
             [
-                Row(json="""{ "a" : 1 }""", path="$.a", newpath="$.z"),
-                Row(json="""{ "b" : 2 }""", path="$.b", newpath="$.z"),
+                Row(json="""{ "a" : 1 }""", path="$.a", newpath="$.z", arr="[1, 2]", arrpath="$"),
+                Row(
+                    json="""{ "b" : 2 }""",
+                    path="$.b",
+                    newpath="$.z",
+                    arr="[[3], 4]",
+                    arrpath="$[0]",
+                ),
             ]
         )
         v = F.parse_json(df.json)
@@ -3534,6 +3540,16 @@ 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}'],
         )
+        arr = F.parse_json(df.arr)
+        check(
+            df.select(F.to_json(F.variant_array_append(arr, "$", F.lit(9)))),
+            ["[1,2,9]", "[[3],4,9]"],
+        )
+        check(df.select(F.to_json(F.variant_array_append(arr, "$", F.lit(None)))), [None, None])
+        check(
+            df.select(F.to_json(F.variant_array_append(arr, df.arrpath, F.lit(9)))),
+            ["[1,2,9]", "[[3,9],4]"],
+        )
         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,43 @@ object functions {
   def variant_insert(v: Column, path: String, value: Column): Column =
     Column.fn("variant_insert", v, lit(path), value)
 
+  /**
+   * Appends a value to the array in a variant at the given JSONPath location. Returns the variant
+   * unchanged if a path key or index is absent, throws an error if a path segment hits a value of
+   * an incompatible type or the target is not an array, and returns NULL if any argument is NULL.
+   *
+   * @param v
+   *   a variant column.
+   * @param path
+   *   the column containing the JSONPath string identifying the target array. A valid path should
+   *   start with `$` and is followed by zero or more segments like `[123]`, `.name`, `['name']`,
+   *   or `["name"]`.
+   * @param value
+   *   the value to append. Any expression castable to variant.
+   * @group variant_funcs
+   * @since 4.3.0
+   */
+  def variant_array_append(v: Column, path: Column, value: Column): Column =
+    Column.fn("variant_array_append", v, path, value)
+
+  /**
+   * Appends a value to the array in a variant at the given JSONPath location. Returns the variant
+   * unchanged if a path key or index is absent, throws an error if a path segment hits a value of
+   * an incompatible type or the target is not an array, and returns NULL if any argument is NULL.
+   *
+   * @param v
+   *   a variant column.
+   * @param path
+   *   the JSONPath identifying the target array. A valid path should start with `$` and is
+   *   followed by zero or more segments like `[123]`, `.name`, `['name']`, or `["name"]`.
+   * @param value
+   *   the value to append. Any expression castable to variant.
+   * @group variant_funcs
+   * @since 4.3.0
+   */
+  def variant_array_append(v: Column, path: String, value: Column): Column =
+    Column.fn("variant_array_append", 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

@@ -987,6 +987,7 @@ object FunctionRegistry {
     expression[IsValidVariant]("is_valid_variant"),
     expression[VariantDelete]("variant_delete"),
     expression[VariantInsert]("variant_insert"),
+    expression[VariantArrayAppend]("variant_array_append"),
 
     // Spatial
     expression[ST_AsBinary]("st_asbinary"),

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

@@ -77,13 +77,17 @@ object VariantExpressionEvalUtils {
 
   /**
    * Parse a JSONPath for a variant manipulation function. Throws `INVALID_VARIANT_PATH` on a
-   * malformed path or on the empty (root `$`) path.
+   * malformed path, or on the empty (root `$`) path unless `allowRoot` is set (as
+   * `variant_array_append` does, since `$` refers to the root array there).
    */
-  def parseVariantPath(pathValue: String, functionName: String): Array[VariantPathSegment] = {
+  def parseVariantPath(
+      pathValue: String,
+      functionName: String,
+      allowRoot: Boolean = false): Array[VariantPathSegment] = {
     val parsed = VariantPathParser.parse(pathValue).getOrElse {
       throw QueryExecutionErrors.invalidVariantPath(pathValue, functionName)
     }
-    if (parsed.isEmpty) {
+    if (!allowRoot && parsed.isEmpty) {
       throw QueryExecutionErrors.invalidVariantPath(pathValue, functionName)
     }
     parsed
@@ -166,6 +170,45 @@ object VariantExpressionEvalUtils {
     insertAtPath(input, javaSegments, pathStr, value, valueDataType, functionName)
   }
 
+  /**
+   * Append `value` to the array in `input` at `javaSegments`. `path` is the source string used in
+   * error messages. `value` is cast to a variant first; a size overflow maps to
+   * `VARIANT_SIZE_LIMIT` and a target that is not an array (or an incompatible path segment) maps
+   * to `VARIANT_PATH_TYPE_MISMATCH`.
+   */
+  def arrayAppendAtPath(
+      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.arrayAppendAtPath(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 arrayAppendAtPath(
+      input: VariantVal,
+      path: UTF8String,
+      value: Any,
+      valueDataType: DataType,
+      functionName: String): VariantVal = {
+    val pathStr = path.toString
+    val javaSegments = toJavaSegments(parseVariantPath(pathStr, functionName, allowRoot = true))
+    arrayAppendAtPath(input, javaSegments, pathStr, value, valueDataType, 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