(spark) branch branch-4.0 updated: [SPARK-51101][ML][PYTHON][CONNECT][TESTS] Add doctest for `pyspark.ml.connect.functions`

[ru...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

ruifengz 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 f923f1525868 [SPARK-51101][ML][PYTHON][CONNECT][TESTS] Add doctest for `pyspark.ml.connect.functions`
f923f1525868 is described below

commit f923f1525868014bc377e779148d29ab653b61d6
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Fri Feb 7 10:27:35 2025 +0800

    [SPARK-51101][ML][PYTHON][CONNECT][TESTS] Add doctest for `pyspark.ml.connect.functions`
    
    ### What changes were proposed in this pull request?
    Add doctest for `pyspark.ml.connect.functions`
    
    ### Why are the changes needed?
    to improve the test coverage
    
    ### Does this PR introduce _any_ user-facing change?
    no, test-only
    
    ### How was this patch tested?
    ci
    
    ### Was this patch authored or co-authored using generative AI tooling?
    no
    
    Closes #49821 from zhengruifeng/ml_connect_function_doc_tests.
    
    Authored-by: Ruifeng Zheng <ru...@apache.org>
    Signed-off-by: Ruifeng Zheng <ru...@apache.org>
    (cherry picked from commit 935c2b0022ea55b09b53bae4fb45fc0ee5ff4a05)
    Signed-off-by: Ruifeng Zheng <ru...@apache.org>
---
 dev/sparktestsupport/modules.py        |  2 ++
 python/pyspark/ml/connect/functions.py | 59 ++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)

diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index 212cd1e10e8a..f1088b212f80 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -1110,6 +1110,8 @@ pyspark_ml_connect = Module(
         "python/pyspark/ml/connect",
     ],
     python_test_goals=[
+        # ml doctests
+        "pyspark.ml.connect.functions",
         # ml unittests
         "pyspark.ml.tests.connect.test_connect_function",
         "pyspark.ml.tests.connect.test_parity_torch_distributor",
diff --git a/python/pyspark/ml/connect/functions.py b/python/pyspark/ml/connect/functions.py
index 6597e6c4118a..e3664db87ae6 100644
--- a/python/pyspark/ml/connect/functions.py
+++ b/python/pyspark/ml/connect/functions.py
@@ -14,10 +14,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+
+from typing import Any, TYPE_CHECKING
+
 from pyspark.ml import functions as PyMLFunctions
 from pyspark.sql.column import Column
 from pyspark.sql.connect.functions.builtin import _invoke_function, _to_col, lit
 
+if TYPE_CHECKING:
+    from pyspark.sql._typing import UserDefinedFunctionLike
+
 
 def vector_to_array(col: Column, dtype: str = "float64") -> Column:
     return _invoke_function("vector_to_array", _to_col(col), lit(dtype))
@@ -31,3 +37,56 @@ def array_to_vector(col: Column) -> Column:
 
 
 array_to_vector.__doc__ = PyMLFunctions.array_to_vector.__doc__
+
+
+def predict_batch_udf(*args: Any, **kwargs: Any) -> "UserDefinedFunctionLike":
+    return PyMLFunctions.predict_batch_udf(*args, **kwargs)
+
+
+predict_batch_udf.__doc__ = PyMLFunctions.predict_batch_udf.__doc__
+
+
+def _test() -> None:
+    import os
+    import sys
+    import doctest
+    from pyspark.sql import SparkSession as PySparkSession
+    import pyspark.ml.connect.functions
+
+    from pyspark.sql.pandas.utils import (
+        require_minimum_pandas_version,
+        require_minimum_pyarrow_version,
+    )
+
+    try:
+        require_minimum_pandas_version()
+        require_minimum_pyarrow_version()
+    except Exception as e:
+        print(
+            f"Skipping pyspark.ml.functions doctests: {e}",
+            file=sys.stderr,
+        )
+        sys.exit(0)
+
+    globs = pyspark.ml.connect.functions.__dict__.copy()
+
+    globs["spark"] = (
+        PySparkSession.builder.appName("ml.connect.functions tests")
+        .remote(os.environ.get("SPARK_CONNECT_TESTING_REMOTE", "local[4]"))
+        .getOrCreate()
+    )
+
+    (failure_count, test_count) = doctest.testmod(
+        pyspark.ml.connect.functions,
+        globs=globs,
+        optionflags=doctest.ELLIPSIS
+        | doctest.NORMALIZE_WHITESPACE
+        | doctest.IGNORE_EXCEPTION_DETAIL,
+    )
+    globs["spark"].stop()
+    if failure_count:
+        sys.exit(-1)
+
+
+if __name__ == "__main__":
+    _test()


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org