datachain 0.3.5__tar.gz → 0.3.7__tar.gz

{datachain-0.3.5/src/datachain.egg-info → datachain-0.3.7}/PKG-INFO

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

{datachain-0.3.5 → datachain-0.3.7}/src/datachain/catalog/catalog.py

@@ -120,13 +120,25 @@ def noop(_: str):
 
 @contextmanager
 def print_and_capture(
-    stream: "IO[str]", callback: Callable[[str], None] = noop
+    stream: "IO[bytes]|IO[str]", callback: Callable[[str], None] = noop
 ) -> "Iterator[list[str]]":
     lines: list[str] = []
     append = lines.append
 
     def loop() -> None:
-        for line in iter(stream.readline, ""):
+        buffer = b""
+        while byt := stream.read(1):  # Read one byte at a time
+            buffer += byt.encode("utf-8") if isinstance(byt, str) else byt
+
+            if byt in (b"\n", b"\r"):  # Check for newline or carriage return
+                line = buffer.decode("utf-8")
+                print(line, end="")
+                callback(line)
+                append(line)
+                buffer = b""  # Clear buffer for next line
+
+        if buffer:  # Handle any remaining data in the buffer
+            line = buffer.decode("utf-8")
             print(line, end="")
             callback(line)
             append(line)
@@ -2128,7 +2140,7 @@ class Catalog:
             stdout=subprocess.PIPE if capture_output else None,
             stderr=subprocess.STDOUT if capture_output else None,
             bufsize=1,
-            text=True,
+            text=False,
             **kwargs,
         ) as proc:
             os.close(w)

{datachain-0.3.5 → datachain-0.3.7}/src/datachain/data_storage/schema.py

@@ -50,7 +50,7 @@ def convert_rows_custom_column_types(
     columns: "ColumnCollection[str, ColumnElement[Any]]",
     rows: Iterator[tuple[Any, ...]],
     dialect: "Dialect",
-):
+) -> Iterator[tuple[Any, ...]]:
     """
     This function converts values of rows columns based on their types which are
     defined in columns. We are only converting column values for which types are

{datachain-0.3.5 → datachain-0.3.7}/src/datachain/data_storage/sqlite.py

@@ -27,10 +27,7 @@ import datachain.sql.sqlite
 from datachain.data_storage import AbstractDBMetastore, AbstractWarehouse
 from datachain.data_storage.db_engine import DatabaseEngine
 from datachain.data_storage.id_generator import AbstractDBIDGenerator
-from datachain.data_storage.schema import (
-    DefaultSchema,
-    convert_rows_custom_column_types,
-)
+from datachain.data_storage.schema import DefaultSchema
 from datachain.dataset import DatasetRecord
 from datachain.error import DataChainError
 from datachain.sql.sqlite import create_user_defined_sql_functions, sqlite_dialect
@@ -209,6 +206,7 @@ class SQLiteDatabaseEngine(DatabaseEngine):
             return cursor.executemany(self.compile(query).string, params)
         return self.db.executemany(self.compile(query).string, params)
 
+    @retry_sqlite_locks
     def execute_str(self, sql: str, parameters=None) -> sqlite3.Cursor:
         if parameters is None:
             return self.db.execute(sql)
@@ -650,12 +648,6 @@ class SQLiteWarehouse(AbstractWarehouse):
         self.db.create_table(table, if_not_exists=if_not_exists)
         return table
 
-    def dataset_rows_select(self, select_query: Select, **kwargs):
-        rows = self.db.execute(select_query, **kwargs)
-        yield from convert_rows_custom_column_types(
-            select_query.selected_columns, rows, sqlite_dialect
-        )
-
     def get_dataset_sources(
         self, dataset: DatasetRecord, version: int
     ) -> list[StorageURI]:

{datachain-0.3.5 → datachain-0.3.7}/src/datachain/data_storage/warehouse.py

@@ -17,6 +17,7 @@ from sqlalchemy.sql.expression import true
 from tqdm import tqdm
 
 from datachain.client import Client
+from datachain.data_storage.schema import convert_rows_custom_column_types
 from datachain.data_storage.serializer import Serializable
 from datachain.dataset import DatasetRecord
 from datachain.node import DirType, DirTypeGroup, Entry, Node, NodeWithPath, get_path
@@ -226,7 +227,7 @@ class AbstractWarehouse(ABC, Serializable):
                     if limit < page_size:
                         paginated_query = paginated_query.limit(None).limit(limit)
 
-                results = self.db.execute(paginated_query.offset(offset))
+                results = self.dataset_rows_select(paginated_query.offset(offset))
 
                 processed = False
                 for row in results:
@@ -309,12 +310,18 @@ class AbstractWarehouse(ABC, Serializable):
         Merge results should not contain duplicates.
         """
 
-    @abstractmethod
-    def dataset_rows_select(self, select_query: sa.sql.selectable.Select, **kwargs):
+    def dataset_rows_select(
+        self,
+        query: sa.sql.selectable.Select,
+        **kwargs,
+    ) -> Iterator[tuple[Any, ...]]:
         """
-        Method for fetching dataset rows from database. This is abstract since
-        in some DBs we need to use special settings
+        Fetch dataset rows from database.
         """
+        rows = self.db.execute(query, **kwargs)
+        yield from convert_rows_custom_column_types(
+            query.selected_columns, rows, self.db.dialect
+        )
 
     @abstractmethod
     def get_dataset_sources(

{datachain-0.3.5 → datachain-0.3.7}/src/datachain/lib/dc.py

@@ -839,6 +839,10 @@ class DataChain(DatasetQuery):
     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.
+
         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.
@@ -859,7 +863,26 @@ class DataChain(DatasetQuery):
                 dist=cosine_distance(embedding_text, embedding_image)
         )
         ```
+
+        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.
+
+        Example:
+        ```py
+         dc.mutate(
+                newkey=Column("oldkey")
+        )
+        ```
         """
+        existing_columns = set(self.signals_schema.values.keys())
+        for col_name in kwargs:
+            if col_name in existing_columns:
+                raise DataChainColumnError(
+                    col_name,
+                    "Cannot modify existing column with mutate(). "
+                    "Use a different name for the new column.",
+                )
         for col_name, expr in kwargs.items():
             if not isinstance(expr, Column) and isinstance(expr.type, NullType):
                 raise DataChainColumnError(

{datachain-0.3.5 → datachain-0.3.7/src/datachain.egg-info}/PKG-INFO

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

{datachain-0.3.5 → datachain-0.3.7}/tests/func/test_datachain.py

@@ -8,10 +8,11 @@ import pandas as pd
 import pytest
 import pytz
 from PIL import Image
+from sqlalchemy import Column
 
 from datachain.data_storage.sqlite import SQLiteWarehouse
 from datachain.dataset import DatasetStats
-from datachain.lib.dc import DataChain
+from datachain.lib.dc import DataChain, DataChainColumnError
 from datachain.lib.file import File, ImageFile
 from tests.utils import images_equal
 
@@ -314,3 +315,16 @@ def test_from_storage_check_rows(tmp_dir, test_session):
             location=None,
             vtype="",
         )
+
+
+def test_mutate_existing_column(catalog):
+    ds = DataChain.from_values(ids=[1, 2, 3])
+
+    with pytest.raises(DataChainColumnError) as excinfo:
+        ds.mutate(ids=Column("ids") + 1)
+
+    assert (
+        str(excinfo.value)
+        == "Error for column ids: Cannot modify existing column with mutate()."
+        " Use a different name for the new column."
+    )