diff --git a/src/datachain/data_storage/warehouse.py b/src/datachain/data_storage/warehouse.py index a7123f605..3b8912748 100644 --- a/src/datachain/data_storage/warehouse.py +++ b/src/datachain/data_storage/warehouse.py @@ -149,25 +149,50 @@ def convert_type( # noqa: PLR0911 if len(val) == 0: return [] - item_python_type = self.python_type(col_type.item_type) - - if item_python_type is not list: - if isinstance(val[0], item_python_type): - # SQLite ARRAY storage expects a list; tuples/sets must be - # converted to lists even when element types already match. - return list(val) - if item_python_type is float and isinstance(val[0], int): - return [float(i) for i in val] - - # Optimization: Reuse these values for each function call within the - # list comprehension. + item_type = col_type.item_type + item_python_type = self.python_type(item_type) item_type_info = ( - col_type.item_type, + item_type, item_python_type, - type(col_type.item_type).__name__, + type(item_type).__name__, col_name, ) - return [self.convert_type(i, *item_type_info) for i in val] + + if type(None) not in map(type, val): + if item_python_type is not list: + if isinstance(val[0], item_python_type): + # SQLite ARRAY storage expects a list; tuples/sets must + # be converted to lists even when element types already + # match. + return list(val) + if item_python_type is float and isinstance(val[0], int): + return [float(i) for i in val] + return [self.convert_type(i, *item_type_info) for i in val] + + # Only an array actually holding a None gets here, so nothing else + # changes shape. Which element came first used to decide the whole + # array's fate, and a None cannot answer for the rest of it. + if item_python_type is dict: + objects = [i for i in val if i is not None] + if objects and all(isinstance(i, dict) for i in objects): + # An array of JSON objects stays objects; normalizing each + # one rather than passing it through is what reaches a model + # nested inside. + return [self._to_jsonable(i) for i in val] + return [self.convert_type(i, *item_type_info) for i in val] + + # A JSON item carries its own null, and already did so in either + # order; only a type with no in-band null needs one kept back. + keep_none = ( + getattr(item_type, "dc_nullable", False) + and type(item_type).__name__ != "JSON" + ) + return [ + i + if (keep_none and i is None) or isinstance(i, item_python_type) + else self.convert_type(i, *item_type_info) + for i in val + ] # Special use case with JSON type as we save it as string if col_python_type is dict or col_type_name == "JSON": diff --git a/src/datachain/lib/convert/flatten.py b/src/datachain/lib/convert/flatten.py index be859a275..02a48d23e 100644 --- a/src/datachain/lib/convert/flatten.py +++ b/src/datachain/lib/convert/flatten.py @@ -85,10 +85,22 @@ def flatten_list(obj_list: list[BaseModel]) -> tuple: def _flatten_list_field(value: list) -> list: assert isinstance(value, list) - if value and ModelStore.is_pydantic(type(value[0])): - return [val.model_dump() for val in value] - if value and isinstance(value[0], list): - return [_flatten_list_field(v) for v in value] + # A None says nothing about what the list holds, so ask the first element + # that carries a type -- but one element cannot answer for the rest, so the + # others have to agree before a shape is assumed. Ordinary values settle it + # on the first element and never walk the list again. + first = next((val for val in value if val is not None), None) + + if ModelStore.is_pydantic(type(first)): + if all(val is None or ModelStore.is_pydantic(type(val)) for val in value): + return [None if val is None else val.model_dump() for val in value] + return value + + if isinstance(first, list) and all( + val is None or isinstance(val, list) for val in value + ): + return [None if val is None else _flatten_list_field(val) for val in value] + return value diff --git a/src/datachain/lib/convert/python_to_sql.py b/src/datachain/lib/convert/python_to_sql.py index f882ff505..257e4b355 100644 --- a/src/datachain/lib/convert/python_to_sql.py +++ b/src/datachain/lib/convert/python_to_sql.py @@ -2,16 +2,15 @@ from datetime import datetime from enum import Enum from types import UnionType -from typing import Annotated, Literal, Union, get_args, get_origin +from typing import Annotated, Any, Literal, Union, get_args, get_origin -from pydantic import BaseModel +from pydantic import BaseModel, RootModel from typing_extensions import Literal as LiteralEx from datachain.lib.data_model import ( NULLABLE_SCALARS, is_mapping_annotation, is_sequence_annotation, - unwrap_optional, ) from datachain.lib.model_store import ModelStore from datachain.sql.types import ( @@ -89,15 +88,136 @@ def _list_to_array(typ, args): if ModelStore.is_pydantic(args0): return Array(JSON()) - list_type = list_of_args_to_type(args) + # A model element is JSON whether or not it admits a None: is_chain_type + # accepts list[Model | None], and the union arm is all that hid the model. + # A fixed tuple keeps every slot in the one column, so every slot has to be + # one -- a slot this cannot read back as a model must still be refused. + slots = [arg for arg in args if arg is not Ellipsis] + if slots and all( + ModelStore.is_pydantic(slot) or _optional_model(slot) is not None + for slot in slots + ): + return Array(JSON()) + + # Resolve what the wrappers hold, not the wrappers: composed ones -- + # Annotated[str, ...] | Literal[None] -- are recognized by neither the scalar + # table nor the union handling on their own. + list_type = list_of_args_to_type(tuple(_unwrapped_for_lookup(a) for a in args)) + # Optional[scalar] elements map to a nullable Array element so None survives - # (ClickHouse: Array(Nullable(T))). - inner, is_optional = unwrap_optional(args0) - if is_optional and inner in NULLABLE_SCALARS: + # (ClickHouse: Array(Nullable(T))). A fixed tuple keeps every slot in the one + # column, so any slot admitting a None decides it, not just the first. + admits_none = any(_peel_optional(arg)[1] for arg in args if arg is not Ellipsis) + if admits_none and _takes_null(list_type): list_type = SQLType.as_nullable(list_type) return Array(list_type) +def _optional_model(annotation: Any) -> Any: + """The model behind ``Model | None``, or None if that is not the shape. + + Deliberately shallow: reading back a nested model is decided elsewhere and + does not look through Annotated, so admitting Annotated[Model, ...] | None + here would store it and then hand back plain dicts. + """ + if get_origin(annotation) not in (Union, UnionType): + return None + + arms = [arm for arm in get_args(annotation) if arm is not type(None)] + if len(arms) != 1 or not ModelStore.is_pydantic(arms[0]): + return None + + # A root model dumps to whatever it wraps, and the reader only rebuilds a + # model from a mapping, so this one would be handed back as its bare value. + if isinstance(arms[0], type) and issubclass(arms[0], RootModel): + return None + return arms[0] + + +def _unwrapped_for_lookup(annotation: Any) -> Any: + """What to resolve an element annotation as. + + Peeling is only useful here when what is left resolves to something. It does + not for a Literal holding nothing but None, for Ellipsis marking a variadic + tuple, or for a bare model that only the union around it made resolvable, and + the annotation as written is what those already mapped to. + """ + if annotation is Ellipsis: + return annotation + + peeled, _ = _peel_optional(annotation) + if peeled is annotation: + return annotation + + try: + resolved = python_to_sql(peeled) + except TypeError: + return annotation + + # An array has nowhere to keep the None that peeling just dropped -- only a + # scalar item is made nullable, and JSON carries its own null. + resolved_cls = resolved if isinstance(resolved, type) else type(resolved) + if issubclass(resolved_cls, Array): + return annotation + return peeled + + +def _peel_optional(annotation: Any) -> tuple[Any, bool]: + """Strip the wrappers around an element type, reporting whether it admits None. + + They nest either way round -- ``Optional[Annotated[int, ...]]`` and + ``Annotated[int | None, ...]`` -- and the None can be spelled as a union arm, + as one of a Literal's values, or as ``Literal[None]`` inside a union. + """ + is_optional = False + while True: + if get_origin(annotation) is Annotated: + annotation = get_args(annotation)[0] + continue + + if get_origin(annotation) in (Literal, LiteralEx): + values = get_args(annotation) + remaining = tuple(v for v in values if v is not None) + if not remaining: + return type(None), True + if len(remaining) != len(values): + is_optional = True + annotation = Literal[remaining] + continue + return annotation, is_optional + + if get_origin(annotation) in (Union, UnionType): + arms = [] + for arm in get_args(annotation): + peeled, arm_optional = _peel_optional(arm) + if arm_optional or peeled is type(None): + is_optional = True + if peeled is not type(None): + arms.append(peeled) + if not arms: + return type(None), True + if len(arms) == 1: + annotation = arms[0] + continue + return Union[tuple(arms)], is_optional # noqa: UP007 + + return annotation, is_optional + + +def _takes_null(sql_type: Any) -> bool: + """Whether a NULL can sit in this column type. + + Compared by subclass: a project or a user may have subclassed the scalar. + """ + sql_cls = sql_type if isinstance(sql_type, type) else type(sql_type) + for scalar in NULLABLE_SCALARS: + nullable = python_to_sql(scalar) + nullable_cls = nullable if isinstance(nullable, type) else type(nullable) + if issubclass(sql_cls, nullable_cls): + return True + return False + + def list_of_args_to_type(args) -> SQLType: first_type = python_to_sql(args[0]) for next_arg in args[1:]: diff --git a/src/datachain/sql/types.py b/src/datachain/sql/types.py index 9894fe5d0..2f8112d5d 100644 --- a/src/datachain/sql/types.py +++ b/src/datachain/sql/types.py @@ -561,6 +561,13 @@ def float64(self, value): def array(self, value, item_type, dialect): if value is None or item_type is None: return value + if getattr(item_type, "dc_nullable", False): + # A nullable element keeps its None, the way a nullable column does; + # float() would otherwise turn it into nan. + return [ + None if x is None else item_type.on_read_convert(x, dialect) + for x in value + ] return [item_type.on_read_convert(x, dialect) for x in value] def json(self, value): diff --git a/tests/func/test_data_storage.py b/tests/func/test_data_storage.py index 7702ebb19..34e1f006a 100644 --- a/tests/func/test_data_storage.py +++ b/tests/func/test_data_storage.py @@ -1,15 +1,16 @@ import uuid from datetime import datetime from decimal import Decimal -from typing import Any +from typing import Annotated, Any, Literal import numpy as np import pandas as pd import pytest import ujson as json -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, RootModel from datachain import json as dcjson +from datachain.lib.convert.python_to_sql import python_to_sql from datachain.sql.types import ( JSON, Array, @@ -19,6 +20,8 @@ Float32, Float64, Int, + Int64, + SQLType, String, ) from tests.utils import ( @@ -355,3 +358,198 @@ def test_a_model_refuses_numpy_that_would_be_stored_as_null(test_session, make): with pytest.raises(ValueError, match="writes NaN and infinities as null"): _model_payload(warehouse, make()) + + +@pytest.mark.parametrize( + "value,item_type,expected", + [ + pytest.param((1, None), Int64, [1, None], id="int-none-last"), + pytest.param((None, 2), Int64, [None, 2], id="int-none-first"), + pytest.param((None, None), Int64, [None, None], id="int-all-none"), + pytest.param([1, None], Float, [1.0, None], id="float-none-last"), + pytest.param([None, 2], Float, [None, 2.0], id="float-none-first"), + pytest.param(["a", None], String, ["a", None], id="str-none-last"), + pytest.param([None, "b"], String, [None, "b"], id="str-none-first"), + ], +) +def test_convert_type_keeps_none_wherever_it_sits_in_an_array( + test_session, value, item_type, expected +): + warehouse = test_session.catalog.warehouse + col_type = Array(SQLType.as_nullable(item_type)) + + converted = warehouse.convert_type( + value, + col_type, + warehouse.python_type(col_type), + "Array", + "test_column", + ) + + assert converted == expected + + +def test_convert_type_stores_a_json_array_the_same_wherever_none_sits(test_session): + warehouse = test_session.catalog.warehouse + col_type = Array(JSON()) + + def to_db(value): + return warehouse.convert_type( + value, col_type, warehouse.python_type(col_type), "Array", "test_column" + ) + + # What an item becomes is the backend's business; that its neighbours do not + # change the answer is not. + assert to_db([{"k": 1}, None]) == list(reversed(to_db([None, {"k": 1}]))) + + +@pytest.mark.parametrize( + "value", + [ + pytest.param([None, 2], id="none-first"), + pytest.param([1, None], id="none-last"), + ], +) +def test_convert_type_refuses_none_in_a_non_nullable_array(test_session, value): + warehouse = test_session.catalog.warehouse + col_type = Array(Int64) + + with pytest.raises(ValueError, match="incompatible"): + warehouse.convert_type( + value, col_type, warehouse.python_type(col_type), "Array", "test_column" + ) + + +class NestedItem(BaseModel): + n: int + + +@pytest.mark.parametrize( + "value", + [ + pytest.param((None, {"a": NestedItem(n=2)}), id="none-first"), + pytest.param(({"a": NestedItem(n=2)}, None), id="none-last"), + ], +) +def test_convert_type_leaves_no_model_in_a_json_array_holding_none(test_session, value): + warehouse = test_session.catalog.warehouse + col_type = Array(JSON()) + + converted = warehouse.convert_type( + value, col_type, warehouse.python_type(col_type), "Array", "test_column" + ) + + # Whatever shape the backend asks for, nothing unserializable may survive. + json.dumps(converted) + + +class SubclassedInt(Int64): + pass + + +@pytest.mark.parametrize( + "annotation,value", + [ + pytest.param(list[int | None], [1, None], id="plain-int"), + pytest.param(list[int | None], [None, 2], id="plain-int-none-first"), + pytest.param( + list[Annotated[int, "meta"] | None], [1, None], id="annotated-int" + ), + pytest.param(list[Literal["a", "b"] | None], ["a", None], id="literal-str"), + pytest.param( + list[Literal["a", "b"] | None], [None, "b"], id="literal-str-none-first" + ), + pytest.param( + list[Annotated[int | None, "meta"]], [1, None], id="optional-in-annotated" + ), + pytest.param( + list[Annotated[int | None, "meta"]], + [None, 2], + id="optional-in-annotated-none-first", + ), + pytest.param( + list[Annotated[Literal["a", None], "meta"]], # noqa: PYI061 + ["a", None], + id="none-among-literal-values", + ), + pytest.param(list[SubclassedInt | None], [1, None], id="subclassed-sql-type"), + pytest.param( + list[str | Literal[None]], # noqa: PYI061 + ["a", None], + id="literal-none-in-a-union", + ), + pytest.param( + list[Annotated[str, "meta"] | Literal[None]], # noqa: PYI061 + ["x", None], + id="annotated-beside-literal-none", + ), + pytest.param( + tuple[str, Literal[None]], # noqa: PYI061 + ("a", None), + id="null-only-literal-slot", + ), + pytest.param( + list[Literal[None]], # noqa: PYI061 + [None], + id="null-only-literal-item", + ), + pytest.param(tuple[int, int | None], (1, None), id="second-tuple-slot"), + pytest.param(tuple[int | None, int], (None, 1), id="first-tuple-slot"), + ], +) +def test_convert_type_keeps_none_for_a_wrapped_nullable_scalar( + test_session, annotation, value +): + warehouse = test_session.catalog.warehouse + col_type = python_to_sql(annotation) + + converted = warehouse.convert_type( + value, col_type, warehouse.python_type(col_type), "Array", "test_column" + ) + + assert converted == list(value) + + +def test_convert_type_json_encodes_an_all_none_array(test_session): + warehouse = test_session.catalog.warehouse + col_type = Array(SQLType.as_nullable(JSON())) + + converted = warehouse.convert_type( + (None, None), col_type, warehouse.python_type(col_type), "Array", "test_column" + ) + + # An array with no object in it is not an array of objects; each None stays + # whatever JSON writes for one, as it did before. + assert converted == [json.dumps(None)] * 2 + + +class SlotA(BaseModel): + x: int + + +class SlotB(BaseModel): + y: int + + +def test_python_to_sql_refuses_a_tuple_slot_that_reads_back_wrong(): + # Every slot of a fixed tuple shares the column, so one that cannot be read + # back as a model has to refuse the whole annotation rather than be stored + # and handed back as a plain dict. + with pytest.raises(TypeError): + python_to_sql(tuple[SlotA | None, Annotated[SlotB, "meta"] | None]) + + assert python_to_sql(tuple[SlotA | None, SlotB | None]).to_dict() == { + "type": "Array", + "item_type": {"type": "JSON"}, + } + + +class RootInt(RootModel[int]): + pass + + +def test_python_to_sql_refuses_an_optional_root_model_element(): + # A root model dumps to its bare value, which the reader cannot rebuild into + # a model, so it must not be admitted by the optional-model shortcut. + with pytest.raises(TypeError): + python_to_sql(list[RootInt | None]) diff --git a/tests/unit/lib/test_datachain.py b/tests/unit/lib/test_datachain.py index 1996f7ada..35bb86c6a 100644 --- a/tests/unit/lib/test_datachain.py +++ b/tests/unit/lib/test_datachain.py @@ -6313,3 +6313,117 @@ def build(): ), ) ] + + +class NullableInts(BaseModel): + vals: list[int | None] + + +class NullableFloats(BaseModel): + vals: list[float | None] + + +@pytest.mark.parametrize( + "vals", + [ + pytest.param([1, None], id="none-last"), + pytest.param([None, 2], id="none-first"), + pytest.param([None, None], id="all-none"), + pytest.param([None, 2, None, 4], id="none-interleaved"), + ], +) +def test_a_nullable_int_list_round_trips_whatever_the_order(test_session, vals): + rows = ( + dc.read_values(i=[1], session=test_session) + .map(h=lambda: NullableInts(vals=vals), output=NullableInts) + .to_list("h.vals") + ) + + assert rows == [(vals,)] + + +@pytest.mark.parametrize( + "vals,expected", + [ + pytest.param([1, None], [1.0, None], id="none-last"), + pytest.param([None, 2], [None, 2.0], id="none-first"), + pytest.param([None, None], [None, None], id="all-none"), + ], +) +def test_a_nullable_float_list_round_trips_whatever_the_order( + test_session, vals, expected +): + rows = ( + dc.read_values(i=[1], session=test_session) + .map(h=lambda: NullableFloats(vals=vals), output=NullableFloats) + .to_list("h.vals") + ) + + assert rows == [(expected,)] + + +class DictItem(BaseModel): + n: int + + +class DictItemHolder(BaseModel): + vals: list[dict[str, DictItem] | None] + + +@pytest.mark.parametrize( + "vals", + [ + pytest.param([{"a": DictItem(n=1)}, None], id="none-last"), + pytest.param([None, {"a": DictItem(n=1)}], id="none-first"), + ], +) +def test_models_inside_dict_items_convert_wherever_none_sits(test_session, vals): + rows = ( + dc.read_values(i=[1], session=test_session) + .map(h=lambda: DictItemHolder(vals=vals), output=DictItemHolder) + .to_list("h.vals") + ) + + assert rows == [(vals,)] + + +class OptionalChild(BaseModel): + x: int + + +class OptionalChildHolder(BaseModel): + vals: list[OptionalChild | None] + + +@pytest.mark.parametrize( + "vals", + [ + pytest.param([OptionalChild(x=1), None], id="none-last"), + pytest.param([None, OptionalChild(x=2)], id="none-first"), + pytest.param([None, None], id="all-none"), + ], +) +def test_a_list_of_models_admitting_none_round_trips(test_session, vals): + rows = ( + dc.read_values(i=[1], session=test_session) + .map(h=lambda: OptionalChildHolder(vals=vals), output=OptionalChildHolder) + .to_list("h.vals") + ) + + assert rows == [(vals,)] + + +class MixedUnionHolder(BaseModel): + vals: list[dict | list[dict] | None] + + +def test_a_list_of_mixed_shapes_is_left_alone(test_session): + value = [None, [{"a": 1}], {"b": 2}] + + rows = ( + dc.read_values(i=[1], session=test_session) + .map(h=lambda: MixedUnionHolder(vals=value), output=MixedUnionHolder) + .to_list("h.vals") + ) + + assert rows == [(value,)]