diff --git a/docs/api-reference/dataframe.md b/docs/api-reference/dataframe.md index 9e12c389f5..84b279c512 100644 --- a/docs/api-reference/dataframe.md +++ b/docs/api-reference/dataframe.md @@ -23,6 +23,7 @@ - is_empty - is_unique - item + - iter_columns - iter_rows - join - join_asof diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py index d08e6c7187..e35a330974 100644 --- a/narwhals/_arrow/dataframe.py +++ b/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]]: diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py index 8d4dc03549..7142372982 100644 --- a/narwhals/_pandas_like/dataframe.py +++ b/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, *, diff --git a/narwhals/_polars/dataframe.py b/narwhals/_polars/dataframe.py index cd18b3167c..30a7590b4c 100644 --- a/narwhals/_polars/dataframe.py +++ b/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 diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py index 347f5e772f..b9a279be57 100644 --- a/narwhals/dataframe.py +++ b/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 = ... diff --git a/tests/frame/columns_test.py b/tests/frame/columns_test.py index d3b3af3497..d444037e76 100644 --- a/tests/frame/columns_test.py +++ b/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