Repository: narwhals-dev/narwhals (python)

Issue:
perf: avoid doing work when mapping ordered categorical to Enum
Currently, when mapping an ordered categorical to Enum, we call `tuple` on the categories, or even `to_pandas()` on the categories for cudf

This should be unnecessary, I think we should only compute the categories as a tuple when it becomes necessary:
- when inspecting `dtype.categories`
- when printing the dtype

Pull request:
perf: Optimize, generalize `_DelayedCategories` -> `_DeferredIterable`
Closes #2412

<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [x] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [x] 🔧 Optimization
- [ ] 📝 Documentation
- [x] ✅ Test
- [ ] 🐳 Other

## Related issues

- #2414
- https://github.com/narwhals-dev/narwhals/pull/2414#issuecomment-2823734330

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below
Related to the name change from `_Delayed*` -> `_Deferred`.
I think that's a more common term for a lazy operation. When I see *delay*, it makes me think of time delay instead

- https://peps.python.org/pep-0749/
- https://github.com/ibis-project/ibis/blob/bd926f0fbda023124453b824aa24dc4a50e27d72/ibis/common/deferred.py
- https://discuss.python.org/t/backquotes-for-deferred-expression/70577

Patch:
diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py
index 3d241de2..9ad4efa2 100644
--- a/narwhals/_pandas_like/utils.py
+++ b/narwhals/_pandas_like/utils.py
@@ -7,6 +7,7 @@ from itertools import chain
 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
@@ -14,12 +15,12 @@ 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)
diff --git a/narwhals/_polars/utils.py b/narwhals/_polars/utils.py
index e9af4c01..8390a540 100644
--- a/narwhals/_polars/utils.py
+++ b/narwhals/_polars/utils.py
@@ -7,11 +7,11 @@ from typing import Iterable
 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
@@ -19,6 +19,7 @@ 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):
diff --git a/narwhals/dtypes.py b/narwhals/dtypes.py
index 0392ca14..e4c0511e 100644
--- a/narwhals/dtypes.py
+++ b/narwhals/dtypes.py
@@ -8,10 +8,10 @@ from typing import TYPE_CHECKING
 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 @@ class Enum(DType):
 
     @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
diff --git a/narwhals/utils.py b/narwhals/utils.py
index 092c58fa..df7b7ff0 100644
--- a/narwhals/utils.py
+++ b/narwhals/utils.py
@@ -14,7 +14,9 @@ from typing import TYPE_CHECKING
 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 @@ if TYPE_CHECKING:
     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 @@ if TYPE_CHECKING:
         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 inherit_doc(
             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)


Test patch:
diff --git a/tests/dtypes_test.py b/tests/dtypes_test.py
index 033f2e44..d1f403b5 100644
--- a/tests/dtypes_test.py
+++ b/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)
 
 
diff --git a/tests/utils_test.py b/tests/utils_test.py
index 72e975af..4339bd91 100644
--- a/tests/utils_test.py
+++ b/tests/utils_test.py
@@ -3,8 +3,12 @@ from __future__ import annotations
 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 @@ import narwhals.stable.v1 as nw_v1
 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 test_requires() -> None:
     )
     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"