chore: Simplify `PandasLikeDataFrame|DaskLazyFrame.join` method #2511

FBruzzesi · 2025-05-08T07:56:20Z

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

Related issues

Closes Simplify *DataFrame.join #2488

Checklist

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

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

Reduces the scope by creating _<join_strategy>_join method, for each join strategy, to be called in join

dangotbanned · 2025-05-09T10:07:14Z

@FBruzzesi would you be open to flipping the method names around?

Most of the private ones that are related share a common prefix.
That helps them show up together in autocomplete 🙂

I do understand why e.g (anti_join, left_join) make a lot of sense in isolation, but another way to think about it is:

DataFrame.join(how="anti") -> CompliantDataFrame._join_anti
DataFrame.join(how="left") -> CompliantDataFrame._join_left

Looking at it that way, we're just embedding the how as a specialization of join via a suffix

`DaskLazyFrame`

`PandasLikeDataFrame`

FBruzzesi · 2025-05-10T07:23:32Z

@dangotbanned moved naming to _join_<method>

dangotbanned · 2025-05-10T10:39:09Z

@dangotbanned moved naming to _join_<method>

Thanks @FBruzzesi will try to do a proper review today 🙏

This should provide the narrowing needed for `@overload`(s)

dangotbanned · 2025-05-10T15:02:57Z

narwhals/dataframe.py

+        if how == "cross":
+            if left_on is not None or right_on is not None or on is not None:
+                msg = "Can not pass `left_on`, `right_on` or `on` keys for cross join"
+                raise ValueError(msg)
+            result = compliant.join(
+                other, how=how, left_on=None, right_on=None, suffix=suffix
            )
-        )
+        elif on is None:
+            if left_on is None or right_on is None:
+                msg = f"Either (`left_on` and `right_on`) or `on` keys should be specified for {how}."
+                raise ValueError(msg)
+            if len(left_on) != len(right_on):
+                msg = "`left_on` and `right_on` must have the same length."
+                raise ValueError(msg)
+            result = compliant.join(
+                other, how=how, left_on=left_on, right_on=right_on, suffix=suffix
+            )
+        else:
+            if left_on is not None or right_on is not None:
+                msg = f"If `on` is specified, `left_on` and `right_on` should be None for {how}."
+                raise ValueError(msg)
+            result = compliant.join(
+                other, how=how, left_on=on, right_on=on, suffix=suffix
+            )
+        return self._with_compliant(result)


We've now got 3 distinct paths, which come from:

how: Literal['cross'], on: None, left_on: None, right_on: None

how: Literal['inner', 'left', 'full', 'semi', 'anti'], on: None, left_on: list[str], right_on: list[str]

how: Literal['inner', 'left', 'full', 'semi', 'anti'], on: list[str], left_on: None, right_on: None

Maybe we can utilize this to avoid the double-dip validation like in (#2000 (comment))?

Edit

I need to remember my own advice 🤦‍♂️

me yesterday on discord 😄

Oh I see!
So a few different ways to solve:

Mutually exclusive @overloads
i. So you'd specify None is only allowed for both *_on when how="cross"
ii. However these can be tricky to get right, especially when there's a lot of parameters

Raise an error instead of asserting when an invariant is broken (529c196)
i. This is the easiest solution, but often feels like you're doing something unnecessary

Split incompatible signatures into distinct methods/functions
i. This is my preferred solution and what I'm usually hinting towards when I suggest adding multiple @classmethods

A follow-up could look at option 3, which would be making some CompliantJoin class, then having:

CompliantJoin.left

CompliantJoin.inner

...

Maybe the calls from BaseFrame would look like:

compliant.join(**kwds).left() compliant.join(**kwds).inner()

Or just have the one join, which makes the other calls internally, to keep it simple from the outside:

compliant.join(**kwds) self.left()

Would probably be easier to share code between pandas and dask that way as well 🙂

A follow-up could look at option 3, which would be making some CompliantJoin class, then having:

CompliantJoin.left

CompliantJoin.inner

...

Maybe the calls from BaseFrame would look like:

compliant.join(**kwds).left() compliant.join(**kwds).inner()

Or just have the one join, which makes the other calls internally, to keep it simple from the outside:

compliant.join(**kwds) self.left()

Would probably be easier to share code between pandas and dask that way as well 🙂

This seems like an interesting idea but I am not sure I get it right and/or I would end up with the expected changes.

Especially:

Where would CompliantJoin exist and how would it interact with other compliant classes?

and

Maybe the calls from BaseFrame would look like:

compliant.join(**kwds).left() compliant.join(**kwds).inner()

Would this require some special path for polars?

@FBruzzesi you could think of it like these:

CompliantGroupBy

CompliantWhen

feat: Improve DataFrame.__getitem__ consistency #2393

Or similar to how the Expr, Series namespace classes work.

Essentially, we define some common parts and then refine from there per-backend

Would this require some special path for polars?

Yeah most likely, polars gets to skip our abstractions in a lot of places and this would be one of those cases I'd assume.

Anyways, this is very much just the kernel of an idea 😅
The goal would be sharing more code between backends 😎

Should we... keep that as a follow up then? 😅

The goal would be sharing more code between backends 😎

Yep this is much loud and clear! Pandas-like and Dask can definitly share some functionalities eventually :)

Should we... keep that as a follow up then? 😅

Yeah for sure!

I started with that in (#2511 (comment)) but clearly got carried away with my rambling 😄

dangotbanned

Thanks @FBruzzesi - looking a lot better! 🥳

I've got a few suggestions on the pandas one, but I think the same would apply for dask as well

narwhals/_pandas_like/dataframe.py

FBruzzesi · 2025-05-10T17:43:04Z

Thanks @FBruzzesi - looking a lot better! 🥳

I've got a few suggestions on the pandas one, but I think the same would apply for dask as well

Thanks @dangotbanned - great work you've done!
I will have time later today to go through each point

dangotbanned · 2025-05-10T18:56:14Z

I was thinking about simplifying BaseFrame.join_asof in a similar way to (1c1daae)

I thought I'd take a look at the polars version(s).

They're simpler, and accept more parameters and types 🤔

Maybe they simplified it at some point?

narwhals/_dask/dataframe.py

dangotbanned · 2025-05-24T13:12:57Z

narwhals/_pandas_like/dataframe.py

+    def _join_semi(
+        self, other: Self, *, left_on: Sequence[str], right_on: Sequence[str]
+    ) -> Self:
+        other_native = self._join_filter_rename(
+            other=other,
+            columns_to_select=list(right_on),
+            columns_mapping=dict(zip(right_on, left_on)),
+        )


This is probably my bad for the huge suggestion in (#2511 (comment)) 😅

So these lines I wrote as:

other_native = self._join_filter_rename(other, left_on, right_on)

And then had the conversion handled inside that method:

... return rename( select_columns_by_name( other.native, list(right_on), backend_version, implementation ), columns=dict(zip(right_on, left_on)), ... )

With that small tweak, we avoid needing to keep this part synced with the same lines here 😎

narwhals/narwhals/_pandas_like/dataframe.py

Lines 668 to 672 in 622a578

other_native = self._join_filter_rename(

other=other,

columns_to_select=list(right_on),

columns_mapping=dict(zip(right_on, left_on)),

)

@dangotbanned please forgive me - I didn't get your point here 🥲

dangotbanned · 2025-05-24T20:33:07Z

narwhals/_pandas_like/dataframe.py

+    def _join_filter_rename(
+        self, other: Self, columns_to_select: list[str], columns_mapping: dict[str, str]
+    ) -> pd.DataFrame:
+        """Helper function to avoid creating extra columns and row duplication.


@FBruzzesi in (#2511 (comment)) I'm saying that this version is more duplicated than what I had in the original suggestion, which was:

def _join_filter_rename( self, other: Self, left_on: Sequence[str], right_on: Sequence[str] ) -> pd.DataFrame: """Rename to avoid creating extra columns in `"anti"`, `"semi"` join, and avoids potential rows duplication.""" implementation = self._implementation backend_version = self._backend_version return rename( select_columns_by_name( other.native, list(right_on), backend_version, implementation ), columns=dict(zip(right_on, left_on)), implementation=implementation, backend_version=backend_version, ).drop_duplicates()

The list(right_on) and dict(zip(right_on, left_on)) parts are defined in one place - rather than repeated each time you call _join_filter_rename

Oh I see, it has been a while since this PR, I didn't spot it in the other comment.
I see what you mean, it a similar conundrum to #2495, the signature with columns_to_select and columns_mapping in my opinion is much more intuitive to understand what happens within the method without jumping into it, on the other hand, yes, we end up creating and passing these objects externally of the method itself.
I really don't know 🤔

FBruzzesi added 5 commits May 5, 2025 22:01

chore: Simplify PandasLikeDataFrame and DaskLazyFrame join

cf7c7fe

format

180a35d

Merge branch 'main' into chore/simplify-join

d9e4f56

cleanup signature, overload still fails

047208a

minors

b17b216

FBruzzesi added the internal label May 8, 2025

move pragma: no cover

b21403a

FBruzzesi marked this pull request as ready for review May 8, 2025 08:25

FBruzzesi and others added 2 commits May 9, 2025 15:27

rename from _<strategy>_join to _join_<strategy>

5e8db52

Merge branch 'main' into chore/simplify-join

0fa2578

FBruzzesi added pandas-like Issue is related to pandas-like backend dask Issue is related to dask backend labels May 10, 2025

dangotbanned self-requested a review May 10, 2025 14:10

dangotbanned added 2 commits May 10, 2025 15:55

refactor: Rearrange BaseFrame.join

1c1daae

This should provide the narrowing needed for `@overload`(s)

Merge remote-tracking branch 'upstream/main' into chore/simplify-join

cf830d3

dangotbanned reviewed May 10, 2025

View reviewed changes

dangotbanned mentioned this pull request May 10, 2025

feat: IbisLazyFrame support #2000

Merged

10 tasks

dangotbanned reviewed May 10, 2025

View reviewed changes

low hanging feedback

ae64d74

FBruzzesi commented May 10, 2025

View reviewed changes

narwhals/_dask/dataframe.py Outdated Show resolved Hide resolved

FBruzzesi and others added 5 commits May 10, 2025 21:33

use left join in anti before filtering

3cd567f

fix Modin anti join

6068da1

factor out with _join_filter_rename helper method

a0fca3d

Merge branch 'main' into chore/simplify-join

01cb3f6

solve conflicts

3ade37f

FBruzzesi requested a review from dangotbanned May 24, 2025 09:12

Merge branch 'main' into chore/simplify-join

622a578

dangotbanned reviewed May 24, 2025

View reviewed changes

FBruzzesi added 2 commits June 1, 2025 20:29

Merge branch 'main' into chore/simplify-join

23adb84

Merge branch 'main' into chore/simplify-join

de09b93

	other_native = self._join_filter_rename(
	other=other,
	columns_to_select=list(right_on),
	columns_mapping=dict(zip(right_on, left_on)),
	)

chore: Simplify PandasLikeDataFrame|DaskLazyFrame.join method #2511

Are you sure you want to change the base?

chore: Simplify PandasLikeDataFrame|DaskLazyFrame.join method #2511

Uh oh!

Conversation

FBruzzesi commented May 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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

Related issues

Checklist

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

Uh oh!

dangotbanned commented May 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

DaskLazyFrame

PandasLikeDataFrame

Uh oh!

FBruzzesi commented May 10, 2025

Uh oh!

dangotbanned commented May 10, 2025

Uh oh!

dangotbanned May 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Edit

Uh oh!

dangotbanned May 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dangotbanned left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

FBruzzesi commented May 10, 2025

Uh oh!

dangotbanned commented May 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

chore: Simplify `PandasLikeDataFrame|DaskLazyFrame.join` method #2511

chore: Simplify `PandasLikeDataFrame|DaskLazyFrame.join` method #2511

FBruzzesi commented May 8, 2025 •

edited

Loading

dangotbanned commented May 9, 2025 •

edited

Loading

`DaskLazyFrame`

`PandasLikeDataFrame`

dangotbanned May 10, 2025 •

edited

Loading

dangotbanned May 10, 2025 •

edited

Loading

dangotbanned commented May 10, 2025 •

edited

Loading