datachain 0.37.8__tar.gz → 0.37.9__tar.gz

{datachain-0.37.8 → datachain-0.37.9}/PKG-INFO

@@ -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 → datachain-0.37.9}/src/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-0.37.9/src/datachain/lib/convert/values_to_tuples.py

@@ -0,0 +1,210 @@
+import itertools
+from collections.abc import Sequence
+from typing import Any
+
+from datachain.lib.data_model import DataType, DataTypeNames, DataValue, is_chain_type
+from datachain.lib.utils import DataChainParamsError
+
+
+class ValuesToTupleError(DataChainParamsError):
+    def __init__(self, ds_name: str, msg: str):
+        if ds_name:
+            ds_name = f"' {ds_name}'"
+        super().__init__(f"Cannot convert signals for dataset{ds_name}: {msg}")
+
+
+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
+
+
+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,
+                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",
+            )
+
+        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():
+        if not isinstance(v, Sequence) or isinstance(v, str):  # type: ignore[unreachable]
+            raise ValuesToTupleError(ds_name, f"signals '{k}' is not a sequence")
+        len_ = len(v)
+
+        if output:
+            if k not in output:  # type: ignore[operator]
+                raise ValuesToTupleError(
+                    ds_name,
+                    f"signal '{k}' is not present in the output",
+                )
+        else:
+            # 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.
+            typ = _infer_type_from_sequence(v, k, ds_name)
+            types_map[k] = typ
+
+        if length < 0:
+            length = len_
+        elif length != len_:
+            raise ValuesToTupleError(
+                ds_name,
+                f"signal '{k}' should have length {length} while {len_} is given",
+            )
+
+    if not output:
+        output = types_map  # type: ignore[assignment]
+
+    if not output:
+        raise ValuesToTupleError(
+            ds_name,
+            "output type must be dict[str, DataType] while empty is given"
+            " and no signals are provided",
+        )
+
+    output_types: list[type] = list(output.values())  # type: ignore[union-attr,call-arg,arg-type]
+    if len(output) > 1:  # type: ignore[arg-type]
+        tuple_type = tuple(output_types)
+        res_type = tuple[tuple_type]  # type: ignore[valid-type]
+        res_values: Sequence[Any] = list(zip(*fr_map.values(), strict=False))
+    else:
+        res_type = output_types[0]  # type: ignore[misc]
+        res_values = next(iter(fr_map.values()))
+
+    return res_type, output, res_values

{datachain-0.37.8 → datachain-0.37.9}/src/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-0.37.8 → datachain-0.37.9}/src/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 → datachain-0.37.9}/src/datachain.egg-info/PKG-INFO

@@ -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 → datachain-0.37.9}/src/datachain.egg-info/SOURCES.txt

@@ -374,6 +374,7 @@ tests/func/test_query.py
 tests/func/test_read_database.py
 tests/func/test_read_dataset_remote.py
 tests/func/test_read_dataset_version_specifiers.py
+tests/func/test_retrieval.py
 tests/func/test_retry.py
 tests/func/test_session.py
 tests/func/test_storage_pattern.py

{datachain-0.37.8 → datachain-0.37.9}/tests/func/test_data_storage.py

@@ -2,6 +2,8 @@ from datetime import datetime
 from typing import Any
 
 import pytest
+import ujson as json
+from pydantic import BaseModel, ConfigDict
 
 from datachain.sql.types import (
     JSON,
@@ -93,14 +95,8 @@ def test_dir_expansion(cloud_test_catalog, version_aware, cloud_type):
     assert to_compare == expected
 
 
-@pytest.mark.parametrize(
-    "cloud_type,version_aware",
-    [("s3", True)],
-    indirect=True,
-)
-def test_convert_type(cloud_test_catalog):
-    ctc = cloud_test_catalog
-    catalog = ctc.catalog
+def test_convert_type(test_session):
+    catalog = test_session.catalog
     warehouse = catalog.warehouse
     now = datetime.now()
 
@@ -137,8 +133,47 @@ def test_convert_type(cloud_test_catalog):
     assert run_convert_type('{"a": 1}', JSON()) == '{"a": 1}'
     assert run_convert_type({"a": 1}, JSON()) == '{"a":1}'
     assert run_convert_type([{"a": 1}], JSON()) == '[{"a":1}]'
-    with pytest.raises(ValueError):
-        run_convert_type(0.5, JSON())
+    assert run_convert_type([[1, 2], [3, 4]], JSON()) == "[[1,2],[3,4]]"
+    assert run_convert_type(None, JSON()) == "null"
+    assert run_convert_type({"a": None}, JSON()) == '{"a":null}'
+    # primitives should serialize to valid JSON
+    assert run_convert_type(0.5, JSON()) == "0.5"
+
+    # JSON with Pydantic models (values and nested)
+    class MyFr(BaseModel):
+        model_config = ConfigDict(frozen=True)
+        nnn: str
+        count: int
+
+    fr1 = MyFr(nnn="x", count=1)
+    fr2 = MyFr(nnn="y", count=2)
+
+    # Pydantic as dict value
+    out = run_convert_type({"a": fr1}, JSON())
+    assert out == '{"a":{"nnn":"x","count":1}}'
+
+    # Pydantic in list
+    out = run_convert_type([fr1, fr2], JSON())
+    assert out == '[{"nnn":"x","count":1},{"nnn":"y","count":2}]'
+
+    # Nested structures with Pydantic
+    out = run_convert_type({"k": [{"inner": fr1}]}, JSON())
+    assert out == '{"k":[{"inner":{"nnn":"x","count":1}}]}'
+
+    # Complex dict key (tuple) becomes a JSON-encoded string key
+    out = run_convert_type({(1, "a"): 3}, JSON())
+    # Decode and compare to expected mapping using encoded key
+    loaded = json.loads(out)
+    assert loaded == {json.dumps([1, "a"]): 3}
+
+    # Pydantic model as dict key
+    key_model = MyFr(nnn="k", count=7)
+    d: dict[Any, Any] = {}
+    d[key_model] = "v"
+    out = run_convert_type(d, JSON())
+    loaded = json.loads(out)
+    expected_key = json.dumps({"nnn": "k", "count": 7})
+    assert loaded == {expected_key: "v"}
 
     # convert array to compatible type
     converted = run_convert_type([1, 2], Array(Float))