datachain 0.30.5__py3-none-any.whl → 0.39.0__py3-none-any.whl

datachain/lib/signal_schema.py

@@ -1,30 +1,31 @@
 import copy
+import hashlib
+import logging
+import math
+import types
 import warnings
-from collections.abc import Iterator, Sequence
+from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
 from dataclasses import dataclass
 from datetime import datetime
 from inspect import isclass
-from typing import (  # noqa: UP035
+from typing import (
     IO,
     TYPE_CHECKING,
     Annotated,
     Any,
-    Callable,
-    Dict,
     Final,
-    List,
     Literal,
-    Mapping,
     Optional,
     Union,
     get_args,
     get_origin,
 )
 
-from pydantic import BaseModel, Field, create_model
+from pydantic import BaseModel, Field, ValidationError, create_model
 from sqlalchemy import ColumnElement
 from typing_extensions import Literal as LiteralEx
 
+from datachain import json
 from datachain.func import literal
 from datachain.func.func import Func
 from datachain.lib.convert.python_to_sql import python_to_sql
@@ -33,7 +34,7 @@ from datachain.lib.convert.unflatten import unflatten_to_json_pos
 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.lib.utils import DataChainColumnError, DataChainParamsError
 from datachain.query.schema import DEFAULT_DELIMITER, C, Column, ColumnMeta
 from datachain.sql.types import SQLType
 
@@ -41,6 +42,8 @@ if TYPE_CHECKING:
     from datachain.catalog import Catalog
 
 
+logger = logging.getLogger(__name__)
+
 NAMES_TO_TYPES = {
     "int": int,
     "str": str,
@@ -69,7 +72,7 @@ class SignalSchemaWarning(RuntimeWarning):
 
 
 class SignalResolvingError(SignalSchemaError):
-    def __init__(self, path: Optional[list[str]], msg: str):
+    def __init__(self, path: list[str] | None, msg: str):
         name = " '" + ".".join(path) + "'" if path else ""
         super().__init__(f"cannot resolve signal name{name}: {msg}")
 
@@ -79,6 +82,55 @@ class SetupError(SignalSchemaError):
         super().__init__(f"cannot setup value '{name}': {msg}")
 
 
+def generate_merge_root_mapping(
+    left_names: Iterable[str],
+    right_names: Sequence[str],
+    *,
+    extract_root: Callable[[str], str],
+    prefix: str,
+) -> dict[str, str]:
+    """Compute root renames for schema merges.
+
+    Returns a mapping from each right-side root to the target root name while
+    preserving the order in which right-side roots first appear. The mapping
+    avoids collisions with roots already present on the left side and among
+    the right-side roots themselves. When a conflict is detected, the
+    ``prefix`` string is used to derive candidate root names until a unique
+    one is found.
+    """
+
+    existing_roots = {extract_root(name) for name in left_names}
+
+    right_root_order: list[str] = []
+    right_roots: set[str] = set()
+    for name in right_names:
+        root = extract_root(name)
+        if root not in right_roots:
+            right_roots.add(root)
+            right_root_order.append(root)
+
+    used_roots = set(existing_roots)
+    root_mapping: dict[str, str] = {}
+
+    for root in right_root_order:
+        if root not in used_roots:
+            root_mapping[root] = root
+            used_roots.add(root)
+            continue
+
+        suffix = 0
+        while True:
+            base = prefix if root in prefix else f"{prefix}{root}"
+            candidate_root = base if suffix == 0 else f"{base}_{suffix}"
+            if candidate_root not in used_roots and candidate_root not in right_roots:
+                root_mapping[root] = candidate_root
+                used_roots.add(candidate_root)
+                break
+            suffix += 1
+
+    return root_mapping
+
+
 class SignalResolvingTypeError(SignalResolvingError):
     def __init__(self, method: str, field):
         super().__init__(
@@ -89,7 +141,7 @@ class SignalResolvingTypeError(SignalResolvingError):
 
 
 class SignalRemoveError(SignalSchemaError):
-    def __init__(self, path: Optional[list[str]], msg: str):
+    def __init__(self, path: list[str] | None, msg: str):
         name = " '" + ".".join(path) + "'" if path else ""
         super().__init__(f"cannot remove signal name{name}: {msg}")
 
@@ -98,8 +150,8 @@ class CustomType(BaseModel):
     schema_version: int = Field(ge=1, le=2, strict=True)
     name: str
     fields: dict[str, str]
-    bases: list[tuple[str, str, Optional[str]]]
-    hidden_fields: Optional[list[str]] = None
+    bases: list[tuple[str, str, str | None]]
+    hidden_fields: list[str] | None = None
 
     @classmethod
     def deserialize(cls, data: dict[str, Any], type_name: str) -> "CustomType":
@@ -119,8 +171,8 @@ class CustomType(BaseModel):
 
 def create_feature_model(
     name: str,
-    fields: Mapping[str, Union[type, None, tuple[type, Any]]],
-    base: Optional[type] = None,
+    fields: Mapping[str, type | tuple[type, Any] | None],
+    base: type | None = None,
 ) -> type[BaseModel]:
     """
     This gets or returns a dynamic feature model for use in restoring a model
@@ -137,7 +189,7 @@ def create_feature_model(
         **{
             field_name: anno if isinstance(anno, tuple) else (anno, None)
             for field_name, anno in fields.items()
-        },
+        },  # type: ignore[arg-type]
     )
 
 
@@ -146,12 +198,12 @@ class SignalSchema:
     values: dict[str, DataType]
     tree: dict[str, Any]
     setup_func: dict[str, Callable]
-    setup_values: Optional[dict[str, Any]]
+    setup_values: dict[str, Any] | None
 
     def __init__(
         self,
         values: dict[str, DataType],
-        setup: Optional[dict[str, Callable]] = None,
+        setup: dict[str, Callable] | None = None,
     ):
         self.values = values
         self.tree = self._build_tree(values)
@@ -190,8 +242,8 @@ class SignalSchema:
         return SignalSchema(signals)
 
     @staticmethod
-    def _get_bases(fr: type) -> list[tuple[str, str, Optional[str]]]:
-        bases: list[tuple[str, str, Optional[str]]] = []
+    def _get_bases(fr: type) -> list[tuple[str, str, str | None]]:
+        bases: list[tuple[str, str, str | None]] = []
         for base in fr.__mro__:
             model_store_name = (
                 ModelStore.get_name(base) if issubclass(base, DataModel) else None
@@ -257,6 +309,11 @@ class SignalSchema:
             signals["_custom_types"] = custom_types
         return signals
 
+    def hash(self) -> str:
+        """Create SHA hash of this schema"""
+        json_str = json.dumps(self.serialize(), sort_keys=True, separators=(",", ":"))
+        return hashlib.sha256(json_str.encode("utf-8")).hexdigest()
+
     @staticmethod
     def _split_subtypes(type_name: str) -> list[str]:
         """This splits a list of subtypes, including proper square bracket handling."""
@@ -283,7 +340,7 @@ class SignalSchema:
     @staticmethod
     def _deserialize_custom_type(
         type_name: str, custom_types: dict[str, Any]
-    ) -> Optional[type]:
+    ) -> type | None:
         """Given a type name like MyType@v1 gets a type from ModelStore or recreates
         it based on the information from the custom types dict that includes fields and
         bases."""
@@ -316,7 +373,7 @@ class SignalSchema:
         return None
 
     @staticmethod
-    def _resolve_type(type_name: str, custom_types: dict[str, Any]) -> Optional[type]:
+    def _resolve_type(type_name: str, custom_types: dict[str, Any]) -> type | None:
         """Convert a string-based type back into a python type."""
         type_name = type_name.strip()
         if not type_name:
@@ -325,7 +382,7 @@ class SignalSchema:
             return None
 
         bracket_idx = type_name.find("[")
-        subtypes: Optional[tuple[Optional[type], ...]] = None
+        subtypes: tuple[type | None, ...] | None = None
         if bracket_idx > -1:
             if bracket_idx == 0:
                 raise ValueError("Type cannot start with '['")
@@ -456,13 +513,33 @@ class SignalSchema:
                 objs.append(self.setup_values.get(name))
             elif (fr := ModelStore.to_pydantic(fr_type)) is not None:
                 j, pos = unflatten_to_json_pos(fr, row, pos)
-                objs.append(fr(**j))
+                try:
+                    obj = fr(**j)
+                except ValidationError as e:
+                    if self._all_values_none(j):
+                        logger.debug("Failed to create input for %s: %s", name, e)
+                        obj = None
+                    else:
+                        raise
+                objs.append(obj)
             else:
                 objs.append(row[pos])
                 pos += 1
         return objs
 
-    def get_file_signal(self) -> Optional[str]:
+    @staticmethod
+    def _all_values_none(value: Any) -> bool:
+        if isinstance(value, dict):
+            return all(SignalSchema._all_values_none(v) for v in value.values())
+        if isinstance(value, (list, tuple, set)):
+            return all(SignalSchema._all_values_none(v) for v in value)
+        if isinstance(value, float):
+            # NaN is used to represent NULL and NaN float values in datachain
+            # Since SQLite does not have a separate NULL type, we need to check for NaN
+            return math.isnan(value) or value is None
+        return value is None
+
+    def get_file_signal(self) -> str | None:
         for signal_name, signal_type in self.values.items():
             if (fr := ModelStore.to_pydantic(signal_type)) is not None and issubclass(
                 fr, File
@@ -472,8 +549,8 @@ class SignalSchema:
 
     def slice(
         self,
-        params: dict[str, Union[DataType, Any]],
-        setup: Optional[dict[str, Callable]] = None,
+        params: dict[str, DataType | Any],
+        setup: dict[str, Callable] | None = None,
         is_batch: bool = False,
     ) -> "SignalSchema":
         """
@@ -497,9 +574,13 @@ class SignalSchema:
             schema_origin = get_origin(schema_type)
             param_origin = get_origin(param_type)
 
-            if schema_origin is Union and type(None) in get_args(schema_type):
+            if schema_origin in (Union, types.UnionType) and type(None) in get_args(
+                schema_type
+            ):
                 schema_type = get_args(schema_type)[0]
-                if param_origin is Union and type(None) in get_args(param_type):
+                if param_origin in (Union, types.UnionType) and type(None) in get_args(
+                    param_type
+                ):
                     param_type = get_args(param_type)[0]
 
             if is_batch:
@@ -535,15 +616,90 @@ class SignalSchema:
         pos = 0
         for fr_cls in self.values.values():
             if (fr := ModelStore.to_pydantic(fr_cls)) is None:
-                res.append(row[pos])
+                value = row[pos]
                 pos += 1
+                converted = self._convert_feature_value(fr_cls, value, catalog, cache)
+                res.append(converted)
             else:
                 json, pos = unflatten_to_json_pos(fr, row, pos)  # type: ignore[union-attr]
-                obj = fr(**json)
-                SignalSchema._set_file_stream(obj, catalog, cache)
+                try:
+                    obj = fr(**json)
+                    SignalSchema._set_file_stream(obj, catalog, cache)
+                except ValidationError as e:
+                    if self._all_values_none(json):
+                        logger.debug("Failed to create feature for %s: %s", fr_cls, e)
+                        obj = None
+                    else:
+                        raise
                 res.append(obj)
         return res
 
+    def _convert_feature_value(
+        self,
+        annotation: DataType,
+        value: Any,
+        catalog: "Catalog",
+        cache: bool,
+    ) -> Any:
+        """Convert raw DB value into declared annotation if needed."""
+        if value is None:
+            return None
+
+        result = value
+        origin = get_origin(annotation)
+
+        if origin in (Union, types.UnionType):
+            non_none_args = [
+                arg for arg in get_args(annotation) if arg is not type(None)
+            ]
+            if len(non_none_args) == 1:
+                annotation = non_none_args[0]
+                origin = get_origin(annotation)
+            else:
+                return result
+
+        if ModelStore.is_pydantic(annotation):
+            if isinstance(value, annotation):
+                obj = value
+            elif isinstance(value, Mapping):
+                obj = annotation(**value)
+            else:
+                return result
+            assert isinstance(obj, BaseModel)
+            SignalSchema._set_file_stream(obj, catalog, cache)
+            result = obj
+        elif origin is list:
+            args = get_args(annotation)
+            if args and isinstance(value, (list, tuple)):
+                item_type = args[0]
+                result = [
+                    self._convert_feature_value(item_type, item, catalog, cache)
+                    if item is not None
+                    else None
+                    for item in value
+                ]
+        elif origin is dict:
+            args = get_args(annotation)
+            if len(args) == 2 and isinstance(value, dict):
+                key_type, val_type = args
+                result = {}
+                for key, val in value.items():
+                    if key_type is str:
+                        converted_key = key
+                    else:
+                        loaded_key = json.loads(key)
+                        converted_key = self._convert_feature_value(
+                            key_type, loaded_key, catalog, cache
+                        )
+                    converted_val = (
+                        self._convert_feature_value(val_type, val, catalog, cache)
+                        if val_type is not Any
+                        else val
+                    )
+                    result[converted_key] = converted_val
+
+        return result
+
     @staticmethod
     def _set_file_stream(
         obj: BaseModel, catalog: "Catalog", cache: bool = False
@@ -572,8 +728,8 @@ class SignalSchema:
         raise SignalResolvingError([col_name], "is not found")
 
     def db_signals(
-        self, name: Optional[str] = None, as_columns=False, include_hidden: bool = True
-    ) -> Union[list[str], list[Column]]:
+        self, name: str | None = None, as_columns=False, include_hidden: bool = True
+    ) -> list[str] | list[Column]:
         """
         Returns DB columns as strings or Column objects with proper types
         Optionally, it can filter results by specific object, returning only his signals
@@ -600,6 +756,35 @@ class SignalSchema:
 
         return signals  # type: ignore[return-value]
 
+    def user_signals(
+        self,
+        *,
+        include_hidden: bool = True,
+        include_sys: bool = False,
+    ) -> list[str]:
+        return [
+            ".".join(path)
+            for path, _, has_subtree, _ in self.get_flat_tree(
+                include_hidden=include_hidden, include_sys=include_sys
+            )
+            if not has_subtree
+        ]
+
+    def compare_signals(
+        self,
+        other: "SignalSchema",
+        *,
+        include_hidden: bool = True,
+        include_sys: bool = False,
+    ) -> tuple[set[str], set[str]]:
+        left = set(
+            self.user_signals(include_hidden=include_hidden, include_sys=include_sys)
+        )
+        right = set(
+            other.user_signals(include_hidden=include_hidden, include_sys=include_sys)
+        )
+        return left - right, right - left
+
     def resolve(self, *names: str) -> "SignalSchema":
         schema = {}
         for field in names:
@@ -733,12 +918,30 @@ class SignalSchema:
         right_schema: "SignalSchema",
         rname: str,
     ) -> "SignalSchema":
-        schema_right = {
-            rname + key if key in self.values else key: type_
-            for key, type_ in right_schema.values.items()
-        }
+        merged_values = dict(self.values)
+
+        right_names = list(right_schema.values.keys())
+        root_mapping = generate_merge_root_mapping(
+            self.values.keys(),
+            right_names,
+            extract_root=self._extract_root,
+            prefix=rname,
+        )
+
+        for key, type_ in right_schema.values.items():
+            root = self._extract_root(key)
+            tail = key.partition(".")[2]
+            mapped_root = root_mapping[root]
+            new_name = mapped_root if not tail else f"{mapped_root}.{tail}"
+            merged_values[new_name] = type_
 
-        return SignalSchema(self.values | schema_right)
+        return SignalSchema(merged_values)
+
+    @staticmethod
+    def _extract_root(name: str) -> str:
+        if "." in name:
+            return name.split(".", 1)[0]
+        return name
 
     def append(self, right: "SignalSchema") -> "SignalSchema":
         missing_schema = {
@@ -758,29 +961,38 @@ class SignalSchema:
         return create_model(
             name,
             __base__=(DataModel,),  # type: ignore[call-overload]
-            **fields,
+            **fields,  # type: ignore[arg-type]
         )
 
     @staticmethod
     def _build_tree(
         values: dict[str, DataType],
-    ) -> dict[str, tuple[DataType, Optional[dict]]]:
+    ) -> dict[str, tuple[DataType, dict | None]]:
         return {
             name: (val, SignalSchema._build_tree_for_type(val))
             for name, val in values.items()
         }
 
     def get_flat_tree(
-        self, include_hidden: bool = True
+        self,
+        include_hidden: bool = True,
+        include_sys: bool = True,
     ) -> Iterator[tuple[list[str], DataType, bool, int]]:
-        yield from self._get_flat_tree(self.tree, [], 0, include_hidden)
+        yield from self._get_flat_tree(self.tree, [], 0, include_hidden, include_sys)
 
     def _get_flat_tree(
-        self, tree: dict, prefix: list[str], depth: int, include_hidden: bool
+        self,
+        tree: dict,
+        prefix: list[str],
+        depth: int,
+        include_hidden: bool,
+        include_sys: bool,
     ) -> Iterator[tuple[list[str], DataType, bool, int]]:
         for name, (type_, substree) in tree.items():
             suffix = name.split(".")
             new_prefix = prefix + suffix
+            if not include_sys and new_prefix and new_prefix[0] == "sys":
+                continue
             hidden_fields = getattr(type_, "_hidden_fields", None)
             if hidden_fields and substree and not include_hidden:
                 substree = {
@@ -793,10 +1005,10 @@ class SignalSchema:
             yield new_prefix, type_, has_subtree, depth
             if substree is not None:
                 yield from self._get_flat_tree(
-                    substree, new_prefix, depth + 1, include_hidden
+                    substree, new_prefix, depth + 1, include_hidden, include_sys
                 )
 
-    def print_tree(self, indent: int = 2, start_at: int = 0, file: Optional[IO] = None):
+    def print_tree(self, indent: int = 2, start_at: int = 0, file: IO | None = None):
         for path, type_, _, depth in self.get_flat_tree():
             total_indent = start_at + depth * indent
             col_name = " " * total_indent + path[-1]
@@ -826,7 +1038,28 @@ class SignalSchema:
         ], max_length
 
     def __or__(self, other):
-        return self.__class__(self.values | other.values)
+        new_values = dict(self.values)
+
+        for name, new_type in other.values.items():
+            if name in new_values:
+                current_type = new_values[name]
+                if current_type != new_type:
+                    raise DataChainColumnError(
+                        name,
+                        "signal already exists with a different type",
+                    )
+                continue
+
+            root = self._extract_root(name)
+            if any(self._extract_root(existing) == root for existing in new_values):
+                raise DataChainColumnError(
+                    name,
+                    "signal root already exists in schema",
+                )
+
+            new_values[name] = new_type
+
+        return self.__class__(new_values)
 
     def __contains__(self, name: str):
         return name in self.values
@@ -835,15 +1068,20 @@ class SignalSchema:
         return self.values.pop(name)
 
     @staticmethod
-    def _type_to_str(type_: Optional[type], subtypes: Optional[list] = None) -> str:  # noqa: PLR0911
+    def _type_to_str(type_: type | None, subtypes: list | None = None) -> str:  # noqa: C901, PLR0911
         """Convert a type to a string-based representation."""
         if type_ is None:
             return "NoneType"
 
         origin = get_origin(type_)
 
-        if origin == Union:
+        if origin in (Union, types.UnionType):
             args = get_args(type_)
+            if len(args) == 2 and type(None) in args:
+                # This is an Optional type.
+                non_none_type = args[0] if args[1] is type(None) else args[1]
+                type_str = SignalSchema._type_to_str(non_none_type, subtypes)
+                return f"Optional[{type_str}]"
             formatted_types = ", ".join(
                 SignalSchema._type_to_str(arg, subtypes) for arg in args
             )
@@ -852,21 +1090,21 @@ class SignalSchema:
             args = get_args(type_)
             type_str = SignalSchema._type_to_str(args[0], subtypes)
             return f"Optional[{type_str}]"
-        if origin in (list, List):  # noqa: UP006
+        if origin is list:
             args = get_args(type_)
+            if len(args) == 0:
+                return "list"
             type_str = SignalSchema._type_to_str(args[0], subtypes)
             return f"list[{type_str}]"
-        if origin in (dict, Dict):  # noqa: UP006
+        if origin is dict:
             args = get_args(type_)
-            type_str = (
-                SignalSchema._type_to_str(args[0], subtypes) if len(args) > 0 else ""
-            )
-            vals = (
-                f", {SignalSchema._type_to_str(args[1], subtypes)}"
-                if len(args) > 1
-                else ""
-            )
-            return f"dict[{type_str}{vals}]"
+            if len(args) == 0:
+                return "dict"
+            key_type = SignalSchema._type_to_str(args[0], subtypes)
+            if len(args) == 1:
+                return f"dict[{key_type}, Any]"
+            val_type = SignalSchema._type_to_str(args[1], subtypes)
+            return f"dict[{key_type}, {val_type}]"
         if origin == Annotated:
             args = get_args(type_)
             return SignalSchema._type_to_str(args[0], subtypes)
@@ -880,7 +1118,7 @@ class SignalSchema:
             # Include this type in the list of all subtypes, if requested.
             subtypes.append(type_)
         if not hasattr(type_, "__name__"):
-            # This can happen for some third-party or custom types, mostly on Python 3.9
+            # This can happen for some third-party or custom types
             warnings.warn(
                 f"Unable to determine name of type '{type_}'.",
                 SignalSchemaWarning,
@@ -895,7 +1133,7 @@ class SignalSchema:
     @staticmethod
     def _build_tree_for_type(
         model: DataType,
-    ) -> Optional[dict[str, tuple[DataType, Optional[dict]]]]:
+    ) -> dict[str, tuple[DataType, dict | None]] | None:
         if (fr := ModelStore.to_pydantic(model)) is not None:
             return SignalSchema._build_tree_for_model(fr)
         return None
@@ -903,8 +1141,8 @@ class SignalSchema:
     @staticmethod
     def _build_tree_for_model(
         model: type[BaseModel],
-    ) -> Optional[dict[str, tuple[DataType, Optional[dict]]]]:
-        res: dict[str, tuple[DataType, Optional[dict]]] = {}
+    ) -> dict[str, tuple[DataType, dict | None]] | None:
+        res: dict[str, tuple[DataType, dict | None]] = {}
 
         for name, f_info in model.model_fields.items():
             anno = f_info.annotation
@@ -953,7 +1191,7 @@ class SignalSchema:
         schema: dict[str, Any] = {}
         schema_custom_types: dict[str, CustomType] = {}
 
-        data_model_bases: Optional[list[tuple[str, str, Optional[str]]]] = None
+        data_model_bases: list[tuple[str, str, str | None]] | None = None
 
         signal_partials: dict[str, str] = {}
         partial_versions: dict[str, int] = {}

datachain/lib/text.py

@@ -1,16 +1,17 @@
-from typing import Any, Callable, Optional, Union
+from collections.abc import Callable
+from typing import Any
 
 import torch
 from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
 
 def convert_text(
-    text: Union[str, list[str]],
-    tokenizer: Optional[Callable] = None,
-    tokenizer_kwargs: Optional[dict[str, Any]] = None,
-    encoder: Optional[Callable] = None,
-    device: Optional[Union[str, torch.device]] = None,
-) -> Union[str, list[str], torch.Tensor]:
+    text: str | list[str],
+    tokenizer: Callable | None = None,
+    tokenizer_kwargs: dict[str, Any] | None = None,
+    encoder: Callable | None = None,
+    device: str | torch.device | None = None,
+) -> str | list[str] | torch.Tensor:
     """
     Tokenize and otherwise transform text.