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 82c3afaaaccf [SPARK-58962][MLLIB] Use zero-copy UDT conversion in
MLUtils
82c3afaaaccf is described below
commit 82c3afaaaccff3dbac15bc23928ca89edd8aa647
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Mon Aug 24 17:07:54 2026 +0800
[SPARK-58962][MLLIB] Use zero-copy UDT conversion in MLUtils
### What changes were proposed in this pull request?
This PR updates `MLUtils` vector and matrix column conversion helpers to use
`unwrap_udt` + `wrap_udt` instead of row-by-row UDF conversion.
It applies to:
- `convertVectorColumnsToML`
- `convertVectorColumnsFromML`
- `convertMatrixColumnsToML`
- `convertMatrixColumnsFromML`
The existing column selection, type validation, and metadata preservation
behavior is unchanged.
### Why are the changes needed?
The old and new MLlib vector/matrix UDTs share the same SQL backing schema.
After SPARK-58875,
Spark can re-wrap the same underlying SQL value with the target UDT
directly, avoiding the
serialization overhead from converting each value through a Scala UDF.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Added nullable vector and matrix column conversion coverage in
`MLUtilsSuite`.
Ran:
```
build/sbt 'mllib/testOnly org.apache.spark.mllib.util.MLUtilsSuite -- -z
"convertVectorColumnsToML" -z "convertVectorColumnsFromML" -z
"convertMatrixColumnsToML" -z "convertMatrixColumnsFromML" -z "convert nullable
vector columns" -z "convert nullable matrix columns"'
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #58240 from zhengruifeng/SPARK-58962-zero-copy-udt-mlutils-dev-3.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../org/apache/spark/mllib/util/MLUtils.scala | 31 +++---------
.../org/apache/spark/mllib/util/MLUtilsSuite.scala | 55 +++++++++++++++++++++-
2 files changed, 59 insertions(+), 27 deletions(-)
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/util/MLUtils.scala
b/mllib/src/main/scala/org/apache/spark/mllib/util/MLUtils.scala
index 5a213e680348..bbfb529bdae6 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/util/MLUtils.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/util/MLUtils.scala
@@ -33,7 +33,8 @@ import org.apache.spark.rdd.{PartitionwiseSampledRDD, RDD}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}
import org.apache.spark.sql.execution.datasources.DataSource
import org.apache.spark.sql.execution.datasources.text.TextFileFormat
-import org.apache.spark.sql.functions._
+import org.apache.spark.sql.functions.{col, length, lit, not, printf,
raise_error, trim,
+ unwrap_udt, when, wrap_udt}
import org.apache.spark.storage.StorageLevel
import org.apache.spark.util.random.BernoulliCellSampler
@@ -350,16 +351,10 @@ object MLUtils extends Logging {
return dataset.toDF()
}
- logWarning("Vector column conversion has serialization overhead. " +
- "Please migrate your datasets and workflows to use the spark.ml
package.")
-
- // TODO: This implementation has performance issues due to unnecessary
serialization.
- // TODO: It is better (but trickier) if we can cast the old vector type to
new type directly.
- val convertToML = udf { v: Vector => v.asML }
val exprs = schema.fields.map { field =>
val c = field.name
if (colSet.contains(c)) {
- convertToML(col(c)).as(c, field.metadata)
+ wrap_udt(unwrap_udt(col(c)), new MLVectorUDT).as(c, field.metadata)
} else {
col(c)
}
@@ -403,16 +398,10 @@ object MLUtils extends Logging {
return dataset.toDF()
}
- logWarning("Vector column conversion has serialization overhead. " +
- "Please migrate your datasets and workflows to use the spark.ml
package.")
-
- // TODO: This implementation has performance issues due to unnecessary
serialization.
- // TODO: It is better (but trickier) if we can cast the new vector type to
old type directly.
- val convertFromML = udf { Vectors.fromML _ }
val exprs = schema.fields.map { field =>
val c = field.name
if (colSet.contains(c)) {
- convertFromML(col(c)).as(c, field.metadata)
+ wrap_udt(unwrap_udt(col(c)), new VectorUDT).as(c, field.metadata)
} else {
col(c)
}
@@ -456,14 +445,10 @@ object MLUtils extends Logging {
return dataset.toDF()
}
- logWarning("Matrix column conversion has serialization overhead. " +
- "Please migrate your datasets and workflows to use the spark.ml
package.")
-
- val convertToML = udf { v: Matrix => v.asML }
val exprs = schema.fields.map { field =>
val c = field.name
if (colSet.contains(c)) {
- convertToML(col(c)).as(c, field.metadata)
+ wrap_udt(unwrap_udt(col(c)), new MLMatrixUDT).as(c, field.metadata)
} else {
col(c)
}
@@ -507,14 +492,10 @@ object MLUtils extends Logging {
return dataset.toDF()
}
- logWarning("Matrix column conversion has serialization overhead. " +
- "Please migrate your datasets and workflows to use the spark.ml
package.")
-
- val convertFromML = udf { Matrices.fromML _ }
val exprs = schema.fields.map { field =>
val c = field.name
if (colSet.contains(c)) {
- convertFromML(col(c)).as(c, field.metadata)
+ wrap_udt(unwrap_udt(col(c)), new MatrixUDT).as(c, field.metadata)
} else {
col(c)
}
diff --git
a/mllib/src/test/scala/org/apache/spark/mllib/util/MLUtilsSuite.scala
b/mllib/src/test/scala/org/apache/spark/mllib/util/MLUtilsSuite.scala
index 5f6ac1b716cc..376a51b87272 100644
--- a/mllib/src/test/scala/org/apache/spark/mllib/util/MLUtilsSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/mllib/util/MLUtilsSuite.scala
@@ -23,13 +23,15 @@ import java.nio.file.Files
import scala.io.Source
import org.apache.spark.{SparkException, SparkFunSuite, SparkRuntimeException}
-import org.apache.spark.mllib.linalg.{DenseVector, Matrices, SparseVector,
Vector, Vectors}
+import org.apache.spark.ml.linalg.{MatrixUDT => MLMatrixUDT, VectorUDT =>
MLVectorUDT}
+import org.apache.spark.mllib.linalg.{DenseVector, Matrices, MatrixUDT =>
OldMatrixUDT,
+ SparseVector, Vector, Vectors, VectorUDT => OldVectorUDT}
import org.apache.spark.mllib.regression.LabeledPoint
import org.apache.spark.mllib.util.MLUtils._
import org.apache.spark.mllib.util.TestingUtils._
import org.apache.spark.sql.Row
import org.apache.spark.sql.functions.col
-import org.apache.spark.sql.types.MetadataBuilder
+import org.apache.spark.sql.types.{IntegerType, MetadataBuilder, StructField,
StructType}
import org.apache.spark.util.Utils
class MLUtilsSuite extends SparkFunSuite with MLlibTestSparkContext {
@@ -307,6 +309,30 @@ class MLUtilsSuite extends SparkFunSuite with
MLlibTestSparkContext {
}
}
+ test("convert nullable vector columns") {
+ val oldVector = Vectors.sparse(2, Array(1), Array(1.0))
+ val oldSchema = StructType(Seq(
+ StructField("id", IntegerType, nullable = false),
+ StructField("x", new OldVectorUDT, nullable = true)))
+ val oldDF = spark.createDataFrame(
+ sc.parallelize(Seq(Row(0, oldVector), Row(1, null))),
+ oldSchema)
+ val toML = convertVectorColumnsToML(oldDF)
+ assert(toML.schema("x").dataType === new MLVectorUDT)
+ assert(toML.collect().toSeq === Seq(Row(0, oldVector.asML), Row(1, null)))
+
+ val mlVector = oldVector.asML
+ val mlSchema = StructType(Seq(
+ StructField("id", IntegerType, nullable = false),
+ StructField("x", new MLVectorUDT, nullable = true)))
+ val mlDF = spark.createDataFrame(
+ sc.parallelize(Seq(Row(0, mlVector), Row(1, null))),
+ mlSchema)
+ val fromML = convertVectorColumnsFromML(mlDF)
+ assert(fromML.schema("x").dataType === new OldVectorUDT)
+ assert(fromML.collect().toSeq === Seq(Row(0, oldVector), Row(1, null)))
+ }
+
test("convertMatrixColumnsToML") {
val x = Matrices.sparse(3, 2, Array(0, 2, 3), Array(0, 2, 1), Array(0.0,
-1.2, 0.0))
val metadata = new MetadataBuilder().putLong("numFeatures", 2L).build()
@@ -357,6 +383,31 @@ class MLUtilsSuite extends SparkFunSuite with
MLlibTestSparkContext {
}
}
+ test("convert nullable matrix columns") {
+ val oldMatrix = Matrices.sparse(
+ 3, 2, Array(0, 2, 3), Array(0, 2, 1), Array(0.0, -1.2, 0.0))
+ val oldSchema = StructType(Seq(
+ StructField("id", IntegerType, nullable = false),
+ StructField("x", new OldMatrixUDT, nullable = true)))
+ val oldDF = spark.createDataFrame(
+ sc.parallelize(Seq(Row(0, oldMatrix), Row(1, null))),
+ oldSchema)
+ val toML = convertMatrixColumnsToML(oldDF)
+ assert(toML.schema("x").dataType === new MLMatrixUDT)
+ assert(toML.collect().toSeq === Seq(Row(0, oldMatrix.asML), Row(1, null)))
+
+ val mlMatrix = oldMatrix.asML
+ val mlSchema = StructType(Seq(
+ StructField("id", IntegerType, nullable = false),
+ StructField("x", new MLMatrixUDT, nullable = true)))
+ val mlDF = spark.createDataFrame(
+ sc.parallelize(Seq(Row(0, mlMatrix), Row(1, null))),
+ mlSchema)
+ val fromML = convertMatrixColumnsFromML(mlDF)
+ assert(fromML.schema("x").dataType === new OldMatrixUDT)
+ assert(fromML.collect().toSeq === Seq(Row(0, oldMatrix), Row(1, null)))
+ }
+
test("kFold with fold column") {
val data = sc.parallelize(1 to 100, 2).map(x => (x, if (x <= 50) 0 else
1)).toDF("i", "fold")
val collectedData = data.collect().map(_.getInt(0)).sorted
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
