[SPARK-34972][PYTHON] Make pandas-on-Spark doctests work.

.github/workflows/build_and_test.yml

@@ -161,6 +161,8 @@ jobs:
             pyspark-sql, pyspark-mllib, pyspark-resource
           - >-
             pyspark-core, pyspark-streaming, pyspark-ml
+          - >-
+            pyspark-pandas
     env:
       MODULES_TO_TEST: ${{ matrix.modules }}
       HADOOP_PROFILE: hadoop3.2

dev/run-tests.py

@@ -123,17 +123,17 @@ def determine_modules_to_test(changed_modules, deduplicated=True):
     >>> [x.name for x in determine_modules_to_test([modules.sql])]
     ... # doctest: +NORMALIZE_WHITESPACE
     ['sql', 'avro', 'hive', 'mllib', 'sql-kafka-0-10', 'examples', 'hive-thriftserver',
-     'pyspark-sql', 'repl', 'sparkr', 'pyspark-mllib', 'pyspark-ml']
+     'pyspark-sql', 'repl', 'sparkr', 'pyspark-mllib', 'pyspark-pandas', 'pyspark-ml']
     >>> sorted([x.name for x in determine_modules_to_test(
     ...     [modules.sparkr, modules.sql], deduplicated=False)])
     ... # doctest: +NORMALIZE_WHITESPACE
     ['avro', 'examples', 'hive', 'hive-thriftserver', 'mllib', 'pyspark-ml',
-     'pyspark-mllib', 'pyspark-sql', 'repl', 'sparkr', 'sql', 'sql-kafka-0-10']
+     'pyspark-mllib', 'pyspark-pandas', 'pyspark-sql', 'repl', 'sparkr', 'sql', 'sql-kafka-0-10']
     >>> sorted([x.name for x in determine_modules_to_test(
     ...     [modules.sql, modules.core], deduplicated=False)])
     ... # doctest: +NORMALIZE_WHITESPACE
     ['avro', 'catalyst', 'core', 'examples', 'graphx', 'hive', 'hive-thriftserver',
-     'mllib', 'mllib-local', 'pyspark-core', 'pyspark-ml', 'pyspark-mllib',
+     'mllib', 'mllib-local', 'pyspark-core', 'pyspark-ml', 'pyspark-mllib', 'pyspark-pandas',
      'pyspark-resource', 'pyspark-sql', 'pyspark-streaming', 'repl', 'root',
      'sparkr', 'sql', 'sql-kafka-0-10', 'streaming', 'streaming-kafka-0-10',
      'streaming-kinesis-asl']

dev/sparktestsupport/modules.py

@@ -573,6 +573,50 @@ def __hash__(self):
     ]
 )
 
+pyspark_pandas = Module(
+    name="pyspark-pandas",
+    dependencies=[pyspark_core, pyspark_sql],
+    source_file_regexes=[
+        "python/pyspark/pandas/"
+    ],
+    python_test_goals=[
+        # doctests
+        "pyspark.pandas.accessors",
+        "pyspark.pandas.base",
+        "pyspark.pandas.categorical",
+        "pyspark.pandas.config",
+        "pyspark.pandas.datetimes",
+        "pyspark.pandas.exceptions",
+        "pyspark.pandas.extensions",
+        "pyspark.pandas.frame",
+        "pyspark.pandas.generic",
+        "pyspark.pandas.groupby",
+        "pyspark.pandas.indexing",
+        "pyspark.pandas.internal",
+        "pyspark.pandas.ml",
+        "pyspark.pandas.mlflow",
+        "pyspark.pandas.namespace",
+        "pyspark.pandas.numpy_compat",
+        "pyspark.pandas.series",
+        "pyspark.pandas.sql_processor",
+        "pyspark.pandas.strings",
+        "pyspark.pandas.utils",
+        "pyspark.pandas.window",
+        "pyspark.pandas.indexes.base",
+        "pyspark.pandas.indexes.category",
+        "pyspark.pandas.indexes.datetimes",
+        "pyspark.pandas.indexes.multi",
+        "pyspark.pandas.indexes.numeric",
+        "pyspark.pandas.spark.accessors",
+        "pyspark.pandas.spark.utils",
+        "pyspark.pandas.typedef.typehints",
+    ],
+    excluded_python_implementations=[
+        "PyPy"  # Skip these tests under PyPy since they require numpy, pandas, and pyarrow and
+                # they aren't available there
+    ]
+)
+
 sparkr = Module(
     name="sparkr",
     dependencies=[hive, mllib],

python/pyspark/pandas/__init__.py

@@ -17,13 +17,24 @@
 import os
 import sys
 from distutils.version import LooseVersion
+import warnings
+
+from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version
+
+try:
+    require_minimum_pandas_version()
+    require_minimum_pyarrow_version()
+except ImportError as e:
+    if os.environ.get("SPARK_TESTING"):
+        warnings.warn(str(e))
+        sys.exit()
+    else:
+        raise
 
 from pyspark.pandas.version import __version__  # noqa: F401
 
 
 def assert_python_version():
-    import warnings
-
     major = 3
     minor = 5
     deprecated_version = (major, minor)
@@ -206,4 +217,4 @@ def _auto_patch_pandas():
 # Import after the usage logger is attached.
 from pyspark.pandas.config import get_option, options, option_context, reset_option, set_option
 from pyspark.pandas.namespace import *  # F405
-from pyspark.pandas.sql import sql
+from pyspark.pandas.sql_processor import sql

python/pyspark/pandas/accessors.py

@@ -928,3 +928,33 @@ def apply_func(pdf):
             ),
             dtype=dtype,
         )
+
+
+def _test():
+    import os
+    import doctest
+    import sys
+    from pyspark.sql import SparkSession
+    import pyspark.pandas.accessors
+
+    os.chdir(os.environ["SPARK_HOME"])
+
+    globs = pyspark.pandas.accessors.__dict__.copy()
+    globs["pp"] = pyspark.pandas
+    spark = (
+        SparkSession.builder.master("local[4]")
+        .appName("pyspark.pandas.accessors tests")
+        .getOrCreate()
+    )
+    (failure_count, test_count) = doctest.testmod(
+        pyspark.pandas.accessors,
+        globs=globs,
+        optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE,
+    )
+    spark.stop()
+    if failure_count:
+        sys.exit(-1)
+
+
+if __name__ == "__main__":
+    _test()

python/pyspark/pandas/base.py

@@ -1993,3 +1993,31 @@ def factorize(
         uniques = pd.Index(uniques_list)
 
         return codes, uniques
+
+
+def _test():
+    import os
+    import doctest
+    import sys
+    from pyspark.sql import SparkSession
+    import pyspark.pandas.base
+
+    os.chdir(os.environ["SPARK_HOME"])
+
+    globs = pyspark.pandas.base.__dict__.copy()
+    globs["pp"] = pyspark.pandas
+    spark = (
+        SparkSession.builder.master("local[4]").appName("pyspark.pandas.base tests").getOrCreate()
+    )
+    (failure_count, test_count) = doctest.testmod(
+        pyspark.pandas.base,
+        globs=globs,
+        optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE,
+    )
+    spark.stop()
+    if failure_count:
+        sys.exit(-1)
+
+
+if __name__ == "__main__":
+    _test()

python/pyspark/pandas/categorical.py

@@ -162,3 +162,33 @@ def set_categories(
         self, new_categories, ordered: bool = None, rename: bool = False, inplace: bool = False
     ):
         raise NotImplementedError()
+
+
+def _test():
+    import os
+    import doctest
+    import sys
+    from pyspark.sql import SparkSession
+    import pyspark.pandas.categorical
+
+    os.chdir(os.environ["SPARK_HOME"])
+
+    globs = pyspark.pandas.categorical.__dict__.copy()
+    globs["pp"] = pyspark.pandas
+    spark = (
+        SparkSession.builder.master("local[4]")
+        .appName("pyspark.pandas.categorical tests")
+        .getOrCreate()
+    )
+    (failure_count, test_count) = doctest.testmod(
+        pyspark.pandas.categorical,
+        globs=globs,
+        optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE,
+    )
+    spark.stop()
+    if failure_count:
+        sys.exit(-1)
+
+
+if __name__ == "__main__":
+    _test()

python/pyspark/pandas/config.py

@@ -440,3 +440,31 @@ def __dir__(self):
 
 
 options = DictWrapper(_options_dict)
+
+
+def _test():
+    import os
+    import doctest
+    import sys
+    from pyspark.sql import SparkSession
+    import pyspark.pandas.config
+
+    os.chdir(os.environ["SPARK_HOME"])
+
+    globs = pyspark.pandas.config.__dict__.copy()
+    globs["pp"] = pyspark.pandas
+    spark = (
+        SparkSession.builder.master("local[4]").appName("pyspark.pandas.config tests").getOrCreate()
+    )
+    (failure_count, test_count) = doctest.testmod(
+        pyspark.pandas.config,
+        globs=globs,
+        optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE,
+    )
+    spark.stop()
+    if failure_count:
+        sys.exit(-1)
+
+
+if __name__ == "__main__":
+    _test()

python/pyspark/pandas/datetimes.py

@@ -848,3 +848,33 @@ def pandas_day_name(s) -> "pp.Series[str]":
             return s.dt.day_name(locale=locale)
 
         return self._data.koalas.transform_batch(pandas_day_name)
+
+
+def _test():
+    import os
+    import doctest
+    import sys
+    from pyspark.sql import SparkSession
+    import pyspark.pandas.datetimes
+
+    os.chdir(os.environ["SPARK_HOME"])
+
+    globs = pyspark.pandas.datetimes.__dict__.copy()
+    globs["pp"] = pyspark.pandas
+    spark = (
+        SparkSession.builder.master("local[4]")
+        .appName("pyspark.pandas.datetimes tests")
+        .getOrCreate()
+    )
+    (failure_count, test_count) = doctest.testmod(
+        pyspark.pandas.datetimes,
+        globs=globs,
+        optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE,
+    )
+    spark.stop()
+    if failure_count:
+        sys.exit(-1)
+
+
+if __name__ == "__main__":
+    _test()

python/pyspark/pandas/exceptions.py

@@ -104,3 +104,33 @@ def __init__(
                     class_name, property_name, reason
                 )
         super().__init__(msg)
+
+
+def _test():
+    import os
+    import doctest
+    import sys
+    from pyspark.sql import SparkSession
+    import pyspark.pandas.exceptions
+
+    os.chdir(os.environ["SPARK_HOME"])
+
+    globs = pyspark.pandas.exceptions.__dict__.copy()
+    globs["pp"] = pyspark.pandas
+    spark = (
+        SparkSession.builder.master("local[4]")
+        .appName("pyspark.pandas.exceptions tests")
+        .getOrCreate()
+    )
+    (failure_count, test_count) = doctest.testmod(
+        pyspark.pandas.exceptions,
+        globs=globs,
+        optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE,
+    )
+    spark.stop()
+    if failure_count:
+        sys.exit(-1)
+
+
+if __name__ == "__main__":
+    _test()

python/pyspark/pandas/extensions.py

@@ -340,3 +340,35 @@ def bar(self):
     from pyspark.pandas import Index
 
     return _register_accessor(name, Index)
+
+
+def _test():
+    import os
+    import doctest
+    import sys
+    import numpy
+    from pyspark.sql import SparkSession
+    import pyspark.pandas.extensions
+
+    os.chdir(os.environ["SPARK_HOME"])
+
+    globs = pyspark.pandas.extensions.__dict__.copy()
+    globs["np"] = numpy
+    globs["pp"] = pyspark.pandas
+    spark = (
+        SparkSession.builder.master("local[4]")
+        .appName("pyspark.pandas.extensions tests")
+        .getOrCreate()
+    )
+    (failure_count, test_count) = doctest.testmod(
+        pyspark.pandas.extensions,
+        globs=globs,
+        optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE,
+    )
+    spark.stop()
+    if failure_count:
+        sys.exit(-1)
+
+
+if __name__ == "__main__":
+    _test()