Check UDF params types in SignalSchema

src/datachain/lib/dc.py

@@ -629,7 +629,8 @@ def jmespath_to_name(s: str):
                 model_name=model_name,
                 jmespath=jmespath,
                 nrows=nrows,
-            )
+            ),
+            "params": {"file": File},
         }
         # disable prefetch if nrows is set
         settings = {"prefetch": 0} if nrows else {}
@@ -1003,8 +1004,9 @@ def _udf_to_obj(
         func: Optional[Union[Callable, UDFObjT]],
         params: Union[None, str, Sequence[str]],
         output: OutputType,
-        signal_map,
+        signal_map: dict[str, Callable],
     ) -> UDFObjT:
+        is_batch = target_class.is_input_batched
         is_generator = target_class.is_output_batched
         name = self.name or ""
 
@@ -1015,7 +1017,9 @@ def _udf_to_obj(
         if self._sys:
             signals_schema = SignalSchema({"sys": Sys}) | signals_schema
 
-        params_schema = signals_schema.slice(sign.params, self._setup)
+        params_schema = signals_schema.slice(
+            sign.params, self._setup, is_batch=is_batch
+        )
 
         return target_class._create(sign, params_schema)
 

src/datachain/lib/meta_formats.py

@@ -10,7 +10,7 @@
 from pydantic import BaseModel, ConfigDict, Field, ValidationError  # noqa: F401
 
 from datachain.lib.data_model import DataModel  # noqa: F401
-from datachain.lib.file import File
+from datachain.lib.file import TextFile
 
 
 class UserModel(BaseModel):
@@ -130,7 +130,7 @@ def read_meta(  # noqa: C901
     #
 
     def parse_data(
-        file: File,
+        file: TextFile,
         data_model=spec,
         format=format,
         jmespath=jmespath,

src/datachain/lib/signal_schema.py

@@ -154,9 +154,9 @@ def __init__(
             if not callable(func):
                 raise SetupError(key, "value must be function or callable class")
 
-    def _init_setup_values(self):
+    def _init_setup_values(self) -> None:
         if self.setup_values is not None:
-            return self.setup_values
+            return
 
         res = {}
         for key, func in self.setup_func.items():
@@ -398,7 +398,7 @@ def deserialize(schema: dict[str, Any]) -> "SignalSchema":
         return SignalSchema(signals)
 
     @staticmethod
-    def get_flatten_hidden_fields(schema):
+    def get_flatten_hidden_fields(schema: dict):
         custom_types = schema.get("_custom_types", {})
         if not custom_types:
             return []
@@ -464,19 +464,61 @@ def contains_file(self) -> bool:
         return False
 
     def slice(
-        self, keys: Sequence[str], setup: Optional[dict[str, Callable]] = None
+        self,
+        params: dict[str, Union[DataType, Any]],
+        setup: Optional[dict[str, Callable]] = None,
+        is_batch: bool = False,
     ) -> "SignalSchema":
-        # Make new schema that combines current schema and setup signals
-        setup = setup or {}
-        setup_no_types = dict.fromkeys(setup.keys(), str)
-        union = SignalSchema(self.values | setup_no_types)
-        # Slice combined schema by keys
-        schema = {}
-        for k in keys:
-            try:
-                schema[k] = union._find_in_tree(k.split("."))
-            except SignalResolvingError:
-                pass
+        """
+        Returns new schema that combines current schema and setup signals.
+        """
+        setup_params = setup.keys() if setup else []
+        schema: dict[str, DataType] = {}
+
+        for param, param_type in params.items():
+            # This is special case for setup params, they are always treated as strings
+            if param in setup_params:
+                schema[param] = str
+                continue
+
+            schema_type = self._find_in_tree(param.split("."))
+
+            if param_type is Any:
+                schema[param] = schema_type
+                continue
+
+            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):
+                schema_type = get_args(schema_type)[0]
+                if param_origin is Union and type(None) in get_args(param_type):
+                    param_type = get_args(param_type)[0]
+
+            if is_batch:
+                if param_type is list:
+                    schema[param] = schema_type
+                    continue
+
+                if param_origin is not list:
+                    raise SignalResolvingError(param.split("."), "is not a list")
+
+                param_type = get_args(param_type)[0]
+
+            if param_type == schema_type or (
+                isclass(param_type)
+                and isclass(schema_type)
+                and issubclass(param_type, File)
+                and issubclass(schema_type, File)
+            ):
+                schema[param] = schema_type
+                continue
+
+            raise SignalResolvingError(
+                param.split("."),
+                f"types mismatch: {param_type} != {schema_type}",
+            )
+
         return SignalSchema(schema, setup)
 
     def row_to_features(

src/datachain/lib/udf.py

@@ -159,6 +159,7 @@ def process(self, file) -> list[float]:
         ```
     """
 
+    is_input_batched = False
     is_output_batched = False
     prefetch: int = 0
 
@@ -395,6 +396,7 @@ def _prepare_rows(udf_inputs) -> "abc.Generator[Sequence[Any], None, None]":
 class BatchMapper(UDFBase):
     """Inherit from this class to pass to `DataChain.batch_map()`."""
 
+    is_input_batched = True
     is_output_batched = True
 
     def run(
@@ -481,6 +483,7 @@ def _process_row(row):
 class Aggregator(UDFBase):
     """Inherit from this class to pass to `DataChain.agg()`."""
 
+    is_input_batched = True
     is_output_batched = True
 
     def run(

src/datachain/lib/udf_signature.py

@@ -1,7 +1,7 @@
 import inspect
 from collections.abc import Generator, Iterator, Sequence
 from dataclasses import dataclass
-from typing import Callable, Union, get_args, get_origin
+from typing import Any, Callable, Union, get_args, get_origin
 
 from datachain.lib.data_model import DataType, DataTypeNames, is_chain_type
 from datachain.lib.signal_schema import SignalSchema
@@ -18,7 +18,7 @@ def __init__(self, chain: str, msg):
 @dataclass
 class UdfSignature:
     func: Union[Callable, UDFBase]
-    params: Sequence[str]
+    params: dict[str, Union[DataType, Any]]
     output_schema: SignalSchema
 
     DEFAULT_RETURN_TYPE = str
@@ -58,15 +58,23 @@ def parse(
         if not isinstance(udf_func, UDFBase) and not callable(udf_func):
             raise UdfSignatureError(chain, f"UDF '{udf_func}' is not callable")
 
-        func_params_map_sign, func_outs_sign, is_iterator = (
-            UdfSignature._func_signature(chain, udf_func)
+        func_params_map_sign, func_outs_sign, is_iterator = cls._func_signature(
+            chain, udf_func
         )
+
+        udf_params: dict[str, Union[DataType, Any]] = {}
         if params:
-            udf_params = [params] if isinstance(params, str) else params
-        elif not func_params_map_sign:
-            udf_params = []
-        else:
-            udf_params = list(func_params_map_sign.keys())
+            udf_params = (
+                {params: Any} if isinstance(params, str) else dict.fromkeys(params, Any)
+            )
+        elif func_params_map_sign:
+            udf_params = {
+                param: (
+                    param_type if param_type is not inspect.Parameter.empty else Any
+                )
+                for param, param_type in func_params_map_sign.items()
+            }
+
         if output:
             udf_output_map = UdfSignature._validate_output(
                 chain, signal_name, func, func_outs_sign, output

tests/func/test_datachain.py

@@ -707,7 +707,7 @@ def __init__(self):
             self.value = MyMapper.DEFAULT_VALUE
             self._had_teardown = False
 
-        def process(self, *args) -> int:
+        def process(self, key) -> int:
             return self.value
 
         def setup(self):

tests/unit/lib/test_datachain_bootstrap.py

@@ -15,7 +15,7 @@ def __init__(self):
             self.value = MyMapper.DEFAULT_VALUE
             self._had_teardown = False
 
-        def process(self, *args) -> int:
+        def process(self, key) -> int:
             return self.value
 
         def setup(self):
@@ -40,7 +40,7 @@ def __init__(self):
             self._had_bootstrap = False
             self._had_teardown = False
 
-        def __call__(self, *args):
+        def __call__(self, key):
             return None
 
         def bootstrap(self):