Unlock 30% off on Manual Testing Annual Plans this Holiday Season.

Copied to Clipboard!

How to use is_extension_dtype method in pandera

Best Python code snippet using pandera_python

utils.py

Source:utils.py

1#2# Licensed to the Apache Software Foundation (ASF) under one or more3# contributor license agreements.  See the NOTICE file distributed with4# this work for additional information regarding copyright ownership.5# The ASF licenses this file to You under the Apache License, Version 2.06# (the "License"); you may not use this file except in compliance with7# the License.  You may obtain a copy of the License at8#9#    http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16#17"""18Commonly used utils in pandas-on-Spark.19"""20import functools21from collections import OrderedDict22from contextlib import contextmanager23import os24from typing import (  # noqa: F401 (SPARK-34943)25    Any,26    Callable,27    Dict,28    Iterable,29    Iterator,30    List,31    Optional,32    Tuple,33    Union,34    TYPE_CHECKING,35    cast,36    no_type_check,37    overload,38)39import warnings40from pyspark.sql import functions as F, Column, DataFrame as SparkDataFrame, SparkSession41from pyspark.sql.types import DoubleType42import pandas as pd43from pandas.api.types import is_list_like44# For running doctests and reference resolution in PyCharm.45from pyspark import pandas as ps  # noqa: F40146from pyspark.pandas._typing import Axis, Label, Name, DataFrameOrSeries47from pyspark.pandas.spark import functions as SF48from pyspark.pandas.typedef.typehints import as_spark_type49if TYPE_CHECKING:50    # This is required in old Python 3.5 to prevent circular reference.51    from pyspark.pandas.base import IndexOpsMixin  # noqa: F401 (SPARK-34943)52    from pyspark.pandas.frame import DataFrame  # noqa: F401 (SPARK-34943)53    from pyspark.pandas.internal import InternalFrame  # noqa: F401 (SPARK-34943)54    from pyspark.pandas.series import Series  # noqa: F401 (SPARK-34943)55ERROR_MESSAGE_CANNOT_COMBINE = (56    "Cannot combine the series or dataframe because it comes from a different dataframe. "57    "In order to allow this operation, enable 'compute.ops_on_diff_frames' option."58)59SPARK_CONF_ARROW_ENABLED = "spark.sql.execution.arrow.pyspark.enabled"60def same_anchor(61    this: Union["DataFrame", "IndexOpsMixin", "InternalFrame"],62    that: Union["DataFrame", "IndexOpsMixin", "InternalFrame"],63) -> bool:64    """65    Check if the anchors of the given DataFrame or Series are the same or not.66    """67    from pyspark.pandas.base import IndexOpsMixin68    from pyspark.pandas.frame import DataFrame69    from pyspark.pandas.internal import InternalFrame70    if isinstance(this, InternalFrame):71        this_internal = this72    else:73        assert isinstance(this, (DataFrame, IndexOpsMixin)), type(this)74        this_internal = this._internal75    if isinstance(that, InternalFrame):76        that_internal = that77    else:78        assert isinstance(that, (DataFrame, IndexOpsMixin)), type(that)79        that_internal = that._internal80    return (81        this_internal.spark_frame is that_internal.spark_frame82        and this_internal.index_level == that_internal.index_level83        and all(84            spark_column_equals(this_scol, that_scol)85            for this_scol, that_scol in zip(86                this_internal.index_spark_columns, that_internal.index_spark_columns87            )88        )89    )90def combine_frames(91    this: "DataFrame",92    *args: DataFrameOrSeries,93    how: str = "full",94    preserve_order_column: bool = False95) -> "DataFrame":96    """97    This method combines `this` DataFrame with a different `that` DataFrame or98    Series from a different DataFrame.99    It returns a DataFrame that has prefix `this_` and `that_` to distinct100    the columns names from both DataFrames101    It internally performs a join operation which can be expensive in general.102    So, if `compute.ops_on_diff_frames` option is False,103    this method throws an exception.104    """105    from pyspark.pandas.config import get_option106    from pyspark.pandas.frame import DataFrame107    from pyspark.pandas.internal import (108        InternalField,109        InternalFrame,110        HIDDEN_COLUMNS,111        NATURAL_ORDER_COLUMN_NAME,112        SPARK_INDEX_NAME_FORMAT,113    )114    from pyspark.pandas.series import Series115    if all(isinstance(arg, Series) for arg in args):116        assert all(117            same_anchor(arg, args[0]) for arg in args118        ), "Currently only one different DataFrame (from given Series) is supported"119        assert not same_anchor(this, args[0]), "We don't need to combine. All series is in this."120        that = args[0]._psdf[list(args)]121    elif len(args) == 1 and isinstance(args[0], DataFrame):122        assert isinstance(args[0], DataFrame)123        assert not same_anchor(124            this, args[0]125        ), "We don't need to combine. `this` and `that` are same."126        that = args[0]127    else:128        raise AssertionError("args should be single DataFrame or " "single/multiple Series")129    if get_option("compute.ops_on_diff_frames"):130        def resolve(internal: InternalFrame, side: str) -> InternalFrame:131            rename = lambda col: "__{}_{}".format(side, col)132            internal = internal.resolved_copy133            sdf = internal.spark_frame134            sdf = internal.spark_frame.select(135                *[136                    scol_for(sdf, col).alias(rename(col))137                    for col in sdf.columns138                    if col not in HIDDEN_COLUMNS139                ],140                *HIDDEN_COLUMNS141            )142            return internal.copy(143                spark_frame=sdf,144                index_spark_columns=[145                    scol_for(sdf, rename(col)) for col in internal.index_spark_column_names146                ],147                index_fields=[148                    field.copy(name=rename(field.name)) for field in internal.index_fields149                ],150                data_spark_columns=[151                    scol_for(sdf, rename(col)) for col in internal.data_spark_column_names152                ],153                data_fields=[field.copy(name=rename(field.name)) for field in internal.data_fields],154            )155        this_internal = resolve(this._internal, "this")156        that_internal = resolve(that._internal, "that")157        this_index_map = list(158            zip(159                this_internal.index_spark_column_names,160                this_internal.index_names,161                this_internal.index_fields,162            )163        )164        that_index_map = list(165            zip(166                that_internal.index_spark_column_names,167                that_internal.index_names,168                that_internal.index_fields,169            )170        )171        assert len(this_index_map) == len(that_index_map)172        join_scols = []173        merged_index_scols = []174        # Note that the order of each element in index_map is guaranteed according to the index175        # level.176        this_and_that_index_map = list(zip(this_index_map, that_index_map))177        this_sdf = this_internal.spark_frame.alias("this")178        that_sdf = that_internal.spark_frame.alias("that")179        # If the same named index is found, that's used.180        index_column_names = []181        index_use_extension_dtypes = []182        for (183            i,184            ((this_column, this_name, this_field), (that_column, that_name, that_field)),185        ) in enumerate(this_and_that_index_map):186            if this_name == that_name:187                # We should merge the Spark columns into one188                # to mimic pandas' behavior.189                this_scol = scol_for(this_sdf, this_column)190                that_scol = scol_for(that_sdf, that_column)191                join_scol = this_scol == that_scol192                join_scols.append(join_scol)193                column_name = SPARK_INDEX_NAME_FORMAT(i)194                index_column_names.append(column_name)195                index_use_extension_dtypes.append(196                    any(field.is_extension_dtype for field in [this_field, that_field])197                )198                merged_index_scols.append(199                    F.when(this_scol.isNotNull(), this_scol).otherwise(that_scol).alias(column_name)200                )201            else:202                raise ValueError("Index names must be exactly matched currently.")203        assert len(join_scols) > 0, "cannot join with no overlapping index names"204        joined_df = this_sdf.join(that_sdf, on=join_scols, how=how)205        if preserve_order_column:206            order_column = [scol_for(this_sdf, NATURAL_ORDER_COLUMN_NAME)]207        else:208            order_column = []209        joined_df = joined_df.select(210            *merged_index_scols,211            *(212                scol_for(this_sdf, this_internal.spark_column_name_for(label))213                for label in this_internal.column_labels214            ),215            *(216                scol_for(that_sdf, that_internal.spark_column_name_for(label))217                for label in that_internal.column_labels218            ),219            *order_column220        )221        index_spark_columns = [scol_for(joined_df, col) for col in index_column_names]222        index_columns = set(index_column_names)223        new_data_columns = [224            col225            for col in joined_df.columns226            if col not in index_columns and col != NATURAL_ORDER_COLUMN_NAME227        ]228        schema = joined_df.select(*index_spark_columns, *new_data_columns).schema229        index_fields = [230            InternalField.from_struct_field(struct_field, use_extension_dtypes=use_extension_dtypes)231            for struct_field, use_extension_dtypes in zip(232                schema.fields[: len(index_spark_columns)], index_use_extension_dtypes233            )234        ]235        data_fields = [236            InternalField.from_struct_field(237                struct_field, use_extension_dtypes=field.is_extension_dtype238            )239            for struct_field, field in zip(240                schema.fields[len(index_spark_columns) :],241                this_internal.data_fields + that_internal.data_fields,242            )243        ]244        level = max(this_internal.column_labels_level, that_internal.column_labels_level)245        def fill_label(label: Optional[Label]) -> List:246            if label is None:247                return ([""] * (level - 1)) + [None]248            else:249                return ([""] * (level - len(label))) + list(label)250        column_labels = [251            tuple(["this"] + fill_label(label)) for label in this_internal.column_labels252        ] + [tuple(["that"] + fill_label(label)) for label in that_internal.column_labels]253        column_label_names = (254            cast(List[Optional[Label]], [None]) * (1 + level - this_internal.column_labels_level)255        ) + this_internal.column_label_names256        return DataFrame(257            InternalFrame(258                spark_frame=joined_df,259                index_spark_columns=index_spark_columns,260                index_names=this_internal.index_names,261                index_fields=index_fields,262                column_labels=column_labels,263                data_spark_columns=[scol_for(joined_df, col) for col in new_data_columns],264                data_fields=data_fields,265                column_label_names=column_label_names,266            )267        )268    else:269        raise ValueError(ERROR_MESSAGE_CANNOT_COMBINE)270def align_diff_frames(271    resolve_func: Callable[272        ["DataFrame", List[Label], List[Label]], Iterator[Tuple["Series", Label]]273    ],274    this: "DataFrame",275    that: "DataFrame",276    fillna: bool = True,277    how: str = "full",278    preserve_order_column: bool = False,279) -> "DataFrame":280    """281    This method aligns two different DataFrames with a given `func`. Columns are resolved and282    handled within the given `func`.283    To use this, `compute.ops_on_diff_frames` should be True, for now.284    :param resolve_func: Takes aligned (joined) DataFrame, the column of the current DataFrame, and285        the column of another DataFrame. It returns an iterable that produces Series.286        >>> from pyspark.pandas.config import set_option, reset_option287        >>>288        >>> set_option("compute.ops_on_diff_frames", True)289        >>>290        >>> psdf1 = ps.DataFrame({'a': [9, 8, 7, 6, 5, 4, 3, 2, 1]})291        >>> psdf2 = ps.DataFrame({'a': [9, 8, 7, 6, 5, 4, 3, 2, 1]})292        >>>293        >>> def func(psdf, this_column_labels, that_column_labels):294        ...    psdf  # conceptually this is A + B.295        ...296        ...    # Within this function, Series from A or B can be performed against `psdf`.297        ...    this_label = this_column_labels[0]  # this is ('a',) from psdf1.298        ...    that_label = that_column_labels[0]  # this is ('a',) from psdf2.299        ...    new_series = (psdf[this_label] - psdf[that_label]).rename(str(this_label))300        ...301        ...    # This new series will be placed in new DataFrame.302        ...    yield (new_series, this_label)303        >>>304        >>>305        >>> align_diff_frames(func, psdf1, psdf2).sort_index()306           a307        0  0308        1  0309        2  0310        3  0311        4  0312        5  0313        6  0314        7  0315        8  0316        >>> reset_option("compute.ops_on_diff_frames")317    :param this: a DataFrame to align318    :param that: another DataFrame to align319    :param fillna: If True, it fills missing values in non-common columns in both `this` and `that`.320        Otherwise, it returns as are.321    :param how: join way. In addition, it affects how `resolve_func` resolves the column conflict.322        - full: `resolve_func` should resolve only common columns from 'this' and 'that' DataFrames.323            For instance, if 'this' has columns A, B, C and that has B, C, D, `this_columns` and324            'that_columns' in this function are B, C and B, C.325        - left: `resolve_func` should resolve columns including that columns.326            For instance, if 'this' has columns A, B, C and that has B, C, D, `this_columns` is327            B, C but `that_columns` are B, C, D.328        - inner: Same as 'full' mode; however, internally performs inner join instead.329    :return: Aligned DataFrame330    """331    from pyspark.pandas.frame import DataFrame332    assert how == "full" or how == "left" or how == "inner"333    this_column_labels = this._internal.column_labels334    that_column_labels = that._internal.column_labels335    common_column_labels = set(this_column_labels).intersection(that_column_labels)336    # 1. Perform the join given two dataframes.337    combined = combine_frames(this, that, how=how, preserve_order_column=preserve_order_column)338    # 2. Apply the given function to transform the columns in a batch and keep the new columns.339    combined_column_labels = combined._internal.column_labels340    that_columns_to_apply = []  # type: List[Label]341    this_columns_to_apply = []  # type: List[Label]342    additional_that_columns = []  # type: List[Label]343    columns_to_keep = []  # type: List[Union[Series, Column]]344    column_labels_to_keep = []  # type: List[Label]345    for combined_label in combined_column_labels:346        for common_label in common_column_labels:347            if combined_label == tuple(["this", *common_label]):348                this_columns_to_apply.append(combined_label)349                break350            elif combined_label == tuple(["that", *common_label]):351                that_columns_to_apply.append(combined_label)352                break353        else:354            if how == "left" and combined_label in [355                tuple(["that", *label]) for label in that_column_labels356            ]:357                # In this case, we will drop `that_columns` in `columns_to_keep` but passes358                # it later to `func`. `func` should resolve it.359                # Note that adding this into a separate list (`additional_that_columns`)360                # is intentional so that `this_columns` and `that_columns` can be paired.361                additional_that_columns.append(combined_label)362            elif fillna:363                columns_to_keep.append(SF.lit(None).cast(DoubleType()).alias(str(combined_label)))364                column_labels_to_keep.append(combined_label)365            else:366                columns_to_keep.append(combined._psser_for(combined_label))367                column_labels_to_keep.append(combined_label)368    that_columns_to_apply += additional_that_columns369    # Should extract columns to apply and do it in a batch in case370    # it adds new columns for example.371    if len(this_columns_to_apply) > 0 or len(that_columns_to_apply) > 0:372        psser_set, column_labels_set = zip(373            *resolve_func(combined, this_columns_to_apply, that_columns_to_apply)374        )375        columns_applied = list(psser_set)  # type: List[Union[Series, Column]]376        column_labels_applied = list(column_labels_set)  # type: List[Label]377    else:378        columns_applied = []379        column_labels_applied = []380    applied = DataFrame(381        combined._internal.with_new_columns(382            columns_applied + columns_to_keep,383            column_labels=column_labels_applied + column_labels_to_keep,384        )385    )  # type: DataFrame386    # 3. Restore the names back and deduplicate columns.387    this_labels = OrderedDict()388    # Add columns in an order of its original frame.389    for this_label in this_column_labels:390        for new_label in applied._internal.column_labels:391            if new_label[1:] not in this_labels and this_label == new_label[1:]:392                this_labels[new_label[1:]] = new_label393    # After that, we will add the rest columns.394    other_labels = OrderedDict()395    for new_label in applied._internal.column_labels:396        if new_label[1:] not in this_labels:397            other_labels[new_label[1:]] = new_label398    psdf = applied[list(this_labels.values()) + list(other_labels.values())]399    psdf.columns = psdf.columns.droplevel()400    return psdf401def is_testing() -> bool:402    """Indicates whether Spark is currently running tests."""403    return "SPARK_TESTING" in os.environ404def default_session(conf: Optional[Dict[str, Any]] = None) -> SparkSession:405    if conf is None:406        conf = dict()407    builder = SparkSession.builder.appName("pandas-on-Spark")408    for key, value in conf.items():409        builder = builder.config(key, value)410    # Currently, pandas-on-Spark is dependent on such join due to 'compute.ops_on_diff_frames'411    # configuration. This is needed with Spark 3.0+.412    builder.config("spark.sql.analyzer.failAmbiguousSelfJoin", False)413    if is_testing():414        builder.config("spark.executor.allowSparkContext", False)415    return builder.getOrCreate()416@contextmanager417def sql_conf(pairs: Dict[str, Any], *, spark: Optional[SparkSession] = None) -> Iterator[None]:418    """419    A convenient context manager to set `value` to the Spark SQL configuration `key` and420    then restores it back when it exits.421    """422    assert isinstance(pairs, dict), "pairs should be a dictionary."423    if spark is None:424        spark = default_session()425    keys = pairs.keys()426    new_values = pairs.values()427    old_values = [spark.conf.get(key, None) for key in keys]428    for key, new_value in zip(keys, new_values):429        spark.conf.set(key, new_value)430    try:431        yield432    finally:433        for key, old_value in zip(keys, old_values):434            if old_value is None:435                spark.conf.unset(key)436            else:437                spark.conf.set(key, old_value)438def validate_arguments_and_invoke_function(439    pobj: Union[pd.DataFrame, pd.Series],440    pandas_on_spark_func: Callable,441    pandas_func: Callable,442    input_args: Dict,443) -> Any:444    """445    Invokes a pandas function.446    This is created because different versions of pandas support different parameters, and as a447    result when we code against the latest version, our users might get a confusing448    "got an unexpected keyword argument" error if they are using an older version of pandas.449    This function validates all the arguments, removes the ones that are not supported if they450    are simply the default value (i.e. most likely the user didn't explicitly specify it). It451    throws a TypeError if the user explicitly specify an argument that is not supported by the452    pandas version available.453    For example usage, look at DataFrame.to_html().454    :param pobj: the pandas DataFrame or Series to operate on455    :param pandas_on_spark_func: pandas-on-Spark function, used to get default parameter values456    :param pandas_func: pandas function, used to check whether pandas supports all the arguments457    :param input_args: arguments to pass to the pandas function, often created by using locals().458                       Make sure locals() call is at the top of the function so it captures only459                       input parameters, rather than local variables.460    :return: whatever pandas_func returns461    """462    import inspect463    # Makes a copy since whatever passed in is likely created by locals(), and we can't delete464    # 'self' key from that.465    args = input_args.copy()466    del args["self"]467    if "kwargs" in args:468        # explode kwargs469        kwargs = args["kwargs"]470        del args["kwargs"]471        args = {**args, **kwargs}472    pandas_on_spark_params = inspect.signature(pandas_on_spark_func).parameters473    pandas_params = inspect.signature(pandas_func).parameters474    for param in pandas_on_spark_params.values():475        if param.name not in pandas_params:476            if args[param.name] == param.default:477                del args[param.name]478            else:479                raise TypeError(480                    (481                        "The pandas version [%s] available does not support parameter '%s' "482                        + "for function '%s'."483                    )484                    % (pd.__version__, param.name, pandas_func.__name__)485                )486    args["self"] = pobj487    return pandas_func(**args)488@no_type_check489def lazy_property(fn: Callable[[Any], Any]) -> property:490    """491    Decorator that makes a property lazy-evaluated.492    Copied from https://stevenloria.com/lazy-properties/493    """494    attr_name = "_lazy_" + fn.__name__495    @property496    @functools.wraps(fn)497    def wrapped_lazy_property(self):498        if not hasattr(self, attr_name):499            setattr(self, attr_name, fn(self))500        return getattr(self, attr_name)501    def deleter(self):502        if hasattr(self, attr_name):503            delattr(self, attr_name)504    return wrapped_lazy_property.deleter(deleter)505def scol_for(sdf: SparkDataFrame, column_name: str) -> Column:506    """Return Spark Column for the given column name."""507    return sdf["`{}`".format(column_name)]508def column_labels_level(column_labels: List[Label]) -> int:509    """Return the level of the column index."""510    if len(column_labels) == 0:511        return 1512    else:513        levels = set(1 if label is None else len(label) for label in column_labels)514        assert len(levels) == 1, levels515        return list(levels)[0]516def name_like_string(name: Optional[Name]) -> str:517    """518    Return the name-like strings from str or tuple of str519    Examples520    --------521    >>> name = 'abc'522    >>> name_like_string(name)523    'abc'524    >>> name = ('abc',)525    >>> name_like_string(name)526    'abc'527    >>> name = ('a', 'b', 'c')528    >>> name_like_string(name)529    '(a, b, c)'530    """531    if name is None:532        label = ("__none__",)  # type: Label533    elif is_list_like(name):534        label = tuple([str(n) for n in name])535    else:536        label = (str(name),)537    return ("(%s)" % ", ".join(label)) if len(label) > 1 else label[0]538def is_name_like_tuple(value: Any, allow_none: bool = True, check_type: bool = False) -> bool:539    """540    Check the given tuple is be able to be used as a name.541    Examples542    --------543    >>> is_name_like_tuple(('abc',))544    True545    >>> is_name_like_tuple((1,))546    True547    >>> is_name_like_tuple(('abc', 1, None))548    True549    >>> is_name_like_tuple(('abc', 1, None), check_type=True)550    True551    >>> is_name_like_tuple((1.0j,))552    True553    >>> is_name_like_tuple(tuple())554    False555    >>> is_name_like_tuple((list('abc'),))556    False557    >>> is_name_like_tuple(('abc', 1, None), allow_none=False)558    False559    >>> is_name_like_tuple((1.0j,), check_type=True)560    False561    """562    if value is None:563        return allow_none564    elif not isinstance(value, tuple):565        return False566    elif len(value) == 0:567        return False568    elif not allow_none and any(v is None for v in value):569        return False570    elif any(is_list_like(v) or isinstance(v, slice) for v in value):571        return False572    elif check_type:573        return all(574            v is None or as_spark_type(type(v), raise_error=False) is not None for v in value575        )576    else:577        return True578def is_name_like_value(579    value: Any, allow_none: bool = True, allow_tuple: bool = True, check_type: bool = False580) -> bool:581    """582    Check the given value is like a name.583    Examples584    --------585    >>> is_name_like_value('abc')586    True587    >>> is_name_like_value(1)588    True589    >>> is_name_like_value(None)590    True591    >>> is_name_like_value(('abc',))592    True593    >>> is_name_like_value(1.0j)594    True595    >>> is_name_like_value(list('abc'))596    False597    >>> is_name_like_value(None, allow_none=False)598    False599    >>> is_name_like_value(('abc',), allow_tuple=False)600    False601    >>> is_name_like_value(1.0j, check_type=True)602    False603    """604    if value is None:605        return allow_none606    elif isinstance(value, tuple):607        return allow_tuple and is_name_like_tuple(608            value, allow_none=allow_none, check_type=check_type609        )610    elif is_list_like(value) or isinstance(value, slice):611        return False612    elif check_type:613        return as_spark_type(type(value), raise_error=False) is not None614    else:615        return True616def validate_axis(axis: Optional[Axis] = 0, none_axis: int = 0) -> int:617    """Check the given axis is valid."""618    # convert to numeric axis619    axis = cast(Dict[Optional[Axis], int], {None: none_axis, "index": 0, "columns": 1}).get(620        axis, axis621    )622    if axis in (none_axis, 0, 1):623        return cast(int, axis)624    else:625        raise ValueError("No axis named {0}".format(axis))626def validate_bool_kwarg(value: Any, arg_name: str) -> Optional[bool]:627    """Ensures that argument passed in arg_name is of type bool."""628    if not (isinstance(value, bool) or value is None):629        raise TypeError(630            'For argument "{}" expected type bool, received '631            "type {}.".format(arg_name, type(value).__name__)632        )633    return value634def validate_how(how: str) -> str:635    """Check the given how for join is valid."""636    if how == "full":637        warnings.warn(638            "Warning: While pandas-on-Spark will accept 'full', you should use 'outer' "639            + "instead to be compatible with the pandas merge API",640            UserWarning,641        )642    if how == "outer":643        # 'outer' in pandas equals 'full' in Spark644        how = "full"645    if how not in ("inner", "left", "right", "full"):646        raise ValueError(647            "The 'how' parameter has to be amongst the following values: ",648            "['inner', 'left', 'right', 'outer']",649        )650    return how651def validate_mode(mode: str) -> str:652    """Check the given mode for writing is valid."""653    if mode in ("w", "w+"):654        # 'w' in pandas equals 'overwrite' in Spark655        # '+' is meaningless for writing methods, but pandas just pass it as 'w'.656        mode = "overwrite"657    if mode in ("a", "a+"):658        # 'a' in pandas equals 'append' in Spark659        # '+' is meaningless for writing methods, but pandas just pass it as 'a'.660        mode = "append"661    if mode not in (662        "w",663        "a",664        "w+",665        "a+",666        "overwrite",667        "append",668        "ignore",669        "error",670        "errorifexists",671    ):672        raise ValueError(673            "The 'mode' parameter has to be amongst the following values: ",674            "['w', 'a', 'w+', 'a+', 'overwrite', 'append', 'ignore', 'error', 'errorifexists']",675        )676    return mode677@overload678def verify_temp_column_name(df: SparkDataFrame, column_name_or_label: str) -> str:679    ...680@overload681def verify_temp_column_name(df: "DataFrame", column_name_or_label: Name) -> Label:682    ...683def verify_temp_column_name(684    df: Union["DataFrame", SparkDataFrame], column_name_or_label: Union[str, Name]685) -> Union[str, Label]:686    """687    Verify that the given column name does not exist in the given pandas-on-Spark or688    Spark DataFrame.689    The temporary column names should start and end with `__`. In addition, `column_name_or_label`690    expects a single string, or column labels when `df` is a pandas-on-Spark DataFrame.691    >>> psdf = ps.DataFrame({("x", "a"): ['a', 'b', 'c']})692    >>> psdf["__dummy__"] = 0693    >>> psdf[("", "__dummy__")] = 1694    >>> psdf  # doctest: +NORMALIZE_WHITESPACE695       x __dummy__696       a           __dummy__697    0  a         0         1698    1  b         0         1699    2  c         0         1700    >>> verify_temp_column_name(psdf, '__tmp__')701    ('__tmp__', '')702    >>> verify_temp_column_name(psdf, ('', '__tmp__'))703    ('', '__tmp__')704    >>> verify_temp_column_name(psdf, '__dummy__')705    Traceback (most recent call last):706    ...707    AssertionError: ... `(__dummy__, )` ...708    >>> verify_temp_column_name(psdf, ('', '__dummy__'))709    Traceback (most recent call last):710    ...711    AssertionError: ... `(, __dummy__)` ...712    >>> verify_temp_column_name(psdf, 'dummy')713    Traceback (most recent call last):714    ...715    AssertionError: ... should be empty or start and end with `__`: ('dummy', '')716    >>> verify_temp_column_name(psdf, ('', 'dummy'))717    Traceback (most recent call last):718    ...719    AssertionError: ... should be empty or start and end with `__`: ('', 'dummy')720    >>> internal = psdf._internal.resolved_copy721    >>> sdf = internal.spark_frame722    >>> sdf.select(internal.data_spark_columns).show()  # doctest: +NORMALIZE_WHITESPACE723    +------+---------+-------------+724    |(x, a)|__dummy__|(, __dummy__)|725    +------+---------+-------------+726    |     a|        0|            1|727    |     b|        0|            1|728    |     c|        0|            1|729    +------+---------+-------------+730    >>> verify_temp_column_name(sdf, '__tmp__')731    '__tmp__'732    >>> verify_temp_column_name(sdf, '__dummy__')733    Traceback (most recent call last):734    ...735    AssertionError: ... `__dummy__` ... '(x, a)', '__dummy__', '(, __dummy__)', ...736    >>> verify_temp_column_name(sdf, ('', '__dummy__'))737    Traceback (most recent call last):738    ...739    AssertionError: <class 'tuple'>740    >>> verify_temp_column_name(sdf, 'dummy')741    Traceback (most recent call last):742    ...743    AssertionError: ... should start and end with `__`: dummy744    """745    from pyspark.pandas.frame import DataFrame746    if isinstance(df, DataFrame):747        if isinstance(column_name_or_label, str):748            column_name = column_name_or_label749            level = df._internal.column_labels_level750            column_name_or_label = tuple([column_name_or_label] + ([""] * (level - 1)))751        else:752            column_name = name_like_string(column_name_or_label)753        assert any(len(label) > 0 for label in column_name_or_label) and all(754            label == "" or (label.startswith("__") and label.endswith("__"))755            for label in column_name_or_label756        ), "The temporary column name should be empty or start and end with `__`: {}".format(757            column_name_or_label758        )759        assert all(760            column_name_or_label != label for label in df._internal.column_labels761        ), "The given column name `{}` already exists in the pandas-on-Spark DataFrame: {}".format(762            name_like_string(column_name_or_label), df.columns763        )764        df = df._internal.resolved_copy.spark_frame765    else:766        assert isinstance(column_name_or_label, str), type(column_name_or_label)767        assert column_name_or_label.startswith("__") and column_name_or_label.endswith(768            "__"769        ), "The temporary column name should start and end with `__`: {}".format(770            column_name_or_label771        )772        column_name = column_name_or_label773    assert isinstance(df, SparkDataFrame), type(df)774    assert (775        column_name not in df.columns776    ), "The given column name `{}` already exists in the Spark DataFrame: {}".format(777        column_name, df.columns778    )779    return column_name_or_label780def spark_column_equals(left: Column, right: Column) -> bool:781    """782    Check both `left` and `right` have the same expressions.783    >>> spark_column_equals(SF.lit(0), SF.lit(0))784    True785    >>> spark_column_equals(SF.lit(0) + 1, SF.lit(0) + 1)786    True787    >>> spark_column_equals(SF.lit(0) + 1, SF.lit(0) + 2)788    False789    >>> sdf1 = ps.DataFrame({"x": ['a', 'b', 'c']}).to_spark()790    >>> spark_column_equals(sdf1["x"] + 1, sdf1["x"] + 1)791    True792    >>> sdf2 = ps.DataFrame({"x": ['a', 'b', 'c']}).to_spark()793    >>> spark_column_equals(sdf1["x"] + 1, sdf2["x"] + 1)794    False795    """796    return left._jc.equals(right._jc)  # type: ignore797def compare_null_first(798    left: Column,799    right: Column,800    comp: Callable[[Column, Column], Column],801) -> Column:802    return (left.isNotNull() & right.isNotNull() & comp(left, right)) | (803        left.isNull() & right.isNotNull()804    )805def compare_null_last(806    left: Column,807    right: Column,808    comp: Callable[[Column, Column], Column],809) -> Column:810    return (left.isNotNull() & right.isNotNull() & comp(left, right)) | (811        left.isNotNull() & right.isNull()812    )813def compare_disallow_null(814    left: Column,815    right: Column,816    comp: Callable[[Column, Column], Column],817) -> Column:818    return left.isNotNull() & right.isNotNull() & comp(left, right)819def compare_allow_null(820    left: Column,821    right: Column,822    comp: Callable[[Column, Column], Column],823) -> Column:824    return left.isNull() | right.isNull() | comp(left, right)825def _test() -> None:826    import os827    import doctest828    import sys829    from pyspark.sql import SparkSession830    import pyspark.pandas.utils831    os.chdir(os.environ["SPARK_HOME"])832    globs = pyspark.pandas.utils.__dict__.copy()833    globs["ps"] = pyspark.pandas834    spark = (835        SparkSession.builder.master("local[4]").appName("pyspark.pandas.utils tests").getOrCreate()836    )837    (failure_count, test_count) = doctest.testmod(838        pyspark.pandas.utils,839        globs=globs,840        optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE,841    )842    spark.stop()843    if failure_count:844        sys.exit(-1)845if __name__ == "__main__":...

pandas_engine.py

Source:pandas_engine.py

...27try:28    from typing import Literal  # type: ignore29except ImportError:30    from typing_extensions import Literal  # type: ignore31def is_extension_dtype(pd_dtype: PandasDataType) -> bool:32    """Check if a value is a pandas extension type or instance of one."""33    return isinstance(pd_dtype, PandasExtensionType) or (34        isinstance(pd_dtype, type)35        and issubclass(pd_dtype, PandasExtensionType)36    )37@immutable(init=True)38class DataType(dtypes.DataType):39    """Base `DataType` for boxing Pandas data types."""40    type: Any = dataclasses.field(repr=False, init=False)41    """Native pandas dtype boxed by the data type."""42    def __init__(self, dtype: Any):43        super().__init__()44        object.__setattr__(self, "type", pd.api.types.pandas_dtype(dtype))45        dtype_cls = dtype if inspect.isclass(dtype) else dtype.__class__46        warnings.warn(47            f"'{dtype_cls}' support is not guaranteed.\n"48            + "Usage Tip: Consider writing a custom "49            + "pandera.dtypes.DataType or opening an issue at "50            + "https://github.com/pandera-dev/pandera"51        )52    def __post_init__(self):53        # this method isn't called if __init__ is defined54        object.__setattr__(55            self, "type", pd.api.types.pandas_dtype(self.type)56        )  # pragma: no cover57    def coerce(self, data_container: PandasObject) -> PandasObject:58        """Pure coerce without catching exceptions."""59        coerced = data_container.astype(self.type)60        if type(data_container).__module__.startswith("modin.pandas"):61            # NOTE: this is a hack to enable catching of errors in modin62            coerced.__str__()63        return coerced64    def try_coerce(self, data_container: PandasObject) -> PandasObject:65        try:66            return self.coerce(data_container)67        except Exception as exc:  # pylint:disable=broad-except68            raise errors.ParserError(69                f"Could not coerce {type(data_container)} data_container "70                f"into type {self.type}",71                failure_cases=utils.numpy_pandas_coerce_failure_cases(72                    data_container, self73                ),74            ) from exc75    def check(self, pandera_dtype: dtypes.DataType) -> bool:76        try:77            pandera_dtype = Engine.dtype(pandera_dtype)78        except TypeError:79            return False80        # attempts to compare pandas native type if possible81        # to let subclass inherit check82        # (super will compare that DataType classes are exactly the same)83        try:84            return self.type == pandera_dtype.type or super().check(85                pandera_dtype86            )87        except TypeError:88            return super().check(pandera_dtype)89    def __str__(self) -> str:90        return str(self.type)91    def __repr__(self) -> str:92        return f"DataType({self})"93class Engine(  # pylint:disable=too-few-public-methods94    metaclass=engine.Engine,95    base_pandera_dtypes=(DataType, numpy_engine.DataType),96):97    """Pandas data type engine."""98    @classmethod99    def dtype(cls, data_type: Any) -> "DataType":100        """Convert input into a pandas-compatible101        Pandera :class:`~pandera.dtypes.DataType` object."""102        try:103            return engine.Engine.dtype(cls, data_type)104        except TypeError:105            if is_extension_dtype(data_type) and isinstance(data_type, type):106                try:107                    np_or_pd_dtype = data_type()108                    # Convert to str here because some pandas dtypes allow109                    # an empty constructor for compatibility but fail on110                    # str(). e.g: PeriodDtype111                    str(np_or_pd_dtype.name)112                except (TypeError, AttributeError) as err:113                    raise TypeError(114                        f" dtype {data_type} cannot be instantiated: {err}\n"115                        "Usage Tip: Use an instance or a string "116                        "representation."117                    ) from None118            else:119                # let pandas transform any acceptable value...

data_manipulation.py

Source:data_manipulation.py

...135    is_int_dtype = pd.api.types.is_integer_dtype136    is_float_dtype = pd.api.types.is_float_dtype137    is_string_dtype = pd.api.types.is_string_dtype138    is_extension_dtype = pd.api.types.is_extension_array_dtype139    if is_extension_dtype(values):140        if is_int_dtype(values):141            downcast_values = values.fillna(-1).astype(int)142        elif is_bool_dtype(values):143            downcast_values = values.fillna(False).astype(np.bool_, errors='raise')144        elif is_float_dtype(values):145            downcast_values = pd.to_numeric(values, errors='coerce')146        elif is_string_dtype(values):147            downcast_values = values.astype(np.object_, errors='raise')148        else:149            raise TypeError('unsupported extended datatype {TYPE} provided'.format(TYPE=values.dtype))150    else:151        downcast_values = values152    return downcast_values153def evaluate_condition(data, expression):154    """155    Evaluate a boolean expression for a set of data.156    Arguments:157        data: dictionary, Series, or DataFrame of data potentially containing one or more variables used in the158            operation.159        expression: string or list of expression components describing a conditional statement to evaluate on the160            provided data.161    Returns:162        results (pd.Series): results of the evaluation for each row of data provided.163    """164    is_bool_dtype = pd.api.types.is_bool_dtype165    is_extension_dtype = pd.api.types.is_extension_array_dtype166    reserved_chars = ('and', 'or', 'in', 'not', '+', '-', '/', '//', '*', '**', '%', '>', '>=', '<', '<=', '==', '!=',167                      '~', ',', '(', ')', '[', ']', '{', '}')168    if isinstance(data, pd.Series):  # single data entry169        df = data.to_frame().T170    elif isinstance(data, pd.DataFrame):  # one or more entries171        df = data.copy()172    elif isinstance(data, dict):  # one or more entries173        try:174            df = pd.DataFrame(data)175        except ValueError:  # single entry was given176            df = pd.DataFrame(data, index=[0])177    else:178        raise ValueError('data must be either a pandas DataFrame or Series')179    header = df.columns.tolist()180    if isinstance(expression, str):181        components = parse_expression(expression)182    elif isinstance(expression, list):183        components = expression184    else:185        raise ValueError('expression {} must be provided as either a string or list'.format(expression))186    if len(components) > 1:187        # Quote non-numeric, static expression variables188        exp_comps = []189        for component in components:190            if component in header:191                col_values = df[component]192                if is_extension_dtype(col_values):193                    try:194                        df.loc[:, component] = downcast_extended(col_values)195                    except TypeError:196                        logger.exception('failed to downcast component {COMP} values'.format(COMP=component))197                        raise198                exp_comp = '`{}`'.format(component)199            elif is_numeric(component) or component in reserved_chars:200                exp_comp = component201            else:  # component is a string202                exp_comp = '"{}"'.format(component)203            exp_comps.append(exp_comp)204        expression = ' '.join(exp_comps)205        try:206            df_match = df.eval(expression)207        except ValueError:208            print(expression)209            print(df)210            print(df.dtypes)211            raise212    else:  # component is a single static value or column values213        expression = ' '.join(components)214        if expression in header:215            values = df[expression]216            if is_bool_dtype(values.dtype):217                df_match = values218            else:219                df_match = ~ values.isna()220        else:221            df_match = df.eval(expression)222    if not isinstance(df_match, pd.Series):223        df_match = pd.Series(df_match, df.index)224    return df_match225def evaluate_operation(data, expression):226    """227    Evaluate a mathematical expression for a set of data.228    Arguments:229        data: dictionary, Series, or DataFrame of data potentially containing one or more variables used in the230            operation.231        expression: string or list of expression components describing an operation to evaluate on the provided data.232    Returns:233        results (pd.Series): results of the evaluation for each row of data provided.234    """235    is_extension_dtype = pd.api.types.is_extension_array_dtype236    if isinstance(data, pd.Series):237        df = data.to_frame().T238    elif isinstance(data, pd.DataFrame):239        df = data.copy()240    elif isinstance(data, dict):241        try:242            df = pd.DataFrame(data)243        except ValueError:  # single entry was given244            df = pd.DataFrame(data, index=[0])245    else:246        raise ValueError('data must be either a pandas DataFrame, Series, or a dictionary')247    header = df.columns.tolist()248    if isinstance(expression, str):249        components = parse_expression(expression)250    elif isinstance(expression, list):251        components = expression252    else:253        components = parse_expression(str(expression))254    if len(components) > 1:255        exp_comps = []256        for component in components:257            if component in header:258                col_values = df[component]259                if is_extension_dtype(col_values):260                    try:261                        df.loc[:, component] = downcast_extended(col_values)262                    except TypeError:263                        logger.exception('failed to downcast component {COMP} values'.format(COMP=component))264                        raise265                exp_comp = '`{}`'.format(component)266            else:267                exp_comp = component268            exp_comps.append(exp_comp)269        expression = ' '.join(exp_comps)270        results = df.eval(expression).squeeze()271    else:  # results are a single static value or the values of a column in the dataframe272        expression = ' '.join(components)273        if expression in header:...

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.