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/03 02:54:00 UTC
(spark) branch master updated: [SPARK-51060][ML][PYTHON][CONNECT] Support `QuantileDiscretizer` on Conenct
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 9ad8d226a8e0 [SPARK-51060][ML][PYTHON][CONNECT] Support `QuantileDiscretizer` on Conenct
9ad8d226a8e0 is described below
commit 9ad8d226a8e0c5c3220448d69290b0898ded2dbd
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Mon Feb 3 10:53:42 2025 +0800
[SPARK-51060][ML][PYTHON][CONNECT] Support `QuantileDiscretizer` on Conenct
### What changes were proposed in this pull request?
Support `QuantileDiscretizer` on Conenct
### Why are the changes needed?
for feature parity
### Does this PR introduce _any_ user-facing change?
yes
### How was this patch tested?
added tests
### Was this patch authored or co-authored using generative AI tooling?
no
Closes #49761 from zhengruifeng/ml_connect_qd.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@apache.org>
---
.../services/org.apache.spark.ml.Estimator | 1 +
python/pyspark/ml/feature.py | 18 +++++
python/pyspark/ml/tests/test_feature.py | 93 ++++++++++++++++++++++
.../org/apache/spark/sql/connect/ml/MLUtils.scala | 1 +
4 files changed, 113 insertions(+)
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 d068303494d7..128938e9eb48 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
@@ -51,6 +51,7 @@ org.apache.spark.ml.recommendation.ALS
org.apache.spark.ml.fpm.FPGrowth
# feature
+org.apache.spark.ml.feature.QuantileDiscretizer
org.apache.spark.ml.feature.RFormula
org.apache.spark.ml.feature.Imputer
org.apache.spark.ml.feature.StandardScaler
diff --git a/python/pyspark/ml/feature.py b/python/pyspark/ml/feature.py
index cc8f19aaf82c..d7754367015f 100755
--- a/python/pyspark/ml/feature.py
+++ b/python/pyspark/ml/feature.py
@@ -53,6 +53,7 @@ from pyspark.ml.param.shared import (
from pyspark.ml.util import JavaMLReadable, JavaMLWritable, try_remote_attribute_relation
from pyspark.ml.wrapper import JavaEstimator, JavaModel, JavaParams, JavaTransformer, _jvm
from pyspark.ml.common import inherit_doc
+from pyspark.sql.utils import is_remote
if TYPE_CHECKING:
from py4j.java_gateway import JavaObject
@@ -3754,6 +3755,23 @@ class QuantileDiscretizer(
"""
Private method to convert the java_model to a Python model.
"""
+ if is_remote():
+ remote_model = JavaModel(java_model)
+ if self.isSet(self.inputCol):
+ return Bucketizer(
+ splits=remote_model._call_java("getSplits"),
+ inputCol=self.getInputCol(),
+ outputCol=self.getOutputCol(),
+ handleInvalid=self.getHandleInvalid(),
+ )
+ else:
+ return Bucketizer(
+ splitsArray=remote_model._call_java("getSplitsArray"),
+ inputCols=self.getInputCols(),
+ outputCols=self.getOutputCols(),
+ handleInvalid=self.getHandleInvalid(),
+ )
+
if self.isSet(self.inputCol):
return Bucketizer(
splits=list(java_model.getSplits()),
diff --git a/python/pyspark/ml/tests/test_feature.py b/python/pyspark/ml/tests/test_feature.py
index cc4783295079..b458a4eb7a25 100644
--- a/python/pyspark/ml/tests/test_feature.py
+++ b/python/pyspark/ml/tests/test_feature.py
@@ -26,6 +26,7 @@ from pyspark.ml.feature import (
DCT,
Binarizer,
Bucketizer,
+ QuantileDiscretizer,
CountVectorizer,
CountVectorizerModel,
OneHotEncoder,
@@ -978,6 +979,98 @@ class FeatureTestsMixin:
binarizer2 = Binarizer.load(d)
self.assertEqual(str(binarizer), str(binarizer2))
+ def test_quantile_discretizer_single_column(self):
+ spark = self.spark
+ values = [(0.1,), (0.4,), (1.2,), (1.5,), (float("nan"),), (float("nan"),)]
+ df = spark.createDataFrame(values, ["values"])
+
+ qds = QuantileDiscretizer(inputCol="values", outputCol="buckets")
+ qds.setNumBuckets(2)
+ qds.setRelativeError(0.01)
+ qds.setHandleInvalid("keep")
+
+ self.assertEqual(qds.getInputCol(), "values")
+ self.assertEqual(qds.getOutputCol(), "buckets")
+ self.assertEqual(qds.getNumBuckets(), 2)
+ self.assertEqual(qds.getRelativeError(), 0.01)
+ self.assertEqual(qds.getHandleInvalid(), "keep")
+
+ bucketizer = qds.fit(df)
+ self.assertIsInstance(bucketizer, Bucketizer)
+ # Bucketizer doesn't inherit uid from QuantileDiscretizer
+ # self.assertEqual(qds.uid, bucketizer.uid)
+
+ # check model coefficients
+ self.assertEqual(bucketizer.getSplits(), [float("-inf"), 0.4, float("inf")])
+
+ output = bucketizer.transform(df)
+ self.assertEqual(output.columns, ["values", "buckets"])
+ self.assertEqual(output.count(), 6)
+
+ # save & load
+ with tempfile.TemporaryDirectory(prefix="quantile_discretizer_single_column") as d:
+ qds.write().overwrite().save(d)
+ qds2 = QuantileDiscretizer.load(d)
+ self.assertEqual(str(qds), str(qds2))
+
+ bucketizer.write().overwrite().save(d)
+ bucketizer2 = Bucketizer.load(d)
+ self.assertEqual(str(bucketizer), str(bucketizer2))
+
+ def test_quantile_discretizer_multiple_columns(self):
+ spark = self.spark
+ inputs = [
+ (0.1, 0.0),
+ (0.4, 1.0),
+ (1.2, 1.3),
+ (1.5, 1.5),
+ (float("nan"), float("nan")),
+ (float("nan"), float("nan")),
+ ]
+ df = spark.createDataFrame(inputs, ["input1", "input2"])
+
+ qds = QuantileDiscretizer(
+ relativeError=0.01,
+ handleInvalid="keep",
+ numBuckets=2,
+ inputCols=["input1", "input2"],
+ outputCols=["output1", "output2"],
+ )
+
+ self.assertEqual(qds.getInputCols(), ["input1", "input2"])
+ self.assertEqual(qds.getOutputCols(), ["output1", "output2"])
+ self.assertEqual(qds.getNumBuckets(), 2)
+ self.assertEqual(qds.getRelativeError(), 0.01)
+ self.assertEqual(qds.getHandleInvalid(), "keep")
+
+ bucketizer = qds.fit(df)
+ self.assertIsInstance(bucketizer, Bucketizer)
+ # Bucketizer doesn't inherit uid from QuantileDiscretizer
+ # self.assertEqual(qds.uid, bucketizer.uid)
+
+ # check model coefficients
+ self.assertEqual(
+ bucketizer.getSplitsArray(),
+ [
+ [float("-inf"), 0.4, float("inf")],
+ [float("-inf"), 1.0, float("inf")],
+ ],
+ )
+
+ output = bucketizer.transform(df)
+ self.assertEqual(output.columns, ["input1", "input2", "output1", "output2"])
+ self.assertEqual(output.count(), 6)
+
+ # save & load
+ with tempfile.TemporaryDirectory(prefix="quantile_discretizer_multiple_columns") as d:
+ qds.write().overwrite().save(d)
+ qds2 = QuantileDiscretizer.load(d)
+ self.assertEqual(str(qds), str(qds2))
+
+ bucketizer.write().overwrite().save(d)
+ bucketizer2 = Bucketizer.load(d)
+ self.assertEqual(str(bucketizer), str(bucketizer2))
+
def test_bucketizer(self):
df = self.spark.createDataFrame(
[
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 6f59af46d0a7..135e7bac826f 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
@@ -634,6 +634,7 @@ private[ml] object MLUtils {
(classOf[FPGrowthModel], Set("associationRules", "freqItemsets")),
// Feature Models
+ (classOf[Bucketizer], Set("getSplits", "getSplitsArray")),
(classOf[ImputerModel], Set("surrogateDF")),
(classOf[StandardScalerModel], Set("mean", "std")),
(classOf[MaxAbsScalerModel], Set("maxAbs")),
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org