This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 9299b130825d [SPARK-58875][SQL][ML][PYTHON][CONNECT] Add WrapUDT
expression for zero-copy linalg conversion
9299b130825d is described below
commit 9299b130825deb4d80b4835001a0cb96f6a09094
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Mon Aug 24 15:06:46 2026 +0800
[SPARK-58875][SQL][ML][PYTHON][CONNECT] Add WrapUDT expression for
zero-copy linalg conversion
### What changes were proposed in this pull request?
This PR adds a Catalyst `WrapUDT` expression that wraps a column with a
target UDT when the
column data type matches the UDT's underlying SQL type. `WrapUDT` is paired
with the existing
`UnwrapUDT`, so callers can convert between a UDT column and its underlying
Catalyst value without
copying the data.
It exposes `wrap_udt` through the Scala SQL functions API and PySpark SQL
functions, including
Spark Connect. The Scala API includes both `wrap_udt(Column,
UserDefinedType[_])` and a
`wrap_udt(Column, Column)` overload so Python can pass the target UDT JSON
as a literal column,
similar to schema handling in `from_json`.
This PR also adds ML-side tests for zero-copy conversions across `.ml` and
`.mllib` vector and
matrix UDTs by composing `unwrap_udt` and `wrap_udt`.
### Why are the changes needed?
Spark ML has parallel `.ml` and `.mllib` linalg UDTs whose serialized SQL
layouts can match even
when the logical UDT classes differ. Existing conversion paths often
deserialize the UDT value and
create a new vector or matrix object even when only the logical UDT wrapper
needs to change.
This UDT deserialization and object-copying cost is a long-standing issue
from the first day of the
DataFrame-based `.ml` APIs: ML pipelines exchange linalg values through
DataFrame UDT columns, and
moving values across `.ml` and `.mllib` UDT boundaries has historically
gone through object
conversion even when the Catalyst payload already has the right layout.
`WrapUDT` provides a typed Catalyst expression for changing the logical UDT
over an already matching
underlying SQL value. This enables zero-copy linalg conversions and avoids
per-row object conversion
for these cases.
### Does this PR introduce _any_ user-facing change?
Yes. It adds a public `wrap_udt` SQL function API in Scala and PySpark,
with Spark Connect support.
The function wraps a column as a target UDT when the column's data type
matches the UDT's underlying
SQL type. Existing `unwrap_udt` behavior is unchanged.
For example, in PySpark:
```python
from pyspark.sql import Row, functions as sf
from pyspark.sql.types import StructField, StructType
from pyspark.ml.linalg import VectorUDT
schema = StructType([StructField("vec", VectorUDT.sqlType(), True)])
df = spark.createDataFrame(
[(Row(type=1, size=None, indices=None, values=[1.0, 2.0, 3.0]),)],
schema)
df.select("*", sf.wrap_udt("vec", VectorUDT())).show()
```
Output:
```text
+--------------------+-------------+
| vec|wrap_udt(vec)|
+--------------------+-------------+
|{1, NULL, NULL, [...|[1.0,2.0,3.0]|
+--------------------+-------------+
```
### How was this patch tested?
Added and updated tests for:
- Catalyst `WrapUDT` positive and negative type checking.
- SQL `wrap_udt` and `unwrap_udt` round trips.
- ML vector and matrix UDT conversions between `.ml` and `.mllib` linalg
types.
- PySpark `wrap_udt` examples for `VectorUDT` and `MatrixUDT`.
- Spark Connect JVM plan generation and golden files for `wrap_udt`.
Ran:
```
git diff --check upstream/master...HEAD
conda run -n spark-dev-313 python -m py_compile
python/pyspark/sql/functions/builtin.py
build/sbt -java-home /usr/lib/jvm/java-17-openjdk-amd64 sql-api/compile
build/sbt -java-home /usr/lib/jvm/java-17-openjdk-amd64
"connect-client-jvm/testOnly org.apache.spark.sql.PlanGenerationTestSuite -- -z
wrap_udt"
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #58124 from zhengruifeng/spark-58875-wrap-udt.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../scala/org/apache/spark/ml/FunctionsSuite.scala | 128 ++++++++++++++++++++-
.../source/reference/pyspark.sql/functions.rst | 1 +
python/pyspark/ml/tests/test_linalg.py | 106 ++++++++++++++++-
python/pyspark/sql/connect/functions/builtin.py | 24 ++++
python/pyspark/sql/functions/__init__.py | 1 +
python/pyspark/sql/functions/builtin.py | 109 ++++++++++++++++++
.../scala/org/apache/spark/sql/functions.scala | 26 +++++
.../sql/catalyst/analysis/FunctionRegistry.scala | 1 +
.../spark/sql/catalyst/expressions/UnwrapUDT.scala | 2 +
.../expressions/{UnwrapUDT.scala => WrapUDT.scala} | 46 ++++++--
.../expressions/WrapUDTExpressionSuite.scala | 115 ++++++++++++++++++
.../apache/spark/sql/PlanGenerationTestSuite.scala | 10 ++
.../test/scala/org/apache/spark/sql/TestUDT.scala | 42 +++++++
.../explain-results/function_unwrap_udt.explain | 2 +
.../explain-results/function_wrap_udt.explain | 2 +
.../query-tests/queries/function_unwrap_udt.json | 127 ++++++++++++++++++++
.../queries/function_unwrap_udt.proto.bin | Bin 0 -> 1281 bytes
.../query-tests/queries/function_wrap_udt.json | 104 +++++++++++++++++
.../queries/function_wrap_udt.proto.bin | Bin 0 -> 1077 bytes
.../test/scala/org/apache/spark/sql/TestUDT.scala | 42 +++++++
.../apache/spark/sql/DataFrameFunctionsSuite.scala | 1 +
.../apache/spark/sql/UserDefinedTypeSuite.scala | 29 ++++-
22 files changed, 903 insertions(+), 15 deletions(-)
diff --git a/mllib/src/test/scala/org/apache/spark/ml/FunctionsSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/FunctionsSuite.scala
index 7fcb1d2fbfbb..4608d9bd683d 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/FunctionsSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/FunctionsSuite.scala
@@ -19,16 +19,50 @@ package org.apache.spark.ml
import org.apache.spark.SparkException
import org.apache.spark.ml.functions._
-import org.apache.spark.ml.linalg.{Vector, Vectors}
+import org.apache.spark.ml.linalg.{Matrices, MatrixUDT, Vector, Vectors,
VectorUDT}
import org.apache.spark.ml.util.MLTest
-import org.apache.spark.mllib.linalg.{Vectors => OldVectors}
-import org.apache.spark.sql.AnalysisException
-import org.apache.spark.sql.functions.col
+import org.apache.spark.mllib.linalg.{Matrices => OldMatrices, MatrixUDT =>
OldMatrixUDT,
+ Vectors => OldVectors, VectorUDT => OldVectorUDT}
+import org.apache.spark.sql.{AnalysisException, DataFrame, Row}
+import org.apache.spark.sql.functions.{col, unwrap_udt, wrap_udt}
+import org.apache.spark.sql.types.{StructField, StructType, UserDefinedType}
class FunctionsSuite extends MLTest {
import testImplicits._
+ private def checkWrapUDTConversion(
+ df: DataFrame,
+ targetUDT: UserDefinedType[_],
+ expected: Any): Unit = {
+ val converted = df.select(wrap_udt(unwrap_udt(col("value")),
targetUDT).as("value"))
+ assert(converted.schema("value").dataType === targetUDT)
+ assert(converted.first().get(0) === expected)
+ }
+
+ private def checkWrapUDTTypeMismatch(df: DataFrame, targetUDT:
UserDefinedType[_]): Unit = {
+ val e = intercept[AnalysisException] {
+ df.select(wrap_udt(unwrap_udt(col("value")),
targetUDT).as("value")).collect()
+ }
+ assert(e.getCondition === "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE")
+ }
+
+ private def checkNullableWrapUDTConversion(
+ value: Any,
+ sourceUDT: UserDefinedType[_],
+ targetUDT: UserDefinedType[_],
+ expected: Any): Unit = {
+ val schema = StructType(Seq(StructField("value", sourceUDT, nullable =
true)))
+ val df = spark.createDataFrame(
+ spark.sparkContext.parallelize(Seq(Row(value), Row(null))),
+ schema)
+ val converted = df.select(wrap_udt(unwrap_udt(col("value")),
targetUDT).as("value"))
+
+ assert(converted.schema("value").dataType === targetUDT)
+ assert(converted.schema("value").nullable)
+ assert(converted.collect().map(_.get(0)).toSeq === Seq(expected, null))
+ }
+
test("test vector_to_array") {
val df = Seq(
(Vectors.dense(1.0, 2.0, 3.0), OldVectors.dense(10.0, 20.0, 30.0)),
@@ -130,4 +164,90 @@ class FunctionsSuite extends MLTest {
val result = df.select(array_argmax(col("arr"))).as[Int].collect()
assert(result === Array(2, 1, 0, 1, 0, -1))
}
+
+ test("wrap and unwrap vector and matrix UDT columns") {
+ val oldVector = OldVectors.sparse(3, Array(1), Array(2.0))
+ val oldVectorDF = Seq(Tuple1(oldVector)).toDF("value")
+ checkWrapUDTConversion(
+ oldVectorDF,
+ new OldVectorUDT,
+ oldVector)
+ checkWrapUDTConversion(
+ oldVectorDF,
+ new VectorUDT,
+ oldVector.asML)
+ checkWrapUDTTypeMismatch(
+ oldVectorDF,
+ new OldMatrixUDT)
+ checkWrapUDTTypeMismatch(
+ oldVectorDF,
+ new MatrixUDT)
+
+ val mlVector = Vectors.dense(1.0, 2.0)
+ val mlVectorDF = Seq(Tuple1(mlVector)).toDF("value")
+ checkWrapUDTConversion(
+ mlVectorDF,
+ new VectorUDT,
+ mlVector)
+ checkWrapUDTConversion(
+ mlVectorDF,
+ new OldVectorUDT,
+ OldVectors.fromML(mlVector))
+ checkWrapUDTTypeMismatch(
+ mlVectorDF,
+ new OldMatrixUDT)
+ checkWrapUDTTypeMismatch(
+ mlVectorDF,
+ new MatrixUDT)
+
+ val oldMatrix = OldMatrices.dense(2, 2, Array(1.0, 2.0, 3.0, 4.0))
+ val oldMatrixDF = Seq(Tuple1(oldMatrix)).toDF("value")
+ checkWrapUDTConversion(
+ oldMatrixDF,
+ new OldMatrixUDT,
+ oldMatrix)
+ checkWrapUDTConversion(
+ oldMatrixDF,
+ new MatrixUDT,
+ oldMatrix.asML)
+ checkWrapUDTTypeMismatch(
+ oldMatrixDF,
+ new OldVectorUDT)
+ checkWrapUDTTypeMismatch(
+ oldMatrixDF,
+ new VectorUDT)
+
+ val mlMatrix = Matrices.dense(2, 2, Array(1.0, 2.0, 3.0, 4.0))
+ val mlMatrixDF = Seq(Tuple1(mlMatrix)).toDF("value")
+ checkWrapUDTConversion(
+ mlMatrixDF,
+ new MatrixUDT,
+ mlMatrix)
+ checkWrapUDTConversion(
+ mlMatrixDF,
+ new OldMatrixUDT,
+ OldMatrices.fromML(mlMatrix))
+ checkWrapUDTTypeMismatch(
+ mlMatrixDF,
+ new OldVectorUDT)
+ checkWrapUDTTypeMismatch(
+ mlMatrixDF,
+ new VectorUDT)
+ }
+
+ test("wrap and unwrap nullable vector UDT columns") {
+ val oldVector = OldVectors.sparse(3, Array(1), Array(2.0))
+ checkNullableWrapUDTConversion(
+ oldVector,
+ new OldVectorUDT,
+ new VectorUDT,
+ oldVector.asML)
+
+ val mlVector = Vectors.dense(1.0, 2.0)
+ checkNullableWrapUDTConversion(
+ mlVector,
+ new VectorUDT,
+ new OldVectorUDT,
+ OldVectors.fromML(mlVector))
+ }
}
diff --git a/python/docs/source/reference/pyspark.sql/functions.rst
b/python/docs/source/reference/pyspark.sql/functions.rst
index 031c48734c49..94c5a9a87e13 100644
--- a/python/docs/source/reference/pyspark.sql/functions.rst
+++ b/python/docs/source/reference/pyspark.sql/functions.rst
@@ -779,6 +779,7 @@ UDF, UDTF and UDT
udf
udtf
unwrap_udt
+ wrap_udt
Table-Valued Functions
diff --git a/python/pyspark/ml/tests/test_linalg.py
b/python/pyspark/ml/tests/test_linalg.py
index 17185ad7b5fd..eea293188ef2 100644
--- a/python/pyspark/ml/tests/test_linalg.py
+++ b/python/pyspark/ml/tests/test_linalg.py
@@ -19,6 +19,7 @@ import array as pyarray
from numpy import arange, array, array_equal, inf, ones, tile, zeros
+from pyspark.errors import AnalysisException
from pyspark.serializers import CPickleSerializer
from pyspark.ml.linalg import (
DenseMatrix,
@@ -32,7 +33,8 @@ from pyspark.ml.linalg import (
)
from pyspark.testing.mllibutils import MLlibTestCase
from pyspark.sql import Row
-from pyspark.sql.functions import unwrap_udt
+from pyspark.sql.functions import lit, unwrap_udt, wrap_udt
+from pyspark.sql.types import StructField, StructType
class VectorTests(MLlibTestCase):
@@ -367,6 +369,49 @@ class VectorUDTTests(MLlibTestCase):
]
self.assertEqual(results, expected)
+ def test_wrap_udt(self):
+ schema = StructType([StructField("vec", VectorUDT.sqlType(), True)])
+ vector_struct = Row("type", "size", "indices", "values")
+ df = self.spark.createDataFrame(
+ [
+ (vector_struct(1, None, None, [1.0, 2.0]),),
+ (vector_struct(0, 2, [1], [2.0]),),
+ ],
+ schema,
+ )
+ wrapped = df.select(wrap_udt("vec", VectorUDT()).alias("vec"))
+
+ self.assertEqual(wrapped.schema["vec"].dataType, VectorUDT())
+ self.assertEqual(wrapped.collect(), [Row(vec=self.dv1),
Row(vec=self.sv1)])
+
+ def test_wrap_udt_type_mismatch(self):
+ schema = StructType([StructField("vec", VectorUDT.sqlType(), True)])
+ vector_struct = Row("type", "size", "indices", "values")
+ df = self.spark.createDataFrame(
+ [(vector_struct(1, None, None, [1.0, 2.0]),)],
+ schema,
+ )
+
+ with self.assertRaises(AnalysisException) as context:
+ df.select(wrap_udt("vec", MatrixUDT())).collect()
+ self.assertEqual(
+ context.exception.getCondition(),
"DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE"
+ )
+
+ def test_wrap_unwrap_udt_round_trip(self):
+ df = self.spark.createDataFrame([(self.dv1,), (self.sv1,)], ["vec"])
+ round_trip = df.select(wrap_udt(unwrap_udt("vec"),
VectorUDT()).alias("vec"))
+
+ self.assertEqual(round_trip.schema["vec"].dataType, VectorUDT())
+ self.assertEqual(round_trip.collect(), [Row(vec=self.dv1),
Row(vec=self.sv1)])
+
+ def test_wrap_unwrap_udt_round_trip_with_udt_column(self):
+ df = self.spark.createDataFrame([(self.dv1,), (self.sv1,)], ["vec"])
+ round_trip = df.select(wrap_udt(unwrap_udt("vec"),
lit(VectorUDT().json())).alias("vec"))
+
+ self.assertEqual(round_trip.schema["vec"].dataType, VectorUDT())
+ self.assertEqual(round_trip.collect(), [Row(vec=self.dv1),
Row(vec=self.sv1)])
+
def test_hashable(self):
_ = hash(VectorUDT())
@@ -404,6 +449,65 @@ class MatrixUDTTests(MLlibTestCase):
else:
raise ValueError("Expected a matrix but got type %r" % type(m))
+ def test_wrap_udt(self):
+ schema = StructType([StructField("mat", MatrixUDT.sqlType(), True)])
+ matrix_struct = Row(
+ "type",
+ "numRows",
+ "numCols",
+ "colPtrs",
+ "rowIndices",
+ "values",
+ "isTransposed",
+ )
+ df = self.spark.createDataFrame(
+ [
+ (matrix_struct(1, 3, 2, None, None, [0.0, 1.0, 4.0, 5.0, 9.0,
10.0], False),),
+ (matrix_struct(0, 1, 1, [0, 1], [0], [2.0], False),),
+ ],
+ schema,
+ )
+ wrapped = df.select(wrap_udt("mat", MatrixUDT()).alias("mat"))
+
+ self.assertEqual(wrapped.schema["mat"].dataType, MatrixUDT())
+ self.assertEqual(wrapped.collect(), [Row(mat=self.dm1),
Row(mat=self.sm1)])
+
+ def test_wrap_udt_type_mismatch(self):
+ schema = StructType([StructField("mat", MatrixUDT.sqlType(), True)])
+ matrix_struct = Row(
+ "type",
+ "numRows",
+ "numCols",
+ "colPtrs",
+ "rowIndices",
+ "values",
+ "isTransposed",
+ )
+ df = self.spark.createDataFrame(
+ [(matrix_struct(1, 3, 2, None, None, [0.0, 1.0, 4.0, 5.0, 9.0,
10.0], False),)],
+ schema,
+ )
+
+ with self.assertRaises(AnalysisException) as context:
+ df.select(wrap_udt("mat", VectorUDT())).collect()
+ self.assertEqual(
+ context.exception.getCondition(),
"DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE"
+ )
+
+ def test_wrap_unwrap_udt_round_trip(self):
+ df = self.spark.createDataFrame([(self.dm1,), (self.sm1,)], ["mat"])
+ round_trip = df.select(wrap_udt(unwrap_udt("mat"),
MatrixUDT()).alias("mat"))
+
+ self.assertEqual(round_trip.schema["mat"].dataType, MatrixUDT())
+ self.assertEqual(round_trip.collect(), [Row(mat=self.dm1),
Row(mat=self.sm1)])
+
+ def test_wrap_unwrap_udt_round_trip_with_udt_column(self):
+ df = self.spark.createDataFrame([(self.dm1,), (self.sm1,)], ["mat"])
+ round_trip = df.select(wrap_udt(unwrap_udt("mat"),
lit(MatrixUDT().json())).alias("mat"))
+
+ self.assertEqual(round_trip.schema["mat"].dataType, MatrixUDT())
+ self.assertEqual(round_trip.collect(), [Row(mat=self.dm1),
Row(mat=self.sm1)])
+
def test_hashable(self):
_ = hash(MatrixUDT())
diff --git a/python/pyspark/sql/connect/functions/builtin.py
b/python/pyspark/sql/connect/functions/builtin.py
index 9d622925bf19..453980fa83e1 100644
--- a/python/pyspark/sql/connect/functions/builtin.py
+++ b/python/pyspark/sql/connect/functions/builtin.py
@@ -65,7 +65,11 @@ from pyspark.sql.types import (
ArrayType,
MapType,
StringType,
+ UserDefinedType as _UserDefinedType,
)
+
+if TYPE_CHECKING:
+ from pyspark.sql.types import UserDefinedType
from pyspark.sql.utils import enum_to_value as _enum_to_value
# The implementations of pandas_udf, arrow_udf and udaf are embedded in
pyspark.sql.functions
@@ -5780,6 +5784,26 @@ def unwrap_udt(col: "ColumnOrName") -> Column:
unwrap_udt.__doc__ = pysparkfuncs.unwrap_udt.__doc__
+def wrap_udt(col: "ColumnOrName", udt: "Union[UserDefinedType, Column]") ->
Column:
+ if isinstance(udt, _UserDefinedType):
+ udt_col = lit(udt.json())
+ elif isinstance(udt, Column):
+ udt_col = udt
+ else:
+ raise PySparkTypeError(
+ errorClass="NOT_EXPECTED_TYPE",
+ messageParameters={
+ "expected_type": "UserDefinedType or Column",
+ "arg_name": "udt",
+ "arg_type": type(udt).__name__,
+ },
+ )
+ return _invoke_function("wrap_udt", _to_col(col), _to_col(udt_col))
+
+
+wrap_udt.__doc__ = pysparkfuncs.wrap_udt.__doc__
+
+
def udf(
f: Optional[Union[Callable[..., Any], "DataTypeOrString"]] = None,
returnType: "DataTypeOrString" = StringType(),
diff --git a/python/pyspark/sql/functions/__init__.py
b/python/pyspark/sql/functions/__init__.py
index 462509a5d7a2..4ef3b1d5067f 100644
--- a/python/pyspark/sql/functions/__init__.py
+++ b/python/pyspark/sql/functions/__init__.py
@@ -628,4 +628,5 @@ __all__ = [ # noqa: F405
"udtf",
"arrow_udtf",
"unwrap_udt",
+ "wrap_udt",
]
diff --git a/python/pyspark/sql/functions/builtin.py
b/python/pyspark/sql/functions/builtin.py
index 5acb3301be4d..1f5d240be04d 100644
--- a/python/pyspark/sql/functions/builtin.py
+++ b/python/pyspark/sql/functions/builtin.py
@@ -52,8 +52,12 @@ from pyspark.sql.types import (
MapType,
NumericType,
_from_numpy_type,
+ UserDefinedType as _UserDefinedType,
)
+if TYPE_CHECKING:
+ from pyspark.sql.types import UserDefinedType
+
# Keep UserDefinedFunction import for backwards compatible import; moved in
SPARK-22409
from pyspark.sql.udf import UserDefinedFunction, _create_py_udf # noqa: F401
from pyspark.sql.udtf import AnalyzeArgument, AnalyzeResult # noqa: F401
@@ -29199,6 +29203,10 @@ def unwrap_udt(col: "ColumnOrName") -> Column:
:class:`~pyspark.sql.Column`
The underlying representation.
+ See Also
+ --------
+ :meth:`pyspark.sql.functions.wrap_udt`
+
Examples
--------
Example 1: Unwrap ML-specific UDT - VectorUDT
@@ -29244,6 +29252,107 @@ def unwrap_udt(col: "ColumnOrName") -> Column:
return _invoke_function("unwrap_udt", _to_java_column(col))
+@_try_remote_functions
+def wrap_udt(col: "ColumnOrName", udt: "Union[UserDefinedType, Column]") ->
Column:
+ """
+ Wrap a column as a user-defined type.
+
+ .. versionadded:: 4.4.0
+
+ Parameters
+ ----------
+ col : :class:`~pyspark.sql.Column` or column name
+ The column to wrap. The column data type must match the UDT's
underlying SQL type.
+ udt : :class:`~pyspark.sql.types.UserDefinedType` or
:class:`~pyspark.sql.Column`
+ The target user-defined type, or a constant string column containing
its JSON
+ representation.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ A column of the target user-defined type.
+
+ See Also
+ --------
+ :meth:`pyspark.sql.functions.unwrap_udt`
+
+ Examples
+ --------
+ Example 1: Wrapping a vector struct as VectorUDT
+
+ >>> from pyspark.sql import functions as sf
+ >>> from pyspark.sql import Row
+ >>> from pyspark.sql.types import StructField, StructType
+ >>> from pyspark.ml.linalg import VectorUDT
+ >>> vector_schema = StructType([StructField("vec", VectorUDT.sqlType(),
True)])
+ >>> df = spark.createDataFrame(
+ ... [(Row(type=1, size=None, indices=None, values=[1.0, 2.0, 3.0]),)],
+ ... vector_schema)
+ >>> df.select("*", sf.wrap_udt("vec", VectorUDT())).show()
+ +--------------------+...+
+ | vec|wrap_udt(vec...|
+ +--------------------+...+
+ |{1, NULL, NULL, [...|...[1.0,2.0,3.0]|
+ +--------------------+...+
+ >>> row = df.select(sf.wrap_udt("vec", VectorUDT())).first()
+ >>> type(row[0])
+ <class 'pyspark.ml.linalg.DenseVector'>
+
+ Example 2: Wrapping a matrix struct as MatrixUDT
+
+ >>> from pyspark.sql import functions as sf
+ >>> from pyspark.sql import Row
+ >>> from pyspark.sql.types import StructField, StructType
+ >>> from pyspark.mllib.linalg import MatrixUDT
+ >>> matrix_schema = StructType([StructField("mat", MatrixUDT.sqlType(),
True)])
+ >>> df = spark.createDataFrame(
+ ... [(
+ ... Row(
+ ... type=1,
+ ... numRows=2,
+ ... numCols=2,
+ ... colPtrs=None,
+ ... rowIndices=None,
+ ... values=[1.0, 2.0, 3.0, 4.0],
+ ... isTransposed=False),
+ ... )],
+ ... matrix_schema)
+ >>> df.select("*", sf.wrap_udt("mat", MatrixUDT())).printSchema()
+ root
+ |-- mat: struct (nullable = true)
+ | |-- type: byte (nullable = false)
+ | |-- numRows: integer (nullable = false)
+ | |-- numCols: integer (nullable = false)
+ | |-- colPtrs: array (nullable = true)
+ | | |-- element: integer (containsNull = false)
+ | |-- rowIndices: array (nullable = true)
+ | | |-- element: integer (containsNull = false)
+ | |-- values: array (nullable = true)
+ | | |-- element: double (containsNull = false)
+ | |-- isTransposed: boolean (nullable = false)
+ |-- wrap_udt(mat...: matrix... (nullable = true)
+ >>> row = df.select(sf.wrap_udt("mat", MatrixUDT())).first()
+ >>> type(row[0])
+ <class 'pyspark.mllib.linalg.DenseMatrix'>
+ """
+ from pyspark.sql.classic.column import _to_java_column
+
+ if isinstance(udt, _UserDefinedType):
+ udt_col = lit(udt.json())
+ elif isinstance(udt, Column):
+ udt_col = udt
+ else:
+ raise PySparkTypeError(
+ errorClass="NOT_EXPECTED_TYPE",
+ messageParameters={
+ "expected_type": "UserDefinedType or Column",
+ "arg_name": "udt",
+ "arg_type": type(udt).__name__,
+ },
+ )
+ return _invoke_function("wrap_udt", _to_java_column(col),
_to_java_column(udt_col))
+
+
# ---------------------- Datasketch functions ------------------------------
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 9c79e1117332..307a82e5a44a 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
@@ -17856,6 +17856,32 @@ object functions {
*/
def unwrap_udt(column: Column): Column = Column.internalFn("unwrap_udt",
column)
+ /**
+ * Wrap a column as a user-defined type.
+ * @param column
+ * the column to wrap. The column data type must match the UDT's
underlying SQL type.
+ * @param udt
+ * the target user-defined type.
+ * @group udf_funcs
+ * @since 4.4.0
+ */
+ def wrap_udt(column: Column, udt: UserDefinedType[_]): Column = {
+ wrap_udt(column, lit(udt.json))
+ }
+
+ /**
+ * Wrap a column as a user-defined type.
+ * @param column
+ * the column to wrap. The column data type must match the UDT's
underlying SQL type.
+ * @param udt
+ * the target user-defined type as a constant JSON string column.
+ * @group udf_funcs
+ * @since 4.4.0
+ */
+ def wrap_udt(column: Column, udt: Column): Column = {
+ Column.internalFn("wrap_udt", column, udt)
+ }
+
// ---------------------- Vector Functions ----------------------
/**
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 dd70ea9d177a..4facf16cefbe 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
@@ -1123,6 +1123,7 @@ object FunctionRegistry {
registerInternalExpression[Days]("days")
registerInternalExpression[Hours]("hours")
registerInternalExpression[UnwrapUDT]("unwrap_udt")
+ registerInternalExpression[WrapUDT]("wrap_udt")
registerInternalExpression[MonotonicallyIncreasingID]("distributed_id",
setAlias = true)
registerInternalExpression[DistributedSequenceID]("distributed_sequence_id")
registerInternalExpression[PandasProduct]("pandas_product")
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/UnwrapUDT.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/UnwrapUDT.scala
index 249e3955a81f..8183e6a447bd 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/UnwrapUDT.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/UnwrapUDT.scala
@@ -25,6 +25,8 @@ import org.apache.spark.sql.types.{DataType, UserDefinedType}
/**
* Unwrap UDT data type column into its underlying type.
+ *
+ * @see [[WrapUDT]] for converting an underlying SQL type column to a UDT.
*/
case class UnwrapUDT(child: Expression) extends UnaryExpression with
NonSQLExpression {
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/UnwrapUDT.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/WrapUDT.scala
similarity index 55%
copy from
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/UnwrapUDT.scala
copy to
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/WrapUDT.scala
index 249e3955a81f..770dd1f8fa47 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/UnwrapUDT.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/WrapUDT.scala
@@ -17,39 +17,69 @@
package org.apache.spark.sql.catalyst.expressions
+import org.apache.spark.sql.AnalysisException
import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch
import org.apache.spark.sql.catalyst.expressions.Cast._
import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext,
ExprCode}
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.catalyst.util.TypeUtils.ordinalNumber
import org.apache.spark.sql.types.{DataType, UserDefinedType}
/**
- * Unwrap UDT data type column into its underlying type.
+ * Wrap a column with a UDT whose underlying SQL type matches the column data
type.
+ *
+ * @see [[UnwrapUDT]] for converting a UDT column to its underlying SQL type.
*/
-case class UnwrapUDT(child: Expression) extends UnaryExpression with
NonSQLExpression {
+case class WrapUDT(child: Expression, udt: UserDefinedType[_])
+ extends UnaryExpression with NonSQLExpression {
- override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode = child.genCode(ctx)
+ def this(child: Expression, udt: Expression) = {
+ this(child, WrapUDT.parseUDT(udt))
+ }
+
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode = {
+ child.genCode(ctx)
+ }
override def checkInputDataTypes(): TypeCheckResult = {
- if (child.dataType.isInstanceOf[UserDefinedType[_]]) {
+ if (DataTypeUtils.sameType(child.dataType, udt.sqlType)) {
TypeCheckResult.TypeCheckSuccess
} else {
DataTypeMismatch(
errorSubClass = "UNEXPECTED_INPUT_TYPE",
messageParameters = Map(
"paramIndex" -> ordinalNumber(0),
- "requiredType" -> toSQLType("UserDefinedType"),
+ "requiredType" -> toSQLType(udt.sqlType),
"inputSql" -> toSQLExpr(child),
"inputType" -> toSQLType(child.dataType)))
}
}
- override def dataType: DataType =
child.dataType.asInstanceOf[UserDefinedType[_]].sqlType
+
+ override def dataType: DataType = udt
override def nullSafeEval(input: Any): Any = input
- override def prettyName: String = "unwrap_udt"
+ override def prettyName: String = "wrap_udt"
- override protected def withNewChildInternal(newChild: Expression): UnwrapUDT
= {
+ override protected def withNewChildInternal(newChild: Expression): WrapUDT =
{
copy(child = newChild)
}
}
+
+object WrapUDT {
+ private def parseUDT(expression: Expression): UserDefinedType[_] = {
+ ExprUtils.evalTypeExpr(expression) match {
+ case udt: UserDefinedType[_] => udt
+ case dataType =>
+ throw new AnalysisException(
+ errorClass = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE",
+ messageParameters = Map(
+ "sqlExpr" -> toSQLExpr(expression),
+ "paramIndex" -> ordinalNumber(1),
+ "requiredType" -> toSQLType("UserDefinedType"),
+ "inputSql" -> toSQLExpr(expression),
+ "inputType" -> toSQLType(dataType)))
+ }
+ }
+}
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/WrapUDTExpressionSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/WrapUDTExpressionSuite.scala
new file mode 100644
index 000000000000..afd14131e5ec
--- /dev/null
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/WrapUDTExpressionSuite.scala
@@ -0,0 +1,115 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions
+
+import org.apache.spark.{SPARK_DOC_ROOT, SparkFunSuite}
+import org.apache.spark.sql.AnalysisException
+import org.apache.spark.sql.catalyst.FunctionIdentifier
+import org.apache.spark.sql.catalyst.analysis.FunctionRegistry
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch
+import org.apache.spark.sql.catalyst.expressions.Cast.{toSQLExpr, toSQLType}
+import org.apache.spark.sql.catalyst.util.GenericArrayData
+import org.apache.spark.sql.catalyst.util.TypeUtils.ordinalNumber
+import org.apache.spark.sql.types.{BooleanType, IntegerType, StringType,
TestUDT}
+
+class WrapUDTExpressionSuite extends SparkFunSuite with ExpressionEvalHelper {
+
+ test("WrapUDT should use target UDT with matching SQL type") {
+ val udt = new TestUDT.MyDenseVectorUDT()
+ val data = new GenericArrayData(Array[Any](1.0, 2.0))
+ val wrapUDTExpression = WrapUDT(Literal.create(data, udt.sqlType), udt)
+
+ assert(wrapUDTExpression.checkInputDataTypes().isSuccess)
+ assert(wrapUDTExpression.dataType == udt)
+ checkEvaluation(wrapUDTExpression, data)
+ }
+
+ test("WrapUDT should parse target UDT from foldable expression") {
+ val udt = new TestUDT.MyDenseVectorUDT()
+ val data = new GenericArrayData(Array[Any](1.0, 2.0))
+ val json = udt.json
+ val target = Concat(Seq(
+ Literal.create(json.substring(0, 8), StringType),
+ Literal.create(json.substring(8), StringType)))
+ val wrapUDTExpression = new WrapUDT(Literal.create(data, udt.sqlType),
target)
+
+ assert(wrapUDTExpression.checkInputDataTypes().isSuccess)
+ assert(wrapUDTExpression.dataType == udt)
+ checkEvaluation(wrapUDTExpression, data)
+ }
+
+ test("WrapUDT target expression should be a UDT") {
+ val target = Literal.create("int", StringType)
+ checkError(
+ exception = intercept[AnalysisException] {
+ new WrapUDT(Literal.create(1, IntegerType), target)
+ },
+ condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE",
+ parameters = Map(
+ "sqlExpr" -> toSQLExpr(target),
+ "paramIndex" -> ordinalNumber(1),
+ "requiredType" -> toSQLType("UserDefinedType"),
+ "inputSql" -> toSQLExpr(target),
+ "inputType" -> toSQLType(IntegerType)))
+ }
+
+ test("WrapUDT target expression should be foldable") {
+ val target = AttributeReference("udt", StringType)()
+ checkError(
+ exception = intercept[AnalysisException] {
+ new WrapUDT(Literal.create(1, IntegerType), target)
+ },
+ condition = "INVALID_SCHEMA.NON_STRING_LITERAL",
+ parameters = Map("inputSchema" -> toSQLExpr(target)))
+ }
+
+ test("WrapUDT should reject wrong number of arguments through
FunctionRegistry") {
+ val expression = Literal.create(1, IntegerType)
+ val target = Literal.create(new TestUDT.MyDenseVectorUDT().json,
StringType)
+
+ Seq(
+ Seq(expression) -> "1",
+ Seq(expression, target, Literal.create("extra", StringType)) ->
"3").foreach {
+ case (arguments, actualNum) =>
+ checkError(
+ exception = intercept[AnalysisException] {
+
FunctionRegistry.internal.lookupFunction(FunctionIdentifier("wrap_udt"),
arguments)
+ },
+ condition = "WRONG_NUM_ARGS.WITHOUT_SUGGESTION",
+ parameters = Map(
+ "functionName" -> "`wrap_udt`",
+ "expectedNum" -> "2",
+ "actualNum" -> actualNum,
+ "docroot" -> SPARK_DOC_ROOT))
+ }
+ }
+
+ test("WrapUDT input type should match target UDT SQL type") {
+ val b1 = Literal.create(false, BooleanType)
+ val udt = new TestUDT.MyDenseVectorUDT()
+ val wrapUDTExpression = WrapUDT(b1, udt)
+ assert(wrapUDTExpression.checkInputDataTypes() ==
+ DataTypeMismatch(
+ errorSubClass = "UNEXPECTED_INPUT_TYPE",
+ messageParameters = Map(
+ "paramIndex" -> ordinalNumber(0),
+ "requiredType" -> toSQLType(udt.sqlType),
+ "inputSql" -> "\"false\"",
+ "inputType" -> "\"BOOLEAN\"")))
+ }
+}
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 6346cf4b5cd8..b117f1a0b5cd 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
@@ -106,6 +106,8 @@ class PlanGenerationTestSuite extends ConnectFunSuite with
Logging {
private val printer = JsonFormat.printer().usingTypeRegistry(registry)
+ private val testUDT = new TestUDT.NewArrayUDT()
+
private var session: SparkSession = _
override protected def beforeAll(): Unit = {
@@ -3300,6 +3302,14 @@ class PlanGenerationTestSuite extends ConnectFunSuite
with Logging {
fn.typeof(fn.col("g"))
}
+ functionTest("wrap_udt") {
+ fn.wrap_udt(fn.array(fn.col("b")), testUDT)
+ }
+
+ functionTest("unwrap_udt") {
+ fn.unwrap_udt(fn.wrap_udt(fn.array(fn.col("b")), testUDT))
+ }
+
functionTest("stack") {
fn.stack(lit(2), fn.col("g"), fn.col("g"), fn.col("g"))
}
diff --git
a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/TestUDT.scala
b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/TestUDT.scala
new file mode 100644
index 000000000000..77704b8a57c3
--- /dev/null
+++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/TestUDT.scala
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql
+
+import org.apache.spark.sql.types.{ArrayType, DataType, DoubleType,
SQLUserDefinedType, UserDefinedType}
+
+object TestUDT {
+
+ @SQLUserDefinedType(udt = classOf[NewArrayUDT])
+ private[sql] class NewArray(val values: Array[Double]) extends Serializable
+
+ private[sql] class NewArrayUDT extends UserDefinedType[NewArray] {
+
+ override def sqlType: DataType = ArrayType(DoubleType, containsNull =
false)
+
+ override def serialize(obj: NewArray): Any = obj.values
+
+ override def deserialize(datum: Any): NewArray = {
+ datum match {
+ case values: Array[_] =>
+ new NewArray(values.map(_.asInstanceOf[Double]))
+ }
+ }
+
+ override def userClass: Class[NewArray] = classOf[NewArray]
+ }
+}
diff --git
a/sql/connect/common/src/test/resources/query-tests/explain-results/function_unwrap_udt.explain
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_unwrap_udt.explain
new file mode 100644
index 000000000000..08b1a3eb011d
--- /dev/null
+++
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_unwrap_udt.explain
@@ -0,0 +1,2 @@
+Project [unwrap_udt(wrap_udt(array(b#0),
org.apache.spark.sql.TestUDT$NewArrayUDT)) AS unwrap_udt(wrap_udt(array(b),
org.apache.spark.sql.TestUDT$NewArrayUDT))#0]
++- LocalRelation <empty>, [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_wrap_udt.explain
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_wrap_udt.explain
new file mode 100644
index 000000000000..a1fb72ba2b77
--- /dev/null
+++
b/sql/connect/common/src/test/resources/query-tests/explain-results/function_wrap_udt.explain
@@ -0,0 +1,2 @@
+Project [wrap_udt(array(b#0), org.apache.spark.sql.TestUDT$NewArrayUDT) AS
wrap_udt(array(b), org.apache.spark.sql.TestUDT$NewArrayUDT)#0]
++- LocalRelation <empty>, [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_unwrap_udt.json
b/sql/connect/common/src/test/resources/query-tests/queries/function_unwrap_udt.json
new file mode 100644
index 000000000000..9a8af401c78d
--- /dev/null
+++
b/sql/connect/common/src/test/resources/query-tests/queries/function_unwrap_udt.json
@@ -0,0 +1,127 @@
+{
+ "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": "unwrap_udt",
+ "arguments": [{
+ "unresolvedFunction": {
+ "functionName": "wrap_udt",
+ "arguments": [{
+ "unresolvedFunction": {
+ "functionName": "array",
+ "arguments": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "b"
+ },
+ "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": "array",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass":
"org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "literal": {
+ "string":
"{\"type\":\"udt\",\"class\":\"org.apache.spark.sql.TestUDT$NewArrayUDT\",\"pyClass\":null,\"sqlType\":{\"type\":\"array\",\"elementType\":\"double\",\"containsNull\":false}}"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "wrap_udt",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass":
"org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }],
+ "isInternal": true
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "wrap_udt",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass":
"org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }],
+ "isInternal": true
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "unwrap_udt",
+ "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_unwrap_udt.proto.bin
b/sql/connect/common/src/test/resources/query-tests/queries/function_unwrap_udt.proto.bin
new file mode 100644
index 000000000000..1ec310e22f77
Binary files /dev/null and
b/sql/connect/common/src/test/resources/query-tests/queries/function_unwrap_udt.proto.bin
differ
diff --git
a/sql/connect/common/src/test/resources/query-tests/queries/function_wrap_udt.json
b/sql/connect/common/src/test/resources/query-tests/queries/function_wrap_udt.json
new file mode 100644
index 000000000000..7376698e8516
--- /dev/null
+++
b/sql/connect/common/src/test/resources/query-tests/queries/function_wrap_udt.json
@@ -0,0 +1,104 @@
+{
+ "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": "wrap_udt",
+ "arguments": [{
+ "unresolvedFunction": {
+ "functionName": "array",
+ "arguments": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "b"
+ },
+ "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": "array",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass":
"org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "literal": {
+ "string":
"{\"type\":\"udt\",\"class\":\"org.apache.spark.sql.TestUDT$NewArrayUDT\",\"pyClass\":null,\"sqlType\":{\"type\":\"array\",\"elementType\":\"double\",\"containsNull\":false}}"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "wrap_udt",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass":
"org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }],
+ "isInternal": true
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "wrap_udt",
+ "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_wrap_udt.proto.bin
b/sql/connect/common/src/test/resources/query-tests/queries/function_wrap_udt.proto.bin
new file mode 100644
index 000000000000..dd9898055ff5
Binary files /dev/null and
b/sql/connect/common/src/test/resources/query-tests/queries/function_wrap_udt.proto.bin
differ
diff --git
a/sql/connect/common/src/test/scala/org/apache/spark/sql/TestUDT.scala
b/sql/connect/common/src/test/scala/org/apache/spark/sql/TestUDT.scala
new file mode 100644
index 000000000000..77704b8a57c3
--- /dev/null
+++ b/sql/connect/common/src/test/scala/org/apache/spark/sql/TestUDT.scala
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql
+
+import org.apache.spark.sql.types.{ArrayType, DataType, DoubleType,
SQLUserDefinedType, UserDefinedType}
+
+object TestUDT {
+
+ @SQLUserDefinedType(udt = classOf[NewArrayUDT])
+ private[sql] class NewArray(val values: Array[Double]) extends Serializable
+
+ private[sql] class NewArrayUDT extends UserDefinedType[NewArray] {
+
+ override def sqlType: DataType = ArrayType(DoubleType, containsNull =
false)
+
+ override def serialize(obj: NewArray): Any = obj.values
+
+ override def deserialize(datum: Any): NewArray = {
+ datum match {
+ case values: Array[_] =>
+ new NewArray(values.map(_.asInstanceOf[Double]))
+ }
+ }
+
+ override def userClass: Class[NewArray] = classOf[NewArray]
+ }
+}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
index f634774faec6..142b4e6dd26e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
@@ -84,6 +84,7 @@ class DataFrameFunctionsSuite extends SharedSparkSession {
"bucket", "days", "hours", "months", "years", // Datasource v2 partition
transformations
"product", // Discussed in https://github.com/apache/spark/pull/30745
"unwrap_udt",
+ "wrap_udt",
"timestamp_add",
"timestamp_diff"
)
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala
index 18adb45c1234..04be47523c58 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala
@@ -293,8 +293,33 @@ class UserDefinedTypeSuite extends SharedSparkSession with
ParquetTest
val unwrappedFeaturesArrays: Array[Array[Double]] =
unwrappedFeatures.collect()
assert(unwrappedFeaturesArrays.length === 2)
- java.util.Arrays.equals(unwrappedFeaturesArrays(0), Array(0.1, 1.0))
- java.util.Arrays.equals(unwrappedFeaturesArrays(1), Array(0.2, 2.0))
+ assert(Arrays.equals(unwrappedFeaturesArrays(0), Array(0.1, 1.0)))
+ assert(Arrays.equals(unwrappedFeaturesArrays(1), Array(0.2, 2.0)))
+ }
+
+ test("Test wrap_udt function") {
+ val udt = new TestUDT.MyDenseVectorUDT()
+ val df = Seq(Array(0.1, 1.0), Array(0.2, 2.0)).toDF("features")
+ val wrappedFeatures = df.select(wrap_udt(col("features"),
udt).as("features"))
+
+ assert(wrappedFeatures.schema("features").dataType === udt)
+ checkAnswer(
+ wrappedFeatures,
+ Seq(
+ Row(new TestUDT.MyDenseVector(Array(0.1, 1.0))),
+ Row(new TestUDT.MyDenseVector(Array(0.2, 2.0)))))
+ }
+
+ test("Test unwrap_udt and wrap_udt round trip") {
+ val udt = new TestUDT.MyDenseVectorUDT()
+ val roundTrip = pointsRDD.select(wrap_udt(unwrap_udt(col("features")),
udt).as("features"))
+
+ assert(roundTrip.schema("features").dataType === udt)
+ checkAnswer(
+ roundTrip,
+ Seq(
+ Row(new TestUDT.MyDenseVector(Array(0.1, 1.0))),
+ Row(new TestUDT.MyDenseVector(Array(0.2, 2.0)))))
}
test("SPARK-46289: UDT ordering") {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
