XL Datasets: Minimal Zarr-only dataset implementation

docs/api/dataset.md

@@ -2,6 +2,12 @@
     options:
         filters: ["!^_"]
 
+--- 
+
+::: polaris.dataset._base.BaseDataset
+    options:
+        filters: ["!^_"]
+
 ---
 
 ::: polaris.dataset.ColumnAnnotation

polaris/benchmark/_base.py

@@ -1,6 +1,6 @@
-from itertools import chain
 import json
 from hashlib import md5
+from itertools import chain
 from typing import Any, Callable, Optional, Union
 
 import fsspec
@@ -18,11 +18,11 @@
 from sklearn.utils.multiclass import type_of_target
 
 from polaris._artifact import BaseArtifactModel
-from polaris.mixins import ChecksumMixin
-from polaris.dataset import Dataset, Subset, CompetitionDataset
+from polaris.dataset import CompetitionDataset, DatasetV1, Subset
 from polaris.evaluate import BenchmarkResults, Metric
 from polaris.evaluate.utils import evaluate_benchmark
 from polaris.hub.settings import PolarisHubSettings
+from polaris.mixins import ChecksumMixin
 from polaris.utils.dict2html import dict2html
 from polaris.utils.errors import InvalidBenchmarkError
 from polaris.utils.misc import listit
@@ -96,7 +96,7 @@ class BenchmarkSpecification(BaseArtifactModel, ChecksumMixin):
 
     # Public attributes
     # Data
-    dataset: Union[Dataset, CompetitionDataset, str, dict[str, Any]]
+    dataset: Union[DatasetV1, CompetitionDataset, str, dict[str, Any]]
     target_cols: ColumnsType
     input_cols: ColumnsType
     split: SplitType
@@ -111,12 +111,11 @@ class BenchmarkSpecification(BaseArtifactModel, ChecksumMixin):
     def _validate_dataset(cls, v):
         """
         Allows either passing a Dataset object or the kwargs to create one
-        TODO (cwognum): Allow multiple datasets to be used as part of a benchmark
         """
         if isinstance(v, dict):
-            v = Dataset(**v)
+            v = DatasetV1(**v)
         elif isinstance(v, str):
-            v = Dataset.from_json(v)
+            v = DatasetV1.from_json(v)
         return v
 
     @field_validator("target_cols", "input_cols")
@@ -345,6 +344,7 @@ def n_classes(self) -> dict[str, int]:
             target_type = self.target_types[target]
             if target_type is None or target_type == TargetType.REGRESSION:
                 continue
+            # TODO: Don't use table attribute
             n_classes[target] = self.dataset.table.loc[:, target].nunique()
         return n_classes
 

polaris/dataset/__init__.py

@@ -1,8 +1,9 @@
-from polaris.dataset._column import ColumnAnnotation, Modality, KnownContentType
-from polaris.dataset._dataset import Dataset
+from polaris.dataset._column import ColumnAnnotation, KnownContentType, Modality
+from polaris.dataset._competition_dataset import CompetitionDataset
+from polaris.dataset._dataset import DatasetV1
+from polaris.dataset._dataset import DatasetV1 as Dataset
 from polaris.dataset._factory import DatasetFactory, create_dataset_from_file, create_dataset_from_files
 from polaris.dataset._subset import Subset
-from polaris.dataset._competition_dataset import CompetitionDataset
 
 __all__ = [
     "ColumnAnnotation",
@@ -14,4 +15,5 @@
     "DatasetFactory",
     "create_dataset_from_file",
     "create_dataset_from_files",
+    "DatasetV1",
 ]