add sizes of dataframes and total sizes on disk (#1503)

Co-authored-by: Riccardo Cappuzzo <7548232+rcap107@users.noreply.github.com>
Co-authored-by: Vincent M <maladiere.vincent@yahoo.fr>
Co-authored-by: Riccardo Cappuzzo <riccardo.cappuzzo@gmail.com>

Author: 3 people

skrub/datasets/_fetching.py

@@ -14,6 +14,7 @@ def fetch_employee_salaries(data_home=None, split="all"):
         active, permanent employees of Montgomery County, MD paid in calendar
         year 2016. This dataset is a copy of https://www.openml.org/d/42125
         where some features are dropped to avoid data leaking.
+        Size on disk: 1.3MB.
 
     .. note::
 
@@ -39,10 +40,12 @@ def fetch_employee_salaries(data_home=None, split="all"):
     bunch : sklearn.utils.Bunch
         A dictionary-like object with the following keys:
 
-        - employee_salaries : pd.DataFrame, the dataframe
-        - X : pd.DataFrame, features, i.e. the dataframe without the target labels
-        - y : pd.DataFrame, target labels
-        - metadata : a dictionary containing the name, description, source and target
+        - ``employee_salaries`` : pd.DataFrame, the dataframe. Shape: (9228, 9)
+        - ``X`` : pd.DataFrame, features, i.e. the dataframe without the target
+          labels. Shape: (9228, 8)
+        - ``y`` : pd.DataFrame, target labels. Shape: (9228, 1)
+        - ``metadata`` : a dictionary containing the name, description, source and
+          target
     """
     if split not in ["train", "test", "all"]:
         raise ValueError(
@@ -75,6 +78,7 @@ def fetch_medical_charge(data_home=None):
         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.
+        Size on disk: 36MB.
 
     Parameters
     ----------
@@ -86,10 +90,12 @@ def fetch_medical_charge(data_home=None):
     bunch : sklearn.utils.Bunch
         A dictionary-like object with the following keys:
 
-        - medical_charge : pd.DataFrame, the dataframe
-        - X : pd.DataFrame, features, i.e. the dataframe without the target labels
-        - y : pd.DataFrame, target labels
-        - metadata : a dictionary containing the name, description, source and target
+        - ``medical_charge`` : pd.DataFrame, the dataframe. Shape: (163065, 12)
+        - ``X`` : pd.DataFrame, features, i.e. the dataframe without the target
+          labels. Shape: (163065, 11)
+        - ``y`` : pd.DataFrame, target labels. Shape: (163065, 1)
+        - ``metadata`` : a dictionary containing the name, description, source and
+          target
     """
     return load_simple_dataset("medical_charge", data_home)
 
@@ -99,7 +105,7 @@ def fetch_midwest_survey(data_home=None):
         https://github.com/skrub-data/skrub-data-files
 
     Description of the dataset:
-        Survey to know if people self-identify as Midwesterners.
+        Survey to know if people self-identify as Midwesterners. Size on disk: 504KB.
 
     Parameters
     ----------
@@ -111,10 +117,12 @@ def fetch_midwest_survey(data_home=None):
     bunch : sklearn.utils.Bunch
         A dictionary-like object with the following keys:
 
-        - midwest_survey : pd.DataFrame, the dataframe
-        - X : pd.DataFrame, features, i.e. the dataframe without the target labels
-        - y : pd.DataFrame, target labels
-        - metadata : a dictionary containing the name, description, source and target
+        - ``midwest_survey`` : pd.DataFrame, the dataframe. Shape: (2494, 29)
+        - ``X`` : pd.DataFrame, features, i.e. the dataframe without the target
+          labels. Shape: (2494, 28)
+        - ``y`` : pd.DataFrame, target labels. Shape: (2494, 1)
+        - ``metadata`` : a dictionary containing the name, description, source and
+          target
     """
     return load_simple_dataset("midwest_survey", data_home)
 
@@ -125,7 +133,7 @@ def fetch_open_payments(data_home=None):
 
     Description of the dataset:
         Payments given by healthcare manufacturing companies to medical doctors
-        or hospitals.
+        or hospitals. Size on disk: 8.7MB.
 
     Parameters
     ----------
@@ -137,10 +145,12 @@ def fetch_open_payments(data_home=None):
     bunch : sklearn.utils.Bunch
         A dictionary-like object with the following keys:
 
-        - open_payments : pd.DataFrame, the dataframe
-        - X : pd.DataFrame, features, i.e. the dataframe without the target labels
-        - y : pd.DataFrame, target labels
-        - metadata : a dictionary containing the name, description, source and target
+        - ``open_payments`` : pd.DataFrame, the dataframe. Shape: (73558, 6)
+        - ``X`` : pd.DataFrame, features, i.e. the dataframe without the target
+          labels. Shape: (73558, 5)
+        - ``y`` : pd.DataFrame, target labels. Shape: (73558, 1)
+        - ``metadata`` : a dictionary containing the name, description, source
+          and target
     """
     return load_simple_dataset("open_payments", data_home)
 
@@ -153,7 +163,7 @@ def fetch_traffic_violations(data_home=None):
         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.
+        the officer issuing the violation will not be published. Size on disk: 736MB.
 
     Parameters
     ----------
@@ -165,10 +175,12 @@ def fetch_traffic_violations(data_home=None):
     bunch : sklearn.utils.Bunch
         A dictionary-like object with the following keys:
 
-        - traffic_violations : pd.DataFrame, the dataframe
-        - X : pd.DataFrame, features, i.e. the dataframe without the target labels
-        - y : pd.DataFrame, target labels
-        - metadata : a dictionary containing the name, description, source and target
+        - ``traffic_violations`` : pd.DataFrame, the dataframe. Shape: (1578154, 43)
+        - ``X`` : pd.DataFrame, features, i.e. the dataframe without the target
+          labels. Shape: (1578154, 42)
+        - ``y`` : pd.DataFrame, target labels. Shape: (1578154, 1)
+        - ``metadata`` : a dictionary containing the name, description, source and
+          target
     """
     return load_simple_dataset("traffic_violations", data_home)
 
@@ -179,7 +191,7 @@ def fetch_drug_directory(data_home=None):
 
     Description of the dataset:
         Product listing data submitted to the U.S. FDA for all unfinished,
-        unapproved drugs.
+        unapproved drugs. Size on disk: 44MB.
 
     Parameters
     ----------
@@ -191,10 +203,12 @@ def fetch_drug_directory(data_home=None):
     bunch : sklearn.utils.Bunch
         A dictionary-like object with the following keys:
 
-        - drug_directory : pd.DataFrame, the dataframe
-        - X : pd.DataFrame, features, i.e. the dataframe without the target labels
-        - y : pd.DataFrame, target labels
-        - metadata : a dictionary containing the name, description, source and target
+        - ``drug_directory`` : pd.DataFrame, the dataframe. Shape: (120215, 21)
+        - ``X`` : pd.DataFrame, features, i.e. the dataframe without the target
+          labels. Shape: (120215, 20)
+        - ``y`` : pd.DataFrame, target labels. Shape: (120215, 1)
+        - ``metadata`` : a dictionary containing the name, description, source and
+          target
     """
     return load_simple_dataset("drug_directory", data_home)
 
@@ -211,6 +225,7 @@ def fetch_credit_fraud(data_home=None, split="train"):
 
     Baskets contain at least one product each, so aggregation then joining operations
     are required to build a design matrix.
+    Size on disk: 16MB.
 
     Parameters
     ----------
@@ -225,10 +240,12 @@ def fetch_credit_fraud(data_home=None, split="train"):
     bunch : sklearn.utils.Bunch
         A dictionary-like object with the following keys:
 
-        - baskets : pd.DataFrame, table containing baskets ID and target
-        - product : pd.DataFrame, table containing features about products contained in
-          baskets
-        - metadata : a dictionary containing the name, description, source and target
+        - ``baskets`` : pd.DataFrame, table containing baskets ID and target.
+        Shape: (92790, 2)
+        - ``product`` : pd.DataFrame, table containing features about products
+          contained in baskets. Shape: (163357, 7)
+        - ``metadata`` : a dictionary containing the name, description, source and
+          target
     """
     if split not in ["train", "test", "all"]:
         raise ValueError(
@@ -254,6 +271,7 @@ def fetch_toxicity(data_home=None):
     consists in only two columns:
     - `text`: the text of the comment
     - `is_toxic`: whether or not the comment is toxic
+    Size on disk: 220KB.
 
     Parameters
     ----------
@@ -265,10 +283,12 @@ def fetch_toxicity(data_home=None):
     bunch : sklearn.utils.Bunch
         A dictionary-like object with the following keys:
 
-        - toxicity : pd.DataFrame, the dataframe
-        - X : pd.DataFrame, features, i.e. the dataframe without the target labels
-        - y : pd.DataFrame, target labels
-        - metadata : a dictionary containing the name, description, source and target
+        - ``toxicity`` : pd.DataFrame, the dataframe. Shape: (1000, 2)
+        - ``X`` : pd.DataFrame, features, i.e. the dataframe without the target
+          labels. Shape: (1000, 1)
+        - ``y`` : pd.DataFrame, target labels. Shape: (1000, 1)
+        - ``metadata`` : a dictionary containing the name, description, source and
+          target
     """
     return load_simple_dataset("toxicity", data_home)
 
@@ -279,7 +299,7 @@ def fetch_videogame_sales(data_home=None):
 
     This is a regression use-case, where the single table contains information
     about videogames such as the publisher and platform, and the goal is to
-    predict the number of sales worldwide.
+    predict the number of sales worldwide. Size on disk: 1.8MB.
 
     .. warning::
 
@@ -298,10 +318,11 @@ def fetch_videogame_sales(data_home=None):
     bunch : sklearn.utils.Bunch
         A dictionary-like object with the following keys:
 
-        - videogame_sales : pd.DataFrame, the full dataframe
-        - X : pd.DataFrame, features, i.e. the dataframe without the target labels
-        - y : pd.DataFrame, target labels
-        - metadata : a dictionary containing the name, source and target
+        - ``videogame_sales`` : pd.DataFrame, the full dataframe. Shape: (16572, 11)
+        - ``X`` : pd.DataFrame, features, i.e. the dataframe without the target
+          labels. Shape: (16572, 5)
+        - ``y`` : pd.DataFrame, target labels. Shape: (16572, 1)
+        - ``metadata`` : a dictionary containing the name, source and target
     """
 
     result = load_simple_dataset("videogame_sales", data_home)
@@ -317,7 +338,7 @@ def fetch_bike_sharing(data_home=None):
 
     This is a regression use-case, where the goal is to predict demand for a
     bike-sharing service. The features are the dates and holiday and weather
-    information.
+    information. Size on disk: 1.3MB.
 
     Parameters
     ----------
@@ -329,10 +350,11 @@ def fetch_bike_sharing(data_home=None):
     bunch : sklearn.utils.Bunch
         A dictionary-like object with the following keys:
 
-        - bike_sharing : pd.DataFrame, the full dataframe
-        - X : pd.DataFrame, features, i.e. the dataframe without the target labels
-        - y : pd.DataFrame, target labels
-        - metadata : a dictionary containing the name and target
+        - ``bike_sharing``: pd.DataFrame, the full dataframe. Shape: (17379, 11)
+        - ``X`` : pd.DataFrame, features, i.e. the dataframe without the target labels.
+          Shape: (17379, 10)
+        - ``y`` : pd.DataFrame, target labels. Shape: (17379, 1)
+        - ``metadata`` : a dictionary containing the name and target
     """
 
     return load_simple_dataset("bike_sharing", data_home)
@@ -344,6 +366,7 @@ def fetch_movielens(data_home=None):
 
     This is a regression use-case, where the goal is to predict movie ratings.
     More details are provided in the output's ``metadata['description']``.
+    Size on disk: 3.6MB.
 
     Parameters
     ----------
@@ -355,9 +378,9 @@ def fetch_movielens(data_home=None):
     bunch : sklearn.utils.Bunch
         A dictionary-like object with the following keys:
 
-        - movies : pd.DataFrame, movie ID, title and genres
-        - ratings: pd.DataFrame, user ID, movie ID, rating
-        - metadata : a dictionary containing the name source and description
+        - ``movies`` : pd.DataFrame, movie ID, title and genres. Shape: (9742, 3)
+        - ``ratings``: pd.DataFrame, user ID, movie ID, rating. Shape: (100836, 4)
+        - ``metadata`` : a dictionary containing the name source and description
     """
 
     return load_dataset_files("movielens", data_home)
@@ -368,6 +391,7 @@ def fetch_flight_delays(data_home=None):
         https://github.com/skrub-data/skrub-data-files
 
     This is a regression use-case, where the goal is to predict flight delays.
+    Size on disk: 657MB.
 
     Parameters
     ----------
@@ -379,20 +403,20 @@ def fetch_flight_delays(data_home=None):
     bunch : sklearn.utils.Bunch
         A dictionary-like object with the following keys:
 
-        - flights: information about the flights, including departure and
-          arrival airports, and delay.
-        - airports: information about airports, such as city and coordinates.
+        - ``flights``: information about the flights, including departure and
+          arrival airports, and delay. Shape: (2370030, 12)
+        - ``airports``: information about airports, such as city and coordinates.
           The airport's ``iata`` can be matched to the flights' ``Origin`` and
-          ``Dest``.
-        - weather: weather data that could be used to help improve the delay
+          ``Dest``. Shape: (3376, 7)
+        - ``weather``: weather data that could be used to help improve the delay
           predictions. Note the weather data is not measured at the airports
           directly but at weather stations, whose location and information is
-          provided in ``stations``.
-        - stations: information about the weather stations. ``weather`` and
+          provided in ``stations``. Shape: (11282238, 5)
+        - ``stations``: information about the weather stations. ``weather`` and
           ``stations`` can be joined on their ``ID`` columns. Weather stations
           can only be matched to the nearest airport based on the latitude and
-          longitude.
-        - metadata : a dictionary containing the name  of the dataset.
+          longitude. Shape: (124245, 9)
+        - ``metadata`` : a dictionary containing the name  of the dataset.
     """
     return load_dataset_files("flight_delays", data_home)
 
@@ -404,7 +428,7 @@ def fetch_country_happiness(data_home=None):
     This is a regression use-case, where the goal is to predict the happiness
     index. The dataset contains data from the `2022 World Happiness Report
     <https://worldhappiness.report/>`_, and from `the World Bank open data
-    platform <https://data.worldbank.org/>`_.
+    platform <https://data.worldbank.org/>`_. Size on disk: 64KB.
 
     Parameters
     ----------
@@ -416,8 +440,12 @@ def fetch_country_happiness(data_home=None):
     bunch : sklearn.utils.Bunch
         A dictionary-like object with the following keys:
 
-        - ``happiness_report``: dataframe, data from the world happiness report
-        - ``GDP_per_capita``, ``life_expectancy``, ``legal_rights_index``:
-          corresponding tables from the World Bank.
+        - ``happiness_report``: dataframe, data from the world happiness report.
+          Shape: (146, 12)
+        - ``GDP_per_capita``: dataframe from the World Bank. Shape: (262, 2)
+        - ``life_expectancy``: dataframe from the World Bank. Shape: (260, 2)
+        - ``legal_rights_index``: dataframe from the World Bank. Shape: (238, 2)
+        - ``metadata`` : a dictionary containing the name of the dataset, a
+          description and the sources.
     """
     return load_dataset_files("country_happiness", data_home)