Repository: feature-engine/feature_engine
Pull request #669 (2023-05-20 02:16:04)

extend MatchVariables functionality to match dtypes as well
The goal of this feature is to ensure that the prediction and training dtypes match.

I have seen this problem when I have trained and saved a model pipeline and then had that model pipeline served in a production API. When that API calls the predict or transform method with some set of feature values, sometimes the dtypes of API call don't match the dtypes used in training the model.

Some examples:

1. The API call contains datetime values that are string dtypes, but the model pipeline is expecting the dtype to be a Pandas datetime dtype.

1. Some downstream user calls the API with some numbers that are saved as strings, e.g. "1234", but the model pipeline is expecting the dtype to be a float or integer.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/feature-engine/feature_engine/pull/669?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine) Report
> Merging [#669](https://app.codecov.io/gh/feature-engine/feature_engine/pull/669?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine) (a5a69f9) into [main](https://app.codecov.io/gh/feature-engine/feature_engine/commit/e73772d7529baf3599193a488a544e3d043ee6d3?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine) (e73772d) will **increase** coverage by `0.07%`.
> The diff coverage is `99.22%`.

```diff
@@            Coverage Diff             @@
##             main     #669      +/-   ##
==========================================
+ Coverage   97.91%   97.99%   +0.07%     
==========================================
  Files         100      100              
  Lines        3748     3834      +86     
  Branches      726      754      +28     
==========================================
+ Hits         3670     3757      +87     
+ Misses         29       28       -1     
  Partials       49       49              
```


| [Impacted Files](https://app.codecov.io/gh/feature-engine/feature_engine/pull/669?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine) | Coverage Δ | |
|---|---|---|
| [feature\_engine/datetime/datetime\_subtraction.py](https://app.codecov.io/gh/feature-engine/feature_engine/pull/669?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine#diff-ZmVhdHVyZV9lbmdpbmUvZGF0ZXRpbWUvZGF0ZXRpbWVfc3VidHJhY3Rpb24ucHk=) | `94.66% <ø> (ø)` | |
| [feature\_engine/tags.py](https://app.codecov.io/gh/feature-engine/feature_engine/pull/669?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine#diff-ZmVhdHVyZV9lbmdpbmUvdGFncy5weQ==) | `100.00% <ø> (ø)` | |
| [feature\_engine/selection/drop\_psi\_features.py](https://app.codecov.io/gh/feature-engine/feature_engine/pull/669?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine#diff-ZmVhdHVyZV9lbmdpbmUvc2VsZWN0aW9uL2Ryb3BfcHNpX2ZlYXR1cmVzLnB5) | `99.40% <98.55%> (+1.08%)` | :arrow_up: |
| [feature\_engine/creation/relative\_features.py](https://app.codecov.io/gh/feature-engine/feature_engine/pull/669?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine#diff-ZmVhdHVyZV9lbmdpbmUvY3JlYXRpb24vcmVsYXRpdmVfZmVhdHVyZXMucHk=) | `100.00% <100.00%> (ø)` | |
| [feature\_engine/encoding/rare\_label.py](https://app.codecov.io/gh/feature-engine/feature_engine/pull/669?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine#diff-ZmVhdHVyZV9lbmdpbmUvZW5jb2RpbmcvcmFyZV9sYWJlbC5weQ==) | `100.00% <100.00%> (ø)` | |
| [feature\_engine/imputation/drop\_missing\_data.py](https://app.codecov.io/gh/feature-engine/feature_engine/pull/669?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine#diff-ZmVhdHVyZV9lbmdpbmUvaW1wdXRhdGlvbi9kcm9wX21pc3NpbmdfZGF0YS5weQ==) | `100.00% <100.00%> (ø)` | |
| [feature\_engine/preprocessing/match\_columns.py](https://app.codecov.io/gh/feature-engine/feature_engine/pull/669?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine#diff-ZmVhdHVyZV9lbmdpbmUvcHJlcHJvY2Vzc2luZy9tYXRjaF9jb2x1bW5zLnB5) | `98.18% <100.00%> (+0.26%)` | :arrow_up: |
| [feature\_engine/selection/shuffle\_features.py](https://app.codecov.io/gh/feature-engine/feature_engine/pull/669?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine#diff-ZmVhdHVyZV9lbmdpbmUvc2VsZWN0aW9uL3NodWZmbGVfZmVhdHVyZXMucHk=) | `100.00% <100.00%> (ø)` | |
| [feature\_engine/transformation/reciprocal.py](https://app.codecov.io/gh/feature-engine/feature_engine/pull/669?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine#diff-ZmVhdHVyZV9lbmdpbmUvdHJhbnNmb3JtYXRpb24vcmVjaXByb2NhbC5weQ==) | `100.00% <100.00%> (ø)` | |
| [...\_engine/variable\_handling/\_variable\_type\_checks.py](https://app.codecov.io/gh/feature-engine/feature_engine/pull/669?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine#diff-ZmVhdHVyZV9lbmdpbmUvdmFyaWFibGVfaGFuZGxpbmcvX3ZhcmlhYmxlX3R5cGVfY2hlY2tzLnB5) | `92.30% <100.00%> (+1.00%)` | :arrow_up: |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine)

kylegilde: @solegalli , What do you think?
kylegilde: @solegalli, Any thoughts or suggestions?


solegalli: Hey @thibaultbl ! We are introducing some minor changes to this class, which, if I remember correctly, you created some time ago.

Looping you in. It would be great to have your feedback if you have a few minutes!

Thanks a lot!
kylegilde: > Hey @kylegilde 
> 
> The implementation looks good to me. Thank you for that.
> 
> And super apologies for the delayed response. I am going on holidays next Thursday, and I kind of needed to get a lot of work done before that date, so I was a bit pressed with things to get done.
> 
> Coming back to this PR; I think the main thing would be to expand the new test you added to cover more scenarios of recasting to make sure that it does not accidentally incur errors. I added a few scenarios in a separate comment. If you can think of any others, please add them.
> 
> If we don't have a test to check that the init raises an error if the new parameter is not a boolean, we need to add that one too.
> 
> In addition, it might be good to add a line to the docstrings under the init, saying that it has the option to match dtypes as well.
> 
> And last but not least, it would be great if you could add a code demo in this file: https://github.com/feature-engine/feature_engine/blob/main/docs/user_guide/preprocessing/MatchVariables.rst so that users can see how to use this transformer.
> 
> Thanks a lot for the suggestion and the implementation. I think overall these are minor changes, and then we are good to go!
> 
> Well done!

I think that I added the necessary docstring in the new line 72 and the bool check in the new line 171. I can add the demo.


kylegilde: > Hey @kylegilde
> 
> Any chance we could have the tests and demo by Thursday morning?
> 
> I am releasing a new version that fixes some bugs and adds some new functionality. it would be great if we can make this one part of the release.
> 
> I know it's short notice and I took ages to review. But I thought I'd try :p

It's unlikely at this point. 
kylegilde: @solegalli I'm not too sure how to fix the decrease in coverage error that my last commit created. 

It has to so with the "if self.verbose" line that I added.
kylegilde: @solegalli I think this is good to go. Let me know! Thanks
kylegilde: I accidentally closed it.
kylegilde: Hmmm, the failures are unrelated to my code.
GiorgioSgl: > Hmmm, the failures are unrelated to my code.

I have same problem in my PR https://github.com/feature-engine/feature_engine/pull/679, it's there some problem with the rare label encoder. 

I have tried it locally with python3.9 and no fail got. 
kylegilde: Posting the failed test here for convenience:

```python

______________________ test_correctly_ignores_nan_in_fit _______________________

df_enc_big =    var_A  ... var_C
0      A  ...     A
1      A  ...     A
2      A  ...     A
3      A  ...     A
4      A  ...     ...    G  ...     G
36     G  ...     G
37     G  ...     G
38     G  ...     G
39     G  ...     G

[40 rows x 3 columns]

    def test_correctly_ignores_nan_in_fit(df_enc_big):
    
        df = df_enc_big.copy()
        df.loc[df["var_C"] == "G", "var_C"] = np.nan
    
        encoder = RareLabelEncoder(
            tol=0.06,
            n_categories=3,
            missing_values="ignore",
        )
        encoder.fit(df)
    
        # expected:
        frequenc_cat = {
            "var_A": ["B", "D", "A", "G", "C"],
            "var_B": ["A", "D", "B", "G", "C"],
            "var_C": ["C", "D", "B", "A"],
        }
>       assert encoder.encoder_dict_ == frequenc_cat
E       AssertionError: assert {'var_A': ['B...C', 'B', 'A']} == {'var_A': ['B...D', 'B', 'A']}
E         Omitting 2 identical items, use -vv to show
E         Differing items:
E         {'var_C': ['D', 'C', 'B', 'A']} != {'var_C': ['C', 'D', 'B', 'A']}
E         Use -v to get more diff

tests/test_encoding/test_rare_label_encoder.py:141: AssertionError
___________ test_correctly_ignores_nan_in_fit_when_var_is_numerical ____________

df_enc_big =    var_A  ... var_C
0      A  ...     A
1      A  ...     A
2      A  ...     A
3      A  ...     A
4      A  ...     ...    G  ...     G
36     G  ...     G
37     G  ...     G
38     G  ...     G
39     G  ...     G

[40 rows x 3 columns]

    def test_correctly_ignores_nan_in_fit_when_var_is_numerical(df_enc_big):
    
        df = df_enc_big.copy()
        df["var_C"] = [
            1,
            1,
            1,
            1,
            2,
            2,
            2,
            2,
            2,
            2,
            3,
            3,
            3,
            3,
            3,
            3,
            3,
            3,
            3,
            3,
            4,
            4,
            4,
            4,
            4,
            4,
            4,
            4,
            4,
            4,
            5,
            5,
            6,
            6,
            np.nan,
            np.nan,
            np.nan,
            np.nan,
            np.nan,
            np.nan,
        ]
    
        encoder = RareLabelEncoder(
            tol=0.06,
            n_categories=3,
            missing_values="ignore",
            ignore_format=True,
        )
        encoder.fit(df)
    
        # expected:
        frequenc_cat = {
            "var_A": ["B", "D", "A", "G", "C"],
            "var_B": ["A", "D", "B", "G", "C"],
            "var_C": [3, 4, 2, 1],
        }
>       assert encoder.encoder_dict_ == frequenc_cat
E       AssertionError: assert {'var_A': ['B....0, 2.0, 1.0]} == {'var_A': ['B... [3, 4, 2, 1]}
E         Omitting 2 identical items, use -vv to show
E         Differing items:
E         {'var_C': [4.0, 3.0, 2.0, 1.0]} != {'var_C': [3, 4, 2, 1]}
E         Use -v to get more diff

tests/test_encoding/test_rare_label_encoder.py:225: AssertionError
kylegilde: @solegalli, 

I wonder if the test failure is related to the changes in https://github.com/feature-engine/feature_engine/pull/665

CCing @GiorgioSgl 
kylegilde: @solegalli,

I think that test failure is related to the changes in https://github.com/feature-engine/feature_engine/pull/665. I don't see that this PR merge passed all of the circleci tests.
solegalli: > @solegalli,
> 
> I think that test failure is related to the changes in #665. I don't see that this PR merge passed all of the circleci tests.

I'll have a look
kylegilde: Hmm, I reran some code where I thought I had seen some data changing, but I wasn't able to reproduce that phenomenon, so I think we are ok.

> Hi @kylegilde
> 
> Thank you so much for adding the demo and all the tests! The demo is very well intertwined with the existing docs!
> 
> I've got some small requests for the tests, but the thing I'd like to discuss is the issue of data changing when changing the data type. What's happening exactly?
> 
> Thanks a lot for your time!


kylegilde: accidentally closed the PR
kylegilde: @solegalli ,

My new feature is ready to go.

I have implemented all of your suggestions and responded to your questions. The CI failures seem related to files I did not change.
solegalli: Hey @kylegilde 

Thanks a lot for the patience till I fixed the failing tests.

I made a PR to your repo, with a rebase that should make the tests pass, plus some minor cosmetic changes to your files: https://github.com/kylegilde/feature_engine/pull/3/

If you merge over there, then it updates here and I can merge and close! 

Thanks a lot for your time and patience!
kylegilde: all set!
solegalli: Hey @kylegilde thank you for the quick turnaround. 

I had a look at the PR I made on your repo, and it is still open, which means that my tiny changes were not incorporated here. Could you please merge in your repo?

Here is the link: https://github.com/kylegilde/feature_engine/pull/3/

Patch:
diff --git a/docs/user_guide/preprocessing/MatchVariables.rst b/docs/user_guide/preprocessing/MatchVariables.rst
index 53251e6..3c5af10 100644
--- a/docs/user_guide/preprocessing/MatchVariables.rst
+++ b/docs/user_guide/preprocessing/MatchVariables.rst
@@ -145,8 +145,54 @@ And now, we transform the data with :class:`MatchVariables()`:
 Now, the transformer simultaneously added the missing columns with NA as values and
 removed the additional columns from the resulting dataset.
 
+
+However, if we look closely, the dtypes for the `sex` variable do not match. This could
+cause issues if other transformations depend upon having the correct dtypes.
+
+.. code:: python
+
+    train.sex.dtype
+
+.. code:: python
+
+    dtype('O')
+
+.. code:: python
+
+    test_tt.sex.dtype
+
+.. code:: python
+
+    dtype('float64')
+
+Set the `match_dtypes` parameter to `True` in order to align the dtypes as well.
+
+.. code:: python
+
+    match_cols_and_dtypes = MatchVariables(missing_values="ignore", match_dtypes=True)
+    match_cols_and_dtypes.fit(train)
+
+    test_ttt = match_cols_and_dtypes.transform(test_t)
+
+.. code:: python
+
+    The following variables are added to the DataFrame: ['sex', 'age']
+    The following variables are dropped from the DataFrame: ['var_b', 'var_a']
+    The sex dtype is changing from  float64 to object
+
+Now the dtype matches.
+
+.. code:: python
+
+    test_ttt.sex.dtype
+
+.. code:: python
+
+    dtype('O')
+
 By default, :class:`MatchVariables()` will print out messages indicating which variables
-were added or removed. We can switch off the messages through the parameter `verbose`.
+were added, removed and altered. We can switch off the messages through the parameter `verbose`.
+
 
 When to use the transformer
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
diff --git a/docs/whats_new/v_160.rst b/docs/whats_new/v_160.rst
index d457e29..554d34a 100644
--- a/docs/whats_new/v_160.rst
+++ b/docs/whats_new/v_160.rst
@@ -9,11 +9,13 @@ Deployed: xx September 2023
 Contributors
 ~~~~~~~~~~~~
 
+- `Kyle Gilde <https://github.com/kylegilde>`_
 - `Soledad Galli <https://github.com/solegalli>`_
 
 New functionality
 ~~~~~~~~~~~~~~~~~
 
+- `MatchVariables()` can now also match the **dtypes** of the variables (`Kyle Gilde <https://github.com/kylegilde>`_)
 - `DatetimeFeatures()` and `DatetimeSubtraction()` can now specify the format of the datetime variables (`Soledad Galli <https://github.com/solegalli>`_)
 
 Bug fixes
diff --git a/feature_engine/preprocessing/match_columns.py b/feature_engine/preprocessing/match_columns.py
index c40243d..7b74ab3 100644
--- a/feature_engine/preprocessing/match_columns.py
+++ b/feature_engine/preprocessing/match_columns.py
@@ -1,4 +1,4 @@
-from typing import List, Union
+from typing import List, Dict, Union
 
 import numpy as np
 import pandas as pd
@@ -69,6 +69,10 @@ class MatchVariables(BaseEstimator, TransformerMixin, GetFeatureNamesOutMixin):
         contain missing values. If 'ignore', missing data will be ignored when learning
         parameters or performing the transformation.
 
+    match_dtypes: bool, default=False
+        Indicates whether the dtypes observed in the train set should be applied to
+        variables in the test set.
+
     verbose: bool, default=True
         If True, the transformer will print out the names of the variables that are
         added and / or removed from the dataset.
@@ -81,6 +85,10 @@ class MatchVariables(BaseEstimator, TransformerMixin, GetFeatureNamesOutMixin):
     n_features_in_:
         The number of features in the train set used in fit.
 
+    dtype_dict_:
+        If `match_dtypes` is set to `True`, then this attribute will exist, and it will
+        contain a dictionary of variables and their corresponding dtypes.
+
     Methods
     -------
     fit:
@@ -150,15 +158,21 @@ class MatchVariables(BaseEstimator, TransformerMixin, GetFeatureNamesOutMixin):
         self,
         fill_value: Union[str, int, float] = np.nan,
         missing_values: str = "raise",
+        match_dtypes: bool = False,
         verbose: bool = True,
     ):
-
         if missing_values not in ["raise", "ignore"]:
             raise ValueError(
                 "missing_values takes only values 'raise' or 'ignore'."
                 f"Got '{missing_values} instead."
             )
 
+        if not isinstance(match_dtypes, bool):
+            raise ValueError(
+                "match_dtypes takes only booleans True and False. "
+                f"Got '{match_dtypes} instead."
+            )
+
         if not isinstance(verbose, bool):
             raise ValueError(
                 "verbose takes only booleans True and False." f"Got '{verbose} instead."
@@ -173,6 +187,7 @@ class MatchVariables(BaseEstimator, TransformerMixin, GetFeatureNamesOutMixin):
 
         self.fill_value = fill_value
         self.missing_values = missing_values
+        self.match_dtypes = match_dtypes
         self.verbose = verbose
 
     def fit(self, X: pd.DataFrame, y: pd.Series = None):
@@ -198,6 +213,9 @@ class MatchVariables(BaseEstimator, TransformerMixin, GetFeatureNamesOutMixin):
 
         self.n_features_in_ = X.shape[1]
 
+        if self.match_dtypes:
+            self.dtype_dict_: Dict = X.dtypes.to_dict()
+
         return self
 
     def transform(self, X: pd.DataFrame) -> pd.DataFrame:
@@ -243,6 +261,23 @@ class MatchVariables(BaseEstimator, TransformerMixin, GetFeatureNamesOutMixin):
 
         X = X.reindex(columns=self.feature_names_in_, fill_value=self.fill_value)
 
+        if self.match_dtypes:
+            _current_dtypes = X.dtypes.to_dict()
+            _columns_to_update = {
+                column: new_dtype
+                for column, new_dtype in self.dtype_dict_.items()
+                if new_dtype != _current_dtypes[column]
+            }
+
+            if self.verbose:
+                for column, new_dtype in _columns_to_update.items():
+                    print(
+                        f"The {column} dtype is changing from ",
+                        f"{_current_dtypes[column]} to {new_dtype}",
+                    )
+
+            X = X.astype(_columns_to_update)
+
         return X
 
     # for the check_estimator tests


Test patch:
diff --git a/tests/test_preprocessing/test_match_columns.py b/tests/test_preprocessing/test_match_columns.py
index eeeaac3..f5e6f35 100644
--- a/tests/test_preprocessing/test_match_columns.py
+++ b/tests/test_preprocessing/test_match_columns.py
@@ -13,9 +13,10 @@ _params_fill_value = [
 ]
 
 _params_allowed = [
-    ([0, 1], "ignore", True),
-    ("nan", "hola", True),
-    ("nan", "ignore", "hallo"),
+    ([0, 1], "ignore", True, True),
+    ("nan", "hola", True, True),
+    ("nan", "ignore", True, "hallo"),
+    ("nan", "ignore", "hallo", True),
 ]
 
 
@@ -59,6 +60,7 @@ def test_drop_and_add_columns(
         assert match_columns.fill_value == fill_value
     assert match_columns.verbose is True
     assert match_columns.missing_values == "ignore"
+    assert match_columns.match_dtypes is False
     # test fit attrs
     assert list(match_columns.feature_names_in_) == list(train.columns)
     assert match_columns.n_features_in_ == 6
@@ -72,7 +74,6 @@ def test_drop_and_add_columns(
 def test_columns_addition_when_more_columns_in_train_than_test(
     fill_value, expected_studies, expected_age, df_vartypes, df_na
 ):
-
     train = df_na.copy()
     test = df_vartypes.copy()
     test = test.drop("Age", axis=1)  # to add more than one column
@@ -103,6 +104,7 @@ def test_columns_addition_when_more_columns_in_train_than_test(
         assert match_columns.fill_value == fill_value
     assert match_columns.verbose is True
     assert match_columns.missing_values == "ignore"
+    assert match_columns.match_dtypes is False
     # test fit attrs
     assert list(match_columns.feature_names_in_) == list(train.columns)
     assert match_columns.n_features_in_ == 6
@@ -126,6 +128,7 @@ def test_drop_columns_when_more_columns_in_test_than_train(df_vartypes, df_na):
     assert match_columns.fill_value is np.nan
     assert match_columns.verbose is True
     assert match_columns.missing_values == "ignore"
+    assert match_columns.match_dtypes is False
     # test fit attrs
     assert list(match_columns.feature_names_in_) == list(train.columns)
     assert match_columns.n_features_in_ == 4
@@ -133,16 +136,150 @@ def test_drop_columns_when_more_columns_in_test_than_train(df_vartypes, df_na):
     pd.testing.assert_frame_equal(expected_result, transformed_df)
 
 
-@pytest.mark.parametrize("fill_value, missing_values, verbose", _params_allowed)
-def test_error_if_param_values_not_allowed(fill_value, missing_values, verbose):
+def test_match_dtypes_string_to_numbers(df_vartypes):
+    train = df_vartypes.copy().select_dtypes("number")
+    test = train.copy().astype("string")
+
+    match_columns = MatchVariables(match_dtypes=True)
+    match_columns.fit(train)
+
+    transformed_df = match_columns.transform(test)
+
+    # test init params
+    assert match_columns.match_dtypes is True
+    # test fit attrs
+    assert match_columns.dtype_dict_ == {
+        "Age": np.dtype("int64"),
+        "Marks": np.dtype("float64"),
+    }
+
+    # test transform output
+    pd.testing.assert_series_equal(train.dtypes, transformed_df.dtypes)
+    pd.testing.assert_frame_equal(transformed_df, train)
+
+
+def test_match_dtypes_numbers_to_string(df_vartypes):
+    train = df_vartypes.copy().select_dtypes("number").astype("string")
+    test = df_vartypes.copy().select_dtypes("number")
+
+    match_columns = MatchVariables(match_dtypes=True)
+    match_columns.fit(train)
+
+    transformed_df = match_columns.transform(test)
+
+    # test init params
+    assert match_columns.match_dtypes is True
+    # test fit attrs
+    assert isinstance(match_columns.dtype_dict_, dict)
+    # test transform output
+    pd.testing.assert_series_equal(train.dtypes, transformed_df.dtypes)
+    pd.testing.assert_frame_equal(transformed_df, train)
+
+
+def test_match_dtypes_string_to_datetime(df_vartypes):
+    train = df_vartypes.copy().loc[:, ["dob"]]
+    test = train.copy().astype("string")
+
+    match_columns = MatchVariables(match_dtypes=True, verbose=False)
+    match_columns.fit(train)
+
+    transformed_df = match_columns.transform(test)
+
+    # test init params
+    assert match_columns.match_dtypes is True
+    assert match_columns.verbose is False
+    # test fit attrs
+    assert match_columns.dtype_dict_ == {"dob": np.dtype("<M8[ns]")}
+    # test transform output
+    pd.testing.assert_series_equal(train.dtypes, transformed_df.dtypes)
+    pd.testing.assert_frame_equal(transformed_df, train)
+
+
+def test_match_dtypes_datetime_to_string(df_vartypes):
+    train = df_vartypes.copy().loc[:, ["dob"]].astype("string")
+    test = df_vartypes.copy().loc[:, ["dob"]]
+
+    match_columns = MatchVariables(match_dtypes=True, verbose=False)
+    match_columns.fit(train)
+
+    transformed_df = match_columns.transform(test)
+
+    # test init params
+    assert match_columns.match_dtypes is True
+    assert match_columns.verbose is False
+    # test fit attrs
+    assert isinstance(match_columns.dtype_dict_, dict)
+    # test transform output
+    pd.testing.assert_series_equal(train.dtypes, transformed_df.dtypes)
+    pd.testing.assert_frame_equal(transformed_df, train)
+
+
+def test_match_dtypes_missing_category(df_vartypes):
+    train = df_vartypes.copy().loc[:, ["Name", "City"]].astype("category")
+    test = df_vartypes.copy().loc[:, ["Name", "City"]].iloc[:-1].astype("category")
+
+    match_columns = MatchVariables(match_dtypes=True, verbose=True)
+    match_columns.fit(train)
+
+    transformed_df = match_columns.transform(test)
+
+    # test init params
+    assert match_columns.match_dtypes is True
+    assert match_columns.verbose is True
+    # test fit attrs
+    assert match_columns.dtype_dict_ == {
+        "Name": pd.CategoricalDtype(
+            categories=["jack", "krish", "nick", "tom"], ordered=False
+        ),
+        "City": pd.CategoricalDtype(
+            categories=["Bristol", "Liverpool", "London", "Manchester"], ordered=False
+        ),
+    }
+    # test transform output
+    pd.testing.assert_series_equal(train.dtypes, transformed_df.dtypes)
+    pd.testing.assert_frame_equal(transformed_df, train.iloc[:-1])
+
+
+def test_match_dtypes_extra_category(df_vartypes):
+    train = df_vartypes.copy().loc[:, ["Name", "City"]].iloc[:-1].astype("category")
+    test = df_vartypes.copy().loc[:, ["Name", "City"]].astype("category")
+
+    match_columns = MatchVariables(match_dtypes=True, verbose=True)
+    match_columns.fit(train)
+
+    transformed_df = match_columns.transform(test)
+
+    # test init params
+    assert match_columns.match_dtypes is True
+    assert match_columns.verbose is True
+    # test fit attrs
+    assert match_columns.dtype_dict_ == {
+        "Name": pd.CategoricalDtype(categories=["krish", "nick", "tom"], ordered=False),
+        "City": pd.CategoricalDtype(
+            categories=["Liverpool", "London", "Manchester"], ordered=False
+        ),
+    }
+
+    # test transform output
+    pd.testing.assert_series_equal(train.dtypes, transformed_df.dtypes)
+
+
+@pytest.mark.parametrize(
+    "fill_value, missing_values, match_dtypes, verbose", _params_allowed
+)
+def test_error_if_param_values_not_allowed(
+    fill_value, missing_values, match_dtypes, verbose
+):
     with pytest.raises(ValueError):
         MatchVariables(
-            fill_value=fill_value, missing_values=missing_values, verbose=verbose
+            fill_value=fill_value,
+            missing_values=missing_values,
+            match_dtypes=match_dtypes,
+            verbose=verbose,
         )
 
 
 def test_verbose_print_out(capfd, df_vartypes, df_na):
-
     match_columns = MatchVariables(missing_values="ignore", verbose=True)
 
     train = df_na.copy()