datachain 0.16.2__tar.gz → 0.16.3__tar.gz

{datachain-0.16.2/src/datachain.egg-info → datachain-0.16.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.16.2
+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.2 → datachain-0.16.3}/pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
   "numpy>=1,<3",
   "pandas>=2.0.0",
   "packaging",
-  "pyarrow",
+  "pyarrow<20",
   "typing-extensions",
   "python-dateutil>=2",
   "attrs>=21.3.0",

{datachain-0.16.2 → datachain-0.16.3}/src/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-0.16.2 → datachain-0.16.3}/src/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-0.16.2 → datachain-0.16.3}/src/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-0.16.2 → datachain-0.16.3}/src/datachain/lib/dc/records.py

@@ -37,6 +37,9 @@ 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 INSERT_BATCH_SIZE, adjust_outputs, get_col_types
     from datachain.sql.types import SQLType

{datachain-0.16.2 → datachain-0.16.3}/src/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-0.16.2 → datachain-0.16.3}/src/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.2 → datachain-0.16.3/src/datachain.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.16.2
+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.2 → datachain-0.16.3}/src/datachain.egg-info/SOURCES.txt

@@ -266,6 +266,7 @@ tests/func/test_dataset_query.py
 tests/func/test_datasets.py
 tests/func/test_feature_pickling.py
 tests/func/test_file.py
+tests/func/test_func.py
 tests/func/test_hf.py
 tests/func/test_hidden_field.py
 tests/func/test_image.py

{datachain-0.16.2 → datachain-0.16.3}/src/datachain.egg-info/requires.txt

@@ -4,7 +4,7 @@ tqdm
 numpy<3,>=1
 pandas>=2.0.0
 packaging
-pyarrow
+pyarrow<20
 typing-extensions
 python-dateutil>=2
 attrs>=21.3.0

datachain-0.16.3/tests/func/test_func.py

@@ -0,0 +1,124 @@
+import math
+
+import datachain as dc
+from datachain import func
+from datachain.sql.types import Float, Int, String
+
+
+def values_almost_equal(a, b):
+    """Compare two values, treating NaNs as equal."""
+    if (
+        isinstance(a, float)
+        and isinstance(b, float)
+        and math.isnan(a)
+        and math.isnan(b)
+    ):
+        return True
+    return a == b
+
+
+def tuples_almost_equal(t1, t2):
+    """Compare two tuples, treating NaN floats as equal."""
+    if len(t1) != len(t2):
+        return False
+    return all(values_almost_equal(x, y) for x, y in zip(t1, t2))
+
+
+def sets_of_tuples_almost_equal(s1, s2):
+    """Compare two sets of tuples, treating NaN floats as equal."""
+    if len(s1) != len(s2):
+        return False
+    unmatched = list(s2)
+    for item1 in s1:
+        for item2 in unmatched:
+            if tuples_almost_equal(item1, item2):
+                unmatched.remove(item2)
+                break
+        else:
+            return False
+    return True
+
+
+def test_array_get_element(test_session):
+    db_dialect = test_session.catalog.warehouse.db.dialect
+
+    class Arr(dc.DataModel):
+        i: list[int]
+        f: list[float]
+
+    ds = (
+        dc.read_values(
+            arr=(
+                Arr(i=[10, 20, 30], f=[1.0, 2.0, 3.0]),
+                Arr(i=[40, 50, 60], f=[4.0, 5.0, 6.0]),
+                Arr(i=[50], f=[5.0]),
+            ),
+            session=test_session,
+        )
+        .mutate(
+            first_i=func.array.get_element("arr.i", 0),
+            second_i=func.array.get_element("arr.i", 1),
+            unknown_i=func.array.get_element("arr.i", 100),
+            first_f=func.array.get_element("arr.f", 0),
+            second_f=func.array.get_element("arr.f", 1),
+            first_f2=func.array.get_element([9.0], 0),
+            first_s=func.array.get_element(["a", "b", "c", "d"], 0),
+            second_s=func.array.get_element(["a", "b", "c", "d"], 1),
+            unknown_s=func.array.get_element(["a", "b", "c", "d"], 100),
+            unknown=func.array.get_element([], 0),
+        )
+        .collect(
+            "first_i",
+            "second_i",
+            "unknown_i",
+            "first_f",
+            "second_f",
+            "first_f2",
+            "first_s",
+            "second_s",
+            "unknown_s",
+            "unknown",
+        )
+    )
+
+    assert sets_of_tuples_almost_equal(
+        set(ds),
+        {
+            (
+                10,
+                20,
+                Int.default_value(db_dialect),
+                1.0,
+                2.0,
+                9.0,
+                "a",
+                "b",
+                String.default_value(db_dialect),
+                String.default_value(db_dialect),
+            ),
+            (
+                40,
+                50,
+                Int.default_value(db_dialect),
+                4.0,
+                5.0,
+                9.0,
+                "a",
+                "b",
+                String.default_value(db_dialect),
+                String.default_value(db_dialect),
+            ),
+            (
+                50,
+                Int.default_value(db_dialect),
+                Int.default_value(db_dialect),
+                5.0,
+                Float.default_value(db_dialect),
+                9.0,
+                "a",
+                "b",
+                String.default_value(db_dialect),
+                String.default_value(db_dialect),
+            ),
+        },
+    )

{datachain-0.16.2 → datachain-0.16.3}/tests/unit/lib/test_datachain.py

@@ -372,10 +372,6 @@ def test_datasets_in_memory():
     ],
 )
 def test_datasets_filtering(test_session, attrs, result):
-    ds = dc.datasets(column="dataset", session=test_session)
-    datasets = [d for d in ds.collect("dataset") if d.name == "fibonacci"]
-    assert len(datasets) == 0
-
     dc.read_values(num=[1, 2, 3], session=test_session).save(
         "primes", attrs=["number", "num=prime", "small"]
     )

{datachain-0.16.2 → datachain-0.16.3}/tests/unit/sql/test_array.py

@@ -5,6 +5,7 @@ from numpy.testing import assert_array_almost_equal
 
 from datachain import func
 from datachain.sql import select
+from datachain.sql.types import Int, String
 
 
 def test_cosine_distance(warehouse):
@@ -70,6 +71,40 @@ def test_length(warehouse):
     assert result == ((4, 5, 2),)
 
 
+def test_get_element(warehouse):
+    db_dialect = warehouse.db.dialect
+
+    query = select(
+        func.array.get_element(["abc", "def", "g", "hi"], 0).label("first1"),
+        func.array.get_element(["abc", "def", "g", "hi"], 1).label("second1"),
+        func.array.get_element([3.0, 5.0, 1.0, 6.0, 1.0], 0).label("first2"),
+        func.array.get_element([3.0, 5.0, 1.0, 6.0, 1.0], 1).label("second2"),
+        func.array.get_element([1, 2, 3, 4, 5, 6], 0).label("first3"),
+        func.array.get_element([1, 2, 3, 4, 5, 6], 1).label("second3"),
+        func.array.get_element([1], 0).label("first4"),
+        func.array.get_element([2.0], 0).label("first5"),
+        func.array.get_element([], 0).label("not_found1"),
+        func.array.get_element([], -1).label("not_found2"),
+        func.array.get_element([1], 2).label("not_found3"),
+    )
+    result = tuple(warehouse.dataset_rows_select(query))
+    assert result == (
+        (
+            "abc",
+            "def",
+            3.0,
+            5.0,
+            1,
+            2,
+            1,
+            2.0,
+            String.default_value(db_dialect),
+            String.default_value(db_dialect),
+            Int.default_value(db_dialect),
+        ),
+    )
+
+
 def test_contains(warehouse):
     query = select(
         func.contains(["abc", "def", "g", "hi"], "abc").label("contains1"),