datachain 0.3.0__tar.gz → 0.3.1__tar.gz

{datachain-0.3.0 → datachain-0.3.1}/.github/workflows/tests.yml

@@ -8,6 +8,7 @@ on:
 
 env:
   FORCE_COLOR: "1"
+  BRANCH: ${{ github.head_ref || github.ref_name }}
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
@@ -144,12 +145,30 @@ jobs:
           - 6379:6379
     steps:
 
+      - name: Studio branch name
+        env:
+          BRANCH: ${{ env.BRANCH }}
+          STUDIO_READ_ACCESS_TOKEN: ${{ secrets.ITERATIVE_STUDIO_READ_ACCESS_TOKEN }}
+        run: |
+          echo "DataChain branch: $BRANCH"
+          if [[ "$BRANCH" == "main" ]]
+          then
+              STUDIO_BRANCH=develop
+          elif git ls-remote --heads https://"$STUDIO_READ_ACCESS_TOKEN"@github.com/iterative/studio.git "$BRANCH" | grep -F "$BRANCH" 2>&1>/dev/null
+          then
+              STUDIO_BRANCH="$BRANCH"
+          else
+              STUDIO_BRANCH=develop
+          fi
+          echo "STUDIO_BRANCH=$STUDIO_BRANCH" >> $GITHUB_ENV
+          echo "Studio branch: $STUDIO_BRANCH"
+
       - name: Check out Studio
         uses: actions/checkout@v4
         with:
           fetch-depth: 0
           repository: iterative/studio
-          ref: develop
+          ref: ${{ env.STUDIO_BRANCH }}
           token: ${{ secrets.ITERATIVE_STUDIO_READ_ACCESS_TOKEN }}
 
       - name: Check out repository

{datachain-0.3.0/src/datachain.egg-info → datachain-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.3.0
+Version: 0.3.1
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0

{datachain-0.3.0 → datachain-0.3.1}/src/datachain/catalog/catalog.py

@@ -676,7 +676,7 @@ class Catalog:
 
     def parse_url(self, uri: str, **config: Any) -> tuple[Client, str]:
         config = config or self.client_config
-        return Client.parse_url(uri, self.metastore, self.cache, **config)
+        return Client.parse_url(uri, self.cache, **config)
 
     def get_client(self, uri: StorageURI, **config: Any) -> Client:
         """

{datachain-0.3.0 → datachain-0.3.1}/src/datachain/client/fsspec.py

@@ -37,7 +37,6 @@ from datachain.storage import StorageURI
 if TYPE_CHECKING:
     from fsspec.spec import AbstractFileSystem
 
-    from datachain.data_storage import AbstractMetastore
 
 logger = logging.getLogger("datachain")
 
@@ -116,13 +115,12 @@ class Client(ABC):
     @staticmethod
     def parse_url(
         source: str,
-        metastore: "AbstractMetastore",
         cache: DataChainCache,
         **kwargs,
     ) -> tuple["Client", str]:
         cls = Client.get_implementation(source)
         storage_url, rel_path = cls.split_url(source)
-        client = cls.from_name(storage_url, metastore, cache, kwargs)
+        client = cls.from_name(storage_url, cache, kwargs)
         return client, rel_path
 
     @classmethod
@@ -136,7 +134,6 @@ class Client(ABC):
     def from_name(
         cls,
         name: str,
-        metastore: "AbstractMetastore",
         cache: DataChainCache,
         kwargs: dict[str, Any],
     ) -> "Client":

{datachain-0.3.0 → datachain-0.3.1}/src/datachain/client/local.py

@@ -2,7 +2,7 @@ import os
 import posixpath
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import Any
 from urllib.parse import urlparse
 
 from fsspec.implementations.local import LocalFileSystem
@@ -12,9 +12,6 @@ from datachain.storage import StorageURI
 
 from .fsspec import Client
 
-if TYPE_CHECKING:
-    from datachain.data_storage import AbstractMetastore
-
 
 class FileClient(Client):
     FS_CLASS = LocalFileSystem
@@ -97,9 +94,7 @@ class FileClient(Client):
         return cls.root_dir(), uri.removeprefix(cls.root_path().as_uri())
 
     @classmethod
-    def from_name(
-        cls, name: str, metastore: "AbstractMetastore", cache, kwargs
-    ) -> "FileClient":
+    def from_name(cls, name: str, cache, kwargs) -> "FileClient":
         use_symlinks = kwargs.pop("use_symlinks", False)
         return cls(name, kwargs, cache, use_symlinks=use_symlinks)
 

{datachain-0.3.0 → datachain-0.3.1}/src/datachain/data_storage/warehouse.py

@@ -17,7 +17,7 @@ from sqlalchemy.sql.expression import true
 
 from datachain.client import Client
 from datachain.data_storage.serializer import Serializable
-from datachain.dataset import DatasetRecord, RowDict
+from datachain.dataset import DatasetRecord
 from datachain.node import DirType, DirTypeGroup, Entry, Node, NodeWithPath, get_path
 from datachain.sql.functions import path as pathfunc
 from datachain.sql.types import Int, SQLType
@@ -201,23 +201,17 @@ class AbstractWarehouse(ABC, Serializable):
     def dataset_select_paginated(
         self,
         query,
-        limit: Optional[int] = None,
-        order_by: tuple["ColumnElement[Any]", ...] = (),
         page_size: int = SELECT_BATCH_SIZE,
-    ) -> Generator[RowDict, None, None]:
+    ) -> Generator[Sequence, None, None]:
         """
         This is equivalent to `db.execute`, but for selecting rows in batches
         """
-        cols = query.selected_columns
-        cols_names = [c.name for c in cols]
+        limit = query._limit
+        paginated_query = query.limit(page_size)
 
-        if not order_by:
-            ordering = [cols.sys__id]
-        else:
-            ordering = order_by  # type: ignore[assignment]
-
-        # reset query order by and apply new order by id
-        paginated_query = query.order_by(None).order_by(*ordering).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
@@ -236,7 +230,7 @@ class AbstractWarehouse(ABC, Serializable):
                 processed = False
                 for row in results:
                     processed = True
-                    yield RowDict(zip(cols_names, row))
+                    yield row
                     num_yielded += 1
 
                 if not processed:

{datachain-0.3.0 → datachain-0.3.1}/src/datachain/lib/dc.py

@@ -1623,7 +1623,7 @@ class DataChain(DatasetQuery):
 
             Using glob to match patterns
             ```py
-            dc.filter(C("file.name").glob("*.jpg))
+            dc.filter(C("file.name").glob("*.jpg"))
             ```
 
             Using `datachain.sql.functions`

{datachain-0.3.0 → datachain-0.3.1}/src/datachain/lib/udf.py

@@ -1,6 +1,5 @@
 import sys
 import traceback
-from collections.abc import Iterable, Iterator
 from typing import TYPE_CHECKING, Callable, Optional
 
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
@@ -14,16 +13,19 @@ from datachain.lib.model_store import ModelStore
 from datachain.lib.signal_schema import SignalSchema
 from datachain.lib.udf_signature import UdfSignature
 from datachain.lib.utils import AbstractUDF, DataChainError, DataChainParamsError
-from datachain.query.batch import RowBatch
+from datachain.query.batch import UDFInputBatch
 from datachain.query.schema import ColumnParameter
 from datachain.query.udf import UDFBase as _UDFBase
-from datachain.query.udf import UDFProperties, UDFResult
+from datachain.query.udf import UDFProperties
 
 if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator, Sequence
+
     from typing_extensions import Self
 
     from datachain.catalog import Catalog
-    from datachain.query.batch import BatchingResult
+    from datachain.query.batch import RowsOutput, UDFInput
+    from datachain.query.udf import UDFResult
 
 
 class UdfError(DataChainParamsError):
@@ -42,22 +44,27 @@ class UDFAdapter(_UDFBase):
 
     def run(
         self,
-        udf_inputs: "Iterable[BatchingResult]",
+        udf_fields: "Sequence[str]",
+        udf_inputs: "Iterable[RowsOutput]",
         catalog: "Catalog",
         is_generator: bool,
         cache: bool,
         download_cb: Callback = DEFAULT_CALLBACK,
         processed_cb: Callback = DEFAULT_CALLBACK,
-    ) -> Iterator[Iterable["UDFResult"]]:
+    ) -> "Iterator[Iterable[UDFResult]]":
         self.inner._catalog = catalog
         if hasattr(self.inner, "setup") and callable(self.inner.setup):
             self.inner.setup()
 
-        for batch in udf_inputs:
-            n_rows = len(batch.rows) if isinstance(batch, RowBatch) else 1
-            output = self.run_once(catalog, batch, is_generator, cache, cb=download_cb)
-            processed_cb.relative_update(n_rows)
-            yield output
+        yield from super().run(
+            udf_fields,
+            udf_inputs,
+            catalog,
+            is_generator,
+            cache,
+            download_cb,
+            processed_cb,
+        )
 
         if hasattr(self.inner, "teardown") and callable(self.inner.teardown):
             self.inner.teardown()
@@ -65,12 +72,12 @@ class UDFAdapter(_UDFBase):
     def run_once(
         self,
         catalog: "Catalog",
-        arg: "BatchingResult",
+        arg: "UDFInput",
         is_generator: bool = False,
         cache: bool = False,
         cb: Callback = DEFAULT_CALLBACK,
-    ) -> Iterable[UDFResult]:
-        if isinstance(arg, RowBatch):
+    ) -> "Iterable[UDFResult]":
+        if isinstance(arg, UDFInputBatch):
             udf_inputs = [
                 self.bind_parameters(catalog, row, cache=cache, cb=cb)
                 for row in arg.rows

{datachain-0.3.0 → datachain-0.3.1}/src/datachain/query/batch.py

@@ -5,21 +5,29 @@ from collections.abc import Generator, Sequence
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Callable, Optional, Union
 
-import sqlalchemy as sa
-
 from datachain.data_storage.schema import PARTITION_COLUMN_ID
 from datachain.data_storage.warehouse import SELECT_BATCH_SIZE
 
 if TYPE_CHECKING:
+    from sqlalchemy import Select
+
     from datachain.dataset import RowDict
 
 
 @dataclass
-class RowBatch:
+class RowsOutputBatch:
+    rows: Sequence[Sequence]
+
+
+RowsOutput = Union[Sequence, RowsOutputBatch]
+
+
+@dataclass
+class UDFInputBatch:
     rows: Sequence["RowDict"]
 
 
-BatchingResult = Union["RowDict", RowBatch]
+UDFInput = Union["RowDict", UDFInputBatch]
 
 
 class BatchingStrategy(ABC):
@@ -28,9 +36,9 @@ class BatchingStrategy(ABC):
     @abstractmethod
     def __call__(
         self,
-        execute: Callable,
-        query: sa.sql.selectable.Select,
-    ) -> Generator[BatchingResult, None, None]:
+        execute: Callable[..., Generator[Sequence, None, None]],
+        query: "Select",
+    ) -> Generator[RowsOutput, None, None]:
         """Apply the provided parameters to the UDF."""
 
 
@@ -42,10 +50,10 @@ class NoBatching(BatchingStrategy):
 
     def __call__(
         self,
-        execute: Callable,
-        query: sa.sql.selectable.Select,
-    ) -> Generator["RowDict", None, None]:
-        return execute(query, limit=query._limit, order_by=query._order_by_clauses)
+        execute: Callable[..., Generator[Sequence, None, None]],
+        query: "Select",
+    ) -> Generator[Sequence, None, None]:
+        return execute(query)
 
 
 class Batch(BatchingStrategy):
@@ -59,31 +67,24 @@ class Batch(BatchingStrategy):
 
     def __call__(
         self,
-        execute: Callable,
-        query: sa.sql.selectable.Select,
-    ) -> Generator[RowBatch, None, None]:
+        execute: Callable[..., Generator[Sequence, None, None]],
+        query: "Select",
+    ) -> Generator[RowsOutputBatch, None, None]:
         # choose page size that is a multiple of the batch size
         page_size = math.ceil(SELECT_BATCH_SIZE / self.count) * self.count
 
         # select rows in batches
-        results: list[RowDict] = []
-
-        with contextlib.closing(
-            execute(
-                query,
-                page_size=page_size,
-                limit=query._limit,
-                order_by=query._order_by_clauses,
-            )
-        ) as rows:
+        results: list[Sequence] = []
+
+        with contextlib.closing(execute(query, page_size=page_size)) as rows:
             for row in rows:
                 results.append(row)
                 if len(results) >= self.count:
                     batch, results = results[: self.count], results[self.count :]
-                    yield RowBatch(batch)
+                    yield RowsOutputBatch(batch)
 
             if len(results) > 0:
-                yield RowBatch(results)
+                yield RowsOutputBatch(results)
 
 
 class Partition(BatchingStrategy):
@@ -95,27 +96,30 @@ class Partition(BatchingStrategy):
 
     def __call__(
         self,
-        execute: Callable,
-        query: sa.sql.selectable.Select,
-    ) -> Generator[RowBatch, None, None]:
+        execute: Callable[..., Generator[Sequence, None, None]],
+        query: "Select",
+    ) -> Generator[RowsOutputBatch, None, None]:
         current_partition: Optional[int] = None
-        batch: list[RowDict] = []
-
-        with contextlib.closing(
-            execute(
-                query,
-                order_by=(PARTITION_COLUMN_ID, "sys__id", *query._order_by_clauses),
-                limit=query._limit,
-            )
-        ) as rows:
+        batch: list[Sequence] = []
+
+        query_fields = [str(c.name) for c in query.selected_columns]
+        partition_column_idx = query_fields.index(PARTITION_COLUMN_ID)
+
+        ordered_query = query.order_by(None).order_by(
+            PARTITION_COLUMN_ID,
+            "sys__id",
+            *query._order_by_clauses,
+        )
+
+        with contextlib.closing(execute(ordered_query)) as rows:
             for row in rows:
-                partition = row[PARTITION_COLUMN_ID]
+                partition = row[partition_column_idx]
                 if current_partition != partition:
                     current_partition = partition
                     if len(batch) > 0:
-                        yield RowBatch(batch)
+                        yield RowsOutputBatch(batch)
                         batch = []
                 batch.append(row)
 
             if len(batch) > 0:
-                yield RowBatch(batch)
+                yield RowsOutputBatch(batch)

{datachain-0.3.0 → datachain-0.3.1}/src/datachain/query/dataset.py

@@ -461,6 +461,8 @@ class UDFStep(Step, ABC):
 
         processes = determine_processes(self.parallel)
 
+        udf_fields = [str(c.name) for c in query.selected_columns]
+
         try:
             if workers:
                 from datachain.catalog.loader import get_distributed_class
@@ -473,6 +475,7 @@ class UDFStep(Step, ABC):
                     query,
                     workers,
                     processes,
+                    udf_fields=udf_fields,
                     is_generator=self.is_generator,
                     use_partitioning=use_partitioning,
                     cache=self.cache,
@@ -489,6 +492,7 @@ class UDFStep(Step, ABC):
                     "warehouse_clone_params": self.catalog.warehouse.clone_params(),
                     "table": udf_table,
                     "query": query,
+                    "udf_fields": udf_fields,
                     "batching": batching,
                     "processes": processes,
                     "is_generator": self.is_generator,
@@ -528,6 +532,7 @@ class UDFStep(Step, ABC):
                     generated_cb = get_generated_callback(self.is_generator)
                     try:
                         udf_results = udf.run(
+                            udf_fields,
                             udf_inputs,
                             self.catalog,
                             self.is_generator,
@@ -1244,21 +1249,23 @@ class DatasetQuery:
         actual_params = [normalize_param(p) for p in params]
         try:
             query = self.apply_steps().select()
+            query_fields = [str(c.name) for c in query.selected_columns]
 
-            def row_iter() -> Generator[RowDict, None, None]:
+            def row_iter() -> Generator[Sequence, None, None]:
                 # warehouse isn't threadsafe, we need to clone() it
                 # in the thread that uses the results
                 with self.catalog.warehouse.clone() as warehouse:
-                    gen = warehouse.dataset_select_paginated(
-                        query, limit=query._limit, order_by=query._order_by_clauses
-                    )
+                    gen = warehouse.dataset_select_paginated(query)
                     with contextlib.closing(gen) as rows:
                         yield from rows
 
-            async def get_params(row: RowDict) -> tuple:
+            async def get_params(row: Sequence) -> tuple:
+                row_dict = RowDict(zip(query_fields, row))
                 return tuple(
                     [
-                        await p.get_value_async(self.catalog, row, mapper, **kwargs)
+                        await p.get_value_async(
+                            self.catalog, row_dict, mapper, **kwargs
+                        )
                         for p in actual_params
                     ]
                 )