Function _get_dataset

plotly/data/__init__.py:379–430

Loads the dataset using the specified backend. Notice that the available backends are 'pandas', 'polars', 'pyarrow' and they all have a `read_csv` function (pyarrow has it via pyarrow.csv). Therefore we can dynamically load the library using `importlib.import_module` and then call

(d, return_type)

def _get_dataset(d, return_type):
    """
    Loads the dataset using the specified backend.

    Notice that the available backends are 'pandas', 'polars', 'pyarrow' and they all have
    a `read_csv` function (pyarrow has it via pyarrow.csv). Therefore we can dynamically
    load the library using `importlib.import_module` and then call
    `backend.read_csv(filepath)`.

    Parameters
    ----------
    d: str
        Name of the dataset to load.

    return_type: {'pandas', 'polars', 'pyarrow', 'modin', 'cudf'}
        Type of the resulting dataframe

    Returns
    -------
    Dataframe of `return_type` type
    """
    filepath = os.path.join(
        os.path.dirname(os.path.dirname(__file__)),
        "package_data",
        "datasets",
        d + ".csv.gz",
    )

    if return_type not in AVAILABLE_BACKENDS:
        msg = (
            f"Unsupported return_type. Found {return_type}, expected one "
            f"of {AVAILABLE_BACKENDS}"
        )
        raise NotImplementedError(msg)

    try:
        if return_type == "pyarrow":
            module_to_load = "pyarrow.csv"
        elif return_type == "modin":
            module_to_load = "modin.pandas"
        else:
            module_to_load = return_type
        backend = import_module(module_to_load)
    except ModuleNotFoundError:
        msg = f"return_type={return_type}, but {return_type} is not installed"
        raise ModuleNotFoundError(msg)

    try:
        return backend.read_csv(filepath)
    except Exception as e:
        msg = f"Unable to read '{d}' dataset due to: {e}"
        raise Exception(msg).with_traceback(e.__traceback__)