feat: Adds DataFrame.iter_columns

[dangotbanned]

What type of PR is this? (check all applicable)

• 💾 Refactor
• ✨ Feature
• 🐛 Bug Fix
• 🔧 Optimization
• 📝 Documentation
• ✅ Test
• 🐳 Other

Related issues

• Closes Add DataFrame.iter_columns #2101

Checklist

• Code follows style guide (ruff)
• Tests added
• Documented the changes

If you have comments or can explain your changes, please do so below

TODO

• Add examples
• Possibly include CompliantDataFrame typing update
  • See (feat: Adds DataFrame.iter_columns #2104 (comment))

[dangotbanned]

I put the test here, following the lead of iter_rows:

narwhals/tests/frame/rows_test.py

Lines 32 to 44 in b362e46

	def test_iter_rows(
	request: Any,
	constructor_eager: ConstructorEager,
	named: bool, # noqa: FBT001
	expected: list[tuple[Any, ...]] | list[dict[str, Any]],
	) -> None:
	if "cudf" in str(constructor_eager):
	request.applymarker(pytest.mark.xfail)
	data = {"a": [1, 3, 2], "_b": [4, 4, 6], "z": [7.0, 8.0, 9.0], "1": [5, 6, 7]}
	df = nw.from_native(constructor_eager(data), eager_only=True)
	result = list(df.iter_rows(named=named))
	assert result == expected

Not sure if we want anything more complex?
Maybe some consideration for a roundtrip through pyarrow and make sure the name is preserved?

[dangotbanned]

See (#2064 (comment)) regarding false-positive (PERF102)

[dangotbanned]

iter_columns and all other methods on DataFrame that include Series[Any] in their return type (e.g get_column) could benefit from a change I have in #2064

narwhals/narwhals/typing.py

Lines 72 to 86 in 1abc05a

	class CompliantDataFrame(Generic[CompliantSeriesT_co], Protocol):
	def __narwhals_dataframe__(self) -> Self: ...
	def __narwhals_namespace__(self) -> Any: ...
	def simple_select(
	self, *column_names: str
	) -> Self: ... # `select` where all args are column names.
	def aggregate(self, *exprs: Any) -> Self:
	... # `select` where all args are aggregations or literals
	# (so, no broadcasting is necessary).
	@property
	def columns(self) -> Sequence[str]: ...
	@property
	def schema(self) -> Mapping[str, DType]: ...
	def get_column(self, name: str) -> CompliantSeriesT_co: ...

On the public side, it would mean adding:

# narwhals.typing.py
SeriesT_co = TypeVar("SeriesT_co", bound="Series[Any]", covariant=True)

And making changes like this:

# narwhals.dataframe.py
from narwhals.typing import SeriesT_co

class DataFrame(BaseFrame[DataFrameT], Generic[SeriesT_co]):
    def get_column(self: Self, name: str) -> SeriesT_co: ...
    def iter_columns(self: Self) -> Iterator[SeriesT_co]: ...
    ...

So then you'd have things like:

DataFrame[pd.DataFrame, pd.Series]
DataFrame[pl.DataFrame, pl.Series]
DataFrame[pa.Table, pa.ChunkedArray]

The link between DataFrame and Series is something I've explored some more in #2055 also

[dangotbanned]

If the v1 backport merges - it might make this process easier.
We wouldn't have the complexity of stable/unstable Series not having a TypeVar free

• feat(typing): Backport generic Series to v1 #2110

[MarcoGorelli]

I would ❤️ love ❤️ to have this

[MarcoGorelli]

thanks!

I wonder - should we just implement this at the narwhals/dataframe.py level and just return

yield from (self[col] for col in self.columns)

?

This is only defined for eager frames anyway so we don't need to worry about invoking self.columns (and Polars does that anyway for this method)

[MarcoGorelli]

ah wait sorry, just remembered that for this feature you worked up to it the other way round (starting at the compliant level, then deciding to expose it)

sure, happy to add it like this then 👍

[MarcoGorelli]

thanks @dangotbanned !