strings.py

import warnings

import numpy as np

from pandas import (
    NA,
    Categorical,
    DataFrame,
    Index,
    Series,
)
from pandas.arrays import StringArray


class Dtypes:
    params = ["str", "string[python]", "string[pyarrow]"]
    param_names = ["dtype"]

    def setup(self, dtype):
        try:
            self.s = Series(
                Index([f"i-{i}" for i in range(10000)], dtype=object)._values,
                dtype=dtype,
            )
        except ImportError as err:
            raise NotImplementedError from err


class Construction:
    params = (
        ["series", "frame", "categorical_series"],
        ["str", "string[python]", "string[pyarrow]"],
    )
    param_names = ["pd_type", "dtype"]
    pd_mapping = {"series": Series, "frame": DataFrame, "categorical_series": Series}
    dtype_mapping = {"str": "str", "string[python]": object, "string[pyarrow]": object}

    def setup(self, pd_type, dtype):
        series_arr = np.array(
            [str(i) * 10 for i in range(100_000)], dtype=self.dtype_mapping[dtype]
        )
        if pd_type == "series":
            self.arr = series_arr
        elif pd_type == "frame":
            self.arr = series_arr.reshape((50_000, 2)).copy()
        elif pd_type == "categorical_series":
            # GH37371. Testing construction of string series/frames from ExtensionArrays
            self.arr = Categorical(series_arr)

    def time_construction(self, pd_type, dtype):
        self.pd_mapping[pd_type](self.arr, dtype=dtype)

    def peakmem_construction(self, pd_type, dtype):
        self.pd_mapping[pd_type](self.arr, dtype=dtype)


class Methods(Dtypes):
    def time_center(self, dtype):
        self.s.str.center(100)

    def time_count(self, dtype):
        self.s.str.count("A")

    def time_endswith(self, dtype):
        self.s.str.endswith("A")

    def time_extract(self, dtype):
        with warnings.catch_warnings(record=True):
            self.s.str.extract("(\\w*)A(\\w*)")

    def time_findall(self, dtype):
        self.s.str.findall("[A-Z]+")

    def time_find(self, dtype):
        self.s.str.find("[A-Z]+")

    def time_rfind(self, dtype):
        self.s.str.rfind("[A-Z]+")

    def time_fullmatch(self, dtype):
        self.s.str.fullmatch("A")

    def time_get(self, dtype):
        self.s.str.get(0)

    def time_len(self, dtype):
        self.s.str.len()

    def time_join(self, dtype):
        self.s.str.join(" ")

    def time_match(self, dtype):
        self.s.str.match("A")

    def time_normalize(self, dtype):
        self.s.str.normalize("NFC")

    def time_pad(self, dtype):
        self.s.str.pad(100, side="both")

    def time_partition(self, dtype):
        self.s.str.partition("A")

    def time_rpartition(self, dtype):
        self.s.str.rpartition("A")

    def time_replace(self, dtype):
        self.s.str.replace("A", "\x01\x01")

    def time_translate(self, dtype):
        self.s.str.translate({"A": "\x01\x01"})

    def time_slice(self, dtype):
        self.s.str.slice(5, 15, 2)

    def time_startswith(self, dtype):
        self.s.str.startswith("A")

    def time_strip(self, dtype):
        self.s.str.strip("A")

    def time_rstrip(self, dtype):
        self.s.str.rstrip("A")

    def time_lstrip(self, dtype):
        self.s.str.lstrip("A")

    def time_title(self, dtype):
        self.s.str.title()

    def time_upper(self, dtype):
        self.s.str.upper()

    def time_lower(self, dtype):
        self.s.str.lower()

    def time_wrap(self, dtype):
        self.s.str.wrap(10)

    def time_zfill(self, dtype):
        self.s.str.zfill(10)

    def time_isalnum(self, dtype):
        self.s.str.isalnum()

    def time_isalpha(self, dtype):
        self.s.str.isalpha()

    def time_isdecimal(self, dtype):
        self.s.str.isdecimal()

    def time_isdigit(self, dtype):
        self.s.str.isdigit()

    def time_islower(self, dtype):
        self.s.str.islower()

    def time_isnumeric(self, dtype):
        self.s.str.isnumeric()

    def time_isspace(self, dtype):
        self.s.str.isspace()

    def time_istitle(self, dtype):
        self.s.str.istitle()

    def time_isupper(self, dtype):
        self.s.str.isupper()


class Repeat:
    params = ["int", "array"]
    param_names = ["repeats"]

    def setup(self, repeats):
        N = 10**5
        self.s = Series(Index([f"i-{i}" for i in range(N)], dtype=object))
        repeat = {"int": 1, "array": np.random.randint(1, 3, N)}
        self.values = repeat[repeats]

    def time_repeat(self, repeats):
        self.s.str.repeat(self.values)


class Cat:
    params = ([0, 3], [None, ","], [None, "-"], [0.0, 0.001, 0.15])
    param_names = ["other_cols", "sep", "na_rep", "na_frac"]

    def setup(self, other_cols, sep, na_rep, na_frac):
        N = 10**5
        mask_gen = lambda: np.random.choice([True, False], N, p=[1 - na_frac, na_frac])
        self.s = Series(Index([f"i-{i}" for i in range(N)], dtype=object)).where(
            mask_gen()
        )
        if other_cols == 0:
            # str.cat self-concatenates only for others=None
            self.others = None
        else:
            self.others = DataFrame(
                {
                    i: Index([f"i-{i}" for i in range(N)], dtype=object).where(
                        mask_gen()
                    )
                    for i in range(other_cols)
                }
            )

    def time_cat(self, other_cols, sep, na_rep, na_frac):
        # before the concatenation (one caller + other_cols columns), the total
        # expected fraction of rows containing any NaN is:
        # reduce(lambda t, _: t + (1 - t) * na_frac, range(other_cols + 1), 0)
        # for other_cols=3 and na_frac=0.15, this works out to ~48%
        self.s.str.cat(others=self.others, sep=sep, na_rep=na_rep)


class Contains(Dtypes):
    params = (Dtypes.params, [True, False])
    param_names = ["dtype", "regex"]

    def setup(self, dtype, regex):
        super().setup(dtype)

    def time_contains(self, dtype, regex):
        self.s.str.contains("A", regex=regex)


class Split(Dtypes):
    params = (Dtypes.params, [True, False])
    param_names = ["dtype", "expand"]

    def setup(self, dtype, expand):
        super().setup(dtype)
        self.s = self.s.str.join("--")

    def time_split(self, dtype, expand):
        self.s.str.split("--", expand=expand)

    def time_rsplit(self, dtype, expand):
        self.s.str.rsplit("--", expand=expand)


class Extract(Dtypes):
    params = (Dtypes.params, [True, False])
    param_names = ["dtype", "expand"]

    def setup(self, dtype, expand):
        super().setup(dtype)

    def time_extract_single_group(self, dtype, expand):
        with warnings.catch_warnings(record=True):
            self.s.str.extract("(\\w*)A", expand=expand)


class Dummies(Dtypes):
    def setup(self, dtype):
        super().setup(dtype)
        N = len(self.s) // 5
        self.s = self.s[:N].str.join("|")

    def time_get_dummies(self, dtype):
        self.s.str.get_dummies("|")


class Encode:
    def setup(self):
        self.ser = Series(Index([f"i-{i}" for i in range(10_000)], dtype=object))

    def time_encode_decode(self):
        self.ser.str.encode("utf-8").str.decode("utf-8")


class Slice:
    def setup(self):
        self.s = Series(["abcdefg", np.nan] * 500000)

    def time_vector_slice(self):
        # GH 2602
        self.s.str[:5]


class Iter(Dtypes):
    def time_iter(self, dtype):
        for i in self.s:
            pass


class StringArrayConstruction:
    def setup(self):
        self.series_arr = np.array([str(i) * 10 for i in range(10**5)], dtype=object)
        self.series_arr_nan = np.concatenate([self.series_arr, np.array([NA] * 1000)])

    def time_string_array_construction(self):
        StringArray(self.series_arr)

    def time_string_array_with_nan_construction(self):
        StringArray(self.series_arr_nan)

    def peakmem_stringarray_construction(self):
        StringArray(self.series_arr)