Feature/mlflow dataset

rationai/mlkit/data/datasets/__init__.py

@@ -1,6 +1,6 @@
+from rationai.mlkit.data.datasets.load_dataset import load_dataset
 from rationai.mlkit.data.datasets.meta_tiled_slides import MetaTiledSlides
 from rationai.mlkit.data.datasets.openslide_tiles_dataset import OpenSlideTilesDataset
-from rationai.mlkit.data.datasets.slides_tiles_loader import SlidesTilesLoader
 
 
-__all__ = ["MetaTiledSlides", "OpenSlideTilesDataset", "SlidesTilesLoader"]
+__all__ = ["MetaTiledSlides", "OpenSlideTilesDataset", "load_dataset"]

rationai/mlkit/data/datasets/load_dataset.py

@@ -0,0 +1,130 @@
+from collections.abc import Iterable
+from concurrent.futures import ThreadPoolExecutor
+from itertools import chain
+from pathlib import Path
+from typing import Any
+
+from datasets import Dataset as HFDataset
+from datasets import DatasetDict as HFDatasetDict
+from datasets import concatenate_datasets
+from datasets import load_dataset as hf_load_dataset
+from mlflow.artifacts import download_artifacts
+
+
+def load_dataset[T: (HFDataset, HFDatasetDict)](
+    *,
+    paths: Iterable[Path | str] | None = None,
+    uris: Iterable[str] | None = None,
+    dataset: T | None = None,
+    entities: Iterable[str] | None = None,
+    hf_kwargs: dict[str, Any] | None = None,
+) -> T:
+    """Load parquet-backed HF datasets from local paths and/or MLflow artifact URIs.
+
+    The result is always a ``DatasetDict`` keyed by split name. If the source
+    data has no split structure, everything falls back to ``"train"``.
+    Multiple sources are concatenated split-by-split.
+
+    Args:
+        paths: Local filesystem paths to search for parquet files.
+        uris: MLflow artifact URIs to download and load.
+        dataset: An already-loaded ``DatasetDict`` to merge with the above.
+        entities: If given, only these entity names are loaded. For URIs this
+            also limits what gets downloaded. ``None`` loads all entities.
+        hf_kwargs: Extra kwargs forwarded to ``datasets.load_dataset``.
+            Defaults to ``{"path": "parquet"}``. Any ``"split"`` key is
+            ignored — use the ``splits`` parameter instead.
+    """
+    if not (paths or uris or dataset):
+        raise ValueError("At least one of paths, uris or dataset must be provided.")
+
+    if hf_kwargs is None:
+        hf_kwargs = {"path": "parquet"}
+
+    _entities = frozenset(entities) if entities is not None else None
+    datasets: list[HFDatasetDict] = [dataset] if dataset is not None else []
+
+    if paths or uris:
+        datasets.extend(_load_sources(paths or [], uris or [], hf_kwargs, _entities))
+
+    types = {type(ds) for ds in datasets}
+    if len(types) > 1:
+        raise TypeError(
+            f"Cannot merge datasets of mixed types: {', '.join(t.__name__ for t in types)}"
+        )
+
+    if all(isinstance(ds, HFDataset) for ds in datasets):
+        return concatenate_datasets(datasets)
+
+    return _merge(datasets)
+
+
+def _merge(datasets: Iterable[HFDatasetDict]) -> HFDatasetDict:
+    """Merge a sequence of DatasetDicts split-by-split.
+
+    Splits present in only a subset of the dicts are still included —
+    they are concatenated from whichever dicts contain them.
+    """
+    datasets = list(datasets)
+    all_splits = set().union(*[dd.keys() for dd in datasets])
+    return HFDatasetDict({
+        split: concatenate_datasets([dd[split] for dd in datasets if split in dd])
+        for split in all_splits
+    })
+
+
+def _download_uri(uri: str, entities: frozenset[str] | None) -> list[Path]:
+    """Download an MLflow artifact URI, returning local paths.
+
+    When ``entities`` is specified only the per-entity subdirectories are
+    downloaded (``{uri}/{entity}``), avoiding pulling unused data.
+    Returns a list because a single URI may expand to multiple local paths
+    (one per entity).
+    """
+    if entities is None:
+        return [Path(download_artifacts(artifact_uri=uri))]
+
+    with ThreadPoolExecutor() as executor:
+        return [
+            Path(p)
+            for p in executor.map(
+                lambda entity: download_artifacts(artifact_uri=str(Path(uri, entity))),
+                entities,
+            )
+        ]
+
+
+def _load_sources(
+    paths: Iterable[str | Path],
+    uris: Iterable[str],
+    hf_kwargs: dict[str, Any],
+    entities: frozenset[str] | None,
+) -> list[HFDataset | HFDatasetDict]:
+    """Load parquet files from local paths and MLflow URIs into DatasetDicts.
+
+    Each source is loaded independently.
+    """
+    with ThreadPoolExecutor() as executor:
+        artifacts_paths = list(
+            chain.from_iterable(
+                executor.map(lambda uri: _download_uri(uri, entities), uris)
+            )
+        )
+
+    resolved = [Path(p) for p in (*paths, *artifacts_paths)]
+
+    if not resolved:
+        return []
+
+    try:
+        return [
+            hf_load_dataset(
+                **hf_kwargs,
+                **({"data_dir": str(p)} if p.is_dir() else {"data_files": str(p)}),
+            )
+            for p in resolved
+        ]
+
+    except Exception as e:
+        msg = "Failed to load Parquet files."
+        raise RuntimeError(msg) from e

rationai/mlkit/data/datasets/meta_tiled_slides.py

@@ -1,13 +1,12 @@
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
-from pathlib import Path
-from typing import Any, TypeVar
+from typing import TypeVar
 
+import numpy as np
+import pyarrow as pa
 from datasets import Dataset as HFDataset
 from torch.utils.data import ConcatDataset, Dataset
 
-from rationai.mlkit.data.datasets.slides_tiles_loader import SlidesTilesLoader
-
 
 T = TypeVar("T", covariant=True)
 
@@ -25,52 +24,41 @@ class MetaTiledSlides(ConcatDataset[T], ABC):
 
     def __init__(
         self,
-        *,
-        paths: Iterable[Path | str] | None = None,
-        uris: Iterable[str] | None = None,
-        slides_and_tiles: tuple[HFDataset, HFDataset] | None = None,
-        hf_kwargs: dict[str, Any] | None = None,
+        slides: HFDataset,
+        tiles: HFDataset
     ) -> None:
         """Load slides and tiles from MLFlow artifacts.
 
         Args:
-            paths: List of directories to load slides and tiles from. Each
-                directory must include either single files (`slides.parquet`
-                and `tiles.parquet`) or subdirectories (`slides/` and `tiles/`)
-                containing chunked Parquet files.
-            uris: List of MLFlow artifact URIs pointing to folders containing
-                either single files (`slides.parquet` and `tiles.parquet`) or
-                subdirectories (`slides/` and `tiles/`) containing chunked
-                Parquet files.
-            slides_and_tiles: Tuple containing the slides and tiles Datasets.
-            hf_kwargs: Additional keyword arguments to pass to HuggingFace's
-                `load_dataset` function. Defaults to `{"path": "parquet", "split": "train"}`.
+            slides: Dataset containing slide metadata.
+            tiles: Dataset containing tile metadata.
         """
-        self._meta = SlidesTilesLoader(
-            paths=paths,
-            uris=uris,
-            slides_and_tiles=slides_and_tiles,
-            hf_kwargs=hf_kwargs,
-        )
-        self.slides = self._meta.slides
-        self.tiles = self._meta.tiles
-        super().__init__(self.generate_datasets())
+        self.slides = slides
+        self.tiles = tiles
 
-    def filter_tiles_by_slide(self, slide_id: str | bytes) -> HFDataset:
-        """Returns a view of the dataset using a slice or indices.
+        self._slide_id_to_indices = self._build_tile_index(self.tiles)
 
-        This function creates a view of the `self.tiles` dataset that contains only
-        the tiles belonging to the specified slide. It uses the precomputed
-        `_slide_id_to_indices` mapping to efficiently retrieve the relevant tiles
-        without copying data.
+        super().__init__(self.generate_datasets())
 
-        Args:
-            slide_id: The ID of the slide to filter tiles.
+    def filter_tiles_by_slide(self, slide_id: str | bytes) -> HFDataset:
+            """Returns a view of the dataset using a slice or indices.
+
+            This function creates a view of the `self.tiles` dataset that contains only
+            the tiles belonging to the specified slide. It uses the precomputed
+            `_slide_id_to_indices` mapping to efficiently retrieve the relevant tiles
+            without copying data.
+
+            Args:
+                slide_id: The ID of the slide to filter tiles.
+
+            Returns:
+                A view of the tiles dataset containing only the tiles for the specified slide.
+            """
+            tile_indices = self._slide_id_to_indices.get(
+                slide_id, pa.scalar([], type=pa.list_(pa.int64()))
+            )
+            return self.tiles.select(tile_indices.values.to_numpy())
 
-        Returns:
-            A view of the tiles dataset containing only the tiles for the specified slide.
-        """
-        return self._meta.filter_tiles_by_slide(slide_id)
 
     @abstractmethod
     def generate_datasets(self) -> Iterable[Dataset[T]]:
@@ -90,3 +78,43 @@ def generate_datasets(self) -> Iterable[Dataset[T]]:
             )
             ```
         """
+
+    @staticmethod
+    def _build_tile_index(tiles: HFDataset) -> dict[str | bytes, pa.ListScalar]:
+        """Creates a fast lookup table for slide indices.
+
+        This function builds a mapping from `slide_id` to the list of indices in the
+        `tiles` dataset that correspond to that slide.
+
+        Args:
+            tiles: A dataset containing a `slide_id` column.
+
+        Returns:
+            A dictionary mapping each `slide_id` to a list of indices in the `tiles` dataset.
+        """
+        if len(tiles) == 0:
+            return {}
+
+        slide_ids = tiles.data.column("slide_id")
+        num_rows = len(slide_ids)
+
+        # group_by requires the "large" variants for string/binary columns
+        current_type = slide_ids.type
+        if pa.types.is_string(current_type):
+            slide_ids = slide_ids.cast(pa.large_string())
+        elif pa.types.is_binary(current_type):
+            slide_ids = slide_ids.cast(pa.large_binary())
+
+        # np.arange is used here because PyArrow can wrap it with zero-copy overhead
+        row_indices = pa.array(np.arange(num_rows, dtype=np.int64))
+        table = pa.Table.from_arrays(
+            [slide_ids, row_indices], names=["slide_id", "idx"]
+        )
+
+        # "list" aggregates all indices for a given slide_id into a single Arrow List scalar
+        grouped = table.group_by("slide_id").aggregate([("idx", "list")])
+
+        # Keep values as PyArrow ListScalars to avoid materializing them in Python
+        keys = grouped.column("slide_id").to_numpy()
+        values_array = grouped.column("idx_list")
+        return {key: values_array[i] for i, key in enumerate(keys)}