datachain 0.16.1__py3-none-any.whl → 0.16.3__py3-none-any.whl

datachain/func/array.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 from datachain.sql.functions import array
 
@@ -178,6 +178,61 @@ def contains(arr: Union[str, Sequence, Func], elem: Any) -> Func:
     return Func("contains", inner=inner, cols=cols, args=args, result_type=int)
 
 
+def get_element(arg: Union[str, Sequence, Func], index: int) -> Func:
+    """
+    Returns the element at the given index from the array.
+    If the index is out of bounds, it returns None or columns default value.
+
+    Args:
+        arg (str | Sequence | Func): Array to get the element from.
+            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.
+        index (int): Index of the element to get from the array.
+
+    Returns:
+        Func: A Func object that represents the array get_element function.
+
+    Example:
+        ```py
+        dc.mutate(
+            first_el=func.array.get_element("signal.values", 0),
+            second_el=func.array.get_element([1, 2, 3, 4, 5], 1),
+        )
+        ```
+
+    Note:
+        - Result column will always be the same type as the elements of the array.
+    """
+
+    def type_from_args(arr, _):
+        if isinstance(arr, list):
+            try:
+                return type(arr[0])
+            except IndexError:
+                return str  # if the array is empty, return str as default type
+        return None
+
+    cols: Optional[Union[str, Sequence, Func]]
+    args: Union[str, Sequence, Func, int]
+
+    if isinstance(arg, (str, Func)):
+        cols = [arg]
+        args = [index]
+    else:
+        cols = None
+        args = [arg, index]
+
+    return Func(
+        "get_element",
+        inner=array.get_element,
+        cols=cols,
+        args=args,
+        from_array=True,
+        type_from_args=type_from_args,
+    )
+
+
 def sip_hash_64(arg: Union[str, Sequence]) -> Func:
     """
     Computes the SipHash-64 hash of the array.

datachain/func/func.py

@@ -1,6 +1,6 @@
 import inspect
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union, get_args, get_origin
 
 from sqlalchemy import BindParameter, Case, ColumnElement, Integer, cast, desc
 from sqlalchemy.sql import func as sa_func
@@ -36,7 +36,9 @@ class Func(Function):
         args: Optional[Sequence[Any]] = None,
         kwargs: Optional[dict[str, Any]] = None,
         result_type: Optional["DataType"] = None,
+        type_from_args: Optional[Callable[..., "DataType"]] = None,
         is_array: bool = False,
+        from_array: bool = False,
         is_window: bool = False,
         window: Optional["Window"] = None,
         label: Optional[str] = None,
@@ -47,7 +49,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 +70,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,
@@ -101,6 +107,20 @@ class Func(Function):
                     "Columns must have the same type to infer result type",
                 )
 
+        if self.from_array:
+            if get_origin(col_type) is list:
+                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]
+            raise DataChainColumnError(
+                str(self),
+                "Array column must be of type list",
+            )
+
         return list[col_type] if self.is_array else col_type  # type: ignore[valid-type]
 
     def __add__(self, other: Union[ColT, float]) -> "Func":
@@ -339,7 +359,9 @@ 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,
@@ -368,6 +390,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",

datachain/lib/dc/database.py

@@ -127,9 +127,11 @@ def read_database(
         ```
 
     Notes:
-        This function works with a variety of databases — including, but not limited to,
-        SQLite, DuckDB, PostgreSQL, and Snowflake, provided the appropriate driver is
-        installed.
+        - This function works with a variety of databases — including,
+        but not limited to, SQLite, DuckDB, PostgreSQL, and Snowflake,
+        provided the appropriate driver is installed.
+        - This call is blocking, and will execute the query and return once the
+          results are saved.
     """
     from datachain.lib.dc.records import read_records
 

datachain/lib/dc/records.py

@@ -37,9 +37,13 @@ def read_records(
         import datachain as dc
         single_record = dc.read_records(dc.DEFAULT_FILE_RECORD)
         ```
+
+    Notes:
+        This call blocks until all records are inserted.
     """
-    from datachain.query.dataset import adjust_outputs, get_col_types
+    from datachain.query.dataset import INSERT_BATCH_SIZE, adjust_outputs, get_col_types
     from datachain.sql.types import SQLType
+    from datachain.utils import batched
 
     from .datasets import read_dataset
 
@@ -89,6 +93,7 @@ def read_records(
         {c.name: c.type for c in columns if isinstance(c.type, SQLType)},
     )
     records = (adjust_outputs(warehouse, record, col_types) for record in to_insert)
-    warehouse.insert_rows(table, records)
+    for chunk in batched(records, INSERT_BATCH_SIZE):
+        warehouse.insert_rows(table, chunk)
     warehouse.insert_rows_done(table)
     return read_dataset(name=dsr.name, session=session, settings=settings)

datachain/sql/functions/array.py

@@ -48,6 +48,16 @@ class contains(GenericFunction):  # noqa: N801
     inherit_cache = True
 
 
+class get_element(GenericFunction):  # noqa: N801
+    """
+    Returns the element at the given index in the array.
+    """
+
+    package = "array"
+    name = "get_element"
+    inherit_cache = True
+
+
 class sip_hash_64(GenericFunction):  # noqa: N801
     """
     Computes the SipHash-64 hash of the array.
@@ -63,4 +73,5 @@ compiler_not_implemented(cosine_distance)
 compiler_not_implemented(euclidean_distance)
 compiler_not_implemented(length)
 compiler_not_implemented(contains)
+compiler_not_implemented(get_element)
 compiler_not_implemented(sip_hash_64)

datachain/sql/sqlite/base.py

@@ -88,6 +88,7 @@ def setup():
     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(array.get_element, "sqlite")(compile_array_get_element)
     compiles(string.length, "sqlite")(compile_string_length)
     compiles(string.split, "sqlite")(compile_string_split)
     compiles(string.regexp_replace, "sqlite")(compile_string_regexp_replace)
@@ -270,6 +271,13 @@ def register_user_defined_sql_functions() -> None:
 
     _registered_function_creators["string_functions"] = create_string_functions
 
+    def create_array_functions(conn):
+        conn.create_function(
+            "json_array_get_element", 2, py_json_array_get_element, deterministic=True
+        )
+
+    _registered_function_creators["array_functions"] = create_array_functions
+
     has_json_extension = functions_exist(["json_array_length", "json_array_contains"])
     if not has_json_extension:
 
@@ -438,6 +446,20 @@ def py_json_array_contains(arr, value, is_json):
     return value in orjson.loads(arr)
 
 
+def py_json_array_get_element(val, idx):
+    arr = orjson.loads(val)
+    try:
+        return arr[idx]
+    except IndexError:
+        return None
+
+
+def compile_array_get_element(element, compiler, **kwargs):
+    return compiler.process(
+        func.json_array_get_element(*element.clauses.clauses), **kwargs
+    )
+
+
 def compile_array_length(element, compiler, **kwargs):
     return compiler.process(func.json_array_length(*element.clauses.clauses), **kwargs)
 

{datachain-0.16.1.dist-info → datachain-0.16.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.16.1
+Version: 0.16.3
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -23,7 +23,7 @@ Requires-Dist: tqdm
 Requires-Dist: numpy<3,>=1
 Requires-Dist: pandas>=2.0.0
 Requires-Dist: packaging
-Requires-Dist: pyarrow
+Requires-Dist: pyarrow<20
 Requires-Dist: typing-extensions
 Requires-Dist: python-dateutil>=2
 Requires-Dist: attrs>=21.3.0

{datachain-0.16.1.dist-info → datachain-0.16.3.dist-info}/RECORD

@@ -56,10 +56,10 @@ datachain/fs/reference.py,sha256=A8McpXF0CqbXPqanXuvpKu50YLB3a2ZXA3YAPxtBXSM,914
 datachain/fs/utils.py,sha256=s-FkTOCGBk-b6TT3toQH51s9608pofoFjUSTc1yy7oE,825
 datachain/func/__init__.py,sha256=CjNLHfJkepdXdRZ6HjJBjNSIjOeFMuMkwPDaPUrM75g,1270
 datachain/func/aggregate.py,sha256=UfxENlw56Qv3UEkj2sZ-JZHmr9q8Rnic9io9_63gF-E,10942
-datachain/func/array.py,sha256=O784_uwmaP5CjZX4VSF4RmS8cmpaForQc8zASxHJB6A,6717
+datachain/func/array.py,sha256=OmmjdK5AQyBXa_7NXFUPY_m3lFlRK4Um4J9NCtYwvak,8394
 datachain/func/base.py,sha256=wA0sBQAVyN9LPxoo7Ox83peS0zUVnyuKxukwAcjGLfY,534
 datachain/func/conditional.py,sha256=HkNamQr9dLyIMDEbIeO6CZR0emQoDqeaWrZ1fECod4M,8062
-datachain/func/func.py,sha256=k8z5tIiabEOPymYWGfz4O7z1qS6zBZnVYRPp_58iU7c,16192
+datachain/func/func.py,sha256=jzEvnc2iN0BAly-uzxhaoMntL_xF4j94DrFuRi2ADSw,17321
 datachain/func/numeric.py,sha256=gMe1Ks0dqQKHkjcpvj7I5S-neECzQ_gltPQLNoaWOyo,5632
 datachain/func/path.py,sha256=mqN_mfkwv44z2II7DMTp_fGGw95hmTCNls_TOFNpr4k,3155
 datachain/func/random.py,sha256=pENOLj9rSmWfGCnOsUIaCsVC5486zQb66qfQvXaz9Z4,452
@@ -96,7 +96,7 @@ datachain/lib/convert/unflatten.py,sha256=ysMkstwJzPMWUlnxn-Z-tXJR3wmhjHeSN_P-sD
 datachain/lib/convert/values_to_tuples.py,sha256=j5yZMrVUH6W7b-7yUvdCTGI7JCUAYUOzHUGPoyZXAB0,4360
 datachain/lib/dc/__init__.py,sha256=HD0NYrdy44u6kkpvgGjJcvGz-UGTHui2azghcT8ZUg0,838
 datachain/lib/dc/csv.py,sha256=asWPAxhMgIoLAdD2dObDlnGL8CTSD3TAuFuM4ci89bQ,4374
-datachain/lib/dc/database.py,sha256=gYKh1iO5hOWMPFTU1vZC5kOXkJzVse14TYTWE4_1iEA,5940
+datachain/lib/dc/database.py,sha256=g5M6NjYR1T0vKte-abV-3Ejnm-HqxTIMir5cRi_SziE,6051
 datachain/lib/dc/datachain.py,sha256=36J8QIB04hKKumQgLvHNTC94Pd7G2yE4slZ9RfwI9zw,76980
 datachain/lib/dc/datasets.py,sha256=u6hlz0Eodh_s39TOW6kz0VIL3nGfadqu8FLoWqDxSJs,6890
 datachain/lib/dc/hf.py,sha256=PJl2wiLjdRsMz0SYbLT-6H8b-D5i2WjeH7li8HHOk_0,2145
@@ -104,7 +104,7 @@ datachain/lib/dc/json.py,sha256=ZUThPDAaP2gBFIL5vsQTwKBcuN_dhvC_O44wdDv0jEc,2683
 datachain/lib/dc/listings.py,sha256=2na9v63xO1vPUNaoBSzA-TSN49V7zQAb-4iS1wOPLFE,1029
 datachain/lib/dc/pandas.py,sha256=ObueUXDUFKJGu380GmazdG02ARpKAHPhSaymfmOH13E,1489
 datachain/lib/dc/parquet.py,sha256=zYcSgrWwyEDW9UxGUSVdIVsCu15IGEf0xL8KfWQqK94,1782
-datachain/lib/dc/records.py,sha256=Z6EWy6c6hf87cWiDlQduvrDgOHMLwqF22g-XksOnXsU,2884
+datachain/lib/dc/records.py,sha256=J1I69J2gFIBjRTGr2LG-5qn_rTVzRLcr2y3tVDrmHdg,3068
 datachain/lib/dc/storage.py,sha256=QLf3-xMV2Gmy3AA8qF9WqAsb7R8Rk87l4s5hBoiCH98,5285
 datachain/lib/dc/utils.py,sha256=VawOAlJSvAtZbsMg33s5tJe21TRx1Km3QggI1nN6tnw,3984
 datachain/lib/dc/values.py,sha256=cBQubhmPNEDMJldUXzGh-UKbdim4P6O2B91Gp39roKw,1389
@@ -138,22 +138,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=ZXPHNQxG-t5tl7Gr0nubu4oc-9fBZmX5q03z9ljWg-I,1443
+datachain/sql/functions/array.py,sha256=LkVDq1Iu3pF2vz9opjxcS0oSvEepoLNYVXcmnOzsmY0,1679
 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=N-cQT0Hpu9ROWe4OiKlkkn_YP1NKCRZZ3xSfTzpyaDA,19651
+datachain/sql/sqlite/base.py,sha256=SktpdtyZmxG9ip_UX_0WL3YxP9o66CYTeMfriRrZzaE,20281
 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=ktGWzY4kyzjWyR86dhvzw-Zhl0lVk_LOX3NciTac6qo,2914
 datachain/torch/__init__.py,sha256=gIS74PoEPy4TB3X6vx9nLO0Y3sLJzsA8ckn8pRWihJM,579
-datachain-0.16.1.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.16.1.dist-info/METADATA,sha256=9YPqP6Sthuf_fuxFX3miQyp9MEjRq8j2DqubLXvZg0k,11328
-datachain-0.16.1.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
-datachain-0.16.1.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.16.1.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.16.1.dist-info/RECORD,,
+datachain-0.16.3.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.16.3.dist-info/METADATA,sha256=3XVhBEDIISei-EeF3LEpAimkbJ2vQ1yHAEKWZqgfVbs,11331
+datachain-0.16.3.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
+datachain-0.16.3.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.16.3.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.16.3.dist-info/RECORD,,

{datachain-0.16.1.dist-info → datachain-0.16.3.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (79.0.0)
+Generator: setuptools (80.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any