datachain 0.18.4__py3-none-any.whl → 0.18.5__py3-none-any.whl

datachain/func/array.py

@@ -1,14 +1,15 @@
 from collections.abc import Sequence
 from typing import Any, Optional, Union
 
+from datachain.query.schema import Column
 from datachain.sql.functions import array
 
 from .func import Func
 
 
-def cosine_distance(*args: Union[str, Sequence]) -> Func:
+def cosine_distance(*args: Union[str, Column, Func, Sequence]) -> Func:
     """
-    Computes the cosine distance between two vectors.
+    Returns the cosine distance between two vectors.
 
     The cosine distance is derived from the cosine similarity, which measures the angle
     between two vectors. This function returns the dissimilarity between the vectors,
@@ -16,29 +17,33 @@ def cosine_distance(*args: Union[str, Sequence]) -> Func:
     indicate higher dissimilarity.
 
     Args:
-        args (str | Sequence): Two vectors to compute the cosine distance between.
+        args (str | Column | Func | Sequence): Two vectors to compute the cosine
+            distance between.
             If a string is provided, it is assumed to be the name of the column vector.
+            If a Column is provided, it is assumed to be an array column.
+            If a Func is provided, it is assumed to be a function returning an array.
             If a sequence is provided, it is assumed to be a vector of values.
 
     Returns:
-        Func: A Func object that represents the cosine_distance function.
+        Func: A `Func` object that represents the cosine_distance function.
 
     Example:
         ```py
         target_embedding = [0.1, 0.2, 0.3]
         dc.mutate(
             cos_dist1=func.cosine_distance("embedding", target_embedding),
-            cos_dist2=func.cosine_distance(target_embedding, [0.4, 0.5, 0.6]),
+            cos_dist2=func.cosine_distance(dc.C("emb1"), "emb2"),
+            cos_dist3=func.cosine_distance(target_embedding, [0.4, 0.5, 0.6]),
         )
         ```
 
     Notes:
         - Ensure both vectors have the same number of elements.
-        - Result column will always be of type float.
+        - The result column will always be of type float.
     """
     cols, func_args = [], []
     for arg in args:
-        if isinstance(arg, str):
+        if isinstance(arg, (str, Column, Func)):
             cols.append(arg)
         else:
             func_args.append(list(arg))
@@ -57,37 +62,41 @@ def cosine_distance(*args: Union[str, Sequence]) -> Func:
     )
 
 
-def euclidean_distance(*args: Union[str, Sequence]) -> Func:
+def euclidean_distance(*args: Union[str, Column, Func, Sequence]) -> Func:
     """
-    Computes the Euclidean distance between two vectors.
+    Returns the Euclidean distance between two vectors.
 
     The Euclidean distance is the straight-line distance between two points
     in Euclidean space. This function returns the distance between the two vectors.
 
     Args:
-        args (str | Sequence): Two vectors to compute the Euclidean distance between.
+        args (str | Column | Func | Sequence): Two vectors to compute the Euclidean
+            distance between.
             If a string is provided, it is assumed to be the name of the column vector.
+            If a Column is provided, it is assumed to be an array column.
+            If a Func is provided, it is assumed to be a function returning an array.
             If a sequence is provided, it is assumed to be a vector of values.
 
     Returns:
-        Func: A Func object that represents the euclidean_distance function.
+        Func: A `Func` object that represents the euclidean_distance function.
 
     Example:
         ```py
         target_embedding = [0.1, 0.2, 0.3]
         dc.mutate(
             eu_dist1=func.euclidean_distance("embedding", target_embedding),
-            eu_dist2=func.euclidean_distance(target_embedding, [0.4, 0.5, 0.6]),
+            eu_dist2=func.euclidean_distance(dc.C("emb1"), "emb2"),
+            eu_dist3=func.euclidean_distance(target_embedding, [0.4, 0.5, 0.6]),
         )
         ```
 
     Notes:
         - Ensure both vectors have the same number of elements.
-        - Result column will always be of type float.
+        - The result column will always be of type float.
     """
     cols, func_args = [], []
     for arg in args:
-        if isinstance(arg, str):
+        if isinstance(arg, (str, Column, Func)):
             cols.append(arg)
         else:
             func_args.append(list(arg))
@@ -106,31 +115,33 @@ def euclidean_distance(*args: Union[str, Sequence]) -> Func:
     )
 
 
-def length(arg: Union[str, Sequence, Func]) -> Func:
+def length(arg: Union[str, Column, Func, Sequence]) -> Func:
     """
     Returns the length of the array.
 
     Args:
-        arg (str | Sequence | Func): Array to compute the length of.
+        arg (str | Column | Func | Sequence): Array to compute the length of.
             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 Column is provided, it is assumed to be an array column.
             If a Func is provided, it is assumed to be a function returning an array.
+            If a sequence is provided, it is assumed to be an array of values.
 
     Returns:
-        Func: A Func object that represents the array length function.
+        Func: A `Func` object that represents the array length function.
 
     Example:
         ```py
         dc.mutate(
             len1=func.array.length("signal.values"),
-            len2=func.array.length([1, 2, 3, 4, 5]),
+            len2=func.array.length(dc.C("signal.values")),
+            len3=func.array.length([1, 2, 3, 4, 5]),
         )
         ```
 
-    Note:
-        - Result column will always be of type int.
+    Notes:
+        - The result column will always be of type int.
     """
-    if isinstance(arg, (str, Func)):
+    if isinstance(arg, (str, Column, Func)):
         cols = [arg]
         args = None
     else:
@@ -140,35 +151,41 @@ 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:
+def contains(arr: Union[str, Column, Func, Sequence], elem: Any) -> Func:
     """
-    Checks whether the `arr` array has the `elem` element.
+    Checks whether the array contains the specified element.
 
     Args:
-        arr (str | Sequence | Func): Array to check for the element.
+        arr (str | Column | Func | Sequence): 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 Column is provided, it is assumed to be an array column.
             If a Func is provided, it is assumed to be a function returning an array.
+            If a sequence is provided, it is assumed to be an array of values.
         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.
+        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),
+            contains2=func.array.contains(dc.C("signal.values"), 7),
+            contains3=func.array.contains([1, 2, 3, 4, 5], 7),
         )
         ```
+
+    Notes:
+        - The result column will always be of type int.
     """
 
     def inner(arg):
         is_json = type(elem) in [list, dict]
         return array.contains(arg, elem, is_json)
 
-    if isinstance(arr, (str, Func)):
+    if isinstance(arr, (str, Column, Func)):
         cols = [arr]
         args = None
     else:
@@ -179,34 +196,38 @@ def contains(arr: Union[str, Sequence, Func], elem: Any) -> Func:
 
 
 def slice(
-    arr: Union[str, Sequence, Func],
+    arr: Union[str, Column, Func, Sequence],
     offset: int,
     length: Optional[int] = None,
 ) -> Func:
     """
-    Returns a slice of the array.
+    Returns a slice of the array starting from the specified offset.
 
     Args:
-        arr (str | Sequence | Func): Array to check for the element.
+        arr (str | Column | Func | Sequence): Array to slice.
             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 Column is provided, it is assumed to be an array column.
             If a Func is provided, it is assumed to be a function returning an array.
-        offset (int): Offset to start the slice from.
-        length (int, optional): Length of the slice. If not provided, the slice will
-            continue to the end of the array.
+            If a sequence is provided, it is assumed to be an array of values.
+        offset (int): Starting position of the slice (0-based).
+        length (int, optional): Number of elements to include in the slice.
+            If not provided, returns all elements from offset to the end.
 
     Returns:
-        Func: A Func object that represents the slice function. Result of the
-            function will be a slice of the array starting from the offset
-            and with the given length.
+        Func: A `Func` object that represents the slice function.
 
     Example:
         ```py
         dc.mutate(
-            contains1=func.array.slice("signal.values", 3),
-            contains2=func.array.slice([1, 2, 3, 4, 5], 1, 3),
+            slice1=func.array.slice("signal.values", 1, 3),
+            slice2=func.array.slice(dc.C("signal.values"), 2),
+            slice3=func.array.slice([1, 2, 3, 4, 5], 1, 2),
         )
         ```
+
+    Notes:
+        - The result column will be of type array with the same element type
+            as the input.
     """
 
     def inner(arg):
@@ -228,11 +249,11 @@ def slice(
             try:
                 return list[element_type(arr[0])]
             except IndexError:
-                # if the array is empty, return list[str] as default type
-                return list[str]
-        return None
+                pass
+        # if not an array or array is empty, return list[str] as default type
+        return list[str]
 
-    if isinstance(arr, (str, Func)):
+    if isinstance(arr, (str, Column, Func)):
         cols = [arr]
         args = None
     else:
@@ -251,37 +272,40 @@ def slice(
 
 
 def join(
-    arr: Union[str, Sequence, Func],
+    arr: Union[str, Column, Func, Sequence],
     sep: str = "",
 ) -> Func:
     """
-    Returns a string that is the concatenation of the elements of the array,
+    Returns a string that is the concatenation of the elements of the array.
 
     Args:
-        arr (str | Sequence | Func): Array to check for the element.
+        arr (str | Column | Func | Sequence): Array to join.
             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 Column is provided, it is assumed to be an array column.
             If a Func is provided, it is assumed to be a function returning an array.
+            If a sequence is provided, it is assumed to be an array of values.
         sep (str): Separator to use for the concatenation. Default is an empty string.
 
     Returns:
-        Func: A Func object that represents the join function. Result of the
-            function will be a string that is the concatenation of the elements
-            of the array, separated by the given separator.
+        Func: A `Func` object that represents the join function.
 
     Example:
         ```py
         dc.mutate(
-            contains1=func.array.join("signal.values", ":"),
-            contains2=func.array.join(["1", "2", "3", "4", "5"], "/"),
+            join1=func.array.join("signal.values", ":"),
+            join2=func.array.join(dc.C("signal.values"), ","),
+            join3=func.array.join(["1", "2", "3", "4", "5"], "/"),
         )
         ```
+
+    Notes:
+        - The result column will always be of type string.
     """
 
     def inner(arg):
         return array.join(arg, sep)
 
-    if isinstance(arr, (str, Func)):
+    if isinstance(arr, (str, Column, Func)):
         cols = [arr]
         args = None
     else:
@@ -298,31 +322,33 @@ def join(
     )
 
 
-def get_element(arg: Union[str, Sequence, Func], index: int) -> Func:
+def get_element(arg: Union[str, Column, Func, Sequence], 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.
+        arg (str | Column | Func | Sequence): 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 Column is provided, it is assumed to be an array column.
             If a Func is provided, it is assumed to be a function returning an array.
+            If a sequence is provided, it is assumed to be an array of values.
         index (int): Index of the element to get from the array.
 
     Returns:
-        Func: A Func object that represents the array get_element function.
+        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),
+            second_el=func.array.get_element(dc.C("signal.values"), 1),
+            third_el=func.array.get_element([1, 2, 3, 4, 5], 2),
         )
         ```
 
-    Note:
-        - Result column will always be the same type as the elements of the array.
+    Notes:
+        - The result column will always be the same type as the elements of the array.
     """
 
     def type_from_args(arr, _):
@@ -333,10 +359,10 @@ def get_element(arg: Union[str, Sequence, Func], index: int) -> Func:
                 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]
+    cols: Optional[Union[str, Column, Func, Sequence]]
+    args: Union[str, Column, Func, Sequence, int]
 
-    if isinstance(arg, (str, Func)):
+    if isinstance(arg, (str, Column, Func)):
         cols = [arg]
         args = [index]
     else:
@@ -353,31 +379,34 @@ def get_element(arg: Union[str, Sequence, Func], index: int) -> Func:
     )
 
 
-def sip_hash_64(arg: Union[str, Sequence]) -> Func:
+def sip_hash_64(arg: Union[str, Column, Func, Sequence]) -> Func:
     """
-    Computes the SipHash-64 hash of the array.
+    Returns the SipHash-64 hash of the array.
 
     Args:
-        arg (str | Sequence): Array to compute the SipHash-64 hash of.
+        arg (str | Column | Func | Sequence): Array to compute the SipHash-64 hash of.
             If a string is provided, it is assumed to be the name of the array column.
+            If a Column is provided, it is assumed to be an array column.
+            If a Func is provided, it is assumed to be a function returning an array.
             If a sequence is provided, it is assumed to be an array of values.
 
     Returns:
-        Func: A Func object that represents the sip_hash_64 function.
+        Func: A `Func` object that represents the sip_hash_64 function.
 
     Example:
         ```py
         dc.mutate(
             hash1=func.sip_hash_64("signal.values"),
-            hash2=func.sip_hash_64([1, 2, 3, 4, 5]),
+            hash2=func.sip_hash_64(dc.C("signal.values")),
+            hash3=func.sip_hash_64([1, 2, 3, 4, 5]),
         )
         ```
 
     Note:
         - This function is only available for the ClickHouse warehouse.
-        - Result column will always be of type int.
+        - The result column will always be of type int.
     """
-    if isinstance(arg, str):
+    if isinstance(arg, (str, Column, Func)):
         cols = [arg]
         args = None
     else: