[PR] [SPARK-51282][ML][CONNECT] Optimize OneVsRestModel transform by eliminating the JVM-Python data exchange [spark]

[via GitHub]

zhengruifeng opened a new pull request, #50041:
URL: https://github.com/apache/spark/pull/50041

   ### What changes were proposed in this pull request?
   Optimize OneVsRestModel transform by eliminating the JVM-Python data exchange
   
   
   ### Why are the changes needed?
   for better performance, using builtin functions to replace the python UDFs, 
then optimize out the python execution
   
   my local test show it is **4.6** X faster than existing implementation
   
   
   ### Does this PR introduce _any_ user-facing change?
   no
   
   
   ### How was this patch tested?
   ci, and manually check
   
   prepare model and test data
   ```
   from pyspark.sql import Row
   from pyspark.ml.classification import *
   from pyspark.ml.linalg import Vectors
   
   data_path = "data/mllib/sample_multiclass_classification_data.txt"
   
   df = spark.read.format("libsvm").load(data_path)
   lr = LogisticRegression(regParam=0.01)
   ovr = OneVsRest(classifier=lr)
   model = ovr.fit(df)
   model.save("/tmp/ovrm")
   
   
   for i in range(10):
       df = df.union(df)
   
   df.write.mode("overwrite").parquet("/tmp/test_data")
   ```
   
   benchmark transform performance
   ```
   import time
   from pyspark.sql import Row
   from pyspark.ml.classification import *
   from pyspark.ml.linalg import Vectors
   
   df = spark.read.parquet("/tmp/test_data")
   
   df.persist()
   df.count()
   df.count()
   
   ovrm = OneVsRestModel.load("/tmp/ovrm")
   
   tic = time.time()
   ovrm.transform(df).write.mode("overwrite").parquet(f"/tmp/pred-{tic}")
   toc = time.time()
   
   toc - tic
   ```
   
   master: 8.527473211288452 sec
   
![image](https://github.com/user-attachments/assets/d270b4c9-3470-4288-923a-c5f0738d1053)
   
   
   this PR: 1.8552227020263672 sec
   
![image](https://github.com/user-attachments/assets/b3953a82-98fd-4483-b0a3-6546adbfde7c)
   
   
   
   ### Was this patch authored or co-authored using generative AI tooling?
   no
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org