datachain 0.3.4__tar.gz → 0.3.6__tar.gz

{datachain-0.3.4 → datachain-0.3.6}/.github/workflows/tests-studio.yml

@@ -17,7 +17,7 @@ concurrency:
 jobs:
   studio:
     if: '!github.event.pull_request.head.repo.fork'
-    runs-on: ubuntu-latest-16-cores
+    runs-on: ubuntu-latest
     strategy:
       matrix:
         pyv: ['3.12']

{datachain-0.3.4 → datachain-0.3.6}/.github/workflows/tests.yml

@@ -62,9 +62,9 @@ jobs:
             pyv: '3.9'
           - os: macos-latest
             pyv: '3.12'
-          - os: windows-latest-8-cores
+          - os: windows-latest
             pyv: '3.9'
-          - os: windows-latest-8-cores
+          - os: windows-latest
             pyv: '3.12'
 
     steps:
@@ -116,9 +116,17 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-latest-16-cores, macos-latest, windows-latest-8-cores]
+        os: [ubuntu-latest, macos-latest, windows-latest]
         pyv: ['3.9', '3.12']
         group: ['get_started', 'llm_and_nlp or computer_vision', 'multimodal']
+        exclude:
+          - {os: ubuntu-latest, pyv: '3.9', group: 'multimodal'}
+          - {os: ubuntu-latest, pyv: '3.12', group: 'multimodal'}
+        include:
+          - {os: ubuntu-latest-4-cores, pyv: "3.9", group: multimodal}
+          - {os: ubuntu-latest-4-cores, pyv: "3.12", group: multimodal}
+
+
     steps:
       - uses: actions/checkout@v4
 

{datachain-0.3.4 → datachain-0.3.6}/.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.5.7'
+    rev: 'v0.6.1'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.3.4/src/datachain.egg-info → datachain-0.3.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.3.4
+Version: 0.3.6
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -41,6 +41,7 @@ Requires-Dist: pydantic<3,>=2
 Requires-Dist: jmespath>=1.0
 Requires-Dist: datamodel-code-generator>=0.25
 Requires-Dist: Pillow<11,>=10.0.0
+Requires-Dist: msgpack<2,>=1.0.4
 Provides-Extra: docs
 Requires-Dist: mkdocs>=1.5.2; extra == "docs"
 Requires-Dist: mkdocs-gen-files>=0.5.0; extra == "docs"
@@ -54,7 +55,6 @@ Requires-Dist: torchvision; extra == "torch"
 Requires-Dist: transformers>=4.36.0; extra == "torch"
 Provides-Extra: remote
 Requires-Dist: lz4; extra == "remote"
-Requires-Dist: msgpack<2,>=1.0.4; extra == "remote"
 Requires-Dist: requests>=2.22.0; extra == "remote"
 Provides-Extra: vector
 Requires-Dist: usearch; extra == "vector"
@@ -87,9 +87,8 @@ Requires-Dist: numpy<2,>=1; extra == "examples"
 Requires-Dist: defusedxml; extra == "examples"
 Requires-Dist: accelerate; extra == "examples"
 Requires-Dist: unstructured[pdf]; extra == "examples"
-Requires-Dist: pdfplumber==0.11.3; extra == "examples"
+Requires-Dist: pdfplumber==0.11.4; extra == "examples"
 Requires-Dist: huggingface_hub[hf_transfer]; extra == "examples"
-Requires-Dist: nltk==3.8.1; extra == "examples"
 
 |PyPI| |Python Version| |Codecov| |Tests|
 

{datachain-0.3.4 → datachain-0.3.6}/examples/get_started/torch-loader.py

@@ -81,6 +81,6 @@ if __name__ == "__main__":
             loss.backward()
             optimizer.step()
 
-            print("[%d, %5d] loss: %.3f" % (epoch + 1, i + 1, loss.item()))
+            print(f"[{epoch + 1}, {i + 1:5d}] loss: {loss.item():.3f}")
 
     print("Finished Training")

{datachain-0.3.4 → datachain-0.3.6}/pyproject.toml

@@ -43,7 +43,8 @@ dependencies = [
   "pydantic>=2,<3",
   "jmespath>=1.0",
   "datamodel-code-generator>=0.25",
-  "Pillow>=10.0.0,<11"
+  "Pillow>=10.0.0,<11",
+  "msgpack>=1.0.4,<2"
 ]
 
 [project.optional-dependencies]
@@ -62,7 +63,6 @@ torch = [
 ]
 remote = [
   "lz4",
-  "msgpack>=1.0.4,<2",
   "requests>=2.22.0"
 ]
 vector = [
@@ -99,9 +99,8 @@ examples = [
   "defusedxml",
   "accelerate",
   "unstructured[pdf]",
-  "pdfplumber==0.11.3",
-  "huggingface_hub[hf_transfer]",
-  "nltk==3.8.1"
+  "pdfplumber==0.11.4",
+  "huggingface_hub[hf_transfer]"
 ]
 
 [project.urls]

{datachain-0.3.4 → datachain-0.3.6}/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.4 → datachain-0.3.6}/src/datachain/data_storage/sqlite.py

@@ -209,6 +209,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)

{datachain-0.3.4 → datachain-0.3.6}/src/datachain/lib/convert/flatten.py

@@ -1,34 +1,6 @@
-from datetime import datetime
-
 from pydantic import BaseModel
 
 from datachain.lib.model_store import ModelStore
-from datachain.sql.types import (
-    JSON,
-    Array,
-    Binary,
-    Boolean,
-    DateTime,
-    Float,
-    Int,
-    Int32,
-    Int64,
-    NullType,
-    String,
-)
-
-DATACHAIN_TO_TYPE = {
-    Int: int,
-    Int32: int,
-    Int64: int,
-    String: str,
-    Float: float,
-    Boolean: bool,
-    DateTime: datetime,
-    Binary: bytes,
-    Array(NullType): list,
-    JSON: dict,
-}
 
 
 def flatten(obj: BaseModel):

{datachain-0.3.4 → datachain-0.3.6}/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(
@@ -1224,14 +1247,11 @@ class DataChain(DatasetQuery):
         """
         headers, max_length = self._effective_signals_schema.get_headers_with_length()
         if flatten or max_length < 2:
-            columns = []
-            if headers:
-                columns = [".".join(filter(None, header)) for header in headers]
-            return pd.DataFrame.from_records(self.to_records(), columns=columns)
+            columns = [".".join(filter(None, header)) for header in headers]
+        else:
+            columns = pd.MultiIndex.from_tuples(map(tuple, headers))
 
-        return pd.DataFrame(
-            self.results(), columns=pd.MultiIndex.from_tuples(map(tuple, headers))
-        )
+        return pd.DataFrame.from_records(self.results(), columns=columns)
 
     def show(
         self,
@@ -1524,6 +1544,7 @@ class DataChain(DatasetQuery):
         to_insert: Optional[Union[dict, list[dict]]],
         session: Optional[Session] = None,
         in_memory: bool = False,
+        schema: Optional[dict[str, DataType]] = None,
     ) -> "DataChain":
         """Create a DataChain from the provided records. This method can be used for
         programmatically generating a chain in contrast of reading data from storages
@@ -1532,10 +1553,10 @@ class DataChain(DatasetQuery):
         Parameters:
             to_insert : records (or a single record) to insert. Each record is
                         a dictionary of signals and theirs values.
+            schema : describes chain signals and their corresponding types
 
         Example:
             ```py
-            empty = DataChain.from_records()
             single_record = DataChain.from_records(DataChain.DEFAULT_FILE_RECORD)
             ```
         """
@@ -1543,11 +1564,27 @@ class DataChain(DatasetQuery):
         catalog = session.catalog
 
         name = session.generate_temp_dataset_name()
-        columns: tuple[sqlalchemy.Column[Any], ...] = tuple(
-            sqlalchemy.Column(name, typ)
-            for name, typ in File._datachain_column_types.items()
+        signal_schema = None
+        columns: list[sqlalchemy.Column] = []
+
+        if schema:
+            signal_schema = SignalSchema(schema)
+            columns = signal_schema.db_signals(as_columns=True)  # type: ignore[assignment]
+        else:
+            columns = [
+                sqlalchemy.Column(name, typ)
+                for name, typ in File._datachain_column_types.items()
+            ]
+
+        dsr = catalog.create_dataset(
+            name,
+            columns=columns,
+            feature_schema=(
+                signal_schema.clone_without_sys_signals().serialize()
+                if signal_schema
+                else None
+            ),
         )
-        dsr = catalog.create_dataset(name, columns=columns)
 
         if isinstance(to_insert, dict):
             to_insert = [to_insert]

{datachain-0.3.4 → datachain-0.3.6}/src/datachain/lib/signal_schema.py

@@ -2,6 +2,7 @@ import copy
 from collections.abc import Iterator, Sequence
 from dataclasses import dataclass
 from datetime import datetime
+from inspect import isclass
 from typing import (
     TYPE_CHECKING,
     Annotated,
@@ -14,10 +15,10 @@ from typing import (
     get_origin,
 )
 
+import sqlalchemy as sa
 from pydantic import BaseModel, create_model
 from typing_extensions import Literal as LiteralEx
 
-from datachain.lib.convert.flatten import DATACHAIN_TO_TYPE
 from datachain.lib.convert.python_to_sql import python_to_sql
 from datachain.lib.convert.sql_to_python import sql_to_python
 from datachain.lib.convert.unflatten import unflatten_to_json_pos
@@ -26,6 +27,7 @@ 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
+from datachain.sql.types import SQLType
 
 if TYPE_CHECKING:
     from datachain.catalog import Catalog
@@ -104,12 +106,15 @@ class SignalSchema:
     def from_column_types(col_types: dict[str, Any]) -> "SignalSchema":
         signals: dict[str, DataType] = {}
         for field, col_type in col_types.items():
-            if (py_type := DATACHAIN_TO_TYPE.get(col_type, None)) is None:
+            if isinstance(col_type, SQLType):
+                signals[field] = col_type.python_type
+            elif isclass(col_type) and issubclass(col_type, SQLType):
+                signals[field] = col_type().python_type
+            else:
                 raise SignalSchemaError(
                     f"signal schema cannot be obtained for column '{field}':"
-                    f" unsupported type '{py_type}'"
+                    f" unsupported type '{col_type}'"
                 )
-            signals[field] = py_type
         return SignalSchema(signals)
 
     def serialize(self) -> dict[str, str]:
@@ -232,7 +237,7 @@ class SignalSchema:
         signals = [
             DEFAULT_DELIMITER.join(path)
             if not as_columns
-            else Column(DEFAULT_DELIMITER.join(path), python_to_sql(_type))
+            else sa.Column(DEFAULT_DELIMITER.join(path), python_to_sql(_type))
             for path, _type, has_subtree, _ in self.get_flat_tree()
             if not has_subtree
         ]

{datachain-0.3.4 → datachain-0.3.6}/src/datachain/query/dataset.py

@@ -878,17 +878,14 @@ class SQLUnion(Step):
         temp_tables.extend(self.query1.temp_table_names)
         q2 = self.query2.apply_steps().select().subquery()
         temp_tables.extend(self.query2.temp_table_names)
-        columns1, columns2 = fill_columns(q1.columns, q2.columns)
+
+        columns1, columns2 = _order_columns(q1.columns, q2.columns)
 
         def q(*columns):
             names = {c.name for c in columns}
             col1 = [c for c in columns1 if c.name in names]
             col2 = [c for c in columns2 if c.name in names]
-            res = (
-                sqlalchemy.select(*col1)
-                .select_from(q1)
-                .union_all(sqlalchemy.select(*col2).select_from(q2))
-            )
+            res = sqlalchemy.select(*col1).union_all(sqlalchemy.select(*col2))
 
             subquery = res.subquery()
             return sqlalchemy.select(*subquery.c).select_from(subquery)
@@ -1021,23 +1018,46 @@ class GroupBy(Step):
         return step_result(q, grouped_query.selected_columns)
 
 
-def fill_columns(
-    *column_iterables: Iterable[ColumnElement],
+def _validate_columns(
+    left_columns: Iterable[ColumnElement], right_columns: Iterable[ColumnElement]
+) -> set[str]:
+    left_names = {c.name for c in left_columns}
+    right_names = {c.name for c in right_columns}
+
+    if left_names == right_names:
+        return left_names
+
+    missing_right = left_names - right_names
+    missing_left = right_names - left_names
+
+    def _prepare_msg_part(missing_columns: set[str], side: str) -> str:
+        return f"{', '.join(sorted(missing_columns))} only present in {side}"
+
+    msg_parts = [
+        _prepare_msg_part(missing_columns, found_side)
+        for missing_columns, found_side in zip(
+            [
+                missing_right,
+                missing_left,
+            ],
+            ["left", "right"],
+        )
+        if missing_columns
+    ]
+    msg = f"Cannot perform union. {'. '.join(msg_parts)}"
+
+    raise ValueError(msg)
+
+
+def _order_columns(
+    left_columns: Iterable[ColumnElement], right_columns: Iterable[ColumnElement]
 ) -> list[list[ColumnElement]]:
-    column_dicts = [{c.name: c for c in columns} for columns in column_iterables]
-    combined_columns = {n: c for col_dict in column_dicts for n, c in col_dict.items()}
-
-    result: list[list[ColumnElement]] = [[] for _ in column_dicts]
-    for n in combined_columns:
-        col = next(col_dict[n] for col_dict in column_dicts if n in col_dict)
-        for col_dict, out in zip(column_dicts, result):
-            if n in col_dict:
-                out.append(col_dict[n])
-            else:
-                # Cast the NULL to ensure all columns are aware of their type
-                # Label it to ensure it's aware of its name
-                out.append(sqlalchemy.cast(sqlalchemy.null(), col.type).label(n))
-    return result
+    column_order = _validate_columns(left_columns, right_columns)
+    column_dicts = [
+        {c.name: c for c in columns} for columns in [left_columns, right_columns]
+    ]
+
+    return [[d[n] for n in column_order] for d in column_dicts]
 
 
 @attrs.define

{datachain-0.3.4 → datachain-0.3.6}/src/datachain/sql/types.py

@@ -20,6 +20,8 @@ from typing import Any, Union
 import sqlalchemy as sa
 from sqlalchemy import TypeDecorator, types
 
+from datachain.lib.data_model import StandardType
+
 _registry: dict[str, "TypeConverter"] = {}
 registry = MappingProxyType(_registry)
 
@@ -91,6 +93,10 @@ class SQLType(TypeDecorator):
     impl: type[types.TypeEngine[Any]] = types.TypeEngine
     cache_ok = True
 
+    @property
+    def python_type(self) -> StandardType:
+        raise NotImplementedError
+
     def to_dict(self) -> dict[str, Any]:
         return {"type": self.__class__.__name__}
 
@@ -103,7 +109,7 @@ class String(SQLType):
     impl = types.String
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return str
 
     def load_dialect_impl(self, dialect):
@@ -125,7 +131,7 @@ class Boolean(SQLType):
     impl = types.Boolean
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return bool
 
     def load_dialect_impl(self, dialect):
@@ -147,7 +153,7 @@ class Int(SQLType):
     impl = types.INTEGER
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return int
 
     def load_dialect_impl(self, dialect):
@@ -217,7 +223,7 @@ class Float(SQLType):
     impl = types.FLOAT
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return float
 
     def load_dialect_impl(self, dialect):
@@ -271,7 +277,7 @@ class Array(SQLType):
     impl = types.ARRAY
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return list
 
     def load_dialect_impl(self, dialect):
@@ -314,7 +320,7 @@ class JSON(SQLType):
     impl = types.JSON
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return dict
 
     def load_dialect_impl(self, dialect):
@@ -336,7 +342,7 @@ class DateTime(SQLType):
     impl = types.DATETIME
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return datetime
 
     def load_dialect_impl(self, dialect):
@@ -358,7 +364,7 @@ class Binary(SQLType):
     impl = types.BINARY
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return bytes
 
     def load_dialect_impl(self, dialect):

{datachain-0.3.4 → datachain-0.3.6/src/datachain.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.3.4
+Version: 0.3.6
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -41,6 +41,7 @@ Requires-Dist: pydantic<3,>=2
 Requires-Dist: jmespath>=1.0
 Requires-Dist: datamodel-code-generator>=0.25
 Requires-Dist: Pillow<11,>=10.0.0
+Requires-Dist: msgpack<2,>=1.0.4
 Provides-Extra: docs
 Requires-Dist: mkdocs>=1.5.2; extra == "docs"
 Requires-Dist: mkdocs-gen-files>=0.5.0; extra == "docs"
@@ -54,7 +55,6 @@ Requires-Dist: torchvision; extra == "torch"
 Requires-Dist: transformers>=4.36.0; extra == "torch"
 Provides-Extra: remote
 Requires-Dist: lz4; extra == "remote"
-Requires-Dist: msgpack<2,>=1.0.4; extra == "remote"
 Requires-Dist: requests>=2.22.0; extra == "remote"
 Provides-Extra: vector
 Requires-Dist: usearch; extra == "vector"
@@ -87,9 +87,8 @@ Requires-Dist: numpy<2,>=1; extra == "examples"
 Requires-Dist: defusedxml; extra == "examples"
 Requires-Dist: accelerate; extra == "examples"
 Requires-Dist: unstructured[pdf]; extra == "examples"
-Requires-Dist: pdfplumber==0.11.3; extra == "examples"
+Requires-Dist: pdfplumber==0.11.4; extra == "examples"
 Requires-Dist: huggingface_hub[hf_transfer]; extra == "examples"
-Requires-Dist: nltk==3.8.1; extra == "examples"
 
 |PyPI| |Python Version| |Codecov| |Tests|
 

{datachain-0.3.4 → datachain-0.3.6}/src/datachain.egg-info/requires.txt

@@ -22,6 +22,7 @@ pydantic<3,>=2
 jmespath>=1.0
 datamodel-code-generator>=0.25
 Pillow<11,>=10.0.0
+msgpack<2,>=1.0.4
 
 [:sys_platform == "win32"]
 numpy<2,>=1
@@ -48,13 +49,11 @@ numpy<2,>=1
 defusedxml
 accelerate
 unstructured[pdf]
-pdfplumber==0.11.3
+pdfplumber==0.11.4
 huggingface_hub[hf_transfer]
-nltk==3.8.1
 
 [remote]
 lz4
-msgpack<2,>=1.0.4
 requests>=2.22.0
 
 [tests]

{datachain-0.3.4 → datachain-0.3.6}/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."
+    )