datachain 0.7.0__py3-none-any.whl → 0.7.2__py3-none-any.whl

datachain/__init__.py

@@ -1,4 +1,3 @@
-from datachain.lib import func, models
 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 (
@@ -35,9 +34,7 @@ __all__ = [
     "Sys",
     "TarVFile",
     "TextFile",
-    "func",
     "is_chain_type",
     "metrics",
-    "models",
     "param",
 ]

datachain/catalog/catalog.py

@@ -54,7 +54,6 @@ from datachain.error import (
     QueryScriptCancelError,
     QueryScriptRunError,
 )
-from datachain.listing import Listing
 from datachain.node import DirType, Node, NodeWithPath
 from datachain.nodes_thread_pool import NodesThreadPool
 from datachain.remote.studio import StudioClient
@@ -76,6 +75,7 @@ if TYPE_CHECKING:
     from datachain.dataset import DatasetVersion
     from datachain.job import Job
     from datachain.lib.file import File
+    from datachain.listing import Listing
 
 logger = logging.getLogger("datachain")
 
@@ -236,7 +236,7 @@ class DatasetRowsFetcher(NodesThreadPool):
 class NodeGroup:
     """Class for a group of nodes from the same source"""
 
-    listing: Listing
+    listing: "Listing"
     sources: list[DataSource]
 
     # The source path within the bucket
@@ -591,8 +591,9 @@ class Catalog:
         client_config=None,
         object_name="file",
         skip_indexing=False,
-    ) -> tuple[Listing, str]:
+    ) -> tuple["Listing", str]:
         from datachain.lib.dc import DataChain
+        from datachain.listing import Listing
 
         DataChain.from_storage(
             source, session=self.session, update=update, object_name=object_name
@@ -660,7 +661,8 @@ class Catalog:
         no_glob: bool = False,
         client_config=None,
     ) -> list[NodeGroup]:
-        from datachain.query import DatasetQuery
+        from datachain.listing import Listing
+        from datachain.query.dataset import DatasetQuery
 
         def _row_to_node(d: dict[str, Any]) -> Node:
             del d["file__source"]
@@ -876,7 +878,7 @@ class Catalog:
     def update_dataset_version_with_warehouse_info(
         self, dataset: DatasetRecord, version: int, rows_dropped=False, **kwargs
     ) -> None:
-        from datachain.query import DatasetQuery
+        from datachain.query.dataset import DatasetQuery
 
         dataset_version = dataset.get_version(version)
 
@@ -1177,7 +1179,7 @@ class Catalog:
     def ls_dataset_rows(
         self, name: str, version: int, offset=None, limit=None
     ) -> list[dict]:
-        from datachain.query import DatasetQuery
+        from datachain.query.dataset import DatasetQuery
 
         dataset = self.get_dataset(name)
 

datachain/cli.py

@@ -957,7 +957,7 @@ def show(
     schema: bool = False,
 ) -> None:
     from datachain.lib.dc import DataChain
-    from datachain.query import DatasetQuery
+    from datachain.query.dataset import DatasetQuery
     from datachain.utils import show_records
 
     dataset = catalog.get_dataset(name)

datachain/client/fsspec.py

@@ -28,7 +28,6 @@ from tqdm import tqdm
 from datachain.cache import DataChainCache
 from datachain.client.fileslice import FileWrapper
 from datachain.error import ClientError as DataChainClientError
-from datachain.lib.file import File
 from datachain.nodes_fetcher import NodesFetcher
 from datachain.nodes_thread_pool import NodeChunk
 
@@ -36,6 +35,7 @@ if TYPE_CHECKING:
     from fsspec.spec import AbstractFileSystem
 
     from datachain.dataset import StorageURI
+    from datachain.lib.file import File
 
 
 logger = logging.getLogger("datachain")
@@ -45,7 +45,7 @@ DELIMITER = "/"  # Path delimiter.
 
 DATA_SOURCE_URI_PATTERN = re.compile(r"^[\w]+:\/\/.*$")
 
-ResultQueue = asyncio.Queue[Optional[Sequence[File]]]
+ResultQueue = asyncio.Queue[Optional[Sequence["File"]]]
 
 
 def _is_win_local_path(uri: str) -> bool:
@@ -212,7 +212,7 @@ class Client(ABC):
 
     async def scandir(
         self, start_prefix: str, method: str = "default"
-    ) -> AsyncIterator[Sequence[File]]:
+    ) -> AsyncIterator[Sequence["File"]]:
         try:
             impl = getattr(self, f"_fetch_{method}")
         except AttributeError:
@@ -317,7 +317,7 @@ class Client(ABC):
         return f"{self.PREFIX}{self.name}/{rel_path}"
 
     @abstractmethod
-    def info_to_file(self, v: dict[str, Any], parent: str) -> File: ...
+    def info_to_file(self, v: dict[str, Any], parent: str) -> "File": ...
 
     def fetch_nodes(
         self,
@@ -354,7 +354,7 @@ class Client(ABC):
             copy2(src, dst)
 
     def open_object(
-        self, file: File, use_cache: bool = True, cb: Callback = DEFAULT_CALLBACK
+        self, file: "File", use_cache: bool = True, cb: Callback = DEFAULT_CALLBACK
     ) -> BinaryIO:
         """Open a file, including files in tar archives."""
         if use_cache and (cache_path := self.cache.get_path(file)):
@@ -362,19 +362,19 @@ class Client(ABC):
         assert not file.location
         return FileWrapper(self.fs.open(self.get_full_path(file.path)), cb)  # type: ignore[return-value]
 
-    def download(self, file: File, *, callback: Callback = DEFAULT_CALLBACK) -> None:
+    def download(self, file: "File", *, callback: Callback = DEFAULT_CALLBACK) -> None:
         sync(get_loop(), functools.partial(self._download, file, callback=callback))
 
-    async def _download(self, file: File, *, callback: "Callback" = None) -> None:
+    async def _download(self, file: "File", *, callback: "Callback" = None) -> None:
         if self.cache.contains(file):
             # Already in cache, so there's nothing to do.
             return
         await self._put_in_cache(file, callback=callback)
 
-    def put_in_cache(self, file: File, *, callback: "Callback" = None) -> None:
+    def put_in_cache(self, file: "File", *, callback: "Callback" = None) -> None:
         sync(get_loop(), functools.partial(self._put_in_cache, file, callback=callback))
 
-    async def _put_in_cache(self, file: File, *, callback: "Callback" = None) -> None:
+    async def _put_in_cache(self, file: "File", *, callback: "Callback" = None) -> None:
         assert not file.location
         if file.etag:
             etag = await self.get_current_etag(file)

datachain/data_storage/schema.py

@@ -12,7 +12,7 @@ import sqlalchemy as sa
 from sqlalchemy.sql import func as f
 from sqlalchemy.sql.expression import false, null, true
 
-from datachain.sql.functions import path
+from datachain.sql.functions import path as pathfunc
 from datachain.sql.types import Int, SQLType, UInt64
 
 if TYPE_CHECKING:
@@ -130,7 +130,7 @@ class DirExpansion:
 
     def query(self, q):
         q = self.base_select(q).cte(recursive=True)
-        parent = path.parent(self.c(q, "path"))
+        parent = pathfunc.parent(self.c(q, "path"))
         q = q.union_all(
             sa.select(
                 sa.literal(-1).label("sys__id"),

datachain/data_storage/sqlite.py

@@ -122,7 +122,9 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         return cls(*cls._connect(db_file=db_file))
 
     @staticmethod
-    def _connect(db_file: Optional[str] = None):
+    def _connect(
+        db_file: Optional[str] = None,
+    ) -> tuple["Engine", "MetaData", sqlite3.Connection, str]:
         try:
             if db_file == ":memory:":
                 # Enable multithreaded usage of the same in-memory db
@@ -130,9 +132,8 @@ class SQLiteDatabaseEngine(DatabaseEngine):
                     _get_in_memory_uri(), uri=True, detect_types=DETECT_TYPES
                 )
             else:
-                db = sqlite3.connect(
-                    db_file or DataChainDir.find().db, detect_types=DETECT_TYPES
-                )
+                db_file = db_file or DataChainDir.find().db
+                db = sqlite3.connect(db_file, detect_types=DETECT_TYPES)
             create_user_defined_sql_functions(db)
             engine = sqlalchemy.create_engine(
                 "sqlite+pysqlite:///", creator=lambda: db, future=True

datachain/data_storage/warehouse.py

@@ -224,28 +224,28 @@ class AbstractWarehouse(ABC, Serializable):
         offset = 0
         num_yielded = 0
 
-        while True:
-            if limit is not None:
-                limit -= num_yielded
-                if limit == 0:
-                    break
-                if limit < page_size:
-                    paginated_query = paginated_query.limit(None).limit(limit)
-
-            # Ensure we're using a thread-local connection
-            with self.clone() as wh:
+        # Ensure we're using a thread-local connection
+        with self.clone() as wh:
+            while True:
+                if limit is not None:
+                    limit -= num_yielded
+                    if limit == 0:
+                        break
+                    if limit < page_size:
+                        paginated_query = paginated_query.limit(None).limit(limit)
+
                 # Cursor results are not thread-safe, so we convert them to a list
                 results = list(wh.dataset_rows_select(paginated_query.offset(offset)))
 
-            processed = False
-            for row in results:
-                processed = True
-                yield row
-                num_yielded += 1
+                processed = False
+                for row in results:
+                    processed = True
+                    yield row
+                    num_yielded += 1
 
-            if not processed:
-                break  # no more results
-            offset += page_size
+                if not processed:
+                    break  # no more results
+                offset += page_size
 
     #
     # Table Name Internal Functions

datachain/func/__init__.py

@@ -0,0 +1,49 @@
+from sqlalchemy import literal
+
+from . import array, path, random, string
+from .aggregate import (
+    any_value,
+    avg,
+    collect,
+    concat,
+    count,
+    dense_rank,
+    first,
+    max,
+    min,
+    rank,
+    row_number,
+    sum,
+)
+from .array import cosine_distance, euclidean_distance, length, sip_hash_64
+from .conditional import greatest, least
+from .random import rand
+from .window import window
+
+__all__ = [
+    "any_value",
+    "array",
+    "avg",
+    "collect",
+    "concat",
+    "cosine_distance",
+    "count",
+    "dense_rank",
+    "euclidean_distance",
+    "first",
+    "greatest",
+    "least",
+    "length",
+    "literal",
+    "max",
+    "min",
+    "path",
+    "rand",
+    "random",
+    "rank",
+    "row_number",
+    "sip_hash_64",
+    "string",
+    "sum",
+    "window",
+]

datachain/{lib/func → func}/aggregate.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 from sqlalchemy import func as sa_func
 
-from datachain.sql import functions as dc_func
+from datachain.sql.functions import aggregate
 
 from .func import Func
 
@@ -31,7 +31,9 @@ def count(col: Optional[str] = None) -> Func:
     Notes:
         - Result column will always be of type int.
     """
-    return Func("count", inner=sa_func.count, col=col, result_type=int)
+    return Func(
+        "count", inner=sa_func.count, cols=[col] if col else None, result_type=int
+    )
 
 
 def sum(col: str) -> Func:
@@ -59,7 +61,7 @@ def sum(col: str) -> Func:
         - 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)
+    return Func("sum", inner=sa_func.sum, cols=[col])
 
 
 def avg(col: str) -> Func:
@@ -87,7 +89,7 @@ def avg(col: str) -> Func:
         - 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)
+    return Func("avg", inner=aggregate.avg, cols=[col], result_type=float)
 
 
 def min(col: str) -> Func:
@@ -115,7 +117,7 @@ def min(col: str) -> Func:
         - 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)
+    return Func("min", inner=sa_func.min, cols=[col])
 
 
 def max(col: str) -> Func:
@@ -143,7 +145,7 @@ def max(col: str) -> Func:
         - 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)
+    return Func("max", inner=sa_func.max, cols=[col])
 
 
 def any_value(col: str) -> Func:
@@ -174,7 +176,7 @@ def any_value(col: str) -> Func:
         - 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)
+    return Func("any_value", inner=aggregate.any_value, cols=[col])
 
 
 def collect(col: str) -> Func:
@@ -203,7 +205,7 @@ def collect(col: str) -> Func:
         - 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)
+    return Func("collect", inner=aggregate.collect, cols=[col], is_array=True)
 
 
 def concat(col: str, separator="") -> Func:
@@ -236,9 +238,9 @@ def concat(col: str, separator="") -> Func:
     """
 
     def inner(arg):
-        return dc_func.aggregate.group_concat(arg, separator)
+        return aggregate.group_concat(arg, separator)
 
-    return Func("concat", inner=inner, col=col, result_type=str)
+    return Func("concat", inner=inner, cols=[col], result_type=str)
 
 
 def row_number() -> Func:
@@ -350,4 +352,4 @@ def first(col: str) -> Func:
           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)
+    return Func("first", inner=sa_func.first_value, cols=[col], is_window=True)

datachain/func/array.py

@@ -0,0 +1,176 @@
+from collections.abc import Sequence
+from typing import Union
+
+from datachain.sql.functions import array
+
+from .func import Func
+
+
+def cosine_distance(*args: Union[str, Sequence]) -> Func:
+    """
+    Computes 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,
+    where 0 indicates identical vectors and values closer to 1
+    indicate higher dissimilarity.
+
+    Args:
+        args (str | 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 sequence is provided, it is assumed to be a vector of values.
+
+    Returns:
+        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]),
+        )
+        ```
+
+    Notes:
+        - Ensure both vectors have the same number of elements.
+        - Result column will always be of type float.
+    """
+    cols, func_args = [], []
+    for arg in args:
+        if isinstance(arg, str):
+            cols.append(arg)
+        else:
+            func_args.append(list(arg))
+
+    if len(cols) + len(func_args) != 2:
+        raise ValueError("cosine_distance() requires exactly two arguments")
+    if not cols and len(func_args[0]) != len(func_args[1]):
+        raise ValueError("cosine_distance() requires vectors of the same length")
+
+    return Func(
+        "cosine_distance",
+        inner=array.cosine_distance,
+        cols=cols,
+        args=func_args,
+        result_type=float,
+    )
+
+
+def euclidean_distance(*args: Union[str, Sequence]) -> Func:
+    """
+    Computes 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.
+            If a string is provided, it is assumed to be the name of the column vector.
+            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.
+
+    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]),
+        )
+        ```
+
+    Notes:
+        - Ensure both vectors have the same number of elements.
+        - Result column will always be of type float.
+    """
+    cols, func_args = [], []
+    for arg in args:
+        if isinstance(arg, str):
+            cols.append(arg)
+        else:
+            func_args.append(list(arg))
+
+    if len(cols) + len(func_args) != 2:
+        raise ValueError("euclidean_distance() requires exactly two arguments")
+    if not cols and len(func_args[0]) != len(func_args[1]):
+        raise ValueError("euclidean_distance() requires vectors of the same length")
+
+    return Func(
+        "euclidean_distance",
+        inner=array.euclidean_distance,
+        cols=cols,
+        args=func_args,
+        result_type=float,
+    )
+
+
+def length(arg: Union[str, Sequence, Func]) -> Func:
+    """
+    Returns the length of the array.
+
+    Args:
+        arg (str | Sequence | Func): 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 Func is provided, it is assumed to be a function returning an array.
+
+    Returns:
+        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]),
+        )
+        ```
+
+    Note:
+        - Result column will always be of type int.
+    """
+    if isinstance(arg, (str, Func)):
+        cols = [arg]
+        args = None
+    else:
+        cols = None
+        args = [arg]
+
+    return Func("length", inner=array.length, cols=cols, args=args, result_type=int)
+
+
+def sip_hash_64(arg: Union[str, Sequence]) -> Func:
+    """
+    Computes the SipHash-64 hash of the array.
+
+    Args:
+        arg (str | 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 sequence is provided, it is assumed to be an array of values.
+
+    Returns:
+        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]),
+        )
+        ```
+
+    Note:
+        - This function is only available for the ClickHouse warehouse.
+        - Result column will always be of type int.
+    """
+    if isinstance(arg, str):
+        cols = [arg]
+        args = None
+    else:
+        cols = None
+        args = [arg]
+
+    return Func(
+        "sip_hash_64", inner=array.sip_hash_64, cols=cols, args=args, result_type=int
+    )

datachain/func/base.py

@@ -0,0 +1,23 @@
+from abc import ABCMeta, abstractmethod
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from sqlalchemy import TableClause
+
+    from datachain.lib.signal_schema import SignalSchema
+    from datachain.query.schema import Column
+
+
+class Function:
+    __metaclass__ = ABCMeta
+
+    name: str
+
+    @abstractmethod
+    def get_column(
+        self,
+        signals_schema: Optional["SignalSchema"] = None,
+        label: Optional[str] = None,
+        table: Optional["TableClause"] = None,
+    ) -> "Column":
+        pass

datachain/func/conditional.py

@@ -0,0 +1,81 @@
+from typing import Union
+
+from datachain.sql.functions import conditional
+
+from .func import ColT, Func
+
+
+def greatest(*args: Union[ColT, float]) -> Func:
+    """
+    Returns the greatest (largest) value from the given input values.
+
+    Args:
+        args (ColT | str | int | float | Sequence): The values to compare.
+            If a string is provided, it is assumed to be the name of the column.
+            If a Func is provided, it is assumed to be a function returning a value.
+            If an int, float, or Sequence is provided, it is assumed to be a literal.
+
+    Returns:
+        Func: A Func object that represents the greatest function.
+
+    Example:
+        ```py
+        dc.mutate(
+            greatest=func.greatest("signal.value", 0),
+        )
+        ```
+
+    Note:
+        - Result column will always be of the same type as the input columns.
+    """
+    cols, func_args = [], []
+
+    for arg in args:
+        if isinstance(arg, (str, Func)):
+            cols.append(arg)
+        else:
+            func_args.append(arg)
+
+    return Func(
+        "greatest",
+        inner=conditional.greatest,
+        cols=cols,
+        args=func_args,
+        result_type=int,
+    )
+
+
+def least(*args: Union[ColT, float]) -> Func:
+    """
+    Returns the least (smallest) value from the given input values.
+
+    Args:
+        args (ColT | str | int | float | Sequence): The values to compare.
+            If a string is provided, it is assumed to be the name of the column.
+            If a Func is provided, it is assumed to be a function returning a value.
+            If an int, float, or Sequence is provided, it is assumed to be a literal.
+
+    Returns:
+        Func: A Func object that represents the least function.
+
+    Example:
+        ```py
+        dc.mutate(
+            least=func.least("signal.value", 0),
+        )
+        ```
+
+    Note:
+        - Result column will always be of the same type as the input columns.
+    """
+    cols, func_args = [], []
+
+    for arg in args:
+        if isinstance(arg, (str, Func)):
+            cols.append(arg)
+        else:
+            func_args.append(arg)
+
+    return Func(
+        "least", inner=conditional.least, cols=cols, args=func_args, result_type=int
+    )