Add support for huggingface datasets to vortex-python

.github/workflows/ci.yml

@@ -118,6 +118,14 @@ jobs:
           uv run --all-packages pytest --benchmark-disable -n auto test/
         working-directory: vortex-python/
 
+      # The datasets integration relies on private `datasets` APIs, so exercise the oldest
+      # supported release in addition to the locked one.
+      - name: Pytest - Vortex (datasets 4.x floor)
+        run: |
+          uv run --all-packages --with 'datasets==4.0.0' --with 'huggingface-hub<1' \
+            pytest --benchmark-disable test/test_hf_datasets.py
+        working-directory: vortex-python/
+
       - name: Setup benchmark environment
         run: sudo bash scripts/setup-benchmark.sh
 

docs/api/python/datasets.rst

@@ -0,0 +1,52 @@
+Hugging Face Datasets
+=====================
+
+Vortex files can be loaded directly as Hugging Face ``datasets`` objects. Install the optional
+dependencies with:
+
+.. code-block:: bash
+
+    pip install vortex-data[hf]
+
+``vortex.datasets.load_dataset`` accepts a local path or glob, a URL readable by Vortex
+(``https://``, ``s3://``, ``gs://``, ``az://``), a Hugging Face Hub dataset repository id, or an
+``hf://`` URI of the form ``hf://datasets/<org>/<name>[@<revision>][/<path-or-glob>]`` (revisions
+containing ``/`` must be percent-encoded, e.g. ``@refs%2Fconvert%2Fparquet``).
+Unlike ``datasets.load_dataset``, it defaults to ``streaming=True`` and returns a streaming
+dataset that keeps Vortex in charge of reading. Column selection, Vortex filter expressions, and
+row limits are pushed down into each Vortex scan before examples are yielded to Hugging Face
+transforms.
+
+Hub repositories are streamed in place: files are read with HTTP range requests, so only the
+projected columns and matching rows are ever transferred. Private and gated repositories
+authenticate with the ``token`` argument or the locally saved login. Files are downloaded (with
+the usual Hub caching) only when ``streaming=False`` or ``local_files_only=True``.
+
+.. code-block:: python
+
+    import vortex as vx
+
+    # Stream a local directory of Vortex files as a Hugging Face IterableDataset.
+    ds = vx.datasets.load_dataset("path/to/dir", split="train")
+
+    # Stream a Hub dataset directly, without downloading it.
+    ds = vx.datasets.load_dataset("username/dataset", split="train")
+
+    # hf:// URIs select a revision and a path or glob within the repository.
+    ds = vx.datasets.load_dataset("hf://datasets/username/dataset@main/data/*.vortex", split="train")
+
+    # Column projection and Vortex filter expressions are pushed into the scan.
+    ds = ds.select_columns(["text", "label"]).filter(vx.expr.column("label") == 1)
+
+    for example in ds.take(100):
+        ...
+
+    # Materialize an in-memory dataset instead, optionally selecting splits via data_files.
+    splits = vx.datasets.load_dataset(
+        "username/dataset",
+        data_files={"train": "train-*.vortex", "validation": "validation.vortex"},
+        streaming=False,
+    )
+
+When loading multiple splits, pass a ``data_files`` mapping from split name to file patterns; a
+list of splits without such a mapping cannot be resolved and raises an error.

docs/api/python/index.rst

@@ -16,7 +16,7 @@ Optional integrations can be installed as extras:
 
 .. code-block:: bash
 
-    pip install vortex-data[polars,pandas,numpy,duckdb,ray]
+    pip install vortex-data[polars,pandas,numpy,duckdb,ray,hf]
 
 
 Compatibility
@@ -66,5 +66,6 @@ API Reference
    io
    store
    dataset
+   datasets
    runtime
    type_aliases

vortex-python/pyproject.toml

@@ -43,6 +43,8 @@ pandas = ["pandas>=2.2.0"]
 numpy = ["numpy>=1.26.0"]
 duckdb = ["duckdb>=1.1.2"]
 ray = ["ray>=2.48"]
+# The datasets integration relies on private `datasets` APIs, so cap the next major version.
+hf = ["datasets>=4.0.0,<6", "huggingface-hub>=0.34.0"]
 
 [project.urls]
 Documentation = "https://docs.vortex.dev"
@@ -72,7 +74,9 @@ include = [
 
 [dependency-groups]
 dev = [
+    "datasets>=4.0.0,<6",
     "duckdb>=1.1.2",
+    "huggingface-hub>=0.34.0",
     "maturin>=1.10.2",
     "numpy>=2.2.2",
     "pandas-stubs>=2.2.3.241126",

vortex-python/python/vortex/__init__.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+import importlib
 import importlib.metadata
 import importlib.util
 
@@ -118,6 +119,17 @@ def cuda_extension_installed() -> bool:
     return importlib.util.find_spec("vortex_cuda") is not None
 
 
+def __getattr__(name: str):
+    # `datasets` is exposed lazily and deliberately kept out of __all__: importing it pulls in the
+    # optional `vortex-data[hf]` dependencies, so it must not be imported by `from vortex import *`
+    # or by merely importing `vortex`.
+    if name == "datasets":
+        module = importlib.import_module(".datasets", __name__)
+        globals()[name] = module
+        return module
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 __all__ = [
     # --- Modules ---
     "arrays",