datachain 0.6.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

datachain/lib/func/aggregate.py

@@ -0,0 +1,353 @@
+from typing import Optional
+
+from sqlalchemy import func as sa_func
+
+from datachain.sql import functions as dc_func
+
+from .func import Func
+
+
+def count(col: Optional[str] = None) -> Func:
+    """
+    Returns the COUNT aggregate SQL function for the given column name.
+
+    The COUNT function returns the number of rows in a table.
+
+    Args:
+        col (str, optional): The name of the column for which to count rows.
+                             If not provided, it defaults to counting all rows.
+
+    Returns:
+        Func: A Func object that represents the COUNT aggregate function.
+
+    Example:
+        ```py
+        dc.group_by(
+            count=func.count(),
+            partition_by="signal.category",
+        )
+        ```
+
+    Notes:
+        - Result column will always be of type int.
+    """
+    return Func("count", inner=sa_func.count, col=col, result_type=int)
+
+
+def sum(col: str) -> Func:
+    """
+    Returns the SUM aggregate SQL function for the given column name.
+
+    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.
+
+    Returns:
+        Func: A Func object that represents the SUM aggregate function.
+
+    Example:
+        ```py
+        dc.group_by(
+            files_size=func.sum("file.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.
+    """
+    return Func("sum", inner=sa_func.sum, col=col)
+
+
+def avg(col: str) -> Func:
+    """
+    Returns the AVG aggregate SQL function for the given column name.
+
+    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.
+
+    Returns:
+        Func: A Func object that represents the AVG aggregate function.
+
+    Example:
+        ```py
+        dc.group_by(
+            avg_file_size=func.avg("file.size"),
+            partition_by="signal.category",
+        )
+        ```
+
+    Notes:
+        - The `avg` function should be used on numeric columns.
+        - Result column will always be of type float.
+    """
+    return Func("avg", inner=dc_func.aggregate.avg, col=col, result_type=float)
+
+
+def min(col: str) -> Func:
+    """
+    Returns the MIN aggregate SQL function for the given column name.
+
+    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.
+
+    Returns:
+        Func: A Func object that represents the MIN aggregate function.
+
+    Example:
+        ```py
+        dc.group_by(
+            smallest_file=func.min("file.size"),
+            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.
+    """
+    return Func("min", inner=sa_func.min, col=col)
+
+
+def max(col: str) -> 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.
+
+    Returns:
+        Func: A Func object that represents the MAX aggregate function.
+
+    Example:
+        ```py
+        dc.group_by(
+            largest_file=func.max("file.size"),
+            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.
+    """
+    return Func("max", inner=sa_func.max, col=col)
+
+
+def any_value(col: str) -> Func:
+    """
+    Returns the ANY_VALUE aggregate SQL function for the given column name.
+
+    The ANY_VALUE function returns an arbitrary value from the specified column.
+    It is useful when you do not care which particular value is returned,
+    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.
+
+    Returns:
+        Func: A Func object that represents the ANY_VALUE aggregate function.
+
+    Example:
+        ```py
+        dc.group_by(
+            file_example=func.any_value("file.name"),
+            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 of `any_value` is non-deterministic,
+          meaning it may return different values for different executions.
+    """
+    return Func("any_value", inner=dc_func.aggregate.any_value, col=col)
+
+
+def collect(col: str) -> Func:
+    """
+    Returns the COLLECT aggregate SQL function for the given column name.
+
+    The COLLECT function gathers all values from the specified column
+    into an array or similar structure. It is useful for combining values from a column
+    into a collection, often for further processing or aggregation.
+
+    Args:
+        col (str): The name of the column from which to collect values.
+
+    Returns:
+        Func: A Func object that represents the COLLECT aggregate function.
+
+    Example:
+        ```py
+        dc.group_by(
+            signals=func.collect("signal"),
+            partition_by="signal.category",
+        )
+        ```
+
+    Notes:
+        - The `collect` function can be used with numeric and string columns.
+        - Result column will have an array type.
+    """
+    return Func("collect", inner=dc_func.aggregate.collect, col=col, is_array=True)
+
+
+def concat(col: str, separator="") -> Func:
+    """
+    Returns the CONCAT aggregate SQL function for the given column name.
+
+    The CONCAT function concatenates values from the specified column
+    into a single string. It is useful for merging text values from multiple rows
+    into a single combined value.
+
+    Args:
+        col (str): The name of the column from which to concatenate values.
+        separator (str, optional): The separator to use between concatenated values.
+                                   Defaults to an empty string.
+
+    Returns:
+        Func: A Func object that represents the CONCAT aggregate function.
+
+    Example:
+        ```py
+        dc.group_by(
+            files=func.concat("file.name", separator=", "),
+            partition_by="signal.category",
+        )
+        ```
+
+    Notes:
+        - The `concat` function can be used with string columns.
+        - Result column will have a string type.
+    """
+
+    def inner(arg):
+        return dc_func.aggregate.group_concat(arg, separator)
+
+    return Func("concat", inner=inner, col=col, result_type=str)
+
+
+def row_number() -> Func:
+    """
+    Returns the ROW_NUMBER window function for SQL queries.
+
+    The ROW_NUMBER function assigns a unique sequential integer to rows
+    within a partition of a result set, starting from 1 for the first row
+    in each partition. It is commonly used to generate row numbers within
+    partitions or ordered results.
+
+    Returns:
+        Func: A Func object that represents the ROW_NUMBER window function.
+
+    Example:
+        ```py
+        window = func.window(partition_by="signal.category", order_by="created_at")
+        dc.mutate(
+            row_number=func.row_number().over(window),
+        )
+        ```
+
+    Note:
+        - The result column will always be of type int.
+    """
+    return Func("row_number", inner=sa_func.row_number, result_type=int, is_window=True)
+
+
+def rank() -> Func:
+    """
+    Returns the RANK window function for SQL queries.
+
+    The RANK function assigns a rank to each row within a partition of a result set,
+    with gaps in the ranking for ties. Rows with equal values receive the same rank,
+    and the next rank is skipped (i.e., if two rows are ranked 1,
+    the next row is ranked 3).
+
+    Returns:
+        Func: A Func object that represents the RANK window function.
+
+    Example:
+        ```py
+        window = func.window(partition_by="signal.category", order_by="created_at")
+        dc.mutate(
+            rank=func.rank().over(window),
+        )
+        ```
+
+    Notes:
+        - The result column will always be of type int.
+        - The RANK function differs from ROW_NUMBER in that rows with the same value
+          in the ordering column(s) receive the same rank.
+    """
+    return Func("rank", inner=sa_func.rank, result_type=int, is_window=True)
+
+
+def dense_rank() -> Func:
+    """
+    Returns the DENSE_RANK window function for SQL queries.
+
+    The DENSE_RANK function assigns a rank to each row within a partition
+    of a result set, without gaps in the ranking for ties. Rows with equal values
+    receive the same rank, but the next rank is assigned consecutively
+    (i.e., if two rows are ranked 1, the next row will be ranked 2).
+
+    Returns:
+        Func: A Func object that represents the DENSE_RANK window function.
+
+    Example:
+        ```py
+        window = func.window(partition_by="signal.category", order_by="created_at")
+        dc.mutate(
+            dense_rank=func.dense_rank().over(window),
+        )
+        ```
+
+    Notes:
+        - The result column will always be of type int.
+        - The DENSE_RANK function differs from RANK in that it does not leave gaps
+          in the ranking for tied values.
+    """
+    return Func("dense_rank", inner=sa_func.dense_rank, result_type=int, is_window=True)
+
+
+def first(col: str) -> Func:
+    """
+    Returns the FIRST_VALUE window function for SQL queries.
+
+    The FIRST_VALUE function returns the first value in an ordered set of values
+    within a partition. The first value is determined by the specified order
+    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.
+
+    Returns:
+        Func: A Func object that represents the FIRST_VALUE window function.
+
+    Example:
+        ```py
+        window = func.window(partition_by="signal.category", order_by="created_at")
+        dc.mutate(
+            first_file=func.first("file.name").over(window),
+        )
+        ```
+
+    Note:
+        - The result of `first_value` will always reflect the value of the first row
+          in the specified order.
+        - The result column will have the same type as the input column.
+    """
+    return Func("first", inner=sa_func.first_value, col=col, is_window=True)

datachain/lib/func/func.py

@@ -0,0 +1,152 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, Optional
+
+from sqlalchemy import desc
+
+from datachain.lib.convert.python_to_sql import python_to_sql
+from datachain.lib.utils import DataChainColumnError, DataChainParamsError
+from datachain.query.schema import Column, ColumnMeta
+
+if TYPE_CHECKING:
+    from datachain import DataType
+    from datachain.lib.signal_schema import SignalSchema
+
+
+@dataclass
+class Window:
+    """Represents a window specification for SQL window functions."""
+
+    partition_by: str
+    order_by: str
+    desc: bool = False
+
+
+def window(partition_by: str, order_by: str, desc: bool = False) -> Window:
+    """
+    Defines a window specification for SQL window functions.
+
+    The `window` function specifies how to partition and order the result set
+    for the associated window function. It is used to define the scope of the rows
+    that the window function will operate on.
+
+    Args:
+        partition_by (str): The column name by which to partition the result set.
+                            Rows with the same value in the partition column
+                            will be grouped together for the window function.
+        order_by (str): The column name by which to order the rows
+                        within each partition. This determines the sequence in which
+                        the window function is applied.
+        desc (bool, optional): If True, the rows will be ordered in descending order.
+                               Defaults to False, which orders the rows
+                               in ascending order.
+
+    Returns:
+        Window: A Window object representing the window specification.
+
+    Example:
+        ```py
+        window = func.window(partition_by="signal.category", order_by="created_at")
+        dc.mutate(
+            row_number=func.row_number().over(window),
+        )
+        ```
+    """
+    return Window(
+        ColumnMeta.to_db_name(partition_by),
+        ColumnMeta.to_db_name(order_by),
+        desc,
+    )
+
+
+class Func:
+    """Represents a function to be applied to a column in a SQL query."""
+
+    def __init__(
+        self,
+        name: str,
+        inner: Callable,
+        col: Optional[str] = None,
+        result_type: Optional["DataType"] = None,
+        is_array: bool = False,
+        is_window: bool = False,
+        window: Optional[Window] = None,
+    ) -> None:
+        self.name = name
+        self.inner = inner
+        self.col = col
+        self.result_type = result_type
+        self.is_array = is_array
+        self.is_window = is_window
+        self.window = window
+
+    def __str__(self) -> str:
+        return self.name + "()"
+
+    def over(self, window: Window) -> "Func":
+        if not self.is_window:
+            raise DataChainParamsError(f"{self} doesn't support window (over())")
+
+        return Func(
+            "over",
+            self.inner,
+            self.col,
+            self.result_type,
+            self.is_array,
+            self.is_window,
+            window,
+        )
+
+    @property
+    def db_col(self) -> Optional[str]:
+        return ColumnMeta.to_db_name(self.col) if self.col else None
+
+    def db_col_type(self, signals_schema: "SignalSchema") -> Optional["DataType"]:
+        if not self.db_col:
+            return None
+        col_type: type = signals_schema.get_column_type(self.db_col)
+        return list[col_type] if self.is_array else col_type  # type: ignore[valid-type]
+
+    def get_result_type(self, signals_schema: "SignalSchema") -> "DataType":
+        if self.result_type:
+            return self.result_type
+
+        if col_type := self.db_col_type(signals_schema):
+            return col_type
+
+        raise DataChainColumnError(
+            str(self),
+            "Column name is required to infer result type",
+        )
+
+    def get_column(
+        self, signals_schema: "SignalSchema", label: Optional[str] = None
+    ) -> Column:
+        col_type = self.get_result_type(signals_schema)
+        sql_type = python_to_sql(col_type)
+
+        if self.col:
+            col = Column(self.db_col, sql_type)
+            func_col = self.inner(col)
+        else:
+            func_col = self.inner()
+
+        if self.is_window:
+            if not self.window:
+                raise DataChainParamsError(
+                    f"Window function {self} requires over() clause with a window spec",
+                )
+            func_col = func_col.over(
+                partition_by=self.window.partition_by,
+                order_by=(
+                    desc(self.window.order_by)
+                    if self.window.desc
+                    else self.window.order_by
+                ),
+            )
+
+        func_col.type = sql_type
+
+        if label:
+            func_col = func_col.label(label)
+
+        return func_col

datachain/lib/listing.py

@@ -1,6 +1,5 @@
 import posixpath
 from collections.abc import Iterator
-from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Callable, Optional, TypeVar
 
 from fsspec.asyn import get_loop
@@ -85,12 +84,13 @@ def parse_listing_uri(uri: str, cache, client_config) -> tuple[str, str, str]:
     storage_uri, path = Client.parse_url(uri)
     telemetry.log_param("client", client.PREFIX)
 
-    # clean path without globs
-    lst_uri_path = (
-        posixpath.dirname(path) if uses_glob(path) or client.fs.isfile(uri) else path
-    )
+    if uses_glob(path) or client.fs.isfile(uri):
+        lst_uri_path = posixpath.dirname(path)
+    else:
+        storage_uri, path = Client.parse_url(f'{uri.rstrip("/")}/')
+        lst_uri_path = path
 
-    lst_uri = f"{storage_uri}/{lst_uri_path.lstrip('/')}"
+    lst_uri = f'{storage_uri}/{lst_uri_path.lstrip("/")}'
     ds_name = (
         f"{LISTING_PREFIX}{storage_uri}/{posixpath.join(lst_uri_path, '').lstrip('/')}"
     )
@@ -108,18 +108,3 @@ def listing_uri_from_name(dataset_name: str) -> str:
     if not is_listing_dataset(dataset_name):
         raise ValueError(f"Dataset {dataset_name} is not a listing")
     return dataset_name.removeprefix(LISTING_PREFIX)
-
-
-def is_listing_expired(created_at: datetime) -> bool:
-    """Checks if listing has expired based on it's creation date"""
-    return datetime.now(timezone.utc) > created_at + timedelta(seconds=LISTING_TTL)
-
-
-def is_listing_subset(ds1_name: str, ds2_name: str) -> bool:
-    """
-    Checks if one listing contains another one by comparing corresponding dataset names
-    """
-    assert ds1_name.endswith("/")
-    assert ds2_name.endswith("/")
-
-    return ds2_name.startswith(ds1_name)

datachain/lib/listing_info.py

@@ -30,3 +30,7 @@ class ListingInfo(DatasetInfo):
     def last_inserted_at(self):
         # TODO we need to add updated_at to dataset version or explicit last_inserted_at
         raise NotImplementedError
+
+    def contains(self, other_name: str) -> bool:
+        """Checks if this listing contains another one"""
+        return other_name.startswith(self.name)

datachain/lib/signal_schema.py

@@ -20,6 +20,7 @@ from typing import (  # noqa: UP035
 )
 
 from pydantic import BaseModel, create_model
+from sqlalchemy import ColumnElement
 from typing_extensions import Literal as LiteralEx
 
 from datachain.lib.convert.python_to_sql import python_to_sql
@@ -27,6 +28,7 @@ from datachain.lib.convert.sql_to_python import sql_to_python
 from datachain.lib.convert.unflatten import unflatten_to_json_pos
 from datachain.lib.data_model import DataModel, DataType, DataValue
 from datachain.lib.file import File
+from datachain.lib.func import Func
 from datachain.lib.model_store import ModelStore
 from datachain.lib.utils import DataChainParamsError
 from datachain.query.schema import DEFAULT_DELIMITER, Column
@@ -400,6 +402,12 @@ class SignalSchema:
             if ModelStore.is_pydantic(finfo.annotation):
                 SignalSchema._set_file_stream(getattr(obj, field), catalog, cache)
 
+    def get_column_type(self, col_name: str) -> DataType:
+        for path, _type, has_subtree, _ in self.get_flat_tree():
+            if not has_subtree and DEFAULT_DELIMITER.join(path) == col_name:
+                return _type
+        raise SignalResolvingError([col_name], "is not found")
+
     def db_signals(
         self, name: Optional[str] = None, as_columns=False
     ) -> Union[list[str], list[Column]]:
@@ -484,13 +492,14 @@ class SignalSchema:
                 # renaming existing signal
                 del new_values[value.name]
                 new_values[name] = self.values[value.name]
-            elif name in self.values:
-                # changing the type of existing signal, e.g File -> ImageFile
-                del new_values[name]
-                new_values[name] = args_map[name]
-            else:
+            elif isinstance(value, Func):
+                # adding new signal with function
+                new_values[name] = value.get_result_type(self)
+            elif isinstance(value, ColumnElement):
                 # adding new signal
-                new_values.update(sql_to_python({name: value}))
+                new_values[name] = sql_to_python(value)
+            else:
+                new_values[name] = value
 
         return SignalSchema(new_values)
 
@@ -534,12 +543,12 @@ class SignalSchema:
             for name, val in values.items()
         }
 
-    def get_flat_tree(self) -> Iterator[tuple[list[str], type, bool, int]]:
+    def get_flat_tree(self) -> Iterator[tuple[list[str], DataType, bool, int]]:
         yield from self._get_flat_tree(self.tree, [], 0)
 
     def _get_flat_tree(
         self, tree: dict, prefix: list[str], depth: int
-    ) -> Iterator[tuple[list[str], type, bool, int]]:
+    ) -> Iterator[tuple[list[str], DataType, bool, int]]:
         for name, (type_, substree) in tree.items():
             suffix = name.split(".")
             new_prefix = prefix + suffix

datachain/lib/udf.py

@@ -11,7 +11,6 @@ from datachain.dataset import RowDict
 from datachain.lib.convert.flatten import flatten
 from datachain.lib.data_model import DataValue
 from datachain.lib.file import File
-from datachain.lib.signal_schema import SignalSchema
 from datachain.lib.utils import AbstractUDF, DataChainError, DataChainParamsError
 from datachain.query.batch import (
     Batch,
@@ -25,6 +24,7 @@ if TYPE_CHECKING:
     from typing_extensions import Self
 
     from datachain.catalog import Catalog
+    from datachain.lib.signal_schema import SignalSchema
     from datachain.lib.udf_signature import UdfSignature
     from datachain.query.batch import RowsOutput
 
@@ -172,7 +172,7 @@ class UDFBase(AbstractUDF):
     def _init(
         self,
         sign: "UdfSignature",
-        params: SignalSchema,
+        params: "SignalSchema",
         func: Optional[Callable],
     ):
         self.params = params
@@ -183,7 +183,7 @@ class UDFBase(AbstractUDF):
     def _create(
         cls,
         sign: "UdfSignature",
-        params: SignalSchema,
+        params: "SignalSchema",
     ) -> "Self":
         if isinstance(sign.func, AbstractUDF):
             if not isinstance(sign.func, cls):  # type: ignore[unreachable]

datachain/lib/utils.py

@@ -23,3 +23,8 @@ class DataChainError(Exception):
 class DataChainParamsError(DataChainError):
     def __init__(self, message):
         super().__init__(message)
+
+
+class DataChainColumnError(DataChainParamsError):
+    def __init__(self, col_name, msg):
+        super().__init__(f"Error for column {col_name}: {msg}")