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

datachain/__init__.py

@@ -1,3 +1,4 @@
+from datachain.lib import func
 from datachain.lib.data_model import DataModel, DataType, is_chain_type
 from datachain.lib.dc import C, Column, DataChain, Sys
 from datachain.lib.file import (
@@ -34,6 +35,7 @@ __all__ = [
     "Sys",
     "TarVFile",
     "TextFile",
+    "func",
     "is_chain_type",
     "metrics",
     "param",

datachain/catalog/catalog.py

@@ -989,13 +989,6 @@ class Catalog:
             c.name: c.type.to_dict() for c in columns if isinstance(c.type, SQLType)
         }
 
-        job_id = job_id or os.getenv("DATACHAIN_JOB_ID")
-        if not job_id:
-            from datachain.query.session import Session
-
-            session = Session.get(catalog=self)
-            job_id = session.job_id
-
         dataset = self.metastore.create_dataset_version(
             dataset,
             version,
@@ -1218,6 +1211,7 @@ class Catalog:
             preview=dataset_version.preview,
             job_id=dataset_version.job_id,
         )
+
         # to avoid re-creating rows table, we are just renaming it for a new version
         # of target dataset
         self.warehouse.rename_dataset_table(
@@ -1325,8 +1319,6 @@ class Catalog:
         if offset:
             q = q.offset(offset)
 
-        q = q.order_by("sys__id")
-
         return q.to_db_records()
 
     def signed_url(self, source: str, path: str, client_config=None) -> str:

datachain/data_storage/sqlite.py

@@ -763,6 +763,14 @@ class SQLiteWarehouse(AbstractWarehouse):
         query: Select,
         progress_cb: Optional[Callable[[int], None]] = None,
     ) -> None:
+        if len(query._group_by_clause) > 0:
+            select_q = query.with_only_columns(
+                *[c for c in query.selected_columns if c.name != "sys__id"]
+            )
+            q = table.insert().from_select(list(select_q.selected_columns), select_q)
+            self.db.execute(q)
+            return
+
         if "sys__id" in query.selected_columns:
             col_id = query.selected_columns.sys__id
         else:

datachain/data_storage/warehouse.py

@@ -215,10 +215,6 @@ class AbstractWarehouse(ABC, Serializable):
         limit = query._limit
         paginated_query = query.limit(page_size)
 
-        if not paginated_query._order_by_clauses:
-            # default order by is order by `sys__id`
-            paginated_query = paginated_query.order_by(query.selected_columns.sys__id)
-
         results = None
         offset = 0
         num_yielded = 0

datachain/lib/convert/sql_to_python.py

@@ -4,15 +4,11 @@ from typing import Any
 from sqlalchemy import ColumnElement
 
 
-def sql_to_python(args_map: dict[str, ColumnElement]) -> dict[str, Any]:
-    res = {}
-    for name, sql_exp in args_map.items():
-        try:
-            type_ = sql_exp.type.python_type
-            if type_ == Decimal:
-                type_ = float
-        except NotImplementedError:
-            type_ = str
-        res[name] = type_
-
-    return res
+def sql_to_python(sql_exp: ColumnElement) -> Any:
+    try:
+        type_ = sql_exp.type.python_type
+        if type_ == Decimal:
+            type_ = float
+    except NotImplementedError:
+        type_ = str
+    return type_

datachain/lib/dc.py

@@ -29,6 +29,7 @@ from datachain.lib.data_model import DataModel, DataType, dict_to_data_model
 from datachain.lib.dataset_info import DatasetInfo
 from datachain.lib.file import ArrowRow, File, get_file_type
 from datachain.lib.file import ExportPlacement as FileExportPlacement
+from datachain.lib.func import Func
 from datachain.lib.listing import (
     is_listing_dataset,
     is_listing_expired,
@@ -42,21 +43,12 @@ from datachain.lib.meta_formats import read_meta, read_schema
 from datachain.lib.model_store import ModelStore
 from datachain.lib.settings import Settings
 from datachain.lib.signal_schema import SignalSchema
-from datachain.lib.udf import (
-    Aggregator,
-    BatchMapper,
-    Generator,
-    Mapper,
-    UDFBase,
-)
+from datachain.lib.udf import Aggregator, BatchMapper, Generator, Mapper, UDFBase
 from datachain.lib.udf_signature import UdfSignature
-from datachain.lib.utils import DataChainParamsError
+from datachain.lib.utils import DataChainColumnError, DataChainParamsError
 from datachain.query import Session
-from datachain.query.dataset import (
-    DatasetQuery,
-    PartitionByType,
-)
-from datachain.query.schema import DEFAULT_DELIMITER, Column, DatasetRow
+from datachain.query.dataset import DatasetQuery, PartitionByType
+from datachain.query.schema import DEFAULT_DELIMITER, Column, ColumnMeta
 from datachain.sql.functions import path as pathfunc
 from datachain.telemetry import telemetry
 from datachain.utils import batched_it, inside_notebook
@@ -149,11 +141,6 @@ class DatasetMergeError(DataChainParamsError):  # noqa: D101
         super().__init__(f"Merge error on='{on_str}'{right_on_str}: {msg}")
 
 
-class DataChainColumnError(DataChainParamsError):  # noqa: D101
-    def __init__(self, col_name, msg):  # noqa: D107
-        super().__init__(f"Error for column {col_name}: {msg}")
-
-
 OutputType = Union[None, DataType, Sequence[str], dict[str, DataType]]
 
 
@@ -982,10 +969,9 @@ class DataChain:
         row is left in the result set.
 
         Example:
-        ```py
-         dc.distinct("file.parent", "file.name")
-        )
-        ```
+            ```py
+            dc.distinct("file.parent", "file.name")
+            ```
         """
         return self._evolve(
             query=self._query.distinct(
@@ -1011,6 +997,60 @@ class DataChain:
             query=self._query.select(*columns), signal_schema=new_schema
         )
 
+    def group_by(
+        self,
+        *,
+        partition_by: Union[str, Sequence[str]],
+        **kwargs: Func,
+    ) -> "Self":
+        """Group rows by specified set of signals and return new signals
+        with aggregated values.
+
+        Example:
+            ```py
+            chain = chain.group_by(
+                cnt=func.count(),
+                partition_by=("file_source", "file_ext"),
+            )
+            ```
+        """
+        if isinstance(partition_by, str):
+            partition_by = [partition_by]
+        if not partition_by:
+            raise ValueError("At least one column should be provided for partition_by")
+
+        if not kwargs:
+            raise ValueError("At least one column should be provided for group_by")
+        for col_name, func in kwargs.items():
+            if not isinstance(func, Func):
+                raise DataChainColumnError(
+                    col_name,
+                    f"Column {col_name} has type {type(func)} but expected Func object",
+                )
+
+        partition_by_columns: list[Column] = []
+        signal_columns: list[Column] = []
+        schema_fields: dict[str, DataType] = {}
+
+        # validate partition_by columns and add them to the schema
+        for col_name in partition_by:
+            col_db_name = ColumnMeta.to_db_name(col_name)
+            col_type = self.signals_schema.get_column_type(col_db_name)
+            col = Column(col_db_name, python_to_sql(col_type))
+            partition_by_columns.append(col)
+            schema_fields[col_db_name] = col_type
+
+        # validate signal columns and add them to the schema
+        for col_name, func in kwargs.items():
+            col = func.get_column(self.signals_schema, label=col_name)
+            signal_columns.append(col)
+            schema_fields[col_name] = func.get_result_type(self.signals_schema)
+
+        return self._evolve(
+            query=self._query.group_by(signal_columns, partition_by_columns),
+            signal_schema=SignalSchema(schema_fields),
+        )
+
     def mutate(self, **kwargs) -> "Self":
         """Create new signals based on existing signals.
 
@@ -1477,12 +1517,6 @@ class DataChain:
         fr_map = {col.lower(): df[col].tolist() for col in df.columns}
 
         for column in fr_map:
-            if column in DatasetRow.schema:
-                raise DatasetPrepareError(
-                    name,
-                    f"import from pandas error - column '{column}' conflicts with"
-                    " default schema",
-                )
             if not column.isidentifier():
                 raise DatasetPrepareError(
                     name,
@@ -1994,6 +2028,8 @@ class DataChain:
             ),
         )
 
+        session.add_dataset_version(dsr, dsr.latest_version)
+
         if isinstance(to_insert, dict):
             to_insert = [to_insert]
         elif not to_insert:

datachain/lib/func/__init__.py

@@ -0,0 +1,14 @@
+from .aggregate import any_value, avg, collect, concat, count, max, min, sum
+from .func import Func
+
+__all__ = [
+    "Func",
+    "any_value",
+    "avg",
+    "collect",
+    "concat",
+    "count",
+    "max",
+    "min",
+    "sum",
+]

datachain/lib/func/aggregate.py

@@ -0,0 +1,42 @@
+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:
+    return Func(inner=sa_func.count, col=col, result_type=int)
+
+
+def sum(col: str) -> Func:
+    return Func(inner=sa_func.sum, col=col)
+
+
+def avg(col: str) -> Func:
+    return Func(inner=dc_func.aggregate.avg, col=col)
+
+
+def min(col: str) -> Func:
+    return Func(inner=sa_func.min, col=col)
+
+
+def max(col: str) -> Func:
+    return Func(inner=sa_func.max, col=col)
+
+
+def any_value(col: str) -> Func:
+    return Func(inner=dc_func.aggregate.any_value, col=col)
+
+
+def collect(col: str) -> Func:
+    return Func(inner=dc_func.aggregate.collect, col=col, is_array=True)
+
+
+def concat(col: str, separator="") -> Func:
+    def inner(arg):
+        return dc_func.aggregate.group_concat(arg, separator)
+
+    return Func(inner=inner, col=col, result_type=str)

datachain/lib/func/func.py

@@ -0,0 +1,64 @@
+from typing import TYPE_CHECKING, Callable, Optional
+
+from datachain.lib.convert.python_to_sql import python_to_sql
+from datachain.lib.utils import DataChainColumnError
+from datachain.query.schema import Column, ColumnMeta
+
+if TYPE_CHECKING:
+    from datachain import DataType
+    from datachain.lib.signal_schema import SignalSchema
+
+
+class Func:
+    def __init__(
+        self,
+        inner: Callable,
+        col: Optional[str] = None,
+        result_type: Optional["DataType"] = None,
+        is_array: bool = False,
+    ) -> None:
+        self.inner = inner
+        self.col = col
+        self.result_type = result_type
+        self.is_array = is_array
+
+    @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":
+        col_type = self.db_col_type(signals_schema)
+
+        if self.result_type:
+            return self.result_type
+
+        if col_type:
+            return col_type
+
+        raise DataChainColumnError(
+            str(self.inner),
+            "Column name is required to infer result type",
+        )
+
+    def get_column(
+        self, signals_schema: "SignalSchema", label: Optional[str] = None
+    ) -> Column:
+        if self.col:
+            if label == "collect":
+                print(label)
+            col_type = self.get_result_type(signals_schema)
+            col = Column(self.db_col, python_to_sql(col_type))
+            func_col = self.inner(col)
+        else:
+            func_col = self.inner()
+
+        if label:
+            func_col = func_col.label(label)
+
+        return func_col

datachain/lib/signal_schema.py

@@ -400,6 +400,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]]:
@@ -490,7 +496,7 @@ class SignalSchema:
                 new_values[name] = args_map[name]
             else:
                 # adding new signal
-                new_values.update(sql_to_python({name: value}))
+                new_values[name] = sql_to_python(value)
 
         return SignalSchema(new_values)
 
@@ -534,12 +540,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/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}")

datachain/query/__init__.py

@@ -1,12 +1,11 @@
 from .dataset import DatasetQuery
 from .params import param
-from .schema import C, DatasetRow, LocalFilename, Object, Stream
+from .schema import C, LocalFilename, Object, Stream
 from .session import Session
 
 __all__ = [
     "C",
     "DatasetQuery",
-    "DatasetRow",
     "LocalFilename",
     "Object",
     "Session",

datachain/query/batch.py

@@ -97,7 +97,6 @@ class Partition(BatchingStrategy):
 
         ordered_query = query.order_by(None).order_by(
             PARTITION_COLUMN_ID,
-            "sys__id",
             *query._order_by_clauses,
         )
 

datachain/query/dataset.py

@@ -591,10 +591,6 @@ class UDFSignal(UDFStep):
             return query, []
         table = self.catalog.warehouse.create_pre_udf_table(query)
         q: Select = sqlalchemy.select(*table.c)
-        if query._order_by_clauses:
-            # we are adding ordering only if it's explicitly added by user in
-            # query part before adding signals
-            q = q.order_by(table.c.sys__id)
         return q, [table]
 
     def create_result_query(
@@ -630,11 +626,6 @@ class UDFSignal(UDFStep):
             else:
                 res = sqlalchemy.select(*cols1).select_from(subq)
 
-            if query._order_by_clauses:
-                # if ordering is used in query part before adding signals, we
-                # will have it as order by id from select from pre-created udf table
-                res = res.order_by(subq.c.sys__id)
-
             if self.partition_by is not None:
                 subquery = res.subquery()
                 res = sqlalchemy.select(*subquery.c).select_from(subquery)
@@ -666,13 +657,6 @@ class RowGenerator(UDFStep):
     def create_result_query(
         self, udf_table, query: Select
     ) -> tuple[QueryGeneratorFunc, list["sqlalchemy.Column"]]:
-        if not query._order_by_clauses:
-            # if we are not selecting all rows in UDF, we need to ensure that
-            # we get the same rows as we got as inputs of UDF since selecting
-            # without ordering can be non deterministic in some databases
-            c = query.selected_columns
-            query = query.order_by(c.sys__id)
-
         udf_table_query = udf_table.select().subquery()
         udf_table_cols: list[sqlalchemy.Label[Any]] = [
             label(c.name, c) for c in udf_table_query.columns
@@ -957,24 +941,24 @@ class SQLJoin(Step):
 
 
 @frozen
-class GroupBy(Step):
-    """Group rows by a specific column."""
-
-    cols: PartitionByType
+class SQLGroupBy(SQLClause):
+    cols: Sequence[Union[str, ColumnElement]]
+    group_by: Sequence[Union[str, ColumnElement]]
 
-    def clone(self) -> "Self":
-        return self.__class__(self.cols)
+    def apply_sql_clause(self, query) -> Select:
+        if not self.cols:
+            raise ValueError("No columns to select")
+        if not self.group_by:
+            raise ValueError("No columns to group by")
 
-    def apply(
-        self, query_generator: QueryGenerator, temp_tables: list[str]
-    ) -> StepResult:
-        query = query_generator.select()
-        grouped_query = query.group_by(*self.cols)
+        subquery = query.subquery()
 
-        def q(*columns):
-            return grouped_query.with_only_columns(*columns)
+        cols = [
+            subquery.c[str(c)] if isinstance(c, (str, C)) else c
+            for c in [*self.group_by, *self.cols]
+        ]
 
-        return step_result(q, grouped_query.selected_columns)
+        return sqlalchemy.select(*cols).select_from(subquery).group_by(*self.group_by)
 
 
 def _validate_columns(
@@ -1130,25 +1114,14 @@ class DatasetQuery:
             query.steps = query.steps[-1:] + query.steps[:-1]
 
         result = query.starting_step.apply()
-        group_by = None
         self.dependencies.update(result.dependencies)
 
         for step in query.steps:
-            if isinstance(step, GroupBy):
-                if group_by is not None:
-                    raise TypeError("only one group_by allowed")
-                group_by = step
-                continue
-
             result = step.apply(
                 result.query_generator, self.temp_table_names
             )  # a chain of steps linked by results
             self.dependencies.update(result.dependencies)
 
-        if group_by:
-            result = group_by.apply(result.query_generator, self.temp_table_names)
-            self.dependencies.update(result.dependencies)
-
         return result.query_generator
 
     @staticmethod
@@ -1410,9 +1383,13 @@ class DatasetQuery:
         return query.as_scalar()
 
     @detach
-    def group_by(self, *cols: ColumnElement) -> "Self":
+    def group_by(
+        self,
+        cols: Sequence[ColumnElement],
+        group_by: Sequence[ColumnElement],
+    ) -> "Self":
         query = self.clone()
-        query.steps.append(GroupBy(cols))
+        query.steps.append(SQLGroupBy(cols, group_by))
         return query
 
     @detach
@@ -1591,6 +1568,8 @@ class DatasetQuery:
             )
             version = version or dataset.latest_version
 
+            self.session.add_dataset_version(dataset=dataset, version=version)
+
             dr = self.catalog.warehouse.dataset_rows(dataset)
 
             self.catalog.warehouse.copy_table(dr.get_table(), query.select())

datachain/query/schema.py

@@ -1,16 +1,13 @@
 import functools
-import json
 from abc import ABC, abstractmethod
-from datetime import datetime, timezone
 from fnmatch import fnmatch
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union
 
 import attrs
 import sqlalchemy as sa
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 
 from datachain.lib.file import File
-from datachain.sql.types import JSON, Boolean, DateTime, Int64, SQLType, String
 
 if TYPE_CHECKING:
     from datachain.catalog import Catalog
@@ -228,61 +225,4 @@ def normalize_param(param: UDFParamSpec) -> UDFParameter:
     raise TypeError(f"Invalid UDF parameter: {param}")
 
 
-class DatasetRow:
-    schema: ClassVar[dict[str, type[SQLType]]] = {
-        "source": String,
-        "path": String,
-        "size": Int64,
-        "location": JSON,
-        "is_latest": Boolean,
-        "last_modified": DateTime,
-        "version": String,
-        "etag": String,
-    }
-
-    @staticmethod
-    def create(
-        path: str,
-        source: str = "",
-        size: int = 0,
-        location: Optional[dict[str, Any]] = None,
-        is_latest: bool = True,
-        last_modified: Optional[datetime] = None,
-        version: str = "",
-        etag: str = "",
-    ) -> tuple[
-        str,
-        str,
-        int,
-        Optional[str],
-        int,
-        bool,
-        datetime,
-        str,
-        str,
-        int,
-    ]:
-        if location:
-            location = json.dumps([location])  # type: ignore [assignment]
-
-        last_modified = last_modified or datetime.now(timezone.utc)
-
-        return (  # type: ignore [return-value]
-            source,
-            path,
-            size,
-            location,
-            is_latest,
-            last_modified,
-            version,
-            etag,
-        )
-
-    @staticmethod
-    def extend(**columns):
-        cols = {**DatasetRow.schema}
-        cols.update(columns)
-        return cols
-
-
 C = Column

datachain/query/session.py

@@ -1,9 +1,9 @@
 import atexit
+import gc
 import logging
-import os
 import re
 import sys
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, ClassVar, Optional
 from uuid import uuid4
 
 from datachain.catalog import get_catalog
@@ -11,6 +11,7 @@ from datachain.error import TableMissingError
 
 if TYPE_CHECKING:
     from datachain.catalog import Catalog
+    from datachain.dataset import DatasetRecord
 
 logger = logging.getLogger("datachain")
 
@@ -39,7 +40,7 @@ class Session:
     """
 
     GLOBAL_SESSION_CTX: Optional["Session"] = None
-    GLOBAL_SESSION: Optional["Session"] = None
+    SESSION_CONTEXTS: ClassVar[list["Session"]] = []
     ORIGINAL_EXCEPT_HOOK = None
 
     DATASET_PREFIX = "session_"
@@ -64,18 +65,21 @@ class Session:
 
         session_uuid = uuid4().hex[: self.SESSION_UUID_LEN]
         self.name = f"{name}_{session_uuid}"
-        self.job_id = os.getenv("DATACHAIN_JOB_ID") or str(uuid4())
         self.is_new_catalog = not catalog
         self.catalog = catalog or get_catalog(
             client_config=client_config, in_memory=in_memory
         )
+        self.dataset_versions: list[tuple[DatasetRecord, int]] = []
 
     def __enter__(self):
+        # Push the current context onto the stack
+        Session.SESSION_CONTEXTS.append(self)
+
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if exc_type:
-            self._cleanup_created_versions(self.name)
+            self._cleanup_created_versions()
 
         self._cleanup_temp_datasets()
         if self.is_new_catalog:
@@ -83,6 +87,12 @@ class Session:
             self.catalog.warehouse.close_on_exit()
             self.catalog.id_generator.close_on_exit()
 
+        if Session.SESSION_CONTEXTS:
+            Session.SESSION_CONTEXTS.pop()
+
+    def add_dataset_version(self, dataset: "DatasetRecord", version: int) -> None:
+        self.dataset_versions.append((dataset, version))
+
     def generate_temp_dataset_name(self) -> str:
         return self.get_temp_prefix() + uuid4().hex[: self.TEMP_TABLE_UUID_LEN]
 
@@ -98,21 +108,15 @@ class Session:
         except TableMissingError:
             pass
 
-    def _cleanup_created_versions(self, job_id: str) -> None:
-        versions = self.catalog.metastore.get_job_dataset_versions(job_id)
-        if not versions:
+    def _cleanup_created_versions(self) -> None:
+        if not self.dataset_versions:
             return
 
-        datasets = {}
-        for dataset_name, version in versions:
-            if dataset_name not in datasets:
-                datasets[dataset_name] = self.catalog.get_dataset(dataset_name)
-            dataset = datasets[dataset_name]
-            logger.info(
-                "Removing dataset version %s@%s due to exception", dataset_name, version
-            )
+        for dataset, version in self.dataset_versions:
             self.catalog.remove_dataset_version(dataset, version)
 
+        self.dataset_versions.clear()
+
     @classmethod
     def get(
         cls,
@@ -125,33 +129,34 @@ class Session:
 
         Parameters:
             session (Session): Optional Session(). If not provided a new session will
-                    be created. It's needed mostly for simplie API purposes.
-            catalog (Catalog): Optional catalog. By default a new catalog is created.
+                    be created. It's needed mostly for simple API purposes.
+            catalog (Catalog): Optional catalog. By default, a new catalog is created.
         """
         if session:
             return session
 
-        if cls.GLOBAL_SESSION is None:
+        # Access the active (most recent) context from the stack
+        if cls.SESSION_CONTEXTS:
+            return cls.SESSION_CONTEXTS[-1]
+
+        if cls.GLOBAL_SESSION_CTX is None:
             cls.GLOBAL_SESSION_CTX = Session(
                 cls.GLOBAL_SESSION_NAME,
                 catalog,
                 client_config=client_config,
                 in_memory=in_memory,
             )
-            cls.GLOBAL_SESSION = cls.GLOBAL_SESSION_CTX.__enter__()
 
             atexit.register(cls._global_cleanup)
             cls.ORIGINAL_EXCEPT_HOOK = sys.excepthook
             sys.excepthook = cls.except_hook
 
-        return cls.GLOBAL_SESSION
+        return cls.GLOBAL_SESSION_CTX
 
     @staticmethod
     def except_hook(exc_type, exc_value, exc_traceback):
+        Session.GLOBAL_SESSION_CTX.__exit__(exc_type, exc_value, exc_traceback)
         Session._global_cleanup()
-        if Session.GLOBAL_SESSION_CTX is not None:
-            job_id = Session.GLOBAL_SESSION_CTX.job_id
-            Session.GLOBAL_SESSION_CTX._cleanup_created_versions(job_id)
 
         if Session.ORIGINAL_EXCEPT_HOOK:
             Session.ORIGINAL_EXCEPT_HOOK(exc_type, exc_value, exc_traceback)
@@ -160,7 +165,6 @@ class Session:
     def cleanup_for_tests(cls):
         if cls.GLOBAL_SESSION_CTX is not None:
             cls.GLOBAL_SESSION_CTX.__exit__(None, None, None)
-            cls.GLOBAL_SESSION = None
             cls.GLOBAL_SESSION_CTX = None
             atexit.unregister(cls._global_cleanup)
 
@@ -171,3 +175,7 @@ class Session:
     def _global_cleanup():
         if Session.GLOBAL_SESSION_CTX is not None:
             Session.GLOBAL_SESSION_CTX.__exit__(None, None, None)
+
+        for obj in gc.get_objects():  # Get all tracked objects
+            if isinstance(obj, Session):  # Cleanup temp dataset for session variables.
+                obj.__exit__(None, None, None)

datachain/sql/functions/__init__.py

@@ -1,7 +1,7 @@
 from sqlalchemy.sql.expression import func
 
 from . import array, path, string
-from .array import avg
+from .aggregate import avg
 from .conditional import greatest, least
 from .random import rand
 

datachain/sql/functions/aggregate.py

@@ -0,0 +1,47 @@
+from sqlalchemy.sql.functions import GenericFunction, ReturnTypeFromArgs
+
+from datachain.sql.types import Float, String
+from datachain.sql.utils import compiler_not_implemented
+
+
+class avg(GenericFunction):  # noqa: N801
+    """
+    Returns the average of the column.
+    """
+
+    type = Float()
+    package = "array"
+    name = "avg"
+    inherit_cache = True
+
+
+class group_concat(GenericFunction):  # noqa: N801
+    """
+    Returns the concatenated string of the column.
+    """
+
+    type = String()
+    package = "array"
+    name = "group_concat"
+    inherit_cache = True
+
+
+class any_value(ReturnTypeFromArgs):  # noqa: N801
+    """
+    Returns first value of the column.
+    """
+
+    inherit_cache = True
+
+
+class collect(ReturnTypeFromArgs):  # noqa: N801
+    """
+    Returns an array of the column.
+    """
+
+    inherit_cache = True
+
+
+compiler_not_implemented(avg)
+compiler_not_implemented(group_concat)
+compiler_not_implemented(any_value)

datachain/sql/functions/array.py

@@ -44,15 +44,7 @@ class sip_hash_64(GenericFunction):  # noqa: N801
     inherit_cache = True
 
 
-class avg(GenericFunction):  # noqa: N801
-    type = Float()
-    package = "array"
-    name = "avg"
-    inherit_cache = True
-
-
 compiler_not_implemented(cosine_distance)
 compiler_not_implemented(euclidean_distance)
 compiler_not_implemented(length)
 compiler_not_implemented(sip_hash_64)
-compiler_not_implemented(avg)

datachain/sql/sqlite/base.py

@@ -14,7 +14,7 @@ from sqlalchemy.sql.elements import literal
 from sqlalchemy.sql.expression import case
 from sqlalchemy.sql.functions import func
 
-from datachain.sql.functions import array, conditional, random, string
+from datachain.sql.functions import aggregate, array, conditional, random, string
 from datachain.sql.functions import path as sql_path
 from datachain.sql.selectable import Values, base_values_compiler
 from datachain.sql.sqlite.types import (
@@ -84,7 +84,10 @@ def setup():
     compiles(conditional.least, "sqlite")(compile_least)
     compiles(Values, "sqlite")(compile_values)
     compiles(random.rand, "sqlite")(compile_rand)
-    compiles(array.avg, "sqlite")(compile_avg)
+    compiles(aggregate.avg, "sqlite")(compile_avg)
+    compiles(aggregate.group_concat, "sqlite")(compile_group_concat)
+    compiles(aggregate.any_value, "sqlite")(compile_any_value)
+    compiles(aggregate.collect, "sqlite")(compile_collect)
 
     if load_usearch_extension(sqlite3.connect(":memory:")):
         compiles(array.cosine_distance, "sqlite")(compile_cosine_distance_ext)
@@ -400,6 +403,21 @@ def compile_avg(element, compiler, **kwargs):
     return compiler.process(func.avg(*element.clauses.clauses), **kwargs)
 
 
+def compile_group_concat(element, compiler, **kwargs):
+    return compiler.process(func.aggregate_strings(*element.clauses.clauses), **kwargs)
+
+
+def compile_any_value(element, compiler, **kwargs):
+    # use bare column to return any value from the group,
+    # this is documented behavior for sqlite,
+    # see https://www.sqlite.org/lang_select.html#bare_columns_in_an_aggregate_query
+    return compiler.process(*element.clauses.clauses, **kwargs)
+
+
+def compile_collect(element, compiler, **kwargs):
+    return compiler.process(func.json_group_array(*element.clauses.clauses), **kwargs)
+
+
 def load_usearch_extension(conn) -> bool:
     try:
         # usearch is part of the vector optional dependencies

{datachain-0.6.0.dist-info → datachain-0.6.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.6.0
+Version: 0.6.1
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -46,7 +46,7 @@ Requires-Dist: iterative-telemetry >=0.0.9
 Requires-Dist: numpy <2,>=1 ; sys_platform == "win32"
 Provides-Extra: dev
 Requires-Dist: datachain[docs,tests] ; extra == 'dev'
-Requires-Dist: mypy ==1.11.2 ; extra == 'dev'
+Requires-Dist: mypy ==1.12.0 ; extra == 'dev'
 Requires-Dist: types-python-dateutil ; extra == 'dev'
 Requires-Dist: types-pytz ; extra == 'dev'
 Requires-Dist: types-PyYAML ; extra == 'dev'

{datachain-0.6.0.dist-info → datachain-0.6.1.dist-info}/RECORD

@@ -1,4 +1,4 @@
-datachain/__init__.py,sha256=ofPJ6B-d-ybSDRrE7J6wqF_ZRAB2W9U8l-eeuBtqPLg,865
+datachain/__init__.py,sha256=OGzc8xZWtwqxiiutjU4AxCRPY0lrX_csgERiTrq4G0o,908
 datachain/__main__.py,sha256=hG3Y4ARGEqe1AWwNMd259rBlqtphx1Wk39YbueQ0yV8,91
 datachain/asyn.py,sha256=Lg3Ck1PQLjQziMx9KU4atzbEnJXTE0924WMYkhgWtGU,8247
 datachain/cache.py,sha256=s0YHN7qurmQv-eC265TjeureK84TebWWAnL07cxchZQ,2997
@@ -18,7 +18,7 @@ datachain/storage.py,sha256=RiSJLYdHUjnrEWkLBKPcETHpAxld_B2WxLg711t0aZI,3733
 datachain/telemetry.py,sha256=0A4IOPPp9VlP5pyW9eBfaTK3YhHGzHl7dQudQjUAx9A,994
 datachain/utils.py,sha256=KeFSRHsiYthnTu4a6bH-rw04mX1m8krTX0f2NqfQGFI,12114
 datachain/catalog/__init__.py,sha256=g2iAAFx_gEIrqshXlhSEbrc8qDaEH11cjU40n3CHDz4,409
-datachain/catalog/catalog.py,sha256=BsMyk2RQibQYHgrmovFZeSEpPVMTwgb_7ntVYdc7t-E,64090
+datachain/catalog/catalog.py,sha256=r5lkwwZDh8cETNniBdzPCY9Ix8G-1RdkehjvUe3d2nE,63834
 datachain/catalog/datasource.py,sha256=D-VWIVDCM10A8sQavLhRXdYSCG7F4o4ifswEF80_NAQ,1412
 datachain/catalog/loader.py,sha256=-6VelNfXUdgUnwInVyA8g86Boxv2xqhTh9xNS-Zlwig,8242
 datachain/client/__init__.py,sha256=T4wiYL9KIM0ZZ_UqIyzV8_ufzYlewmizlV4iymHNluE,86
@@ -36,14 +36,14 @@ datachain/data_storage/job.py,sha256=w-7spowjkOa1P5fUVtJou3OltT0L48P0RYWZ9rSJ9-s
 datachain/data_storage/metastore.py,sha256=HfCxk4lmDUg2Q4WsFNQGMWxllP0mToA00fxkFTwdNIE,52919
 datachain/data_storage/schema.py,sha256=AGbjyEir5UmRZXI3m0jChZogUh5wd8csj6-YlUWaAxQ,8383
 datachain/data_storage/serializer.py,sha256=6G2YtOFqqDzJf1KbvZraKGXl2XHZyVml2krunWUum5o,927
-datachain/data_storage/sqlite.py,sha256=fW08P7AbJ0cDbTbcTKuAGpvMXvBjg-QkGsKT_Dslyws,28383
-datachain/data_storage/warehouse.py,sha256=fXhVfao3NfWFGbbG5uJ-Ga4bX1FiKVfcbDyQgECYfk8,32122
+datachain/data_storage/sqlite.py,sha256=V8fGRPjSwIT7kdw1qyQfUfdqGjXB8dE68npkyXfKW0o,28702
+datachain/data_storage/warehouse.py,sha256=Ea0wVcWxe7Bu-8V8eqrPJ8Ov5-DT1dvv1MgxMINettc,31931
 datachain/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/lib/arrow.py,sha256=0R2CYsN82nNa5_03iS6jVix9EKeeqNZNAMgpSQP2hfo,9482
 datachain/lib/clip.py,sha256=lm5CzVi4Cj1jVLEKvERKArb-egb9j1Ls-fwTItT6vlI,6150
 datachain/lib/data_model.py,sha256=ECTbvlnzM98hp2mZ4fo82Yi0-MuoqTIQasQKGIyd89I,2040
 datachain/lib/dataset_info.py,sha256=srPPhI2UHf6hFPBecyFEVw2SS5aPisIIMsvGgKqi7ss,2366
-datachain/lib/dc.py,sha256=XmAFU9k79wUHIh0gYab8j-wF4vIlyW6opJcOy8fmoVc,76666
+datachain/lib/dc.py,sha256=wEqBDCENfBmeow0-uu8R4qJhQa8taEIzveUiNdr2CyY,78341
 datachain/lib/file.py,sha256=LjTW_-PDAnoUhvyB4bJ8Y8n__XGqrxvmd9mDOF0Gir8,14875
 datachain/lib/hf.py,sha256=cPnmLuprr0pYABH7KqA5FARQ1JGlywdDwD3yDzVAm4k,5920
 datachain/lib/image.py,sha256=AMXYwQsmarZjRbPCZY3M1jDsM2WAB_b3cTY4uOIuXNU,2675
@@ -53,30 +53,33 @@ datachain/lib/meta_formats.py,sha256=3f-0vpMTesagS9iMd3y9-u9r-7g0eqYsxmK4fVfNWlw
 datachain/lib/model_store.py,sha256=DNIv8Y6Jtk1_idNLzIpsThOsdW2BMAudyUCbPUcgcxk,2515
 datachain/lib/pytorch.py,sha256=W-ARi2xH1f1DUkVfRuerW-YWYgSaJASmNCxtz2lrJGI,6072
 datachain/lib/settings.py,sha256=39thOpYJw-zPirzeNO6pmRC2vPrQvt4eBsw1xLWDFsw,2344
-datachain/lib/signal_schema.py,sha256=gj45dRQuOsKDmaKaJxb5j63HYVGw-Ks1fyAS1FpyOWA,24145
+datachain/lib/signal_schema.py,sha256=6fgQIZz4jFvuiaL1mqK5Cq6yr4WC57o2ptHxk36MRNY,24438
 datachain/lib/tar.py,sha256=3WIzao6yD5fbLqXLTt9GhPGNonbFIs_fDRu-9vgLgsA,1038
 datachain/lib/text.py,sha256=UNHm8fhidk7wdrWqacEWaA6I9ykfYqarQ2URby7jc7M,1261
 datachain/lib/udf.py,sha256=GvhWLCXZUY7sz1QMRBj1AJDSzzhyj15xs3Ia9hjJrJE,12697
 datachain/lib/udf_signature.py,sha256=GXw24A-Olna6DWCdgy2bC-gZh_gLGPQ-KvjuI6pUjC0,7281
-datachain/lib/utils.py,sha256=5-kJlAZE0D9nXXweAjo7-SP_AWGo28feaDByONYaooQ,463
+datachain/lib/utils.py,sha256=12elAX6eTFgMGKIf2UfZ4IW07kRwjK6wz8yGE41RtNM,618
 datachain/lib/vfile.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/lib/webdataset.py,sha256=o7SHk5HOUWsZ5Ln04xOM04eQqiBHiJNO7xLgyVBrwo8,6924
 datachain/lib/webdataset_laion.py,sha256=aGMWeFmeYNK75ewO9JTA11iB1i3QtTzUfenQA5jajfo,2535
 datachain/lib/convert/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/lib/convert/flatten.py,sha256=Uebc5CeqCsacp-nr6IG9i6OGuUavXqdqnoGctZBk3RQ,1384
 datachain/lib/convert/python_to_sql.py,sha256=40SAOdoOgikZRhn8iomCPDRoxC3RFxjJLivEAA9MHDU,2880
-datachain/lib/convert/sql_to_python.py,sha256=lGnKzSF_tz9Y_5SSKkrIU95QEjpcDzvOxIRkEKTQag0,443
+datachain/lib/convert/sql_to_python.py,sha256=XXCBYDQFUXJIBNWkjEP944cnCfJ8GF2Tji0DLF3A_zQ,315
 datachain/lib/convert/unflatten.py,sha256=Ogvh_5wg2f38_At_1lN0D_e2uZOOpYEvwvB2xdq56Tw,2012
 datachain/lib/convert/values_to_tuples.py,sha256=varRCnSMT_pZmHznrd2Yi05qXLLz_v9YH_pOCpHSkdc,3921
-datachain/query/__init__.py,sha256=0NBOZVgIDpCcj1Ci883dQ9A0iiwe03xzmotkOCFbxYc,293
-datachain/query/batch.py,sha256=3QlwshhpUc1amZRtXWVXEEuq47hEQgQlY0Ji48DR6hg,3508
-datachain/query/dataset.py,sha256=MF_E7yjbFQV6NcP4gKbJFXiWuoQkpQ7-Jmxa59FxenE,53630
+datachain/lib/func/__init__.py,sha256=ucJ15J_Q5Hy--boKV-tPuhKagVD3NpnuUPhLtDp7doI,230
+datachain/lib/func/aggregate.py,sha256=B5VV6WoSYYiO_9uN4_nXPMkF9OOkgyE6suJ7XD-JiPI,938
+datachain/lib/func/func.py,sha256=kFhVZlWZzgAfM7-DpkpZWf5zzdEutp_3NxIFWxXww_I,1956
+datachain/query/__init__.py,sha256=7DhEIjAA8uZJfejruAVMZVcGFmvUpffuZJwgRqNwe-c,263
+datachain/query/batch.py,sha256=5fEhORFe7li12SdYddaSK3LyqksMfCHhwN1_A6TfsA4,3485
+datachain/query/dataset.py,sha256=-J8t8XGUQveh-4aM5HrnbYx9xLfMQ8p6P9sKmBaTpLU,52683
 datachain/query/dispatch.py,sha256=wjjTWw6sFQbB9SKRh78VbfvwSMgJXCfqJklS3-9KnCU,12025
 datachain/query/metrics.py,sha256=r5b0ygYhokbXp8Mg3kCH8iFSRw0jxzyeBe-C-J_bKFc,938
 datachain/query/params.py,sha256=O_j89mjYRLOwWNhYZl-z7mi-rkdP7WyFmaDufsdTryE,863
 datachain/query/queue.py,sha256=waqM_KzavU8C-G95-4211Nd4GXna_u2747Chgwtgz2w,3839
-datachain/query/schema.py,sha256=I8zLWJuWl5N332ni9mAzDYtcxMJupVPgWkSDe8spNEk,8019
-datachain/query/session.py,sha256=kpFFJMfWBnxaMPojMGhJRbk-BOsSYI8Ckl6vvqnx7d0,5787
+datachain/query/schema.py,sha256=b_KnVy6B26Ol4nYG0LqNNpeQ1QYPk95YRGUjXfdaQWs,6606
+datachain/query/session.py,sha256=50SOdLNCjqHHKI-L4xGXyzTVxzMWfANqKqjeYre-c2k,5959
 datachain/remote/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/remote/studio.py,sha256=f5s6qSZ9uB4URGUoU_8_W1KZRRQQVSm6cgEBkBUEfuE,7226
 datachain/sql/__init__.py,sha256=A2djrbQwSMUZZEIKGnm-mnRA-NDSbiDJNpAmmwGNyIo,303
@@ -85,20 +88,21 @@ datachain/sql/types.py,sha256=3aXpoxkmCYbw0Dlta5J1enwS8_FuvjfSqyrNZO-dWj4,13383
 datachain/sql/utils.py,sha256=rzlJw08etivdrcuQPqNVvVWhuVSyUPUQEEc6DOhu258,818
 datachain/sql/default/__init__.py,sha256=XQ2cEZpzWiABqjV-6yYHUBGI9vN_UHxbxZENESmVAWw,45
 datachain/sql/default/base.py,sha256=QD-31C6JnyOXzogyDx90sUhm7QvgXIYpeHEASH84igU,628
-datachain/sql/functions/__init__.py,sha256=Ioyy7nSetrTLVnHGcGcmZU99HxUFcx-5PFbrh2dPNH0,396
-datachain/sql/functions/array.py,sha256=EB7nJSncUc1PuxlHyzU2gVhF8DuXaxpGlxb5e8X2KFY,1297
+datachain/sql/functions/__init__.py,sha256=-vIkU0AqwOW5FX6P89xYl-uBIUdt46CEnCtshmN85gM,400
+datachain/sql/functions/aggregate.py,sha256=3AQdA8YHPFdtCEfwZKQXTT8SlQWdG9gD5PBtGN3Odqs,944
+datachain/sql/functions/array.py,sha256=rvH27SWN9gdh_mFnp0GIiXuCrNW6n8ZbY4I_JUS-_e0,1140
 datachain/sql/functions/conditional.py,sha256=q7YUKfunXeEldXaxgT-p5pUTcOEVU_tcQ2BJlquTRPs,207
 datachain/sql/functions/path.py,sha256=zixpERotTFP6LZ7I4TiGtyRA8kXOoZmH1yzH9oRW0mg,1294
 datachain/sql/functions/random.py,sha256=vBwEEj98VH4LjWixUCygQ5Bz1mv1nohsCG0-ZTELlVg,271
 datachain/sql/functions/string.py,sha256=DYgiw8XSk7ge7GXvyRI1zbaMruIizNeI-puOjriQGZQ,1148
 datachain/sql/sqlite/__init__.py,sha256=TAdJX0Bg28XdqPO-QwUVKy8rg78cgMileHvMNot7d04,166
-datachain/sql/sqlite/base.py,sha256=3gDMLKSWkxnbiZ1dykYa5VuHSSlg5sLY9ihMqcH_o1M,13578
+datachain/sql/sqlite/base.py,sha256=aHSZVvh4XSVkvZ07h3jMoRlHI4sWD8y3SnmGs9xMG9Y,14375
 datachain/sql/sqlite/types.py,sha256=yzvp0sXSEoEYXs6zaYC_2YubarQoZH-MiUNXcpuEP4s,1573
 datachain/sql/sqlite/vector.py,sha256=ncW4eu2FlJhrP_CIpsvtkUabZlQdl2D5Lgwy_cbfqR0,469
 datachain/torch/__init__.py,sha256=gIS74PoEPy4TB3X6vx9nLO0Y3sLJzsA8ckn8pRWihJM,579
-datachain-0.6.0.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.6.0.dist-info/METADATA,sha256=4nxP9eUg6o9ymkwy-hz4DsqRM5IBtqhInNE7vsE0lxY,17156
-datachain-0.6.0.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-datachain-0.6.0.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.6.0.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.6.0.dist-info/RECORD,,
+datachain-0.6.1.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.6.1.dist-info/METADATA,sha256=kOEDXkaNjPHB-A1fLt60s_EJvnjuLIU3xdfp5UhflUA,17156
+datachain-0.6.1.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+datachain-0.6.1.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.6.1.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.6.1.dist-info/RECORD,,