datachain 0.37.8__py3-none-any.whl → 0.37.9__py3-none-any.whl

datachain/data_storage/warehouse.py

@@ -18,6 +18,7 @@ from datachain.data_storage.schema import convert_rows_custom_column_types
 from datachain.data_storage.serializer import Serializable
 from datachain.dataset import DatasetRecord, StorageURI
 from datachain.lib.file import File
+from datachain.lib.model_store import ModelStore
 from datachain.lib.signal_schema import SignalSchema
 from datachain.node import DirType, DirTypeGroup, Node, NodeWithPath, get_path
 from datachain.query.batch import RowsOutput
@@ -76,6 +77,29 @@ class AbstractWarehouse(ABC, Serializable):
     def cleanup_for_tests(self):
         """Cleanup for tests."""
 
+    def _to_jsonable(self, obj: Any) -> Any:
+        """Recursively convert Python/Pydantic structures into JSON-serializable
+        objects.
+        """
+
+        if ModelStore.is_pydantic(type(obj)):
+            return obj.model_dump()
+
+        if isinstance(obj, dict):
+            out: dict[str, Any] = {}
+            for k, v in obj.items():
+                if not isinstance(k, str):
+                    key_str = json.dumps(self._to_jsonable(k), ensure_ascii=False)
+                else:
+                    key_str = k
+                out[key_str] = self._to_jsonable(v)
+            return out
+
+        if isinstance(obj, (list, tuple, set)):
+            return [self._to_jsonable(i) for i in obj]
+
+        return obj
+
     def convert_type(  # noqa: PLR0911
         self,
         val: Any,
@@ -122,11 +146,13 @@ class AbstractWarehouse(ABC, Serializable):
             if col_python_type is dict or col_type_name == "JSON":
                 if value_type is str:
                     return val
-                if value_type in (dict, list):
-                    return json.dumps(val, ensure_ascii=False)
-                raise ValueError(
-                    f"Cannot convert value {val!r} with type {value_type} to JSON"
-                )
+                try:
+                    json_ready = self._to_jsonable(val)
+                    return json.dumps(json_ready, ensure_ascii=False)
+                except Exception as e:
+                    raise ValueError(
+                        f"Cannot convert value {val!r} with type {value_type} to JSON"
+                    ) from e
 
             if isinstance(val, col_python_type):
                 return val

datachain/lib/convert/values_to_tuples.py

@@ -13,41 +13,153 @@ class ValuesToTupleError(DataChainParamsError):
         super().__init__(f"Cannot convert signals for dataset{ds_name}: {msg}")
 
 
-def values_to_tuples(  # noqa: C901, PLR0912
-    ds_name: str = "",
-    output: DataType | Sequence[str] | dict[str, DataType] | None = None,
-    **fr_map: Sequence[DataValue],
-) -> tuple[Any, Any, Any]:
-    if output:
-        if not isinstance(output, (Sequence, str, dict)):
-            if len(fr_map) != 1:
-                raise ValuesToTupleError(
-                    ds_name,
-                    f"only one output type was specified, {len(fr_map)} expected",
-                )
-            if not isinstance(output, type):
-                raise ValuesToTupleError(
-                    ds_name,
-                    f"output must specify a type while '{output}' was given",
-                )
+def _find_first_non_none(sequence: Sequence[Any]) -> Any | None:
+    """Find the first non-None element in a sequence."""
+    try:
+        return next(itertools.dropwhile(lambda i: i is None, sequence))
+    except StopIteration:
+        return None
+
+
+def _infer_list_item_type(lst: list) -> type:
+    """Infer the item type of a list, handling None values and nested lists."""
+    if len(lst) == 0:
+        # Default to str when list is empty to avoid generic list
+        return str
+
+    first_item = _find_first_non_none(lst)
+    if first_item is None:
+        # Default to str when all items are None
+        return str
+
+    item_type = type(first_item)
+
+    # Handle nested lists one level deep
+    if isinstance(first_item, list) and len(first_item) > 0:
+        nested_item = _find_first_non_none(first_item)
+        if nested_item is not None:
+            return list[type(nested_item)]  # type: ignore[misc, return-value]
+        # Default to str for nested lists with all None
+        return list[str]  # type: ignore[return-value]
+
+    return item_type
+
+
+def _infer_dict_value_type(dct: dict) -> type:
+    """Infer the value type of a dict, handling None values and list values."""
+    if len(dct) == 0:
+        # Default to str when dict is empty to avoid generic dict values
+        return str
+
+    # Find first non-None value
+    first_value = None
+    for val in dct.values():
+        if val is not None:
+            first_value = val
+            break
+
+    if first_value is None:
+        # Default to str when all values are None
+        return str
+
+    # Handle list values
+    if isinstance(first_value, list) and len(first_value) > 0:
+        list_item = _find_first_non_none(first_value)
+        if list_item is not None:
+            return list[type(list_item)]  # type: ignore[misc, return-value]
+        # Default to str for lists with all None
+        return list[str]  # type: ignore[return-value]
+
+    return type(first_value)
+
+
+def _infer_type_from_sequence(
+    sequence: Sequence[DataValue], signal_name: str, ds_name: str
+) -> type:
+    """
+    Infer the type from a sequence of values.
+
+    Returns str if all values are None, otherwise infers from the first non-None value.
+    Handles lists and dicts with proper type inference for nested structures.
+    """
+    first_element = _find_first_non_none(sequence)
+
+    if first_element is None:
+        # Default to str if column is empty or all values are None
+        return str
+
+    typ = type(first_element)
+
+    if not is_chain_type(typ):
+        raise ValuesToTupleError(
+            ds_name,
+            f"signal '{signal_name}' has unsupported type '{typ.__name__}'."
+            f" Please use DataModel types: {DataTypeNames}",
+        )
+
+    if isinstance(first_element, list):
+        item_type = _infer_list_item_type(first_element)
+        return list[item_type]  # type: ignore[valid-type, return-value]
+
+    if isinstance(first_element, dict):
+        # If the first dict is empty, use str as default key/value types
+        if len(first_element) == 0:
+            return dict[str, str]  # type: ignore[return-value]
+        first_key = next(iter(first_element.keys()))
+        value_type = _infer_dict_value_type(first_element)
+        return dict[type(first_key), value_type]  # type: ignore[misc, return-value]
+
+    return typ
 
-            key: str = next(iter(fr_map.keys()))
-            output = {key: output}  # type: ignore[dict-item]
 
-        if not isinstance(output, dict):
+def _validate_and_normalize_output(
+    output: DataType | Sequence[str] | dict[str, DataType] | None,
+    fr_map: dict[str, Sequence[DataValue]],
+    ds_name: str,
+) -> dict[str, DataType] | None:
+    """Validate and normalize the output parameter to a dict format."""
+    if not output:
+        return None
+
+    if not isinstance(output, (Sequence, str, dict)):
+        if len(fr_map) != 1:
             raise ValuesToTupleError(
                 ds_name,
-                "output type must be dict[str, DataType] while "
-                f"'{type(output).__name__}' is given",
+                f"only one output type was specified, {len(fr_map)} expected",
             )
-
-        if len(output) != len(fr_map):
+        if not isinstance(output, type):
             raise ValuesToTupleError(
                 ds_name,
-                f"number of outputs '{len(output)}' should match"
-                f" number of signals '{len(fr_map)}'",
+                f"output must specify a type while '{output}' was given",
             )
 
+        key: str = next(iter(fr_map.keys()))
+        return {key: output}  # type: ignore[dict-item]
+
+    if not isinstance(output, dict):
+        raise ValuesToTupleError(
+            ds_name,
+            "output type must be dict[str, DataType] while "
+            f"'{type(output).__name__}' is given",
+        )
+
+    if len(output) != len(fr_map):
+        raise ValuesToTupleError(
+            ds_name,
+            f"number of outputs '{len(output)}' should match"
+            f" number of signals '{len(fr_map)}'",
+        )
+
+    return output  # type: ignore[return-value]
+
+
+def values_to_tuples(
+    ds_name: str = "",
+    output: DataType | Sequence[str] | dict[str, DataType] | None = None,
+    **fr_map: Sequence[DataValue],
+) -> tuple[Any, Any, Any]:
+    output = _validate_and_normalize_output(output, fr_map, ds_name)
+
     types_map: dict[str, type] = {}
     length = -1
     for k, v in fr_map.items():
@@ -65,23 +177,7 @@ def values_to_tuples(  # noqa: C901, PLR0912
             # FIXME: Stops as soon as it finds the first non-None value.
             # If a non-None value appears early, it won't check the remaining items for
             # `None` values.
-            try:
-                first_not_none_element = next(
-                    itertools.dropwhile(lambda i: i is None, v)
-                )
-            except StopIteration:
-                # set default type to `str` if column is empty or all values are `None`
-                typ = str
-            else:
-                typ = type(first_not_none_element)  # type: ignore[assignment]
-                if not is_chain_type(typ):
-                    raise ValuesToTupleError(
-                        ds_name,
-                        f"signal '{k}' has unsupported type '{typ.__name__}'."
-                        f" Please use DataModel types: {DataTypeNames}",
-                    )
-                if isinstance(first_not_none_element, list):
-                    typ = list[type(first_not_none_element[0])]  # type: ignore[assignment, misc]
+            typ = _infer_type_from_sequence(v, k, ds_name)
             types_map[k] = typ
 
         if length < 0:

datachain/lib/data_model.py

@@ -64,6 +64,9 @@ def is_chain_type(t: type) -> bool:
     if orig is list and len(args) == 1:
         return is_chain_type(get_args(t)[0])
 
+    if orig is dict and len(args) == 2:
+        return is_chain_type(args[0]) and is_chain_type(args[1])
+
     if orig in (Union, types.UnionType) and len(args) == 2 and (type(None) in args):
         return is_chain_type(args[0] if args[1] is type(None) else args[1])
 

datachain/lib/signal_schema.py

@@ -1,6 +1,5 @@
 import copy
 import hashlib
-import json
 import logging
 import math
 import types
@@ -14,9 +13,7 @@ from typing import (
     TYPE_CHECKING,
     Annotated,
     Any,
-    Dict,  # type: ignore[UP035]
     Final,
-    List,  # type: ignore[UP035]
     Literal,
     Optional,
     Union,
@@ -24,6 +21,7 @@ from typing import (
     get_origin,
 )
 
+import ujson as json
 from pydantic import BaseModel, Field, ValidationError, create_model
 from sqlalchemy import ColumnElement
 from typing_extensions import Literal as LiteralEx
@@ -569,8 +567,10 @@ 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]
                 try:
@@ -585,6 +585,72 @@ class SignalSchema:
                 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
@@ -898,13 +964,13 @@ 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_)
             if len(args) == 0:
                 return "dict"

{datachain-0.37.8.dist-info → datachain-0.37.9.dist-info}/METADATA

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

{datachain-0.37.8.dist-info → datachain-0.37.9.dist-info}/RECORD

@@ -58,7 +58,7 @@ datachain/data_storage/metastore.py,sha256=DFyTkKLJN5-nFXXc7ln_rGj-FLctj0nrhXJxu
 datachain/data_storage/schema.py,sha256=3fAgiE11TIDYCW7EbTdiOm61SErRitvsLr7YPnUlVm0,9801
 datachain/data_storage/serializer.py,sha256=oL8i8smyAeVUyDepk8Xhf3lFOGOEHMoZjA5GdFzvfGI,3862
 datachain/data_storage/sqlite.py,sha256=o9TR6N27JB52M9rRXdM9uwdBektGucWtJi9UnmLGh0A,29669
-datachain/data_storage/warehouse.py,sha256=Zhf_HzhiEpsI0IuinAK-sF4ZMH66rV_ZDSOx-UFHv5o,34771
+datachain/data_storage/warehouse.py,sha256=_TGfMOtpltHA-G1KgoeIc_FFUomSmpAr94p-9AWNYIE,35642
 datachain/diff/__init__.py,sha256=lGrygGzdWSSYJ1DgX4h2q_ko5QINEW8PKfxOwE9ZFnI,9394
 datachain/fs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/fs/reference.py,sha256=A8McpXF0CqbXPqanXuvpKu50YLB3a2ZXA3YAPxtBXSM,914
@@ -78,7 +78,7 @@ datachain/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/lib/arrow.py,sha256=eCZtqbjAzkL4aemY74f_XkIJ_FWwXugJNjIFOwDa9w0,10815
 datachain/lib/audio.py,sha256=hHG29vqrV389im152wCjh80d0xqXGGvFnUpUwkzZejQ,7385
 datachain/lib/clip.py,sha256=nF8-N6Uz0MbAsPJBY2iXEYa3DPLo80OOer5SRNAtcGM,6149
-datachain/lib/data_model.py,sha256=H-bagx24-cLlC7ngSP6Dby4mB6kSxxV7KDiHxQjzwlg,3798
+datachain/lib/data_model.py,sha256=srz0pfFohSXwFnt5OMi1fNjSbKkFq8vzkcO0n4PHxlQ,3904
 datachain/lib/dataset_info.py,sha256=Ym7yYcGpfUmPLrfdxueijCVRP2Go6KbyuLk_fmzYgDU,3273
 datachain/lib/file.py,sha256=YO4QUaZVZ0TVW9fahERZ3HJXPNXjB4oYzvLQntQYT9s,47501
 datachain/lib/hf.py,sha256=jmyqRDXdksojUJCiU_2XFSIoMzzDJAZQs9xr-sEwEJc,7281
@@ -91,7 +91,7 @@ datachain/lib/namespaces.py,sha256=d4Zt2mYdGFctkA20SkB1woUxrNI4JwSxruxUGKwfauc,3
 datachain/lib/projects.py,sha256=FfBfGoWvy1SccCQW2ITKdDA6V03FbnRCusOeHdPHr6Y,4059
 datachain/lib/pytorch.py,sha256=gDJiUGoSaraW3JDPr5JW2a3SqT7KwgIMMpDTAC0L1_Y,7792
 datachain/lib/settings.py,sha256=maMtywOUetJvEApDiMVfTTq-oaRNvUIfDCrqZwFL2GE,7559
-datachain/lib/signal_schema.py,sha256=HeACY2i1bp9HOuaNm4_DvJh54Xnay3-qjcdERUBwFTU,41180
+datachain/lib/signal_schema.py,sha256=k43MncD1eew3zS6h_OYujg3jbvR6WH4Sj2mbrGvvvhc,43554
 datachain/lib/tar.py,sha256=MLcVjzIgBqRuJacCNpZ6kwSZNq1i2tLyROc8PVprHsA,999
 datachain/lib/text.py,sha256=uZom8qXfrv9QYvuDrvd0PuvPmj6qCsjVUwZSNr60BI4,1242
 datachain/lib/udf.py,sha256=51qgPO5s5MA5ccwl7IIPxbkEZ4IKZe4tzihcpZ8ufX0,18618
@@ -105,7 +105,7 @@ datachain/lib/convert/flatten.py,sha256=_5rjGFnN6t1KCX5ftL5rG7tiiNat7j0SdNqajO15
 datachain/lib/convert/python_to_sql.py,sha256=wfnqJ2vRL5UydNPQHshd82hUONsDBa4XyobCSTGqcEo,3187
 datachain/lib/convert/sql_to_python.py,sha256=Gxc4FylWC_Pvvuawuc2MKZIiuAWI7wje8pyeN1MxRrU,670
 datachain/lib/convert/unflatten.py,sha256=ysMkstwJzPMWUlnxn-Z-tXJR3wmhjHeSN_P-sDcLS6s,2010
-datachain/lib/convert/values_to_tuples.py,sha256=Sxj0ojeMSpAwM_NNoXa1dMR_2L_cQ6Xw_bAaNkEoNhU,4342
+datachain/lib/convert/values_to_tuples.py,sha256=nOn7dkzScYERZH-2vgUxkQawRQ1KgdIuSDIicvqZkc0,7171
 datachain/lib/dc/__init__.py,sha256=UrUzmDH6YyVl8fxM5iXTSFtl5DZTUzEYm1MaazK4vdQ,900
 datachain/lib/dc/csv.py,sha256=fIfj5-2Ix4z5D5yZueagd5WUWw86pusJ9JJKD-U3KGg,4407
 datachain/lib/dc/database.py,sha256=Wqob3dQc9Mol_0vagzVEXzteCKS9M0E3U5130KVmQKg,14629
@@ -165,9 +165,9 @@ datachain/sql/sqlite/vector.py,sha256=ncW4eu2FlJhrP_CIpsvtkUabZlQdl2D5Lgwy_cbfqR
 datachain/toolkit/__init__.py,sha256=eQ58Q5Yf_Fgv1ZG0IO5dpB4jmP90rk8YxUWmPc1M2Bo,68
 datachain/toolkit/split.py,sha256=9HHZl0fGs5Zj8b9l2L3IKf0AiiVNL9SnWbc2rfDiXRA,3710
 datachain/torch/__init__.py,sha256=gIS74PoEPy4TB3X6vx9nLO0Y3sLJzsA8ckn8pRWihJM,579
-datachain-0.37.8.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.37.8.dist-info/METADATA,sha256=6MLsgOSmSsxKXzbiOqTs9yQXaPhFu1QwgSqN_OmuQQM,13763
-datachain-0.37.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-datachain-0.37.8.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.37.8.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.37.8.dist-info/RECORD,,
+datachain-0.37.9.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.37.9.dist-info/METADATA,sha256=iZmFzvJMHOE2j4t9zGX2eliujOaRIcD0E39Cx1IXSXg,13763
+datachain-0.37.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+datachain-0.37.9.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.37.9.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.37.9.dist-info/RECORD,,