| # |
| # 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. |
| # |
| |
| """ |
| Generate 'Supported pandas APIs' documentation file |
| """ |
| import warnings |
| from enum import Enum, unique |
| from inspect import getmembers, isclass, isfunction, signature |
| from typing import Any, Dict, List, NamedTuple, Set, TextIO, Tuple |
| from types import FunctionType |
| |
| import pyspark.pandas as ps |
| import pyspark.pandas.groupby as psg |
| import pyspark.pandas.window as psw |
| import pandas as pd |
| import pandas.core.groupby as pdg |
| import pandas.core.window as pdw |
| |
| from pyspark.loose_version import LooseVersion |
| from pyspark.pandas.exceptions import PandasNotImplementedError |
| |
| # Constants |
| MAX_MISSING_PARAMS_SIZE = 5 |
| COMMON_PARAMETER_SET = {"kwargs", "args", "cls"} |
| MODULE_GROUP_MATCH = [(pd, ps), (pdw, psw), (pdg, psg)] |
| PANDAS_LATEST_VERSION = "2.3.2" |
| |
| RST_HEADER = """ |
| ===================== |
| Supported pandas API |
| ===================== |
| |
| .. currentmodule:: pyspark.pandas |
| |
| The following table shows the pandas APIs that implemented or non-implemented from pandas API on |
| Spark. Some pandas API do not implement full parameters, so the third column shows missing |
| parameters for each API. |
| |
| * 'Y' in the second column means it's implemented including its whole parameter. |
| * 'N' means it's not implemented yet. |
| * 'P' means it's partially implemented with the missing of some parameters. |
| |
| All API in the list below computes the data with distributed execution except the ones that require |
| the local execution by design. For example, `DataFrame.to_numpy() <https://spark.apache.org/docs/ |
| latest/api/python/reference/pyspark.pandas/api/pyspark.pandas.DataFrame.to_numpy.html>`__ |
| requires to collect the data to the driver side. |
| |
| If there is non-implemented pandas API or parameter you want, you can create an `Apache Spark |
| JIRA <https://issues.apache.org/jira/projects/SPARK/summary>`__ to request or to contribute by |
| your own. |
| |
| The API list is updated based on the `latest pandas official API reference |
| <https://pandas.pydata.org/docs/reference/index.html#>`__. |
| |
| """ |
| |
| |
| @unique |
| class Implemented(Enum): |
| """ |
| Enumeration of implementation statuses. |
| """ |
| |
| IMPLEMENTED = "Y" |
| NOT_IMPLEMENTED = "N" |
| PARTIALLY_IMPLEMENTED = "P" |
| |
| |
| class SupportedStatus(NamedTuple): |
| """ |
| Defines a supported status for specific pandas API. |
| """ |
| |
| implemented: str |
| missing: str |
| |
| |
| def generate_supported_api(output_rst_file_path: str) -> None: |
| """ |
| Generate the supported APIs status dictionary and write it to an RST file. |
| |
| Parameters |
| ---------- |
| output_rst_file_path : str |
| The path to the document file in RST format. |
| """ |
| _check_pandas_version() |
| all_supported_status = _collect_supported_status() |
| _write_rst(output_rst_file_path, all_supported_status) |
| |
| |
| def _check_pandas_version() -> None: |
| """ |
| Check if the installed pandas version matches the expected version. |
| """ |
| # Work around pandas version string issue, |
| # see https://github.com/pandas-dev/pandas/issues/61579. |
| if LooseVersion(pd.__version__.split("+")[0]) != LooseVersion(PANDAS_LATEST_VERSION): |
| msg = ( |
| f"Warning: pandas {PANDAS_LATEST_VERSION} is required; your version is {pd.__version__}" |
| ) |
| warnings.warn(msg, UserWarning) |
| raise ImportError(msg) |
| |
| |
| def _collect_supported_status() -> Dict[Tuple[str, str], Dict[str, SupportedStatus]]: |
| """ |
| Collect the supported status across multiple module paths. |
| """ |
| all_supported_status: Dict[Tuple[str, str], Dict[str, SupportedStatus]] = {} |
| for pd_module_group, ps_module_group in MODULE_GROUP_MATCH: |
| pd_modules = _get_pd_modules(pd_module_group) |
| _update_all_supported_status( |
| all_supported_status, pd_modules, pd_module_group, ps_module_group |
| ) |
| return all_supported_status |
| |
| |
| def _get_pd_modules(pd_module_group: Any) -> List[str]: |
| """ |
| Get sorted list of pandas member names from a pandas module. |
| |
| Parameters |
| ---------- |
| pd_module_group : Any |
| Importable pandas module. |
| |
| Returns |
| ------- |
| List[str] |
| Sorted list of member names. |
| """ |
| return sorted(m[0] for m in getmembers(pd_module_group, isclass) if not m[0].startswith("_")) |
| |
| |
| def _update_all_supported_status( |
| all_supported_status: Dict[Tuple[str, str], Dict[str, SupportedStatus]], |
| pd_modules: List[str], |
| pd_module_group: Any, |
| ps_module_group: Any, |
| ) -> None: |
| """ |
| Update the supported status dictionary with status from multiple modules. |
| |
| Parameters |
| ---------- |
| all_supported_status : Dict[Tuple[str, str], Dict[str, SupportedStatus]] |
| The dictionary to update with supported statuses. |
| pd_modules : List[str] |
| List of module names in pandas. |
| pd_module_group : Any |
| Importable pandas module group. |
| ps_module_group : Any |
| Corresponding pyspark.pandas module group. |
| """ |
| pd_modules.append("") # Include General Function APIs |
| for module_name in pd_modules: |
| supported_status = _create_supported_by_module( |
| module_name, pd_module_group, ps_module_group |
| ) |
| if supported_status: |
| all_supported_status[(module_name, ps_module_group.__name__)] = supported_status |
| |
| |
| def _create_supported_by_module( |
| module_name: str, pd_module_group: Any, ps_module_group: Any |
| ) -> Dict[str, SupportedStatus]: |
| """ |
| Create a dictionary of supported status for a specific pandas module. |
| |
| Parameters |
| ---------- |
| module_name : str |
| Name of the module in pandas. |
| pd_module_group : Any |
| Importable pandas module. |
| ps_module_group : Any |
| Corresponding pyspark.pandas module. |
| |
| Returns |
| ------- |
| Dict[str, SupportedStatus] |
| Dictionary of supported status for the module. |
| """ |
| pd_module = getattr(pd_module_group, module_name) if module_name else pd_module_group |
| try: |
| ps_module = getattr(ps_module_group, module_name) if module_name else ps_module_group |
| except (AttributeError, PandasNotImplementedError): |
| # module not implemented |
| return {} |
| |
| pd_funcs = dict([m for m in getmembers(pd_module, isfunction) if not m[0].startswith("_")]) |
| if not pd_funcs: |
| return {} |
| |
| ps_funcs = dict([m for m in getmembers(ps_module, isfunction) if not m[0].startswith("_")]) |
| |
| return _organize_by_implementation_status( |
| module_name, pd_funcs, ps_funcs, pd_module_group, ps_module_group |
| ) |
| |
| |
| def _organize_by_implementation_status( |
| module_name: str, |
| pd_funcs: Dict[str, FunctionType], |
| ps_funcs: Dict[str, FunctionType], |
| pd_module_group: Any, |
| ps_module_group: Any, |
| ) -> Dict[str, SupportedStatus]: |
| """ |
| Organize functions by implementation status between pandas and pyspark.pandas. |
| |
| Parameters |
| ---------- |
| module_name : str |
| Class name that exists in the path of the module. |
| pd_funcs: Dict[str, Callable] |
| function name and function object mapping of pandas module. |
| ps_funcs: Dict[str, Callable] |
| function name and function object mapping of pyspark.pandas module. |
| pd_module_group : Any |
| Specific path of importable pandas module. |
| ps_module_group: Any |
| Specific path of importable pyspark.pandas module. |
| |
| Returns |
| ------- |
| Dict[str, SupportedStatus] |
| Dictionary of implementation status. |
| """ |
| pd_dict = {} |
| for pd_func_name, pd_func in pd_funcs.items(): |
| ps_func = ps_funcs.get(pd_func_name) |
| if ps_func: |
| missing_set = ( |
| set(signature(pd_func).parameters) |
| - set(signature(ps_func).parameters) |
| - COMMON_PARAMETER_SET |
| ) |
| if missing_set: |
| # partially implemented |
| pd_dict[pd_func_name] = SupportedStatus( |
| implemented=Implemented.PARTIALLY_IMPLEMENTED.value, |
| missing=_transform_missing( |
| module_name, |
| pd_func_name, |
| missing_set, |
| pd_module_group.__name__, |
| ps_module_group.__name__, |
| ), |
| ) |
| else: |
| # implemented including it's whole parameter |
| pd_dict[pd_func_name] = SupportedStatus( |
| implemented=Implemented.IMPLEMENTED.value, missing="" |
| ) |
| else: |
| # not implemented yet |
| pd_dict[pd_func_name] = SupportedStatus( |
| implemented=Implemented.NOT_IMPLEMENTED.value, missing="" |
| ) |
| return pd_dict |
| |
| |
| def _transform_missing( |
| module_name: str, |
| pd_func_name: str, |
| missing_set: Set[str], |
| pd_module_path: str, |
| ps_module_path: str, |
| ) -> str: |
| """ |
| Transform missing parameters into a formatted string for table display. |
| |
| Parameters |
| ---------- |
| module_name : str |
| Class name that exists in the path of the module. |
| pd_func_name : str |
| Name of pandas API. |
| missing_set : Set[str] |
| A set of parameters not yet implemented. |
| pd_module_path : str |
| Path string of pandas module. |
| ps_module_path : str |
| Path string of pyspark.pandas module. |
| |
| Returns |
| ------- |
| str |
| Formatted string representing missing parameters. |
| |
| Examples |
| -------- |
| >>> _transform_missing("DataFrame", "add", {"axis", "fill_value", "level"}, |
| ... "pandas.DataFrame", "pyspark.pandas.DataFrame") |
| '``axis`` , ``fill_value`` , ``level``' |
| """ |
| missing_str = " , ".join("``%s``" % x for x in sorted(missing_set)[:MAX_MISSING_PARAMS_SIZE]) |
| if len(missing_set) > MAX_MISSING_PARAMS_SIZE: |
| module_dot_func = "%s.%s" % (module_name, pd_func_name) if module_name else pd_func_name |
| additional_str = ( |
| " and more. See the " |
| + "`%s.%s " % (pd_module_path, module_dot_func) |
| + "<https://pandas.pydata.org/docs/reference/api/" |
| + "%s.%s.html>`__ and " % (pd_module_path, module_dot_func) |
| + "`%s.%s " % (ps_module_path, module_dot_func) |
| + "<https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/api/" |
| + "%s.%s.html>`__ for detail." % (ps_module_path, module_dot_func) |
| ) |
| missing_str += additional_str |
| return missing_str |
| |
| |
| def _write_table( |
| module_name: str, |
| module_path: str, |
| supported_status: Dict[str, SupportedStatus], |
| w_fd: TextIO, |
| ) -> None: |
| """ |
| Write the support status in a table format using Sphinx list-table directive. |
| |
| Parameters |
| ---------- |
| module_name : str |
| The name of the module whose support status is being documented. |
| module_path : str |
| The import path of the module in the documentation. |
| supported_status : Dict[str, SupportedStatus] |
| A dictionary mapping each function name to its support status. |
| w_fd : TextIO |
| An open file descriptor where the table will be written. |
| """ |
| lines = [] |
| if module_name: |
| lines.append(module_name) |
| else: |
| lines.append("General Function") |
| lines.append(" API\n") |
| lines.append("-" * 100) |
| lines.append("\n") |
| lines.append(".. currentmodule:: %s" % module_path) |
| if module_name: |
| lines.append(".%s\n" % module_name) |
| else: |
| lines.append("\n") |
| lines.append("\n") |
| lines.append(".. list-table::\n") |
| lines.append(" :header-rows: 1\n") |
| lines.append("\n") |
| lines.append(" * - API\n") |
| lines.append(" - Implemented\n") |
| lines.append(" - Missing parameters\n") |
| for func_str, status in supported_status.items(): |
| func_str = _escape_func_str(func_str) |
| if status.implemented == Implemented.NOT_IMPLEMENTED.value: |
| lines.append(" * - %s\n" % func_str) |
| else: |
| lines.append(" * - :func:`%s`\n" % func_str) |
| lines.append(" - %s\n" % status.implemented) |
| lines.append(" - \n") if not status.missing else lines.append( |
| " - %s\n" % status.missing |
| ) |
| w_fd.writelines(lines) |
| |
| |
| def _escape_func_str(func_str: str) -> str: |
| """ |
| Escape function names to conform to RST format. |
| |
| Parameters |
| ---------- |
| func_str : str |
| Function name to escape. |
| |
| Returns |
| ------- |
| str |
| Escaped function name. |
| """ |
| # TODO: Take into account that this function can create links incorrectly |
| # We can create alias links or links to parent methods |
| if func_str.endswith("_"): |
| return func_str[:-1] + "\_" # noqa: W605 |
| else: |
| return func_str |
| |
| |
| def _write_rst( |
| output_rst_file_path: str, |
| all_supported_status: Dict[Tuple[str, str], Dict[str, SupportedStatus]], |
| ) -> None: |
| """ |
| Write the final RST file with the collected support status. |
| |
| Parameters |
| ---------- |
| output_rst_file_path : str |
| Path to the output RST file. |
| all_supported_status : Dict |
| Collected support status data. |
| """ |
| with open(output_rst_file_path, "w") as w_fd: |
| w_fd.write(RST_HEADER) |
| for module_info, supported_status in all_supported_status.items(): |
| module, module_path = module_info |
| if supported_status: |
| _write_table(module, module_path, supported_status, w_fd) |
| w_fd.write("\n") |
| |
| |
| def _test() -> None: |
| import doctest |
| import sys |
| |
| import pyspark.pandas.supported_api_gen |
| |
| globs = pyspark.pandas.supported_api_gen.__dict__.copy() |
| (failure_count, test_count) = doctest.testmod(pyspark.pandas.supported_api_gen, globs=globs) |
| if failure_count: |
| sys.exit(-1) |
| |
| |
| if __name__ == "__main__": |
| _test() |
