-
Notifications
You must be signed in to change notification settings - Fork 28.5k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[SPARK-49907][ML][CONNECT] Support spark.ml on Connect
### What changes were proposed in this pull request? This PR which derived from #40479 authored by WeichenXu123 enables running spark.ml on Connect. Currently, this PR supports the following functionalities: - Fit operation in LogisticRegression - Transform/predict operation in LogisticRegressionModel - Retrieving attributes in LogisticRegressionModel - Retrieving summary and its attributes in LogisticRegressionModel - Read/write operations for LogisticRegressionModel - Read/write operations for LogisticRegression - Evaluate a dataset in LogisticRegressionModel and return a summary and retrieve attributes from it. ### Why are the changes needed? It's a new feature that makes spark.ml run on connect environment. ### Does this PR introduce _any_ user-facing change? Yes, new feature. ### How was this patch tested? Make sure the CI (especially the newly added tests) pass. And we can manually run below code without any exception. ``` python (pyspark) userbobby:~ $ pyspark --remote sc://localhost Python 3.11.10 (main, Oct 3 2024, 07:29:13) [GCC 11.2.0] on linux Type "help", "copyright", "credits" or "license" for more information. Welcome to ____ __ / __/__ ___ _____/ /__ _\ \/ _ \/ _ `/ __/ '_/ /__ / .__/\_,_/_/ /_/\_\ version 4.0.0.dev0 /_/ Using Python version 3.11.10 (main, Oct 3 2024 07:29:13) Client connected to the Spark Connect server at localhost SparkSession available as 'spark'. >>> from pyspark.ml.classification import (LogisticRegression, LogisticRegressionModel) >>> from pyspark.ml.linalg import Vectors >>> >>> df = spark.createDataFrame([ ... (Vectors.dense([1.0, 2.0]), 1), ... (Vectors.dense([2.0, -1.0]), 1), ... (Vectors.dense([-3.0, -2.0]), 0), ... (Vectors.dense([-1.0, -2.0]), 0), ... ], schema=['features', 'label']) >>> lr = LogisticRegression() >>> lr.setMaxIter(30) LogisticRegression_27d6a4e4f39d >>> lr.setThreshold(0.8) LogisticRegression_27d6a4e4f39d >>> lr.write().overwrite().save("/tmp/connect-ml-demo/estimator") >>> loaded_lr = LogisticRegression.load("/tmp/connect-ml-demo/estimator") >>> assert (loaded_lr.getThreshold() == 0.8) >>> assert loaded_lr.getMaxIter() == 30 >>> >>> model: LogisticRegressionModel = lr.fit(df) >>> assert (model.getThreshold() == 0.8) >>> assert model.getMaxIter() == 30 >>> model.predictRaw(Vectors.dense([1.0, 2.0])) DenseVector([-21.1048, 21.1048]) >>> model.summary.roc.show() +---+---+ |FPR|TPR| +---+---+ |0.0|0.0| |0.0|0.5| |0.0|1.0| |0.5|1.0| |1.0|1.0| |1.0|1.0| +---+---+ >>> model.summary.weightedRecall 1.0 >>> model.summary.recallByLabel [1.0, 1.0] >>> model.coefficients DenseVector([10.3964, 4.513]) >>> model.intercept 1.682348909633995 >>> model.transform(df).show() +-----------+-----+--------------------+--------------------+----------+ | features|label| rawPrediction| probability|prediction| +-----------+-----+--------------------+--------------------+----------+ | [1.0,2.0]| 1|[-21.104818251026...|[6.82800596289009...| 1.0| | [2.0,-1.0]| 1|[-17.962094978515...|[1.58183529116627...| 1.0| |[-3.0,-2.0]| 0|[38.5329050234205...| [1.0,0.0]| 0.0| |[-1.0,-2.0]| 0|[17.7401204317581...|[0.99999998025016...| 0.0| +-----------+-----+--------------------+--------------------+----------+ >>> model.write().overwrite().save("/tmp/connect-ml-demo/model") >>> loaded_model = LogisticRegressionModel.load("/tmp/connect-ml-demo/model") >>> assert loaded_model.getMaxIter() == 30 >>> loaded_model.transform(df).show() +-----------+-----+--------------------+--------------------+----------+ | features|label| rawPrediction| probability|prediction| +-----------+-----+--------------------+--------------------+----------+ | [1.0,2.0]| 1|[-21.104818251026...|[6.82800596289009...| 1.0| | [2.0,-1.0]| 1|[-17.962094978515...|[1.58183529116627...| 1.0| |[-3.0,-2.0]| 0|[38.5329050234205...| [1.0,0.0]| 0.0| |[-1.0,-2.0]| 0|[17.7401204317581...|[0.99999998025016...| 0.0| +-----------+-----+--------------------+--------------------+----------+ >>> >>> summary = loaded_model.evaluate(df) >>> summary.weightCol 'weightCol' >>> summary.recallByLabel [1.0, 1.0] >>> summary.accuracy 1.0 >>> summary.predictions.show() +-----------+-----+--------------------+--------------------+----------+ | features|label| rawPrediction| probability|prediction| +-----------+-----+--------------------+--------------------+----------+ | [1.0,2.0]| 1|[-21.104818251026...|[6.82800596289009...| 1.0| | [2.0,-1.0]| 1|[-17.962094978515...|[1.58183529116627...| 1.0| |[-3.0,-2.0]| 0|[38.5329050234205...| [1.0,0.0]| 0.0| |[-1.0,-2.0]| 0|[17.7401204317581...|[0.99999998025016...| 0.0| +-----------+-----+--------------------+--------------------+----------+ ``` ### Was this patch authored or co-authored using generative AI tooling? No Closes #48791 from wbo4958/connect-ml. Authored-by: Bobby Wang <[email protected]> Signed-off-by: Ruifeng Zheng <[email protected]>
- Loading branch information
1 parent
f96417f
commit fafe43c
Showing
47 changed files
with
4,384 additions
and
640 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
20 changes: 20 additions & 0 deletions
20
mllib/src/main/resources/META-INF/services/org.apache.spark.ml.Estimator
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,20 @@ | ||
# | ||
# 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. | ||
# | ||
|
||
# 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. | ||
org.apache.spark.ml.classification.LogisticRegression |
20 changes: 20 additions & 0 deletions
20
mllib/src/main/resources/META-INF/services/org.apache.spark.ml.Transformer
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,20 @@ | ||
# | ||
# 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. | ||
# | ||
|
||
# Spark Connect ML uses ServiceLoader to find out the supported Spark Ml non-model transformer. | ||
# So register the supported transformer here if you're trying to add a new one. | ||
org.apache.spark.ml.feature.VectorAssembler |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
28 changes: 28 additions & 0 deletions
28
mllib/src/main/scala/org/apache/spark/ml/util/Summary.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,28 @@ | ||
/* | ||
* 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. | ||
*/ | ||
|
||
package org.apache.spark.ml.util | ||
|
||
import org.apache.spark.annotation.Since | ||
|
||
/** | ||
* Trait for the Summary | ||
* All the summaries should extend from this Summary in order to | ||
* support connect. | ||
*/ | ||
@Since("4.0.0") | ||
private[spark] trait Summary |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,16 @@ | ||
# | ||
# 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. | ||
# |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,76 @@ | ||
# | ||
# 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 typing import Optional, TYPE_CHECKING, List | ||
|
||
import pyspark.sql.connect.proto as pb2 | ||
from pyspark.sql.connect.plan import LogicalPlan | ||
|
||
if TYPE_CHECKING: | ||
from pyspark.sql.connect.client import SparkConnectClient | ||
|
||
|
||
class TransformerRelation(LogicalPlan): | ||
"""A logical plan for transforming of a transformer which could be a cached model | ||
or a non-model transformer like VectorAssembler.""" | ||
|
||
def __init__( | ||
self, | ||
child: Optional["LogicalPlan"], | ||
name: str, | ||
ml_params: pb2.MlParams, | ||
uid: str = "", | ||
is_model: bool = True, | ||
) -> None: | ||
super().__init__(child) | ||
self._name = name | ||
self._ml_params = ml_params | ||
self._uid = uid | ||
self._is_model = is_model | ||
|
||
def plan(self, session: "SparkConnectClient") -> pb2.Relation: | ||
assert self._child is not None | ||
plan = self._create_proto_relation() | ||
plan.ml_relation.transform.input.CopyFrom(self._child.plan(session)) | ||
|
||
if self._is_model: | ||
plan.ml_relation.transform.obj_ref.CopyFrom(pb2.ObjectRef(id=self._name)) | ||
else: | ||
plan.ml_relation.transform.transformer.CopyFrom( | ||
pb2.MlOperator(name=self._name, uid=self._uid, type=pb2.MlOperator.TRANSFORMER) | ||
) | ||
|
||
if self._ml_params is not None: | ||
plan.ml_relation.transform.params.CopyFrom(self._ml_params) | ||
|
||
return plan | ||
|
||
|
||
class AttributeRelation(LogicalPlan): | ||
"""A logical plan used in ML to represent an attribute of an instance, which | ||
could be a model or a summary. This attribute returns a DataFrame. | ||
""" | ||
|
||
def __init__(self, ref_id: str, methods: List[pb2.Fetch.Method]) -> None: | ||
super().__init__(None) | ||
self._ref_id = ref_id | ||
self._methods = methods | ||
|
||
def plan(self, session: "SparkConnectClient") -> pb2.Relation: | ||
plan = self._create_proto_relation() | ||
plan.ml_relation.fetch.obj_ref.CopyFrom(pb2.ObjectRef(id=self._ref_id)) | ||
plan.ml_relation.fetch.methods.extend(self._methods) | ||
return plan |
Oops, something went wrong.