datachain 0.30.3__tar.gz → 0.30.4__tar.gz

{datachain-0.30.3 → datachain-0.30.4}/.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.4}/PKG-INFO

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

{datachain-0.30.3 → datachain-0.30.4}/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.4}/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.4}/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.4}/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.4}/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.4}/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.4}/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.4}/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.4}/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.4}/src/datachain/lib/dc/datachain.py

@@ -67,6 +67,7 @@ from .utils import (
     Sys,
     _get_merge_error_str,
     _validate_merge_on,
+    is_studio,
     resolve_columns,
 )
 
@@ -609,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
@@ -1184,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()
@@ -1221,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`
         )
         ```
         """
@@ -1262,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

{datachain-0.30.3 → datachain-0.30.4}/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.4}/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.4}/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.4}/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)

{datachain-0.30.3 → datachain-0.30.4}/src/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-0.30.3 → datachain-0.30.4}/src/datachain/query/dataset.py

@@ -10,7 +10,6 @@ from abc import ABC, abstractmethod
 from collections.abc import Generator, Iterable, Iterator, Sequence
 from copy import copy
 from functools import wraps
-from secrets import token_hex
 from types import GeneratorType
 from typing import (
     TYPE_CHECKING,
@@ -29,7 +28,7 @@ from attrs import frozen
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback, TqdmCallback
 from sqlalchemy import Column
 from sqlalchemy.sql import func as f
-from sqlalchemy.sql.elements import ColumnClause, ColumnElement
+from sqlalchemy.sql.elements import ColumnClause, ColumnElement, Label
 from sqlalchemy.sql.expression import label
 from sqlalchemy.sql.schema import TableClause
 from sqlalchemy.sql.selectable import Select
@@ -46,6 +45,7 @@ from datachain.dataset import DatasetDependency, DatasetStatus, RowDict
 from datachain.error import DatasetNotFoundError, QueryScriptCancelError
 from datachain.func.base import Function
 from datachain.lib.listing import is_listing_dataset, listing_dataset_expired
+from datachain.lib.signal_schema import SignalSchema
 from datachain.lib.udf import UDFAdapter, _get_cache
 from datachain.progress import CombinedDownloadCallback, TqdmCombinedDownloadCallback
 from datachain.project import Project
@@ -795,28 +795,32 @@ class SQLSelectExcept(SQLClause):
 
 @frozen
 class SQLMutate(SQLClause):
-    args: tuple[Union[Function, ColumnElement], ...]
+    args: tuple[Label, ...]
+    new_schema: SignalSchema
 
     def apply_sql_clause(self, query: Select) -> Select:
         original_subquery = query.subquery()
-        args = [
-            original_subquery.c[str(c)] if isinstance(c, (str, C)) else c
-            for c in self.parse_cols(self.args)
-        ]
-        to_mutate = {c.name for c in args}
+        to_mutate = {c.name for c in self.args}
 
-        prefix = f"mutate{token_hex(8)}_"
-        cols = [
-            c.label(prefix + c.name) if c.name in to_mutate else c
+        # Drop the original versions to avoid name collisions, exclude renamed
+        # columns. Always keep system columns (sys__*) if they exist in original query
+        new_schema_columns = set(self.new_schema.db_signals())
+        base_cols = [
+            c
             for c in original_subquery.c
+            if c.name not in to_mutate
+            and (c.name in new_schema_columns or c.name.startswith("sys__"))
         ]
-        # this is needed for new column to be used in clauses
-        # like ORDER BY, otherwise new column is not recognized
-        subquery = (
-            sqlalchemy.select(*cols, *args).select_from(original_subquery).subquery()
+
+        # Create intermediate subquery to properly handle window functions
+        intermediate_query = sqlalchemy.select(*base_cols, *self.args).select_from(
+            original_subquery
         )
+        intermediate_subquery = intermediate_query.subquery()
 
-        return sqlalchemy.select(*subquery.c).select_from(subquery)
+        return sqlalchemy.select(*intermediate_subquery.c).select_from(
+            intermediate_subquery
+        )
 
 
 @frozen
@@ -1470,7 +1474,7 @@ class DatasetQuery:
         return query
 
     @detach
-    def mutate(self, *args, **kwargs) -> "Self":
+    def mutate(self, *args, new_schema, **kwargs) -> "Self":
         """
         Add new columns to this query.
 
@@ -1482,7 +1486,7 @@ class DatasetQuery:
         """
         query_args = [v.label(k) for k, v in dict(args, **kwargs).items()]
         query = self.clone()
-        query.steps.append(SQLMutate((*query_args,)))
+        query.steps.append(SQLMutate((*query_args,), new_schema))
         return query
 
     @detach

{datachain-0.30.3 → datachain-0.30.4}/src/datachain/query/schema.py

@@ -36,6 +36,10 @@ class ColumnMeta(type):
     def __getattr__(cls, name: str):
         return cls(ColumnMeta.to_db_name(name))
 
+    @staticmethod
+    def is_nested(name: str) -> bool:
+        return DEFAULT_DELIMITER in name
+
 
 class Column(sa.ColumnClause, metaclass=ColumnMeta):
     inherit_cache: Optional[bool] = True

{datachain-0.30.3 → datachain-0.30.4}/src/datachain/utils.py

@@ -531,3 +531,10 @@ def safe_closing(thing: T) -> Iterator[T]:
     finally:
         if hasattr(thing, "close"):
             thing.close()
+
+
+def getenv_bool(name: str, default: bool = False) -> bool:
+    val = os.getenv(name)
+    if val is None:
+        return default
+    return val.lower() in ("1", "true", "yes", "on")

{datachain-0.30.3 → datachain-0.30.4}/src/datachain.egg-info/PKG-INFO

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

{datachain-0.30.3 → datachain-0.30.4}/src/datachain.egg-info/SOURCES.txt

@@ -307,6 +307,7 @@ tests/func/test_ls.py
 tests/func/test_meta_formats.py
 tests/func/test_metastore.py
 tests/func/test_metrics.py
+tests/func/test_mutate.py
 tests/func/test_pull.py
 tests/func/test_pytorch.py
 tests/func/test_query.py

{datachain-0.30.3 → datachain-0.30.4}/tests/conftest.py

@@ -8,7 +8,6 @@ from datetime import datetime
 from pathlib import PosixPath
 from time import sleep
 from typing import NamedTuple
-from unittest.mock import PropertyMock, patch
 
 import attrs
 import pytest
@@ -21,7 +20,6 @@ from datachain.catalog import Catalog
 from datachain.catalog.loader import get_metastore, get_warehouse
 from datachain.cli.utils import CommaSeparatedArgs
 from datachain.config import Config, ConfigLevel
-from datachain.data_storage.metastore import AbstractMetastore
 from datachain.data_storage.sqlite import (
     SQLiteDatabaseEngine,
     SQLiteMetastore,
@@ -543,42 +541,16 @@ def cloud_test_catalog_tmpfile(
 
 
 @pytest.fixture
-def allow_create_project():
+def is_studio():
     return True
 
 
-@pytest.fixture
-def allow_create_namespace():
-    return True
-
-
-@pytest.fixture(autouse=True)
-def mock_allowed_to_create_project(allow_create_project):
-    if not allow_create_project:
-        yield
-    else:
-        with patch.object(
-            AbstractMetastore, "project_allowed_to_create", new_callable=PropertyMock
-        ) as mock_metastore:
-            mock_metastore.return_value = True
-            yield
-
-
 @pytest.fixture(autouse=True)
-def mock_allowed_to_create_namespace(allow_create_namespace):
-    if not allow_create_namespace:
+def mock_is_studio(monkeypatch, is_studio):
+    if not is_studio:
         yield
     else:
-        with patch.object(
-            AbstractMetastore, "namespace_allowed_to_create", new_callable=PropertyMock
-        ) as mock_metastore:
-            mock_metastore.return_value = True
-            yield
-
-
-@pytest.fixture
-def mock_is_local_dataset():
-    with patch.object(AbstractMetastore, "is_local_dataset", return_value=True):
+        monkeypatch.setenv("DATACHAIN_IS_STUDIO", True)
         yield