(spark) branch master updated: [SPARK-58962][MLLIB] Use zero-copy UDT conversion in MLUtils

[ru...@apache.org]

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 <ru...@apache.org>
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 <ru...@apache.org>
    Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
 .../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: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org