fix: keep None in an array wherever it sits

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":

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
 
 

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:]:

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):