datachain 0.30.3__tar.gz → 0.30.5__tar.gz

{datachain-0.30.3 → datachain-0.30.5}/.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.12.9'
+    rev: 'v0.12.10'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.30.3 → datachain-0.30.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.30.3
+Version: 0.30.5
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -22,6 +22,7 @@ Requires-Dist: tomlkit
 Requires-Dist: tqdm
 Requires-Dist: numpy<3,>=1
 Requires-Dist: pandas>=2.0.0
+Requires-Dist: ujson>=5.10.0
 Requires-Dist: packaging
 Requires-Dist: pyarrow
 Requires-Dist: typing-extensions
@@ -38,7 +39,6 @@ Requires-Dist: shtab<2,>=1.3.4
 Requires-Dist: sqlalchemy>=2
 Requires-Dist: multiprocess==0.70.16
 Requires-Dist: cloudpickle
-Requires-Dist: orjson>=3.10.5
 Requires-Dist: pydantic
 Requires-Dist: jmespath>=1.0
 Requires-Dist: datamodel-code-generator>=0.25
@@ -92,7 +92,7 @@ Requires-Dist: pytest-mock>=3.12.0; extra == "tests"
 Requires-Dist: pytest-servers[all]>=0.5.9; extra == "tests"
 Requires-Dist: pytest-benchmark[histogram]; extra == "tests"
 Requires-Dist: pytest-xdist>=3.3.1; extra == "tests"
-Requires-Dist: pytest-env>=1.1.0; extra == "tests"
+Requires-Dist: pytest-dotenv; extra == "tests"
 Requires-Dist: virtualenv; extra == "tests"
 Requires-Dist: dulwich; extra == "tests"
 Requires-Dist: hypothesis; extra == "tests"

datachain-0.30.5/examples/get_started/nested_datamodel.py

@@ -0,0 +1,70 @@
+"""Example: Nested DataModels with parallel execution.
+
+Demonstrates mapping a function that returns a nested DataModel (a DataModel
+containing other DataModels).
+
+The example keeps things minimal: we persist a tiny dataset, run a parallel map
+that returns a nested DataModel, and display the result.
+"""
+
+from typing import Optional
+
+from pydantic import Field
+
+import datachain as dc
+
+
+class Metric(dc.DataModel):
+    """Represents a single computed metric with quality metadata."""
+
+    value: Optional[float] = Field(default=None, description="Computed metric value")
+    confidence: Optional[float] = Field(
+        default=None, description="Confidence / quality score"
+    )
+    status: Optional[str] = Field(default=None, description="Processing status label")
+    metric_error: Optional[str] = Field(
+        default=None, description="Error message if metric computation failed"
+    )
+
+
+class SampleMetrics(dc.DataModel):
+    """Container for two illustrative nested metrics.
+
+    Each sub-field is its own DataModel instance to demonstrate nested schemas
+    """
+
+    metric_primary: Metric = Field(
+        default_factory=lambda: Metric(), description="Primary metric"
+    )
+    metric_secondary: Metric = Field(
+        default_factory=lambda: Metric(), description="Secondary metric"
+    )
+
+
+def generate_sample_metrics() -> SampleMetrics:
+    """Synthesize a pair of metrics.
+
+    In real scenarios you'd compute these values; here we just return constants
+    to keep the example deterministic.
+    """
+
+    return SampleMetrics(
+        metric_primary=Metric(value=50.0, confidence=0.95, status="ok"),
+    )
+
+
+def main():
+    (
+        dc.read_values(record_id=[1, 2])
+        .settings(parallel=2)  # Keep it parallel to test serialization
+        .map(metrics=generate_sample_metrics)
+        .save("nested_datamodel")
+    )
+
+    dc.read_dataset("nested_datamodel").show()
+
+    print(dc.read_dataset("nested_datamodel").to_values("metrics"))
+
+
+if __name__ == "__main__":
+    main()

{datachain-0.30.3 → datachain-0.30.5}/pyproject.toml

@@ -26,6 +26,7 @@ dependencies = [
   "tqdm",
   "numpy>=1,<3",
   "pandas>=2.0.0",
+  "ujson>=5.10.0",
   "packaging",
   "pyarrow",
   "typing-extensions",
@@ -42,7 +43,6 @@ dependencies = [
   "sqlalchemy>=2",
   "multiprocess==0.70.16",
   "cloudpickle",
-  "orjson>=3.10.5",
   "pydantic",
   "jmespath>=1.0",
   "datamodel-code-generator>=0.25",
@@ -108,7 +108,7 @@ tests = [
   "pytest-servers[all]>=0.5.9",
   "pytest-benchmark[histogram]",
   "pytest-xdist>=3.3.1",
-  "pytest-env>=1.1.0",
+  "pytest-dotenv",
   "virtualenv",
   "dulwich",
   "hypothesis",

{datachain-0.30.3 → datachain-0.30.5}/src/datachain/__init__.py

@@ -6,6 +6,7 @@ from datachain.lib.dc import (
     Sys,
     datasets,
     delete_dataset,
+    is_studio,
     listings,
     move_dataset,
     read_csv,
@@ -74,6 +75,7 @@ __all__ = [
     "datasets",
     "delete_dataset",
     "is_chain_type",
+    "is_studio",
     "listings",
     "metrics",
     "move_dataset",

{datachain-0.30.3 → datachain-0.30.5}/src/datachain/catalog/__init__.py

@@ -3,6 +3,7 @@ from .catalog import (
     QUERY_SCRIPT_CANCELED_EXIT_CODE,
     QUERY_SCRIPT_INVALID_LAST_STATEMENT_EXIT_CODE,
     Catalog,
+    is_namespace_local,
 )
 from .loader import get_catalog
 
@@ -12,4 +13,5 @@ __all__ = [
     "QUERY_SCRIPT_INVALID_LAST_STATEMENT_EXIT_CODE",
     "Catalog",
     "get_catalog",
+    "is_namespace_local",
 ]

{datachain-0.30.3 → datachain-0.30.5}/src/datachain/catalog/catalog.py

@@ -113,6 +113,11 @@ else:
     SIGINT = signal.SIGINT
 
 
+def is_namespace_local(namespace_name) -> bool:
+    """Checks if namespace is from local environment, i.e. is `local`"""
+    return namespace_name == "local"
+
+
 def shutdown_process(
     proc: subprocess.Popen,
     interrupt_timeout: Optional[int] = None,
@@ -1121,6 +1126,8 @@ class Catalog:
         pull_dataset: bool = False,
         update: bool = False,
     ) -> DatasetRecord:
+        from datachain.lib.dc.utils import is_studio
+
         # Intentionally ignore update flag is version is provided. Here only exact
         # version can be provided and update then doesn't make sense.
         # It corresponds to a query like this for example:
@@ -1129,7 +1136,12 @@ class Catalog:
         if version:
             update = False
 
-        if self.metastore.is_local_dataset(namespace_name) or not update:
+        # we don't do Studio fallback is script is already ran in Studio, or if we try
+        # to fetch dataset with local namespace as that one cannot
+        # exist in Studio in the first place
+        no_fallback = is_studio() or is_namespace_local(namespace_name)
+
+        if no_fallback or not update:
             try:
                 ds = self.get_dataset(
                     name,
@@ -1141,7 +1153,7 @@ class Catalog:
             except (NamespaceNotFoundError, ProjectNotFoundError, DatasetNotFoundError):
                 pass
 
-        if self.metastore.is_local_dataset(namespace_name):
+        if no_fallback:
             raise DatasetNotFoundError(
                 f"Dataset {name}"
                 + (f" version {version} " if version else " ")

{datachain-0.30.3 → datachain-0.30.5}/src/datachain/catalog/loader.py

@@ -127,7 +127,8 @@ def get_udf_distributor_class() -> Optional[type["AbstractUDFDistributor"]]:
 
 
 def get_catalog(
-    client_config: Optional[dict[str, Any]] = None, in_memory: bool = False
+    client_config: Optional[dict[str, Any]] = None,
+    in_memory: bool = False,
 ) -> "Catalog":
     """
     Function that creates Catalog instance with appropriate metastore
@@ -142,8 +143,9 @@ def get_catalog(
     """
     from datachain.catalog import Catalog
 
+    metastore = get_metastore(in_memory=in_memory)
     return Catalog(
-        metastore=get_metastore(in_memory=in_memory),
+        metastore=metastore,
         warehouse=get_warehouse(in_memory=in_memory),
         client_config=client_config,
         in_memory=in_memory,

{datachain-0.30.3 → datachain-0.30.5}/src/datachain/cli/__init__.py

@@ -6,6 +6,7 @@ from multiprocessing import freeze_support
 from typing import Optional
 
 from datachain.cli.utils import get_logging_level
+from datachain.error import DataChainError as DataChainError
 
 from .commands import (
     clear_cache,

{datachain-0.30.3 → datachain-0.30.5}/src/datachain/cli/commands/datasets.py

@@ -6,6 +6,7 @@ from tabulate import tabulate
 if TYPE_CHECKING:
     from datachain.catalog import Catalog
 
+from datachain.catalog import is_namespace_local
 from datachain.cli.utils import determine_flavors
 from datachain.config import Config
 from datachain.error import DataChainError, DatasetNotFoundError
@@ -138,15 +139,18 @@ def rm_dataset(
 ):
     namespace_name, project_name, name = catalog.get_full_dataset_name(name)
 
-    if not catalog.metastore.is_local_dataset(namespace_name) and studio:
+    if studio:
+        # removing Studio dataset from CLI
         from datachain.studio import remove_studio_dataset
 
-        token = Config().read().get("studio", {}).get("token")
-        if not token:
+        if Config().read().get("studio", {}).get("token"):
+            remove_studio_dataset(
+                team, name, namespace_name, project_name, version, force
+            )
+        else:
             raise DataChainError(
                 "Not logged in to Studio. Log in with 'datachain auth login'."
             )
-        remove_studio_dataset(team, name, namespace_name, project_name, version, force)
     else:
         try:
             project = catalog.metastore.get_project(project_name, namespace_name)
@@ -163,9 +167,11 @@ def edit_dataset(
     attrs: Optional[list[str]] = None,
     team: Optional[str] = None,
 ):
+    from datachain.lib.dc.utils import is_studio
+
     namespace_name, project_name, name = catalog.get_full_dataset_name(name)
 
-    if catalog.metastore.is_local_dataset(namespace_name):
+    if is_studio() or is_namespace_local(namespace_name):
         try:
             catalog.edit_dataset(
                 name, catalog.metastore.default_project, new_name, description, attrs
@@ -175,11 +181,11 @@ def edit_dataset(
     else:
         from datachain.studio import edit_studio_dataset
 
-        token = Config().read().get("studio", {}).get("token")
-        if not token:
+        if Config().read().get("studio", {}).get("token"):
+            edit_studio_dataset(
+                team, name, namespace_name, project_name, new_name, description, attrs
+            )
+        else:
             raise DataChainError(
                 "Not logged in to Studio. Log in with 'datachain auth login'."
             )
-        edit_studio_dataset(
-            team, name, namespace_name, project_name, new_name, description, attrs
-        )

{datachain-0.30.3 → datachain-0.30.5}/src/datachain/data_storage/metastore.py

@@ -145,23 +145,6 @@ class AbstractMetastore(ABC, Serializable):
     def list_namespaces(self, conn=None) -> list[Namespace]:
         """Gets a list of all namespaces"""
 
-    @property
-    @abstractmethod
-    def is_studio(self) -> bool:
-        """Returns True if this code is ran in Studio"""
-
-    def is_local_dataset(self, dataset_namespace: str) -> bool:
-        """
-        Returns True if this is local dataset i.e. not pulled from Studio but
-        created locally. This is False if we ran code in CLI mode but using dataset
-        names that are present in Studio.
-        """
-        return self.is_studio or dataset_namespace == Namespace.default()
-
-    @property
-    def namespace_allowed_to_create(self):
-        return self.is_studio
-
     #
     # Projects
     #
@@ -215,10 +198,6 @@ class AbstractMetastore(ABC, Serializable):
     def list_projects(self, namespace_id: Optional[int], conn=None) -> list[Project]:
         """Gets list of projects in some namespace or in general (in all namespaces)"""
 
-    @property
-    def project_allowed_to_create(self):
-        return self.is_studio
-
     #
     # Datasets
     #

{datachain-0.30.3 → datachain-0.30.5}/src/datachain/data_storage/sqlite.py

@@ -542,10 +542,6 @@ class SQLiteMetastore(AbstractDBMetastore):
     def _jobs_insert(self) -> "Insert":
         return sqlite.insert(self._jobs)
 
-    @property
-    def is_studio(self) -> bool:
-        return False
-
     #
     # Namespaces
     #

{datachain-0.30.3 → datachain-0.30.5}/src/datachain/data_storage/warehouse.py

@@ -1,5 +1,4 @@
 import glob
-import json
 import logging
 import posixpath
 import random
@@ -11,6 +10,7 @@ from urllib.parse import urlparse
 
 import attrs
 import sqlalchemy as sa
+import ujson as json
 from sqlalchemy.sql.expression import true
 
 from datachain.client import Client
@@ -122,7 +122,7 @@ class AbstractWarehouse(ABC, Serializable):
                 if value_type is str:
                     return val
                 if value_type in (dict, list):
-                    return json.dumps(val)
+                    return json.dumps(val, ensure_ascii=False)
                 raise ValueError(
                     f"Cannot convert value {val!r} with type {value_type} to JSON"
                 )

{datachain-0.30.3 → datachain-0.30.5}/src/datachain/lib/arrow.py

@@ -2,8 +2,8 @@ from collections.abc import Sequence
 from itertools import islice
 from typing import TYPE_CHECKING, Any, Optional
 
-import orjson
 import pyarrow as pa
+import ujson as json
 from pyarrow._csv import ParseOptions
 from pyarrow.dataset import CsvFileFormat, dataset
 from tqdm.auto import tqdm
@@ -269,7 +269,7 @@ def _get_hf_schema(
 def _get_datachain_schema(schema: "pa.Schema") -> Optional[SignalSchema]:
     """Return a restored SignalSchema from parquet metadata, if any is found."""
     if schema.metadata and DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY in schema.metadata:
-        serialized_signal_schema = orjson.loads(
+        serialized_signal_schema = json.loads(
             schema.metadata[DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY]
         )
         return SignalSchema.deserialize(serialized_signal_schema)

{datachain-0.30.3 → datachain-0.30.5}/src/datachain/lib/dc/__init__.py

@@ -9,7 +9,7 @@ from .pandas import read_pandas
 from .parquet import read_parquet
 from .records import read_records
 from .storage import read_storage
-from .utils import DatasetMergeError, DatasetPrepareError, Sys
+from .utils import DatasetMergeError, DatasetPrepareError, Sys, is_studio
 from .values import read_values
 
 __all__ = [
@@ -21,6 +21,7 @@ __all__ = [
     "Sys",
     "datasets",
     "delete_dataset",
+    "is_studio",
     "listings",
     "move_dataset",
     "read_csv",

{datachain-0.30.3 → datachain-0.30.5}/src/datachain/lib/dc/datachain.py

@@ -19,8 +19,8 @@ from typing import (
     overload,
 )
 
-import orjson
 import sqlalchemy
+import ujson as json
 from pydantic import BaseModel
 from sqlalchemy.sql.elements import ColumnElement
 from tqdm import tqdm
@@ -67,6 +67,7 @@ from .utils import (
     Sys,
     _get_merge_error_str,
     _validate_merge_on,
+    is_studio,
     resolve_columns,
 )
 
@@ -461,8 +462,6 @@ class DataChain:
         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
@@ -609,7 +608,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
@@ -1184,17 +1183,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()
@@ -1221,13 +1216,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`
         )
         ```
         """
@@ -1262,8 +1264,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
@@ -2123,9 +2127,9 @@ class DataChain:
             fsspec_fs = client.create_fs(**fs_kwargs)
 
         _partition_cols = list(partition_cols) if partition_cols else None
-        signal_schema_metadata = orjson.dumps(
-            self._effective_signals_schema.serialize()
-        )
+        signal_schema_metadata = json.dumps(
+            self._effective_signals_schema.serialize(), ensure_ascii=False
+        ).encode("utf-8")
 
         column_names, column_chunks = self.to_columnar_data_with_names(chunk_size)
 
@@ -2272,7 +2276,11 @@ class DataChain:
                         f.write(b"\n")
                 else:
                     is_first = False
-                f.write(orjson.dumps(row_to_nested_dict(headers, row)))
+                f.write(
+                    json.dumps(
+                        row_to_nested_dict(headers, row), ensure_ascii=False
+                    ).encode("utf-8")
+                )
             if include_outer_list:
                 # This makes the file JSON instead of JSON lines.
                 f.write(b"\n]\n")

{datachain-0.30.3 → datachain-0.30.5}/src/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
         )
@@ -418,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-0.30.3 → datachain-0.30.5}/src/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-0.30.3 → datachain-0.30.5}/src/datachain/lib/model_store.py

@@ -89,3 +89,15 @@ class ModelStore:
             and ModelStore.is_pydantic(parent_type)
             and "@" in ModelStore.get_name(parent_type)
         )
+
+    @classmethod
+    def rebuild_all(cls) -> None:
+        """Ensure pydantic schemas are (re)built for all registered models.
+
+        Uses ``force=True`` to avoid subtle cases where a deserialized class
+        (e.g. from by-value cloudpickle in workers) reports built state but
+        nested model field schemas aren't fully resolved yet.
+        """
+        for versions in cls.store.values():
+            for model in versions.values():
+                model.model_rebuild(force=True)

{datachain-0.30.3 → datachain-0.30.5}/src/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-0.30.3 → datachain-0.30.5}/src/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)