[SPARK-51975][SQL] Add variant_from_arrays and variant_from_entries

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

@@ -605,6 +605,8 @@ VARIANT Functions
     variant_array_append
     try_variant_array_append
     variant_delete
+    variant_from_arrays
+    variant_from_entries
     variant_get
     variant_insert
     try_variant_insert

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

@@ -2185,6 +2185,20 @@ def to_variant_object(col: "ColumnOrName") -> Column:
 to_variant_object.__doc__ = pysparkfuncs.to_variant_object.__doc__
 
 
+def variant_from_arrays(keys: "ColumnOrName", values: "ColumnOrName") -> Column:
+    return _invoke_function_over_columns("variant_from_arrays", keys, values)
+
+
+variant_from_arrays.__doc__ = pysparkfuncs.variant_from_arrays.__doc__
+
+
+def variant_from_entries(entries: "ColumnOrName") -> Column:
+    return _invoke_function_over_columns("variant_from_entries", entries)
+
+
+variant_from_entries.__doc__ = pysparkfuncs.variant_from_entries.__doc__
+
+
 def parse_json(col: "ColumnOrName") -> Column:
     return _invoke_function("parse_json", _to_col(col))
 

python/pyspark/sql/functions/__init__.py

@@ -483,6 +483,8 @@
     "variant_array_append",
     "try_variant_array_append",
     "variant_delete",
+    "variant_from_arrays",
+    "variant_from_entries",
     "variant_get",
     "variant_insert",
     "try_variant_insert",

python/pyspark/sql/functions/builtin.py

@@ -22499,6 +22499,75 @@ def to_variant_object(
     return _invoke_function("to_variant_object", _to_java_column(col))
 
 
+@_try_remote_functions
+def variant_from_arrays(keys: "ColumnOrName", values: "ColumnOrName") -> Column:
+    """
+    Creates a variant object from the given arrays of keys and values. The keys must be non-null
+    strings and the two arrays must have the same length.
+
+    .. versionadded:: 4.3.0
+
+    Parameters
+    ----------
+    keys : :class:`~pyspark.sql.Column` or column name
+        an array of string keys.
+    values : :class:`~pyspark.sql.Column` or column name
+        an array of values.
+
+    Returns
+    -------
+    :class:`~pyspark.sql.Column`
+        a new column of VariantType.
+
+    See Also
+    --------
+    :meth:`pyspark.sql.functions.variant_from_entries`
+    :meth:`pyspark.sql.functions.to_variant_object`
+
+    Examples
+    --------
+    >>> from pyspark.sql import functions as sf
+    >>> df = spark.sql("SELECT array('a', 'b') AS keys, array(1, 2) AS values")
+    >>> df.select(sf.variant_from_arrays("keys", "values").cast("string").alias("r")).collect()
+    [Row(r='{"a":1,"b":2}')]
+    """
+    return _invoke_function_over_columns("variant_from_arrays", keys, values)
+
+
+@_try_remote_functions
+def variant_from_entries(entries: "ColumnOrName") -> Column:
+    """
+    Creates a variant object from an array of key/value struct entries. The keys must be non-null
+    strings.
+
+    .. versionadded:: 4.3.0
+
+    Parameters
+    ----------
+    entries : :class:`~pyspark.sql.Column` or column name
+        an array of key/value structs, where the first field is a string key and the second field
+        is the value.
+
+    Returns
+    -------
+    :class:`~pyspark.sql.Column`
+        a new column of VariantType.
+
+    See Also
+    --------
+    :meth:`pyspark.sql.functions.variant_from_arrays`
+    :meth:`pyspark.sql.functions.to_variant_object`
+
+    Examples
+    --------
+    >>> from pyspark.sql import functions as sf
+    >>> df = spark.sql("SELECT array(struct('a', 1), struct('b', 2)) AS entries")
+    >>> df.select(sf.variant_from_entries("entries").cast("string").alias("r")).collect()
+    [Row(r='{"a":1,"b":2}')]
+    """
+    return _invoke_function_over_columns("variant_from_entries", entries)
+
+
 @_try_remote_functions
 def parse_json(
     col: "ColumnOrName",

python/pyspark/sql/tests/test_functions.py

@@ -3661,6 +3661,23 @@ def test_to_variant_object(self):
         ).collect()
         self.assertEqual("""{"a":1}""", actual[0]["var"])
 
+    def test_variant_from_arrays_and_entries(self):
+        df = self.spark.createDataFrame(
+            [(["a", "b"], [1, 2])], "keys array<string>, values array<int>"
+        )
+        actual = df.select(
+            F.to_json(F.variant_from_arrays("keys", "values")).alias("var"),
+        ).collect()
+        self.assertEqual("""{"a":1,"b":2}""", actual[0]["var"])
+
+        df2 = self.spark.createDataFrame(
+            [([("a", 1), ("b", 2)],)], "entries array<struct<k string, v int>>"
+        )
+        actual2 = df2.select(
+            F.to_json(F.variant_from_entries("entries")).alias("var"),
+        ).collect()
+        self.assertEqual("""{"a":1,"b":2}""", actual2[0]["var"])
+
     def test_schema_of_csv(self):
         with self.assertRaises(PySparkTypeError) as pe:
             F.schema_of_csv(1)

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

@@ -14081,6 +14081,36 @@ object functions {
    */
   def to_variant_object(col: Column): Column = Column.fn("to_variant_object", col)
 
+  /**
+   * Creates a variant object from the given arrays of keys and values. The keys must be non-null
+   * strings and the two arrays must have the same length.
+   *
+   * @param keys
+   *   a column that evaluates to an array of string keys.
+   * @param values
+   *   a column that evaluates to an array of values.
+   * @group variant_funcs
+   * @since 4.3.0
+   * @return
+   *   Returns a column that evaluates to a variant.
+   */
+  def variant_from_arrays(keys: Column, values: Column): Column =
+    Column.fn("variant_from_arrays", keys, values)
+
+  /**
+   * Creates a variant object from an array of key/value struct entries. The keys must be non-null
+   * strings.
+   *
+   * @param entries
+   *   a column that evaluates to an array of key/value structs.
+   * @group variant_funcs
+   * @since 4.3.0
+   * @return
+   *   Returns a column that evaluates to a variant.
+   */
+  def variant_from_entries(entries: Column): Column =
+    Column.fn("variant_from_entries", entries)
+
   /**
    * Check if a variant value is a variant null. Returns true if and only if the input is a
    * variant null and false otherwise (including in the case of SQL NULL).

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

@@ -985,6 +985,8 @@ object FunctionRegistry {
     expression[SchemaOfVariant]("schema_of_variant"),
     expression[SchemaOfVariantAgg]("schema_of_variant_agg"),
     expression[ToVariantObject]("to_variant_object"),
+    expression[VariantFromArrays]("variant_from_arrays"),
+    expression[VariantFromEntries]("variant_from_entries"),
     expression[IsValidVariant]("is_valid_variant"),
     expression[VariantDelete]("variant_delete"),
     expressionBuilder("variant_insert", VariantInsertExpressionBuilder),

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

@@ -274,6 +274,61 @@ object VariantExpressionEvalUtils {
     new VariantVal(v.getValue, v.getMetadata)
   }
 
+  /**
+   * Build a variant object directly from a keys array and a values array, without materializing an
+   * intermediate map. Keys must be non-null strings and the two arrays must have equal length. A
+   * null key raises `NULL_MAP_KEY`, a duplicate key raises `VARIANT_DUPLICATE_KEY` (matching
+   * to_variant_object), and null values are kept as variant null.
+   */
+  def variantFromArrays(keys: ArrayData, values: ArrayData, valueType: DataType): VariantVal = {
+    if (keys.numElements() != values.numElements()) {
+      // Reuse the same error map_from_arrays raises for a keys/values length mismatch.
+      throw QueryExecutionErrors.mapDataKeyArrayLengthDiffersFromValueArrayLengthError()
+    }
+    val builder = new VariantBuilder(false)
+    val start = builder.getWritePos
+    val fields = new java.util.ArrayList[VariantBuilder.FieldEntry](keys.numElements())
+    for (i <- 0 until keys.numElements()) {
+      if (keys.isNullAt(i)) {
+        throw QueryExecutionErrors.nullAsMapKeyNotAllowedError()
+      }
+      val key = keys.getUTF8String(i).toString
+      val id = builder.addKey(key)
+      fields.add(new VariantBuilder.FieldEntry(key, id, builder.getWritePos - start))
+      val value = if (values.isNullAt(i)) null else values.get(i, valueType)
+      buildVariant(builder, value, valueType)
+    }
+    builder.finishWritingObject(start, fields)
+    val v = builder.result()
+    new VariantVal(v.getValue, v.getMetadata)
+  }
+
+  /**
+   * Build a variant object directly from an array of key/value struct entries, without an
+   * intermediate map. Keys must be non-null strings. A null key raises `NULL_MAP_KEY`, a
+   * duplicate key raises `VARIANT_DUPLICATE_KEY`, null values are kept as variant null, and null
+   * entries are handled by the caller.
+   */
+  def variantFromEntries(entries: ArrayData, valueType: DataType): VariantVal = {
+    val builder = new VariantBuilder(false)
+    val start = builder.getWritePos
+    val fields = new java.util.ArrayList[VariantBuilder.FieldEntry](entries.numElements())
+    for (i <- 0 until entries.numElements()) {
+      val entry = entries.getStruct(i, 2)
+      if (entry.isNullAt(0)) {
+        throw QueryExecutionErrors.nullAsMapKeyNotAllowedError()
+      }
+      val key = entry.getUTF8String(0).toString
+      val id = builder.addKey(key)
+      fields.add(new VariantBuilder.FieldEntry(key, id, builder.getWritePos - start))
+      val value = if (entry.isNullAt(1)) null else entry.get(1, valueType)
+      buildVariant(builder, value, valueType)
+    }
+    builder.finishWritingObject(start, fields)
+    val v = builder.result()
+    new VariantVal(v.getValue, v.getMetadata)
+  }
+
   /** Returns `true` if a data type is or has a child variant type. */
   def typeContainsVariant(dt: DataType): Boolean = dt match {
     case _: VariantType => true