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/01/16 09:53:44 UTC
(spark) branch master updated: [SPARK-50825][ML][PYTHON][CONNECT] Support Tree Classifiers in ML Connect
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 d12bb23e4992 [SPARK-50825][ML][PYTHON][CONNECT] Support Tree Classifiers in ML Connect
d12bb23e4992 is described below
commit d12bb23e49928ae36b57ac5fbeaa492dadf7d28f
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Thu Jan 16 17:53:24 2025 +0800
[SPARK-50825][ML][PYTHON][CONNECT] Support Tree Classifiers in ML Connect
### What changes were proposed in this pull request?
Support Tree Classifiers in ML Connect:
- DecisionTreeClassifier
- RandomForestClassifier
- GBTClassifier
### Why are the changes needed?
for parity with spark clasic
### Does this PR introduce _any_ user-facing change?
yes, new algorithms supported on ml
### How was this patch tested?
added tests
### Was this patch authored or co-authored using generative AI tooling?
no
Closes #49525 from zhengruifeng/py_ml_dt.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@apache.org>
---
dev/sparktestsupport/modules.py | 2 +-
.../services/org.apache.spark.ml.Estimator | 5 +
...sification.py => test_parity_classification.py} | 4 +-
python/pyspark/ml/tests/test_classification.py | 331 +++++++++++++++++++++
.../org/apache/spark/sql/connect/ml/MLUtils.scala | 9 +
5 files changed, 348 insertions(+), 3 deletions(-)
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index f785a72e6a1f..f41767c75324 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -1116,7 +1116,7 @@ pyspark_ml_connect = Module(
"pyspark.ml.tests.connect.test_connect_classification",
"pyspark.ml.tests.connect.test_connect_pipeline",
"pyspark.ml.tests.connect.test_connect_tuning",
- "pyspark.ml.tests.connect.test_connect_spark_ml_classification",
+ "pyspark.ml.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/mllib/src/main/resources/META-INF/services/org.apache.spark.ml.Estimator b/mllib/src/main/resources/META-INF/services/org.apache.spark.ml.Estimator
index e6902f62c4d6..f010e48e8cab 100644
--- a/mllib/src/main/resources/META-INF/services/org.apache.spark.ml.Estimator
+++ b/mllib/src/main/resources/META-INF/services/org.apache.spark.ml.Estimator
@@ -17,4 +17,9 @@
# Spark Connect ML uses ServiceLoader to find out the supported Spark Ml estimators.
# So register the supported estimator here if you're trying to add a new one.
+
+# classification
org.apache.spark.ml.classification.LogisticRegression
+org.apache.spark.ml.classification.DecisionTreeClassifier
+org.apache.spark.ml.classification.RandomForestClassifier
+org.apache.spark.ml.classification.GBTClassifier
diff --git a/python/pyspark/ml/tests/connect/test_connect_spark_ml_classification.py b/python/pyspark/ml/tests/connect/test_parity_classification.py
similarity index 89%
rename from python/pyspark/ml/tests/connect/test_connect_spark_ml_classification.py
rename to python/pyspark/ml/tests/connect/test_parity_classification.py
index 2000a38d9e61..ae358f70b184 100644
--- a/python/pyspark/ml/tests/connect/test_connect_spark_ml_classification.py
+++ b/python/pyspark/ml/tests/connect/test_parity_classification.py
@@ -22,7 +22,7 @@ from pyspark.ml.tests.test_classification import ClassificationTestsMixin
from pyspark.sql import SparkSession
-class ClassificationTestsOnConnect(ClassificationTestsMixin, unittest.TestCase):
+class ClassificationParityTests(ClassificationTestsMixin, unittest.TestCase):
def setUp(self) -> None:
self.spark = SparkSession.builder.remote(
os.environ.get("SPARK_CONNECT_TESTING_REMOTE", "local[2]")
@@ -38,7 +38,7 @@ class ClassificationTestsOnConnect(ClassificationTestsMixin, unittest.TestCase):
if __name__ == "__main__":
- from pyspark.ml.tests.connect.test_connect_spark_ml_classification import * # noqa: F401
+ from pyspark.ml.tests.connect.test_parity_classification import * # noqa: F401
try:
import xmlrunner # type: ignore[import]
diff --git a/python/pyspark/ml/tests/test_classification.py b/python/pyspark/ml/tests/test_classification.py
index ee72e0394e3a..a29794344745 100644
--- a/python/pyspark/ml/tests/test_classification.py
+++ b/python/pyspark/ml/tests/test_classification.py
@@ -29,6 +29,16 @@ from pyspark.ml.classification import (
LogisticRegressionModel,
LogisticRegressionSummary,
BinaryLogisticRegressionSummary,
+ DecisionTreeClassifier,
+ DecisionTreeClassificationModel,
+ RandomForestClassifier,
+ RandomForestClassificationModel,
+ RandomForestClassificationSummary,
+ RandomForestClassificationTrainingSummary,
+ BinaryRandomForestClassificationSummary,
+ BinaryRandomForestClassificationTrainingSummary,
+ GBTClassifier,
+ GBTClassificationModel,
)
@@ -283,6 +293,327 @@ class ClassificationTestsMixin:
except OSError:
pass
+ def test_decision_tree_classifier(self):
+ df = (
+ self.spark.createDataFrame(
+ [
+ (1.0, 1.0, Vectors.dense(0.0, 5.0)),
+ (0.0, 2.0, Vectors.dense(1.0, 2.0)),
+ (1.0, 3.0, Vectors.dense(2.0, 1.0)),
+ (0.0, 4.0, Vectors.dense(3.0, 3.0)),
+ ],
+ ["label", "weight", "features"],
+ )
+ .coalesce(1)
+ .sortWithinPartitions("weight")
+ )
+
+ dt = DecisionTreeClassifier(
+ maxDepth=2,
+ labelCol="label",
+ leafCol="leaf",
+ seed=1,
+ )
+ self.assertEqual(dt.getMaxDepth(), 2)
+ self.assertEqual(dt.getSeed(), 1)
+ self.assertEqual(dt.getLabelCol(), "label")
+ self.assertEqual(dt.getLeafCol(), "leaf")
+
+ # Estimator save & load
+ with tempfile.TemporaryDirectory(prefix="decision_tree_classifier") as d:
+ dt.write().overwrite().save(d)
+ dt2 = DecisionTreeClassifier.load(d)
+ self.assertEqual(str(dt), str(dt2))
+
+ model = dt.fit(df)
+ self.assertEqual(model.numClasses, 2)
+ self.assertEqual(model.numFeatures, 2)
+ self.assertEqual(model.depth, 2)
+ self.assertEqual(model.numNodes, 5)
+ self.assertEqual(model.featureImportances, Vectors.sparse(2, [0], [1.0]))
+ self.assertTrue("depth=2, numNodes=5, numClasses=2, numFeatures=2" in model.toDebugString)
+ self.assertTrue("If (feature 0 <= 2.5)" in model.toDebugString)
+
+ vec = Vectors.dense(0.0, 5.0)
+ self.assertEqual(model.predict(vec), 1.0)
+ self.assertEqual(model.predictRaw(vec), Vectors.dense(0.0, 1.0))
+ self.assertEqual(model.predictProbability(vec), Vectors.dense(0.0, 1.0))
+ self.assertEqual(model.predictLeaf(vec), 0.0)
+
+ output = model.transform(df)
+ expected_cols = [
+ "label",
+ "weight",
+ "features",
+ "rawPrediction",
+ "probability",
+ "prediction",
+ "leaf",
+ ]
+ self.assertEqual(output.columns, expected_cols)
+ self.assertEqual(output.count(), 4)
+
+ # Model save & load
+ with tempfile.TemporaryDirectory(prefix="decision_tree_model") as d:
+ model.write().overwrite().save(d)
+ model2 = DecisionTreeClassificationModel.load(d)
+ self.assertEqual(str(model), str(model2))
+ self.assertEqual(model.toDebugString, model2.toDebugString)
+
+ def test_gbt_classifier(self):
+ df = (
+ self.spark.createDataFrame(
+ [
+ (1.0, 1.0, Vectors.dense(0.0, 5.0)),
+ (0.0, 2.0, Vectors.dense(1.0, 2.0)),
+ (1.0, 3.0, Vectors.dense(2.0, 1.0)),
+ (0.0, 4.0, Vectors.dense(3.0, 3.0)),
+ ],
+ ["label", "weight", "features"],
+ )
+ .coalesce(1)
+ .sortWithinPartitions("weight")
+ )
+
+ gbt = GBTClassifier(
+ maxIter=3,
+ maxDepth=2,
+ labelCol="label",
+ leafCol="leaf",
+ seed=1,
+ )
+ self.assertEqual(gbt.getMaxIter(), 3)
+ self.assertEqual(gbt.getMaxDepth(), 2)
+ self.assertEqual(gbt.getSeed(), 1)
+ self.assertEqual(gbt.getLabelCol(), "label")
+ self.assertEqual(gbt.getLeafCol(), "leaf")
+
+ # Estimator save & load
+ with tempfile.TemporaryDirectory(prefix="gbt_classifier") as d:
+ gbt.write().overwrite().save(d)
+ gbt2 = GBTClassifier.load(d)
+ self.assertEqual(str(gbt), str(gbt2))
+
+ model = gbt.fit(df)
+ self.assertEqual(model.numClasses, 2)
+ self.assertEqual(model.numFeatures, 2)
+ # TODO(SPARK-50843): Support access submodel in TreeEnsembleModel
+ # model.trees
+ self.assertEqual(model.treeWeights, [1.0, 0.1, 0.1])
+ self.assertEqual(model.totalNumNodes, 15)
+ self.assertTrue(np.allclose(model.featureImportances, [0.3333, 0.6667], atol=1e-4))
+ self.assertTrue("numTrees=3, numClasses=2, numFeatures=2" in model.toDebugString)
+ self.assertTrue("If (feature 0 <= 0.5)" in model.toDebugString)
+
+ vec = Vectors.dense(0.0, 5.0)
+ self.assertEqual(model.predict(vec), 1.0)
+ self.assertTrue(np.allclose(model.predictRaw(vec), [-1.0915, 1.0915], atol=1e-4))
+ self.assertTrue(np.allclose(model.predictProbability(vec), [0.1013, 0.8987], atol=1e-4))
+ self.assertEqual(model.predictLeaf(vec), Vectors.dense(0.0, 0.0, 0.0))
+
+ # GBT-specific method: evaluateEachIteration
+ self.assertTrue(
+ np.allclose(
+ model.evaluateEachIteration(df),
+ [0.253856022085945, 0.23205304779013333, 0.21358401299568353],
+ atol=1e-4,
+ )
+ )
+
+ output = model.transform(df)
+ expected_cols = [
+ "label",
+ "weight",
+ "features",
+ "rawPrediction",
+ "probability",
+ "prediction",
+ "leaf",
+ ]
+ self.assertEqual(output.columns, expected_cols)
+ self.assertEqual(output.count(), 4)
+
+ # Model save & load
+ with tempfile.TemporaryDirectory(prefix="gbt_classification_model") as d:
+ model.write().overwrite().save(d)
+ model2 = GBTClassificationModel.load(d)
+ self.assertEqual(str(model), str(model2))
+ self.assertEqual(model.toDebugString, model2.toDebugString)
+
+ def test_binary_random_forest_classifier(self):
+ df = (
+ self.spark.createDataFrame(
+ [
+ (1.0, 1.0, Vectors.dense(0.0, 5.0)),
+ (0.0, 2.0, Vectors.dense(1.0, 2.0)),
+ (1.0, 3.0, Vectors.dense(2.0, 1.0)),
+ (0.0, 4.0, Vectors.dense(3.0, 3.0)),
+ ],
+ ["label", "weight", "features"],
+ )
+ .coalesce(1)
+ .sortWithinPartitions("weight")
+ )
+
+ rf = RandomForestClassifier(
+ numTrees=3,
+ maxDepth=2,
+ labelCol="label",
+ leafCol="leaf",
+ seed=1,
+ )
+ self.assertEqual(rf.getNumTrees(), 3)
+ self.assertEqual(rf.getMaxDepth(), 2)
+ self.assertEqual(rf.getSeed(), 1)
+ self.assertEqual(rf.getLabelCol(), "label")
+ self.assertEqual(rf.getLeafCol(), "leaf")
+
+ # Estimator save & load
+ with tempfile.TemporaryDirectory(prefix="binary_random_forest_classifier") as d:
+ rf.write().overwrite().save(d)
+ rf2 = RandomForestClassifier.load(d)
+ self.assertEqual(str(rf), str(rf2))
+
+ model = rf.fit(df)
+ self.assertEqual(model.numClasses, 2)
+ self.assertEqual(model.numFeatures, 2)
+ # TODO(SPARK-50843): Support access submodel in TreeEnsembleModel
+ # model.trees
+ self.assertEqual(model.treeWeights, [1.0, 1.0, 1.0])
+ self.assertEqual(model.totalNumNodes, 9)
+ self.assertTrue(np.allclose(model.featureImportances, [0.7778, 0.2222], atol=1e-4))
+ self.assertTrue("numTrees=3, numClasses=2, numFeatures=2" in model.toDebugString)
+ self.assertTrue("If (feature 0 <= 0.5)" in model.toDebugString)
+
+ vec = Vectors.dense(0.0, 5.0)
+ self.assertEqual(model.predict(vec), 1.0)
+ self.assertEqual(model.predictRaw(vec), Vectors.dense(1.0, 2.0))
+ self.assertTrue(np.allclose(model.predictProbability(vec), [0.3333, 0.6667], atol=1e-4))
+ self.assertEqual(model.predictLeaf(vec), Vectors.dense(0.0, 0.0, 1.0))
+
+ output = model.transform(df)
+ expected_cols = [
+ "label",
+ "weight",
+ "features",
+ "rawPrediction",
+ "probability",
+ "prediction",
+ "leaf",
+ ]
+ self.assertEqual(output.columns, expected_cols)
+ self.assertEqual(output.count(), 4)
+
+ # model summary
+ summary = model.summary
+ self.assertTrue(isinstance(summary, BinaryRandomForestClassificationSummary))
+ self.assertTrue(isinstance(summary, BinaryRandomForestClassificationTrainingSummary))
+ self.assertEqual(summary.labels, [0.0, 1.0])
+ self.assertEqual(summary.accuracy, 0.75)
+ self.assertEqual(summary.areaUnderROC, 0.875)
+ self.assertEqual(summary.predictions.columns, expected_cols)
+
+ summary2 = model.evaluate(df)
+ self.assertTrue(isinstance(summary2, BinaryRandomForestClassificationSummary))
+ self.assertFalse(isinstance(summary2, BinaryRandomForestClassificationTrainingSummary))
+ self.assertEqual(summary2.labels, [0.0, 1.0])
+ self.assertEqual(summary2.accuracy, 0.75)
+ self.assertEqual(summary.areaUnderROC, 0.875)
+ self.assertEqual(summary2.predictions.columns, expected_cols)
+
+ # Model save & load
+ with tempfile.TemporaryDirectory(prefix="binary_random_forest_model") as d:
+ model.write().overwrite().save(d)
+ model2 = RandomForestClassificationModel.load(d)
+ self.assertEqual(str(model), str(model2))
+ self.assertEqual(model.toDebugString, model2.toDebugString)
+
+ def test_multiclass_random_forest_classifier(self):
+ df = (
+ self.spark.createDataFrame(
+ [
+ (1.0, 1.0, Vectors.dense(0.0, 5.0)),
+ (0.0, 2.0, Vectors.dense(1.0, 2.0)),
+ (1.0, 3.0, Vectors.dense(2.0, 1.0)),
+ (2.0, 4.0, Vectors.dense(3.0, 3.0)),
+ ],
+ ["label", "weight", "features"],
+ )
+ .coalesce(1)
+ .sortWithinPartitions("weight")
+ )
+
+ rf = RandomForestClassifier(
+ numTrees=3,
+ maxDepth=2,
+ labelCol="label",
+ leafCol="leaf",
+ seed=1,
+ )
+ self.assertEqual(rf.getNumTrees(), 3)
+ self.assertEqual(rf.getMaxDepth(), 2)
+ self.assertEqual(rf.getSeed(), 1)
+ self.assertEqual(rf.getLabelCol(), "label")
+ self.assertEqual(rf.getLeafCol(), "leaf")
+
+ # Estimator save & load
+ with tempfile.TemporaryDirectory(prefix="binary_random_forest_classifier") as d:
+ rf.write().overwrite().save(d)
+ rf2 = RandomForestClassifier.load(d)
+ self.assertEqual(str(rf), str(rf2))
+
+ model = rf.fit(df)
+ self.assertEqual(model.numClasses, 3)
+ self.assertEqual(model.numFeatures, 2)
+ # TODO(SPARK-50843): Support access submodel in TreeEnsembleModel
+ # model.trees
+ self.assertEqual(model.treeWeights, [1.0, 1.0, 1.0])
+ self.assertEqual(model.totalNumNodes, 9)
+ self.assertEqual(model.featureImportances, Vectors.sparse(2, [0], [1.0]))
+ self.assertTrue("numTrees=3, numClasses=3, numFeatures=2" in model.toDebugString)
+ self.assertTrue("If (feature 0 <= 1.5)" in model.toDebugString)
+
+ vec = Vectors.dense(0.0, 5.0)
+ self.assertEqual(model.predict(vec), 0.0)
+ self.assertEqual(model.predictRaw(vec), Vectors.dense(1.5, 1.5, 0.0))
+ self.assertTrue(np.allclose(model.predictProbability(vec), [0.5, 0.5, 0.0], atol=1e-4))
+ self.assertEqual(model.predictLeaf(vec), Vectors.dense(0.0, 0.0, 0.0))
+
+ output = model.transform(df)
+ expected_cols = [
+ "label",
+ "weight",
+ "features",
+ "rawPrediction",
+ "probability",
+ "prediction",
+ "leaf",
+ ]
+ self.assertEqual(output.columns, expected_cols)
+ self.assertEqual(output.count(), 4)
+
+ # model summary
+ summary = model.summary
+ self.assertTrue(isinstance(summary, RandomForestClassificationSummary))
+ self.assertTrue(isinstance(summary, RandomForestClassificationTrainingSummary))
+ self.assertEqual(summary.labels, [0.0, 1.0, 2.0])
+ self.assertEqual(summary.accuracy, 0.5)
+ self.assertEqual(summary.predictions.columns, expected_cols)
+
+ summary2 = model.evaluate(df)
+ self.assertTrue(isinstance(summary2, RandomForestClassificationSummary))
+ self.assertFalse(isinstance(summary2, RandomForestClassificationTrainingSummary))
+ self.assertEqual(summary2.labels, [0.0, 1.0, 2.0])
+ self.assertEqual(summary2.accuracy, 0.5)
+ self.assertEqual(summary2.predictions.columns, expected_cols)
+
+ # Model save & load
+ with tempfile.TemporaryDirectory(prefix="multiclass_random_forest_model") as d:
+ model.write().overwrite().save(d)
+ model2 = RandomForestClassificationModel.load(d)
+ self.assertEqual(str(model), str(model2))
+ self.assertEqual(model.toDebugString, model2.toDebugString)
+
class ClassificationTests(ClassificationTestsMixin, unittest.TestCase):
def setUp(self) -> None:
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLUtils.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLUtils.scala
index 72c86401eb63..f3db36c7647d 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLUtils.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLUtils.scala
@@ -279,9 +279,17 @@ private[ml] object MLUtils {
// The attributes could be retrieved from the corresponding python class
private lazy val ALLOWED_ATTRIBUTES = HashSet(
"toString",
+ "toDebugString",
"numFeatures",
"predict", // PredictionModel
+ "predictLeaf", // Tree models
"numClasses",
+ "depth", // DecisionTreeClassificationModel
+ "numNodes", // Tree models
+ "totalNumNodes", // Tree models
+ "javaTreeWeights", // Tree models
+ "treeWeights", // Tree models
+ "featureImportances", // Tree models
"predictRaw", // ClassificationModel
"predictProbability", // ProbabilisticClassificationModel
"coefficients",
@@ -291,6 +299,7 @@ private[ml] object MLUtils {
"summary",
"hasSummary",
"evaluate", // LogisticRegressionModel
+ "evaluateEachIteration", // GBTClassificationModel
"predictions",
"predictionCol",
"labelCol",
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org