datachain 0.8.12__py3-none-any.whl → 0.8.13__py3-none-any.whl

datachain/func/__init__.py

@@ -15,7 +15,7 @@ from .aggregate import (
     row_number,
     sum,
 )
-from .array import cosine_distance, euclidean_distance, length, sip_hash_64
+from .array import contains, cosine_distance, euclidean_distance, length, sip_hash_64
 from .conditional import case, greatest, ifelse, isnone, least
 from .numeric import bit_and, bit_hamming_distance, bit_or, bit_xor, int_hash_64
 from .random import rand
@@ -34,6 +34,7 @@ __all__ = [
     "case",
     "collect",
     "concat",
+    "contains",
     "cosine_distance",
     "count",
     "dense_rank",

datachain/func/array.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import Union
+from typing import Any, Union
 
 from datachain.sql.functions import array
 
@@ -140,6 +140,44 @@ def length(arg: Union[str, Sequence, Func]) -> Func:
     return Func("length", inner=array.length, cols=cols, args=args, result_type=int)
 
 
+def contains(arr: Union[str, Sequence, Func], elem: Any) -> Func:
+    """
+    Checks whether the `arr` array has the `elem` element.
+
+    Args:
+        arr (str | Sequence | Func): Array to check for the element.
+            If a string is provided, it is assumed to be the name of the array column.
+            If a sequence is provided, it is assumed to be an array of values.
+            If a Func is provided, it is assumed to be a function returning an array.
+        elem (Any): Element to check for in the array.
+
+    Returns:
+        Func: A Func object that represents the contains function. Result of the
+            function will be 1 if the element is present in the array, and 0 otherwise.
+
+    Example:
+        ```py
+        dc.mutate(
+            contains1=func.array.contains("signal.values", 3),
+            contains2=func.array.contains([1, 2, 3, 4, 5], 7),
+        )
+        ```
+    """
+
+    def inner(arg):
+        is_json = type(elem) in [list, dict]
+        return array.contains(arg, elem, is_json)
+
+    if isinstance(arr, (str, Func)):
+        cols = [arr]
+        args = None
+    else:
+        cols = None
+        args = [arr]
+
+    return Func("contains", inner=inner, cols=cols, args=args, result_type=int)
+
+
 def sip_hash_64(arg: Union[str, Sequence]) -> Func:
     """
     Computes the SipHash-64 hash of the array.

datachain/sql/functions/array.py

@@ -1,6 +1,6 @@
 from sqlalchemy.sql.functions import GenericFunction
 
-from datachain.sql.types import Float, Int64
+from datachain.sql.types import Boolean, Float, Int64
 from datachain.sql.utils import compiler_not_implemented
 
 
@@ -37,6 +37,17 @@ class length(GenericFunction):  # noqa: N801
     inherit_cache = True
 
 
+class contains(GenericFunction):  # noqa: N801
+    """
+    Checks if element is in the array.
+    """
+
+    type = Boolean()
+    package = "array"
+    name = "contains"
+    inherit_cache = True
+
+
 class sip_hash_64(GenericFunction):  # noqa: N801
     """
     Computes the SipHash-64 hash of the array.
@@ -51,4 +62,5 @@ class sip_hash_64(GenericFunction):  # noqa: N801
 compiler_not_implemented(cosine_distance)
 compiler_not_implemented(euclidean_distance)
 compiler_not_implemented(length)
+compiler_not_implemented(contains)
 compiler_not_implemented(sip_hash_64)

datachain/sql/sqlite/base.py

@@ -87,6 +87,7 @@ def setup():
     compiles(sql_path.file_stem, "sqlite")(compile_path_file_stem)
     compiles(sql_path.file_ext, "sqlite")(compile_path_file_ext)
     compiles(array.length, "sqlite")(compile_array_length)
+    compiles(array.contains, "sqlite")(compile_array_contains)
     compiles(string.length, "sqlite")(compile_string_length)
     compiles(string.split, "sqlite")(compile_string_split)
     compiles(string.regexp_replace, "sqlite")(compile_string_regexp_replace)
@@ -269,13 +270,16 @@ def register_user_defined_sql_functions() -> None:
 
     _registered_function_creators["string_functions"] = create_string_functions
 
-    has_json_extension = functions_exist(["json_array_length"])
+    has_json_extension = functions_exist(["json_array_length", "json_array_contains"])
     if not has_json_extension:
 
         def create_json_functions(conn):
             conn.create_function(
                 "json_array_length", 1, py_json_array_length, deterministic=True
             )
+            conn.create_function(
+                "json_array_contains", 3, py_json_array_contains, deterministic=True
+            )
 
         _registered_function_creators["json_functions"] = create_json_functions
 
@@ -428,10 +432,22 @@ def py_json_array_length(arr):
     return len(orjson.loads(arr))
 
 
+def py_json_array_contains(arr, value, is_json):
+    if is_json:
+        value = orjson.loads(value)
+    return value in orjson.loads(arr)
+
+
 def compile_array_length(element, compiler, **kwargs):
     return compiler.process(func.json_array_length(*element.clauses.clauses), **kwargs)
 
 
+def compile_array_contains(element, compiler, **kwargs):
+    return compiler.process(
+        func.json_array_contains(*element.clauses.clauses), **kwargs
+    )
+
+
 def compile_string_length(element, compiler, **kwargs):
     return compiler.process(func.length(*element.clauses.clauses), **kwargs)
 

datachain/sql/sqlite/types.py

@@ -31,6 +31,10 @@ def adapt_array(arr):
     return orjson.dumps(arr).decode("utf-8")
 
 
+def adapt_dict(dct):
+    return orjson.dumps(dct).decode("utf-8")
+
+
 def convert_array(arr):
     return orjson.loads(arr)
 
@@ -52,6 +56,7 @@ def adapt_np_generic(val):
 
 def register_type_converters():
     sqlite3.register_adapter(list, adapt_array)
+    sqlite3.register_adapter(dict, adapt_dict)
     sqlite3.register_converter("ARRAY", convert_array)
     if numpy_imported:
         sqlite3.register_adapter(np.ndarray, adapt_np_array)

{datachain-0.8.12.dist-info → datachain-0.8.13.dist-info}/METADATA

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

{datachain-0.8.12.dist-info → datachain-0.8.13.dist-info}/RECORD

@@ -50,9 +50,9 @@ datachain/data_storage/serializer.py,sha256=6G2YtOFqqDzJf1KbvZraKGXl2XHZyVml2kru
 datachain/data_storage/sqlite.py,sha256=KJ8hI0Hrwv9eAA-nLUlw2AYCQxiAAZ12a-ftUBtroNQ,24545
 datachain/data_storage/warehouse.py,sha256=ovdH9LmOWLfCrvf0UvXnrNC-CrdAjns3EmXEgFdz4KM,30824
 datachain/diff/__init__.py,sha256=OapNRBsyGDOQHelefUEoXoFHRWCJuBnhvD0ibebKvBc,10486
-datachain/func/__init__.py,sha256=qaSjakSaTsRtnU7Hcb4lJk71tbwk7M0oWmjRqXExCLA,1099
+datachain/func/__init__.py,sha256=DDpkbK6Kg53ONVGZo8szaNEFJrcepCAG9ev27DC6_7E,1125
 datachain/func/aggregate.py,sha256=7_IPrIwb2XSs3zG4iOr1eTvzn6kNVe2mkzvNzjusDHk,10942
-datachain/func/array.py,sha256=zHDNWuWLA7HVa9FEvQeHhVi00_xqenyleTqcLwkXWBI,5477
+datachain/func/array.py,sha256=O784_uwmaP5CjZX4VSF4RmS8cmpaForQc8zASxHJB6A,6717
 datachain/func/base.py,sha256=wA0sBQAVyN9LPxoo7Ox83peS0zUVnyuKxukwAcjGLfY,534
 datachain/func/conditional.py,sha256=g46zwW-i87uA45zWJnPHtHaqr6qOXSg6xLb4p9W3Gtk,6400
 datachain/func/func.py,sha256=PnwTRAiEJUus3e4NYdQ-hldqLzKS9hY0FjiyBMZhsSo,16183
@@ -120,22 +120,22 @@ datachain/sql/default/__init__.py,sha256=XQ2cEZpzWiABqjV-6yYHUBGI9vN_UHxbxZENESm
 datachain/sql/default/base.py,sha256=QD-31C6JnyOXzogyDx90sUhm7QvgXIYpeHEASH84igU,628
 datachain/sql/functions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/sql/functions/aggregate.py,sha256=3AQdA8YHPFdtCEfwZKQXTT8SlQWdG9gD5PBtGN3Odqs,944
-datachain/sql/functions/array.py,sha256=Zq59CaMHf_hFapU4kxvy2mwteH344k5Wksxja4MfBks,1204
+datachain/sql/functions/array.py,sha256=ZXPHNQxG-t5tl7Gr0nubu4oc-9fBZmX5q03z9ljWg-I,1443
 datachain/sql/functions/conditional.py,sha256=q7YUKfunXeEldXaxgT-p5pUTcOEVU_tcQ2BJlquTRPs,207
 datachain/sql/functions/numeric.py,sha256=BK2KCiPSgM2IveCq-9M_PG3CtPBlztaS9TTn1LGzyLs,1250
 datachain/sql/functions/path.py,sha256=zixpERotTFP6LZ7I4TiGtyRA8kXOoZmH1yzH9oRW0mg,1294
 datachain/sql/functions/random.py,sha256=vBwEEj98VH4LjWixUCygQ5Bz1mv1nohsCG0-ZTELlVg,271
 datachain/sql/functions/string.py,sha256=E-T9OIzUR-GKaLgjZsEtg5CJrY_sLf1lt1awTvY7w2w,1426
 datachain/sql/sqlite/__init__.py,sha256=TAdJX0Bg28XdqPO-QwUVKy8rg78cgMileHvMNot7d04,166
-datachain/sql/sqlite/base.py,sha256=bPrYfj2ZF9hFZFs0chgH7J5l_tdXI4VMZMgkuBjf7Ng,19070
-datachain/sql/sqlite/types.py,sha256=lPXS1XbkmUtlkkiRxy_A_UzsgpPv2VSkXYOD4zIHM4w,1734
+datachain/sql/sqlite/base.py,sha256=Rfemu8pj7V0aWhWwryDghhnbiMFfQS5X9FCihGuplb8,19593
+datachain/sql/sqlite/types.py,sha256=cH6oge2E_YWFy22wY-txPJH8gxoQFSpCthtZR8PZjpo,1849
 datachain/sql/sqlite/vector.py,sha256=ncW4eu2FlJhrP_CIpsvtkUabZlQdl2D5Lgwy_cbfqR0,469
 datachain/toolkit/__init__.py,sha256=eQ58Q5Yf_Fgv1ZG0IO5dpB4jmP90rk8YxUWmPc1M2Bo,68
 datachain/toolkit/split.py,sha256=z3zRJNzjWrpPuRw-zgFbCOBKInyYxJew8ygrYQRQLNc,2930
 datachain/torch/__init__.py,sha256=gIS74PoEPy4TB3X6vx9nLO0Y3sLJzsA8ckn8pRWihJM,579
-datachain-0.8.12.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.8.12.dist-info/METADATA,sha256=C1vaFTVw44GIVe32CcfLthfCi5nbbqTgS7HL61iSFGg,10880
-datachain-0.8.12.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-datachain-0.8.12.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.8.12.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.8.12.dist-info/RECORD,,
+datachain-0.8.13.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.8.13.dist-info/METADATA,sha256=ugNjNgfRl-dnR4tKM3AqiRHEmPGFXSYjA-7Bl1BRgOA,10880
+datachain-0.8.13.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+datachain-0.8.13.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.8.13.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.8.13.dist-info/RECORD,,