feat: Adds DataFrame.iter_columns (#2104)

Author: dangotbanned

docs/api-reference/dataframe.md

@@ -23,6 +23,7 @@
         - is_empty
         - is_unique
         - item
+        - iter_columns
         - iter_rows
         - join
         - join_asof

narwhals/_arrow/dataframe.py

@@ -153,6 +153,17 @@ def rows(self: Self, *, named: bool) -> list[tuple[Any, ...]] | list[dict[str, A
             return list(self.iter_rows(named=False, buffer_size=512))  # type: ignore[return-value]
         return self._native_frame.to_pylist()
 
+    def iter_columns(self) -> Iterator[ArrowSeries]:
+        from narwhals._arrow.series import ArrowSeries
+
+        for name, series in zip(self.columns, self._native_frame.itercolumns()):
+            yield ArrowSeries(
+                series,
+                name=name,
+                backend_version=self._backend_version,
+                version=self._version,
+            )
+
     def iter_rows(
         self: Self, *, named: bool, buffer_size: int
     ) -> Iterator[tuple[Any, ...]] | Iterator[dict[str, Any]]:

narwhals/_pandas_like/dataframe.py

@@ -341,6 +341,15 @@ def rows(self: Self, *, named: bool) -> list[tuple[Any, ...]] | list[dict[str, A
 
         return self._native_frame.to_dict(orient="records")
 
+    def iter_columns(self) -> Iterator[PandasLikeSeries]:
+        for _name, series in self._native_frame.items():  # noqa: PERF102
+            yield PandasLikeSeries(
+                series,
+                implementation=self._implementation,
+                backend_version=self._backend_version,
+                version=self._version,
+            )
+
     def iter_rows(
         self: Self,
         *,

narwhals/_polars/dataframe.py

@@ -2,6 +2,7 @@
 
 from typing import TYPE_CHECKING
 from typing import Any
+from typing import Iterator
 from typing import Literal
 from typing import Sequence
 from typing import overload
@@ -232,6 +233,14 @@ def get_column(self: Self, name: str) -> PolarsSeries:
             version=self._version,
         )
 
+    def iter_columns(self) -> Iterator[PolarsSeries]:
+        from narwhals._polars.series import PolarsSeries
+
+        for series in self._native_frame.iter_columns():
+            yield PolarsSeries(
+                series, backend_version=self._backend_version, version=self._version
+            )
+
     @property
     def columns(self: Self) -> list[str]:
         return self._native_frame.columns

narwhals/dataframe.py

@@ -1235,6 +1235,37 @@ def rows(
         """
         return self._compliant_frame.rows(named=named)  # type: ignore[no-any-return]
 
+    def iter_columns(self: Self) -> Iterator[Series[Any]]:
+        """Returns an iterator over the columns of this DataFrame.
+
+        Yields:
+            A Narwhals Series, backed by a native series.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import narwhals as nw
+            >>> df_native = pd.DataFrame({"foo": [1, 2], "bar": [6.0, 7.0]})
+            >>> iter_columns = nw.from_native(df_native).iter_columns()
+            >>> next(iter_columns)
+            ┌───────────────────────┐
+            |    Narwhals Series    |
+            |-----------------------|
+            |0    1                 |
+            |1    2                 |
+            |Name: foo, dtype: int64|
+            └───────────────────────┘
+            >>> next(iter_columns)
+            ┌─────────────────────────┐
+            |     Narwhals Series     |
+            |-------------------------|
+            |0    6.0                 |
+            |1    7.0                 |
+            |Name: bar, dtype: float64|
+            └─────────────────────────┘
+        """
+        for series in self._compliant_frame.iter_columns():
+            yield self._series(series, level=self._level)
+
     @overload
     def iter_rows(
         self: Self, *, named: Literal[False], buffer_size: int = ...

tests/frame/columns_test.py

@@ -8,6 +8,9 @@
 
 if TYPE_CHECKING:
     from tests.utils import Constructor
+    from tests.utils import ConstructorEager
+
+data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8.0, 9.0]}
 
 
 @pytest.mark.filterwarnings("ignore:Determining|Resolving.*")
@@ -17,3 +20,10 @@ def test_columns(constructor: Constructor) -> None:
     result = df.columns
     expected = ["a", "b", "z"]
     assert result == expected
+
+
+def test_iter_columns(constructor_eager: ConstructorEager) -> None:
+    df = nw.from_native(constructor_eager(data), eager_only=True)
+    expected = df.to_dict(as_series=True)
+    result = {series.name: series for series in df.iter_columns()}
+    assert result == expected