Repository: dirty-cat/dirty_cat

Issue:
Fetching functions' style is not great
![image](https://user-images.githubusercontent.com/33222654/190182237-17b1a4c5-a0fb-404d-8af8-05fdf30f8d79.png)

Two issue I notice:
- The title is overlapping with the logo (linked to #349)
- The quote doesn't work
- There is no link to `DatasetInfoOnly` and `DatasetAll` (linked to #330)


Discussion:
GaelVaroquaux: Can you resolve the conflicts please, so that CI runs and we can check the generated docs
GaelVaroquaux: Darn, I made a comment that got lost (probably from my mobile phone).

So, one problem is that the links of DatasetAll and DatasetInfoOnly points nowhere. For them to point somewhere, there needs to be a page where these objects are documented in the docs (as the GapEnoder is documented on the front page). I would prefer not to add them on the front page, and to create a dedicated page for them. This page does not need to be linked anywhere and should focus on documenting those data structures with only a sentence in the beginning of the page.

Patch:
diff --git a/.github/workflows/codeql-analysis.yml b/.github/workflows/codeql-analysis.yml
index 40e31c3..268b600 100644
--- a/.github/workflows/codeql-analysis.yml
+++ b/.github/workflows/codeql-analysis.yml
@@ -13,10 +13,10 @@ name: "CodeQL"
 
 on:
   push:
-    branches: [ "master" ]
+    branches: [ "main" ]
   pull_request:
     # The branches below must be a subset of the branches above
-    branches: [ "master" ]
+    branches: [ "main" ]
   schedule:
     - cron: '23 0 * * 4'
 
diff --git a/CHANGES.rst b/CHANGES.rst
index d3a8f13..96fbf2e 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -11,15 +11,12 @@ Major changes
 * **datasets.fetching**: contains a new function :func:`fetch_world_bank_indicator` that can be
     used to download any indicator from the World Bank Open Data platform. It only needs the
     indicator ID that can be found on the website.
-
-Major changes
--------------
-
 * Unnecessary API has been made private: everything (files, functions, classes)
   starting with an underscore shouldn't be imported in your code.
 
 Minor changes
 -------------
+* Removed example `Fitting scalable, non-linear models on data with dirty categories`.
 
 Bug fixes
 ---------
diff --git a/README.rst b/README.rst
index 26a7a14..c3165d4 100644
--- a/README.rst
+++ b/README.rst
@@ -11,8 +11,8 @@ dirty_cat
 .. |py_ver| image:: https://img.shields.io/pypi/pyversions/dirty_cat
 .. |pypi_var| image:: https://img.shields.io/pypi/v/dirty_cat?color=informational
 .. |pypi_dl| image:: https://img.shields.io/pypi/dm/dirty_cat
-.. |codecov| image:: https://img.shields.io/codecov/c/github/dirty-cat/dirty_cat/master
-.. |circleci| image:: https://img.shields.io/circleci/build/github/dirty-cat/dirty_cat/master?label=CircleCI
+.. |codecov| image:: https://img.shields.io/codecov/c/github/dirty-cat/dirty_cat/main
+.. |circleci| image:: https://img.shields.io/circleci/build/github/dirty-cat/dirty_cat/main?label=CircleCI
 .. |Black| image:: https://img.shields.io/badge/code%20style-black-000000.svg
 
 dirty_cat is a Python module for machine-learning on dirty categorical variables.
diff --git a/build_tools/circle/build_doc.sh b/build_tools/circle/build_doc.sh
index 1229ff9..8fea51e 100755
--- a/build_tools/circle/build_doc.sh
+++ b/build_tools/circle/build_doc.sh
@@ -132,7 +132,7 @@ python setup.py develop
 
 #if [[ "$CIRCLE_BRANCH" =~ ^master$ && -z "$CI_PULL_REQUEST" ]]
 #then
-#    # List available documentation versions if on master
+#    # List available documentation versions if on default branch
 #    python build_tools/circle/list_versions.py > doc/versions.rst
 #fi
 
diff --git a/build_tools/circle/checkout_merge_commit.sh b/build_tools/circle/checkout_merge_commit.sh
index 010a6a0..d9860b0 100755
--- a/build_tools/circle/checkout_merge_commit.sh
+++ b/build_tools/circle/checkout_merge_commit.sh
@@ -1,9 +1,9 @@
 #!/bin/bash
 
 
-# Add `master` branch to the update list.
+# Add `main` branch to the update list.
 # Otherwise CircleCI will give us a cached one.
-FETCH_REFS="+master:master"
+FETCH_REFS="+main:main"
 
 # Update PR refs for testing.
 if [[ -n "${CIRCLE_PR_NUMBER}" ]]
@@ -20,13 +20,13 @@ if [[ -n "${CIRCLE_PR_NUMBER}" ]]
 then
     git checkout -qf "pr/${CIRCLE_PR_NUMBER}/merge" || (
         echo Could not fetch merge commit. >&2
-        echo There may be conflicts in merging PR \#${CIRCLE_PR_NUMBER} with master. >&2;
+        echo There may be conflicts in merging PR \#${CIRCLE_PR_NUMBER} with main. >&2;
         exit 1)
 fi
 
 # Check for merge conflicts.
 if [[ -n "${CIRCLE_PR_NUMBER}" ]]
 then
-    git branch --merged | grep master > /dev/null
+    git branch --merged | grep main > /dev/null
     git branch --merged | grep "pr/${CIRCLE_PR_NUMBER}/head" > /dev/null
 fi
diff --git a/codecov.yml b/codecov.yml
index 726fcaf..8faf740 100644
--- a/codecov.yml
+++ b/codecov.yml
@@ -5,7 +5,7 @@ coverage:
   status:
     project:
       default:
-        # Commits pushed to master should not make the overall
+        # Commits pushed to the main branch should not make the overall
         # project coverage decrease by more than 1%:
         target: auto
         threshold: 1%
diff --git a/dirty_cat/_fuzzy_join.py b/dirty_cat/_fuzzy_join.py
index 1330ba2..67d2b75 100644
--- a/dirty_cat/_fuzzy_join.py
+++ b/dirty_cat/_fuzzy_join.py
@@ -24,16 +24,16 @@ from sklearn.neighbors import NearestNeighbors
 def fuzzy_join(
     left: pd.DataFrame,
     right: pd.DataFrame,
+    how: Literal["left", "right"] = "left",
     left_on: Union[str, None] = None,
     right_on: Union[str, None] = None,
     on: Union[str, None] = None,
-    how: Literal["left", "right", "all"] = "all",
-    return_score: bool = False,
     analyzer: Literal["word", "char", "char_wb"] = "char_wb",
     ngram_range: Tuple[int, int] = (2, 4),
+    return_score: bool = False,
     match_score: float = 0,
-    drop_unmatched=False,
-    suffixes: Tuple[str, str] = ("_l", "_r"),
+    drop_unmatched: bool = False,
+    suffixes: Tuple[str, str] = ("_x", "_y"),
 ) -> pd.DataFrame:
     """
     Join two tables categorical string columns based on approximate
@@ -46,6 +46,11 @@ def fuzzy_join(
         A table to merge.
     right : pandas.DataFrame
         A table used to merge with.
+    how: typing.Literal["left", "right"], default=`left`
+        Type of merge to be performed. Note that unlike pandas' merge, 
+        only "left" and "right" are supported so far, as the fuzzy-join comes
+        with its own mechanism to resolve lack of correspondence between
+        left and right tables.
     left_on : typing.Union[str, None]
         Name of left table column to join.
     right_on : typing.Union[str, None]
@@ -55,12 +60,6 @@ def fuzzy_join(
         Name of common left and right table join key columns.
         Must be found in both DataFrames. Use only if `left_on`
         and `right_on` parameters are not specified.
-    how : typing.Literal['left', 'right', 'all'], default='all'
-        Keep the join key columns from the left, right or
-        all tables.
-    return_score : boolean, default=True
-        Wheter to return matching score based on the distance between
-        nearest matched categories.
     analyzer : typing.Literal["word", "char", "char_wb"], default=`char_wb`
         Analyzer parameter for the CountVectorizer used for the string
         similarities.
@@ -72,6 +71,9 @@ def fuzzy_join(
         The lower and upper boundary of the range of n-values for different
         n-grams used in the string similarity. All values of n such
         that min_n <= n <= max_n will be used.
+    return_score : boolean, default=True
+        Whether to return matching score based on the distance between
+        nearest matched categories.
     match_score : float, default=0
         Distance score between the closest matches that will be accepted.
         In a [0, 1] interval. Closer to 1 means the matches need to be very
@@ -109,7 +111,7 @@ def fuzzy_join(
     Examples
     --------
     >>> df1 = pd.DataFrame({'a': ['ana', 'lala', 'nana'], 'b': [1, 2, 3]})
-    >>> df2 = pd.DataFrame({'a': ['anna', 'lala', 'ana', 'sana'], 'c': [5, 6, 7, 8]})
+    >>> df2 = pd.DataFrame({'a': ['anna', 'lala', 'ana', 'nnana'], 'c': [5, 6, 7, 8]})
 
     >>> df1
         a  b
@@ -118,25 +120,25 @@ def fuzzy_join(
     2  nana  3
 
     >>> df2
-        a  c
+        a    c
     0  anna  5
     1  lala  6
-    2   ana  7
-    3  sana  8
+    2  ana   7
+    3  nnana 8
 
     To do a simple join based on the nearest match:
 
     >>> fuzzy_join(df1, df2, on='a')
-        a_l  b   a_r    c
-    0   ana  1   ana   7
-    1  lala  2  lala   6
-    2  nana  3  sana   8
+        a_x  b   a_y    c
+    0   ana  1   ana    7
+    1  lala  2  lala    6
+    2  nana  3  nnana   8
 
     When we want to accept only a certain match precison,
     we can use the `match_score` argument:
 
     >>> fuzzy_join(df1, df2, on='a', match_score=1, return_score=True)
-        a_l  b   a_r    c  matching_score
+        a_x  b   a_y    c  matching_score
     0   ana  1   ana  7.0  1.000000
     1  lala  2  lala  6.0  1.000000
     2  nana  3   NaN  NaN  0.532717
@@ -152,71 +154,57 @@ def fuzzy_join(
             f"analyzer should be either 'char', 'word' or 'char_wb', got {analyzer!r}",
         )
 
-    if how not in ["left", "right", "all"]:
+    if how not in ["left", "right"]:
         raise ValueError(
-            f"how should be either 'left', 'right' or 'all', got {how!r}",
+            f"how should be either 'left' or 'right', got {how!r}",
         )
 
-    if len(suffixes) != 2:
-        raise ValueError(f"Invalid number of suffixes: expected 2, got {len(suffixes)}")
-    lsuffix, rsuffix = suffixes
-
     for param in [on, left_on, right_on]:
         if param is not None and not isinstance(param, str):
-            raise ValueError(
+            raise KeyError(
                 "Parameter 'left_on', 'right_on' or 'on' has invalid type, expected"
                 " string"
             )
 
-    left_table_clean = left.reset_index(drop=True).copy()
-    right_table_clean = right.reset_index(drop=True).copy()
-
-    overlap_cols = left_table_clean._info_axis.intersection(
-        right_table_clean._info_axis
-    )
-    if len(overlap_cols) > 0:
-        if suffixes[0] == "" and suffixes[1] == "":
-            raise ValueError(f"Columns overlap but no suffix specified: {overlap_cols}")
-        for i in range(len(overlap_cols)):
-            new_name_l = overlap_cols[i] + lsuffix
-            new_name_r = overlap_cols[i] + rsuffix
-            left_table_clean.rename(columns={overlap_cols[i]: new_name_l}, inplace=True)
-            right_table_clean.rename(
-                columns={overlap_cols[i]: new_name_r}, inplace=True
-            )
-            if left_on is not None and overlap_cols[i] in left_on:
-                left_on = new_name_l
-            if right_on is not None and overlap_cols[i] in right_on:
-                right_on = new_name_r
+    if how == "left":
+        left_table_clean = left.reset_index(drop=True).copy()
+        right_table_clean = right.reset_index(drop=True).copy()
+    else:
+        # We inverse the process so that the join is performed on the right column
+        left_table_clean = right.reset_index(drop=True).copy()
+        right_table_clean = left.reset_index(drop=True).copy()
 
     if on is not None:
-        left_col = on + lsuffix
-        right_col = on + rsuffix
+        left_col = on
+        right_col = on
     elif left_on is not None and right_on is not None:
-        left_col = left_on
-        right_col = right_on
+        if how == "left":
+            left_col = left_on
+            right_col = right_on
+        else:
+            left_col = right_on
+            right_col = left_on
+    else:
+        raise KeyError(
+            "Required parameter missing: either parameter"
+            "'on' or the pair 'left_on', 'right_on' should be specified."
+        )
 
     # Drop missing values in key columns
     left_table_clean.dropna(subset=[left_col], inplace=True)
     right_table_clean.dropna(subset=[right_col], inplace=True)
 
     # Make sure that the column types are string and categorical:
-    left_table_clean[left_col] = (
-        left_table_clean[left_col].astype(str).astype("category")
-    )
-    right_table_clean[right_col] = (
-        right_table_clean[right_col].astype(str).astype("category")
-    )
+    left_col_clean = left_table_clean[left_col].astype(str)
+    right_col_clean = right_table_clean[right_col].astype(str)
 
     enc = CountVectorizer(analyzer=analyzer, ngram_range=ngram_range)
 
-    all_cats = pd.concat(
-        [left_table_clean[left_col], right_table_clean[right_col]], axis=0
-    )
+    all_cats = pd.concat([left_col_clean, right_col_clean], axis=0).unique()
 
     enc_cv = enc.fit(all_cats)
-    left_enc = enc_cv.transform(left_table_clean[left_col])
-    right_enc = enc_cv.transform(right_table_clean[right_col])
+    left_enc = enc_cv.transform(left_col_clean)
+    right_enc = enc_cv.transform(right_col_clean)
 
     all_enc = vstack((left_enc, right_enc))
 
@@ -230,49 +218,31 @@ def fuzzy_join(
     distance, neighbors = neigh.kneighbors(left_enc, return_distance=True)
     idx_closest = np.ravel(neighbors)
 
+    left_table_clean["fj_idx"] = idx_closest
+    right_table_clean["fj_idx"] = right_table_clean.index
+
     norm_distance = 1 - (distance / 2)
+    if drop_unmatched:
+        left_table_clean = left_table_clean[match_score <= norm_distance]
+        norm_distance = norm_distance[match_score <= norm_distance]
+    else:
+        left_table_clean.loc[np.ravel(match_score > norm_distance), "fj_nan"] = 1
 
-    left_array = np.array(left_table_clean)
-    right_array = np.array(right_table_clean)
-    joined = np.append(
-        left_array,
-        np.array(
-            [
-                right_array[idx_closest[idr]]
-                if norm_distance[idr] >= match_score
-                else np.tile(np.nan, (right_array.shape[1],))
-                for idr in left_table_clean.index
-            ]
-        ),
-        axis=1,
+    df_joined = pd.merge(
+        left_table_clean, right_table_clean, on="fj_idx", suffixes=suffixes, how="left"
     )
 
-    cols = list(left_table_clean.columns) + list(right_table_clean.columns)
-    df_joined = pd.DataFrame(joined, columns=cols).replace(r"^\s*$", np.nan, regex=True)
-
-    duplicate_names = df_joined.columns.duplicated(keep=False)
-    if sum(duplicate_names) > 0:
-        warnings.warn("Column names overlaps. Please set appropriate suffixes.")
-        idx_to_keep = list(np.where(~duplicate_names)[0])
+    if drop_unmatched:
+        df_joined.drop(columns=["fj_idx"], inplace=True)
+    else:
+        idx = df_joined.index[df_joined["fj_nan"] == 1]
+        if len(idx) != 0:
+            df_joined.iloc[idx, df_joined.columns.get_loc("fj_idx") :] = np.NaN
+        df_joined.drop(columns=["fj_idx", "fj_nan"], inplace=True)
 
     if return_score:
         df_joined = pd.concat(
             [df_joined, pd.DataFrame(norm_distance, columns=["matching_score"])], axis=1
         )
-    if drop_unmatched:
-        df_joined.dropna(subset=[left_col, right_col], inplace=True)
-    if how == "left":
-        if sum(duplicate_names) > 0:
-            idx_to_keep.append(np.where(duplicate_names)[0][0])
-            idx_to_keep.sort()
-            df_joined = df_joined.iloc[:, idx_to_keep]
-        else:
-            df_joined.drop(columns=[right_col], inplace=True)
-    elif how == "right":
-        if sum(duplicate_names) > 0:
-            idx_to_keep.append(np.where(duplicate_names)[0][1])
-            idx_to_keep.sort()
-            df_joined = df_joined.iloc[:, idx_to_keep]
-        else:
-            df_joined.drop(columns=[left_col], inplace=True)
+
     return df_joined
diff --git a/dirty_cat/datasets/__init__.py b/dirty_cat/datasets/__init__.py
index 0a7bdd1..ff8ccc4 100644
--- a/dirty_cat/datasets/__init__.py
+++ b/dirty_cat/datasets/__init__.py
@@ -1,4 +1,6 @@
 from ._fetching import (
+    DatasetAll,
+    DatasetInfoOnly,
     fetch_drug_directory,
     fetch_employee_salaries,
     fetch_medical_charge,
@@ -11,13 +13,15 @@ from ._fetching import (
 )
 
 __all__ = [
-    "get_data_dir",
+    "DatasetAll",
+    "DatasetInfoOnly",
+    "fetch_drug_directory",
+    "fetch_employee_salaries",
     "fetch_medical_charge",
     "fetch_midwest_survey",
-    "fetch_employee_salaries",
-    "fetch_road_safety",
     "fetch_open_payments",
-    "fetch_drug_directory",
+    "fetch_road_safety",
     "fetch_traffic_violations",
     "fetch_world_bank_indicator",
+    "get_data_dir",
 ]
diff --git a/dirty_cat/datasets/_fetching.py b/dirty_cat/datasets/_fetching.py
index f5f8a3b..d0549b6 100644
--- a/dirty_cat/datasets/_fetching.py
+++ b/dirty_cat/datasets/_fetching.py
@@ -64,16 +64,38 @@ class Features:
 
 @dataclass(unsafe_hash=True)
 class DatasetAll:
+    """
+    Represents a dataset and its information.
+    With this state, the dataset is loaded in memory as a pandas DataFrame
+    (``X`` and ``y``).
+    Additional information such as `path` and `read_csv_kwargs` are provided
+    in case the dataframe has to be read from disk, as such:
+    .. code:: python
+        ds = fetch_employee_salaries(load_dataframe=False)
+        df = pd.read_csv(ds.path, **ds.read_csv_kwargs)
+    """
+
     name: str
     description: str
+    source: str
+    target: str
     X: pd.DataFrame
     y: pd.Series
-    source: str
     path: Path
+    read_csv_kwargs: Dict[str, Any]
 
 
 @dataclass(unsafe_hash=True)
 class DatasetInfoOnly:
+    """
+    Represents a dataset and its information.
+    With this state, the dataset is NOT loaded in memory, but can be read
+    with ``path`` and ``read_csv_kwargs``, as such:
+    .. code:: python
+        ds = fetch_employee_salaries(load_dataframe=False)
+        df = pd.read_csv(ds.path, **ds.read_csv_kwargs)
+    """
+
     name: str
     description: str
     source: str
@@ -431,10 +453,12 @@ def fetch_dataset_as_dataclass(
         dataset = DatasetAll(
             name=dataset_name,
             description=info["description"],
+            source=info["source"],
+            target=target,
             X=X,
             y=y,
-            source=info["source"],
             path=info["path"],
+            read_csv_kwargs=read_csv_kwargs,
         )
     else:
         dataset = DatasetInfoOnly(
@@ -462,9 +486,9 @@ def fetch_employee_salaries(
     https://openml.org/d/42125
 
     Description of the dataset:
-    > Annual salary information including gross pay and overtime pay for all
-    active, permanent employees of Montgomery County, MD paid in calendar
-    year 2016. This information will be published annually each year.
+        Annual salary information including gross pay and overtime pay for all
+        active, permanent employees of Montgomery County, MD paid in calendar
+        year 2016. This information will be published annually each year.
 
     Parameters
     ----------
@@ -513,10 +537,10 @@ def fetch_road_safety(
     https://openml.org/d/42803
 
     Description of the dataset:
-    > Data reported to the police about the circumstances of personal injury
-    road accidents in Great Britain from 1979, and the maker and model
-    information of vehicles involved in the respective accident. This version
-    includes data up to 2015.
+        Data reported to the police about the circumstances of personal injury
+        road accidents in Great Britain from 1979, and the maker and model
+        information of vehicles involved in the respective accident.
+        This version includes data up to 2015.
 
     Returns
     -------
@@ -544,14 +568,14 @@ def fetch_medical_charge(
     https://openml.org/d/42720
 
     Description of the dataset:
-    > The Inpatient Utilization and Payment Public Use File (Inpatient PUF)
-    provides information on inpatient discharges for Medicare fee-for-service
-    beneficiaries. The Inpatient PUF includes information on utilization,
-    payment (total payment and Medicare payment), and hospital-specific charges
-    for the more than 3,000 U.S. hospitals that receive Medicare Inpatient
-    Prospective Payment System (IPPS) payments. The PUF is organized by
-    hospital and Medicare Severity Diagnosis Related Group (MS-DRG) and
-    covers Fiscal Year (FY) 2011 through FY 2016.
+        The Inpatient Utilization and Payment Public Use File (Inpatient PUF)
+        provides information on inpatient discharges for Medicare
+        fee-for-service beneficiaries. The Inpatient PUF includes information
+        on utilization, payment (total payment and Medicare payment), and
+        hospital-specific charges for the more than 3,000 U.S. hospitals that
+        receive Medicare Inpatient Prospective Payment System (IPPS) payments.
+        The PUF is organized by hospital and Medicare Severity Diagnosis
+        Related Group (MS-DRG) and covers Fiscal Year (FY) 2011 through FY 2016.
 
     Returns
     -------
@@ -580,7 +604,7 @@ def fetch_midwest_survey(
     https://openml.org/d/42805
 
     Description of the dataset:
-    > Survey to know if people self-identify as Midwesterners.
+        Survey to know if people self-identify as Midwesterners.
 
     Returns
     -------
@@ -609,8 +633,8 @@ def fetch_open_payments(
     https://openml.org/d/42738
 
     Description of the dataset:
-    > Payments given by healthcare manufacturing companies to medical doctors
-    or hospitals.
+        Payments given by healthcare manufacturing companies to medical doctors
+        or hospitals.
 
     Returns
     -------
@@ -640,10 +664,10 @@ def fetch_traffic_violations(
     https://openml.org/d/42132
 
     Description of the dataset:
-    > This dataset contains traffic violation information from all electronic
-    traffic violations issued in the Montgomery County, MD. Any information
-    that can be used to uniquely identify the vehicle, the vehicle owner or
-    the officer issuing the violation will not be published.
+        This dataset contains traffic violation information from all electronic
+        traffic violations issued in the Montgomery County, MD. Any information
+        that can be used to uniquely identify the vehicle, the vehicle owner or
+        the officer issuing the violation will not be published.
 
     Returns
     -------
@@ -673,8 +697,8 @@ def fetch_drug_directory(
     https://openml.org/d/43044
 
     Description of the dataset:
-    > Product listing data submitted to the U.S. FDA for all unfinished,
-    unapproved drugs.
+        Product listing data submitted to the U.S. FDA for all unfinished,
+        unapproved drugs.
 
     Returns
     -------
diff --git a/doc/conf.py b/doc/conf.py
index 25f37a3..1b3b545 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -23,9 +23,9 @@ import shutil
 #
 # We avoid duplicating the information, but we do not use symlinks to be
 # able to build the docs on windows
-shutil.copyfile('../RELEASE_PROCESS.rst', 'RELEASE_PROCESS.rst')
-shutil.copyfile('../CHANGES.rst', 'CHANGES.rst')
-shutil.copyfile('../CONTRIBUTING.rst', 'CONTRIBUTING.rst')
+shutil.copyfile("../RELEASE_PROCESS.rst", "RELEASE_PROCESS.rst")
+shutil.copyfile("../CHANGES.rst", "CHANGES.rst")
+shutil.copyfile("../CONTRIBUTING.rst", "CONTRIBUTING.rst")
 
 # -- General configuration ------------------------------------------------
 
@@ -265,7 +265,7 @@ sphinx_gallery_conf = {
         "org": "dirty-cat",
         "repo": "dirty-cat.github.io",
         "binderhub_url": "https://mybinder.org",
-        "branch": "master",
+        "branch": "main",
         "dependencies": ["../requirements.txt", "binder/requirements.txt"],
         "notebooks_dir": "dev",
     },
diff --git a/doc/index.rst b/doc/index.rst
index 8db5a9a..09c8654 100644
--- a/doc/index.rst
+++ b/doc/index.rst
@@ -153,6 +153,7 @@ Data download
    datasets.fetch_open_payments
    datasets.fetch_road_safety
    datasets.fetch_traffic_violations
+   datasets.fetch_drug_directory
    datasets.fetch_world_bank_indicator
    datasets.get_data_dir
 
diff --git a/doc/unlisted.rst b/doc/unlisted.rst
new file mode 100644
index 0000000..5a71ff5
--- /dev/null
+++ b/doc/unlisted.rst
@@ -0,0 +1,17 @@
+========
+Unlisted
+========
+
+This document is here to trigger documentation generation for internal
+functions and classes which should not be on the index page or anywhere else,
+but still have a generated page.
+
+.. currentmodule:: dirty_cat
+
+.. autosummary::
+   :toctree: generated/
+   :template: class.rst
+   :nosignatures:
+
+   datasets.DatasetAll
+   datasets.DatasetInfoOnly
diff --git a/examples/06_dimension_reduction_and_performance.py b/examples/05_dimension_reduction_and_performance.py
similarity index 79%
rename from examples/06_dimension_reduction_and_performance.py
rename to examples/05_dimension_reduction_and_performance.py
index f37c069..43b1bf2 100644
--- a/examples/06_dimension_reduction_and_performance.py
+++ b/examples/05_dimension_reduction_and_performance.py
@@ -46,9 +46,8 @@ def resource_used(func):
         out = func(*args, **kwargs)
         size, peak = tracemalloc.get_traced_memory()
         tracemalloc.stop()
-        peak /= (1024 ** 2)  # Convert to megabytes
-        print(f"Run time: {perf_counter() - t0:.2f}s | "
-              f"Memory used: {peak:.2f}MB. ")
+        peak /= 1024**2  # Convert to megabytes
+        print(f"Run time: {perf_counter() - t0:.2f}s | Memory used: {peak:.2f}MB. ")
         return out
 
     return wrapped_func
@@ -84,13 +83,13 @@ y = y[~na_mask].reset_index(drop=True)
 ###############################################################################
 # We'll write down which columns are clean and which are dirty
 clean_columns = [
-    'Applicable_Manufacturer_or_Applicable_GPO_Making_Payment_Name',
-    'Dispute_Status_for_Publication',
-    'Physician_Specialty',
+    "Applicable_Manufacturer_or_Applicable_GPO_Making_Payment_Name",
+    "Dispute_Status_for_Publication",
+    "Physician_Specialty",
 ]
 dirty_columns = [
-    'Name_of_Associated_Covered_Device_or_Medical_Supply1',
-    'Name_of_Associated_Covered_Drug_or_Biological1',
+    "Name_of_Associated_Covered_Device_or_Medical_Supply1",
+    "Name_of_Associated_Covered_Drug_or_Biological1",
 ]
 
 ###############################################################################
@@ -113,23 +112,19 @@ from sklearn.compose import ColumnTransformer
 from dirty_cat import SimilarityEncoder
 
 clean_col_transformer = [
-    ('one_hot',
-     OneHotEncoder(sparse=False, handle_unknown='ignore'),
-     clean_columns),
+    ("one_hot", OneHotEncoder(sparse=False, handle_unknown="ignore"), clean_columns),
 ]
 
 column_trans = ColumnTransformer(
-    transformers=clean_col_transformer + [
-        ('sim_enc',
-         SimilarityEncoder(),
-         dirty_columns)
-    ],
-    remainder='drop')
+    transformers=clean_col_transformer
+    + [("sim_enc", SimilarityEncoder(), dirty_columns)],
+    remainder="drop",
+)
 
 t0 = perf_counter()
 X_enc = column_trans.fit_transform(X)
 t1 = perf_counter()
-print(f'Time to vectorize: {t1 - t0:.3f}s')
+print(f"Time to vectorize: {t1 - t0:.3f}s")
 
 ###############################################################################
 # Let's now run a cross-validation!
@@ -146,9 +141,9 @@ print(f"Cross-validation score: {results['test_score']}")
 ###############################################################################
 # Store results for later
 scores = dict()
-scores['Default options'] = results['test_score']
+scores["Default options"] = results["test_score"]
 times = dict()
-times['Default options'] = results['fit_time']
+times["Default options"] = results["fit_time"]
 
 ###############################################################################
 # Most frequent strategy to define prototypes
@@ -160,12 +155,16 @@ times['Default options'] = results['fit_time']
 # Here, we arbitrarily choose 100 as the number of prototypes we want to use.
 
 column_trans = ColumnTransformer(
-    transformers=clean_col_transformer + [
-        ('sim_enc',
-         SimilarityEncoder(categories='most_frequent', n_prototypes=100),
-         dirty_columns)
+    transformers=clean_col_transformer
+    + [
+        (
+            "sim_enc",
+            SimilarityEncoder(categories="most_frequent", n_prototypes=100),
+            dirty_columns,
+        )
     ],
-    remainder='drop')
+    remainder="drop",
+)
 
 ###############################################################################
 # Check that the prediction is still as good
@@ -175,8 +174,8 @@ print(f"Cross-validation score: {results['test_score']}")
 
 ###############################################################################
 # Store results for later
-scores['Most frequent'] = results['test_score']
-times['Most frequent'] = results['fit_time']
+scores["Most frequent"] = results["test_score"]
+times["Most frequent"] = results["fit_time"]
 
 ###############################################################################
 # KMeans strategy to define prototypes
@@ -188,23 +187,27 @@ times['Most frequent'] = results['fit_time']
 # we chose here is arbitrary.
 
 column_trans = ColumnTransformer(
-    transformers=clean_col_transformer + [
-        ('sim_enc',
-         SimilarityEncoder(categories='k-means', n_prototypes=100),
-         dirty_columns)
+    transformers=clean_col_transformer
+    + [
+        (
+            "sim_enc",
+            SimilarityEncoder(categories="k-means", n_prototypes=100),
+            dirty_columns,
+        )
     ],
-    remainder='drop')
+    remainder="drop",
+)
 
 ###############################################################################
 # Check that the prediction is still as good
 model = pipeline.make_pipeline(column_trans, log_reg)
 results = resource_used(model_selection.cross_validate)(model, X, y)
-print("Cross-validation score: %s" % results['test_score'])
+print("Cross-validation score: %s" % results["test_score"])
 
 ###############################################################################
 # Store results for later
-scores['KMeans'] = results['test_score']
-times['KMeans'] = results['fit_time']
+scores["KMeans"] = results["test_score"]
+times["KMeans"] = results["fit_time"]
 
 ###############################################################################
 # Summary
@@ -213,11 +216,11 @@ import seaborn
 import matplotlib.pyplot as plt
 
 _, (ax1, ax2) = plt.subplots(nrows=2, figsize=(4, 3))
-seaborn.boxplot(data=pd.DataFrame(scores), orient='h', ax=ax1)
-ax1.set_xlabel('Prediction accuracy', size=16)
+seaborn.boxplot(data=pd.DataFrame(scores), orient="h", ax=ax1)
+ax1.set_xlabel("Prediction accuracy", size=16)
 [t.set(size=16) for t in ax1.get_yticklabels()]
 
-seaborn.boxplot(data=pd.DataFrame(times), orient='h', ax=ax2)
-ax2.set_xlabel('Computation time', size=16)
+seaborn.boxplot(data=pd.DataFrame(times), orient="h", ax=ax2)
+ax2.set_xlabel("Computation time", size=16)
 [t.set(size=16) for t in ax2.get_yticklabels()]
 plt.tight_layout()
diff --git a/examples/05_scaling_non_linear_models.py b/examples/05_scaling_non_linear_models.py
deleted file mode 100644
index cb71275..0000000
--- a/examples/05_scaling_non_linear_models.py
+++ /dev/null
@@ -1,539 +0,0 @@
-r"""
-Fitting scalable, non-linear models on data with dirty categories
-=================================================================
-
-When training a machine learning model, it is often hard to choose between
-linear and a non-linear models.
-
-On one hand, linear models are well studied and understood. They are generally
-fast and easy to optimize, and generate interpretable results.
-They do however reach their expressiveness limits when there is a complex
-relation between the input and the output: whatever the number of samples the
-training set may have, past some point, their precision won't get any better.
-
-On the other hand, non-linear models tend to scale better with sample size:
-they are able to digest the massive amount of information and get a better
-estimate of the link between the input and the output.
-
-Non-linear models form a very large class. Among others, it includes:
-
-* Neural Networks
-* Tree-based methods such as Random Forests, and the very powerful Gradient
-  Boosting Machines [#xgboost]_
-* Kernel Methods
-
-However, reaching the phase where non-linear models outperform the linear ones
-can be complicated. Indeed, a more complex model often means a longer
-fitting/tuning process
-
-* Neural networks often need extended model tuning time in order to
-  achieve good optimization and architecture.
-* Gradient Boosting Machines tend to not scale very well with increasing
-  sample size, as all the data needs to be loaded in the main memory.
-* For kernel methods, parameter fitting requires the inversion of a gram matrix
-  of size :math:`n \times n` (:math:`n` being the number of samples), yielding
-  a quadratic dependency (with :math:`n`) as the final computation time.
-
-
-In order to get the best out of a non-linear model, one has to **make it
-scalable**. Quite notably, for kernel methods, approximation algorithms that
-drop the quadratic dependency with the sample size while ensuring capacity
-approximation exist.
-
-In this example, you will learn how to:
-    1. Build a machine learning pipeline that uses a kernel method.
-    2. Make this pipeline scalable, by using online algorithms and dimension
-       reduction methods.
-
-
-.. note::
-   This example assumes the reader is familiar with similarity encoding and
-   its use-cases.
-
-   * For an introduction to dirty categories, see :ref:`this example<https://dirty-cat.github.io/stable/auto_examples/02_investigating_dirty_categories.html#sphx-glr-auto-examples-02-investigating-dirty-categories-py>`.
-   * To learn with dirty categories using the SimilarityEncoder, see :ref:`this example<https://dirty-cat.github.io/stable/auto_examples/01_dirty_categories.html#sphx-glr-auto-examples-01-dirty-categories-py>`.
-
-
-.. |NYS| replace:: :class:`Nystroem <sklearn.kernel_approximation.Nystroem>`
-
-.. |NYS_EXAMPLE|
-    replace:: :ref:`scikit-learn documentation <nystroem_kernel_approx>`
-
-.. |RBF| replace:: :class:`~sklearn.kernel_approximation.RBFSampler`
-
-.. |SVC| replace:: :class:`SupportVectorClassifier <sklearn.svm.SVC>`
-
-.. |SE| replace:: :class:`~dirty_cat.SimilarityEncoder`
-
-.. |SGDClassifier| replace:: :class:`~sklearn.linear_model.SGDClassifier`
-
-.. |APS| replace:: :func:`~sklearn.metrics.average_precision_score`
-
-.. |OHE| replace:: :class:`~sklearn.preprocessing.OneHotEncoder`
-
-.. |ColumnTransformer| replace:: :class:`~sklearn.compose.ColumnTransformer`
-
-.. |LabelEncoder| replace:: :class:`~sklearn.preprocessing.LabelEncoder`
-
-.. |SGDClassifier_partialfit| replace::
-    :meth:`~sklearn.linear_model.SGDClassifier.partial_fit`
-
-.. |Pipeline| replace:: :class:`~sklearn.pipeline.Pipeline`
-
-.. |pd_read_csv| replace:: :func:`pandas.read_csv`
-
-.. |sklearn| replace:: :std:doc:`scikit-learn <sklearn:index>`
-
-.. |transform| replace:: :std:term:`transform <sklearn:transform>`
-
-.. |fit| replace:: :std:term:`fit <sklearn:fit>`
-
-"""
-###############################################################################
-# Training a simple pipeline
-# --------------------------
-# We will train our model on the :code:`drug_directory` dataset.
-from dirty_cat.datasets import fetch_drug_directory
-
-# We'll only gather the information, and not load the dataset in memory for now
-drug_directory = fetch_drug_directory(load_dataframe=False)
-print(drug_directory.description)
-
-###############################################################################
-# .. topic:: Problem Setting:
-#
-#    .. centered::
-#       **predict the drug type given its composition**
-#
-#
-# The :code:`NONPROPRIETARYNAME` column, is composed of textual observations
-# describing each drug's composition. The :code:`PRODUCTTYPENAME` column
-# consists of categorical values; we therefore a classification problem.
-#
-# You can have a glimpse of the values here:
-import pandas as pd
-
-df = pd.read_csv(
-    drug_directory.path,
-    nrows=10,
-    **drug_directory.read_csv_kwargs,
-).astype(str)
-print(df[['NONPROPRIETARYNAME', 'PRODUCTTYPENAME']].head())
-# This will be useful further down in the example.
-columns_names = df.columns
-
-###############################################################################
-# Estimators construction
-# -----------------------
-# Our input is categorical, thus needs to be encoded.
-#
-# As observations often consist in variations around a few concepts
-# - for instance, "Amlodipine Besylate" and
-# "Amlodipine besylate and atorvastatin calcium"
-# have one ingredient in common - we need an encoding able to
-# capture similarities between observations. We choose to use a |SE|.
-#
-# Two other columns are used to predict the output: ``DOSAGEFORMNAME`` and
-# ``ROUTENAME``.
-# They are both categorical and can be encoded with a |OHE|.
-# We use a |ColumnTransformer| to stack both.
-# We can now choose a kernel method. Here, we'll use a |SVC| to fit
-# the encoded inputs.
-from dirty_cat import SimilarityEncoder
-from sklearn.preprocessing import OneHotEncoder
-from sklearn.compose import make_column_transformer
-
-column_transformer = make_column_transformer(
-    (SimilarityEncoder(), ['NONPROPRIETARYNAME']),
-    (OneHotEncoder(handle_unknown='ignore'), ['DOSAGEFORMNAME', 'ROUTENAME']),
-    sparse_threshold=1,
-)
-
-###############################################################################
-# We will bundle the ColumnTransformer and the classifier into a |Pipeline|
-from sklearn.svm import SVC
-from sklearn.pipeline import Pipeline
-classifier = SVC(kernel='rbf', random_state=42, gamma=1)
-steps = [('transformer', column_transformer), ('classifier', classifier)]
-model = Pipeline(steps)
-
-###############################################################################
-# Data Loading and Preprocessing
-# ------------------------------
-#
-# As usual in most machine learning setups, the data has to be split into 2
-# exclusive parts:
-#
-# As usual in machine learning settings, we will split our data in two
-# exclusive sets:
-#
-# * One for model training
-# * One for model testing
-#
-# We create a simple wrapper around |pd_read_csv|, that extracts :code:`X`,
-# and :code:`y` from the dataset.
-#
-# .. topic:: Note about class imbalance:
-#
-#    The :code:`y` labels are composed of 7 unique classes. However, the two
-#    classes ``HUMAN OTC DRUG`` and ``HUMAN PRESCRIPTION DRUG`` represent
-#    around 97% of the data. The other 5 classes are much rarer.
-#    Dealing with class imbalance is out of the scope of this example, so the
-#    models will be trained only on the two most common classes.
-#
-# Classifiers in |sklearn| often require :code:`y` to be integer labels.
-# Additionally, |APS| requires a binary version of the labels.
-#
-# We therefore create for these respective constraints a |LabelEncoder| and
-# a |OHE|
-import numpy as np
-from typing import Tuple
-from sklearn.preprocessing import LabelEncoder, OneHotEncoder
-
-
-label_encoder = LabelEncoder()
-# Pre-fit on the known classes of y
-label_encoder.fit(['HUMAN OTC DRUG', 'HUMAN PRESCRIPTION DRUG'])
-
-one_hot_encoder = OneHotEncoder(categories="auto", sparse=False)
-# Pre-fit on the integer labels
-one_hot_encoder.fit([[0], [1]])
-
-
-def preprocess(df: pd.DataFrame) -> Tuple[pd.DataFrame, np.array]:
-    df = df.loc[df['PRODUCTTYPENAME'].isin(
-        ['HUMAN OTC DRUG', 'HUMAN PRESCRIPTION DRUG'])]
-
-    df = df[['NONPROPRIETARYNAME', 'DOSAGEFORMNAME',
-             'ROUTENAME', 'PRODUCTTYPENAME']]
-    df = df.dropna()
-
-    X = df[['NONPROPRIETARYNAME', 'DOSAGEFORMNAME', 'ROUTENAME']]
-    y = df[['PRODUCTTYPENAME']].values
-
-    y_int = label_encoder.transform(np.squeeze(y))
-
-    return X, y_int
-
-
-def get_X_y(**kwargs) -> Tuple[pd.DataFrame, np.array]:
-    """Simple wrapper around pd.read_csv that extracts features and labels.
-
-    Some systematic preprocessing is also carried out to avoid doing this
-    transformation repeatedly in the code.
-    """
-    df = pd.read_csv(drug_directory.path, names=columns_names,
-                     **drug_directory.read_csv_kwargs, **kwargs)
-    return preprocess(df)
-
-
-###############################################################################
-# .. warning::
-#
-#    In the following steps, we assume the dataset to be shuffled prior to its
-#    loading. Be careful as this is not the case for all datasets!
-#    This is important here because we will take the :math:`n` first
-#    observations for the training set and the next :math:`m` observations for
-#    the test set repeatably.
-#
-#
-# Finally, :code:`X` and :code:`y` are loaded.
-#
-#
-# .. topic:: Note: offsetting the test set
-#
-#    For the online-learning procedures of this example, we would like to
-#    compare the accuracy at each iteration. In order to do that, we need to
-#    keep the same test set. As such, we "reserve" the first 100k rows for
-#    the training phase, and use the rest as the test set.
-train_set_size = 5000
-test_set_size = 10000
-offset = 100000
-
-X_train, y_train = get_X_y(skiprows=1, nrows=train_set_size)
-
-X_test, y_test = get_X_y(skiprows=offset, nrows=test_set_size)
-
-###############################################################################
-# Evaluating time and sample complexity
-# -------------------------------------
-#
-# Let's get an idea of model precision and performance depending on the number
-# of the samples used in the train set.
-# The |Pipeline| is trained over different training set sizes. For this,
-# :code:`X_train` and :code:`y_train` get sliced into subsets of increasing
-# size, while :code:`X_test` and :code:`y_test` do not change.
-import time
-from sklearn.metrics import average_precision_score
-
-# Define the different train set sizes
-train_set_sizes = [train_set_size // 10, train_set_size // 3, train_set_size]
-
-train_times_svc, test_scores_svc = [], []
-
-for n in train_set_sizes:
-
-    t0 = time.perf_counter()
-    model.fit(X_train[:n], y_train[:n])
-    train_time = time.perf_counter() - t0
-
-    y_pred = model.predict(X_test)
-
-    y_pred_onehot = one_hot_encoder.transform(y_pred.reshape(-1, 1))
-    y_test_onehot = one_hot_encoder.transform(y_test.reshape(-1, 1))
-
-    test_score = average_precision_score(y_test_onehot, y_pred_onehot)
-
-    train_times_svc.append(train_time)
-    test_scores_svc.append(test_score)
-
-    print(
-        f"Using {n:>5} samples: model fitting took {train_time:.1f}s, "
-        f"test accuracy of {test_score:.3f}. "
-    )
-
-
-###############################################################################
-# Increasing the training set size clearly improves model accuracy.
-# However, the training time and the input size increase quadratically with it.
-#
-# Indeed, kernel methods need to process an entire :math:`n \times n`
-# matrix at once. In order for this matrix to be loaded into memory, :math:`n`
-# has to remain low: If we used 32-bit floats, and with 30k observations, the
-# total size of this 30k x 30k matrix would be around
-# :math:`30000^2 \times 32 = 2.8 \times 10^{10} \text{bits} = 4\text{GB}`
-
-###############################################################################
-# Reducing input dimension using kernel approximation methods
-# -----------------------------------------------------------
-#
-# The main scalability issues with kernels methods is the processing of a large
-# square matrix. To understand where this matrix comes from, we need to delve a
-# bit deeper into these methods.
-#
-# .. topic:: Kernel methods
-#
-#    Kernel methods address non-linear problems by leveraging similarities
-#    between each pair of inputs. Using a similarity matrix to solve a machine
-#    learning problem allows to catch complex, non-linear relationships within
-#    the data.  But it requires inverting this matrix, which can be a
-#    computational burden when the sample size increases.
-
-
-###############################################################################
-# Kernel approximation methods
-# ----------------------------
-#
-# From what we saw above, two criterion are limiting kernel algorithms from
-# scaling:
-#
-# * It processes a matrix, whose size increases quadratically with the number
-#   of samples.
-# * During training, this matrix is **inverted**, meaning it has to be
-#   loaded into main memory.
-#
-# Kernel approximation methods such as |RBF| or |NYS| [#nys_ref]_ try to
-# approximate this similarity matrix, without actually creating it. By allowing
-# the program to not compute the perfect similarity matrix, **the problem
-# complexity becomes linear**! Furthermore, the samples don't need to be
-# processed at once.
-#
-# We are not bound to use a |SVC| anymore, and can instead use a |RBF| in an
-# online optimization setting to process the input by batch.
-#
-# .. topic:: Online algorithms
-#
-#    An online algorithm [#online_ref]_ is an algorithm that treats its input
-#    piece by piece in a serial fashion.
-#    A famous example is the stochastic gradient descent [#sgd_ref]_,
-#    where an estimation of the objective function's gradient is computed
-#    on a batch of the data at each step.
-#
-
-###############################################################################
-# Reducing the transformers dimensionality
-# ----------------------------------------
-#
-# There is one last scalability issue in our pipeline: the |SE| and the |RBF|
-# both implement the |fit| method. How to adapt those methods to an online
-# setting, where the data is never loaded as a whole?
-#
-# A simple solution is to partially fit the |SE| and the |RBF| on a subset of
-# the data, prior to the online fitting step.
-
-from sklearn.kernel_approximation import RBFSampler
-n_out_encoder = 1000
-n_out_rbf = 5000
-n_samples_encoder = 10000
-
-X_encoder, _ = get_X_y(nrows=n_samples_encoder)
-
-# Fit the rbf_sampler with the similarity matrix.
-column_transformer = make_column_transformer(
-    (
-        SimilarityEncoder(categories='most_frequent',
-                          n_prototypes=n_out_encoder,
-                          random_state=42, ngram_range=(2, 4)),
-        ['NONPROPRIETARYNAME']
-    ), (
-        OneHotEncoder(handle_unknown='ignore'),
-        ['DOSAGEFORMNAME', 'ROUTENAME']
-    ),
-    sparse_threshold=1)
-
-transformed_categories = column_transformer.fit_transform(X_encoder)
-
-# gamma is a parameter of the rbf function, that sets how fast the similarity
-# between two points should decrease as the distance between them rises. It
-# is data-specific, and needs to be chosen carefully, for example using
-# cross-validation.
-rbf_sampler = RBFSampler(
-    gamma=0.5, n_components=n_out_rbf, random_state=42)
-rbf_sampler.fit(transformed_categories)
-
-
-def encode(X, y_int, one_hot_encoder, column_transformer, rbf_sampler):
-    X_sim_encoded = column_transformer.transform(X)
-
-    X_highdim = rbf_sampler.transform(X_sim_encoded.toarray())
-
-    y_onehot = one_hot_encoder.transform(y_int.reshape(-1, 1))
-
-    return X_highdim, y_onehot
-
-
-# The inputs and labels of the val and test sets have to be pre-processed the
-# same way the training set was processed:
-X_test_kernel_approx, y_true_test_onehot = encode(
-    X_test, y_test, one_hot_encoder, column_transformer, rbf_sampler)
-
-
-###############################################################################
-# Online training for out-of-memory data
-# --------------------------------------
-# We now have all the elements to create a non-linear, online kernel method.
-import warnings
-from sklearn.linear_model import SGDClassifier
-
-online_train_set_size = 100000
-# Filter warnings on max_iter and tol
-warnings.filterwarnings('ignore', module='sklearn.linear_model')
-sgd_classifier = SGDClassifier(
-    max_iter=1, tol=None, random_state=42, average=10)
-
-
-###############################################################################
-# We can now start the training by looping over batches one by one.
-# Note that only one pass over the whole dataset is done.
-# It may be worth doing several passes, but for very large sample sizes,
-# the increase in test accuracy is likely to be marginal.
-batchsize = 1000
-test_scores_rbf = []
-train_times_rbf = []
-online_train_set_sizes = []
-t0 = time.perf_counter()
-
-iter_csv = pd.read_csv(
-    drug_directory.path, nrows=online_train_set_size, chunksize=batchsize,
-    skiprows=1, names=columns_names, **drug_directory.read_csv_kwargs)
-
-for batch_no, batch in enumerate(iter_csv):
-    X_batch, y_batch = preprocess(batch)
-    # Skip iteration if batch is empty after preprocessing
-    if len(y_batch) == 0:
-        continue
-    X_batch_kernel_approx, y_batch_onehot = encode(
-        X_batch, y_batch, one_hot_encoder, column_transformer, rbf_sampler)
-
-    # Make one pass of stochastic gradient descent over the batch.
-    sgd_classifier.partial_fit(
-        X_batch_kernel_approx, y_batch, classes=[0, 1])
-
-    # Print train/test accuracy metrics every 5 batch
-    if (batch_no % 5) == 0:
-        message = f"batch {batch_no:>4} "
-        for origin, X, y_true_onehot in zip(
-                ('train', 'val'),
-                (X_batch_kernel_approx, X_test_kernel_approx),
-                (y_batch_onehot, y_true_test_onehot)):
-
-            y_pred = sgd_classifier.predict(X)
-
-            # Preprocess correctly the labels and prediction to match
-            # average_precision_score expectations
-            y_pred_onehot = one_hot_encoder.transform(
-                y_pred.reshape(-1, 1))
-
-            score = average_precision_score(y_true_onehot, y_pred_onehot)
-            message += f"{origin} precision: {score:.4f}  "
-            if origin == 'val':
-                test_scores_rbf.append(score)
-                train_times_rbf.append(time.perf_counter() - t0)
-                online_train_set_sizes.append((batch_no + 1) * batchsize)
-
-        print(message)
-
-###############################################################################
-# So far, we fitted two kinds of models: an exact kernel algorithm, and an
-# approximate, online one.
-#
-# Let's compare both the accuracies and the number of visited samples
-# for each model as we increase our time budget:
-import matplotlib.pyplot as plt
-
-f, axs = plt.subplots(2, 1, sharex=True, figsize=(10, 7))
-ax_score, ax_capacity = axs
-
-ax_score.set_ylabel('score')
-ax_capacity.set_ylabel('training set size')
-ax_capacity.set_xlabel('time')
-
-ax_score.plot(train_times_svc, test_scores_svc, 'b-', label='exact')
-ax_score.plot(train_times_rbf, test_scores_rbf, 'r-', label='online')
-
-ax_capacity.plot(train_times_svc, train_set_sizes, 'b-', label='exact')
-ax_capacity.plot(train_times_rbf, online_train_set_sizes, 'r-', label='online')
-ax_capacity.set_yscale('log')
-
-ax_score.legend(bbox_to_anchor=(0., 1.02, 1., .102),
-                loc=3, ncol=2, mode='expand',
-                borderaxespad=0.)
-
-# Compare the two methods in the common time range
-ax_score.set_xlim(0, min(train_times_svc[-1], train_times_rbf[-1]))
-
-f.suptitle("Test set accuracy and number of samples visited"
-           "samples seen by the SGDClassifier")
-
-f.show()
-
-
-###############################################################################
-# This plot shows us that for the time budget, the online model will eventually
-# process more samples, be faster and reach a far higher test accuracy that the
-# non-online, exact kernel method.
-#
-# Our online model also outperforms online **linear** models (for instance,
-# |SE| + |SGDClassifier|).
-#
-# .. note::
-#
-#    We did not fit two online models here for the sake of simplicity, but to
-#    train an online linear model, simply comment out the line
-#    :code:`X_highdim = rbf_sampler.transform(X_sim_encoded.toarray())`
-#    in the function :code:`encode` and change it with, for example,
-#    :code:`X_highdim = X_sim_encoded`.
-#
-# In particular, this hierarchy between the linear and the non-linear models
-# shows that there were some significant non-linear relationships
-# between the input and the output. By scaling a kernel method, we successfully
-# took this non-linearity into account in our model, which was far from a
-# trivial task at the beginning of this example!
-
-###############################################################################
-# .. rubric:: Footnotes
-# .. [#xgboost] `Slides on gradient boosting by Tianqi Chen, the founder of XGBoost <https://web.njit.edu/~usman/courses/cs675_summer20/BoostedTree.pdf>`_
-# .. [#online_ref] `Wikipedia article on online algorithms <https://en.wikipedia.org/wiki/Online_algorithm>`_
-# .. [#sgd_ref] `Leon Bouttou's article on stochastic gradient descent <http://khalilghorbal.info/assets/spa/papers/ML_GradDescent.pdf>`_
-# .. [#nys_ref] |NYS_EXAMPLE|
-# .. [#dual_ref] `Introduction to duality <https://pdfs.semanticscholar.org/0373/e7289a1978108d6455218160a529c85842c0.pdf>`_
diff --git a/examples/07_joining_tables_with_FuzzyJoin.py b/examples/06_joining_tables_with_FuzzyJoin.py
similarity index 84%
rename from examples/07_joining_tables_with_FuzzyJoin.py
rename to examples/06_joining_tables_with_FuzzyJoin.py
index 57e1c72..2fb087b 100755
--- a/examples/07_joining_tables_with_FuzzyJoin.py
+++ b/examples/06_joining_tables_with_FuzzyJoin.py
@@ -41,13 +41,11 @@ df.head(3)
 # Generosity etc.
 #
 # For more information, read the `World Happiness Report <https://worldhappiness.report/>`_.
-X = df[["Country"]]
-y = df[["Happiness score"]]
 
 #######################################################################
-# We keep the country names in our X table and we create
-# the y table with the happiness score (our prediction target, or
-# exogenous variable).
+# For the sake of this example, we only keep the country names and our
+# variable of interest: the happiness score
+df = df[["Country", "Happiness score"]]
 
 #############################################################################
 # Additional tables from other sources
@@ -81,15 +79,16 @@ legal_rights.head(3)
 # ------------------------
 #
 # Alas, the entries for countries do not perfectly match between our
-# original table (X), and those that we downloaded from the worldbank:
+# original table (df), and those that we downloaded from the worldbank
+# (gdppc):
 
-X.sort_values(by='Country').tail(7)
+df.sort_values(by="Country").tail(7)
 
 #######################################################################
-gdppc.sort_values(by='Country Name').tail(7)
+gdppc.sort_values(by="Country Name").tail(7)
 
 #######################################################################
-# We can see that Yemen is written "Yemen*" on one side, and 
+# We can see that Yemen is written "Yemen*" on one side, and
 # "Yemen, Rep." on the other.
 #
 # We also have entries that probably do not have correspondances: "World"
@@ -99,7 +98,7 @@ gdppc.sort_values(by='Country Name').tail(7)
 # Joining tables with imperfect correspondance
 # --------------------------------------------
 #
-# We will now join our initial table, X, with the 3 additional ones that
+# We will now join our initial table, df, with the 3 additional ones that
 # we have extracted.
 #
 
@@ -110,15 +109,15 @@ gdppc.sort_values(by='Country Name').tail(7)
 # To join them with dirty_cat, we only need to do the following:
 from dirty_cat import fuzzy_join
 
-X1 = fuzzy_join(
-    X,  # our table to join
+df1 = fuzzy_join(
+    df,  # our table to join
     gdppc,  # the table to join with
     left_on="Country",  # the first join key column
     right_on="Country Name",  # the second join key column
     return_score=True,
 )
 
-X1.head(20)
+df1.head(20)
 # We merged the first WB table to our initial one.
 
 #################################################################
@@ -166,7 +165,7 @@ def print_worst_matches(joined_table, n=5):
 # Let's print the four worst matches, which will give
 # us an overview of the situation:
 
-print_worst_matches(X1, n=4)
+print_worst_matches(df1, n=4)
 
 #################################################################
 # We see that some matches were unsuccesful
@@ -178,25 +177,30 @@ print_worst_matches(X1, n=4)
 # In this case, it is better to use the threshold parameter
 # so as to include only precise-enough matches:
 #
-X1 = fuzzy_join(
-    X,
+df1 = fuzzy_join(
+    df,
     gdppc,
     left_on="Country",
     right_on="Country Name",
     match_score=0.35,
     return_score=True,
 )
-print_worst_matches(X1, n=4)
+print_worst_matches(df1, n=4)
 
 #################################################################
 # Matches that are not available (or precise enough) are marked as `NaN`.
-# We will remove them, as well as missing or unused information:
+# We will remove them using the drop_unmatched parameter:
 
-mask = X1["GDP per capita (current US$)"].notna()
-X1 = X1[mask]
-y = np.ravel(y[mask])
+df1 = fuzzy_join(
+    df,
+    gdppc,
+    left_on="Country",
+    right_on="Country Name",
+    match_score=0.35,
+    drop_unmatched=True,
+)
 
-X1.drop(["matching_score"], axis=1, inplace=True)
+df1.drop(columns=["Country Name"], inplace=True)
 
 #################################################################
 #
@@ -208,7 +212,9 @@ import seaborn as sns
 sns.set_context("notebook")
 
 plt.figure(figsize=(4, 3))
-ax = sns.regplot(x=X1["GDP per capita (current US$)"], y=y, lowess=True)
+ax = sns.regplot(
+    data=df1, x="GDP per capita (current US$)", y="Happiness score", lowess=True
+)
 ax.set_ylabel("Happiness index")
 ax.set_title("Is a higher GDP per capita linked to happiness?")
 plt.tight_layout()
@@ -226,21 +232,27 @@ plt.show()
 #
 # Now let's include other information that may be relevant, such as
 # life expectancy table:
-X2 = fuzzy_join(
-    X1,
+df2 = fuzzy_join(
+    df1,
     life_exp,
     left_on="Country",
     right_on="Country Name",
     match_score=0.45,
-    how="left",
 )
 
-X2.head(3)
+df2.drop(columns=["Country Name"], inplace=True)
+
+df2.head(3)
 
 #################################################
 # Let's plot this relation:
 plt.figure(figsize=(4, 3))
-fig = sns.regplot(x=X2["Life expectancy at birth, total (years)"], y=y, lowess=True)
+fig = sns.regplot(
+    data=df2,
+    x="Life expectancy at birth, total (years)",
+    y="Happiness score",
+    lowess=True,
+)
 fig.set_ylabel("Happiness index")
 fig.set_title("Is a higher life expectancy linked to happiness?")
 plt.tight_layout()
@@ -260,22 +272,26 @@ plt.show()
 #    column, so as not to have unnecessary overlaping column with country names.
 #
 # And the table with a measure of legal rights strength in the country:
-X3 = fuzzy_join(
-    X2,
+df3 = fuzzy_join(
+    df2,
     legal_rights,
     left_on="Country",
     right_on="Country Name",
     match_score=0.45,
-    how="left",
 )
 
-X3.head(3)
+df3.drop(columns=["Country Name"], inplace=True)
+
+df3.head(3)
 
 #################################################
 # Let's take a look at their correspondance in a figure:
 plt.figure(figsize=(4, 3))
 fig = sns.regplot(
-    x=X3["Strength of legal rights index (0=weak to 12=strong)"], y=y, lowess=True
+    data=df3,
+    x="Strength of legal rights index (0=weak to 12=strong)",
+    y="Happiness score",
+    lowess=True,
 )
 fig.set_ylabel("Happiness index")
 fig.set_title("Does a country's legal rights strength lead to happiness?")
@@ -294,6 +310,12 @@ plt.show()
 # Prediction model
 # -----------------
 #
+# We now separate our covariates (X), from the target (or exogenous)
+# variables: y
+X = df3.drop("Happiness score", axis=1).select_dtypes(exclude=object)
+y = df3[["Happiness score"]]
+
+###################################################################
 # Let us now define the model that will be used to predict the happiness score:
 from sklearn import __version__ as sklearn_version
 
@@ -312,19 +334,17 @@ cv = KFold(n_splits=2, shuffle=True, random_state=0)
 # Let's finally assess the results of our models:
 from sklearn.model_selection import cross_validate
 
-cv_results_t = cross_validate(
-    hgdb, X3.select_dtypes(exclude=object), y, cv=cv, scoring="r2"
-)
+cv_results_t = cross_validate(hgdb, X, y, cv=cv, scoring="r2")
 
 cv_r2_t = cv_results_t["test_score"]
 
 print(
-    f"Mean R2 score with {len(X3.columns) - 2} feature columns is"
+    f"Mean R2 score with {len(X.columns) - 2} feature columns is"
     f" {cv_r2_t.mean():.2f} +- {cv_r2_t.std():.2f}"
 )
 
 #################################################################
-# We have a satisfying first result: an R2 of 0.63!
+# We have a satisfying first result: an R2 of 0.66!
 #
 # Data cleaning varies from dataset to dataset: there are as
 # many ways to clean a table as there are errors. :func:`fuzzy_join`


Test patch:
diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml
index d3126e8..8c4fe28 100644
--- a/.github/workflows/testing.yml
+++ b/.github/workflows/testing.yml
@@ -3,7 +3,7 @@ name: 'build'
 on:
   push:
     branches:
-      - "master"
+      - "main"
   pull_request:
     branches:
       - '*'
diff --git a/dirty_cat/datasets/tests/test_fetching.py b/dirty_cat/datasets/tests/test_fetching.py
index 053d09f..01a2b11 100644
--- a/dirty_cat/datasets/tests/test_fetching.py
+++ b/dirty_cat/datasets/tests/test_fetching.py
@@ -329,6 +329,7 @@ def test_import_all_datasets(
         name="Example dataset",
         description="This is a dataset.",
         source="https://www.openml.org/",
+        target="To_predict",
         path=Path("/path/to/file.csv"),
         X=pd.DataFrame([1, 2, 3, 4]),
         y=pd.Series(
@@ -336,6 +337,7 @@ def test_import_all_datasets(
                 5,
             ]
         ),
+        read_csv_kwargs={"a": "b"},
     )
 
     expected_return_value_info_only = _fetching.DatasetInfoOnly(
diff --git a/dirty_cat/tests/test_fuzzy_join.py b/dirty_cat/tests/test_fuzzy_join.py
index c4791b1..d981ecf 100644
--- a/dirty_cat/tests/test_fuzzy_join.py
+++ b/dirty_cat/tests/test_fuzzy_join.py
@@ -5,11 +5,8 @@ import pytest
 from dirty_cat import fuzzy_join
 
 
-@pytest.mark.parametrize(
-    "analyzer, how",
-    [("char", "left"), ("char_wb", "right"), ("word", "all")],
-)
-def test_fuzzy_join(analyzer, how):
+@pytest.mark.parametrize("analyzer", ["char", "char_wb", "word"])
+def test_fuzzy_join(analyzer):
     """Testing if fuzzy_join results are as expected."""
 
     df1 = pd.DataFrame({"a1": ["ana", "lala", "nana et sana", np.NaN]})
@@ -20,7 +17,7 @@ def test_fuzzy_join(analyzer, how):
         right=df2,
         left_on="a1",
         right_on="a2",
-        match_score=0.5,
+        match_score=0.45,
         return_score=True,
         analyzer=analyzer,
     )
@@ -32,9 +29,10 @@ def test_fuzzy_join(analyzer, how):
     df_joined2 = fuzzy_join(
         df2,
         df1,
+        how="left",
         left_on="a2",
         right_on="a1",
-        match_score=0.45,
+        match_score=0.35,
         return_score=True,
         analyzer=analyzer,
     )
@@ -42,11 +40,12 @@ def test_fuzzy_join(analyzer, how):
     assert df_joined2.shape == (len(df2.dropna()), n_cols)
 
     df_joined3 = fuzzy_join(
-        df2,
         df1,
-        left_on="a2",
-        right_on="a1",
-        match_score=0.45,
+        df2,
+        how="right",
+        right_on="a2",
+        left_on="a1",
+        match_score=0.35,
         return_score=True,
         analyzer=analyzer,
     )
@@ -54,41 +53,11 @@ def test_fuzzy_join(analyzer, how):
 
     df1["a2"] = 1
 
-    df_how = fuzzy_join(
-        df1,
-        df2,
-        left_on="a1",
-        right_on="a2",
-        analyzer=analyzer,
-        how=how,
-        suffixes=("", "r"),
-    )
-
-    if how == "left":
-        pd.testing.assert_frame_equal(df_how, df1.dropna())
-    if how == "right":
-        assert df_how.shape == (len(df1.dropna()), df1.shape[1])
-
     df_on = fuzzy_join(df_joined, df1, on="a1", analyzer=analyzer, suffixes=("1", "2"))
     assert ("a11" and "a12") in df_on.columns
 
     df2["a1"] = 1
 
-    with pytest.raises(
-        ValueError,
-        match=r"Columns overlap but no suffix",
-    ):
-        fuzzy_join(df1, df2, left_on="a1", right_on="a2", suffixes=("", ""))
-
-    if analyzer == "word":
-        with pytest.warns(
-            UserWarning,
-            match=r"Column names overlaps.",
-        ):
-            fuzzy_join(
-                df1, df2, left_on="a1", right_on="a2", suffixes=("a", "a"), how=how
-            )
-
     df = fuzzy_join(
         df1,
         df2,
@@ -100,24 +69,24 @@ def test_fuzzy_join(analyzer, how):
     )
     assert ("a1l" and "a1r") in df.columns
 
-    df_dropna = fuzzy_join(
-        left=df1,
-        right=df2,
-        left_on="a1",
-        right_on="a2",
-        analyzer=analyzer,
-        match_score=0.5,
-        return_score=True,
-        drop_unmatched=True,
-    )
-    assert df_dropna.shape[0] == df_joined.dropna().shape[0]
+
+def test_fuzzy_join_dtypes():
+    """
+    Test that the dtypes of dataframes are maintained after join
+    """
+    a = pd.DataFrame({"col1": ["aaa", "bbb"], "col2": [1, 2]})
+    b = pd.DataFrame({"col1": ["aaa_", "bbb_"], "col3": [1, 2]})
+    c = fuzzy_join(a, b, on="col1")
+    assert a.dtypes["col2"].kind == "i"
+    assert c.dtypes["col2"] == a.dtypes["col2"]
+    assert c.dtypes["col3"] == b.dtypes["col3"]
 
 
 @pytest.mark.parametrize(
-    "analyzer, how, suffixes, on",
-    [("a_blabla", "k_blabla", ["a", "b", "c"], ["a"]), (1, 34, [1, 2, 3], 3)],
+    "analyzer, on, how",
+    [("a_blabla", ["a"], "left"), (1, 3, "right")],
 )
-def test_parameters_error(analyzer, how, suffixes, on):
+def test_parameters_error(analyzer, on, how):
     """Testing if correct errors are raised when wrong parameter values are given."""
     df1 = pd.DataFrame({"a": ["ana", "lala", "nana"], "b": [1, 2, 3]})
     df2 = pd.DataFrame({"a": ["anna", "lala", "ana", "sana"], "c": [5, 6, 7, 8]})
@@ -127,44 +96,57 @@ def test_parameters_error(analyzer, how, suffixes, on):
             f"analyzer should be either 'char', 'word' or 'char_wb', got {analyzer!r}"
         ),
     ):
-        fuzzy_join(df1, df2, on="a", analyzer=analyzer)
-    with pytest.raises(
-        ValueError,
-        match=f"how should be either 'left', 'right' or 'all', got {how!r}",
-    ):
-        fuzzy_join(df1, df2, on="a", how=how)
+        fuzzy_join(df1, df2, on="a", analyzer=analyzer, how=how)
     with pytest.raises(
-        ValueError, match="Invalid number of suffixes: expected 2, got 3"
+        KeyError,
+        match=r"invalid type",
     ):
-        fuzzy_join(df1, df2, on="a", suffixes=suffixes)
+        fuzzy_join(df1, df2, on=on, how=how)
+
+
+def test_missing_keys():
+    a = pd.DataFrame({"col1": ["aaa", "bbb"], "col2": [1, 2]})
+    b = pd.DataFrame({"col1": ["aaa_", "bbb_"], "col3": [1, 2]})
     with pytest.raises(
-        ValueError,
-        match=(
-            "Parameter 'left_on', 'right_on' or 'on' has invalid type, expected string"
-        ),
+        KeyError,
+        match=r"Required parameter missing",
     ):
-        fuzzy_join(df1, df2, on=on)
+        fuzzy_join(a, b, left_on="col1")
 
 
-def test_duplicate_columns_drop():
-    df1 = pd.DataFrame(
-        {"a1": ["ana", "lala", "nana et sana", np.NaN], "b": [14, 13, 51, np.NaN]}
-    )
-    df2 = pd.DataFrame(
+def test_drop_unmatched():
+    a = pd.DataFrame({"col1": ["aaaa", "bbb", "ddd dd"], "col2": [1, 2, 3]})
+    b = pd.DataFrame({"col1": ["aaa_", "bbb_", "cc ccc"], "col3": [1, 2, 3]})
+    c = fuzzy_join(a, b, on="col1", match_score=0.5, drop_unmatched=True)
+    assert c.shape == (2, 4)
+
+    a = pd.DataFrame({"col1": ["aaaa", "bbb", "ddd dd"], "col2": [1, 2, 3]})
+    b = pd.DataFrame({"col1": ["aaa_", "bbb_", "cc ccc"], "col3": [1, 2, 3]})
+    c = fuzzy_join(a, b, on="col1", match_score=0.5)
+    assert sum(c["col3"].isna()) > 0
+
+
+def test_how_param():
+    """
+    Test correct shape of left and right joins. 
+    Also test if an error is raised when an incorrect parameter value is passed.
+    """
+    a = pd.DataFrame({"col1": ["aaaa", "bbb", "ddd dd"], "col2": [1, 2, 3]})
+    b = pd.DataFrame(
         {
-            "a2": ["anna", "lala et nana", "lana", "sana", np.NaN],
-            "b": [1, 3, 6, 7, np.NaN],
+            "col2": ["aaa_", "bbb_", "ccc", "ddd", "eee", "fff"],
+            "col3": [1, 2, 3, 4, 5, 6],
         }
     )
 
-    for h in ["left", "right"]:
-        df_joined = fuzzy_join(
-            left=df1,
-            right=df2,
-            left_on="a1",
-            right_on="a2",
-            analyzer="char",
-            suffixes=("1", "1"),
-            how=h,
-        )
-        pd.testing.assert_frame_equal(df_joined, df_joined.dropna())
+    c = fuzzy_join(a, b, left_on="col1", right_on="col2", how="left")
+    assert c.shape == (len(a), 4)
+
+    c = fuzzy_join(a, b, left_on="col1", right_on="col2", how="right")
+    assert c.shape == (len(b), 4)
+
+    with pytest.raises(
+        ValueError,
+        match=r"how should be either 'left' or 'right', got",
+    ):
+        c = fuzzy_join(a, b, how="inner")