You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ru...@apache.org on 2025/02/26 12:41:48 UTC
(spark) branch master updated: [SPARK-51282][ML][PYTHON][CONNECT] Optimize OneVsRestModel transform by eliminating the JVM-Python data exchange
This is an automated email from the ASF dual-hosted git repository.
ruifengz 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 87fd092b19be [SPARK-51282][ML][PYTHON][CONNECT] Optimize OneVsRestModel transform by eliminating the JVM-Python data exchange
87fd092b19be is described below
commit 87fd092b19bef8ba853ebd53d64fc9a7e9d1d717
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Wed Feb 26 20:41:28 2025 +0800
[SPARK-51282][ML][PYTHON][CONNECT] Optimize OneVsRestModel transform by eliminating the JVM-Python data exchange
### 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 data exchange
my local test shows that 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
this PR: 1.8552227020263672 sec
### Was this patch authored or co-authored using generative AI tooling?
no
Closes #50041 from zhengruifeng/ml_connect_ovr_builtin.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@apache.org>
---
python/pyspark/ml/classification.py | 43 ++++++++-----------------------------
python/pyspark/sql/internal.py | 41 +++++++++++++++++++++++++++++++++++
2 files changed, 50 insertions(+), 34 deletions(-)
diff --git a/python/pyspark/ml/classification.py b/python/pyspark/ml/classification.py
index 4a2c213325c3..21e4ee4f6d0e 100644
--- a/python/pyspark/ml/classification.py
+++ b/python/pyspark/ml/classification.py
@@ -16,7 +16,6 @@
#
import os
-import operator
import sys
import uuid
import warnings
@@ -27,7 +26,6 @@ from typing import (
Any,
Dict,
Generic,
- Iterable,
List,
Optional,
Type,
@@ -41,7 +39,7 @@ from typing import (
)
from pyspark import keyword_only, since, inheritable_thread_target
-from pyspark.ml import Estimator, Predictor, PredictionModel, Model
+from pyspark.ml import Estimator, Predictor, PredictionModel, Model, functions as MF
from pyspark.ml.param.shared import (
HasRawPredictionCol,
HasProbabilityCol,
@@ -94,10 +92,9 @@ from pyspark.ml.util import (
)
from pyspark.ml.wrapper import JavaParams, JavaPredictor, JavaPredictionModel, JavaWrapper
from pyspark.ml.common import inherit_doc
-from pyspark.ml.linalg import Matrix, Vector, Vectors, VectorUDT
-from pyspark.sql import DataFrame, Row, SparkSession
-from pyspark.sql.functions import udf, when
-from pyspark.sql.types import ArrayType, DoubleType
+from pyspark.ml.linalg import Matrix, Vector
+from pyspark.sql import DataFrame, Row, SparkSession, functions as F
+from pyspark.sql.internal import InternalFunction as SF
from pyspark.storagelevel import StorageLevel
from pyspark.sql.utils import is_remote
@@ -3591,7 +3588,7 @@ class OneVsRest(
binaryLabelCol = "mc2b$" + str(index)
trainingDataset = multiclassLabeled.withColumn(
binaryLabelCol,
- when(multiclassLabeled[labelCol] == float(index), 1.0).otherwise(0.0),
+ F.when(multiclassLabeled[labelCol] == float(index), 1.0).otherwise(0.0),
)
paramMap = dict(
[
@@ -3835,14 +3832,12 @@ class OneVsRestModel(
)
def _transform(self, dataset: DataFrame) -> DataFrame:
- # TODO(SPARK-48515): Use Arrow Python UDF
# determine the input columns: these need to be passed through
origCols = dataset.columns
# add an accumulator column to store predictions of all the models
accColName = "mbc$acc" + str(uuid.uuid4())
- initUDF = udf(lambda _: [], ArrayType(DoubleType()), useArrow=False)
- newDataset = dataset.withColumn(accColName, initUDF(dataset[origCols[0]]))
+ newDataset = dataset.withColumn(accColName, F.array().cast("array<double>"))
# persist if underlying dataset is not persistent.
handlePersistence = dataset.storageLevel == StorageLevel(False, False, False, False)
@@ -3858,15 +3853,10 @@ class OneVsRestModel(
# add temporary column to store intermediate scores and update
tmpColName = "mbc$tmp" + str(uuid.uuid4())
- updateUDF = udf(
- lambda predictions, prediction: predictions + [prediction.tolist()[1]],
- ArrayType(DoubleType()),
- useArrow=False,
- )
transformedDataset = model.transform(aggregatedDataset).select(*columns)
updatedDataset = transformedDataset.withColumn(
tmpColName,
- updateUDF(transformedDataset[accColName], transformedDataset[rawPredictionCol]),
+ F.array_append(accColName, SF.vector_get(F.col(rawPredictionCol), F.lit(1))),
)
newColumns = origCols + [tmpColName]
@@ -3879,29 +3869,14 @@ class OneVsRestModel(
newDataset.unpersist()
if self.getRawPredictionCol():
-
- def func(predictions: Iterable[float]) -> Vector:
- predArray: List[float] = []
- for x in predictions:
- predArray.append(x)
- return Vectors.dense(predArray)
-
- rawPredictionUDF = udf(func, VectorUDT(), useArrow=False)
aggregatedDataset = aggregatedDataset.withColumn(
- self.getRawPredictionCol(), rawPredictionUDF(aggregatedDataset[accColName])
+ self.getRawPredictionCol(), MF.array_to_vector(F.col(accColName))
)
if self.getPredictionCol():
# output the index of the classifier with highest confidence as prediction
- labelUDF = udf(
- lambda predictions: float(
- max(enumerate(predictions), key=operator.itemgetter(1))[0]
- ),
- DoubleType(),
- useArrow=False,
- )
aggregatedDataset = aggregatedDataset.withColumn(
- self.getPredictionCol(), labelUDF(aggregatedDataset[accColName])
+ self.getPredictionCol(), SF.array_argmax(F.col(accColName)).cast("double")
)
return aggregatedDataset.drop(accColName)
diff --git a/python/pyspark/sql/internal.py b/python/pyspark/sql/internal.py
index eb4a2153cdaa..3007b28b0044 100644
--- a/python/pyspark/sql/internal.py
+++ b/python/pyspark/sql/internal.py
@@ -130,3 +130,44 @@ class InternalFunction:
"SECOND": "secs",
}
return F.make_interval(**{unit_mapping[unit]: F.lit(e)})
+
+ @staticmethod
+ def vector_get(vec: Column, idx: Column) -> Column:
+ unwrapped = F.unwrap_udt(vec)
+ is_dense = unwrapped.getField("type") == F.lit(1)
+ values = unwrapped.getField("values")
+ size = F.when(is_dense, F.array_size(values)).otherwise(unwrapped.getField("size"))
+ sparse_idx = InternalFunction.array_binary_search(unwrapped.getField("indices"), idx)
+ value = (
+ F.when(is_dense, F.get(values, idx))
+ .when(sparse_idx >= 0, F.get(values, sparse_idx))
+ .otherwise(F.lit(0.0))
+ )
+
+ return F.when((0 <= idx) & (idx < size), value).otherwise(
+ F.raise_error(F.printf(F.lit("Vector index must be in [0, %s), but got %s"), size, idx))
+ )
+
+ # The main different between this function and the `array_max` + `array_position`:
+ # This function ignores NaN/NULL values, while `array_max` treats NaN values as largest values.
+ @staticmethod
+ def array_argmax(arr: Column) -> Column:
+ def merge(acc: Column, vv: Column) -> Column:
+ v = acc.getField("v")
+ i = acc.getField("i")
+ j = acc.getField("j")
+ return F.when(
+ (~vv.isNaN()) & (~vv.isNull()) & (vv > v),
+ F.struct(vv.alias("v"), j.alias("i"), j + 1),
+ ).otherwise(F.struct(v.alias("v"), i.alias("i"), j + 1))
+
+ return F.aggregate(
+ arr,
+ F.struct(
+ F.lit(float("-inf")).alias("v"), # max value
+ F.lit(-1).alias("i"), # index of max value
+ F.lit(0).alias("j"), # current index
+ ),
+ merge,
+ lambda acc: acc.getField("i"),
+ )
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org