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

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,18 +122,19 @@ 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.
 
@@ -128,7 +142,8 @@ def max(col: str) -> Func:
     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),
         )
         ```