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

datachain/query/dataset.py

@@ -43,9 +43,10 @@ from datachain.data_storage.schema import (
 )
 from datachain.dataset import DatasetStatus, RowDict
 from datachain.error import DatasetNotFoundError, QueryScriptCancelError
+from datachain.func.base import Function
 from datachain.lib.udf import UDFAdapter
 from datachain.progress import CombinedDownloadCallback
-from datachain.sql.functions import rand
+from datachain.sql.functions.random import rand
 from datachain.utils import (
     batched,
     determine_processes,
@@ -65,15 +66,16 @@ if TYPE_CHECKING:
     from datachain.catalog import Catalog
     from datachain.data_storage import AbstractWarehouse
     from datachain.dataset import DatasetRecord
-
-    from .udf import UDFResult
+    from datachain.lib.udf import UDFResult
 
     P = ParamSpec("P")
 
 
 INSERT_BATCH_SIZE = 10000
 
-PartitionByType = Union[ColumnElement, Sequence[ColumnElement]]
+PartitionByType = Union[
+    Function, ColumnElement, Sequence[Union[Function, ColumnElement]]
+]
 JoinPredicateType = Union[str, ColumnClause, ColumnElement]
 DatasetDependencyType = tuple[str, int]
 
@@ -457,18 +459,15 @@ class UDFStep(Step, ABC):
                 # Run the UDFDispatcher in another process to avoid needing
                 # if __name__ == '__main__': in user scripts
                 exec_cmd = get_datachain_executable()
+                cmd = [*exec_cmd, "internal-run-udf"]
                 envs = dict(os.environ)
                 envs.update({"PYTHONPATH": os.getcwd()})
                 process_data = filtered_cloudpickle_dumps(udf_info)
-                result = subprocess.run(  # noqa: S603
-                    [*exec_cmd, "internal-run-udf"],
-                    input=process_data,
-                    check=False,
-                    env=envs,
-                )
-                if result.returncode != 0:
-                    raise RuntimeError("UDF Execution Failed!")
 
+                with subprocess.Popen(cmd, env=envs, stdin=subprocess.PIPE) as process:  # noqa: S603
+                    process.communicate(process_data)
+                    if process.poll():
+                        raise RuntimeError("UDF Execution Failed!")
             else:
                 # Otherwise process single-threaded (faster for smaller UDFs)
                 warehouse = self.catalog.warehouse
@@ -520,13 +519,17 @@ class UDFStep(Step, ABC):
         else:
             list_partition_by = [self.partition_by]
 
+        partition_by = [
+            p.get_column() if isinstance(p, Function) else p for p in list_partition_by
+        ]
+
         # create table with partitions
         tbl = self.catalog.warehouse.create_udf_table(partition_columns())
 
         # fill table with partitions
         cols = [
             query.selected_columns.sys__id,
-            f.dense_rank().over(order_by=list_partition_by).label(PARTITION_COLUMN_ID),
+            f.dense_rank().over(order_by=partition_by).label(PARTITION_COLUMN_ID),
         ]
         self.catalog.warehouse.db.execute(
             tbl.insert().from_select(cols, query.with_only_columns(*cols))
@@ -683,6 +686,12 @@ class SQLClause(Step, ABC):
 
         return step_result(q, new_query.selected_columns)
 
+    def parse_cols(
+        self,
+        cols: Sequence[Union[Function, ColumnElement]],
+    ) -> tuple[ColumnElement, ...]:
+        return tuple(c.get_column() if isinstance(c, Function) else c for c in cols)
+
     @abstractmethod
     def apply_sql_clause(self, query):
         pass
@@ -690,12 +699,14 @@ class SQLClause(Step, ABC):
 
 @frozen
 class SQLSelect(SQLClause):
-    args: tuple[Union[str, ColumnElement], ...]
+    args: tuple[Union[Function, ColumnElement], ...]
 
     def apply_sql_clause(self, query) -> Select:
         subquery = query.subquery()
-
-        args = [subquery.c[str(c)] if isinstance(c, (str, C)) else c for c in self.args]
+        args = [
+            subquery.c[str(c)] if isinstance(c, (str, C)) else c
+            for c in self.parse_cols(self.args)
+        ]
         if not args:
             args = subquery.c
 
@@ -704,22 +715,25 @@ class SQLSelect(SQLClause):
 
 @frozen
 class SQLSelectExcept(SQLClause):
-    args: tuple[str, ...]
+    args: tuple[Union[Function, ColumnElement], ...]
 
     def apply_sql_clause(self, query: Select) -> Select:
         subquery = query.subquery()
-        names = set(self.args)
-        args = [c for c in subquery.c if c.name not in names]
+        args = [c for c in subquery.c if c.name not in set(self.parse_cols(self.args))]
         return sqlalchemy.select(*args).select_from(subquery)
 
 
 @frozen
 class SQLMutate(SQLClause):
-    args: tuple[ColumnElement, ...]
+    args: tuple[Union[Function, ColumnElement], ...]
 
     def apply_sql_clause(self, query: Select) -> Select:
         original_subquery = query.subquery()
-        to_mutate = {c.name for c in self.args}
+        args = [
+            original_subquery.c[str(c)] if isinstance(c, (str, C)) else c
+            for c in self.parse_cols(self.args)
+        ]
+        to_mutate = {c.name for c in args}
 
         prefix = f"mutate{token_hex(8)}_"
         cols = [
@@ -729,9 +743,7 @@ class SQLMutate(SQLClause):
         # this is needed for new column to be used in clauses
         # like ORDER BY, otherwise new column is not recognized
         subquery = (
-            sqlalchemy.select(*cols, *self.args)
-            .select_from(original_subquery)
-            .subquery()
+            sqlalchemy.select(*cols, *args).select_from(original_subquery).subquery()
         )
 
         return sqlalchemy.select(*subquery.c).select_from(subquery)
@@ -739,21 +751,24 @@ class SQLMutate(SQLClause):
 
 @frozen
 class SQLFilter(SQLClause):
-    expressions: tuple[ColumnElement, ...]
+    expressions: tuple[Union[Function, ColumnElement], ...]
 
     def __and__(self, other):
-        return self.__class__(self.expressions + other)
+        expressions = self.parse_cols(self.expressions)
+        return self.__class__(expressions + other)
 
     def apply_sql_clause(self, query: Select) -> Select:
-        return query.filter(*self.expressions)
+        expressions = self.parse_cols(self.expressions)
+        return query.filter(*expressions)
 
 
 @frozen
 class SQLOrderBy(SQLClause):
-    args: tuple[ColumnElement, ...]
+    args: tuple[Union[Function, ColumnElement], ...]
 
     def apply_sql_clause(self, query: Select) -> Select:
-        return query.order_by(*self.args)
+        args = self.parse_cols(self.args)
+        return query.order_by(*args)
 
 
 @frozen
@@ -948,8 +963,8 @@ class SQLJoin(Step):
 
 @frozen
 class SQLGroupBy(SQLClause):
-    cols: Sequence[Union[str, ColumnElement]]
-    group_by: Sequence[Union[str, ColumnElement]]
+    cols: Sequence[Union[str, Function, ColumnElement]]
+    group_by: Sequence[Union[str, Function, ColumnElement]]
 
     def apply_sql_clause(self, query) -> Select:
         if not self.cols:
@@ -959,12 +974,20 @@ class SQLGroupBy(SQLClause):
 
         subquery = query.subquery()
 
+        group_by = [
+            c.get_column() if isinstance(c, Function) else c for c in self.group_by
+        ]
+
         cols = [
-            subquery.c[str(c)] if isinstance(c, (str, C)) else c
-            for c in [*self.group_by, *self.cols]
+            c.get_column()
+            if isinstance(c, Function)
+            else subquery.c[str(c)]
+            if isinstance(c, (str, C))
+            else c
+            for c in (*group_by, *self.cols)
         ]
 
-        return sqlalchemy.select(*cols).select_from(subquery).group_by(*self.group_by)
+        return sqlalchemy.select(*cols).select_from(subquery).group_by(*group_by)
 
 
 def _validate_columns(

datachain/sql/__init__.py

@@ -1,13 +1,11 @@
 from sqlalchemy.sql.elements import literal
 from sqlalchemy.sql.expression import column
 
-from . import functions
 from .default import setup as default_setup
 from .selectable import select, values
 
 __all__ = [
     "column",
-    "functions",
     "literal",
     "select",
     "values",

datachain/sql/functions/__init__.py

@@ -1,26 +0,0 @@
-from sqlalchemy.sql.expression import func
-
-from . import array, path, string
-from .aggregate import avg
-from .conditional import greatest, least
-from .random import rand
-
-count = func.count
-sum = func.sum
-min = func.min
-max = func.max
-
-__all__ = [
-    "array",
-    "avg",
-    "count",
-    "func",
-    "greatest",
-    "least",
-    "max",
-    "min",
-    "path",
-    "rand",
-    "string",
-    "sum",
-]

datachain/sql/selectable.py

@@ -9,7 +9,9 @@ class Values(selectable.Values):
             columns = [expression.column(f"c{i}") for i in range(1, num_columns + 1)]
         else:
             columns = [
-                expression.column(c) if isinstance(c, str) else c for c in columns
+                process_column_expression(c)
+                for c in columns
+                # expression.column(c) if isinstance(c, str) else c for c in columns
             ]
         super().__init__(*columns, **kwargs)
         self._data += tuple(data)
@@ -19,13 +21,17 @@ def values(data, columns=None, **kwargs) -> Values:
     return Values(data, columns=columns, **kwargs)
 
 
-def process_column_expressions(columns):
-    return [expression.column(c) if isinstance(c, str) else c for c in columns]
+def process_column_expression(col):
+    if hasattr(col, "get_column"):
+        return col.get_column()
+    if isinstance(col, str):
+        return expression.column(col)
+    return col
 
 
 def select(*columns, **kwargs) -> "expression.Select":
-    columns = process_column_expressions(columns)
-    return expression.select(*columns, **kwargs)
+    columns_processed = [process_column_expression(c) for c in columns]
+    return expression.select(*columns_processed, **kwargs)
 
 
 def base_values_compiler(column_name_func, element, compiler, **kwargs):

datachain/sql/sqlite/base.py

@@ -1,6 +1,7 @@
 import logging
 import re
 import sqlite3
+import warnings
 from collections.abc import Iterable
 from datetime import MAXYEAR, MINYEAR, datetime, timezone
 from types import MappingProxyType
@@ -418,14 +419,22 @@ def compile_collect(element, compiler, **kwargs):
     return compiler.process(func.json_group_array(*element.clauses.clauses), **kwargs)
 
 
-def load_usearch_extension(conn) -> bool:
+def load_usearch_extension(conn: sqlite3.Connection) -> bool:
     try:
         # usearch is part of the vector optional dependencies
         # we use the extension's cosine and euclidean distance functions
         from usearch import sqlite_path
 
         conn.enable_load_extension(True)
-        conn.load_extension(sqlite_path())
+
+        with warnings.catch_warnings():
+            # usearch binary is not available for Windows, see: https://github.com/unum-cloud/usearch/issues/427.
+            # and, sometimes fail to download the binary in other platforms
+            # triggering UserWarning.
+
+            warnings.filterwarnings("ignore", category=UserWarning, module="usearch")
+            conn.load_extension(sqlite_path())
+
         conn.enable_load_extension(False)
         return True
 

datachain/toolkit/split.py

@@ -58,10 +58,14 @@ def train_test_split(dc: DataChain, weights: list[float]) -> list[DataChain]:
 
     weights_normalized = [weight / sum(weights) for weight in weights]
 
+    resolution = 2**31 - 1  # Maximum positive value for a 32-bit signed integer.
+
     return [
         dc.filter(
-            C("sys__rand") % 1000 >= round(sum(weights_normalized[:index]) * 1000),
-            C("sys__rand") % 1000 < round(sum(weights_normalized[: index + 1]) * 1000),
+            C("sys__rand") % resolution
+            >= round(sum(weights_normalized[:index]) * resolution),
+            C("sys__rand") % resolution
+            < round(sum(weights_normalized[: index + 1]) * resolution),
         )
         for index, _ in enumerate(weights_normalized)
     ]

{datachain-0.7.0.dist-info → datachain-0.7.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.7.0
+Version: 0.7.2
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -19,85 +19,86 @@ License-File: LICENSE
 Requires-Dist: pyyaml
 Requires-Dist: tomlkit
 Requires-Dist: tqdm
-Requires-Dist: numpy <3,>=1
-Requires-Dist: pandas >=2.0.0
+Requires-Dist: numpy<3,>=1
+Requires-Dist: pandas>=2.0.0
 Requires-Dist: pyarrow
 Requires-Dist: typing-extensions
-Requires-Dist: python-dateutil >=2
-Requires-Dist: attrs >=21.3.0
-Requires-Dist: s3fs >=2024.2.0
-Requires-Dist: gcsfs >=2024.2.0
-Requires-Dist: adlfs >=2024.2.0
-Requires-Dist: dvc-data <4,>=3.10
-Requires-Dist: dvc-objects <6,>=4
-Requires-Dist: shtab <2,>=1.3.4
-Requires-Dist: sqlalchemy >=2
-Requires-Dist: multiprocess ==0.70.16
+Requires-Dist: python-dateutil>=2
+Requires-Dist: attrs>=21.3.0
+Requires-Dist: s3fs>=2024.2.0
+Requires-Dist: gcsfs>=2024.2.0
+Requires-Dist: adlfs>=2024.2.0
+Requires-Dist: dvc-data<4,>=3.10
+Requires-Dist: dvc-objects<6,>=4
+Requires-Dist: shtab<2,>=1.3.4
+Requires-Dist: sqlalchemy>=2
+Requires-Dist: multiprocess==0.70.16
 Requires-Dist: cloudpickle
-Requires-Dist: orjson >=3.10.5
-Requires-Dist: pydantic <3,>=2
-Requires-Dist: jmespath >=1.0
-Requires-Dist: datamodel-code-generator >=0.25
-Requires-Dist: Pillow <12,>=10.0.0
-Requires-Dist: msgpack <2,>=1.0.4
+Requires-Dist: orjson>=3.10.5
+Requires-Dist: pydantic<3,>=2
+Requires-Dist: jmespath>=1.0
+Requires-Dist: datamodel-code-generator>=0.25
+Requires-Dist: Pillow<12,>=10.0.0
+Requires-Dist: msgpack<2,>=1.0.4
 Requires-Dist: psutil
-Requires-Dist: huggingface-hub
-Requires-Dist: iterative-telemetry >=0.0.9
+Requires-Dist: huggingface_hub
+Requires-Dist: iterative-telemetry>=0.0.9
 Requires-Dist: platformdirs
-Requires-Dist: dvc-studio-client <1,>=0.21
+Requires-Dist: dvc-studio-client<1,>=0.21
 Requires-Dist: tabulate
-Provides-Extra: dev
-Requires-Dist: datachain[docs,tests] ; extra == 'dev'
-Requires-Dist: mypy ==1.13.0 ; extra == 'dev'
-Requires-Dist: types-python-dateutil ; extra == 'dev'
-Requires-Dist: types-pytz ; extra == 'dev'
-Requires-Dist: types-PyYAML ; extra == 'dev'
-Requires-Dist: types-requests ; extra == 'dev'
-Requires-Dist: types-tabulate ; extra == 'dev'
 Provides-Extra: docs
-Requires-Dist: mkdocs >=1.5.2 ; extra == 'docs'
-Requires-Dist: mkdocs-gen-files >=0.5.0 ; extra == 'docs'
-Requires-Dist: mkdocs-material >=9.3.1 ; extra == 'docs'
-Requires-Dist: mkdocs-section-index >=0.3.6 ; extra == 'docs'
-Requires-Dist: mkdocstrings-python >=1.6.3 ; extra == 'docs'
-Requires-Dist: mkdocs-literate-nav >=0.6.1 ; extra == 'docs'
-Provides-Extra: examples
-Requires-Dist: datachain[tests] ; extra == 'examples'
-Requires-Dist: numpy <2,>=1 ; extra == 'examples'
-Requires-Dist: defusedxml ; extra == 'examples'
-Requires-Dist: accelerate ; extra == 'examples'
-Requires-Dist: unstructured[embed-huggingface,pdf] <0.16.0 ; extra == 'examples'
-Requires-Dist: pdfplumber ==0.11.4 ; extra == 'examples'
-Requires-Dist: huggingface-hub[hf_transfer] ; extra == 'examples'
-Requires-Dist: onnx ==1.16.1 ; extra == 'examples'
-Provides-Extra: hf
-Requires-Dist: numba >=0.60.0 ; extra == 'hf'
-Requires-Dist: datasets[audio,vision] >=2.21.0 ; extra == 'hf'
-Provides-Extra: remote
-Requires-Dist: lz4 ; extra == 'remote'
-Requires-Dist: requests >=2.22.0 ; extra == 'remote'
-Provides-Extra: tests
-Requires-Dist: datachain[hf,remote,torch,vector] ; extra == 'tests'
-Requires-Dist: pytest <9,>=8 ; extra == 'tests'
-Requires-Dist: pytest-sugar >=0.9.6 ; extra == 'tests'
-Requires-Dist: pytest-cov >=4.1.0 ; extra == 'tests'
-Requires-Dist: pytest-mock >=3.12.0 ; extra == 'tests'
-Requires-Dist: pytest-servers[all] >=0.5.8 ; extra == 'tests'
-Requires-Dist: pytest-benchmark[histogram] ; extra == 'tests'
-Requires-Dist: pytest-xdist >=3.3.1 ; extra == 'tests'
-Requires-Dist: virtualenv ; extra == 'tests'
-Requires-Dist: dulwich ; extra == 'tests'
-Requires-Dist: hypothesis ; extra == 'tests'
-Requires-Dist: open-clip-torch ; extra == 'tests'
-Requires-Dist: aiotools >=1.7.0 ; extra == 'tests'
-Requires-Dist: requests-mock ; extra == 'tests'
-Requires-Dist: scipy ; extra == 'tests'
+Requires-Dist: mkdocs>=1.5.2; extra == "docs"
+Requires-Dist: mkdocs-gen-files>=0.5.0; extra == "docs"
+Requires-Dist: mkdocs-material>=9.3.1; extra == "docs"
+Requires-Dist: mkdocs-section-index>=0.3.6; extra == "docs"
+Requires-Dist: mkdocstrings-python>=1.6.3; extra == "docs"
+Requires-Dist: mkdocs-literate-nav>=0.6.1; extra == "docs"
 Provides-Extra: torch
-Requires-Dist: torch >=2.1.0 ; extra == 'torch'
-Requires-Dist: torchvision ; extra == 'torch'
-Requires-Dist: transformers >=4.36.0 ; extra == 'torch'
+Requires-Dist: torch>=2.1.0; extra == "torch"
+Requires-Dist: torchvision; extra == "torch"
+Requires-Dist: transformers>=4.36.0; extra == "torch"
+Provides-Extra: remote
+Requires-Dist: lz4; extra == "remote"
+Requires-Dist: requests>=2.22.0; extra == "remote"
 Provides-Extra: vector
-Requires-Dist: usearch ; extra == 'vector'
+Requires-Dist: usearch; extra == "vector"
+Provides-Extra: hf
+Requires-Dist: numba>=0.60.0; extra == "hf"
+Requires-Dist: datasets[audio,vision]>=2.21.0; extra == "hf"
+Provides-Extra: tests
+Requires-Dist: datachain[hf,remote,torch,vector]; extra == "tests"
+Requires-Dist: pytest<9,>=8; extra == "tests"
+Requires-Dist: pytest-sugar>=0.9.6; extra == "tests"
+Requires-Dist: pytest-cov>=4.1.0; extra == "tests"
+Requires-Dist: pytest-mock>=3.12.0; extra == "tests"
+Requires-Dist: pytest-servers[all]>=0.5.8; extra == "tests"
+Requires-Dist: pytest-benchmark[histogram]; extra == "tests"
+Requires-Dist: pytest-xdist>=3.3.1; extra == "tests"
+Requires-Dist: virtualenv; extra == "tests"
+Requires-Dist: dulwich; extra == "tests"
+Requires-Dist: hypothesis; extra == "tests"
+Requires-Dist: open_clip_torch; extra == "tests"
+Requires-Dist: aiotools>=1.7.0; extra == "tests"
+Requires-Dist: requests-mock; extra == "tests"
+Requires-Dist: scipy; extra == "tests"
+Provides-Extra: dev
+Requires-Dist: datachain[docs,tests]; extra == "dev"
+Requires-Dist: mypy==1.13.0; extra == "dev"
+Requires-Dist: types-python-dateutil; extra == "dev"
+Requires-Dist: types-pytz; extra == "dev"
+Requires-Dist: types-PyYAML; extra == "dev"
+Requires-Dist: types-requests; extra == "dev"
+Requires-Dist: types-tabulate; extra == "dev"
+Provides-Extra: examples
+Requires-Dist: datachain[tests]; extra == "examples"
+Requires-Dist: numpy<2,>=1; extra == "examples"
+Requires-Dist: defusedxml; extra == "examples"
+Requires-Dist: accelerate; extra == "examples"
+Requires-Dist: unstructured[embed-huggingface,pdf]<0.16.0; extra == "examples"
+Requires-Dist: pdfplumber==0.11.4; extra == "examples"
+Requires-Dist: huggingface_hub[hf_transfer]; extra == "examples"
+Requires-Dist: onnx==1.16.1; extra == "examples"
+Requires-Dist: ultralytics==8.3.29; extra == "examples"
 
 ================
 |logo| DataChain