feat: return dataset entity when creating a dataset (#768)

* return dataset id when upload

* return dataset entity instead

* fix test

* add back missing diff

Author: nankolena

docs/reference/dataset/index.md

@@ -13,5 +13,5 @@
         show_root_heading: false
 ::: kolena._api.v2.dataset
     options:
-        members: ["Filters", "GeneralFieldFilter"]
+        members: ["Filters", "GeneralFieldFilter", "DatasetEntity"]
         show_root_heading: false

kolena/_api/v2/dataset.py

@@ -90,11 +90,19 @@ class LoadDatasetByNameRequest:
 
 
 @dataclass(frozen=True)
-class EntityData:
+class DatasetEntity:
+    """
+    The descriptor of a dataset on Kolena.
+    """
+
     id: int
+    """ID of the dataset."""
     name: str
+    """Name of the dataset."""
     description: str
+    """Description of the dataset."""
     id_fields: List[str]
+    """ID fields of the dataset."""
 
 
 @dataclass(frozen=True)

kolena/_experimental/quality_standard.py

@@ -35,7 +35,7 @@
 from kolena._api.v2._testing import StratifyFieldSpec
 from kolena._api.v2._testing import TestingRequest
 from kolena._api.v2._testing import TestingResponse
-from kolena._api.v2.dataset import EntityData
+from kolena._api.v2.dataset import DatasetEntity
 from kolena._api.v2.model import ModelWithEvalConfig
 from kolena._api.v2.quality_standard import CopyQualityStandardRequest
 from kolena._api.v2.quality_standard import Path
@@ -150,7 +150,7 @@ def _download_quality_standard(
 
 def _calculate_moe_map(
     qs_result: pd.DataFrame,
-    dataset_entity: EntityData,
+    dataset_entity: DatasetEntity,
     confidence_level: float,
     qs: QualityStandardResponse,
 ) -> Dict[Tuple[str, Any], float]:

kolena/dataset/__init__.py

@@ -18,6 +18,7 @@
 from kolena.dataset.evaluation import download_results
 from kolena.dataset.evaluation import EvalConfigResults
 from kolena.dataset.dataset import list_datasets
+from kolena.dataset.dataset import DatasetEntity
 from kolena.dataset.evaluation import ModelEntity
 from kolena.dataset.evaluation import get_models
 from kolena.dataset.embeddings import upload_dataset_embeddings
@@ -33,6 +34,7 @@
     "download_results",
     "EvalConfigResults",
     "list_datasets",
+    "DatasetEntity",
     "ModelEntity",
     "get_models",
     "upload_dataset_embeddings",

kolena/dataset/dataset.py

@@ -30,7 +30,7 @@
 
 from kolena._api.v1.event import EventAPI
 from kolena._api.v2.dataset import CommitData
-from kolena._api.v2.dataset import EntityData
+from kolena._api.v2.dataset import DatasetEntity
 from kolena._api.v2.dataset import Filters
 from kolena._api.v2.dataset import ListCommitHistoryRequest
 from kolena._api.v2.dataset import ListCommitHistoryResponse
@@ -189,7 +189,7 @@ def _upload_dataset_chunk(df: pd.DataFrame, load_uuid: str, id_fields: List[str]
     upload_data_frame(df=df_serialized, load_uuid=load_uuid)
 
 
-def _load_dataset_metadata(name: str, raise_error_if_not_found: bool = True) -> Optional[EntityData]:
+def _load_dataset_metadata(name: str, raise_error_if_not_found: bool = True) -> Optional[DatasetEntity]:
     """
     Load the metadata of a given dataset.
 
@@ -210,13 +210,13 @@ def _load_dataset_metadata(name: str, raise_error_if_not_found: bool = True) ->
             return None
     response.raise_for_status()
 
-    return from_dict(EntityData, response.json())
+    return from_dict(DatasetEntity, response.json())
 
 
 def _resolve_id_fields(
     df: pd.DataFrame,
     id_fields: Optional[List[str]],
-    existing_dataset: Optional[EntityData],
+    existing_dataset: Optional[DatasetEntity],
 ) -> List[str]:
     existing_id_fields = []
     if existing_dataset:
@@ -269,7 +269,7 @@ def _send_upload_dataset_request(
     commit_tags: Optional[List[str]] = None,
     dataset_tags: Optional[List[str]] = None,
     description: Optional[str] = None,
-) -> EntityData:
+) -> DatasetEntity:
     request = RegisterRequest(
         name=name,
         id_fields=id_fields,
@@ -282,8 +282,8 @@ def _send_upload_dataset_request(
     )
     response = krequests.post(Path.REGISTER, json=asdict(request))
     krequests.raise_for_status(response)
-    data = from_dict(EntityData, response.json())
-    return data
+    dataset_entity = from_dict(DatasetEntity, response.json())
+    return dataset_entity
 
 
 def _upload_dataset(
@@ -296,10 +296,10 @@ def _upload_dataset(
     commit_tags: Optional[List[str]] = None,
     dataset_tags: Optional[List[str]] = None,
     description: Optional[str] = None,
-) -> None:
+) -> DatasetEntity:
     prepared_id_fields, load_uuid = _prepare_upload_dataset_request(name, df, id_fields=id_fields)
 
-    data = _send_upload_dataset_request(
+    dataset_entity = _send_upload_dataset_request(
         name,
         prepared_id_fields,
         load_uuid,
@@ -309,7 +309,8 @@ def _upload_dataset(
         dataset_tags=dataset_tags,
         description=description,
     )
-    log.info(f"uploaded dataset '{name}' ({get_dataset_url(dataset_id=data.id)})")
+    log.info(f"uploaded dataset '{name}' ({get_dataset_url(dataset_id=dataset_entity.id)})")
+    return dataset_entity
 
 
 @with_event(event_name=EventAPI.Event.REGISTER_DATASET)
@@ -322,7 +323,7 @@ def upload_dataset(
     dataset_tags: Optional[List[str]] = None,
     append_only: bool = False,
     description: Optional[str] = None,
-) -> None:
+) -> DatasetEntity:
     """
     Create or update a dataset with the contents of the provided DataFrame `df`.
 
@@ -343,8 +344,10 @@ def upload_dataset(
         datapoints from the input dataframe will be added, and existing datapoints will be modified if present in the
         input dataframe, but no datapoints will be deleted from the datasets. This behaves like an `UPSERT` operation.
     :param description: Optionally specify the description of the dataset.
+
+    :return: The dataset as a [`DatasetEntity`][kolena.dataset.DatasetEntity] object.
     """
-    _upload_dataset(
+    return _upload_dataset(
         name,
         df,
         id_fields=id_fields,

tests/integration/dataset/test_dataset.py

@@ -110,7 +110,9 @@ def test__upload_dataset() -> None:
     ]
     columns = ["locator", "width", "height", "city", "bboxes", "time_str", "time_num"]
 
-    upload_dataset(name, pd.DataFrame(datapoints[:10], columns=columns), id_fields=["locator"])
+    dataset_entity = upload_dataset(name, pd.DataFrame(datapoints[:10], columns=columns), id_fields=["locator"])
+    assert dataset_entity.id == _load_dataset_metadata(name).id
+    assert dataset_entity.name == name
 
     loaded_datapoints = download_dataset(name).sort_values("width", ignore_index=True).reindex(columns=columns)
     expected = pd.DataFrame(expected_datapoints[:10], columns=columns)

tests/unit/_experimental/trace/test_trace.py

@@ -15,7 +15,7 @@
 from unittest.mock import Mock
 from unittest.mock import patch
 
-from kolena._api.v2.dataset import EntityData
+from kolena._api.v2.dataset import DatasetEntity
 from kolena._experimental.trace import kolena_trace
 from kolena._experimental.trace.trace import _Trace
 
@@ -126,7 +126,7 @@ def predict(data, request_id):  # type: ignore
         assert str(e) == "Id Field request_id cannot be None in datapoint input"
 
     with patch("kolena._experimental.trace.trace._load_dataset_metadata") as mock_load_dataset_metadata:
-        mock_load_dataset_metadata.return_value = EntityData(
+        mock_load_dataset_metadata.return_value = DatasetEntity(
             id=1,
             name=dataset_name,
             description="test",

tests/unit/dataset/test_dataset.py

@@ -20,7 +20,7 @@
 import pytest
 from pandas.testing import assert_frame_equal
 
-from kolena._api.v2.dataset import EntityData
+from kolena._api.v2.dataset import DatasetEntity
 from kolena._utils.datatypes import DATA_TYPE_FIELD
 from kolena.dataset._common import COL_DATAPOINT
 from kolena.dataset._common import COL_RESULT
@@ -357,7 +357,7 @@ def test__infer_id_fields__error(input_df: pd.DataFrame) -> None:
 
 def test__resolve_id_fields() -> None:
     df = pd.DataFrame(dict(user_dp=["a", "b", "c"], new_user_dp=["d", "e", "f"]))
-    dataset = EntityData(id=1, name="foo", description="", id_fields=["user_dp"])
+    dataset = DatasetEntity(id=1, name="foo", description="", id_fields=["user_dp"])
     inferrable_df = pd.DataFrame(dict(locator=["x", "y", "z"]))
 
     # new dataset without id_fields
@@ -371,7 +371,7 @@ def test__resolve_id_fields() -> None:
     assert _resolve_id_fields(
         inferrable_df,
         None,
-        EntityData(id=1, name="foo", description="", id_fields=["locator"]),
+        DatasetEntity(id=1, name="foo", description="", id_fields=["locator"]),
     ) == ["locator"]
 
     # new dataset with explicit id_fields should resolve to explicit id_fields