chore: refactor EVERYTHING to do with broadcasting

[MarcoGorelli]

This...took some effort

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

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

Related issues

• Related issue #<issue number>
• Closes #<issue number>

Checklist

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

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

[MarcoGorelli]

but strings_are_column_names is the inverse of str_as_lit right?

yup! that's a fair point, thanks, might as well align with Polars' internal name

[dangotbanned]

but strings_are_column_names is the inverse of str_as_lit right?

yup! that's a fair point, thanks, might as well align with Polars' internal name

No worries
You can always make a TypeAlias with a doc if you wanna add more detail that follows the parameter around.
E.g.:

StrAsLit: TypeAlias = bool
"""<Explain me>"""

def func(str_as_lit: StrAsLit):...

Personal preference I suppose, but I like being able to get more info - but not having it always there if I'm familiar. These things seem small but do add up over time

[MarcoGorelli]

great suggestions, thanks - fancy adding a commit with them?

[dangotbanned]

great suggestions, thanks - fancy adding a commit with them?

Sure thing

[dangotbanned]

Suggested edit:

The import was blocking me using the normal suggestion.
This should be a perf improvement as well, since broadcast = False skips the rest

diff --git a/narwhals/_expression_parsing.py b/narwhals/_expression_parsing.py
index 38300ad8..09419a5d 100644
--- a/narwhals/_expression_parsing.py
+++ b/narwhals/_expression_parsing.py
@@ -8,6 +8,7 @@ from enum import auto
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
+from typing import Iterable
 from typing import Sequence
 from typing import TypedDict
 from typing import TypeVar
@@ -445,25 +446,27 @@ def infer_kind(obj: IntoExpr | _1DArray | object, *, str_as_lit: bool) -> ExprKi
     return ExprKind.LITERAL
 
 
+def _iter_n_ary_operation(
+    compliant_exprs: Iterable[CompliantExpr[Any] | object], kinds: Sequence[ExprKind]
+) -> Iterable[CompliantExpr[Any] | object]:
+    broadcast = any(kind is ExprKind.TRANSFORM for kind in kinds)
+    broadcast_kind = {ExprKind.AGGREGATION, ExprKind.LITERAL}
+    if not broadcast:
+        yield from compliant_exprs
+    else:
+        for compliant_expr, kind in zip(compliant_exprs, kinds):
+            if kind in broadcast_kind and is_compliant_expr(compliant_expr):
+                yield compliant_expr.broadcast(kind)
+            else:
+                yield compliant_expr
+
+
 def apply_n_ary_operation(
-    plx: CompliantNamespace,
-    function: Any,
-    *comparands: IntoExpr,
-    str_as_lit: bool,
+    plx: CompliantNamespace, function: Any, *comparands: IntoExpr, str_as_lit: bool
 ) -> CompliantExpr[Any]:
     compliant_exprs = (
         extract_compliant(plx, comparand, str_as_lit=str_as_lit)
         for comparand in comparands
     )
     kinds = [infer_kind(comparand, str_as_lit=str_as_lit) for comparand in comparands]
-
-    broadcast = any(kind is ExprKind.TRANSFORM for kind in kinds)
-    compliant_exprs = (
-        compliant_expr.broadcast(kind)
-        if broadcast
-        and kind in (ExprKind.AGGREGATION, ExprKind.LITERAL)
-        and is_compliant_expr(compliant_expr)
-        else compliant_expr
-        for compliant_expr, kind in zip(compliant_exprs, kinds)
-    )
-    return function(*compliant_exprs)
+    return function(*_iter_n_ary_operation(compliant_exprs, kinds))

[dangotbanned]

#2005 (comment)

@MarcoGorelli ah I forgot about how singular aliases inherit their docs (unlike unions):

I'm not as sold on my own idea now 😅

Also I was thinking this did the same thing as attribute docs, but it isn't visible everywhere like:

[dangotbanned]

Maybe you'd want some of these sets defined on ExprKind like you've got for Implementation?

narwhals/narwhals/utils.py

Lines 238 to 256 in 390d6e0

	def is_pandas_like(self: Self) -> bool:
	"""Return whether implementation is pandas, Modin, or cuDF.
	Returns:
	Boolean.
	Examples:
	>>> import pandas as pd
	>>> import narwhals as nw
	>>> df_native = pd.DataFrame({"a": [1, 2, 3]})
	>>> df = nw.from_native(df_native)
	>>> df.implementation.is_pandas_like()
	True
	"""
	return self in {
	Implementation.PANDAS,
	Implementation.MODIN,
	Implementation.CUDF,
	}

I think you'd be able to make a lot of little ergonomic changes - if there's more cases that are repeated a lot

[dangotbanned]

Could split out into is_transform(obj: Any, str_as_lit: bool)

[dangotbanned]

This is identicial to https://github.com/MarcoGorelli/narwhals/blob/5ee5ce276d54d676371a518521b7b180504dd83b/narwhals/dataframe.py#L139-L144

Maybe move into a function they both use?

[dangotbanned]

Have you thought about making this a @classmethod for ExprKind?

Maybe called from_expr or I guess technically from_into_expr (but thats a mouthful)

[EdAbati]

Note for a follow up:
This reminds me that we should fix the typing for CompliantSeriesT_co
for Spark we use Column, in duckdb Expression. They are not compliant series :/

[dangotbanned]

Opened as (#2044)

[MarcoGorelli]

Thanks all for comments, much appreciated 🙏

If people are OK with the logic of the changes, OK if we get this in and then get further refactors in as follow-ups?

[dangotbanned]

Great work @MarcoGorelli

[MarcoGorelli]

Thanks for your review and for taking such a close look, legend mate!

Merging then so we can proceed with further improvements (e.g. other typing PRs, order-dependent lazy ops, and more) before too many merge conflicts crop up