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)

Status

Waiting on the next vega-datasets release.
Once there is a stable datapackage.json available - there is quite a lot of tools/datasets that can be simplified/removed.

• 3.0.0 Release vega-datasets#654

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, 4 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/vega-datasets@v2.9.0/data/flights-200k.json

Examples

These all come from the docstrings of:

• Loader
• Loader.from_backend
• Loader.__call__

from altair.datasets import Loader

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

cars = load("cars")

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

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

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

load = Loader.from_backend("pandas[pyarrow]")
cars = load("cars")

>>> 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                timestamp[ns][pyarrow]
Origin                     string[pyarrow]
dtype: object

load = Loader.from_backend("pandas")
source = load("stocks")

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

load = Loader.from_backend("pyarrow")
source = load("stocks")

>>> 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]]

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… │
└───────────────┴────────┴─────────┴───────────────────┴─────────────────────────────────┘