datachain 0.7.7__tar.gz → 0.7.8__tar.gz

{datachain-0.7.7 → datachain-0.7.8}/.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.8.0'
+    rev: 'v0.8.1'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.7.7/src/datachain.egg-info → datachain-0.7.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.7.7
+Version: 0.7.8
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0

{datachain-0.7.7 → datachain-0.7.8}/src/datachain/func/__init__.py

@@ -17,6 +17,7 @@ from .aggregate import (
 )
 from .array import cosine_distance, euclidean_distance, length, sip_hash_64
 from .conditional import greatest, least
+from .numeric import bit_and, bit_or, bit_xor, int_hash_64
 from .random import rand
 from .window import window
 
@@ -24,6 +25,9 @@ __all__ = [
     "any_value",
     "array",
     "avg",
+    "bit_and",
+    "bit_or",
+    "bit_xor",
     "case",
     "collect",
     "concat",
@@ -33,6 +37,7 @@ __all__ = [
     "euclidean_distance",
     "first",
     "greatest",
+    "int_hash_64",
     "least",
     "length",
     "literal",

{datachain-0.7.7 → datachain-0.7.8}/src/datachain/func/func.py

@@ -2,13 +2,15 @@ import inspect
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Callable, Optional, Union
 
-from sqlalchemy import BindParameter, Case, ColumnElement, desc
+from sqlalchemy import BindParameter, Case, ColumnElement, Integer, cast, desc
 from sqlalchemy.ext.hybrid import Comparator
+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.utils import DataChainColumnError, DataChainParamsError
 from datachain.query.schema import Column, ColumnMeta
+from datachain.sql.functions import numeric
 
 from .base import Function
 
@@ -98,94 +100,232 @@ class Func(Function):
         return list[col_type] if self.is_array else col_type  # type: ignore[valid-type]
 
     def __add__(self, other: Union[ColT, float]) -> "Func":
-        return math_add(self, other)
+        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":
-        return math_add(other, self)
+        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":
-        return math_sub(self, other)
+        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":
-        return math_sub(other, self)
+        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":
-        return math_mul(self, other)
+        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":
-        return math_mul(other, self)
+        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":
-        return math_truediv(self, other)
+        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":
-        return math_truediv(other, self)
+        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":
-        return math_floordiv(self, other)
+        if isinstance(other, (int, float)):
+            return Func(
+                "floordiv", lambda a: _floordiv(a, other), [self], result_type=int
+            )
+        return Func(
+            "floordiv", lambda a1, a2: _floordiv(a1, a2), [self, other], result_type=int
+        )
 
     def __rfloordiv__(self, other: Union[ColT, float]) -> "Func":
-        return math_floordiv(other, self)
+        if isinstance(other, (int, float)):
+            return Func(
+                "floordiv", lambda a: _floordiv(other, a), [self], result_type=int
+            )
+        return Func(
+            "floordiv", lambda a1, a2: _floordiv(a1, a2), [other, self], result_type=int
+        )
 
     def __mod__(self, other: Union[ColT, float]) -> "Func":
-        return math_mod(self, other)
+        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":
-        return math_mod(other, self)
-
-    def __pow__(self, other: Union[ColT, float]) -> "Func":
-        return math_pow(self, other)
-
-    def __rpow__(self, other: Union[ColT, float]) -> "Func":
-        return math_pow(other, self)
-
-    def __lshift__(self, other: Union[ColT, float]) -> "Func":
-        return math_lshift(self, other)
-
-    def __rlshift__(self, other: Union[ColT, float]) -> "Func":
-        return math_lshift(other, self)
-
-    def __rshift__(self, other: Union[ColT, float]) -> "Func":
-        return math_rshift(self, other)
-
-    def __rrshift__(self, other: Union[ColT, float]) -> "Func":
-        return math_rshift(other, self)
+        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":
-        return math_and(self, other)
+        if isinstance(other, (int, float)):
+            return Func(
+                "and", lambda a: numeric.bit_and(a, other), [self], result_type=int
+            )
+        return Func(
+            "and",
+            lambda a1, a2: numeric.bit_and(a1, a2),
+            [self, other],
+            result_type=int,
+        )
 
     def __rand__(self, other: Union[ColT, float]) -> "Func":
-        return math_and(other, self)
+        if isinstance(other, (int, float)):
+            return Func(
+                "and", lambda a: numeric.bit_and(other, a), [self], result_type=int
+            )
+        return Func(
+            "and",
+            lambda a1, a2: numeric.bit_and(a1, a2),
+            [other, self],
+            result_type=int,
+        )
 
     def __or__(self, other: Union[ColT, float]) -> "Func":
-        return math_or(self, other)
+        if isinstance(other, (int, float)):
+            return Func(
+                "or", lambda a: numeric.bit_or(a, other), [self], result_type=int
+            )
+        return Func(
+            "or", lambda a1, a2: numeric.bit_or(a1, a2), [self, other], result_type=int
+        )
 
     def __ror__(self, other: Union[ColT, float]) -> "Func":
-        return math_or(other, self)
+        if isinstance(other, (int, float)):
+            return Func(
+                "or", lambda a: numeric.bit_or(other, a), [self], result_type=int
+            )
+        return Func(
+            "or", lambda a1, a2: numeric.bit_or(a1, a2), [other, self], result_type=int
+        )
 
     def __xor__(self, other: Union[ColT, float]) -> "Func":
-        return math_xor(self, other)
+        if isinstance(other, (int, float)):
+            return Func(
+                "xor", lambda a: numeric.bit_xor(a, other), [self], result_type=int
+            )
+        return Func(
+            "xor",
+            lambda a1, a2: numeric.bit_xor(a1, a2),
+            [self, other],
+            result_type=int,
+        )
 
     def __rxor__(self, other: Union[ColT, float]) -> "Func":
-        return math_xor(other, self)
+        if isinstance(other, (int, float)):
+            return Func(
+                "xor", lambda a: numeric.bit_xor(other, a), [self], result_type=int
+            )
+        return Func(
+            "xor",
+            lambda a1, a2: numeric.bit_xor(a1, a2),
+            [other, self],
+            result_type=int,
+        )
+
+    def __rshift__(self, other: Union[ColT, float]) -> "Func":
+        if isinstance(other, (int, float)):
+            return Func(
+                "rshift",
+                lambda a: numeric.bit_rshift(a, other),
+                [self],
+                result_type=int,
+            )
+        return Func(
+            "rshift",
+            lambda a1, a2: numeric.bit_rshift(a1, a2),
+            [self, other],
+            result_type=int,
+        )
+
+    def __rrshift__(self, other: Union[ColT, float]) -> "Func":
+        if isinstance(other, (int, float)):
+            return Func(
+                "rshift",
+                lambda a: numeric.bit_rshift(other, a),
+                [self],
+                result_type=int,
+            )
+        return Func(
+            "rshift",
+            lambda a1, a2: numeric.bit_rshift(a1, a2),
+            [other, self],
+            result_type=int,
+        )
+
+    def __lshift__(self, other: Union[ColT, float]) -> "Func":
+        if isinstance(other, (int, float)):
+            return Func(
+                "lshift",
+                lambda a: numeric.bit_lshift(a, other),
+                [self],
+                result_type=int,
+            )
+        return Func(
+            "lshift",
+            lambda a1, a2: numeric.bit_lshift(a1, a2),
+            [self, other],
+            result_type=int,
+        )
+
+    def __rlshift__(self, other: Union[ColT, float]) -> "Func":
+        if isinstance(other, (int, float)):
+            return Func(
+                "lshift",
+                lambda a: numeric.bit_lshift(other, a),
+                [self],
+                result_type=int,
+            )
+        return Func(
+            "lshift",
+            lambda a1, a2: numeric.bit_lshift(a1, a2),
+            [other, self],
+            result_type=int,
+        )
 
     def __lt__(self, other: Union[ColT, float]) -> "Func":
-        return math_lt(self, other)
+        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":
-        return math_le(self, other)
+        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)
 
     def __eq__(self, other):
-        return math_eq(self, other)
+        if isinstance(other, (int, float)):
+            return Func("eq", lambda a: a == other, [self], result_type=bool)
+        return Func("eq", lambda a1, a2: a1 == a2, [self, other], result_type=bool)
 
     def __ne__(self, other):
-        return math_ne(self, other)
+        if isinstance(other, (int, float)):
+            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":
-        return math_gt(self, other)
+        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":
-        return math_ge(self, other)
+        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)
 
     def label(self, label: str) -> "Func":
         return Func(
@@ -283,107 +423,12 @@ def get_db_col_type(signals_schema: "SignalSchema", col: ColT) -> "DataType":
     )
 
 
-def math_func(
-    name: str,
-    inner: Callable,
-    params: Sequence[Union[ColT, float]],
-    result_type: Optional["DataType"] = None,
-) -> Func:
-    """Returns math function from the columns."""
-    cols, args = [], []
-    for arg in params:
-        if isinstance(arg, (int, float)):
-            args.append(arg)
-        else:
-            cols.append(arg)
-    return Func(name, inner, cols=cols, args=args, result_type=result_type)
-
-
-def math_add(*args: Union[ColT, float]) -> Func:
-    """Computes the sum of the column."""
-    return math_func("add", lambda a1, a2: a1 + a2, args)
-
-
-def math_sub(*args: Union[ColT, float]) -> Func:
-    """Computes the diff of the column."""
-    return math_func("sub", lambda a1, a2: a1 - a2, args)
-
-
-def math_mul(*args: Union[ColT, float]) -> Func:
-    """Computes the product of the column."""
-    return math_func("mul", lambda a1, a2: a1 * a2, args)
-
-
-def math_truediv(*args: Union[ColT, float]) -> Func:
-    """Computes the division of the column."""
-    return math_func("div", lambda a1, a2: a1 / a2, args, result_type=float)
-
-
-def math_floordiv(*args: Union[ColT, float]) -> Func:
-    """Computes the floor division of the column."""
-    return math_func("floordiv", lambda a1, a2: a1 // a2, args, result_type=float)
-
-
-def math_mod(*args: Union[ColT, float]) -> Func:
-    """Computes the modulo of the column."""
-    return math_func("mod", lambda a1, a2: a1 % a2, args, result_type=float)
-
-
-def math_pow(*args: Union[ColT, float]) -> Func:
-    """Computes the power of the column."""
-    return math_func("pow", lambda a1, a2: a1**a2, args, result_type=float)
-
-
-def math_lshift(*args: Union[ColT, float]) -> Func:
-    """Computes the left shift of the column."""
-    return math_func("lshift", lambda a1, a2: a1 << a2, args, result_type=int)
-
-
-def math_rshift(*args: Union[ColT, float]) -> Func:
-    """Computes the right shift of the column."""
-    return math_func("rshift", lambda a1, a2: a1 >> a2, args, result_type=int)
-
-
-def math_and(*args: Union[ColT, float]) -> Func:
-    """Computes the logical AND of the column."""
-    return math_func("and", lambda a1, a2: a1 & a2, args, result_type=bool)
-
-
-def math_or(*args: Union[ColT, float]) -> Func:
-    """Computes the logical OR of the column."""
-    return math_func("or", lambda a1, a2: a1 | a2, args, result_type=bool)
-
-
-def math_xor(*args: Union[ColT, float]) -> Func:
-    """Computes the logical XOR of the column."""
-    return math_func("xor", lambda a1, a2: a1 ^ a2, args, result_type=bool)
-
-
-def math_lt(*args: Union[ColT, float]) -> Func:
-    """Computes the less than comparison of the column."""
-    return math_func("lt", lambda a1, a2: a1 < a2, args, result_type=bool)
-
-
-def math_le(*args: Union[ColT, float]) -> Func:
-    """Computes the less than or equal comparison of the column."""
-    return math_func("le", lambda a1, a2: a1 <= a2, args, result_type=bool)
-
-
-def math_eq(*args: Union[ColT, float]) -> Func:
-    """Computes the equality comparison of the column."""
-    return math_func("eq", lambda a1, a2: a1 == a2, args, result_type=bool)
-
-
-def math_ne(*args: Union[ColT, float]) -> Func:
-    """Computes the inequality comparison of the column."""
-    return math_func("ne", lambda a1, a2: a1 != a2, args, result_type=bool)
-
-
-def math_gt(*args: Union[ColT, float]) -> Func:
-    """Computes the greater than comparison of the column."""
-    return math_func("gt", lambda a1, a2: a1 > a2, args, result_type=bool)
+def _truediv(a, b):
+    # Using sqlalchemy.sql.func.divide here instead of / operator
+    # because of a bug in ClickHouse SQLAlchemy dialect
+    # See https://github.com/xzkostyan/clickhouse-sqlalchemy/issues/335
+    return sa_func.divide(a, b)
 
 
-def math_ge(*args: Union[ColT, float]) -> Func:
-    """Computes the greater than or equal comparison of the column."""
-    return math_func("ge", lambda a1, a2: a1 >= a2, args, result_type=bool)
+def _floordiv(a, b):
+    return cast(_truediv(a, b), Integer)

datachain-0.7.8/src/datachain/func/numeric.py

@@ -0,0 +1,162 @@
+from typing import Union
+
+from datachain.sql.functions import numeric
+
+from .func import ColT, Func
+
+
+def bit_and(*args: Union[ColT, int]) -> Func:
+    """
+    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.
+            If an integer is provided, it is assumed to be a constant value.
+
+    Returns:
+        Func: A Func object that represents the bitwise AND function.
+
+    Example:
+        ```py
+        dc.mutate(
+            xor1=func.bit_and("signal.values", 0x0F),
+        )
+        ```
+
+    Notes:
+        - Result column will always be of type int.
+    """
+    cols, func_args = [], []
+    for arg in args:
+        if isinstance(arg, int):
+            func_args.append(arg)
+        else:
+            cols.append(arg)
+
+    if len(cols) + len(func_args) != 2:
+        raise ValueError("bit_and() requires exactly two arguments")
+
+    return Func(
+        "bit_and",
+        inner=numeric.bit_and,
+        cols=cols,
+        args=func_args,
+        result_type=int,
+    )
+
+
+def bit_or(*args: Union[ColT, int]) -> Func:
+    """
+    Computes the bitwise AND 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.
+            If an integer is provided, it is assumed to be a constant value.
+
+    Returns:
+        Func: A Func object that represents the bitwise OR function.
+
+    Example:
+        ```py
+        dc.mutate(
+            xor1=func.bit_or("signal.values", 0x0F),
+        )
+        ```
+
+    Notes:
+        - Result column will always be of type int.
+    """
+    cols, func_args = [], []
+    for arg in args:
+        if isinstance(arg, int):
+            func_args.append(arg)
+        else:
+            cols.append(arg)
+
+    if len(cols) + len(func_args) != 2:
+        raise ValueError("bit_or() requires exactly two arguments")
+
+    return Func(
+        "bit_or",
+        inner=numeric.bit_or,
+        cols=cols,
+        args=func_args,
+        result_type=int,
+    )
+
+
+def bit_xor(*args: Union[ColT, int]) -> Func:
+    """
+    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.
+            If an integer is provided, it is assumed to be a constant value.
+
+    Returns:
+        Func: A Func object that represents the bitwise XOR function.
+
+    Example:
+        ```py
+        dc.mutate(
+            xor1=func.bit_xor("signal.values", 0x0F),
+        )
+        ```
+
+    Notes:
+        - Result column will always be of type int.
+    """
+    cols, func_args = [], []
+    for arg in args:
+        if isinstance(arg, int):
+            func_args.append(arg)
+        else:
+            cols.append(arg)
+
+    if len(cols) + len(func_args) != 2:
+        raise ValueError("bit_xor() requires exactly two arguments")
+
+    return Func(
+        "bit_xor",
+        inner=numeric.bit_xor,
+        cols=cols,
+        args=func_args,
+        result_type=int,
+    )
+
+
+def int_hash_64(col: Union[ColT, int]) -> Func:
+    """
+    Returns the 64-bit hash of an integer.
+
+    Args:
+        col (str | int): String 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 Func is provided, it is assumed to be a function returning an int.
+
+    Returns:
+        Func: A Func object that represents the 64-bit hash function.
+
+    Example:
+        ```py
+        dc.mutate(
+            val_hash=func.int_hash_64("val"),
+        )
+        ```
+
+    Note:
+        - Result column will always be of type int.
+    """
+    cols, args = [], []
+    if isinstance(col, int):
+        args.append(col)
+    else:
+        cols.append(col)
+
+    return Func(
+        "int_hash_64", inner=numeric.int_hash_64, cols=cols, args=args, result_type=int
+    )

{datachain-0.7.7 → datachain-0.7.8}/src/datachain/sql/functions/array.py

@@ -38,6 +38,10 @@ class length(GenericFunction):  # noqa: N801
 
 
 class sip_hash_64(GenericFunction):  # noqa: N801
+    """
+    Computes the SipHash-64 hash of the array.
+    """
+
     type = Int64()
     package = "hash"
     name = "sip_hash_64"

datachain-0.7.8/src/datachain/sql/functions/numeric.py

@@ -0,0 +1,43 @@
+from sqlalchemy.sql.functions import GenericFunction, ReturnTypeFromArgs
+
+from datachain.sql.types import Int64
+from datachain.sql.utils import compiler_not_implemented
+
+
+class bit_and(ReturnTypeFromArgs):  # noqa: N801
+    inherit_cache = True
+
+
+class bit_or(ReturnTypeFromArgs):  # noqa: N801
+    inherit_cache = True
+
+
+class bit_xor(ReturnTypeFromArgs):  # noqa: N801
+    inherit_cache = True
+
+
+class bit_rshift(ReturnTypeFromArgs):  # noqa: N801
+    inherit_cache = True
+
+
+class bit_lshift(ReturnTypeFromArgs):  # noqa: N801
+    inherit_cache = True
+
+
+class int_hash_64(GenericFunction):  # noqa: N801
+    """
+    Computes the 64-bit hash of an integer.
+    """
+
+    type = Int64()
+    package = "hash"
+    name = "int_hash_64"
+    inherit_cache = True
+
+
+compiler_not_implemented(bit_and)
+compiler_not_implemented(bit_or)
+compiler_not_implemented(bit_xor)
+compiler_not_implemented(bit_rshift)
+compiler_not_implemented(bit_lshift)
+compiler_not_implemented(int_hash_64)