You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2023/06/06 13:54:35 UTC
[spark] branch master updated: [SPARK-43097][ML] New pyspark ML logistic regression estimator implemented on top of distributor
This is an automated email from the ASF dual-hosted git repository.
weichenxu123 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 2a82b42bcb1 [SPARK-43097][ML] New pyspark ML logistic regression estimator implemented on top of distributor
2a82b42bcb1 is described below
commit 2a82b42bcb100fade622d3d08ef5d316425d3e5a
Author: Weichen Xu <we...@databricks.com>
AuthorDate: Tue Jun 6 21:53:46 2023 +0800
[SPARK-43097][ML] New pyspark ML logistic regression estimator implemented on top of distributor
### What changes were proposed in this pull request?
This PR takes over https://github.com/apache/spark/pull/40748
Closes https://github.com/apache/spark/pull/40748
### Why are the changes needed?
Distributed ML on spark connect project.
### Does this PR introduce _any_ user-facing change?
Yes.
### How was this patch tested?
Unit tests.
Manually testing code:
Start local pyspark shell by:
```
bin/pyspark --remote "local[2]" # spark connect mode
```
or
```
bin/pyspark --master "local[2]" # legacy mode
```
to launch pyspark shell,
Then, paste following code to pyspark shell:
```
from pyspark.mlv2.classification import LogisticRegression as LORV2
lorv2 = LORV2(maxIter=2, numTrainWorkers=2)
from pyspark.ml.linalg import Vectors
df = spark.createDataFrame(
[
(1.0, Vectors.dense(0.0, 5.0)),
(0.0, Vectors.dense(1.0, 2.0)),
(1.0, Vectors.dense(2.0, 1.0)),
(0.0, Vectors.dense(3.0, 3.0)),
] * 100,
["label", "features"],
)
model.transform(df).show(truncate=False)
model.transform(df.toPandas())
model.set(model.probabilityCol, "")
model.transform(df).show(truncate=False)
model.transform(df.toPandas())
```
Closes #41383 from WeichenXu123/lor-torch-2.
Authored-by: Weichen Xu <we...@databricks.com>
Signed-off-by: Weichen Xu <we...@databricks.com>
---
dev/sparktestsupport/modules.py | 2 +
python/pyspark/ml/param/_shared_params_code_gen.py | 24 ++
python/pyspark/ml/param/shared.py | 89 ++++++
python/pyspark/ml/torch/data.py | 13 +-
python/pyspark/ml/torch/distributor.py | 13 +-
python/pyspark/mlv2/base.py | 75 ++++-
python/pyspark/mlv2/classification.py | 306 +++++++++++++++++++++
.../tests/connect/test_parity_classification.py | 41 +++
python/pyspark/mlv2/tests/test_classification.py | 139 ++++++++++
9 files changed, 696 insertions(+), 6 deletions(-)
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index aa755033aa4..ecc471fd700 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -609,6 +609,7 @@ pyspark_ml = Module(
"pyspark.mlv2.tests.test_summarizer",
"pyspark.mlv2.tests.test_evaluation",
"pyspark.mlv2.tests.test_feature",
+ "pyspark.mlv2.tests.test_classification",
],
excluded_python_implementations=[
"PyPy" # Skip these tests under PyPy since they require numpy and it isn't available there
@@ -823,6 +824,7 @@ pyspark_connect = Module(
"pyspark.mlv2.tests.connect.test_parity_summarizer",
"pyspark.mlv2.tests.connect.test_parity_evaluation",
"pyspark.mlv2.tests.connect.test_parity_feature",
+ "pyspark.mlv2.tests.connect.test_parity_classification",
],
excluded_python_implementations=[
"PyPy" # Skip these tests under PyPy since they require numpy, pandas, and pyarrow and
diff --git a/python/pyspark/ml/param/_shared_params_code_gen.py b/python/pyspark/ml/param/_shared_params_code_gen.py
index 5df1782084a..2bec3a5053f 100644
--- a/python/pyspark/ml/param/_shared_params_code_gen.py
+++ b/python/pyspark/ml/param/_shared_params_code_gen.py
@@ -332,6 +332,30 @@ if __name__ == "__main__":
"0.0",
"TypeConverters.toFloat",
),
+ (
+ "numTrainWorkers",
+ "number of training workers",
+ "1",
+ "TypeConverters.toInt",
+ ),
+ (
+ "batchSize",
+ "number of training batch size",
+ None,
+ "TypeConverters.toInt",
+ ),
+ (
+ "learningRate",
+ "learning rate for training",
+ None,
+ "TypeConverters.toFloat",
+ ),
+ (
+ "momentum",
+ "momentum for training optimizer",
+ None,
+ "TypeConverters.toFloat",
+ ),
]
code = []
diff --git a/python/pyspark/ml/param/shared.py b/python/pyspark/ml/param/shared.py
index fcfced2e566..d61d206d219 100644
--- a/python/pyspark/ml/param/shared.py
+++ b/python/pyspark/ml/param/shared.py
@@ -787,3 +787,92 @@ class HasMaxBlockSizeInMB(Params):
Gets the value of maxBlockSizeInMB or its default value.
"""
return self.getOrDefault(self.maxBlockSizeInMB)
+
+
+class HasNumTrainWorkers(Params):
+ """
+ Mixin for param numTrainWorkers: number of training workers
+ """
+
+ numTrainWorkers: "Param[int]" = Param(
+ Params._dummy(),
+ "numTrainWorkers",
+ "number of training workers",
+ typeConverter=TypeConverters.toInt,
+ )
+
+ def __init__(self) -> None:
+ super(HasNumTrainWorkers, self).__init__()
+ self._setDefault(numTrainWorkers=1)
+
+ def getNumTrainWorkers(self) -> int:
+ """
+ Gets the value of numTrainWorkers or its default value.
+ """
+ return self.getOrDefault(self.numTrainWorkers)
+
+
+class HasBatchSize(Params):
+ """
+ Mixin for param batchSize: number of training batch size
+ """
+
+ batchSize: "Param[int]" = Param(
+ Params._dummy(),
+ "batchSize",
+ "number of training batch size",
+ typeConverter=TypeConverters.toInt,
+ )
+
+ def __init__(self) -> None:
+ super(HasBatchSize, self).__init__()
+
+ def getBatchSize(self) -> int:
+ """
+ Gets the value of batchSize or its default value.
+ """
+ return self.getOrDefault(self.batchSize)
+
+
+class HasLearningRate(Params):
+ """
+ Mixin for param learningRate: learning rate for training
+ """
+
+ learningRate: "Param[float]" = Param(
+ Params._dummy(),
+ "learningRate",
+ "learning rate for training",
+ typeConverter=TypeConverters.toFloat,
+ )
+
+ def __init__(self) -> None:
+ super(HasLearningRate, self).__init__()
+
+ def getLearningRate(self) -> float:
+ """
+ Gets the value of learningRate or its default value.
+ """
+ return self.getOrDefault(self.learningRate)
+
+
+class HasMomentum(Params):
+ """
+ Mixin for param momentum: momentum for training optimizer
+ """
+
+ momentum: "Param[float]" = Param(
+ Params._dummy(),
+ "momentum",
+ "momentum for training optimizer",
+ typeConverter=TypeConverters.toFloat,
+ )
+
+ def __init__(self) -> None:
+ super(HasMomentum, self).__init__()
+
+ def getMomentum(self) -> float:
+ """
+ Gets the value of momentum or its default value.
+ """
+ return self.getOrDefault(self.momentum)
diff --git a/python/pyspark/ml/torch/data.py b/python/pyspark/ml/torch/data.py
index d421683c16d..0a5597fbd24 100644
--- a/python/pyspark/ml/torch/data.py
+++ b/python/pyspark/ml/torch/data.py
@@ -46,7 +46,18 @@ class _SparkPartitionTorchDataset(torch.utils.data.IterableDataset):
sparse_array[value["indices"]] = value["values"]
return sparse_array
- elif field_type in ["float", "double", "int", "bigint", "smallint"]:
+ elif field_type in [
+ "float",
+ "double",
+ "int",
+ "bigint",
+ "smallint",
+ "array<float>",
+ "array<double>",
+ "array<int>",
+ "array<bigint>",
+ "array<smallint>",
+ ]:
def converter(value: Any) -> Any:
return value
diff --git a/python/pyspark/ml/torch/distributor.py b/python/pyspark/ml/torch/distributor.py
index 2af92f8399f..be49dc147c0 100644
--- a/python/pyspark/ml/torch/distributor.py
+++ b/python/pyspark/ml/torch/distributor.py
@@ -968,7 +968,9 @@ class TorchDistributor(Distributor):
)
-def _get_spark_partition_data_loader(num_samples: int, batch_size: int, prefetch: int = 2) -> Any:
+def _get_spark_partition_data_loader(
+ num_samples: int, batch_size: int, num_workers: int = 1, prefetch_factor: Optional[int] = 2
+) -> Any:
"""
This function must be called inside the `train_function` where `train_function`
is the input argument of `TorchDistributor.train_on_dataframe`.
@@ -985,8 +987,11 @@ def _get_spark_partition_data_loader(num_samples: int, batch_size: int, prefetch
it wraps round back to the first row.
batch_size:
How many samples per batch to load.
- prefetch:
- Number of batches loaded in advance.
+ num_workers:
+ How many subprocesses to use for data loading.
+ 0 means that the data will be loaded in the main process.
+ prefetch_factor:
+ Number of batches loaded in advance by each worker
"""
from pyspark.sql.types import StructType
from pyspark.ml.torch.data import _SparkPartitionTorchDataset
@@ -1000,4 +1005,4 @@ def _get_spark_partition_data_loader(num_samples: int, batch_size: int, prefetch
dataset = _SparkPartitionTorchDataset(arrow_file, schema, num_samples)
- return DataLoader(dataset, batch_size, num_workers=1, prefetch_factor=prefetch)
+ return DataLoader(dataset, batch_size, num_workers=num_workers, prefetch_factor=prefetch_factor)
diff --git a/python/pyspark/mlv2/base.py b/python/pyspark/mlv2/base.py
index 70547d01966..6088866b8ea 100644
--- a/python/pyspark/mlv2/base.py
+++ b/python/pyspark/mlv2/base.py
@@ -35,7 +35,11 @@ from pyspark import since
from pyspark.ml.common import inherit_doc
from pyspark.sql.dataframe import DataFrame
from pyspark.ml.param import Params
-
+from pyspark.ml.param.shared import (
+ HasLabelCol,
+ HasFeaturesCol,
+ HasPredictionCol,
+)
from pyspark.mlv2.util import transform_dataframe_column
if TYPE_CHECKING:
@@ -257,3 +261,72 @@ class Model(Transformer, metaclass=ABCMeta):
"""
pass
+
+
+@inherit_doc
+class _PredictorParams(HasLabelCol, HasFeaturesCol, HasPredictionCol):
+ """
+ Params for :py:class:`Predictor` and :py:class:`PredictorModel`.
+
+ .. versionadded:: 3.5.0
+ """
+
+ pass
+
+
+@inherit_doc
+class Predictor(Estimator[M], _PredictorParams, metaclass=ABCMeta):
+ """
+ Estimator for prediction tasks (regression and classification).
+ """
+
+ @since("3.5.0")
+ def setLabelCol(self, value: str) -> "Predictor":
+ """
+ Sets the value of :py:attr:`labelCol`.
+ """
+ return self._set(labelCol=value)
+
+ @since("3.5.0")
+ def setFeaturesCol(self, value: str) -> "Predictor":
+ """
+ Sets the value of :py:attr:`featuresCol`.
+ """
+ return self._set(featuresCol=value)
+
+ @since("3.5.0")
+ def setPredictionCol(self, value: str) -> "Predictor":
+ """
+ Sets the value of :py:attr:`predictionCol`.
+ """
+ return self._set(predictionCol=value)
+
+
+@inherit_doc
+class PredictionModel(Model, _PredictorParams, metaclass=ABCMeta):
+ """
+ Model for prediction tasks (regression and classification).
+ """
+
+ @since("3.5.0")
+ def setFeaturesCol(self, value: str) -> "PredictionModel":
+ """
+ Sets the value of :py:attr:`featuresCol`.
+ """
+ return self._set(featuresCol=value)
+
+ @since("3.5.0")
+ def setPredictionCol(self, value: str) -> "PredictionModel":
+ """
+ Sets the value of :py:attr:`predictionCol`.
+ """
+ return self._set(predictionCol=value)
+
+ @property # type: ignore[misc]
+ @abstractmethod
+ @since("3.5.0")
+ def numFeatures(self) -> int:
+ """
+ Returns the number of features the model was trained on. If unknown, returns -1
+ """
+ raise NotImplementedError()
diff --git a/python/pyspark/mlv2/classification.py b/python/pyspark/mlv2/classification.py
new file mode 100644
index 00000000000..a72fe89c01b
--- /dev/null
+++ b/python/pyspark/mlv2/classification.py
@@ -0,0 +1,306 @@
+#
+# 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.
+#
+
+from pyspark.mlv2.base import _PredictorParams
+
+from pyspark.ml.param.shared import HasProbabilityCol
+
+from typing import Any, Union, List, Tuple, Callable
+import numpy as np
+import pandas as pd
+import math
+
+from pyspark.sql import DataFrame
+from pyspark.ml.common import inherit_doc
+from pyspark.ml.torch.distributor import TorchDistributor
+from pyspark.ml.param.shared import (
+ HasMaxIter,
+ HasFitIntercept,
+ HasTol,
+ HasWeightCol,
+ HasSeed,
+ HasNumTrainWorkers,
+ HasBatchSize,
+ HasLearningRate,
+ HasMomentum,
+)
+from pyspark.mlv2.base import Predictor, PredictionModel
+from pyspark.sql.functions import lit, count, countDistinct
+
+import torch
+import torch.nn as torch_nn
+
+
+class _LogisticRegressionParams(
+ _PredictorParams,
+ HasMaxIter,
+ HasFitIntercept,
+ HasTol,
+ HasWeightCol,
+ HasNumTrainWorkers,
+ HasBatchSize,
+ HasLearningRate,
+ HasMomentum,
+ HasProbabilityCol,
+ HasSeed,
+):
+ """
+ Params for :py:class:`LogisticRegression` and :py:class:`LogisticRegressionModel`.
+
+ .. versionadded:: 3.0.0
+ """
+
+ pass
+
+
+class _LinearNet(torch_nn.Module):
+ def __init__(self, num_features: int, num_classes: int, bias: bool) -> None:
+ super(_LinearNet, self).__init__()
+ output_dim = num_classes
+ self.fc = torch_nn.Linear(num_features, output_dim, bias=bias, dtype=torch.float32)
+
+ def forward(self, x: Any) -> Any:
+ output = self.fc(x)
+ return output
+
+
+def _train_logistic_regression_model_worker_fn(
+ num_samples_per_worker: int,
+ num_features: int,
+ batch_size: int,
+ max_iter: int,
+ num_classes: int,
+ learning_rate: float,
+ momentum: float,
+ fit_intercept: bool,
+ seed: int,
+) -> Any:
+ from pyspark.ml.torch.distributor import _get_spark_partition_data_loader
+ from torch.nn.parallel import DistributedDataParallel as DDP
+ import torch.distributed
+ import torch.optim as optim
+
+ # TODO: add a setting seed param.
+ torch.manual_seed(seed)
+
+ # TODO: support training on GPU
+ # TODO: support L1 / L2 regularization
+ torch.distributed.init_process_group("gloo")
+
+ ddp_model = DDP(
+ _LinearNet(num_features=num_features, num_classes=num_classes, bias=fit_intercept)
+ )
+
+ loss_fn = torch_nn.CrossEntropyLoss()
+
+ optimizer = optim.SGD(ddp_model.parameters(), lr=learning_rate, momentum=momentum)
+ data_loader = _get_spark_partition_data_loader(
+ num_samples_per_worker,
+ batch_size,
+ num_workers=0,
+ prefetch_factor=None,
+ )
+ for i in range(max_iter):
+ ddp_model.train()
+
+ step_count = 0
+
+ loss_sum = 0.0
+ for x, target in data_loader:
+ optimizer.zero_grad()
+ output = ddp_model(x.to(torch.float32))
+ loss = loss_fn(output, target.to(torch.long))
+ loss.backward()
+ loss_sum += loss.detach().numpy()
+ optimizer.step()
+ step_count += 1
+
+ # TODO: early stopping
+ # When each epoch ends, computes loss on validation dataset and compare
+ # current epoch validation loss with last epoch validation loss, if
+ # less than provided `tol`, stop training.
+
+ if torch.distributed.get_rank() == 0:
+ print(f"Progress: train epoch {i + 1} completes, train loss = {loss_sum / step_count}")
+
+ if torch.distributed.get_rank() == 0:
+ return ddp_model.module.state_dict()
+
+ return None
+
+
+@inherit_doc
+class LogisticRegression(Predictor["LogisticRegressionModel"], _LogisticRegressionParams):
+ """
+ Logistic regression estimator.
+
+ .. versionadded:: 3.5.0
+ """
+
+ def __init__(
+ self,
+ *,
+ featuresCol: str = "features",
+ labelCol: str = "label",
+ predictionCol: str = "prediction",
+ probabilityCol: str = "probability",
+ maxIter: int = 100,
+ tol: float = 1e-6,
+ numTrainWorkers: int = 1,
+ batchSize: int = 32,
+ learningRate: float = 0.001,
+ momentum: float = 0.9,
+ seed: int = 0,
+ ):
+ super(_LogisticRegressionParams, self).__init__()
+ self._set(
+ featuresCol=featuresCol,
+ labelCol=labelCol,
+ predictionCol=predictionCol,
+ probabilityCol=probabilityCol,
+ maxIter=maxIter,
+ tol=tol,
+ numTrainWorkers=numTrainWorkers,
+ batchSize=batchSize,
+ learningRate=learningRate,
+ momentum=momentum,
+ seed=seed,
+ )
+
+ def _fit(self, dataset: Union[DataFrame, pd.DataFrame]) -> "LogisticRegressionModel":
+ if isinstance(dataset, pd.DataFrame):
+ # TODO: support pandas dataframe fitting
+ raise NotImplementedError("Fitting pandas dataframe is not supported yet.")
+
+ num_train_workers = self.getNumTrainWorkers()
+ batch_size = self.getBatchSize()
+
+ # We don't need to persist the dataset because the shuffling result from the repartition
+ # has been cached.
+ dataset = dataset.select(self.getFeaturesCol(), self.getLabelCol()).repartition(
+ num_train_workers
+ )
+
+ # TODO: check label values are in range of [0, num_classes)
+ num_rows, num_classes = dataset.agg(
+ count(lit(1)), countDistinct(self.getLabelCol())
+ ).head() # type: ignore[misc]
+
+ num_batches_per_worker = math.ceil(num_rows / num_train_workers / batch_size)
+ num_samples_per_worker = num_batches_per_worker * batch_size
+
+ num_features = len(dataset.select(self.getFeaturesCol()).head()[0]) # type: ignore[index]
+
+ if num_classes < 2:
+ raise ValueError("Training dataset distinct labels must >= 2.")
+
+ # TODO: support GPU.
+ distributor = TorchDistributor(
+ local_mode=False, use_gpu=False, num_processes=num_train_workers
+ )
+ model_state_dict = distributor._train_on_dataframe(
+ _train_logistic_regression_model_worker_fn,
+ dataset,
+ num_samples_per_worker=num_samples_per_worker,
+ num_features=num_features,
+ batch_size=batch_size,
+ max_iter=self.getMaxIter(),
+ num_classes=num_classes,
+ learning_rate=self.getLearningRate(),
+ momentum=self.getMomentum(),
+ fit_intercept=self.getFitIntercept(),
+ seed=self.getSeed(),
+ )
+
+ dataset.unpersist()
+
+ torch_model = _LinearNet(
+ num_features=num_features, num_classes=num_classes, bias=self.getFitIntercept()
+ )
+ torch_model.load_state_dict(model_state_dict)
+
+ lor_model = LogisticRegressionModel(
+ torch_model, num_features=num_features, num_classes=num_classes
+ )
+ lor_model._resetUid(self.uid)
+ return self._copyValues(lor_model)
+
+
+@inherit_doc
+class LogisticRegressionModel(PredictionModel, _LogisticRegressionParams):
+ """
+ Model fitted by LogisticRegression.
+
+ .. versionadded:: 3.5.0
+ """
+
+ def __init__(self, torch_model: Any, num_features: int, num_classes: int):
+ super().__init__()
+ self.torch_model = torch_model
+ self.num_features = num_features
+ self.num_classes = num_classes
+
+ def numFeatures(self) -> int:
+ return self.num_features
+
+ def numClasses(self) -> int:
+ return self.num_classes
+
+ def _input_columns(self) -> List[str]:
+ return [self.getOrDefault(self.featuresCol)]
+
+ def _output_columns(self) -> List[Tuple[str, str]]:
+ output_cols = [(self.getOrDefault(self.predictionCol), "bigint")]
+ prob_col = self.getOrDefault(self.probabilityCol)
+ if prob_col:
+ output_cols += [(prob_col, "array<double>")]
+ return output_cols
+
+ def _get_transform_fn(self) -> Callable[["pd.Series"], Any]:
+ model_state_dict = self.torch_model.state_dict()
+ num_features = self.num_features
+ num_classes = self.num_classes
+ fit_intercept = self.getFitIntercept()
+
+ def transform_fn(input_series: Any) -> Any:
+ torch_model = _LinearNet(
+ num_features=num_features, num_classes=num_classes, bias=fit_intercept
+ )
+ # TODO: Use spark broadast for `model_state_dict`,
+ # it can improve performance when model is large.
+ torch_model.load_state_dict(model_state_dict)
+
+ input_array = np.stack(input_series.values)
+
+ with torch.inference_mode():
+ result = torch_model(torch.tensor(input_array, dtype=torch.float32))
+ predictions = torch.argmax(result, dim=1).numpy()
+
+ if self.getProbabilityCol():
+ probabilities = torch.softmax(result, dim=1).numpy()
+
+ return pd.DataFrame(
+ {
+ self.getPredictionCol(): list(predictions),
+ self.getProbabilityCol(): list(probabilities),
+ },
+ index=input_series.index.copy(),
+ )
+ else:
+ return pd.Series(data=list(predictions), index=input_series.index.copy())
+
+ return transform_fn
diff --git a/python/pyspark/mlv2/tests/connect/test_parity_classification.py b/python/pyspark/mlv2/tests/connect/test_parity_classification.py
new file mode 100644
index 00000000000..16bc1af1a64
--- /dev/null
+++ b/python/pyspark/mlv2/tests/connect/test_parity_classification.py
@@ -0,0 +1,41 @@
+# -*- coding: utf-8 -*-
+#
+# 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.
+#
+
+import unittest
+from pyspark.sql import SparkSession
+from pyspark.mlv2.tests.test_classification import ClassificationTestsMixin
+
+
+class FeatureTestsOnConnect(ClassificationTestsMixin, unittest.TestCase):
+ def setUp(self) -> None:
+ self.spark = SparkSession.builder.remote("local[2]").getOrCreate()
+
+ def tearDown(self) -> None:
+ self.spark.stop()
+
+
+if __name__ == "__main__":
+ from pyspark.mlv2.tests.connect.test_parity_classification import * # noqa: F401,F403
+
+ try:
+ import xmlrunner # type: ignore[import]
+
+ testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+ except ImportError:
+ testRunner = None
+ unittest.main(testRunner=testRunner, verbosity=2)
diff --git a/python/pyspark/mlv2/tests/test_classification.py b/python/pyspark/mlv2/tests/test_classification.py
new file mode 100644
index 00000000000..b9c112ef094
--- /dev/null
+++ b/python/pyspark/mlv2/tests/test_classification.py
@@ -0,0 +1,139 @@
+# -*- coding: utf-8 -*-
+#
+# 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.
+#
+
+import unittest
+import numpy as np
+from pyspark.mlv2.classification import LogisticRegression as LORV2
+from pyspark.sql import SparkSession
+
+
+have_torch = True
+try:
+ import torch # noqa: F401
+except ImportError:
+ have_torch = False
+
+
+class ClassificationTestsMixin:
+ @staticmethod
+ def _check_result(result_dataframe, expected_predictions, expected_probabilities=None):
+ np.testing.assert_array_equal(list(result_dataframe.prediction), expected_predictions)
+ if "probability" in result_dataframe.columns:
+ np.testing.assert_allclose(
+ list(result_dataframe.probability),
+ expected_probabilities,
+ rtol=1e-2,
+ )
+
+ def test_binary_classes_logistic_regression(self):
+ df1 = self.spark.createDataFrame(
+ [
+ (1.0, [0.0, 5.0]),
+ (0.0, [1.0, 2.0]),
+ (1.0, [2.0, 1.0]),
+ (0.0, [3.0, 3.0]),
+ ]
+ * 100,
+ ["label", "features"],
+ )
+ eval_df1 = self.spark.createDataFrame(
+ [
+ ([0.0, 2.0],),
+ ([3.5, 3.0],),
+ ],
+ ["features"],
+ )
+
+ lorv2 = LORV2(maxIter=200, numTrainWorkers=2, learningRate=0.001)
+ assert lorv2.getMaxIter() == 200
+ assert lorv2.getNumTrainWorkers() == 2
+ assert lorv2.getOrDefault(lorv2.learningRate) == 0.001
+
+ model = lorv2.fit(df1)
+ assert model.uid == lorv2.uid
+
+ expected_predictions = [1, 0]
+ expected_probabilities = [
+ [0.217875, 0.782125],
+ [0.839615, 0.160385],
+ ]
+
+ result = model.transform(eval_df1).toPandas()
+ self._check_result(result, expected_predictions, expected_probabilities)
+ local_transform_result = model.transform(eval_df1.toPandas())
+ self._check_result(local_transform_result, expected_predictions, expected_probabilities)
+
+ model.set(model.probabilityCol, "")
+ result_without_prob = model.transform(eval_df1).toPandas()
+ assert "probability" not in result_without_prob.columns
+ self._check_result(result_without_prob, expected_predictions, None)
+
+ def test_multi_classes_logistic_regression(self):
+ df1 = self.spark.createDataFrame(
+ [
+ (1.0, [1.0, 5.0]),
+ (2.0, [1.0, -2.0]),
+ (0.0, [-2.0, 1.5]),
+ ]
+ * 100,
+ ["label", "features"],
+ )
+ eval_df1 = self.spark.createDataFrame(
+ [
+ ([1.5, 5.0],),
+ ([1.0, -2.5],),
+ ([-2.0, 1.0],),
+ ],
+ ["features"],
+ )
+
+ lorv2 = LORV2(maxIter=200, numTrainWorkers=2, learningRate=0.001)
+
+ model = lorv2.fit(df1)
+
+ expected_predictions = [1, 2, 0]
+ expected_probabilities = [
+ [5.526459e-03, 9.943553e-01, 1.183146e-04],
+ [4.629959e-03, 8.141352e-03, 9.872288e-01],
+ [9.624363e-01, 3.080821e-02, 6.755549e-03],
+ ]
+
+ result = model.transform(eval_df1).toPandas()
+ self._check_result(result, expected_predictions, expected_probabilities)
+ local_transform_result = model.transform(eval_df1.toPandas())
+ self._check_result(local_transform_result, expected_predictions, expected_probabilities)
+
+
+class ClassificationTests(ClassificationTestsMixin, unittest.TestCase):
+ def setUp(self) -> None:
+ self.spark = SparkSession.builder.master("local[2]").getOrCreate()
+
+ def tearDown(self) -> None:
+ self.spark.stop()
+
+
+if __name__ == "__main__":
+ from pyspark.mlv2.tests.test_classification import * # noqa: F401,F403
+
+ try:
+ import xmlrunner # type: ignore[import]
+
+ testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+ except ImportError:
+ testRunner = None
+ unittest.main(testRunner=testRunner, verbosity=2)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org