perf: Optimize, generalize _DelayedCategories -> _DeferredIterable (#2421)

Author: dangotbanned

narwhals/_pandas_like/utils.py

@@ -7,19 +7,20 @@
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
+from typing import Literal
 from typing import Sequence
 from typing import Sized
 from typing import TypeVar
 
 import pandas as pd
 
 from narwhals._compliant.series import EagerSeriesNamespace
-from narwhals.dtypes import _DelayedCategories
 from narwhals.exceptions import ColumnNotFoundError
 from narwhals.exceptions import DuplicateError
 from narwhals.exceptions import ShapeError
 from narwhals.utils import Implementation
 from narwhals.utils import Version
+from narwhals.utils import _DeferredIterable
 from narwhals.utils import isinstance_or_issubclass
 
 T = TypeVar("T", bound=Sized)
@@ -255,9 +256,7 @@ def non_object_native_to_narwhals_dtype(native_dtype: Any, version: Version) ->
     if dtype.startswith("dictionary<"):
         return dtypes.Categorical()
     if dtype == "category":
-        return native_categorical_to_narwhals_dtype(
-            native_dtype, version, lambda: tuple(native_dtype.categories)
-        )
+        return native_categorical_to_narwhals_dtype(native_dtype, version)
     if (match_ := PATTERN_PD_DATETIME.match(dtype)) or (
         match_ := PATTERN_PA_DATETIME.match(dtype)
     ):
@@ -305,20 +304,33 @@ def object_native_to_narwhals_dtype(
 def native_categorical_to_narwhals_dtype(
     native_dtype: pd.CategoricalDtype,
     version: Version,
-    get_categories: Callable[[], tuple[str, ...]],
+    implementation: Literal[Implementation.CUDF] | None = None,
 ) -> DType:
     dtypes = version.dtypes
     if version is Version.V1:
         return dtypes.Categorical()
     if native_dtype.ordered:
-        return dtypes.Enum(_DelayedCategories(get_categories))
+        into_iter = (
+            _cudf_categorical_to_list(native_dtype)
+            if implementation is Implementation.CUDF
+            else native_dtype.categories.to_list
+        )
+        return dtypes.Enum(_DeferredIterable(into_iter))
     return dtypes.Categorical()
 
 
-def native_to_narwhals_dtype(
+def _cudf_categorical_to_list(
     native_dtype: Any,
-    version: Version,
-    implementation: Implementation,
+) -> Callable[[], list[Any]]:  # pragma: no cover
+    # NOTE: https://docs.rapids.ai/api/cudf/stable/user_guide/api_docs/api/cudf.core.dtypes.categoricaldtype/#cudf.core.dtypes.CategoricalDtype
+    def fn() -> list[Any]:
+        return native_dtype.categories.to_arrow().to_pylist()
+
+    return fn
+
+
+def native_to_narwhals_dtype(
+    native_dtype: Any, version: Version, implementation: Implementation
 ) -> DType:
     str_dtype = str(native_dtype)
 
@@ -333,8 +345,9 @@ def native_to_narwhals_dtype(
         return arrow_native_to_narwhals_dtype(native_dtype.pyarrow_dtype, version)
     if str_dtype == "category" and implementation.is_cudf():
         # https://github.com/rapidsai/cudf/issues/18536
+        # https://github.com/rapidsai/cudf/issues/14027
         return native_categorical_to_narwhals_dtype(
-            native_dtype, version, lambda: tuple(native_dtype.categories.to_pandas())
+            native_dtype, version, Implementation.CUDF
         )
     if str_dtype != "object":
         return non_object_native_to_narwhals_dtype(native_dtype, version)

narwhals/_polars/utils.py

@@ -7,18 +7,19 @@
 from typing import Iterator
 from typing import Mapping
 from typing import TypeVar
+from typing import cast
 from typing import overload
 
 import polars as pl
 
-from narwhals.dtypes import _DelayedCategories
 from narwhals.exceptions import ColumnNotFoundError
 from narwhals.exceptions import ComputeError
 from narwhals.exceptions import DuplicateError
 from narwhals.exceptions import InvalidOperationError
 from narwhals.exceptions import NarwhalsError
 from narwhals.exceptions import ShapeError
 from narwhals.utils import Version
+from narwhals.utils import _DeferredIterable
 from narwhals.utils import isinstance_or_issubclass
 
 if TYPE_CHECKING:
@@ -101,7 +102,12 @@ def native_to_narwhals_dtype(
     if isinstance_or_issubclass(dtype, pl.Enum):
         if version is Version.V1:
             return dtypes.Enum()  # type: ignore[call-arg]
-        return dtypes.Enum(_DelayedCategories(lambda: tuple(dtype.categories)))
+        categories = _DeferredIterable(
+            dtype.categories.to_list
+            if backend_version >= (0, 20, 4)
+            else lambda: cast("list[str]", dtype.categories)
+        )
+        return dtypes.Enum(categories)
     if dtype == pl.Date:
         return dtypes.Date()
     if isinstance_or_issubclass(dtype, pl.Datetime):

narwhals/dtypes.py

@@ -8,10 +8,10 @@
 from typing import Iterable
 from typing import Mapping
 
+from narwhals.utils import _DeferredIterable
 from narwhals.utils import isinstance_or_issubclass
 
 if TYPE_CHECKING:
-    from typing import Callable
     from typing import Iterator
     from typing import Sequence
 
@@ -461,22 +461,6 @@ class Categorical(DType):
     """
 
 
-class _DelayedCategories:
-    """Store callable which produces tupleified version of Enum's categories.
-
-    !!! warning
-        This class is not meant to be instantiated directly and exists
-        for internal usage only.
-    """
-
-    def __init__(self, get_categories: Callable[[], tuple[str, ...]]) -> None:
-        self.get_categories = get_categories
-
-    def __iter__(self) -> Iterator[str]:  # pragma: no cover
-        msg = "This is only provided for type-checking and should not be called"
-        raise AssertionError(msg)
-
-
 class Enum(DType):
     """A fixed categorical encoding of a unique set of strings.
 
@@ -490,10 +474,10 @@ class Enum(DType):
     """
 
     def __init__(self, categories: Iterable[str] | type[enum.Enum]) -> None:
-        self._delayed_categories: _DelayedCategories | None = None
+        self._delayed_categories: _DeferredIterable[str] | None = None
         self._cached_categories: tuple[str, ...] | None = None
 
-        if isinstance(categories, _DelayedCategories):
+        if isinstance(categories, _DeferredIterable):
             self._delayed_categories = categories
         elif isinstance(categories, type) and issubclass(categories, enum.Enum):
             self._cached_categories = tuple(member.value for member in categories)
@@ -502,10 +486,14 @@ def __init__(self, categories: Iterable[str] | type[enum.Enum]) -> None:
 
     @property
     def categories(self) -> tuple[str, ...]:
-        if self._cached_categories is None:
-            assert self._delayed_categories is not None  # noqa: S101
-            self._cached_categories = self._delayed_categories.get_categories()
-        return self._cached_categories
+        if cached := self._cached_categories:
+            return cached
+        elif delayed := self._delayed_categories:
+            self._cached_categories = delayed.to_tuple()
+            return self._cached_categories
+        else:  # pragma: no cover
+            msg = f"Internal structure of {type(self).__name__!r} is invalid."
+            raise TypeError(msg)
 
     def __eq__(self, other: object) -> bool:
         # allow comparing object instances to class

narwhals/utils.py

@@ -14,7 +14,9 @@
 from typing import Any
 from typing import Callable
 from typing import Container
+from typing import Generic
 from typing import Iterable
+from typing import Iterator
 from typing import Literal
 from typing import Protocol
 from typing import Sequence
@@ -102,7 +104,7 @@
     FrameOrSeriesT = TypeVar(
         "FrameOrSeriesT", bound=Union[LazyFrame[Any], DataFrame[Any], Series[Any]]
     )
-    _T = TypeVar("_T")
+
     _T1 = TypeVar("_T1")
     _T2 = TypeVar("_T2")
     _T3 = TypeVar("_T3")
@@ -150,6 +152,7 @@ class _StoresColumns(Protocol):
         def columns(self) -> Sequence[str]: ...
 
 
+_T = TypeVar("_T")
 NativeT_co = TypeVar("NativeT_co", covariant=True)
 CompliantT_co = TypeVar("CompliantT_co", covariant=True)
 _ContextT = TypeVar("_ContextT", bound="_FullContext")
@@ -1946,3 +1949,18 @@ def decorate(init_child: _Constructor[_T, P, R2], /) -> _Constructor[_T, P, R2]:
             raise TypeError(msg)
 
     return decorate
+
+
+class _DeferredIterable(Generic[_T]):
+    """Store a callable producing an iterable to defer collection until we need it."""
+
+    def __init__(self, into_iter: Callable[[], Iterable[_T]], /) -> None:
+        self._into_iter: Callable[[], Iterable[_T]] = into_iter
+
+    def __iter__(self) -> Iterator[_T]:
+        yield from self._into_iter()
+
+    def to_tuple(self) -> tuple[_T, ...]:
+        # Collect and return as a `tuple`.
+        it = self._into_iter()
+        return it if isinstance(it, tuple) else tuple(it)

tests/dtypes_test.py

@@ -435,6 +435,8 @@ def test_enum_categories_immutable() -> None:
     dtype = nw.Enum(["a", "b"])
     with pytest.raises(TypeError, match="does not support item assignment"):
         dtype.categories[0] = "c"  # type: ignore[index]
+    with pytest.raises(AttributeError):
+        dtype.categories = "a", "b", "c"  # type: ignore[misc]
 
 
 def test_enum_repr_pd() -> None:
@@ -444,7 +446,8 @@ def test_enum_repr_pd() -> None:
         )
     )
     dtype = df.schema["a"]
-    assert dtype.categories == ("broccoli", "cabbage")  # type: ignore[attr-defined]
+    assert isinstance(dtype, nw.Enum)
+    assert dtype.categories == ("broccoli", "cabbage")
     assert "Enum(categories=['broccoli', 'cabbage'])" in str(dtype)
 
 
@@ -458,7 +461,8 @@ def test_enum_repr_pl() -> None:
         )
     )
     dtype = df.schema["a"]
-    assert dtype.categories == ("broccoli", "cabbage")  # type: ignore[attr-defined]
+    assert isinstance(dtype, nw.Enum)
+    assert dtype.categories == ("broccoli", "cabbage")
     assert "Enum(categories=['broccoli', 'cabbage'])" in repr(dtype)
 
 

tests/utils_test.py

@@ -3,8 +3,12 @@
 import re
 import string
 from dataclasses import dataclass
+from itertools import chain
 from typing import TYPE_CHECKING
 from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import Iterator
 from typing import Protocol
 from typing import cast
 
@@ -22,6 +26,7 @@
 from narwhals.exceptions import ColumnNotFoundError
 from narwhals.utils import Implementation
 from narwhals.utils import Version
+from narwhals.utils import _DeferredIterable
 from narwhals.utils import check_column_exists
 from narwhals.utils import deprecate_native_namespace
 from narwhals.utils import parse_version
@@ -557,3 +562,21 @@ def repeat(self, n: int) -> str:
     )
     with pytest.raises(NotImplementedError, match=pattern):
         v_05.concat("never")
+
+
+def test_deferred_iterable() -> None:
+    def to_upper(it: Iterable[str]) -> Callable[[], Iterator[str]]:
+        def fn() -> Iterator[str]:
+            for el in it:
+                yield el.capitalize()
+
+        return fn
+
+    iterable = list("hello")
+    deferred_1 = _DeferredIterable(iterable.copy)
+    deferred_2 = _DeferredIterable(to_upper(iterable))
+
+    assert deferred_1.to_tuple() == tuple("hello")
+    assert next(iter(deferred_1)) == "h"
+    assert list(deferred_1) == list("hello")
+    assert "".join(chain(deferred_1, deferred_2)) == "helloHELLO"