Skip to content

Commit c8ae828

Browse files
solegalliclaude
solegalli
and
claude
authored
Migrate MatchVariables to narwhals, add polars support (#1063)
* Migrate MatchVariables to narwhals, add polars support fit() and transform() take any narwhals-supported dataframe and return the same backend. pandas adds, drops and reorders the columns with a single reindex (faster than the previous drop/setitem/select and than narwhals at 10k-500k rows); other backends use with_columns + select. With match_dtypes, pandas keeps its dtypes and astype, since narwhals dtypes don't hold the categories of pandas categoricals. Other backends store narwhals dtypes and cast with narwhals, turning values outside Enum categories into nulls and parsing strings into dates, as pandas does. With polars, np.nan fill values add null Float64 columns and integer fill values add Int64 columns. The verbose messages list the variables in a fixed order: training order for added variables, input order for dropped ones. Init error messages now end with "Got {param} instead.". Rewrite the tests to the make_df conventions and add a polars example to the docstring and the user guide, refreshing its outputs. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Keep the dtypes learned by MatchVariables private and allow missing data in integer and boolean variables Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Mark polars output blocks in the user guide as text Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
1 parent b0ec8b5 commit c8ae828

3 files changed

Lines changed: 631 additions & 403 deletions

File tree

‎docs/user_guide/preprocessing/MatchVariables.rst‎

Lines changed: 77 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -86,11 +86,11 @@ We see that `sex` and `age` are no longer in the dataframe:
8686
.. code:: python
8787
8888
pclass survived sibsp parch fare cabin embarked
89-
1000 3 1 0 0 7.7500 n Q
90-
1001 3 1 2 0 23.2500 n Q
91-
1002 3 1 2 0 23.2500 n Q
92-
1003 3 1 2 0 23.2500 n Q
93-
1004 3 1 0 0 7.7875 n Q
89+
1000 3 1 0 0 7.7500 NaN Q
90+
1001 3 1 2 0 23.2500 NaN Q
91+
1002 3 1 2 0 23.2500 NaN Q
92+
1003 3 1 2 0 23.2500 NaN Q
93+
1004 3 1 0 0 7.7875 NaN Q
9494
9595
If we transform the dataframe with the dropped columns using :class:`MatchVariables()`,
9696
we see that the new dataframe contains all the variables, and those that were missing
@@ -107,13 +107,13 @@ Indeed, `sex` and `age` are back, filled with missing values:
107107

108108
.. code:: python
109109
110-
The following variables are added to the DataFrame: ['age', 'sex']
110+
The following variables are added to the DataFrame: ['sex', 'age']
111111
pclass survived sex age sibsp parch fare cabin embarked
112-
1000 3 1 NaN NaN 0 0 7.7500 n Q
113-
1001 3 1 NaN NaN 2 0 23.2500 n Q
114-
1002 3 1 NaN NaN 2 0 23.2500 n Q
115-
1003 3 1 NaN NaN 2 0 23.2500 n Q
116-
1004 3 1 NaN NaN 0 0 7.7875 n Q
112+
1000 3 1 NaN NaN 0 0 7.7500 NaN Q
113+
1001 3 1 NaN NaN 2 0 23.2500 NaN Q
114+
1002 3 1 NaN NaN 2 0 23.2500 NaN Q
115+
1003 3 1 NaN NaN 2 0 23.2500 NaN Q
116+
1004 3 1 NaN NaN 0 0 7.7875 NaN Q
117117
118118
Note how the missing columns were added back to the transformed test set, with
119119
missing values, in the position (i.e., order) in which they were in the train set.
@@ -133,11 +133,11 @@ We now have 2 extra columns, `var_a` and `var_b`, that were not present in the t
133133
.. code:: python
134134
135135
pclass survived sibsp parch fare cabin embarked var_a var_b
136-
1000 3 1 0 0 7.7500 n Q 0 0
137-
1001 3 1 2 0 23.2500 n Q 0 0
138-
1002 3 1 2 0 23.2500 n Q 0 0
139-
1003 3 1 2 0 23.2500 n Q 0 0
140-
1004 3 1 0 0 7.7875 n Q 0 0
136+
1000 3 1 0 0 7.7500 NaN Q 0 0
137+
1001 3 1 2 0 23.2500 NaN Q 0 0
138+
1002 3 1 2 0 23.2500 NaN Q 0 0
139+
1003 3 1 2 0 23.2500 NaN Q 0 0
140+
1004 3 1 0 0 7.7875 NaN Q 0 0
141141
142142
And now, we transform the data with :class:`MatchVariables()`:
143143

@@ -152,14 +152,14 @@ the additional columns from the resulting dataset:
152152

153153
.. code:: python
154154
155-
The following variables are added to the DataFrame: ['age', 'sex']
156-
The following variables are dropped from the DataFrame: ['var_b', 'var_a']
155+
The following variables are added to the DataFrame: ['sex', 'age']
156+
The following variables are dropped from the DataFrame: ['var_a', 'var_b']
157157
pclass survived sex age sibsp parch fare cabin embarked
158-
1000 3 1 NaN NaN 0 0 7.7500 n Q
159-
1001 3 1 NaN NaN 2 0 23.2500 n Q
160-
1002 3 1 NaN NaN 2 0 23.2500 n Q
161-
1003 3 1 NaN NaN 2 0 23.2500 n Q
162-
1004 3 1 NaN NaN 0 0 7.7875 n Q
158+
1000 3 1 NaN NaN 0 0 7.7500 NaN Q
159+
1001 3 1 NaN NaN 2 0 23.2500 NaN Q
160+
1002 3 1 NaN NaN 2 0 23.2500 NaN Q
161+
1003 3 1 NaN NaN 2 0 23.2500 NaN Q
162+
1004 3 1 NaN NaN 0 0 7.7875 NaN Q
163163
164164
However, if we look closely, the dtypes for the `sex` variable do not match. This could
165165
cause issues if other transformations depend upon having the correct dtypes. This is the
@@ -173,7 +173,7 @@ Which is:
173173

174174
.. code:: python
175175
176-
dtype('O')
176+
<StringDtype(storage='python', na_value=nan)>
177177
178178
And this is the dtype in the transformed test set:
179179

@@ -201,8 +201,8 @@ We see in the messages that the `sex` dtype was changed to match that of the tra
201201
.. code:: python
202202
203203
The following variables are added to the DataFrame: ['sex', 'age']
204-
The following variables are dropped from the DataFrame: ['var_b', 'var_a']
205-
The sex dtype is changing from float64 to object
204+
The following variables are dropped from the DataFrame: ['var_a', 'var_b']
205+
The sex dtype is changing from float64 to str
206206
207207
Now the dtype matches:
208208

@@ -214,11 +214,61 @@ Which is:
214214

215215
.. code:: python
216216
217-
dtype('O')
217+
<StringDtype(storage='python', na_value=nan)>
218218
219219
By default, :class:`MatchVariables()` will print out messages indicating which variables
220220
were added, removed and altered. We can switch off the messages through the parameter `verbose`.
221221

222+
Working with polars
223+
^^^^^^^^^^^^^^^^^^^
224+
225+
:class:`MatchVariables()` also works with polars dataframes, and returns a polars
226+
dataframe. With polars, the variables added with the default `fill_value` contain
227+
nulls, which is how polars represents missing data:
228+
229+
.. code:: python
230+
231+
import polars as pl
232+
from feature_engine.preprocessing import MatchVariables
233+
234+
train = pl.DataFrame({
235+
"pclass": [1, 1, 3, 2],
236+
"sex": ["female", "male", "male", "female"],
237+
"age": [29.0, 2.0, 30.0, 25.0],
238+
"fare": [211.34, 151.55, 7.75, 26.0],
239+
})
240+
241+
test = pl.DataFrame({
242+
"fare": [7.75, 23.25, 7.78],
243+
"pclass": [3, 3, 3],
244+
"var_a": [0, 0, 0],
245+
})
246+
247+
match_cols = MatchVariables(match_dtypes=True)
248+
match_cols.fit(train)
249+
250+
test_t = match_cols.transform(test)
251+
print(test_t)
252+
253+
The transformer added `sex` and `age`, removed `var_a`, sorted the variables as in the
254+
train set and cast `sex` to the string dtype that it had in the train set:
255+
256+
.. code:: text
257+
258+
The following variables are added to the DataFrame: ['sex', 'age']
259+
The following variables are dropped from the DataFrame: ['var_a']
260+
The sex dtype is changing from Float64 to String
261+
shape: (3, 4)
262+
┌────────┬──────┬──────┬───────┐
263+
│ pclass ┆ sex ┆ age ┆ fare │
264+
│ --- ┆ --- ┆ --- ┆ --- │
265+
│ i64 ┆ str ┆ f64 ┆ f64 │
266+
╞════════╪══════╪══════╪═══════╡
267+
│ 3 ┆ null ┆ null ┆ 7.75 │
268+
│ 3 ┆ null ┆ null ┆ 23.25 │
269+
│ 3 ┆ null ┆ null ┆ 7.78 │
270+
└────────┴──────┴──────┴───────┘
271+
222272
223273
When to use the transformer
224274
^^^^^^^^^^^^^^^^^^^^^^^^^^^

0 commit comments

Comments
 (0)