datachain 0.6.8__py3-none-any.whl → 0.6.10__py3-none-any.whl

datachain/catalog/catalog.py

@@ -603,9 +603,10 @@ class Catalog:
         )
 
         lst = Listing(
+            self.metastore.clone(),
             self.warehouse.clone(),
             Client.get_client(list_uri, self.cache, **self.client_config),
-            self.get_dataset(list_ds_name),
+            dataset_name=list_ds_name,
             object_name=object_name,
         )
 
@@ -698,9 +699,13 @@ class Catalog:
 
                     client = self.get_client(source, **client_config)
                     uri = client.uri
-                    st = self.warehouse.clone()
                     dataset_name, _, _, _ = DataChain.parse_uri(uri, self.session)
-                    listing = Listing(st, client, self.get_dataset(dataset_name))
+                    listing = Listing(
+                        self.metastore.clone(),
+                        self.warehouse.clone(),
+                        client,
+                        dataset_name=dataset_name,
+                    )
                     rows = DatasetQuery(
                         name=dataset.name, version=ds_version, catalog=self
                     ).to_db_records()
@@ -769,6 +774,7 @@ class Catalog:
         create_rows: Optional[bool] = True,
         validate_version: Optional[bool] = True,
         listing: Optional[bool] = False,
+        uuid: Optional[str] = None,
     ) -> "DatasetRecord":
         """
         Creates new dataset of a specific version.
@@ -816,6 +822,7 @@ class Catalog:
             query_script=query_script,
             create_rows_table=create_rows,
             columns=columns,
+            uuid=uuid,
         )
 
     def create_new_dataset_version(
@@ -832,6 +839,7 @@ class Catalog:
         script_output="",
         create_rows_table=True,
         job_id: Optional[str] = None,
+        uuid: Optional[str] = None,
     ) -> DatasetRecord:
         """
         Creates dataset version if it doesn't exist.
@@ -855,6 +863,7 @@ class Catalog:
             schema=schema,
             job_id=job_id,
             ignore_if_exists=True,
+            uuid=uuid,
         )
 
         if create_rows_table:
@@ -1350,6 +1359,13 @@ class Catalog:
             # we will create new one if it doesn't exist
             pass
 
+        if dataset and version and dataset.has_version(version):
+            """No need to communicate with Studio at all"""
+            dataset_uri = create_dataset_uri(remote_dataset_name, version)
+            print(f"Local copy of dataset {dataset_uri} already present")
+            _instantiate_dataset()
+            return
+
         remote_dataset = self.get_remote_dataset(remote_dataset_name)
         # if version is not specified in uri, take the latest one
         if not version:
@@ -1400,6 +1416,7 @@ class Catalog:
             columns=columns,
             feature_schema=remote_dataset_version.feature_schema,
             validate_version=False,
+            uuid=remote_dataset_version.uuid,
         )
 
         # asking remote to export dataset rows table to s3 and to return signed

datachain/client/fsspec.py

@@ -358,7 +358,7 @@ class Client(ABC):
     ) -> BinaryIO:
         """Open a file, including files in tar archives."""
         if use_cache and (cache_path := self.cache.get_path(file)):
-            return open(cache_path, mode="rb")  # noqa: SIM115
+            return open(cache_path, mode="rb")
         assert not file.location
         return FileWrapper(self.fs.open(self.get_full_path(file.path)), cb)  # type: ignore[return-value]
 

datachain/data_storage/metastore.py

@@ -138,6 +138,7 @@ class AbstractMetastore(ABC, Serializable):
         size: Optional[int] = None,
         preview: Optional[list[dict]] = None,
         job_id: Optional[str] = None,
+        uuid: Optional[str] = None,
     ) -> DatasetRecord:
         """Creates new dataset version."""
 
@@ -352,6 +353,7 @@ class AbstractDBMetastore(AbstractMetastore):
         """Datasets versions table columns."""
         return [
             Column("id", Integer, primary_key=True),
+            Column("uuid", Text, nullable=False, default=uuid4()),
             Column(
                 "dataset_id",
                 Integer,
@@ -545,6 +547,7 @@ class AbstractDBMetastore(AbstractMetastore):
         size: Optional[int] = None,
         preview: Optional[list[dict]] = None,
         job_id: Optional[str] = None,
+        uuid: Optional[str] = None,
         conn=None,
     ) -> DatasetRecord:
         """Creates new dataset version."""
@@ -555,6 +558,7 @@ class AbstractDBMetastore(AbstractMetastore):
 
         query = self._datasets_versions_insert().values(
             dataset_id=dataset.id,
+            uuid=uuid or str(uuid4()),
             version=version,
             status=status,
             feature_schema=json.dumps(feature_schema or {}),

datachain/data_storage/sqlite.py

@@ -747,8 +747,12 @@ class SQLiteWarehouse(AbstractWarehouse):
 
         ids = self.db.execute(select_ids).fetchall()
 
-        select_q = query.with_only_columns(
-            *[c for c in query.selected_columns if c.name != "sys__id"]
+        select_q = (
+            query.with_only_columns(
+                *[c for c in query.selected_columns if c.name != "sys__id"]
+            )
+            .offset(None)
+            .limit(None)
         )
 
         for batch in batched_it(ids, 10_000):

datachain/dataset.py

@@ -163,6 +163,7 @@ class DatasetStatus:
 @dataclass
 class DatasetVersion:
     id: int
+    uuid: str
     dataset_id: int
     version: int
     status: int
@@ -184,6 +185,7 @@ class DatasetVersion:
     def parse(  # noqa: PLR0913
         cls: type[V],
         id: int,
+        uuid: str,
         dataset_id: int,
         version: int,
         status: int,
@@ -203,6 +205,7 @@ class DatasetVersion:
     ):
         return cls(
             id,
+            uuid,
             dataset_id,
             version,
             status,
@@ -306,6 +309,7 @@ class DatasetRecord:
         query_script: str,
         schema: str,
         version_id: int,
+        version_uuid: str,
         version_dataset_id: int,
         version: int,
         version_status: int,
@@ -331,6 +335,7 @@ class DatasetRecord:
 
         dataset_version = DatasetVersion.parse(
             version_id,
+            version_uuid,
             version_dataset_id,
             version,
             version_status,

datachain/lib/dataset_info.py

@@ -1,6 +1,7 @@
 import json
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Optional, Union
+from uuid import uuid4
 
 from pydantic import Field, field_validator
 
@@ -15,6 +16,7 @@ if TYPE_CHECKING:
 
 class DatasetInfo(DataModel):
     name: str
+    uuid: str = Field(default=str(uuid4()))
     version: int = Field(default=1)
     status: int = Field(default=DatasetStatus.CREATED)
     created_at: datetime = Field(default=TIME_ZERO)
@@ -60,6 +62,7 @@ class DatasetInfo(DataModel):
         job: Optional[Job],
     ) -> "Self":
         return cls(
+            uuid=version.uuid,
             name=dataset.name,
             version=version.version,
             status=version.status,

datachain/lib/dc.py

@@ -30,7 +30,7 @@ from datachain.client.local import FileClient
 from datachain.dataset import DatasetRecord
 from datachain.lib.convert.python_to_sql import python_to_sql
 from datachain.lib.convert.values_to_tuples import values_to_tuples
-from datachain.lib.data_model import DataModel, DataType, dict_to_data_model
+from datachain.lib.data_model import DataModel, DataType, DataValue, dict_to_data_model
 from datachain.lib.dataset_info import DatasetInfo
 from datachain.lib.file import ArrowRow, File, get_file_type
 from datachain.lib.file import ExportPlacement as FileExportPlacement
@@ -642,6 +642,59 @@ class DataChain:
         }
         return chain.gen(**signal_dict)  # type: ignore[misc, arg-type]
 
+    def explode(
+        self,
+        col: str,
+        model_name: Optional[str] = None,
+        object_name: Optional[str] = None,
+    ) -> "DataChain":
+        """Explodes a column containing JSON objects (dict or str DataChain type) into
+           individual columns based on the schema of the JSON. Schema is inferred from
+           the first row of the column.
+
+        Args:
+            col: the name of the column containing JSON to be exploded.
+            model_name: optional generated model name.  By default generates the name
+                automatically.
+            object_name: optional generated object column name. By default generates the
+                name automatically.
+
+        Returns:
+            DataChain: A new DataChain instance with the new set of columns.
+        """
+        import json
+
+        import pyarrow as pa
+
+        from datachain.lib.arrow import schema_to_output
+
+        json_value = next(self.limit(1).collect(col))
+        json_dict = (
+            json.loads(json_value) if isinstance(json_value, str) else json_value
+        )
+
+        if not isinstance(json_dict, dict):
+            raise TypeError(f"Column {col} should be a string or dict type with JSON")
+
+        schema = pa.Table.from_pylist([json_dict]).schema
+        output = schema_to_output(schema, None)
+
+        if not model_name:
+            model_name = f"{col.title()}ExplodedModel"
+
+        model = dict_to_data_model(model_name, output)
+
+        def json_to_model(json_value: Union[str, dict]):
+            json_dict = (
+                json.loads(json_value) if isinstance(json_value, str) else json_value
+            )
+            return model.model_validate(json_dict)
+
+        if not object_name:
+            object_name = f"{col}_expl"
+
+        return self.map(json_to_model, params=col, output={object_name: model})
+
     @classmethod
     def datasets(
         cls,
@@ -895,7 +948,7 @@ class DataChain:
         2. Group-based UDF function input: Instead of individual rows, the function
            receives a list all rows within each group defined by `partition_by`.
 
-        Example:
+        Examples:
             ```py
             chain = chain.agg(
                 total=lambda category, amount: [sum(amount)],
@@ -904,6 +957,26 @@ class DataChain:
             )
             chain.save("new_dataset")
             ```
+
+            An alternative syntax, when you need to specify a more complex function:
+
+            ```py
+            # It automatically resolves which columns to pass to the function
+            # by looking at the function signature.
+            def agg_sum(
+                file: list[File], amount: list[float]
+            ) -> Iterator[tuple[File, float]]:
+                yield file[0], sum(amount)
+
+            chain = chain.agg(
+                agg_sum,
+                output={"file": File, "total": float},
+                # Alternative syntax is to use `C` (short for Column) to specify
+                # a column name or a nested column, e.g. C("file.path").
+                partition_by=C("category"),
+            )
+            chain.save("new_dataset")
+            ```
         """
         udf_obj = self._udf_to_obj(Aggregator, func, params, output, signal_map)
         return self._evolve(
@@ -1242,15 +1315,15 @@ class DataChain:
         return self.results(row_factory=to_dict)
 
     @overload
-    def collect(self) -> Iterator[tuple[DataType, ...]]: ...
+    def collect(self) -> Iterator[tuple[DataValue, ...]]: ...
 
     @overload
-    def collect(self, col: str) -> Iterator[DataType]: ...  # type: ignore[overload-overlap]
+    def collect(self, col: str) -> Iterator[DataValue]: ...
 
     @overload
-    def collect(self, *cols: str) -> Iterator[tuple[DataType, ...]]: ...
+    def collect(self, *cols: str) -> Iterator[tuple[DataValue, ...]]: ...
 
-    def collect(self, *cols: str) -> Iterator[Union[DataType, tuple[DataType, ...]]]:  # type: ignore[overload-overlap,misc]
+    def collect(self, *cols: str) -> Iterator[Union[DataValue, tuple[DataValue, ...]]]:  # type: ignore[overload-overlap,misc]
         """Yields rows of values, optionally limited to the specified columns.
 
         Args:

datachain/lib/meta_formats.py

@@ -114,6 +114,7 @@ def read_meta(  # noqa: C901
             )
         )
         (model_output,) = chain.collect("meta_schema")
+        assert isinstance(model_output, str)
         if print_schema:
             print(f"{model_output}")
         # Below 'spec' should be a dynamically converted DataModel from Pydantic

datachain/lib/models/__init__.py

@@ -1,5 +1,6 @@
-from . import yolo
-from .bbox import BBox
+from . import ultralytics
+from .bbox import BBox, OBBox
 from .pose import Pose, Pose3D
+from .segment import Segments
 
-__all__ = ["BBox", "Pose", "Pose3D", "yolo"]
+__all__ = ["BBox", "OBBox", "Pose", "Pose3D", "Segments", "ultralytics"]

datachain/lib/models/bbox.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 from pydantic import Field
 
 from datachain.lib.data_model import DataModel
@@ -11,10 +9,7 @@ class BBox(DataModel):
 
     Attributes:
         title (str): The title of the bounding box.
-        x1 (float): The x-coordinate of the top-left corner of the bounding box.
-        y1 (float): The y-coordinate of the top-left corner of the bounding box.
-        x2 (float): The x-coordinate of the bottom-right corner of the bounding box.
-        y2 (float): The y-coordinate of the bottom-right corner of the bounding box.
+        coords (list[int]): The coordinates of the bounding box.
 
     The bounding box is defined by two points:
         - (x1, y1): The top-left corner of the box.
@@ -22,24 +17,100 @@ class BBox(DataModel):
     """
 
     title: str = Field(default="")
-    x1: float = Field(default=0)
-    y1: float = Field(default=0)
-    x2: float = Field(default=0)
-    y2: float = Field(default=0)
+    coords: list[int] = Field(default=None)
+
+    @staticmethod
+    def from_list(coords: list[float], title: str = "") -> "BBox":
+        assert len(coords) == 4, "Bounding box coordinates must be a list of 4 floats."
+        assert all(
+            isinstance(value, (int, float)) for value in coords
+        ), "Bounding box coordinates must be integers or floats."
+        return BBox(
+            title=title,
+            coords=[round(c) for c in coords],
+        )
+
+    @staticmethod
+    def from_dict(coords: dict[str, float], title: str = "") -> "BBox":
+        assert (
+            len(coords) == 4
+        ), "Bounding box coordinates must be a dictionary of 4 floats."
+        assert set(coords) == {
+            "x1",
+            "y1",
+            "x2",
+            "y2",
+        }, "Bounding box coordinates must contain keys with coordinates."
+        assert all(
+            isinstance(value, (int, float)) for value in coords.values()
+        ), "Bounding box coordinates must be integers or floats."
+        return BBox(
+            title=title,
+            coords=[
+                round(coords["x1"]),
+                round(coords["y1"]),
+                round(coords["x2"]),
+                round(coords["y2"]),
+            ],
+        )
+
+
+class OBBox(DataModel):
+    """
+    A data model for representing oriented bounding boxes.
+
+    Attributes:
+        title (str): The title of the oriented bounding box.
+        coords (list[int]): The coordinates of the oriented bounding box.
+
+    The oriented bounding box is defined by four points:
+        - (x1, y1): The first corner of the box.
+        - (x2, y2): The second corner of the box.
+        - (x3, y3): The third corner of the box.
+        - (x4, y4): The fourth corner of the box.
+    """
+
+    title: str = Field(default="")
+    coords: list[int] = Field(default=None)
+
+    @staticmethod
+    def from_list(coords: list[float], title: str = "") -> "OBBox":
+        assert (
+            len(coords) == 8
+        ), "Oriented bounding box coordinates must be a list of 8 floats."
+        assert all(
+            isinstance(value, (int, float)) for value in coords
+        ), "Oriented bounding box coordinates must be integers or floats."
+        return OBBox(
+            title=title,
+            coords=[round(c) for c in coords],
+        )
 
     @staticmethod
-    def from_xywh(bbox: list[float], title: Optional[str] = None) -> "BBox":
-        """
-        Converts a bounding box in (x, y, width, height) format
-        to a BBox data model instance.
-
-        Args:
-            bbox (list[float]): A bounding box, represented as a list
-                                of four floats [x, y, width, height].
-
-        Returns:
-            BBox2D: An instance of the BBox data model.
-        """
-        assert len(bbox) == 4, f"Bounding box must have 4 elements, got f{len(bbox)}"
-        x, y, w, h = bbox
-        return BBox(title=title or "", x1=x, y1=y, x2=x + w, y2=y + h)
+    def from_dict(coords: dict[str, float], title: str = "") -> "OBBox":
+        assert set(coords) == {
+            "x1",
+            "y1",
+            "x2",
+            "y2",
+            "x3",
+            "y3",
+            "x4",
+            "y4",
+        }, "Oriented bounding box coordinates must contain keys with coordinates."
+        assert all(
+            isinstance(value, (int, float)) for value in coords.values()
+        ), "Oriented bounding box coordinates must be integers or floats."
+        return OBBox(
+            title=title,
+            coords=[
+                round(coords["x1"]),
+                round(coords["y1"]),
+                round(coords["x2"]),
+                round(coords["y2"]),
+                round(coords["x3"]),
+                round(coords["y3"]),
+                round(coords["x4"]),
+                round(coords["y4"]),
+            ],
+        )

datachain/lib/models/pose.py

@@ -8,15 +8,48 @@ class Pose(DataModel):
     A data model for representing pose keypoints.
 
     Attributes:
-        x (list[float]): The x-coordinates of the keypoints.
-        y (list[float]): The y-coordinates of the keypoints.
+        x (list[int]): The x-coordinates of the keypoints.
+        y (list[int]): The y-coordinates of the keypoints.
 
     The keypoints are represented as lists of x and y coordinates, where each index
     corresponds to a specific body part.
     """
 
-    x: list[float] = Field(default=None)
-    y: list[float] = Field(default=None)
+    x: list[int] = Field(default=None)
+    y: list[int] = Field(default=None)
+
+    @staticmethod
+    def from_list(points: list[list[float]]) -> "Pose":
+        assert len(points) == 2, "Pose coordinates must be a list of 2 lists."
+        points_x, points_y = points
+        assert (
+            len(points_x) == len(points_y) == 17
+        ), "Pose x and y coordinates must have the same length of 17."
+        assert all(
+            isinstance(value, (int, float)) for value in [*points_x, *points_y]
+        ), "Pose coordinates must be integers or floats."
+        return Pose(
+            x=[round(coord) for coord in points_x],
+            y=[round(coord) for coord in points_y],
+        )
+
+    @staticmethod
+    def from_dict(points: dict[str, list[float]]) -> "Pose":
+        assert set(points) == {
+            "x",
+            "y",
+        }, "Pose coordinates must contain keys 'x' and 'y'."
+        points_x, points_y = points["x"], points["y"]
+        assert (
+            len(points_x) == len(points_y) == 17
+        ), "Pose x and y coordinates must have the same length of 17."
+        assert all(
+            isinstance(value, (int, float)) for value in [*points_x, *points_y]
+        ), "Pose coordinates must be integers or floats."
+        return Pose(
+            x=[round(coord) for coord in points_x],
+            y=[round(coord) for coord in points_y],
+        )
 
 
 class Pose3D(DataModel):
@@ -24,14 +57,52 @@ class Pose3D(DataModel):
     A data model for representing 3D pose keypoints.
 
     Attributes:
-        x (list[float]): The x-coordinates of the keypoints.
-        y (list[float]): The y-coordinates of the keypoints.
+        x (list[int]): The x-coordinates of the keypoints.
+        y (list[int]): The y-coordinates of the keypoints.
         visible (list[float]): The visibility of the keypoints.
 
     The keypoints are represented as lists of x, y, and visibility values,
     where each index corresponds to a specific body part.
     """
 
-    x: list[float] = Field(default=None)
-    y: list[float] = Field(default=None)
+    x: list[int] = Field(default=None)
+    y: list[int] = Field(default=None)
     visible: list[float] = Field(default=None)
+
+    @staticmethod
+    def from_list(points: list[list[float]]) -> "Pose3D":
+        assert len(points) == 3, "Pose coordinates must be a list of 3 lists."
+        points_x, points_y, points_v = points
+        assert (
+            len(points_x) == len(points_y) == len(points_v) == 17
+        ), "Pose x, y, and visibility coordinates must have the same length of 17."
+        assert all(
+            isinstance(value, (int, float))
+            for value in [*points_x, *points_y, *points_v]
+        ), "Pose coordinates must be integers or floats."
+        return Pose3D(
+            x=[round(coord) for coord in points_x],
+            y=[round(coord) for coord in points_y],
+            visible=points_v,
+        )
+
+    @staticmethod
+    def from_dict(points: dict[str, list[float]]) -> "Pose3D":
+        assert set(points) == {
+            "x",
+            "y",
+            "visible",
+        }, "Pose coordinates must contain keys 'x', 'y', and 'visible'."
+        points_x, points_y, points_v = points["x"], points["y"], points["visible"]
+        assert (
+            len(points_x) == len(points_y) == len(points_v) == 17
+        ), "Pose x, y, and visibility coordinates must have the same length of 17."
+        assert all(
+            isinstance(value, (int, float))
+            for value in [*points_x, *points_y, *points_v]
+        ), "Pose coordinates must be integers or floats."
+        return Pose3D(
+            x=[round(coord) for coord in points_x],
+            y=[round(coord) for coord in points_y],
+            visible=points_v,
+        )

datachain/lib/models/segment.py

@@ -0,0 +1,53 @@
+from pydantic import Field
+
+from datachain.lib.data_model import DataModel
+
+
+class Segments(DataModel):
+    """
+    A data model for representing segments.
+
+    Attributes:
+        title (str): The title of the segments.
+        x (list[int]): The x-coordinates of the segments.
+        y (list[int]): The y-coordinates of the segments.
+
+    The segments are represented as lists of x and y coordinates, where each index
+    corresponds to a specific segment.
+    """
+
+    title: str = Field(default="")
+    x: list[int] = Field(default=None)
+    y: list[int] = Field(default=None)
+
+    @staticmethod
+    def from_list(points: list[list[float]], title: str = "") -> "Segments":
+        assert len(points) == 2, "Segments coordinates must be a list of 2 lists."
+        points_x, points_y = points
+        assert len(points_x) == len(
+            points_y
+        ), "Segments x and y coordinates must have the same length."
+        assert all(
+            isinstance(value, (int, float)) for value in [*points_x, *points_y]
+        ), "Segments coordinates must be integers or floats."
+        return Segments(
+            title=title,
+            x=[round(coord) for coord in points_x],
+            y=[round(coord) for coord in points_y],
+        )
+
+    @staticmethod
+    def from_dict(points: dict[str, list[float]], title: str = "") -> "Segments":
+        assert set(points) == {
+            "x",
+            "y",
+        }, "Segments coordinates must contain keys 'x' and 'y'."
+        points_x, points_y = points["x"], points["y"]
+        assert all(
+            isinstance(value, (int, float)) for value in [*points_x, *points_y]
+        ), "Segments coordinates must be integers or floats."
+        return Segments(
+            title=title,
+            x=[round(coord) for coord in points_x],
+            y=[round(coord) for coord in points_y],
+        )

datachain/lib/models/ultralytics/__init__.py

@@ -0,0 +1,14 @@
+from .bbox import YoloBBox, YoloBBoxes, YoloOBBox, YoloOBBoxes
+from .pose import YoloPose, YoloPoses
+from .segment import YoloSegment, YoloSegments
+
+__all__ = [
+    "YoloBBox",
+    "YoloBBoxes",
+    "YoloOBBox",
+    "YoloOBBoxes",
+    "YoloPose",
+    "YoloPoses",
+    "YoloSegment",
+    "YoloSegments",
+]