typing: restore NativeSeriesT to EagerDataFrame and EagerNamespace, avoid unnecessary broadcast in eager when/then/otherwise

[MarcoGorelli]

This addresses part of #2666

Related

• typ: restore (if possible) NativeSeriesT to EagerNamespace #2666
• remove NativeSeriesT from EagerDataFrame / EagerNamespace

Description

EagerNamespace leaves me in covariance/contravariance hell, I'm trying to put together a smaller repro to isolate the issue i'm facing

OK, I've finally fully (I think) made sense of the issues with EagerNamespace and why I wasn't able to put NativeSeriesT there too

• If you define a protocol P generic in T and T doesn't appear in any of P's method's arguments, then type checkers expect P to be covariant in T
• If you define a method foo(self, a: T) -> None, then type checkers expect P to be contravariant in T. But, if you instead define it to be foo(self, a: T | Any) -> None, then that turns off the type checking for a and so type checkers expect P to be covariant in T.
• If you then overload foo and one of the overloads matches just T, then you'll run into a bad place, as type checkers simultaneously expect P to be contravariant in T and also covariant

In this case, EagerNamespace.from_native accepts NativeSeriesT | Any, and it's the | Any that's the really problematic part. It might as well just be Any, at which point we don't need EagerNamespace to be generic in NativeSeriesT if it's not used anywhere

Having said that, there might also be a bug in MyPy? https://gist.github.com/mypy-play/021d9f12a5e02e1e837767f0e6fd0e4a

I think possible solutions on that side are:

• keep EagerNamespace non-generic in NativeSeriesT (i.e. status-quo)
• make EagerNamespace generic in NativeSeriesT_contra, and then type: ignore a couple of lines in _from_native_impl when ModinSeries can't be assigned to NativeSeries

[dangotbanned]

Thanks @MarcoGorelli - love to see all this detail in the description! 😍

Will do my best to review tomorrow, just wanted to say I really appreciate this

[dangotbanned]

I haven't tried these out, but two ideas I had

Modin

Change these defs

narwhals/narwhals/_namespace.py

Lines 113 to 117 in f7d9451

	class _ModinDataFrame(Protocol):
	_pandas_class: type[pd.DataFrame]
	class _ModinSeries(Protocol):
	_pandas_class: type[pd.Series[Any]]

To this, which should make them compatible with the Native* protocols

    class _ModinDataFrame(NativeFrame, Protocol):
        _pandas_class: type[pd.DataFrame]

    class _ModinSeries(NativeSeries, Protocol):
        _pandas_class: type[pd.Series[Any]]

Variance

For the TypeVar(s) or Protocol(s) that are problematic, trying defining the TypeVar(s) using infer_variance

from typing import Protocol
from narwhals._typing_compat import TypeVar

class CompliantWhat(Protocol):
    def hello(self) -> str: ...

CompliantWhatT = TypeVar("CompliantWhatT", bound=CompliantWhat, infer_variance=True)

[MarcoGorelli]

thanks - i tried the first one but it didn't resolve the issue

[dangotbanned]

thanks - i tried the first one but it didn't resolve the issue

No worries!

I'm just moving this to draft while I'm experimenting (#2693 (comment))

I really should've followed up this TODO I left (e19d5e1#r2139770243)
A large chunk of the issue is that

[dangotbanned]

What I'm planning to do is de-deduplicate the broadcasting semantics that we have in these two, that was introduced in (#2662)

narwhals/narwhals/_pandas_like/namespace.py

Lines 321 to 327 in f7d9451

	if otherwise is None:
	when, then = align_series_full_broadcast(when, then)
	res_native = then.native.where(when.native)
	else:
	when, then, otherwise = align_series_full_broadcast(when, then, otherwise)
	res_native = then.native.where(when.native, otherwise.native)
	return then._with_native(res_native)

narwhals/narwhals/_arrow/namespace.py

Lines 274 to 282 in f7d9451

	if otherwise is None:
	when, then = align_series_full_broadcast(when, then)
	res_native = pc.if_else(
	when.native, then.native, pa.nulls(len(when.native), then.native.type)
	)
	else:
	when, then, otherwise = align_series_full_broadcast(when, then, otherwise)
	res_native = pc.if_else(when.native, then.native, otherwise.native)
	return then._with_native(res_native)

AFAICT, all we needed to do in (#2662) was replace the use of _extract_comparand - with some generic version of align_series_full_broadcast

narwhals/narwhals/_compliant/dataframe.py

Lines 432 to 434 in f7d9451

	def _extract_comparand(self, other: EagerSeriesT, /) -> Any:
	"""Extract native Series, broadcasting to `len(self)` if necessary."""
	...

Before we had something like this:

        otherwise = pa.nulls(len(when), then.type) if otherwise is None else otherwise
        return pc.if_else(when, then, otherwise)

So the step can just happen in the implemented EagerWhen.__call__ and depend on the uninplemented Eager*._align_series_full_broadcast

[dangotbanned]

I'm leaving the coverage to fail, since I really don't want to forget about not merging that method 😅

https://github.com/narwhals-dev/narwhals/actions/runs/15735541100/job/44347151400?pr=2693

Hopefully will be able to dip back in today

[dangotbanned]

Went down a bit of a rabbit hole, but hopefully we've now got:

• More functional typing
• Better docs on broadcasting
• More shared code
• Hopefully more performant broadcasting (by using (np|pa).repeat)

[MarcoGorelli]

thanks @dangotbanned , awesome additions and rewrites!

this helped notice that there was a further chance to improve by avoiding an unnecessary broadcast if otherwise_value is a scalar

[dangotbanned]

Thanks @MarcoGorelli

[dangotbanned]

Yes please 😍

[dangotbanned]

The PR probably needs a rename, but other than that 🚢 🚢 🚢

[dangotbanned]

I did this kind of aliasing on almost every use.
Haven't come up with a nice short name, that still conveys the message well

In this PR we make this change, which I like because the series part is embedded in the type you call it on

utils.align_series_full_broadcast
EagerSeries._align_full_broadcast

I dunno maybe something like one of these could work:

EagerSeries.align_all
EagerSeries().align_with  # <-- instance method
EagerSeries.broadcast

I wonder if the rust impl of polars has any similar utils with a name we can steal? 😄

[MarcoGorelli]

thanks Dan!