feat(RFC): Adds altair.datasets

[dangotbanned]

Related

• `vega_datasets` source #3150
• feat: allow for creating new_series / from_dict without specifying a backend narwhals-dev/narwhals#876 (comment)

Description

Providing a minimal, but up-to-date source for https://github.com/vega/vega-datasets.

This PR takes a different approach to that of https://github.com/altair-viz/vega_datasets, notably:

• No datasets are included in the package
  • Instead, included is only a single 18.7KB file metadata.parquet
  • The file describes all versions of all datasets
    • provided they are accessible via both npm and github
• Strong support for typing
  • Annotations are generated from the metadata itself
  • https://github.com/vega/altair/blob/9e9deeb95668d2c4e7d30311e85a8f9f6acdc88c/altair/datasets/_typing.py
• So far, 5 backends have been implemented, instead of only pandas
  • These provide precise IDE completions, with a lot of help from https://github.com/narwhals-dev/narwhals
• Users can opt-in to caching remote dataset requests
  • With the "polars" backend, the slowest I've had on a cache-hit is 0.1s to load
    • https://cdn.jsdelivr.net/npm/[email protected]/data/flights-200k.json

Examples

These all come from the docstrings of:

• Loader
• Loader.from_backend
• Loader.__call__

import altair as alt
from altair.datasets import Loader

data = Loader.from_backend("polars")
>>> data
Loader[polars]

cars = data("cars")

>>> type(cars)
polars.dataframe.frame.DataFrame

data = Loader.from_backend("pandas")
cars = data("cars")

>>> type(cars)
pandas.core.frame.DataFrame

data = Loader.from_backend("pandas[pyarrow]")
cars = data("cars", tag="v1.29.0")

>>> type(cars)
pandas.core.frame.DataFrame

>>> cars.dtypes
Name                string[pyarrow]
Miles_per_Gallon    double[pyarrow]
Cylinders            int64[pyarrow]
Displacement        double[pyarrow]
Horsepower           int64[pyarrow]
Weight_in_lbs        int64[pyarrow]
Acceleration        double[pyarrow]
Year                string[pyarrow]
Origin              string[pyarrow]
dtype: object

data = Loader.from_backend("pandas")
source = data("stocks", tag="v2.10.0")

>>> source.columns
Index(['symbol', 'date', 'price'], dtype='object')

data = Loader.from_backend("pyarrow")
source = data("stocks", tag="v2.10.0")

>>> source.column_names
['symbol', 'date', 'price']

Tasks

• Add test suite
  • (https://github.com/vega/altair/blob/b70aef883721ce1ce905e1ec8e82938eb4859257/tests/test_datasets.py)
  • Test/add handling for no remote connection (90428a6)
• Finish docs
  • altair.datasets
    • Loader.__call__
      • Parameters
      • Examples
    • _readers._Reader
    • _typing.Metadata (Align with revised descriptions from (c572180))
  • tools.datasets
    • Application
    • models.py
    • github.py
    • npm.py
    • semver.py
• Mitigating _PyArrowReader JSON limitation (3fbc759), (4f5b4de)

Resolved

Investigate bundling metadata

• Investigating bundling metadata (22a5039), (1792340)
  • Depending on how well the compression scales, it might be reasonable to include this for some number of versions
  • Deliberately including redundant info early on - can always chip away at it later

npm does not have every version available GitHub

• Sources
  • npm/vega-datasets
    • Fixed with: https://data.jsdelivr.com/v1/packages/npm/vega-datasets
  • https://github.com/vega/vega-datasets/tags
• Known missing
  • https://github.com/vega/vega-datasets/releases/tag/v1.9.0
  • https://github.com/vega/vega-datasets/releases/tag/v1.6.0
  • https://github.com/vega/vega-datasets/releases/tag/v1.4.0
• feat(DRAFT): Add a source for available npm versions
• Need to add some handling to invalidate these entries returned from list-repository-tags once confirmed they cannot be requested from npm
  • Can technically request from github, but during testing this was much slower
  • Also, these versions would not have been available from https://github.com/altair-viz/vega_datasets, since that only used npm

Plan strategy for user-configurable dataset cache

• Everything so far has been building the tools for a compact bundled index
  • 1, 2, 3, 4, 5
  • Refreshing the index would not be included in altair, each release would simply ship with changes baked in
• Trying to avoid bloating altair package size with datasets
• User-facing
  • Goal of requesting each unique dataset version once
    • The user cache would not need to be updated between altair versions
  • Some kind of opt-in config to say store the datasets in this directory please
    • Basic solution would be defining an env variable like ALTAIR_DATASETS_DIR
    • When not provided, always perform remote requests
      • User motivation would be that it would be faster to enable caching

Deferred

Reducing cache footprint

• e.g. storing the .(csv|tsv|json) files as .parquet
• Need to do more testing on this though to ensure
  • the shape of each dataset is preserved
  • where relevant - intentional errors remain intact

Investigate providing a decorator to add a backend

• Will be trivial for the user-side, since they don't need to be concerned about imports
• Just need to provide these attributes:
  • _name: LiteralString
  • _read_fn: dict[Extension, Callable[..., IntoDataFrameT]]
  • _scan_fn: dict[_ExtensionScan, Callable[..., IntoFrameT]]

Investigate ways to utilize (https://github.com/vega/vega-datasets/blob/main/SOURCES.md)

• Can use a similar approach to feat: Generate expr method signatures, docs #3600
• Explored a little locally, finishing it off in this PR would add quite a lot of markdown parsing code

Provide more meaningful info on the state of ALTAIR_DATASETS_DIR

• How many datasets, size (per & total)?
• What version range does a given sha cover?
• Blocked: Running into issues with
  • pandas/pyarrow group_by warnings
  • min and max return all nulls in pl.Enum pola-rs/polars#18394
  • Missing nw.Expr.(first|last)
  • nw.Expr.(head|tail)(1) not equivalent in a group_by().agg(...) context
    • pandas -> scalar
    • polars -> list
  • pl.Enum translating to non-ordered pd.Categorical

polars-native solution

from __future__ import annotations

from pathlib import Path

import polars as pl
from altair.datasets import Loader, _readers

data = Loader.from_backend("polars")

# NOTE: Enable caching, populate with some responses
data.cache_dir = Path.home() / ".altair_cache"
data("cars")
data("cars", tag="v1.5.0")
data("movies")
data("movies", tag="v1.24.0")
data("jobs")


if cache_dir := data.cache_dir:
    cached_stems: tuple[str, ...] = tuple(fp.stem for fp in cache_dir.iterdir())
else:
    msg = "Datasets cache unset"
    raise TypeError(msg)

# NOTE: Lots of redundancies, many urls point to the same data (sha)
>>> pl.read_parquet(_readers._METADATA).shape
# (2879, 9)

# NOTE: Version range per sha
tag_sort: pl.Expr = pl.col("tag").sort()
tag_range: pl.Expr = pl.concat_str(tag_sort.first(), tag_sort.last(), separator=" - ")

# NOTE: Producing a name only when the file is already in the cache
file_name: pl.Expr = pl.when(pl.col("sha").is_in(cached_stems)).then(
    pl.concat_str("sha", "suffix")
)

cache_summary: pl.DataFrame = (
    pl.scan_parquet(_readers._METADATA)
    .group_by("dataset_name", "suffix", "size", "sha")
    .agg(tag_range=tag_range)
    .select(pl.exclude("sha"), file_name=file_name)
    .sort("dataset_name", "size")
    .collect()
)

>>> cache_summary.shape
# (116, 5)

>>> cache_summary.head(10)

shape: (10, 5)
┌───────────────┬────────┬─────────┬───────────────────┬─────────────────────────────────┐
│ dataset_name  ┆ suffix ┆ size    ┆ tag_range         ┆ file_name                       │
│ ---           ┆ ---    ┆ ---     ┆ ---               ┆ ---                             │
│ str           ┆ str    ┆ i64     ┆ str               ┆ str                             │
╞═══════════════╪════════╪═════════╪═══════════════════╪═════════════════════════════════╡
│ 7zip          ┆ .png   ┆ 3969    ┆ v1.5.0 - v2.10.0  ┆ null                            │
│ airports      ┆ .csv   ┆ 210365  ┆ v1.5.0 - v2.10.0  ┆ 608ba6d51fa70584c3fa1d31eb9453… │
│ annual-precip ┆ .json  ┆ 266265  ┆ v1.29.0 - v2.10.0 ┆ null                            │
│ anscombe      ┆ .json  ┆ 1703    ┆ v1.5.0 - v2.10.0  ┆ null                            │
│ barley        ┆ .json  ┆ 8487    ┆ v1.5.0 - v2.10.0  ┆ 8dc50de2509b6e197ce95c24c98f90… │
│ birdstrikes   ┆ .csv   ┆ 1223329 ┆ v2.0.0 - v2.10.0  ┆ null                            │
│ birdstrikes   ┆ .json  ┆ 4183924 ┆ v1.5.0 - v1.31.1  ┆ null                            │
│ budget        ┆ .json  ┆ 374289  ┆ v1.5.0 - v2.8.1   ┆ null                            │
│ budget        ┆ .json  ┆ 391353  ┆ v2.9.0 - v2.10.0  ┆ null                            │
│ budgets       ┆ .json  ┆ 18079   ┆ v1.5.0 - v2.10.0  ┆ 8a909e24f698a3b0f6c637c30ec95e… │
└───────────────┴────────┴─────────┴───────────────────┴─────────────────────────────────┘

[mattijn]

I'm reviewing as an average user of Altair and for this use-case it is probably an associate professor who will need to update all her/his lecture materials at the evening before the semester starts.

What would it be great if we could say:

# old way (this is deprecated)
from vega_datasets import data

# new way (this will be awesome)
from altair.datasets import data

And everything else is still functioning. So this still works:

source_url = data.cars.url
source_pandas = data.cars()

But, the awesome thing that we provide with this PR is:

source_polars = data.cars(backend="polars")  # or `engine=`

Or polars with pyarrow dtypes:

source_pl_pa = data.cars(backend="polars[pyarrow]")  # or `engine=`

If it is like this than I'm fine with engine or backend as argument name. And than within this function we call the agnostic Loader using the dataset and backend choice. All in all, in my humble opinion, awesomeness.

[dangotbanned]

@jonmmease I just tried updating this branch, seems to be some vegafusion issues?

9d97096 (#3631)

Update

Resolved in #3702