datachain 0.14.2__py3-none-any.whl → 0.39.0__py3-none-any.whl

datachain/func/func.py

@@ -1,12 +1,13 @@
 import inspect
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Any, Union, get_args, get_origin
 
 from sqlalchemy import BindParameter, Case, ColumnElement, Integer, cast, desc
 from sqlalchemy.sql import func as sa_func
 
 from datachain.lib.convert.python_to_sql import python_to_sql
 from datachain.lib.convert.sql_to_python import sql_to_python
+from datachain.lib.model_store import ModelStore
 from datachain.lib.utils import DataChainColumnError, DataChainParamsError
 from datachain.query.schema import Column, ColumnMeta
 from datachain.sql.functions import numeric
@@ -22,24 +23,29 @@ if TYPE_CHECKING:
     from .window import Window
 
 
-ColT = Union[str, ColumnElement, "Func", tuple]
+ColT = Union[str, tuple, Column, ColumnElement, "Func"]
 
 
-class Func(Function):
+class Func(Function):  # noqa: PLW1641
     """Represents a function to be applied to a column in a SQL query."""
 
+    cols: Sequence[ColT]
+    args: Sequence[Any]
+
     def __init__(
         self,
         name: str,
         inner: Callable,
-        cols: Optional[Sequence[ColT]] = None,
-        args: Optional[Sequence[Any]] = None,
-        kwargs: Optional[dict[str, Any]] = None,
-        result_type: Optional["DataType"] = None,
+        cols: Sequence[ColT] | None = None,
+        args: Sequence[Any] | None = None,
+        kwargs: dict[str, Any] | None = None,
+        result_type: "DataType | None" = None,
+        type_from_args: Callable[..., "DataType"] | None = None,
         is_array: bool = False,
+        from_array: bool = False,
         is_window: bool = False,
-        window: Optional["Window"] = None,
-        label: Optional[str] = None,
+        window: "Window | None" = None,
+        label: str | None = None,
     ) -> None:
         self.name = name
         self.inner = inner
@@ -47,7 +53,9 @@ class Func(Function):
         self.args = args or []
         self.kwargs = kwargs or {}
         self.result_type = result_type
+        self.type_from_args = type_from_args
         self.is_array = is_array
+        self.from_array = from_array
         self.is_window = is_window
         self.window = window
         self.col_label = label
@@ -66,7 +74,9 @@ class Func(Function):
             self.args,
             self.kwargs,
             self.result_type,
+            self.type_from_args,
             self.is_array,
+            self.from_array,
             self.is_window,
             window,
             self.col_label,
@@ -89,7 +99,7 @@ class Func(Function):
             else []
         )
 
-    def _db_col_type(self, signals_schema: "SignalSchema") -> Optional["DataType"]:
+    def _db_col_type(self, signals_schema: "SignalSchema") -> "DataType | None":
         if not self._db_cols:
             return None
 
@@ -101,53 +111,69 @@ class Func(Function):
                     "Columns must have the same type to infer result type",
                 )
 
+        if self.from_array:
+            if get_origin(col_type) is not list:
+                raise DataChainColumnError(
+                    str(self),
+                    "Array column must be of type list",
+                )
+            if self.is_array:
+                return col_type
+            col_args = get_args(col_type)
+            if len(col_args) != 1:
+                raise DataChainColumnError(
+                    str(self),
+                    "Array column must have a single type argument",
+                )
+            return col_args[0]
+
         return list[col_type] if self.is_array else col_type  # type: ignore[valid-type]
 
-    def __add__(self, other: Union[ColT, float]) -> "Func":
+    def __add__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("add", lambda a: a + other, [self])
         return Func("add", lambda a1, a2: a1 + a2, [self, other])
 
-    def __radd__(self, other: Union[ColT, float]) -> "Func":
+    def __radd__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("add", lambda a: other + a, [self])
         return Func("add", lambda a1, a2: a1 + a2, [other, self])
 
-    def __sub__(self, other: Union[ColT, float]) -> "Func":
+    def __sub__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("sub", lambda a: a - other, [self])
         return Func("sub", lambda a1, a2: a1 - a2, [self, other])
 
-    def __rsub__(self, other: Union[ColT, float]) -> "Func":
+    def __rsub__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("sub", lambda a: other - a, [self])
         return Func("sub", lambda a1, a2: a1 - a2, [other, self])
 
-    def __mul__(self, other: Union[ColT, float]) -> "Func":
+    def __mul__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("mul", lambda a: a * other, [self])
         return Func("mul", lambda a1, a2: a1 * a2, [self, other])
 
-    def __rmul__(self, other: Union[ColT, float]) -> "Func":
+    def __rmul__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("mul", lambda a: other * a, [self])
         return Func("mul", lambda a1, a2: a1 * a2, [other, self])
 
-    def __truediv__(self, other: Union[ColT, float]) -> "Func":
+    def __truediv__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("div", lambda a: _truediv(a, other), [self], result_type=float)
         return Func(
             "div", lambda a1, a2: _truediv(a1, a2), [self, other], result_type=float
         )
 
-    def __rtruediv__(self, other: Union[ColT, float]) -> "Func":
+    def __rtruediv__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("div", lambda a: _truediv(other, a), [self], result_type=float)
         return Func(
             "div", lambda a1, a2: _truediv(a1, a2), [other, self], result_type=float
         )
 
-    def __floordiv__(self, other: Union[ColT, float]) -> "Func":
+    def __floordiv__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "floordiv", lambda a: _floordiv(a, other), [self], result_type=int
@@ -156,7 +182,7 @@ class Func(Function):
             "floordiv", lambda a1, a2: _floordiv(a1, a2), [self, other], result_type=int
         )
 
-    def __rfloordiv__(self, other: Union[ColT, float]) -> "Func":
+    def __rfloordiv__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "floordiv", lambda a: _floordiv(other, a), [self], result_type=int
@@ -165,17 +191,17 @@ class Func(Function):
             "floordiv", lambda a1, a2: _floordiv(a1, a2), [other, self], result_type=int
         )
 
-    def __mod__(self, other: Union[ColT, float]) -> "Func":
+    def __mod__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("mod", lambda a: a % other, [self], result_type=int)
         return Func("mod", lambda a1, a2: a1 % a2, [self, other], result_type=int)
 
-    def __rmod__(self, other: Union[ColT, float]) -> "Func":
+    def __rmod__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("mod", lambda a: other % a, [self], result_type=int)
         return Func("mod", lambda a1, a2: a1 % a2, [other, self], result_type=int)
 
-    def __and__(self, other: Union[ColT, float]) -> "Func":
+    def __and__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "and", lambda a: numeric.bit_and(a, other), [self], result_type=int
@@ -187,7 +213,7 @@ class Func(Function):
             result_type=int,
         )
 
-    def __rand__(self, other: Union[ColT, float]) -> "Func":
+    def __rand__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "and", lambda a: numeric.bit_and(other, a), [self], result_type=int
@@ -199,7 +225,7 @@ class Func(Function):
             result_type=int,
         )
 
-    def __or__(self, other: Union[ColT, float]) -> "Func":
+    def __or__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "or", lambda a: numeric.bit_or(a, other), [self], result_type=int
@@ -208,7 +234,7 @@ class Func(Function):
             "or", lambda a1, a2: numeric.bit_or(a1, a2), [self, other], result_type=int
         )
 
-    def __ror__(self, other: Union[ColT, float]) -> "Func":
+    def __ror__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "or", lambda a: numeric.bit_or(other, a), [self], result_type=int
@@ -217,7 +243,7 @@ class Func(Function):
             "or", lambda a1, a2: numeric.bit_or(a1, a2), [other, self], result_type=int
         )
 
-    def __xor__(self, other: Union[ColT, float]) -> "Func":
+    def __xor__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "xor", lambda a: numeric.bit_xor(a, other), [self], result_type=int
@@ -229,7 +255,7 @@ class Func(Function):
             result_type=int,
         )
 
-    def __rxor__(self, other: Union[ColT, float]) -> "Func":
+    def __rxor__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "xor", lambda a: numeric.bit_xor(other, a), [self], result_type=int
@@ -241,7 +267,7 @@ class Func(Function):
             result_type=int,
         )
 
-    def __rshift__(self, other: Union[ColT, float]) -> "Func":
+    def __rshift__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "rshift",
@@ -256,7 +282,7 @@ class Func(Function):
             result_type=int,
         )
 
-    def __rrshift__(self, other: Union[ColT, float]) -> "Func":
+    def __rrshift__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "rshift",
@@ -271,7 +297,7 @@ class Func(Function):
             result_type=int,
         )
 
-    def __lshift__(self, other: Union[ColT, float]) -> "Func":
+    def __lshift__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "lshift",
@@ -286,7 +312,7 @@ class Func(Function):
             result_type=int,
         )
 
-    def __rlshift__(self, other: Union[ColT, float]) -> "Func":
+    def __rlshift__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func(
                 "lshift",
@@ -301,12 +327,12 @@ class Func(Function):
             result_type=int,
         )
 
-    def __lt__(self, other: Union[ColT, float]) -> "Func":
+    def __lt__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("lt", lambda a: a < other, [self], result_type=bool)
         return Func("lt", lambda a1, a2: a1 < a2, [self, other], result_type=bool)
 
-    def __le__(self, other: Union[ColT, float]) -> "Func":
+    def __le__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("le", lambda a: a <= other, [self], result_type=bool)
         return Func("le", lambda a1, a2: a1 <= a2, [self, other], result_type=bool)
@@ -321,12 +347,12 @@ class Func(Function):
             return Func("ne", lambda a: a != other, [self], result_type=bool)
         return Func("ne", lambda a1, a2: a1 != a2, [self, other], result_type=bool)
 
-    def __gt__(self, other: Union[ColT, float]) -> "Func":
+    def __gt__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("gt", lambda a: a > other, [self], result_type=bool)
         return Func("gt", lambda a1, a2: a1 > a2, [self, other], result_type=bool)
 
-    def __ge__(self, other: Union[ColT, float]) -> "Func":
+    def __ge__(self, other: ColT | float) -> "Func":
         if isinstance(other, (int, float)):
             return Func("ge", lambda a: a >= other, [self], result_type=bool)
         return Func("ge", lambda a1, a2: a1 >= a2, [self, other], result_type=bool)
@@ -339,13 +365,15 @@ class Func(Function):
             self.args,
             self.kwargs,
             self.result_type,
+            self.type_from_args,
             self.is_array,
+            self.from_array,
             self.is_window,
             self.window,
             label,
         )
 
-    def get_col_name(self, label: Optional[str] = None) -> str:
+    def get_col_name(self, label: str | None = None) -> str:
         if label:
             return label
         if self.col_label:
@@ -360,7 +388,7 @@ class Func(Function):
         return self.name
 
     def get_result_type(
-        self, signals_schema: Optional["SignalSchema"] = None
+        self, signals_schema: "SignalSchema | None" = None
     ) -> "DataType":
         if self.result_type:
             return self.result_type
@@ -368,6 +396,15 @@ class Func(Function):
         if signals_schema and (col_type := self._db_col_type(signals_schema)):
             return col_type
 
+        if (
+            self.type_from_args
+            and (self.cols is None or self.cols == [])
+            and self.args is not None
+            and len(self.args) > 0
+            and (result_type := self.type_from_args(*self.args)) is not None
+        ):
+            return result_type
+
         raise DataChainColumnError(
             str(self),
             "Column name is required to infer result type",
@@ -375,10 +412,24 @@ class Func(Function):
 
     def get_column(
         self,
-        signals_schema: Optional["SignalSchema"] = None,
-        label: Optional[str] = None,
-        table: Optional["TableClause"] = None,
+        signals_schema: "SignalSchema | None" = None,
+        label: str | None = None,
+        table: "TableClause | None" = None,
     ) -> Column:
+        # Guard against using complex (pydantic) object columns in SQL funcs
+        if signals_schema and self._db_cols:
+            for arg in self._db_cols:
+                # _db_cols normalizes known columns to strings; skip non-string args
+                if not isinstance(arg, str):
+                    continue
+                t_with_sub = signals_schema.get_column_type(arg, with_subtree=True)
+                if ModelStore.is_pydantic(t_with_sub):
+                    raise DataChainParamsError(
+                        f"Function {self.name} doesn't support complex object "
+                        f"columns like '{arg}'. Use a leaf field (e.g., "
+                        f"'{arg}.path') or use UDFs to operate on complex objects."
+                    )
+
         col_type = self.get_result_type(signals_schema)
         sql_type = python_to_sql(col_type)
 
@@ -398,6 +449,7 @@ class Func(Function):
             return col
 
         cols = [get_col(col) for col in self._db_cols]
+
         kwargs = {k: get_col(v, string_as_literal=True) for k, v in self.kwargs.items()}
         func_col = self.inner(*cols, *self.args, **kwargs)
 
@@ -434,9 +486,8 @@ def get_db_col_type(signals_schema: "SignalSchema", col: ColT) -> "DataType":
     if isinstance(col, ColumnElement) and not hasattr(col, "name"):
         return sql_to_python(col)
 
-    return signals_schema.get_column_type(
-        col.name if isinstance(col, ColumnElement) else col  # type: ignore[arg-type]
-    )
+    name = col.name if isinstance(col, ColumnElement) else col  # type: ignore[assignment]
+    return signals_schema.get_column_type(name)  # type: ignore[arg-type]
 
 
 def _truediv(a, b):

datachain/func/numeric.py

@@ -1,31 +1,34 @@
-from typing import Union
-
+from datachain.query.schema import Column
 from datachain.sql.functions import numeric
 
-from .func import ColT, Func
+from .func import Func
 
 
-def bit_and(*args: Union[ColT, int]) -> Func:
+def bit_and(*args: str | Column | Func | int) -> Func:
     """
-    Computes the bitwise AND operation between two values.
+    Returns a function that computes the bitwise AND operation between two values.
 
     Args:
-        args (str | int): Two values to compute the bitwise AND operation between.
-            If a string is provided, it is assumed to be the name of the column vector.
+        args (str | Column | Func | int): Two values to compute
+            the bitwise AND operation between.
+            If a string is provided, it is assumed to be the name of the column.
+            If a Column is provided, it is assumed to be a column.
+            If a Func is provided, it is assumed to be a function returning an int.
             If an integer is provided, it is assumed to be a constant value.
 
     Returns:
-        Func: A Func object that represents the bitwise AND function.
+        Func: A `Func` object that represents the bitwise AND function.
 
     Example:
         ```py
         dc.mutate(
-            xor1=func.bit_and("signal.values", 0x0F),
+            and1=func.bit_and("signal.value", 0x0F),
+            and2=func.bit_and(dc.C("signal.value1"), "signal.value2"),
         )
         ```
 
     Notes:
-        - Result column will always be of type int.
+        - The result column will always be of type int.
     """
     cols, func_args = [], []
     for arg in args:
@@ -46,27 +49,31 @@ def bit_and(*args: Union[ColT, int]) -> Func:
     )
 
 
-def bit_or(*args: Union[ColT, int]) -> Func:
+def bit_or(*args: str | Column | Func | int) -> Func:
     """
-    Computes the bitwise AND operation between two values.
+    Returns a function that computes the bitwise OR operation between two values.
 
     Args:
-        args (str | int): Two values to compute the bitwise OR operation between.
-            If a string is provided, it is assumed to be the name of the column vector.
+        args (str | Column | Func | int): Two values to compute
+            the bitwise OR operation between.
+            If a string is provided, it is assumed to be the name of the column.
+            If a Column is provided, it is assumed to be a column.
+            If a Func is provided, it is assumed to be a function returning an int.
             If an integer is provided, it is assumed to be a constant value.
 
     Returns:
-        Func: A Func object that represents the bitwise OR function.
+        Func: A `Func` object that represents the bitwise OR function.
 
     Example:
         ```py
         dc.mutate(
-            xor1=func.bit_or("signal.values", 0x0F),
+            or1=func.bit_or("signal.value", 0x0F),
+            or2=func.bit_or(dc.C("signal.value1"), "signal.value2"),
         )
         ```
 
     Notes:
-        - Result column will always be of type int.
+        - The result column will always be of type int.
     """
     cols, func_args = [], []
     for arg in args:
@@ -87,27 +94,31 @@ def bit_or(*args: Union[ColT, int]) -> Func:
     )
 
 
-def bit_xor(*args: Union[ColT, int]) -> Func:
+def bit_xor(*args: str | Column | Func | int) -> Func:
     """
-    Computes the bitwise XOR operation between two values.
+    Returns a function that computes the bitwise XOR operation between two values.
 
     Args:
-        args (str | int): Two values to compute the bitwise XOR operation between.
-            If a string is provided, it is assumed to be the name of the column vector.
+        args (str | Column | Func | int): Two values to compute
+            the bitwise XOR operation between.
+            If a string is provided, it is assumed to be the name of the column.
+            If a Column is provided, it is assumed to be a column.
+            If a Func is provided, it is assumed to be a function returning an int.
             If an integer is provided, it is assumed to be a constant value.
 
     Returns:
-        Func: A Func object that represents the bitwise XOR function.
+        Func: A `Func` object that represents the bitwise XOR function.
 
     Example:
         ```py
         dc.mutate(
-            xor1=func.bit_xor("signal.values", 0x0F),
+            xor1=func.bit_xor("signal.value", 0x0F),
+            xor2=func.bit_xor(dc.C("signal.value1"), "signal.value2"),
         )
         ```
 
     Notes:
-        - Result column will always be of type int.
+        - The result column will always be of type int.
     """
     cols, func_args = [], []
     for arg in args:
@@ -128,28 +139,30 @@ def bit_xor(*args: Union[ColT, int]) -> Func:
     )
 
 
-def int_hash_64(col: Union[ColT, int]) -> Func:
+def int_hash_64(col: str | Column | Func | int) -> Func:
     """
-    Returns the 64-bit hash of an integer.
+    Returns a function that computes the 64-bit hash of an integer.
 
     Args:
-        col (str | int): String to compute the hash of.
+        col (str | Column | Func | int): Integer to compute the hash of.
             If a string is provided, it is assumed to be the name of the column.
-            If a int is provided, it is assumed to be an int literal.
+            If a Column is provided, it is assumed to be a column.
             If a Func is provided, it is assumed to be a function returning an int.
+            If an int is provided, it is assumed to be an int literal.
 
     Returns:
-        Func: A Func object that represents the 64-bit hash function.
+        Func: A `Func` object that represents the 64-bit hash function.
 
     Example:
         ```py
         dc.mutate(
             val_hash=func.int_hash_64("val"),
+            val_hash2=func.int_hash_64(dc.C("val2")),
         )
         ```
 
-    Note:
-        - Result column will always be of type int.
+    Notes:
+        - The result column will always be of type int.
     """
     cols, args = [], []
     if isinstance(col, int):
@@ -162,9 +175,9 @@ def int_hash_64(col: Union[ColT, int]) -> Func:
     )
 
 
-def bit_hamming_distance(*args: Union[ColT, int]) -> Func:
+def bit_hamming_distance(*args: str | Column | Func | int) -> Func:
     """
-    Computes the Hamming distance between the bit representations of two integer values.
+    Returns a function that computes the Hamming distance between two integers.
 
     The Hamming distance is the number of positions at which the corresponding bits
     are different. This function returns the dissimilarity between the integers,
@@ -172,22 +185,26 @@ def bit_hamming_distance(*args: Union[ColT, int]) -> Func:
     in the integer indicate higher dissimilarity.
 
     Args:
-        args (str | int): Two integers to compute the Hamming distance between.
-            If a str is provided, it is assumed to be the name of the column.
+        args (str | Column | Func | int): Two integers to compute
+            the Hamming distance between.
+            If a string is provided, it is assumed to be the name of the column.
+            If a Column is provided, it is assumed to be a column.
+            If a Func is provided, it is assumed to be a function returning an int.
             If an int is provided, it is assumed to be an integer literal.
 
     Returns:
-        Func: A Func object that represents the Hamming distance function.
+        Func: A `Func` object that represents the Hamming distance function.
 
     Example:
         ```py
         dc.mutate(
-            ham_dist=func.bit_hamming_distance("embed1", 123456),
+            hd1=func.bit_hamming_distance("signal.value1", "signal.value2"),
+            hd2=func.bit_hamming_distance(dc.C("signal.value1"), 0x0F),
         )
         ```
 
     Notes:
-        - Result column will always be of type int.
+        - The result column will always be of type int.
     """
     cols, func_args = [], []
     for arg in args: