Skip to content

ENH polars support for InterpolationJoiner#1016

Merged
jeromedockes merged 47 commits into
skrub-data:mainfrom
TheooJ:interpolation_joiner_polars
Jul 26, 2024
Merged

ENH polars support for InterpolationJoiner#1016
jeromedockes merged 47 commits into
skrub-data:mainfrom
TheooJ:interpolation_joiner_polars

Conversation

@TheooJ

@TheooJ TheooJ commented Jul 25, 2024

Copy link
Copy Markdown
Contributor

Closes #897

>>> import polars as pl
>>> from sklearn.neighbors import KNeighborsRegressor
>>> from skrub import InterpolationJoiner

>>> buildings = pl.DataFrame(
...     {"latitude": [1.0, 2.0], "longitude": [1.0, 2.0], "n_stories": [3, 7]}
... )
>>> annual_avg_temp = pl.DataFrame(
...     {
...         "latitude": [1.2, 0.9, 1.4, 1.7, 5.0],
...         "longitude": [0.8, 1.1, 1.8, 1.8, 5.0],
...         "avg_temp": [10.0, 11.0, 15.0, 16.0, 20.0],
...     }
... )

>>> joiner = InterpolationJoiner(
...     annual_avg_temp,
...     key=["latitude", "longitude"],
...     regressor=KNeighborsRegressor(2),
... )
>>> joiner.fit_transform(buildings)

┌──────────┬───────────┬───────────┬──────────┐
│ latitude ┆ longitude ┆ n_stories ┆ avg_temp │
│ ---      ┆ ---       ┆ ---       ┆ ---      │
│ f64      ┆ f64       ┆ i64       ┆ f64      │
╞══════════╪═══════════╪═══════════╪══════════╡
│ 1.0      ┆ 1.0       ┆ 3         ┆ 10.5     │
│ 2.0      ┆ 2.0       ┆ 7         ┆ 15.5     │
└──────────┴───────────┴───────────┴──────────┘

As discussed with @jeromedockes, if the joined dataframe could contain duplicate column names, we add random suffixes and don't raise -- to be consistent with the behavior of the other joiners

@TheooJ TheooJ marked this pull request as draft July 25, 2024 09:49
@TheooJ TheooJ marked this pull request as ready for review July 25, 2024 14:36

@jeromedockes jeromedockes left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

a first pass

Comment thread skrub/_dataframe/tests/test_common.py
Comment thread skrub/_dataframe/_common.py
Comment thread skrub/tests/test_interpolation_join.py
Comment thread skrub/tests/test_interpolation_join.py
).fit_transform(main)
assert_array_equal(ns.to_list(ns.col(join, "C")), [10, 11, 12])

join = InterpolationJoiner(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

here why is this check not necessary anymore?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We were just fitting on two different column names ("A" and "rB"). Now the test checks that the double fitting returns unique column names. Other tests check fitting on left + right keys

assert_array_equal(ns.column_names(join), ["A", "B", "B_aux"])


def test_mismatched_indexes(df_module):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why do we remove this test?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The case where we join [0, 1] on [1, 0] is tested in test_fit_on_none. I can separate into two tests if necessary

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I readded a test (for pandas only). Note that because of concat_horizontal, the initial index is reset. We could store it and reapply it in the case of pandas

It's possible to store the index of the first arg of concat_horizontal and set the output to have this index, WDYT ?

@jeromedockes

Copy link
Copy Markdown
Member

@TheooJ could you do git commit --allow-empty -m '[doc build]' and push? so that we can see the rendered example

@jeromedockes jeromedockes left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks a lot @TheooJ I think it's almost ready

Comment thread skrub/_dataframe/_common.py
Comment thread skrub/_dataframe/tests/test_common.py
Comment thread skrub/_interpolation_joiner.py
Comment thread skrub/_interpolation_joiner.py Outdated
Comment thread skrub/_interpolation_joiner.py Outdated
Comment thread skrub/_interpolation_joiner.py Outdated
Comment thread skrub/_interpolation_joiner.py Outdated

@jeromedockes jeromedockes left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

very cool, thanks @TheooJ !

@jeromedockes jeromedockes merged commit 917e6a2 into skrub-data:main Jul 26, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

InterpolationJoiner - polars

2 participants