More func tests (#1119)

* Add func tests for sql functions
* Update and improve docstrings for consistency

Author: dreadatour

src/datachain/func/aggregate.py

@@ -1,78 +1,89 @@
-from typing import Optional
+from typing import Optional, Union
 
 from sqlalchemy import func as sa_func
 
+from datachain.query.schema import Column
 from datachain.sql.functions import aggregate
 
 from .func import Func
 
 
-def count(col: Optional[str] = None) -> Func:
+def count(col: Optional[Union[str, Column]] = None) -> Func:
     """
-    Returns the COUNT aggregate SQL function for the given column name.
+    Returns a COUNT aggregate SQL function for the specified column.
 
-    The COUNT function returns the number of rows in a table.
+    The COUNT function returns the number of rows, optionally filtered
+    by a specific column.
 
     Args:
-        col (str, optional): The name of the column for which to count rows.
-                             If not provided, it defaults to counting all rows.
+        col (str | Column, optional): The column to count.
+            If omitted, counts all rows.
+            The column can be specified as a string or a `Column` object.
 
     Returns:
-        Func: A Func object that represents the COUNT aggregate function.
+        Func: A `Func` object representing the COUNT aggregate function.
 
     Example:
         ```py
         dc.group_by(
-            count=func.count(),
+            count1=func.count(),
+            count2=func.count("signal.id"),
+            count3=func.count(dc.C("signal.category")),
             partition_by="signal.category",
         )
         ```
 
     Notes:
-        - Result column will always be of type int.
+        - The result column will always have an integer type.
     """
     return Func(
-        "count", inner=sa_func.count, cols=[col] if col else None, result_type=int
+        "count",
+        inner=sa_func.count,
+        cols=[col] if col is not None else None,
+        result_type=int,
     )
 
 
-def sum(col: str) -> Func:
+def sum(col: Union[str, Column]) -> Func:
     """
-    Returns the SUM aggregate SQL function for the given column name.
+    Returns the SUM aggregate SQL function for the specified column.
 
     The SUM function returns the total sum of a numeric column in a table.
     It sums up all the values for the specified column.
 
     Args:
-        col (str): The name of the column for which to calculate the sum.
+        col (str | Column): The name of the column for which to calculate the sum.
+            The column can be specified as a string or a `Column` object.
 
     Returns:
-        Func: A Func object that represents the SUM aggregate function.
+        Func: A `Func` object that represents the SUM aggregate function.
 
     Example:
         ```py
         dc.group_by(
             files_size=func.sum("file.size"),
+            total_size=func.sum(dc.C("size")),
             partition_by="signal.category",
         )
         ```
 
     Notes:
         - The `sum` function should be used on numeric columns.
-        - Result column type will be the same as the input column type.
+        - The result column type will be the same as the input column type.
     """
     return Func("sum", inner=sa_func.sum, cols=[col])
 
 
-def avg(col: str) -> Func:
+def avg(col: Union[str, Column]) -> Func:
     """
-    Returns the AVG aggregate SQL function for the given column name.
+    Returns the AVG aggregate SQL function for the specified column.
 
     The AVG function returns the average of a numeric column in a table.
     It calculates the mean of all values in the specified column.
 
     Args:
-        col (str): The name of the column for which to calculate the average.
+        col (str | Column): The name of the column for which to calculate the average.
+            Column can be specified as a string or a `Column` object.
 
     Returns:
         Func: A Func object that represents the AVG aggregate function.
@@ -81,26 +92,28 @@ def avg(col: str) -> Func:
         ```py
         dc.group_by(
             avg_file_size=func.avg("file.size"),
+            avg_signal_value=func.avg(dc.C("signal.value")),
             partition_by="signal.category",
         )
         ```
 
     Notes:
         - The `avg` function should be used on numeric columns.
-        - Result column will always be of type float.
+        - The result column will always be of type float.
     """
     return Func("avg", inner=aggregate.avg, cols=[col], result_type=float)
 
 
-def min(col: str) -> Func:
+def min(col: Union[str, Column]) -> Func:
     """
-    Returns the MIN aggregate SQL function for the given column name.
+    Returns the MIN aggregate SQL function for the specified column.
 
     The MIN function returns the smallest value in the specified column.
     It can be used on both numeric and non-numeric columns to find the minimum value.
 
     Args:
-        col (str): The name of the column for which to find the minimum value.
+        col (str | Column): The name of the column for which to find the minimum value.
+            Column can be specified as a string or a `Column` object.
 
     Returns:
         Func: A Func object that represents the MIN aggregate function.
@@ -109,26 +122,28 @@ def min(col: str) -> Func:
         ```py
         dc.group_by(
             smallest_file=func.min("file.size"),
+            min_signal=func.min(dc.C("signal")),
             partition_by="signal.category",
         )
         ```
 
     Notes:
         - The `min` function can be used with numeric, date, and string columns.
-        - Result column will have the same type as the input column.
+        - The result column will have the same type as the input column.
     """
     return Func("min", inner=sa_func.min, cols=[col])
 
 
-def max(col: str) -> Func:
+def max(col: Union[str, Column]) -> Func:
     """
     Returns the MAX aggregate SQL function for the given column name.
 
     The MAX function returns the smallest value in the specified column.
     It can be used on both numeric and non-numeric columns to find the maximum value.
 
     Args:
-        col (str): The name of the column for which to find the maximum value.
+        col (str | Column): The name of the column for which to find the maximum value.
+            Column can be specified as a string or a `Column` object.
 
     Returns:
         Func: A Func object that represents the MAX aggregate function.
@@ -137,18 +152,19 @@ def max(col: str) -> Func:
         ```py
         dc.group_by(
             largest_file=func.max("file.size"),
+            max_signal=func.max(dc.C("signal")),
             partition_by="signal.category",
         )
         ```
 
     Notes:
         - The `max` function can be used with numeric, date, and string columns.
-        - Result column will have the same type as the input column.
+        - The result column will have the same type as the input column.
     """
     return Func("max", inner=sa_func.max, cols=[col])
 
 
-def any_value(col: str) -> Func:
+def any_value(col: Union[str, Column]) -> Func:
     """
     Returns the ANY_VALUE aggregate SQL function for the given column name.
 
@@ -157,7 +173,9 @@ def any_value(col: str) -> Func:
     as long as it comes from one of the rows in the group.
 
     Args:
-        col (str): The name of the column from which to return an arbitrary value.
+        col (str | Column): The name of the column from which to return
+            an arbitrary value.
+            Column can be specified as a string or a `Column` object.
 
     Returns:
         Func: A Func object that represents the ANY_VALUE aggregate function.
@@ -166,20 +184,21 @@ def any_value(col: str) -> Func:
         ```py
         dc.group_by(
             file_example=func.any_value("file.path"),
+            signal_example=func.any_value(dc.C("signal.value")),
             partition_by="signal.category",
         )
         ```
 
     Notes:
         - The `any_value` function can be used with any type of column.
-        - Result column will have the same type as the input column.
+        - The result column will have the same type as the input column.
         - The result of `any_value` is non-deterministic,
           meaning it may return different values for different executions.
     """
     return Func("any_value", inner=aggregate.any_value, cols=[col])
 
 
-def collect(col: str) -> Func:
+def collect(col: Union[str, Column]) -> Func:
     """
     Returns the COLLECT aggregate SQL function for the given column name.
 
@@ -188,7 +207,8 @@ def collect(col: str) -> Func:
     into a collection, often for further processing or aggregation.
 
     Args:
-        col (str): The name of the column from which to collect values.
+        col (str | Column): The name of the column from which to collect values.
+            Column can be specified as a string or a `Column` object.
 
     Returns:
         Func: A Func object that represents the COLLECT aggregate function.
@@ -197,18 +217,19 @@ def collect(col: str) -> Func:
         ```py
         dc.group_by(
             signals=func.collect("signal"),
+            file_paths=func.collect(dc.C("file.path")),
             partition_by="signal.category",
         )
         ```
 
     Notes:
         - The `collect` function can be used with numeric and string columns.
-        - Result column will have an array type.
+        - The result column will have an array type.
     """
     return Func("collect", inner=aggregate.collect, cols=[col], is_array=True)
 
 
-def concat(col: str, separator="") -> Func:
+def concat(col: Union[str, Column], separator="") -> Func:
     """
     Returns the CONCAT aggregate SQL function for the given column name.
 
@@ -217,9 +238,10 @@ def concat(col: str, separator="") -> Func:
     into a single combined value.
 
     Args:
-        col (str): The name of the column from which to concatenate values.
+        col (str | Column): The name of the column from which to concatenate values.
+            Column can be specified as a string or a `Column` object.
         separator (str, optional): The separator to use between concatenated values.
-                                   Defaults to an empty string.
+            Defaults to an empty string.
 
     Returns:
         Func: A Func object that represents the CONCAT aggregate function.
@@ -228,13 +250,14 @@ def concat(col: str, separator="") -> Func:
         ```py
         dc.group_by(
             files=func.concat("file.path", separator=", "),
+            signals=func.concat(dc.C("signal.name"), separator=" | "),
             partition_by="signal.category",
         )
         ```
 
     Notes:
         - The `concat` function can be used with string columns.
-        - Result column will have a string type.
+        - The result column will have a string type.
     """
 
     def inner(arg):
@@ -325,7 +348,7 @@ def dense_rank() -> Func:
     return Func("dense_rank", inner=sa_func.dense_rank, result_type=int, is_window=True)
 
 
-def first(col: str) -> Func:
+def first(col: Union[str, Column]) -> Func:
     """
     Returns the FIRST_VALUE window function for SQL queries.
 
@@ -334,7 +357,9 @@ def first(col: str) -> Func:
     and can be useful for retrieving the leading value in a group of rows.
 
     Args:
-        col (str): The name of the column from which to retrieve the first value.
+        col (str | Column): The name of the column from which to retrieve
+            the first value.
+            Column can be specified as a string or a `Column` object.
 
     Returns:
         Func: A Func object that represents the FIRST_VALUE window function.
@@ -344,6 +369,7 @@ def first(col: str) -> Func:
         window = func.window(partition_by="signal.category", order_by="created_at")
         dc.mutate(
             first_file=func.first("file.path").over(window),
+            first_signal=func.first(dc.C("signal.value")).over(window),
         )
         ```