datachain 0.15.0__tar.gz → 0.16.0__tar.gz

{datachain-0.15.0/src/datachain.egg-info → datachain-0.16.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.15.0
+Version: 0.16.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.15.0 → datachain-0.16.0}/src/datachain/catalog/catalog.py

@@ -776,7 +776,7 @@ class Catalog:
         listing: Optional[bool] = False,
         uuid: Optional[str] = None,
         description: Optional[str] = None,
-        labels: Optional[list[str]] = None,
+        attrs: Optional[list[str]] = None,
     ) -> "DatasetRecord":
         """
         Creates new dataset of a specific version.
@@ -794,16 +794,16 @@ class Catalog:
             dataset = self.get_dataset(name)
             default_version = dataset.next_version
 
-            if (description or labels) and (
-                dataset.description != description or dataset.labels != labels
+            if (description or attrs) and (
+                dataset.description != description or dataset.attrs != attrs
             ):
                 description = description or dataset.description
-                labels = labels or dataset.labels
+                attrs = attrs or dataset.attrs
 
                 self.update_dataset(
                     dataset,
                     description=description,
-                    labels=labels,
+                    attrs=attrs,
                 )
 
         except DatasetNotFoundError:
@@ -817,7 +817,7 @@ class Catalog:
                 schema=schema,
                 ignore_if_exists=True,
                 description=description,
-                labels=labels,
+                attrs=attrs,
             )
 
         version = version or default_version
@@ -1334,15 +1334,15 @@ class Catalog:
         name: str,
         new_name: Optional[str] = None,
         description: Optional[str] = None,
-        labels: Optional[list[str]] = None,
+        attrs: Optional[list[str]] = None,
     ) -> DatasetRecord:
         update_data = {}
         if new_name:
             update_data["name"] = new_name
         if description is not None:
             update_data["description"] = description
-        if labels is not None:
-            update_data["labels"] = labels  # type: ignore[assignment]
+        if attrs is not None:
+            update_data["attrs"] = attrs  # type: ignore[assignment]
 
         dataset = self.get_dataset(name)
         return self.update_dataset(dataset, **update_data)

{datachain-0.15.0 → datachain-0.16.0}/src/datachain/cli/__init__.py

@@ -149,7 +149,7 @@ def handle_dataset_command(args, catalog):
             args.name,
             new_name=args.new_name,
             description=args.description,
-            labels=args.labels,
+            attrs=args.attrs,
             studio=args.studio,
             local=args.local,
             all=args.all,

{datachain-0.15.0 → datachain-0.16.0}/src/datachain/cli/commands/datasets.py

@@ -154,7 +154,7 @@ def edit_dataset(
     name: str,
     new_name: Optional[str] = None,
     description: Optional[str] = None,
-    labels: Optional[list[str]] = None,
+    attrs: Optional[list[str]] = None,
     studio: bool = False,
     local: bool = False,
     all: bool = True,
@@ -167,9 +167,9 @@ def edit_dataset(
 
     if all or local:
         try:
-            catalog.edit_dataset(name, new_name, description, labels)
+            catalog.edit_dataset(name, new_name, description, attrs)
         except DatasetNotFoundError:
             print("Dataset not found in local", file=sys.stderr)
 
     if (all or studio) and token:
-        edit_studio_dataset(team, name, new_name, description, labels)
+        edit_studio_dataset(team, name, new_name, description, attrs)

{datachain-0.15.0 → datachain-0.16.0}/src/datachain/cli/commands/show.py

@@ -42,8 +42,8 @@ def show(
     print("Name: ", name)
     if dataset.description:
         print("Description: ", dataset.description)
-    if dataset.labels:
-        print("Labels: ", ",".join(dataset.labels))
+    if dataset.attrs:
+        print("Attributes: ", ",".join(dataset.attrs))
     print("\n")
 
     show_records(records, collapse_columns=not no_collapse, hidden_fields=hidden_fields)

{datachain-0.15.0 → datachain-0.16.0}/src/datachain/cli/parser/__init__.py

@@ -217,9 +217,9 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help="Dataset description",
     )
     parse_edit_dataset.add_argument(
-        "--labels",
+        "--attrs",
         nargs="+",
-        help="Dataset labels",
+        help="Dataset attributes",
     )
     parse_edit_dataset.add_argument(
         "--studio",

{datachain-0.15.0 → datachain-0.16.0}/src/datachain/data_storage/metastore.py

@@ -120,7 +120,7 @@ class AbstractMetastore(ABC, Serializable):
         schema: Optional[dict[str, Any]] = None,
         ignore_if_exists: bool = False,
         description: Optional[str] = None,
-        labels: Optional[list[str]] = None,
+        attrs: Optional[list[str]] = None,
     ) -> DatasetRecord:
         """Creates new dataset."""
 
@@ -326,7 +326,7 @@ class AbstractDBMetastore(AbstractMetastore):
             Column("id", Integer, primary_key=True),
             Column("name", Text, nullable=False),
             Column("description", Text),
-            Column("labels", JSON, nullable=True),
+            Column("attrs", JSON, nullable=True),
             Column("status", Integer, nullable=False),
             Column("feature_schema", JSON, nullable=True),
             Column("created_at", DateTime(timezone=True)),
@@ -521,7 +521,7 @@ class AbstractDBMetastore(AbstractMetastore):
         schema: Optional[dict[str, Any]] = None,
         ignore_if_exists: bool = False,
         description: Optional[str] = None,
-        labels: Optional[list[str]] = None,
+        attrs: Optional[list[str]] = None,
         **kwargs,  # TODO registered = True / False
     ) -> DatasetRecord:
         """Creates new dataset."""
@@ -538,7 +538,7 @@ class AbstractDBMetastore(AbstractMetastore):
             query_script=query_script,
             schema=json.dumps(schema or {}),
             description=description,
-            labels=json.dumps(labels or []),
+            attrs=json.dumps(attrs or []),
         )
         if ignore_if_exists and hasattr(query, "on_conflict_do_nothing"):
             # SQLite and PostgreSQL both support 'on_conflict_do_nothing',
@@ -621,7 +621,7 @@ class AbstractDBMetastore(AbstractMetastore):
         dataset_values = {}
         for field, value in kwargs.items():
             if field in self._dataset_fields[1:]:
-                if field in ["labels", "schema"]:
+                if field in ["attrs", "schema"]:
                     values[field] = json.dumps(value) if value else None
                 else:
                     values[field] = value

{datachain-0.15.0 → datachain-0.16.0}/src/datachain/dataset.py

@@ -329,7 +329,7 @@ class DatasetRecord:
     id: int
     name: str
     description: Optional[str]
-    labels: list[str]
+    attrs: list[str]
     schema: dict[str, Union[SQLType, type[SQLType]]]
     feature_schema: dict
     versions: list[DatasetVersion]
@@ -357,7 +357,7 @@ class DatasetRecord:
         id: int,
         name: str,
         description: Optional[str],
-        labels: str,
+        attrs: str,
         status: int,
         feature_schema: Optional[str],
         created_at: datetime,
@@ -387,7 +387,7 @@ class DatasetRecord:
         version_schema: str,
         version_job_id: Optional[str] = None,
     ) -> "DatasetRecord":
-        labels_lst: list[str] = json.loads(labels) if labels else []
+        attrs_lst: list[str] = json.loads(attrs) if attrs else []
         schema_dct: dict[str, Any] = json.loads(schema) if schema else {}
         version_schema_dct: dict[str, str] = (
             json.loads(version_schema) if version_schema else {}
@@ -418,7 +418,7 @@ class DatasetRecord:
             id,
             name,
             description,
-            labels_lst,
+            attrs_lst,
             cls.parse_schema(schema_dct),  # type: ignore[arg-type]
             json.loads(feature_schema) if feature_schema else {},
             [dataset_version],
@@ -562,7 +562,7 @@ class DatasetListRecord:
     id: int
     name: str
     description: Optional[str]
-    labels: list[str]
+    attrs: list[str]
     versions: list[DatasetListVersion]
     created_at: Optional[datetime] = None
 
@@ -572,7 +572,7 @@ class DatasetListRecord:
         id: int,
         name: str,
         description: Optional[str],
-        labels: str,
+        attrs: str,
         created_at: datetime,
         version_id: int,
         version_uuid: str,
@@ -588,7 +588,7 @@ class DatasetListRecord:
         version_query_script: Optional[str],
         version_job_id: Optional[str] = None,
     ) -> "DatasetListRecord":
-        labels_lst: list[str] = json.loads(labels) if labels else []
+        attrs_lst: list[str] = json.loads(attrs) if attrs else []
 
         dataset_version = DatasetListVersion.parse(
             version_id,
@@ -610,7 +610,7 @@ class DatasetListRecord:
             id,
             name,
             description,
-            labels_lst,
+            attrs_lst,
             [dataset_version],
             created_at,
         )

{datachain-0.15.0 → datachain-0.16.0}/src/datachain/lib/dataset_info.py

@@ -32,11 +32,28 @@ class DatasetInfo(DataModel):
     metrics: dict[str, Any] = Field(default={})
     error_message: str = Field(default="")
     error_stack: str = Field(default="")
+    attrs: list[str] = Field(default=[])
 
     @property
     def is_temp(self) -> bool:
         return Session.is_temp_dataset(self.name)
 
+    def has_attr(self, attr: str) -> bool:
+        s = attr.split("=")
+        if len(s) == 1:
+            return attr in self.attrs
+
+        name = s[0]
+        value = s[1]
+        for a in self.attrs:
+            s = a.split("=")
+            if value == "*" and s[0] == name:
+                return True
+            if len(s) == 2 and s[0] == name and s[1] == value:
+                return True
+
+        return False
+
     @staticmethod
     def _validate_dict(
         v: Optional[Union[str, dict]],
@@ -83,4 +100,5 @@ class DatasetInfo(DataModel):
             metrics=job.metrics if job else {},
             error_message=version.error_message,
             error_stack=version.error_stack,
+            attrs=dataset.attrs,
         )

{datachain-0.15.0 → datachain-0.16.0}/src/datachain/lib/dc/datachain.py

@@ -459,7 +459,7 @@ class DataChain:
         name: str,
         version: Optional[int] = None,
         description: Optional[str] = None,
-        labels: Optional[list[str]] = None,
+        attrs: Optional[list[str]] = None,
         **kwargs,
     ) -> "Self":
         """Save to a Dataset. It returns the chain itself.
@@ -468,7 +468,8 @@ class DataChain:
             name : dataset name.
             version : version of a dataset. Default - the last version that exist.
             description : description of a dataset.
-            labels : labels of a dataset.
+            attrs : attributes of a dataset. They can be without value, e.g "NLP",
+                or with a value, e.g "location=US".
         """
         schema = self.signals_schema.clone_without_sys_signals().serialize()
         return self._evolve(
@@ -476,7 +477,7 @@ class DataChain:
                 name=name,
                 version=version,
                 description=description,
-                labels=labels,
+                attrs=attrs,
                 feature_schema=schema,
                 **kwargs,
             )

{datachain-0.15.0 → datachain-0.16.0}/src/datachain/lib/dc/datasets.py

@@ -102,6 +102,7 @@ def datasets(
     column: Optional[str] = None,
     include_listing: bool = False,
     studio: bool = False,
+    attrs: Optional[list[str]] = None,
 ) -> "DataChain":
     """Generate chain with list of registered datasets.
 
@@ -114,6 +115,10 @@ def datasets(
         include_listing: If True, includes listing datasets. Defaults to False.
         studio: If True, returns datasets from Studio only,
             otherwise returns all local datasets. Defaults to False.
+        attrs: Optional list of attributes to filter datasets on. It can be just
+            attribute without value e.g "NLP", or attribute with value
+            e.g "location=US". Attribute with value can also accept "*" to target
+            all that have specific name e.g "location=*"
 
     Returns:
         DataChain: A new DataChain instance containing dataset information.
@@ -139,6 +144,10 @@ def datasets(
     ]
     datasets_values = [d for d in datasets_values if not d.is_temp]
 
+    if attrs:
+        for attr in attrs:
+            datasets_values = [d for d in datasets_values if d.has_attr(attr)]
+
     if not column:
         # flattening dataset fields
         schema = {

{datachain-0.15.0 → datachain-0.16.0}/src/datachain/lib/udf.py

@@ -474,8 +474,9 @@ class Generator(UDFBase):
             remove_prefetched=bool(self.prefetch) and not cache,
         )
         with closing(prepared_inputs):
-            for row in processed_cb.wrap(prepared_inputs):
+            for row in prepared_inputs:
                 yield _process_row(row)
+                processed_cb.relative_update(1)
 
         self.teardown()
 

{datachain-0.15.0 → datachain-0.16.0}/src/datachain/query/dataset.py

@@ -1680,7 +1680,7 @@ class DatasetQuery:
         version: Optional[int] = None,
         feature_schema: Optional[dict] = None,
         description: Optional[str] = None,
-        labels: Optional[list[str]] = None,
+        attrs: Optional[list[str]] = None,
         **kwargs,
     ) -> "Self":
         """Save the query as a dataset."""
@@ -1714,7 +1714,7 @@ class DatasetQuery:
                 feature_schema=feature_schema,
                 columns=columns,
                 description=description,
-                labels=labels,
+                attrs=attrs,
                 **kwargs,
             )
             version = version or dataset.latest_version

{datachain-0.15.0 → datachain-0.16.0}/src/datachain/remote/studio.py

@@ -290,13 +290,13 @@ class StudioClient:
         name: str,
         new_name: Optional[str] = None,
         description: Optional[str] = None,
-        labels: Optional[list[str]] = None,
+        attrs: Optional[list[str]] = None,
     ) -> Response[DatasetInfoData]:
         body = {
             "new_name": new_name,
             "dataset_name": name,
             "description": description,
-            "labels": labels,
+            "attrs": attrs,
         }
 
         return self._send_request(

{datachain-0.15.0 → datachain-0.16.0}/src/datachain/studio.py

@@ -187,10 +187,10 @@ def edit_studio_dataset(
     name: str,
     new_name: Optional[str] = None,
     description: Optional[str] = None,
-    labels: Optional[list[str]] = None,
+    attrs: Optional[list[str]] = None,
 ):
     client = StudioClient(team=team_name)
-    response = client.edit_dataset(name, new_name, description, labels)
+    response = client.edit_dataset(name, new_name, description, attrs)
     if not response.ok:
         raise DataChainError(response.message)
 

{datachain-0.15.0 → datachain-0.16.0/src/datachain.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.15.0
+Version: 0.16.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.15.0 → datachain-0.16.0}/tests/conftest.py

@@ -550,7 +550,7 @@ def animal_dataset(listed_bucket, cloud_test_catalog):
     src_uri = cloud_test_catalog.src_uri
     dataset = catalog.create_dataset_from_sources(name, [src_uri], recursive=True)
     return catalog.update_dataset(
-        dataset, {"description": "animal dataset", "labels": ["cats", "dogs"]}
+        dataset, {"description": "animal dataset", "attrs": ["cats", "dogs"]}
     )
 
 
@@ -563,7 +563,7 @@ def dogs_dataset(listed_bucket, cloud_test_catalog):
         name, [f"{src_uri}/dogs/*"], recursive=True
     )
     return catalog.update_dataset(
-        dataset, {"description": "dogs dataset", "labels": ["dogs", "dataset"]}
+        dataset, {"description": "dogs dataset", "attrs": ["dogs", "dataset"]}
     )
 
 
@@ -576,7 +576,7 @@ def cats_dataset(listed_bucket, cloud_test_catalog):
         name, [f"{src_uri}/cats/*"], recursive=True
     )
     return catalog.update_dataset(
-        dataset, {"description": "cats dataset", "labels": ["cats", "dataset"]}
+        dataset, {"description": "cats dataset", "attrs": ["cats", "dataset"]}
     )
 
 
@@ -586,7 +586,7 @@ def dataset_record():
         id=1,
         name=f"ds_{uuid.uuid4().hex}",
         description="",
-        labels=[],
+        attrs=[],
         versions=[],
         status=1,
         schema={},
@@ -651,7 +651,7 @@ def studio_datasets(requests_mock, studio_token):
         "id": 1,
         "name": "dogs",
         "description": "dogs dataset",
-        "labels": ["dogs", "dataset"],
+        "attrs": ["dogs", "dataset"],
         "versions": [
             {
                 "version": 1,
@@ -676,7 +676,7 @@ def studio_datasets(requests_mock, studio_token):
             "id": 2,
             "name": "cats",
             "description": "cats dataset",
-            "labels": ["cats", "dataset"],
+            "attrs": ["cats", "dataset"],
             "versions": [
                 {
                     "version": 1,
@@ -691,7 +691,7 @@ def studio_datasets(requests_mock, studio_token):
             "id": 3,
             "name": "both",
             "description": "both dataset",
-            "labels": ["both", "dataset"],
+            "attrs": ["both", "dataset"],
             "versions": [
                 {
                     "version": 1,

{datachain-0.15.0 → datachain-0.16.0}/tests/func/test_datachain.py

@@ -560,23 +560,23 @@ def test_save(test_session):
         name="new_name",
         version=1,
         description="new description",
-        labels=["new_label", "old_label"],
+        attrs=["new_label", "old_label"],
     )
 
     ds = test_session.catalog.get_dataset("new_name")
     assert ds.name == "new_name"
     assert ds.description == "new description"
-    assert ds.labels == ["new_label", "old_label"]
+    assert ds.attrs == ["new_label", "old_label"]
 
     chain.save(
         name="new_name",
         description="updated description",
-        labels=["new_label", "old_label", "new_label2"],
+        attrs=["new_label", "old_label", "new_label2"],
     )
     ds = test_session.catalog.get_dataset("new_name")
     assert ds.name == "new_name"
     assert ds.description == "updated description"
-    assert ds.labels == ["new_label", "old_label", "new_label2"]
+    assert ds.attrs == ["new_label", "old_label", "new_label2"]
 
 
 def test_show_nested_empty(capsys, test_session):

{datachain-0.15.0 → datachain-0.16.0}/tests/func/test_datasets.py

@@ -170,7 +170,7 @@ def test_create_dataset_from_sources(listed_bucket, cloud_test_catalog):
     assert dataset.name == dataset_name
     assert dataset.description is None
     assert dataset.versions_values == [1]
-    assert dataset.labels == []
+    assert dataset.attrs == []
     assert dataset.status == DatasetStatus.COMPLETE
 
     assert dataset_version.status == DatasetStatus.COMPLETE
@@ -207,7 +207,7 @@ def test_create_dataset_from_sources_dataset(cloud_test_catalog, dogs_dataset):
     assert dataset.name == dataset_name
     assert dataset.description is None
     assert dataset.versions_values == [1]
-    assert dataset.labels == []
+    assert dataset.attrs == []
     assert dataset.status == DatasetStatus.COMPLETE
 
     assert dataset_version.status == DatasetStatus.COMPLETE
@@ -546,14 +546,14 @@ def test_edit_dataset(cloud_test_catalog, dogs_dataset):
         dogs_dataset.name,
         new_name=dataset_new_name,
         description="new description",
-        labels=["cats", "birds"],
+        attrs=["cats", "birds"],
     )
 
     dataset = catalog.get_dataset(dataset_new_name)
     assert dataset.versions_values == [1]
     assert dataset.name == dataset_new_name
     assert dataset.description == "new description"
-    assert dataset.labels == ["cats", "birds"]
+    assert dataset.attrs == ["cats", "birds"]
 
     # check if dataset tables are renamed correctly
     old_dataset_table_name = catalog.warehouse.dataset_table_name(dataset_old_name, 1)
@@ -589,7 +589,7 @@ def test_edit_dataset_same_name(cloud_test_catalog, dogs_dataset):
     )
 
 
-def test_edit_dataset_remove_labels_and_description(cloud_test_catalog, dogs_dataset):
+def test_edit_dataset_remove_attrs_and_description(cloud_test_catalog, dogs_dataset):
     dataset_new_name = uuid.uuid4().hex
     catalog = cloud_test_catalog.catalog
 
@@ -597,14 +597,14 @@ def test_edit_dataset_remove_labels_and_description(cloud_test_catalog, dogs_dat
         dogs_dataset.name,
         new_name=dataset_new_name,
         description="",
-        labels=[],
+        attrs=[],
     )
 
     dataset = catalog.get_dataset(dataset_new_name)
     assert dataset.versions_values == [1]
     assert dataset.name == dataset_new_name
     assert dataset.description == ""
-    assert dataset.labels == []
+    assert dataset.attrs == []
 
 
 def test_ls_dataset_rows(cloud_test_catalog, dogs_dataset):

{datachain-0.15.0 → datachain-0.16.0}/tests/func/test_pull.py

@@ -118,7 +118,7 @@ def remote_dataset(remote_dataset_version, schema):
         "id": 1,
         "name": "dogs",
         "description": "",
-        "labels": [],
+        "attrs": [],
         "schema": schema,
         "status": 4,
         "feature_schema": {},

{datachain-0.15.0 → datachain-0.16.0}/tests/test_cli_studio.py

@@ -247,7 +247,7 @@ def test_studio_edit_dataset(capsys, mocker):
             "new_name": "new-name",
             "team_name": "team_name",
             "description": None,
-            "labels": None,
+            "attrs": None,
         }
 
         # With all arguments
@@ -261,8 +261,8 @@ def test_studio_edit_dataset(capsys, mocker):
                     "new-name",
                     "--description",
                     "description",
-                    "--labels",
-                    "label1",
+                    "--attrs",
+                    "attr1",
                     "--team",
                     "team_name",
                     "--studio",
@@ -275,7 +275,7 @@ def test_studio_edit_dataset(capsys, mocker):
             "dataset_name": "name",
             "new_name": "new-name",
             "description": "description",
-            "labels": ["label1"],
+            "attrs": ["attr1"],
             "team_name": "team_name",
         }
 

{datachain-0.15.0 → datachain-0.16.0}/tests/unit/lib/test_datachain.py

@@ -356,6 +356,41 @@ def test_datasets_in_memory():
     assert datasets[0].num_objects == 6
 
 
+@pytest.mark.parametrize(
+    "attrs,result",
+    [
+        (["number"], ["evens", "primes"]),
+        (["num=prime"], ["primes"]),
+        (["num=even"], ["evens"]),
+        (["num=*"], ["evens", "primes"]),
+        (["num=*", "small"], ["primes"]),
+        (["letter"], ["letters"]),
+        (["missing"], []),
+        (["num=*", "missing"], []),
+        (None, ["evens", "letters", "primes"]),
+        ([], ["evens", "letters", "primes"]),
+    ],
+)
+def test_datasets_filtering(test_session, attrs, result):
+    ds = dc.datasets(column="dataset", session=test_session)
+    datasets = [d for d in ds.collect("dataset") if d.name == "fibonacci"]
+    assert len(datasets) == 0
+
+    dc.read_values(num=[1, 2, 3], session=test_session).save(
+        "primes", attrs=["number", "num=prime", "small"]
+    )
+
+    dc.read_values(num=[2, 4, 6], session=test_session).save(
+        "evens", attrs=["number", "num=even"]
+    )
+
+    dc.read_values(letter=["a", "b", "c"], session=test_session).save(
+        "letters", attrs=["letter"]
+    )
+
+    assert sorted(dc.datasets(attrs=attrs).collect("name")) == sorted(result)
+
+
 def test_listings(test_session, tmp_dir):
     df = pd.DataFrame(DF_DATA)
     df.to_parquet(tmp_dir / "df.parquet")