You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2025/01/16 08:14:54 UTC
(spark) branch branch-4.0 updated: [SPARK-49907][ML][CONNECT][FOLLOW-UP] Move pyspark.ml.remote to pyspark.ml.connect
This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch branch-4.0
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.0 by this push:
new 71c98295fe6d [SPARK-49907][ML][CONNECT][FOLLOW-UP] Move pyspark.ml.remote to pyspark.ml.connect
71c98295fe6d is described below
commit 71c98295fe6d42dd9a063601942f0d4f3b872ff2
Author: Hyukjin Kwon <gu...@apache.org>
AuthorDate: Thu Jan 16 17:14:34 2025 +0900
[SPARK-49907][ML][CONNECT][FOLLOW-UP] Move pyspark.ml.remote to pyspark.ml.connect
### What changes were proposed in this pull request?
This PR is a retry of https://github.com/apache/spark/pull/49522 that refactors pyspark.ml.remote to pyspark.ml.connect
### Why are the changes needed?
To move the connect code into one place.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests should cover
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #49524 from HyukjinKwon/ml-rename.
Lead-authored-by: Hyukjin Kwon <gu...@apache.org>
Co-authored-by: Bobby Wang <wb...@gmail.com>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
(cherry picked from commit 97241eb0cfc98f8911dc205d7a19f31a0eaad5fb)
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
python/packaging/classic/setup.py | 1 -
python/packaging/connect/setup.py | 1 -
python/pyspark/ml/classification.py | 2 +-
python/pyspark/ml/{remote => connect}/proto.py | 0
python/pyspark/ml/{remote => connect}/readwrite.py | 2 +-
python/pyspark/ml/{remote => connect}/serialize.py | 0
python/pyspark/ml/connect/util.py | 15 +-
python/pyspark/ml/remote/__init__.py | 16 --
python/pyspark/ml/remote/util.py | 293 ---------------------
python/pyspark/ml/util.py | 267 ++++++++++++++++++-
python/pyspark/ml/wrapper.py | 2 +-
11 files changed, 283 insertions(+), 316 deletions(-)
diff --git a/python/packaging/classic/setup.py b/python/packaging/classic/setup.py
index 211252a73d10..ae20fc1efdef 100755
--- a/python/packaging/classic/setup.py
+++ b/python/packaging/classic/setup.py
@@ -269,7 +269,6 @@ try:
"pyspark.mllib.stat",
"pyspark.ml",
"pyspark.ml.connect",
- "pyspark.ml.remote",
"pyspark.ml.linalg",
"pyspark.ml.param",
"pyspark.ml.torch",
diff --git a/python/packaging/connect/setup.py b/python/packaging/connect/setup.py
index c7fbe98c06e7..4c4afab7c403 100755
--- a/python/packaging/connect/setup.py
+++ b/python/packaging/connect/setup.py
@@ -147,7 +147,6 @@ try:
"pyspark.mllib.stat",
"pyspark.ml",
"pyspark.ml.connect",
- "pyspark.ml.remote",
"pyspark.ml.linalg",
"pyspark.ml.param",
"pyspark.ml.torch",
diff --git a/python/pyspark/ml/classification.py b/python/pyspark/ml/classification.py
index e003ba43ec7c..fa51d8828340 100644
--- a/python/pyspark/ml/classification.py
+++ b/python/pyspark/ml/classification.py
@@ -62,7 +62,7 @@ from pyspark.ml.param.shared import (
HasSolver,
HasParallelism,
)
-from pyspark.ml.remote.util import try_remote_attribute_relation
+from pyspark.ml.util import try_remote_attribute_relation
from pyspark.ml.tree import (
_DecisionTreeModel,
_DecisionTreeParams,
diff --git a/python/pyspark/ml/remote/proto.py b/python/pyspark/ml/connect/proto.py
similarity index 100%
rename from python/pyspark/ml/remote/proto.py
rename to python/pyspark/ml/connect/proto.py
diff --git a/python/pyspark/ml/remote/readwrite.py b/python/pyspark/ml/connect/readwrite.py
similarity index 98%
rename from python/pyspark/ml/remote/readwrite.py
rename to python/pyspark/ml/connect/readwrite.py
index 9149ab3bfd45..feb92e0a36f8 100644
--- a/python/pyspark/ml/remote/readwrite.py
+++ b/python/pyspark/ml/connect/readwrite.py
@@ -18,7 +18,7 @@
from typing import cast, Type, TYPE_CHECKING
import pyspark.sql.connect.proto as pb2
-from pyspark.ml.remote.serialize import serialize_ml_params, deserialize, deserialize_param
+from pyspark.ml.connect.serialize import serialize_ml_params, deserialize, deserialize_param
from pyspark.ml.util import MLWriter, MLReader, RL
from pyspark.ml.wrapper import JavaWrapper
diff --git a/python/pyspark/ml/remote/serialize.py b/python/pyspark/ml/connect/serialize.py
similarity index 100%
rename from python/pyspark/ml/remote/serialize.py
rename to python/pyspark/ml/connect/serialize.py
diff --git a/python/pyspark/ml/connect/util.py b/python/pyspark/ml/connect/util.py
index 63f9a30381d2..1c77baeba5f8 100644
--- a/python/pyspark/ml/connect/util.py
+++ b/python/pyspark/ml/connect/util.py
@@ -15,13 +15,16 @@
# limitations under the License.
#
-from typing import Any, Union, List, Tuple, Callable, Iterable
+from typing import Any, TypeVar, Callable, List, Tuple, Union, Iterable
import pandas as pd
from pyspark import cloudpickle
from pyspark.sql import DataFrame
from pyspark.sql.functions import col, pandas_udf
+import pyspark.sql.connect.proto as pb2
+
+FuncT = TypeVar("FuncT", bound=Callable[..., Any])
def aggregate_dataframe(
@@ -173,3 +176,13 @@ def transform_dataframe_column(
).drop(output_col_name)
else:
return result_spark_df
+
+
+def _extract_id_methods(obj_identifier: str) -> Tuple[List["pb2.Fetch.Method"], str]:
+ """Extract the obj reference id and the methods. Eg, model.summary"""
+ method_chain = obj_identifier.split(".")
+ obj_ref = method_chain[0]
+ methods: List["pb2.Fetch.Method"] = []
+ if len(method_chain) > 1:
+ methods = [pb2.Fetch.Method(method=m) for m in method_chain[1:]]
+ return methods, obj_ref
diff --git a/python/pyspark/ml/remote/__init__.py b/python/pyspark/ml/remote/__init__.py
deleted file mode 100644
index cce3acad34a4..000000000000
--- a/python/pyspark/ml/remote/__init__.py
+++ /dev/null
@@ -1,16 +0,0 @@
-#
-# 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.
-#
diff --git a/python/pyspark/ml/remote/util.py b/python/pyspark/ml/remote/util.py
deleted file mode 100644
index cb34dae16586..000000000000
--- a/python/pyspark/ml/remote/util.py
+++ /dev/null
@@ -1,293 +0,0 @@
-#
-# 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 functools
-import os
-from typing import Any, cast, TypeVar, Callable, TYPE_CHECKING, Type, List, Tuple
-
-import pyspark.sql.connect.proto as pb2
-from pyspark.ml.remote.serialize import serialize_ml_params, serialize, deserialize
-from pyspark.sql import is_remote
-
-if TYPE_CHECKING:
- from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame
- from pyspark.ml.wrapper import JavaWrapper, JavaEstimator
- from pyspark.ml.util import JavaMLReadable, JavaMLWritable
-
-FuncT = TypeVar("FuncT", bound=Callable[..., Any])
-
-
-def _extract_id_methods(obj_identifier: str) -> Tuple[List[pb2.Fetch.Method], str]:
- """Extract the obj reference id and the methods. Eg, model.summary"""
- method_chain = obj_identifier.split(".")
- obj_ref = method_chain[0]
- methods: List[pb2.Fetch.Method] = []
- if len(method_chain) > 1:
- methods = [pb2.Fetch.Method(method=m) for m in method_chain[1:]]
- return methods, obj_ref
-
-
-def try_remote_intermediate_result(f: FuncT) -> FuncT:
- """Mark the function/property that returns the intermediate result of the remote call.
- Eg, model.summary"""
-
- @functools.wraps(f)
- def wrapped(self: "JavaWrapper") -> Any:
- if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
- return f"{self._java_obj}.{f.__name__}"
- else:
- return f(self)
-
- return cast(FuncT, wrapped)
-
-
-def try_remote_attribute_relation(f: FuncT) -> FuncT:
- """Mark the function/property that returns a Relation.
- Eg, model.summary.roc"""
-
- @functools.wraps(f)
- def wrapped(self: "JavaWrapper", *args: Any, **kwargs: Any) -> Any:
- if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
- # The attribute returns a dataframe, we need to wrap it
- # in the AttributeRelation
- from pyspark.ml.remote.proto import AttributeRelation
- from pyspark.sql.connect.session import SparkSession
- from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame
-
- session = SparkSession.getActiveSession()
- assert session is not None
-
- assert isinstance(self._java_obj, str)
-
- methods, obj_ref = _extract_id_methods(self._java_obj)
- methods.append(
- pb2.Fetch.Method(method=f.__name__, args=serialize(session.client, *args))
- )
- plan = AttributeRelation(obj_ref, methods)
- return ConnectDataFrame(plan, session)
- else:
- return f(self, *args, **kwargs)
-
- return cast(FuncT, wrapped)
-
-
-def try_remote_fit(f: FuncT) -> FuncT:
- """Mark the function that fits a model."""
-
- @functools.wraps(f)
- def wrapped(self: "JavaEstimator", dataset: "ConnectDataFrame") -> Any:
- if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
- client = dataset.sparkSession.client
- input = dataset._plan.plan(client)
- assert isinstance(self._java_obj, str)
- estimator = pb2.MlOperator(
- name=self._java_obj, uid=self.uid, type=pb2.MlOperator.ESTIMATOR
- )
- command = pb2.Command()
- command.ml_command.fit.CopyFrom(
- pb2.MlCommand.Fit(
- estimator=estimator,
- params=serialize_ml_params(self, client),
- dataset=input,
- )
- )
- (_, properties, _) = client.execute_command(command)
- model_info = deserialize(properties)
- client.add_ml_cache(model_info.obj_ref.id)
- return model_info.obj_ref.id
- else:
- return f(self, dataset)
-
- return cast(FuncT, wrapped)
-
-
-def try_remote_transform_relation(f: FuncT) -> FuncT:
- """Mark the function/property that returns a relation for model transform."""
-
- @functools.wraps(f)
- def wrapped(self: "JavaWrapper", dataset: "ConnectDataFrame") -> Any:
- if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
- from pyspark.ml import Model, Transformer
- from pyspark.sql.connect.session import SparkSession
- from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame
-
- session = SparkSession.getActiveSession()
- assert session is not None
- # Model is also a Transformer, so we much match Model first
- if isinstance(self, Model):
- params = serialize_ml_params(self, session.client)
- from pyspark.ml.remote.proto import TransformerRelation
-
- assert isinstance(self._java_obj, str)
- return ConnectDataFrame(
- TransformerRelation(
- child=dataset._plan, name=self._java_obj, ml_params=params, is_model=True
- ),
- session,
- )
- elif isinstance(self, Transformer):
- params = serialize_ml_params(self, session.client)
- from pyspark.ml.remote.proto import TransformerRelation
-
- assert isinstance(self._java_obj, str)
- return ConnectDataFrame(
- TransformerRelation(
- child=dataset._plan,
- name=self._java_obj,
- ml_params=params,
- uid=self.uid,
- is_model=False,
- ),
- session,
- )
- else:
- raise RuntimeError(f"Unsupported {self}")
- else:
- return f(self, dataset)
-
- return cast(FuncT, wrapped)
-
-
-def try_remote_call(f: FuncT) -> FuncT:
- """Mark the function/property for the remote call.
- Eg, model.coefficients"""
-
- @functools.wraps(f)
- def wrapped(self: "JavaWrapper", name: str, *args: Any) -> Any:
- if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
- """Launch a remote call if possible"""
- from pyspark.sql.connect.session import SparkSession
-
- session = SparkSession.getActiveSession()
- assert session is not None
- assert isinstance(self._java_obj, str)
- methods, obj_ref = _extract_id_methods(self._java_obj)
- methods.append(pb2.Fetch.Method(method=name, args=serialize(session.client, *args)))
- command = pb2.Command()
- command.ml_command.fetch.CopyFrom(
- pb2.Fetch(obj_ref=pb2.ObjectRef(id=obj_ref), methods=methods)
- )
- (_, properties, _) = session.client.execute_command(command)
- ml_command_result = properties["ml_command_result"]
- if ml_command_result.HasField("summary"):
- summary = ml_command_result.summary
- session.client.add_ml_cache(summary)
- return summary
- else:
- return deserialize(properties)
- else:
- return f(self, name, *args)
-
- return cast(FuncT, wrapped)
-
-
-def try_remote_del(f: FuncT) -> FuncT:
- """Mark the function/property to delete a model on the server side."""
-
- @functools.wraps(f)
- def wrapped(self: "JavaWrapper") -> Any:
- try:
- in_remote = is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ
- except Exception:
- return
-
- if in_remote:
- # Delete the model if possible
- model_id = self._java_obj
- if model_id is not None and "." not in model_id:
- try:
- from pyspark.sql.connect.session import SparkSession
-
- session = SparkSession.getActiveSession()
- if session is not None:
- session.client.remove_ml_cache(model_id)
- return
- except Exception:
- # SparkSession's down.
- return
- else:
- return f(self)
-
- return cast(FuncT, wrapped)
-
-
-def try_remote_return_java_class(f: FuncT) -> FuncT:
- """Mark the function/property that returns none."""
-
- @functools.wraps(f)
- def wrapped(java_class: str, *args: Any) -> Any:
- if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
- return java_class
- else:
- return f(java_class, *args)
-
- return cast(FuncT, wrapped)
-
-
-def try_remote_write(f: FuncT) -> FuncT:
- """Mark the function that write an estimator/model or evaluator"""
-
- @functools.wraps(f)
- def wrapped(self: "JavaMLWritable") -> Any:
- if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
- from pyspark.ml.remote.readwrite import RemoteMLWriter
-
- return RemoteMLWriter(self)
- else:
- return f(self)
-
- return cast(FuncT, wrapped)
-
-
-def try_remote_read(f: FuncT) -> FuncT:
- """Mark the function to read an estimator/model or evaluator"""
-
- @functools.wraps(f)
- def wrapped(cls: Type["JavaMLReadable"]) -> Any:
- if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
- from pyspark.ml.remote.readwrite import RemoteMLReader
-
- return RemoteMLReader(cls)
- else:
- return f(cls)
-
- return cast(FuncT, wrapped)
-
-
-def try_remote_intercept(f: FuncT) -> FuncT:
- """Mark the function/property that returns none."""
-
- @functools.wraps(f)
- def wrapped(java_class: str, *args: Any) -> Any:
- if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
- return None
- else:
- return f(java_class, *args)
-
- return cast(FuncT, wrapped)
-
-
-def try_remote_not_supporting(f: FuncT) -> FuncT:
- """Mark the function/property that has not been supported yet"""
-
- @functools.wraps(f)
- def wrapped(*args: Any) -> Any:
- if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
- raise NotImplementedError("")
- else:
- return f(*args)
-
- return cast(FuncT, wrapped)
diff --git a/python/pyspark/ml/util.py b/python/pyspark/ml/util.py
index 3fe97f44619c..d161db04fad2 100644
--- a/python/pyspark/ml/util.py
+++ b/python/pyspark/ml/util.py
@@ -37,7 +37,6 @@ from typing import (
from pyspark import since
from pyspark.ml.common import inherit_doc
-from pyspark.ml.remote.util import try_remote_intermediate_result, try_remote_write, try_remote_read
from pyspark.sql import SparkSession
from pyspark.sql.utils import is_remote
from pyspark.util import VersionUtils
@@ -48,6 +47,8 @@ if TYPE_CHECKING:
from pyspark.ml.base import Params
from pyspark.ml.wrapper import JavaWrapper
from pyspark.core.context import SparkContext
+ from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame
+ from pyspark.ml.wrapper import JavaWrapper, JavaEstimator
T = TypeVar("T")
RW = TypeVar("RW", bound="BaseReadWrite")
@@ -59,6 +60,270 @@ JR = TypeVar("JR", bound="JavaMLReader")
FuncT = TypeVar("FuncT", bound=Callable[..., Any])
+def try_remote_intermediate_result(f: FuncT) -> FuncT:
+ """Mark the function/property that returns the intermediate result of the remote call.
+ Eg, model.summary"""
+
+ @functools.wraps(f)
+ def wrapped(self: "JavaWrapper") -> Any:
+ if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+ return f"{self._java_obj}.{f.__name__}"
+ else:
+ return f(self)
+
+ return cast(FuncT, wrapped)
+
+
+def try_remote_attribute_relation(f: FuncT) -> FuncT:
+ """Mark the function/property that returns a Relation.
+ Eg, model.summary.roc"""
+
+ @functools.wraps(f)
+ def wrapped(self: "JavaWrapper", *args: Any, **kwargs: Any) -> Any:
+ if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+ import pyspark.sql.connect.proto as pb2
+ from pyspark.ml.connect.util import _extract_id_methods
+ from pyspark.ml.connect.serialize import serialize
+
+ # The attribute returns a dataframe, we need to wrap it
+ # in the AttributeRelation
+ from pyspark.ml.connect.proto import AttributeRelation
+ from pyspark.sql.connect.session import SparkSession
+ from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame
+
+ session = SparkSession.getActiveSession()
+ assert session is not None
+
+ assert isinstance(self._java_obj, str)
+
+ methods, obj_ref = _extract_id_methods(self._java_obj)
+ methods.append(
+ pb2.Fetch.Method(method=f.__name__, args=serialize(session.client, *args))
+ )
+ plan = AttributeRelation(obj_ref, methods)
+ return ConnectDataFrame(plan, session)
+ else:
+ return f(self, *args, **kwargs)
+
+ return cast(FuncT, wrapped)
+
+
+def try_remote_fit(f: FuncT) -> FuncT:
+ """Mark the function that fits a model."""
+
+ @functools.wraps(f)
+ def wrapped(self: "JavaEstimator", dataset: "ConnectDataFrame") -> Any:
+ if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+ import pyspark.sql.connect.proto as pb2
+ from pyspark.ml.connect.serialize import serialize_ml_params, deserialize
+
+ client = dataset.sparkSession.client
+ input = dataset._plan.plan(client)
+ assert isinstance(self._java_obj, str)
+ estimator = pb2.MlOperator(
+ name=self._java_obj, uid=self.uid, type=pb2.MlOperator.ESTIMATOR
+ )
+ command = pb2.Command()
+ command.ml_command.fit.CopyFrom(
+ pb2.MlCommand.Fit(
+ estimator=estimator,
+ params=serialize_ml_params(self, client),
+ dataset=input,
+ )
+ )
+ (_, properties, _) = client.execute_command(command)
+ model_info = deserialize(properties)
+ client.add_ml_cache(model_info.obj_ref.id)
+ return model_info.obj_ref.id
+ else:
+ return f(self, dataset)
+
+ return cast(FuncT, wrapped)
+
+
+def try_remote_transform_relation(f: FuncT) -> FuncT:
+ """Mark the function/property that returns a relation for model transform."""
+
+ @functools.wraps(f)
+ def wrapped(self: "JavaWrapper", dataset: "ConnectDataFrame") -> Any:
+ if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+ from pyspark.ml import Model, Transformer
+ from pyspark.sql.connect.session import SparkSession
+ from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame
+ from pyspark.ml.connect.serialize import serialize_ml_params
+
+ session = SparkSession.getActiveSession()
+ assert session is not None
+ # Model is also a Transformer, so we much match Model first
+ if isinstance(self, Model):
+ params = serialize_ml_params(self, session.client)
+ from pyspark.ml.connect.proto import TransformerRelation
+
+ assert isinstance(self._java_obj, str)
+ return ConnectDataFrame(
+ TransformerRelation(
+ child=dataset._plan, name=self._java_obj, ml_params=params, is_model=True
+ ),
+ session,
+ )
+ elif isinstance(self, Transformer):
+ params = serialize_ml_params(self, session.client)
+ from pyspark.ml.connect.proto import TransformerRelation
+
+ assert isinstance(self._java_obj, str)
+ return ConnectDataFrame(
+ TransformerRelation(
+ child=dataset._plan,
+ name=self._java_obj,
+ ml_params=params,
+ uid=self.uid,
+ is_model=False,
+ ),
+ session,
+ )
+ else:
+ raise RuntimeError(f"Unsupported {self}")
+ else:
+ return f(self, dataset)
+
+ return cast(FuncT, wrapped)
+
+
+def try_remote_call(f: FuncT) -> FuncT:
+ """Mark the function/property for the remote call.
+ Eg, model.coefficients"""
+
+ @functools.wraps(f)
+ def wrapped(self: "JavaWrapper", name: str, *args: Any) -> Any:
+ if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+ # Launch a remote call if possible
+ import pyspark.sql.connect.proto as pb2
+ from pyspark.sql.connect.session import SparkSession
+ from pyspark.ml.connect.util import _extract_id_methods
+ from pyspark.ml.connect.serialize import serialize, deserialize
+
+ session = SparkSession.getActiveSession()
+ assert session is not None
+ assert isinstance(self._java_obj, str)
+ methods, obj_ref = _extract_id_methods(self._java_obj)
+ methods.append(pb2.Fetch.Method(method=name, args=serialize(session.client, *args)))
+ command = pb2.Command()
+ command.ml_command.fetch.CopyFrom(
+ pb2.Fetch(obj_ref=pb2.ObjectRef(id=obj_ref), methods=methods)
+ )
+ (_, properties, _) = session.client.execute_command(command)
+ ml_command_result = properties["ml_command_result"]
+ if ml_command_result.HasField("summary"):
+ summary = ml_command_result.summary
+ session.client.add_ml_cache(summary)
+ return summary
+ else:
+ return deserialize(properties)
+ else:
+ return f(self, name, *args)
+
+ return cast(FuncT, wrapped)
+
+
+def try_remote_del(f: FuncT) -> FuncT:
+ """Mark the function/property to delete a model on the server side."""
+
+ @functools.wraps(f)
+ def wrapped(self: "JavaWrapper") -> Any:
+ try:
+ in_remote = is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ
+ except Exception:
+ return
+
+ if in_remote:
+ # Delete the model if possible
+ model_id = self._java_obj
+ if model_id is not None and "." not in model_id:
+ try:
+ from pyspark.sql.connect.session import SparkSession
+
+ session = SparkSession.getActiveSession()
+ if session is not None:
+ session.client.remove_ml_cache(model_id)
+ return
+ except Exception:
+ # SparkSession's down.
+ return
+ else:
+ return f(self)
+
+ return cast(FuncT, wrapped)
+
+
+def try_remote_return_java_class(f: FuncT) -> FuncT:
+ """Mark the function/property that returns none."""
+
+ @functools.wraps(f)
+ def wrapped(java_class: str, *args: Any) -> Any:
+ if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+ return java_class
+ else:
+ return f(java_class, *args)
+
+ return cast(FuncT, wrapped)
+
+
+def try_remote_write(f: FuncT) -> FuncT:
+ """Mark the function that write an estimator/model or evaluator"""
+
+ @functools.wraps(f)
+ def wrapped(self: "JavaMLWritable") -> Any:
+ if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+ from pyspark.ml.connect.readwrite import RemoteMLWriter
+
+ return RemoteMLWriter(self)
+ else:
+ return f(self)
+
+ return cast(FuncT, wrapped)
+
+
+def try_remote_read(f: FuncT) -> FuncT:
+ """Mark the function to read an estimator/model or evaluator"""
+
+ @functools.wraps(f)
+ def wrapped(cls: Type["JavaMLReadable"]) -> Any:
+ if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+ from pyspark.ml.connect.readwrite import RemoteMLReader
+
+ return RemoteMLReader(cls)
+ else:
+ return f(cls)
+
+ return cast(FuncT, wrapped)
+
+
+def try_remote_intercept(f: FuncT) -> FuncT:
+ """Mark the function/property that returns none."""
+
+ @functools.wraps(f)
+ def wrapped(java_class: str, *args: Any) -> Any:
+ if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+ return None
+ else:
+ return f(java_class, *args)
+
+ return cast(FuncT, wrapped)
+
+
+def try_remote_not_supporting(f: FuncT) -> FuncT:
+ """Mark the function/property that has not been supported yet"""
+
+ @functools.wraps(f)
+ def wrapped(*args: Any) -> Any:
+ if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+ raise NotImplementedError("")
+ else:
+ return f(*args)
+
+ return cast(FuncT, wrapped)
+
+
def _jvm() -> "JavaGateway":
"""
Returns the JVM view associated with SparkContext. Must be called
diff --git a/python/pyspark/ml/wrapper.py b/python/pyspark/ml/wrapper.py
index e2bf25386c77..7b5c4f783a65 100644
--- a/python/pyspark/ml/wrapper.py
+++ b/python/pyspark/ml/wrapper.py
@@ -19,7 +19,7 @@ from abc import ABCMeta, abstractmethod
from typing import Any, Generic, Optional, List, Type, TypeVar, TYPE_CHECKING
from pyspark import since
-from pyspark.ml.remote.util import (
+from pyspark.ml.util import (
try_remote_transform_relation,
try_remote_call,
try_remote_fit,
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org