diff --git a/asv_bench/benchmarks/accessors.py b/asv_bench/benchmarks/accessors.py index 229e5f7215d..47412d63733 100644 --- a/asv_bench/benchmarks/accessors.py +++ b/asv_bench/benchmarks/accessors.py @@ -22,3 +22,42 @@ def time_year(self, calendar): def time_floor(self, calendar): _ = self.da.time.dt.floor("D") + + +NSTRINGS = 100_000 + + +@parameterize({"dtype": ("U", "S", "O")}) +class StringAccessor: + def setup(self, dtype): + rng = np.random.default_rng(0) + chars = np.array(list("abcdefgh XYZ")) + strings = rng.choice(chars, size=(NSTRINGS, 12)).view("U12").ravel() + self.da = xr.DataArray(strings.astype(dtype), dims="x") + + def time_len(self, dtype): + self.da.str.len() + + def time_isalpha(self, dtype): + self.da.str.isalpha() + + def time_startswith(self, dtype): + self.da.str.startswith("ab") + + def time_find(self, dtype): + self.da.str.find("ab") + + def time_contains(self, dtype): + self.da.str.contains("ab", regex=False) + + def time_strip(self, dtype): + self.da.str.strip() + + def time_center(self, dtype): + self.da.str.center(16) + + def time_replace(self, dtype): + self.da.str.replace("ab", "-", regex=False) + + def time_repeat(self, dtype): + self.da.str.repeat(2) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 0ebce396662..5aa10185c44 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -147,6 +147,11 @@ Deprecations Bug Fixes ~~~~~~~~~ +- Results of the string accessor methods ``strip``, ``pad``, ``center``, + ``ljust``, ``rjust``, ``zfill`` and ``replace`` on dask arrays are no longer + truncated to a single character, and ``repeat`` no longer fails on dask arrays + (:pull:`11746`). + By `Michael Niklas `_. - Encoding a string variable whose ``char_dim_name`` has characters after its last number (e.g. ``"n4_chars"``) no longer drops them. When the length does not match, only the number is replaced (``"n3_chars"``) (:pull:`11757`). @@ -386,6 +391,14 @@ Documentation Performance ~~~~~~~~~~~ +- The string accessor uses the vectorized functions of ``np.strings`` for + ``len``, ``startswith``, ``endswith``, ``find``, ``rfind``, ``index``, + ``rindex``, ``strip``, ``lstrip``, ``rstrip``, ``pad``, ``center``, + ``ljust``, ``rjust``, ``zfill``, ``repeat``, the ``is*`` methods and + ``contains`` and ``replace`` with ``regex=False`` and case-sensitive matching, + which makes them 5-80 times faster for arrays with a string dtype + (:pull:`11746`). + By `Michael Niklas `_. - Speed up rolling reductions of dask arrays without bottleneck, e.g. ``sum``, ``mean`` and ``count``, by about 7x with about 12x fewer tasks for a window of 20. Counting the values in each window no longer splits the chunks of the other diff --git a/xarray/core/accessor_str.py b/xarray/core/accessor_str.py index 765db262f15..a51c3822bc2 100644 --- a/xarray/core/accessor_str.py +++ b/xarray/core/accessor_str.py @@ -44,6 +44,7 @@ import textwrap from collections.abc import Callable, Hashable, Mapping from functools import reduce +from operator import methodcaller from operator import or_ as set_union from re import Pattern from typing import TYPE_CHECKING, Any @@ -53,6 +54,7 @@ from xarray.core import duck_array_ops from xarray.core.types import T_DataArray +from xarray.namedarray.pycompat import is_chunked_array if TYPE_CHECKING: from numpy.typing import DTypeLike @@ -117,6 +119,7 @@ def _apply_str_ufunc( output_sizes: Mapping[Any, int] | None = None, func_args: tuple = (), func_kwargs: Mapping = {}, + vectorize: bool = True, ) -> Any: # TODO handling of na values ? if dtype is None: @@ -132,7 +135,7 @@ def _apply_str_ufunc( func, obj, *func_args, - vectorize=True, + vectorize=vectorize, dask="parallelized", output_dtypes=[dtype], output_core_dims=output_core_dims, @@ -233,12 +236,22 @@ def _apply( self, *, func: Callable, + np_func: Callable | None = None, dtype: DTypeLike | None = None, output_core_dims: list | tuple = ((),), output_sizes: Mapping[Any, int] | None = None, func_args: tuple = (), func_kwargs: Mapping = {}, ) -> T_DataArray: + """Apply ``func`` to each element, or ``np_func`` to the whole array. + + ``np_func`` is the vectorized equivalent from ``np.strings``, used if it + supports the dtypes. + """ + vectorize = True + if np_func is not None and self._supports_np_strings(func_args): + func = np_func + vectorize = False return _apply_str_ufunc( obj=self._obj, func=func, @@ -247,6 +260,20 @@ def _apply( output_sizes=output_sizes, func_args=func_args, func_kwargs=func_kwargs, + vectorize=vectorize, + ) + + def _supports_np_strings(self, func_args: tuple) -> bool: + """Whether the functions of ``np.strings`` support the array and arguments. + + They support string dtypes, but not object arrays of strings. Empty + arrays use the elementwise path, as e.g. ``np.strings.center`` fails on + them. + """ + return ( + self._obj.dtype.kind in "UST" + and self._obj.size > 0 + and not any(getattr(arg, "dtype", None) == np.object_ for arg in func_args) ) def _re_compile( @@ -296,7 +323,7 @@ def len(self) -> T_DataArray: ------- lengths array : array of int """ - return self._apply(func=len, dtype=int) + return self._apply(func=len, np_func=np.strings.str_len, dtype=int) def __getitem__( self, @@ -867,7 +894,9 @@ def isalnum(self) -> T_DataArray: array([ True, False]) Dimensions without coordinates: x """ - return self._apply(func=lambda x: x.isalnum(), dtype=bool) + return self._apply( + func=lambda x: x.isalnum(), np_func=np.strings.isalnum, dtype=bool + ) def isalpha(self) -> T_DataArray: """Check whether all characters in each string are alphabetic. @@ -890,7 +919,9 @@ def isalpha(self) -> T_DataArray: array([ True, False, False]) Dimensions without coordinates: x """ - return self._apply(func=lambda x: x.isalpha(), dtype=bool) + return self._apply( + func=lambda x: x.isalpha(), np_func=np.strings.isalpha, dtype=bool + ) def isdecimal(self) -> T_DataArray: """Check whether all characters in each string are decimal. @@ -913,7 +944,9 @@ def isdecimal(self) -> T_DataArray: array([False, True, True]) Dimensions without coordinates: x """ - return self._apply(func=lambda x: x.isdecimal(), dtype=bool) + return self._apply( + func=lambda x: x.isdecimal(), np_func=np.strings.isdecimal, dtype=bool + ) def isdigit(self) -> T_DataArray: """Check whether all characters in each string are digits. @@ -936,7 +969,9 @@ def isdigit(self) -> T_DataArray: array([ True, False, True, False, False]) Dimensions without coordinates: x """ - return self._apply(func=lambda x: x.isdigit(), dtype=bool) + return self._apply( + func=lambda x: x.isdigit(), np_func=np.strings.isdigit, dtype=bool + ) def islower(self) -> T_DataArray: """Check whether all characters in each string are lowercase. @@ -960,7 +995,9 @@ def islower(self) -> T_DataArray: array([ True, False, False]) Dimensions without coordinates: x """ - return self._apply(func=lambda x: x.islower(), dtype=bool) + return self._apply( + func=lambda x: x.islower(), np_func=np.strings.islower, dtype=bool + ) def isnumeric(self) -> T_DataArray: """Check whether all characters in each string are numeric. @@ -983,7 +1020,9 @@ def isnumeric(self) -> T_DataArray: array([ True, False, False, False, False]) Dimensions without coordinates: x """ - return self._apply(func=lambda x: x.isnumeric(), dtype=bool) + return self._apply( + func=lambda x: x.isnumeric(), np_func=np.strings.isnumeric, dtype=bool + ) def isspace(self) -> T_DataArray: r"""Check whether all characters in each string are spaces. @@ -1006,7 +1045,9 @@ def isspace(self) -> T_DataArray: array([False, True, True, True]) Dimensions without coordinates: x """ - return self._apply(func=lambda x: x.isspace(), dtype=bool) + return self._apply( + func=lambda x: x.isspace(), np_func=np.strings.isspace, dtype=bool + ) def istitle(self) -> T_DataArray: """Check whether all characters in each string are titlecase. @@ -1037,7 +1078,9 @@ def istitle(self) -> T_DataArray: array([ True, False, False]) Dimensions without coordinates: title """ - return self._apply(func=lambda x: x.istitle(), dtype=bool) + return self._apply( + func=lambda x: x.istitle(), np_func=np.strings.istitle, dtype=bool + ) def isupper(self) -> T_DataArray: """Check whether all characters in each string are uppercase. @@ -1060,7 +1103,9 @@ def isupper(self) -> T_DataArray: array([ True, False, False]) Dimensions without coordinates: x """ - return self._apply(func=lambda x: x.isupper(), dtype=bool) + return self._apply( + func=lambda x: x.isupper(), np_func=np.strings.isupper, dtype=bool + ) def count( self, pat: str | bytes | Pattern | Any, flags: int = 0, case: bool | None = None @@ -1172,7 +1217,9 @@ def startswith(self, pat: str | bytes | Any) -> T_DataArray: """ pat = self._stringify(pat) func = lambda x, y: x.startswith(y) - return self._apply(func=func, func_args=(pat,), dtype=bool) + return self._apply( + func=func, np_func=np.strings.startswith, func_args=(pat,), dtype=bool + ) def endswith(self, pat: str | bytes | Any) -> T_DataArray: """Test if the end of each string in the array matches a pattern. @@ -1207,7 +1254,9 @@ def endswith(self, pat: str | bytes | Any) -> T_DataArray: """ pat = self._stringify(pat) func = lambda x, y: x.endswith(y) - return self._apply(func=func, func_args=(pat,), dtype=bool) + return self._apply( + func=func, np_func=np.strings.endswith, func_args=(pat,), dtype=bool + ) def pad( self, @@ -1312,18 +1361,24 @@ def _padder( self, *, func: Callable, + np_func: Callable, width: int | Any, fillchar: str | bytes | Any = " ", ) -> T_DataArray: """Handle padding operations.""" fillchar = self._stringify(fillchar) + # check here, np.strings raises a different error + if isinstance(fillchar, str | bytes) and len(fillchar) != 1: + raise TypeError("fillchar must be a character, not str") + # np.strings requires integer widths + width = width.astype(int) if hasattr(width, "astype") else int(width) def overfunc(x, iwidth, ifillchar): if len(ifillchar) != 1: raise TypeError("fillchar must be a character, not str") - return func(x, int(iwidth), ifillchar) + return func(x, iwidth, ifillchar) - return self._apply(func=overfunc, func_args=(width, fillchar)) + return self._apply(func=overfunc, np_func=np_func, func_args=(width, fillchar)) def center( self, width: int | Any, fillchar: str | bytes | Any = " " @@ -1347,7 +1402,9 @@ def center( filled : same type as values """ func = self._method("center") - return self._padder(func=func, width=width, fillchar=fillchar) + return self._padder( + func=func, np_func=np.strings.center, width=width, fillchar=fillchar + ) def ljust( self, @@ -1373,7 +1430,9 @@ def ljust( filled : same type as values """ func = self._method("ljust") - return self._padder(func=func, width=width, fillchar=fillchar) + return self._padder( + func=func, np_func=np.strings.ljust, width=width, fillchar=fillchar + ) def rjust( self, @@ -1399,7 +1458,9 @@ def rjust( filled : same type as values """ func = self._method("rjust") - return self._padder(func=func, width=width, fillchar=fillchar) + return self._padder( + func=func, np_func=np.strings.rjust, width=width, fillchar=fillchar + ) def zfill(self, width: int | Any) -> T_DataArray: """Pad each string in the array by prepending '0' characters. @@ -1471,6 +1532,7 @@ def contains( "Must use regular expression matching for regular expression object." ) + np_func = None if regex: if not is_compiled_re: pat = self._re_compile(pat=pat, flags=flags, case=case) @@ -1484,12 +1546,13 @@ def func(x, ipat): pat = self._stringify(pat) if case or case is None: func = lambda x, ipat: ipat in x + np_func = lambda x, ipat: np.strings.find(x, ipat) >= 0 elif self._obj.dtype.kind == "S": func = lambda x, ipat: ipat.upper() in x.upper() else: func = lambda x, ipat: ipat.casefold() in x.casefold() - return self._apply(func=func, func_args=(pat,), dtype=bool) + return self._apply(func=func, np_func=np_func, func_args=(pat,), dtype=bool) def match( self, @@ -1551,19 +1614,36 @@ def strip( ------- stripped : same type as values """ - if to_strip is not None: - to_strip = self._stringify(to_strip) - if side == "both": - func = lambda x, y: x.strip(y) + method = "strip" elif side == "left": - func = lambda x, y: x.lstrip(y) + method = "lstrip" elif side == "right": - func = lambda x, y: x.rstrip(y) + method = "rstrip" else: # pragma: no cover raise ValueError("Invalid side") - return self._apply(func=func, func_args=(to_strip,)) + if to_strip is None: + # don't pass None as argument, dask turns it into an object array + func: Callable = methodcaller(method) + func_args: tuple = () + else: + func = lambda x, y: getattr(x, method)(y) + func_args = (self._stringify(to_strip),) + + result = self._apply( + func=func, np_func=getattr(np.strings, method), func_args=func_args + ) + if ( + result.dtype.kind in "US" + and result.size > 0 + and not is_chunked_array(result.data) + ): + # np.strings keeps the width of the input dtype, shrink it to the + # longest string like the elementwise path + width = max(int(np.strings.str_len(result.data).max()), 1) + result = result.astype(f"{result.dtype.kind}{width}") + return result def lstrip(self, to_strip: str | bytes | Any = None) -> T_DataArray: """Remove leading characters. @@ -1675,7 +1755,7 @@ def repeat( Array of repeated string objects. """ func = lambda x, y: x * y - return self._apply(func=func, func_args=(repeats,)) + return self._apply(func=func, np_func=np.strings.multiply, func_args=(repeats,)) def find( self, @@ -1711,6 +1791,9 @@ def find( found : array of int """ sub = self._stringify(sub) + # don't pass None as argument, dask turns it into an object array + if end is None: + end = np.iinfo(np.int64).max if side == "left": method = "find" @@ -1720,7 +1803,12 @@ def find( raise ValueError("Invalid side") func = lambda x, isub, istart, iend: getattr(x, method)(isub, istart, iend) - return self._apply(func=func, func_args=(sub, start, end), dtype=int) + return self._apply( + func=func, + np_func=getattr(np.strings, method), + func_args=(sub, start, end), + dtype=int, + ) def rfind( self, @@ -1794,6 +1882,9 @@ def index( substring is not found """ sub = self._stringify(sub) + # don't pass None as argument, dask turns it into an object array + if end is None: + end = np.iinfo(np.int64).max if side == "left": method = "index" @@ -1803,7 +1894,12 @@ def index( raise ValueError("Invalid side") func = lambda x, isub, istart, iend: getattr(x, method)(isub, istart, iend) - return self._apply(func=func, func_args=(sub, start, end), dtype=int) + return self._apply( + func=func, + np_func=getattr(np.strings, method), + func_args=(sub, start, end), + dtype=int, + ) def rindex( self, @@ -1905,6 +2001,7 @@ def replace( if not regex and callable(repl): raise ValueError("Cannot use a callable replacement when regex=False") + np_func = None if regex: pat = self._re_compile(pat=pat, flags=flags, case=case) # ``re.sub`` interprets ``count=0`` as "replace every occurrence", @@ -1930,7 +2027,8 @@ def replace( ) else: func = lambda x, ipat, irepl, i_n: x.replace(ipat, irepl, i_n) - return self._apply(func=func, func_args=(pat, repl, n)) + np_func = np.strings.replace + return self._apply(func=func, np_func=np_func, func_args=(pat, repl, n)) def extract( self, diff --git a/xarray/tests/test_accessor_str.py b/xarray/tests/test_accessor_str.py index e71558aaca1..d49f8bca2c9 100644 --- a/xarray/tests/test_accessor_str.py +++ b/xarray/tests/test_accessor_str.py @@ -65,6 +65,30 @@ def test_dask() -> None: assert_equal(result, expected) +@requires_dask +@pytest.mark.parametrize( + "func", + [ + pytest.param(lambda s: s.strip(), id="strip"), + pytest.param(lambda s: s.rstrip(" c"), id="rstrip"), + pytest.param(lambda s: s.center(9, "*"), id="center"), + pytest.param(lambda s: s.zfill(6), id="zfill"), + pytest.param(lambda s: s.replace("b", "BB", regex=False), id="replace"), + pytest.param(lambda s: s.repeat(2), id="repeat"), + pytest.param(lambda s: s.find("b", 1), id="find"), + pytest.param(lambda s: s.contains("b", regex=False), id="contains"), + ], +) +def test_dask_string_results(dtype, func: Callable) -> None: + # results used to be truncated to a single character with dask + values = xr.DataArray([" abc ", "Hello World", "ab", ""], dims="x").astype(dtype) + + expected = func(values.str) + result = func(values.chunk(x=2).str).compute() + + assert_equal(result, expected) + + def test_count(dtype) -> None: values = xr.DataArray(["foo", "foofoo", "foooofooofommmfoo"]).astype(dtype) pat_str = dtype(r"f[o]+") @@ -2012,8 +2036,7 @@ def test_index(dtype) -> None: assert_equal(result_0, expected) assert_equal(result_1, expected) - matchtype = "subsection" if dtype == np.bytes_ else "substring" - with pytest.raises(ValueError, match=f"{matchtype} not found"): + with pytest.raises(ValueError, match="substring not found"): s.str.index("DE")