datachain 0.30.2__py3-none-any.whl → 0.30.4__py3-none-any.whl

datachain/lib/dc/database.py

@@ -77,12 +77,15 @@ def to_database(
     on_conflict: Optional[str] = None,
     conflict_columns: Optional[list[str]] = None,
     column_mapping: Optional[dict[str, Optional[str]]] = None,
-) -> None:
+) -> int:
     """
     Implementation function for exporting DataChain to database tables.
 
     This is the core implementation that handles the actual database operations.
     For user-facing documentation, see DataChain.to_database() method.
+
+    Returns:
+        int: Number of rows affected (inserted/updated).
     """
     if on_conflict and on_conflict not in ("ignore", "update"):
         raise ValueError(
@@ -101,11 +104,16 @@ def to_database(
         all_columns, normalized_column_mapping
     )
 
+    normalized_conflict_columns = _normalize_conflict_columns(
+        conflict_columns, normalized_column_mapping
+    )
+
     with _connect(connection) as conn:
         metadata = sqlalchemy.MetaData()
         table = sqlalchemy.Table(table_name, metadata, *columns)
 
         table_existed_before = False
+        total_rows_affected = 0
         try:
             with conn.begin():
                 # Check if table exists to determine if we should clean up on error.
@@ -117,14 +125,18 @@ def to_database(
 
                 rows_iter = chain._leaf_values()
                 for batch in batched(rows_iter, batch_rows):
-                    _process_batch(
+                    rows_affected = _process_batch(
                         conn,
                         table,
                         batch,
                         on_conflict,
-                        conflict_columns,
+                        normalized_conflict_columns,
                         column_indices_and_names,
                     )
+                    if rows_affected < 0 or total_rows_affected < 0:
+                        total_rows_affected = -1
+                    else:
+                        total_rows_affected += rows_affected
         except Exception:
             if not table_existed_before:
                 try:
@@ -134,6 +146,8 @@ def to_database(
                     pass
             raise
 
+    return total_rows_affected
+
 
 def _normalize_column_mapping(
     column_mapping: dict[str, Optional[str]],
@@ -174,6 +188,30 @@ def _normalize_column_mapping(
     return normalized_mapping
 
 
+def _normalize_conflict_columns(
+    conflict_columns: Optional[list[str]], column_mapping: dict[str, Optional[str]]
+) -> Optional[list[str]]:
+    """
+    Normalize conflict_columns by converting DataChain format to database format
+    and applying column mapping.
+    """
+    if not conflict_columns:
+        return None
+
+    normalized_columns = []
+    for col in conflict_columns:
+        db_col = ColumnMeta.to_db_name(col)
+
+        if db_col in column_mapping or hasattr(column_mapping, "default_factory"):
+            mapped_name = column_mapping[db_col]
+            if mapped_name:
+                normalized_columns.append(mapped_name)
+        else:
+            normalized_columns.append(db_col)
+
+    return normalized_columns
+
+
 def _prepare_columns(all_columns, column_mapping):
     """Prepare column mapping and column definitions."""
     column_indices_and_names = []  # List of (index, target_name) tuples
@@ -192,8 +230,12 @@ def _prepare_columns(all_columns, column_mapping):
 
 def _process_batch(
     conn, table, batch, on_conflict, conflict_columns, column_indices_and_names
-):
-    """Process a batch of rows with conflict resolution."""
+) -> int:
+    """Process a batch of rows with conflict resolution.
+
+    Returns:
+        int: Number of rows affected by the insert operation.
+    """
 
     def prepare_row(row_values):
         """Convert a row tuple to a dictionary with proper DB column names."""
@@ -206,6 +248,7 @@ def _process_batch(
 
     supports_conflict = on_conflict and conn.engine.name in ("postgresql", "sqlite")
 
+    insert_stmt: Any  # Can be PostgreSQL, SQLite, or regular insert statement
     if supports_conflict:
         # Use dialect-specific insert for conflict resolution
         if conn.engine.name == "postgresql":
@@ -249,7 +292,8 @@ def _process_batch(
             stacklevel=2,
         )
 
-    conn.execute(insert_stmt, rows_to_insert)
+    result = conn.execute(insert_stmt, rows_to_insert)
+    return result.rowcount
 
 
 def read_database(

datachain/lib/dc/datachain.py

@@ -67,6 +67,7 @@ from .utils import (
     Sys,
     _get_merge_error_str,
     _validate_merge_on,
+    is_studio,
     resolve_columns,
 )
 
@@ -284,7 +285,11 @@ class DataChain:
         """Underlying dataset, if there is one."""
         if not self.name:
             return None
-        return self.session.catalog.get_dataset(self.name, self._query.project)
+        return self.session.catalog.get_dataset(
+            self.name,
+            namespace_name=self._query.project.namespace.name,
+            project_name=self._query.project.name,
+        )
 
     def __or__(self, other: "Self") -> "Self":
         """Return `self.union(other)`."""
@@ -605,7 +610,7 @@ class DataChain:
             project = self.session.catalog.metastore.get_project(
                 project_name,
                 namespace_name,
-                create=self.session.catalog.metastore.project_allowed_to_create,
+                create=is_studio(),
             )
         except ProjectNotFoundError as e:
             # not being able to create it as creation is not allowed
@@ -1180,17 +1185,13 @@ class DataChain:
         )
 
     def mutate(self, **kwargs) -> "Self":
-        """Create new signals based on existing signals.
-
-        This method cannot modify existing columns. If you need to modify an
-        existing column, use a different name for the new column and then use
-        `select()` to choose which columns to keep.
+        """Create or modify signals based on existing signals.
 
         This method is vectorized and more efficient compared to map(), and it does not
         extract or download any data from the internal database. However, it can only
         utilize predefined built-in functions and their combinations.
 
-        The supported functions:
+        Supported functions:
            Numerical:   +, -, *, /, rand(), avg(), count(), func(),
                         greatest(), least(), max(), min(), sum()
            String:      length(), split(), replace(), regexp_replace()
@@ -1217,13 +1218,20 @@ class DataChain:
         ```
 
         This method can be also used to rename signals. If the Column("name") provided
-        as value for the new signal - the old column will be dropped. Otherwise a new
-        column is created.
+        as value for the new signal - the old signal will be dropped. Otherwise a new
+        signal is created. Exception, if the old signal is nested one (e.g.
+        `C("file.path")`), it will be kept to keep the object intact.
 
         Example:
         ```py
          dc.mutate(
-            newkey=Column("oldkey")
+            newkey=Column("oldkey") # drops oldkey
+        )
+        ```
+
+        ```py
+         dc.mutate(
+            size=Column("file.size") # keeps `file.size`
         )
         ```
         """
@@ -1258,8 +1266,10 @@ class DataChain:
                 # adding new signal
                 mutated[name] = value
 
+        new_schema = schema.mutate(kwargs)
         return self._evolve(
-            query=self._query.mutate(**mutated), signal_schema=schema.mutate(kwargs)
+            query=self._query.mutate(new_schema=new_schema, **mutated),
+            signal_schema=new_schema,
         )
 
     @property
@@ -2298,13 +2308,17 @@ class DataChain:
         on_conflict: Optional[str] = None,
         conflict_columns: Optional[list[str]] = None,
         column_mapping: Optional[dict[str, Optional[str]]] = None,
-    ) -> None:
+    ) -> int:
         """Save chain to a database table using a given database connection.
 
         This method exports all DataChain records to a database table, creating the
         table if it doesn't exist and appending data if it does. The table schema
         is automatically inferred from the DataChain's signal schema.
 
+        For PostgreSQL, tables are created in the schema specified by the connection's
+        search_path (defaults to 'public'). Use URL parameters to target specific
+        schemas.
+
         Parameters:
             table_name: Name of the database table to create/write to.
             connection: SQLAlchemy connectable, str, or a sqlite3 connection
@@ -2328,20 +2342,26 @@ class DataChain:
                 - Set values to None to skip columns entirely, or use `defaultdict` to
                   skip all columns except those specified.
 
+        Returns:
+            int: Number of rows affected (inserted/updated). -1 if DB driver doesn't
+                 support telemetry.
+
         Examples:
             Basic usage with PostgreSQL:
             ```py
-            import sqlalchemy as sa
             import datachain as dc
 
-            chain = dc.read_storage("s3://my-bucket/")
-            engine = sa.create_engine("postgresql://user:pass@localhost/mydb")
-            chain.to_database("files_table", engine)
+            rows_affected = (dc
+              .read_storage("s3://my-bucket/")
+              .to_database("files_table", "postgresql://user:pass@localhost/mydb")
+            )
+            print(f"Inserted/updated {rows_affected} rows")
             ```
 
             Using SQLite with connection string:
             ```py
-            chain.to_database("my_table", "sqlite:///data.db")
+            rows_affected = chain.to_database("my_table", "sqlite:///data.db")
+            print(f"Affected {rows_affected} rows")
             ```
 
             Column mapping and renaming:
@@ -2360,7 +2380,9 @@ class DataChain:
             chain.to_database("my_table", engine, on_conflict="ignore")
 
             # Update existing records
-            chain.to_database("my_table", engine, on_conflict="update")
+            chain.to_database(
+               "my_table", engine, on_conflict="update", conflict_columns=["id"]
+            )
             ```
 
             Working with different databases:
@@ -2372,10 +2394,16 @@ class DataChain:
             # SQLite in-memory
             chain.to_database("temp_table", "sqlite:///:memory:")
             ```
+
+            PostgreSQL with schema support:
+            ```py
+            pg_url = "postgresql://user:pass@host/db?options=-c search_path=analytics"
+            chain.to_database("processed_data", pg_url)
+            ```
         """
         from .database import to_database
 
-        to_database(
+        return to_database(
             self,
             table_name,
             connection,

datachain/lib/dc/datasets.py

@@ -13,7 +13,7 @@ from datachain.lib.signal_schema import SignalSchema
 from datachain.query import Session
 from datachain.query.dataset import DatasetQuery
 
-from .utils import Sys
+from .utils import Sys, is_studio
 from .values import read_values
 
 if TYPE_CHECKING:
@@ -343,7 +343,7 @@ def delete_dataset(
         namespace_name=namespace,
     )
 
-    if not catalog.metastore.is_local_dataset(namespace_name) and studio:
+    if not is_studio() and studio:
         return remove_studio_dataset(
             None, name, namespace_name, project_name, version=version, force=force
         )
@@ -357,7 +357,14 @@ def delete_dataset(
         ) from None
 
     if not force:
-        version = version or catalog.get_dataset(name, ds_project).latest_version
+        version = (
+            version
+            or catalog.get_dataset(
+                name,
+                namespace_name=ds_project.namespace.name,
+                project_name=ds_project.name,
+            ).latest_version
+        )
     else:
         version = None
     catalog.remove_dataset(name, ds_project, version=version, force=force)
@@ -403,9 +410,7 @@ def move_dataset(
     namespace, project, name = catalog.get_full_dataset_name(src)
     dest_namespace, dest_project, dest_name = catalog.get_full_dataset_name(dest)
 
-    dataset = catalog.get_dataset(
-        name, catalog.metastore.get_project(project, namespace)
-    )
+    dataset = catalog.get_dataset(name, namespace_name=namespace, project_name=project)
 
     catalog.update_dataset(
         dataset,
@@ -413,6 +418,6 @@ def move_dataset(
         project_id=catalog.metastore.get_project(
             dest_project,
             dest_namespace,
-            create=catalog.metastore.project_allowed_to_create,
+            create=is_studio(),
         ).id,
     )

datachain/lib/dc/utils.py

@@ -15,6 +15,7 @@ from datachain.func.base import Function
 from datachain.lib.data_model import DataModel, DataType
 from datachain.lib.utils import DataChainParamsError
 from datachain.query.schema import DEFAULT_DELIMITER
+from datachain.utils import getenv_bool
 
 if TYPE_CHECKING:
     from typing_extensions import Concatenate, ParamSpec
@@ -26,6 +27,10 @@ if TYPE_CHECKING:
 D = TypeVar("D", bound="DataChain")
 
 
+def is_studio() -> bool:
+    return getenv_bool("DATACHAIN_IS_STUDIO", default=False)
+
+
 def resolve_columns(
     method: "Callable[Concatenate[D, P], D]",
 ) -> "Callable[Concatenate[D, P], D]":

datachain/lib/namespaces.py

@@ -28,7 +28,9 @@ def create(
     """
     session = Session.get(session)
 
-    if not session.catalog.metastore.namespace_allowed_to_create:
+    from datachain.lib.dc.utils import is_studio
+
+    if not is_studio():
         raise NamespaceCreateNotAllowedError("Creating namespace is not allowed")
 
     Namespace.validate_name(name)

datachain/lib/projects.py

@@ -32,7 +32,9 @@ def create(
     """
     session = Session.get(session)
 
-    if not session.catalog.metastore.project_allowed_to_create:
+    from datachain.lib.dc.utils import is_studio
+
+    if not is_studio():
         raise ProjectCreateNotAllowedError("Creating project is not allowed")
 
     Project.validate_name(name)

datachain/lib/signal_schema.py

@@ -34,7 +34,7 @@ from datachain.lib.data_model import DataModel, DataType, DataValue
 from datachain.lib.file import File
 from datachain.lib.model_store import ModelStore
 from datachain.lib.utils import DataChainParamsError
-from datachain.query.schema import DEFAULT_DELIMITER, Column, ColumnMeta
+from datachain.query.schema import DEFAULT_DELIMITER, C, Column, ColumnMeta
 from datachain.sql.types import SQLType
 
 if TYPE_CHECKING:
@@ -680,35 +680,46 @@ class SignalSchema:
         primitives = (bool, str, int, float)
 
         for name, value in args_map.items():
+            current_type = None
+
+            if C.is_nested(name):
+                try:
+                    current_type = self.get_column_type(name)
+                except SignalResolvingError as err:
+                    msg = f"Creating new nested columns directly is not allowed: {name}"
+                    raise ValueError(msg) from err
+
             if isinstance(value, Column) and value.name in self.values:
                 # renaming existing signal
+                # Note: it won't touch nested signals here (e.g. file__path)
+                # we don't allow removing nested columns to keep objects consistent
                 del new_values[value.name]
                 new_values[name] = self.values[value.name]
-                continue
-            if isinstance(value, Column):
+            elif isinstance(value, Column):
                 # adding new signal from existing signal field
-                try:
-                    new_values[name] = self.get_column_type(
-                        value.name, with_subtree=True
-                    )
-                    continue
-                except SignalResolvingError:
-                    pass
-            if isinstance(value, Func):
+                new_values[name] = self.get_column_type(value.name, with_subtree=True)
+            elif isinstance(value, Func):
                 # adding new signal with function
                 new_values[name] = value.get_result_type(self)
-                continue
-            if isinstance(value, primitives):
+            elif isinstance(value, primitives):
                 # For primitives, store the type, not the value
                 val = literal(value)
                 val.type = python_to_sql(type(value))()
                 new_values[name] = sql_to_python(val)
-                continue
-            if isinstance(value, ColumnElement):
+            elif isinstance(value, ColumnElement):
                 # adding new signal
                 new_values[name] = sql_to_python(value)
-                continue
-            new_values[name] = value
+            else:
+                new_values[name] = value
+
+            if C.is_nested(name):
+                if current_type != new_values[name]:
+                    msg = (
+                        f"Altering nested column type is not allowed: {name}, "
+                        f"current type: {current_type}, new type: {new_values[name]}"
+                    )
+                    raise ValueError(msg)
+                del new_values[name]
 
         return SignalSchema(new_values)
 

datachain/listing.py

@@ -65,17 +65,13 @@ class Listing:
 
     @cached_property
     def dataset(self) -> "DatasetRecord":
-        from datachain.error import DatasetNotFoundError
-
         assert self.dataset_name
         project = self.metastore.listing_project
-        try:
-            return self.metastore.get_dataset(self.dataset_name, project.id)
-        except DatasetNotFoundError:
-            raise DatasetNotFoundError(
-                f"Dataset {self.dataset_name} not found in namespace"
-                f" {project.namespace.name} and project {project.name}"
-            ) from None
+        return self.metastore.get_dataset(
+            self.dataset_name,
+            namespace_name=project.namespace.name,
+            project_name=project.name,
+        )
 
     @cached_property
     def dataset_rows(self):

datachain/model/ultralytics/bbox.py

@@ -31,11 +31,11 @@ class YoloBBox(DataModel):
         if not summary:
             return YoloBBox(box=BBox())
         name = summary[0].get("name", "")
-        box = (
-            BBox.from_dict(summary[0]["box"], title=name)
-            if summary[0].get("box")
-            else BBox()
-        )
+        if summary[0].get("box"):
+            assert isinstance(summary[0]["box"], dict)
+            box = BBox.from_dict(summary[0]["box"], title=name)
+        else:
+            box = BBox()
         return YoloBBox(
             cls=summary[0]["class"],
             name=name,
@@ -70,7 +70,8 @@ class YoloBBoxes(DataModel):
                 names.append(name)
                 confidence.append(s["confidence"])
                 if s.get("box"):
-                    box.append(BBox.from_dict(s.get("box"), title=name))
+                    assert isinstance(s["box"], dict)
+                    box.append(BBox.from_dict(s["box"], title=name))
         return YoloBBoxes(
             cls=cls,
             name=names,
@@ -101,11 +102,11 @@ class YoloOBBox(DataModel):
         if not summary:
             return YoloOBBox(box=OBBox())
         name = summary[0].get("name", "")
-        box = (
-            OBBox.from_dict(summary[0]["box"], title=name)
-            if summary[0].get("box")
-            else OBBox()
-        )
+        if summary[0].get("box"):
+            assert isinstance(summary[0]["box"], dict)
+            box = OBBox.from_dict(summary[0]["box"], title=name)
+        else:
+            box = OBBox()
         return YoloOBBox(
             cls=summary[0]["class"],
             name=name,
@@ -140,7 +141,8 @@ class YoloOBBoxes(DataModel):
                 names.append(name)
                 confidence.append(s["confidence"])
                 if s.get("box"):
-                    box.append(OBBox.from_dict(s.get("box"), title=name))
+                    assert isinstance(s["box"], dict)
+                    box.append(OBBox.from_dict(s["box"], title=name))
         return YoloOBBoxes(
             cls=cls,
             name=names,

datachain/model/ultralytics/pose.py

@@ -56,16 +56,16 @@ class YoloPose(DataModel):
         if not summary:
             return YoloPose(box=BBox(), pose=Pose3D())
         name = summary[0].get("name", "")
-        box = (
-            BBox.from_dict(summary[0]["box"], title=name)
-            if summary[0].get("box")
-            else BBox()
-        )
-        pose = (
-            Pose3D.from_dict(summary[0]["keypoints"])
-            if summary[0].get("keypoints")
-            else Pose3D()
-        )
+        if summary[0].get("box"):
+            assert isinstance(summary[0]["box"], dict)
+            box = BBox.from_dict(summary[0]["box"], title=name)
+        else:
+            box = BBox()
+        if summary[0].get("keypoints"):
+            assert isinstance(summary[0]["keypoints"], dict)
+            pose = Pose3D.from_dict(summary[0]["keypoints"])
+        else:
+            pose = Pose3D()
         return YoloPose(
             cls=summary[0]["class"],
             name=name,
@@ -103,9 +103,11 @@ class YoloPoses(DataModel):
                 names.append(name)
                 confidence.append(s["confidence"])
                 if s.get("box"):
-                    box.append(BBox.from_dict(s.get("box"), title=name))
+                    assert isinstance(s["box"], dict)
+                    box.append(BBox.from_dict(s["box"], title=name))
                 if s.get("keypoints"):
-                    pose.append(Pose3D.from_dict(s.get("keypoints")))
+                    assert isinstance(s["keypoints"], dict)
+                    pose.append(Pose3D.from_dict(s["keypoints"]))
         return YoloPoses(
             cls=cls,
             name=names,

datachain/model/ultralytics/segment.py

@@ -34,16 +34,16 @@ class YoloSegment(DataModel):
         if not summary:
             return YoloSegment(box=BBox(), segment=Segment())
         name = summary[0].get("name", "")
-        box = (
-            BBox.from_dict(summary[0]["box"], title=name)
-            if summary[0].get("box")
-            else BBox()
-        )
-        segment = (
-            Segment.from_dict(summary[0]["segments"], title=name)
-            if summary[0].get("segments")
-            else Segment()
-        )
+        if summary[0].get("box"):
+            assert isinstance(summary[0]["box"], dict)
+            box = BBox.from_dict(summary[0]["box"], title=name)
+        else:
+            box = BBox()
+        if summary[0].get("segments"):
+            assert isinstance(summary[0]["segments"], dict)
+            segment = Segment.from_dict(summary[0]["segments"], title=name)
+        else:
+            segment = Segment()
         return YoloSegment(
             cls=summary[0]["class"],
             name=summary[0]["name"],
@@ -81,9 +81,11 @@ class YoloSegments(DataModel):
                 names.append(name)
                 confidence.append(s["confidence"])
                 if s.get("box"):
-                    box.append(BBox.from_dict(s.get("box"), title=name))
+                    assert isinstance(s["box"], dict)
+                    box.append(BBox.from_dict(s["box"], title=name))
                 if s.get("segments"):
-                    segment.append(Segment.from_dict(s.get("segments"), title=name))
+                    assert isinstance(s["segments"], dict)
+                    segment.append(Segment.from_dict(s["segments"], title=name))
         return YoloSegments(
             cls=cls,
             name=names,