[SPARK-51101][ML][PYTHON][CONNECT][TESTS] Add doctest for pyspark.ml.connect.functions

dev/sparktestsupport/modules.py

@@ -1110,6 +1110,8 @@ def __hash__(self):
         "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",

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()