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

datachain/lib/signal_schema.py

@@ -1,30 +1,32 @@
 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
 from datachain.lib.convert.sql_to_python import sql_to_python
@@ -32,14 +34,16 @@ 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.query.schema import DEFAULT_DELIMITER, Column
+from datachain.lib.utils import DataChainColumnError, DataChainParamsError
+from datachain.query.schema import DEFAULT_DELIMITER, C, Column, ColumnMeta
 from datachain.sql.types import SQLType
 
 if TYPE_CHECKING:
     from datachain.catalog import Catalog
 
 
+logger = logging.getLogger(__name__)
+
 NAMES_TO_TYPES = {
     "int": int,
     "str": str,
@@ -68,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}")
 
@@ -78,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__(
@@ -87,12 +140,18 @@ class SignalResolvingTypeError(SignalResolvingError):
         )
 
 
+class SignalRemoveError(SignalSchemaError):
+    def __init__(self, path: list[str] | None, msg: str):
+        name = " '" + ".".join(path) + "'" if path else ""
+        super().__init__(f"cannot remove signal name{name}: {msg}")
+
+
 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":
@@ -112,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
@@ -130,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]
     )
 
 
@@ -139,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)
@@ -183,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
@@ -250,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."""
@@ -276,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."""
@@ -309,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:
@@ -318,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 '['")
@@ -439,35 +503,54 @@ class SignalSchema:
                 res[db_name] = python_to_sql(type_)
         return res
 
-    def row_to_objs(self, row: Sequence[Any]) -> list[DataValue]:
+    def row_to_objs(self, row: Sequence[Any]) -> list[Any]:
         self._init_setup_values()
 
-        objs: list[DataValue] = []
+        objs: list[Any] = []
         pos = 0
         for name, fr_type in self.values.items():
-            if self.setup_values and (val := self.setup_values.get(name, None)):
-                objs.append(val)
+            if self.setup_values and name in self.setup_values:
+                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 contains_file(self) -> bool:
-        for type_ in self.values.values():
-            if (fr := ModelStore.to_pydantic(type_)) is not None and issubclass(
+    @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
             ):
-                return True
-
-        return False
+                return signal_name
+        return None
 
     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":
         """
@@ -491,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:
@@ -529,22 +616,97 @@ 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
     ) -> None:
         if isinstance(obj, File):
             obj._set_stream(catalog, caching_enabled=cache)
-        for field, finfo in obj.model_fields.items():
+        for field, finfo in type(obj).model_fields.items():
             if ModelStore.is_pydantic(finfo.annotation):
                 SignalSchema._set_file_stream(getattr(obj, field), catalog, cache)
 
@@ -566,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
@@ -583,6 +745,9 @@ class SignalSchema:
         ]
 
         if name:
+            if "." in name:
+                name = ColumnMeta.to_db_name(name)
+
             signals = [
                 s
                 for s in signals
@@ -591,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:
@@ -601,37 +795,60 @@ class SignalSchema:
         return SignalSchema(schema)
 
     def _find_in_tree(self, path: list[str]) -> DataType:
+        if val := self.tree.get(".".join(path)):
+            # If the path is a single string, we can directly access it
+            # without traversing the tree.
+            return val[0]
+
         curr_tree = self.tree
         curr_type = None
         i = 0
         while curr_tree is not None and i < len(path):
             if val := curr_tree.get(path[i]):
                 curr_type, curr_tree = val
-            elif i == 0 and len(path) > 1 and (val := curr_tree.get(".".join(path))):
-                curr_type, curr_tree = val
-                break
             else:
                 curr_type = None
+                break
             i += 1
 
-        if curr_type is None:
+        if curr_type is None or i < len(path):
+            # If we reached the end of the path and didn't find a type,
+            # or if we didn't traverse the entire path, raise an error.
             raise SignalResolvingError(path, "is not found")
 
         return curr_type
 
+    def group_by(
+        self, partition_by: Sequence[str], new_column: Sequence[Column]
+    ) -> "SignalSchema":
+        orig_schema = SignalSchema(copy.deepcopy(self.values))
+        schema = orig_schema.to_partial(*partition_by)
+
+        vals = {c.name: sql_to_python(c) for c in new_column}
+        return SignalSchema(schema.values | vals)
+
     def select_except_signals(self, *args: str) -> "SignalSchema":
-        schema = copy.deepcopy(self.values)
-        for field in args:
-            if not isinstance(field, str):
-                raise SignalResolvingTypeError("select_except()", field)
+        def has_signal(signal: str):
+            signal = signal.replace(".", DEFAULT_DELIMITER)
+            return any(signal == s for s in self.db_signals())
 
-            if field not in self.values:
+        schema = copy.deepcopy(self.values)
+        for signal in args:
+            if not isinstance(signal, str):
+                raise SignalResolvingTypeError("select_except()", signal)
+
+            if signal not in self.values:
+                if has_signal(signal):
+                    raise SignalRemoveError(
+                        signal.split("."),
+                        "select_except() error - removing nested signal would"
+                        " break parent schema, which isn't supported.",
+                    )
                 raise SignalResolvingError(
-                    field.split("."),
-                    "select_except() error - the feature name does not exist or "
-                    "inside of feature (not supported)",
+                    signal.split("."),
+                    "select_except() error - the signal does not exist",
                 )
-            del schema[field]
+            del schema[signal]
 
         return SignalSchema(schema)
 
@@ -645,31 +862,49 @@ class SignalSchema:
 
     def mutate(self, args_map: dict) -> "SignalSchema":
         new_values = self.values.copy()
+        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, ColumnElement):
+            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)
+            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)
 
@@ -683,12 +918,37 @@ 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,
+        )
 
-        return SignalSchema(self.values | schema_right)
+        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(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 = {
+            key: right.values[key]
+            for key in [k for k in right.values if k not in self.values]
+        }
+        return SignalSchema(self.values | missing_schema)
 
     def get_signals(self, target_type: type[DataModel]) -> Iterator[str]:
         for path, type_, has_subtree, _ in self.get_flat_tree():
@@ -701,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 = {
@@ -736,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]
@@ -769,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
@@ -778,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
             )
@@ -795,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)
@@ -823,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,
@@ -838,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
@@ -846,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
@@ -859,7 +1154,7 @@ class SignalSchema:
 
         return res
 
-    def to_partial(self, *columns: str) -> "SignalSchema":
+    def to_partial(self, *columns: str) -> "SignalSchema":  # noqa: C901
         """
         Convert the schema to a partial schema with only the specified columns.
 
@@ -896,15 +1191,21 @@ 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] = {}
 
         def _type_name_to_partial(signal_name: str, type_name: str) -> str:
-            if "@" not in type_name:
+            # Check if we need to create a partial for this type
+            # Only create partials for custom types that are in the custom_types dict
+            if type_name not in custom_types:
                 return type_name
-            model_name, _ = ModelStore.parse_name_version(type_name)
+
+            if "@" in type_name:
+                model_name, _ = ModelStore.parse_name_version(type_name)
+            else:
+                model_name = type_name
 
             if signal_name not in signal_partials:
                 partial_versions.setdefault(model_name, 0)
@@ -928,6 +1229,14 @@ class SignalSchema:
                     parent_type_partial = _type_name_to_partial(signal, parent_type)
 
                     schema[signal] = parent_type_partial
+
+                    # If this is a complex signal without field specifier (just "file")
+                    # and it's a custom type, include the entire complex signal
+                    if len(column_parts) == 1 and parent_type in custom_types:
+                        # Include the entire complex signal - no need to create partial
+                        schema[signal] = parent_type
+                        continue
+
                     continue
 
                 if parent_type not in custom_types:
@@ -942,6 +1251,20 @@ class SignalSchema:
                         f"Field {signal} not found in custom type {parent_type}"
                     )
 
+                # Check if this is the last part and if the column type is a complex
+                is_last_part = i == len(column_parts) - 1
+                is_complex_signal = signal_type in custom_types
+
+                if is_last_part and is_complex_signal:
+                    schema[column] = signal_type
+                    # Also need to remove the partial schema entry we created for the
+                    # parent since we're promoting the nested complex column to root
+                    parent_signal = column_parts[0]
+                    schema.pop(parent_signal, None)
+                    # Don't create partial types for this case
+                    break
+
+                # Create partial type for this field
                 partial_type = _type_name_to_partial(
                     ".".join(column_parts[: i + 1]),
                     signal_type,