refactor: Generic CompliantSelector

[dangotbanned]

Tracking

• Tracking: Drop support for python 3.8 #2084
• chore(typing): enable typing checks for pyspark #2051
  • After that merges, move these into (chore(typing): Generic CompliantDataFrame #2115)
    • refactor: Generic CompliantSelector #2064 (comment)
    • refactor: Generic CompliantSelector #2064 (comment)
    • refactor: Generic CompliantSelector #2064 (comment)
    • refactor: Generic CompliantSelector #2064 (comment)
  • That will support moving these fixes out as well
    • refactor: Generic CompliantSelector #2064 (comment)
• Add DataFrame.iter_columns #2101

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

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

Related issues

• chore(typing): add CompliantExpr._version #2060
• chore: remove kwargs from selector function #2058
• refactor: deduplicate *.selectors #2057
• fix TypeVar used in (SparkLike|DuckDB)Expr  #2044
• Commitment to keep package size down #1942

Checklist

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

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

Expecting quite a nice reduction in copy/pasted code

• _arrow
• _dask
• _duckdb
• _pandas_like
• _spark_like
• Investigate fix 3.8 protocol bug

[dangotbanned]

FIXME

Seems like Protocol(s) with a constructor are not triggered for <3.11?

• 3.10
• 3.11
• 3.13

[dangotbanned]

Ah this was fixed in:

• typing.Protocol silently overrides __init__ method of delivered class python/cpython#88970

[dangotbanned]

This part was resolved in (5a36c81)

Seems 3.8 has a different issue
https://github.com/narwhals-dev/narwhals/actions/runs/13462903656/job/37622190616?pr=2064

[dangotbanned]

Important

I'm ignoring 3.8 issues for the time being.
Absolute worst-case - I'll just use regular classes.

[dangotbanned]

Looking for a solution w/

• (a4edcb7)
• (cc00fdc)
• (b650c78)

[dangotbanned]

This could be the way to go:

• refactor: try reusing Protocol's Generic)

Will be attempting something like (python/typing#572 (comment))

[dangotbanned]

Well it certainly is a solution (#2064 (comment))

😞

[dangotbanned]

@MarcoGorelli I'm looking at the lazy backends atm

Are selectors implemented for them in a lazy way?
I would've expected they'd be handled quite a bit differently to avoid materialization

narwhals/narwhals/_dask/selectors.py

Lines 41 to 51 in f0a5fbb

	def by_dtype(self: Self, dtypes: Iterable[DType | type[DType]]) -> DaskSelector:
	def func(df: DaskLazyFrame) -> list[dx.Series]:
	return [
	df._native_frame[col] for col in df.columns if df.schema[col] in dtypes
	]
	def evaluate_output_names(df: DaskLazyFrame) -> Sequence[str]:
	return [col for col in df.columns if df.schema[col] in dtypes]
	return selector(self, func, evaluate_output_names)

I'm still getting to know these backends, so maybe this is cheap to do?

[dangotbanned]

I'd probably move these to nw.typing and add docs to them.

There's lots of other places they could be used to simplify these two types:

narwhals/narwhals/_pandas_like/expr.py

Line 53 in 9387ecd

call: Callable[[PandasLikeDataFrame], Sequence[PandasLikeSeries]],

narwhals/narwhals/_pandas_like/expr.py

Line 57 in 9387ecd

evaluate_output_names: Callable[[PandasLikeDataFrame], Sequence[str]],

IMO this isn't as noisy:

narwhals/narwhals/_pandas_like/selectors.py

Lines 26 to 27 in ec5dd0c

	call: EvalSeries[PandasLikeDataFrame, PandasLikeSeries],
	evaluate_output_names: EvalNames[PandasLikeDataFrame],

[MarcoGorelli]

love this, will do a final check later

[dangotbanned]

I was originally planning to use utils.is_tracks_depth more than just in __repr__

utils.is_tracks_depth

narwhals/narwhals/utils.py

Lines 1392 to 1394 in d03cc22

	def is_tracks_depth(obj: Implementation, /) -> TypeIs[_TracksDepth]: # pragma: no cover
	# Return `True` for implementations that utilize `CompliantExpr._depth`.
	return obj.is_pandas_like() or obj in {Implementation.PYARROW, Implementation.DASK}

The idea I had was to move most of the implementation of these methods up into the protocols.

You'd only need to declare the types (e.g. ArrowSelector & ArrowExpr):

*SelectorNamespace._selector

narwhals/narwhals/_arrow/selectors.py

Lines 20 to 34 in d03cc22

	def _selector(
	self,
	call: EvalSeries[ArrowDataFrame, ArrowSeries],
	evaluate_output_names: EvalNames[ArrowDataFrame],
	/,
	) -> ArrowSelector:
	return ArrowSelector(
	call,
	depth=0,
	function_name="selector",
	evaluate_output_names=evaluate_output_names,
	alias_output_names=None,
	backend_version=self._backend_version,
	version=self._version,
	)

*Selector._to_expr

narwhals/narwhals/_arrow/selectors.py

Lines 43 to 52 in d03cc22

	def _to_expr(self: Self) -> ArrowExpr:
	return ArrowExpr(
	self._call,
	depth=self._depth,
	function_name=self._function_name,
	evaluate_output_names=self._evaluate_output_names,
	alias_output_names=self._alias_output_names,
	backend_version=self._backend_version,
	version=self._version,
	)

But I stopped short of that - because it should probably live in *Expr.

You'd have an even greater reduction in repetition, as it would account for all the arg mapping in:

narwhals/narwhals/_arrow/expr.py

Lines 34 to 151 in 140833c

	class ArrowExpr(CompliantExpr["ArrowDataFrame", ArrowSeries]):
	_implementation: Implementation = Implementation.PYARROW
	def __init__(
	self: Self,
	call: Callable[[ArrowDataFrame], Sequence[ArrowSeries]],
	*,
	depth: int,
	function_name: str,
	evaluate_output_names: Callable[[ArrowDataFrame], Sequence[str]],
	alias_output_names: Callable[[Sequence[str]], Sequence[str]] | None,
	backend_version: tuple[int, ...],
	version: Version,
	call_kwargs: dict[str, Any] | None = None,
	) -> None:
	self._call = call
	self._depth = depth
	self._function_name = function_name
	self._depth = depth
	self._evaluate_output_names = evaluate_output_names
	self._alias_output_names = alias_output_names
	self._backend_version = backend_version
	self._version = version
	self._call_kwargs = call_kwargs or {}
	def __repr__(self: Self) -> str: # pragma: no cover
	return f"ArrowExpr(depth={self._depth}, function_name={self._function_name}, "
	def __call__(self: Self, df: ArrowDataFrame) -> Sequence[ArrowSeries]:
	return self._call(df)
	def broadcast(self, kind: Literal[ExprKind.AGGREGATION, ExprKind.LITERAL]) -> Self:
	# Mark the resulting ArrowSeries with `_broadcast = True`.
	# Then, when extracting native objects, `extract_native` will
	# know what to do.
	def func(df: ArrowDataFrame) -> list[ArrowSeries]:
	results = []
	for result in self(df):
	result._broadcast = True
	results.append(result)
	return results
	return self.__class__(
	func,
	depth=self._depth,
	function_name=self._function_name,
	evaluate_output_names=self._evaluate_output_names,
	alias_output_names=self._alias_output_names,
	backend_version=self._backend_version,
	version=self._version,
	call_kwargs=self._call_kwargs,
	)
	@classmethod
	def from_column_names(
	cls: type[Self],
	*column_names: str,
	backend_version: tuple[int, ...],
	version: Version,
	) -> Self:
	from narwhals._arrow.series import ArrowSeries
	def func(df: ArrowDataFrame) -> list[ArrowSeries]:
	try:
	return [
	ArrowSeries(
	df._native_frame[column_name],
	name=column_name,
	backend_version=df._backend_version,
	version=df._version,
	)
	for column_name in column_names
	]
	except KeyError as e:
	missing_columns = [x for x in column_names if x not in df.columns]
	raise ColumnNotFoundError.from_missing_and_available_column_names(
	missing_columns=missing_columns, available_columns=df.columns
	) from e
	return cls(
	func,
	depth=0,
	function_name="col",
	evaluate_output_names=lambda _df: column_names,
	alias_output_names=None,
	backend_version=backend_version,
	version=version,
	)
	@classmethod
	def from_column_indices(
	cls: type[Self],
	*column_indices: int,
	backend_version: tuple[int, ...],
	version: Version,
	) -> Self:
	from narwhals._arrow.series import ArrowSeries
	def func(df: ArrowDataFrame) -> list[ArrowSeries]:
	return [
	ArrowSeries(
	df._native_frame[column_index],
	name=df._native_frame.column_names[column_index],
	backend_version=df._backend_version,
	version=df._version,
	)
	for column_index in column_indices
	]
	return cls(
	func,
	depth=0,
	function_name="nth",
	evaluate_output_names=lambda df: [df.columns[i] for i in column_indices],
	alias_output_names=None,
	backend_version=backend_version,
	version=version,
	)

[dangotbanned]

This PR revealed some patterns that I think could be better encoded into the Protocol definitions:

Single or multi-Implementation

• Single defines on the class
• Multi defines on the instance, and passes it around to keep track
  • Currently the *Like* classes
  • Although I think that should be revisted
    • chore(typing): enable typing checks for pyspark #2051 (comment)
    • https://discord.com/channels/1235257048170762310/1344629165407600661/1344689702753665054

Tracks depth

• As mentioned, we need that for the dask special-case
• All other lazy backends don't track depth

(Lazy|Eager)*

• I've made a start with that for *Selector
• Planning the same for *Expr in
  • feat: Spec'd-out CompliantExpr #2119
  • fix TypeVar used in (SparkLike|DuckDB)Expr  #2044 (comment)

Reuse*

• I mentioned in this PR that I was thinking about renaming to EagerOnly*
  • refactor: adding Compliant* sub-protocols #2055
• I think it would be worth having the distinction between those cases and Polars*

[dangotbanned]

@MarcoGorelli feel free to move this into another issue - just trying to connect the dots between a lot of the stuff I've been working on 🙂

[MarcoGorelli]

thanks - i'm not really sure about having _depth long-term, i think we'll need a better way to do this in the future, so i'd suggest not spending too long on that specifically

[MarcoGorelli]

thanks @dangotbanned - just a comment about moving a comment from just being the PR to putting in the code (a practice I'd encourage in general), then feel free to merge

[dangotbanned]

thanks @dangotbanned - just a comment about moving a comment from just being the PR to putting in the code (a practice I'd encourage in general), then feel free to merge

Thanks for the review @MarcoGorelli