python/pyspark/pandas/__init__.py - master - githubmirror/spark

python/pyspark/pandas/__init__.py
145 строк4 KB
Devin Petersohn
[SPARK-40373][PS] Implement ps.show_versions
11 май 2026, 00:52
11 май 2026, 00:5235f02b1
#
# 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.
#

"""
.. versionadded:: 3.2.0
    pandas API on Spark
"""

import os
import sys
import warnings
from typing import Any

from pyspark.pandas.missing.general_functions import MissingPandasLikeGeneralFunctions
from pyspark.pandas.missing.scalars import MissingPandasLikeScalars
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.frame import DataFrame
from pyspark.pandas.indexes.base import Index
from pyspark.pandas.indexes.category import CategoricalIndex
from pyspark.pandas.indexes.datetimes import DatetimeIndex
from pyspark.pandas.indexes.multi import MultiIndex
from pyspark.pandas.indexes.timedelta import TimedeltaIndex
from pyspark.pandas.series import Series
from pyspark.pandas.groupby import NamedAgg

__all__ = [  # noqa: F405
    "read_csv",
    "read_parquet",
    "to_datetime",
    "date_range",
    "from_pandas",
    "get_dummies",
    "DataFrame",
    "Series",
    "Index",
    "MultiIndex",
    "CategoricalIndex",
    "DatetimeIndex",
    "TimedeltaIndex",
    "sql",
    "range",
    "concat",
    "melt",
    "get_option",
    "set_option",
    "reset_option",
    "read_sql_table",
    "read_sql_query",
    "read_sql",
    "options",
    "option_context",
    "NamedAgg",
    "show_versions",
]


def _auto_patch_spark() -> None:
    import os
    import logging

    # Attach a usage logger. 'KOALAS_USAGE_LOGGER' is legacy, and it's for compatibility.
    logger_module = os.getenv("PYSPARK_PANDAS_USAGE_LOGGER", os.getenv("KOALAS_USAGE_LOGGER", ""))
    if logger_module != "":
        try:
            from pyspark.pandas import usage_logging

            usage_logging.attach(logger_module)
        except Exception as e:
            logger = logging.getLogger("pyspark.pandas.usage_logger")
            logger.warning(
                "Tried to attach usage logger `{}`, but an exception was raised: {}".format(
                    logger_module, str(e)
                )
            )


_frame_has_class_getitem = False
_series_has_class_getitem = False


def _auto_patch_pandas() -> None:
    import pandas as pd

    # In order to use it in test cases.
    global _frame_has_class_getitem
    global _series_has_class_getitem

    _frame_has_class_getitem = hasattr(pd.DataFrame, "__class_getitem__")
    _series_has_class_getitem = hasattr(pd.Series, "__class_getitem__")

    # Just in case pandas implements '__class_getitem__' later.
    if not _frame_has_class_getitem:
        pd.DataFrame.__class_getitem__ = (  # type: ignore[attr-defined]
            lambda params: DataFrame.__class_getitem__(params)
        )

    if not _series_has_class_getitem:
        pd.Series.__class_getitem__ = (  # type: ignore[attr-defined]
            lambda params: Series.__class_getitem__(params)
        )


_auto_patch_spark()
_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 *  # noqa: F403
from pyspark.pandas.sql_formatter import sql


def __getattr__(key: str) -> Any:
    if key.startswith("__"):
        raise AttributeError(key)
    if hasattr(MissingPandasLikeScalars, key):
        raise getattr(MissingPandasLikeScalars, key)
    if hasattr(MissingPandasLikeGeneralFunctions, key):
        return getattr(MissingPandasLikeGeneralFunctions, key)
    else:
        raise AttributeError("module 'pyspark.pandas' has no attribute '%s'" % (key))