datachain 0.3.16__py3-none-any.whl → 0.3.18__py3-none-any.whl

datachain/query/dataset.py

@@ -3,7 +3,6 @@ import inspect
 import logging
 import os
 import random
-import re
 import string
 import subprocess
 import sys
@@ -36,7 +35,6 @@ from sqlalchemy.sql.selectable import Select
 
 from datachain.asyn import ASYNC_WORKERS, AsyncMapper, OrderedMapper
 from datachain.catalog import QUERY_SCRIPT_CANCELED_EXIT_CODE, get_catalog
-from datachain.client import Client
 from datachain.data_storage.schema import (
     PARTITION_COLUMN_ID,
     partition_col_names,
@@ -46,7 +44,6 @@ from datachain.dataset import DatasetStatus, RowDict
 from datachain.error import DatasetNotFoundError, QueryScriptCancelError
 from datachain.progress import CombinedDownloadCallback
 from datachain.sql.functions import rand
-from datachain.storage import Storage, StorageURI
 from datachain.utils import (
     batched,
     determine_processes,
@@ -56,7 +53,7 @@ from datachain.utils import (
 
 from .schema import C, UDFParamSpec, normalize_param
 from .session import Session
-from .udf import UDFBase, UDFClassWrapper, UDFFactory, UDFType
+from .udf import UDFBase
 
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import ClauseElement
@@ -77,9 +74,7 @@ INSERT_BATCH_SIZE = 10000
 
 PartitionByType = Union[ColumnElement, Sequence[ColumnElement]]
 JoinPredicateType = Union[str, ColumnClause, ColumnElement]
-# dependency can be either dataset_name + dataset_version tuple or just storage uri
-# depending what type of dependency we are adding
-DatasetDependencyType = Union[tuple[str, int], StorageURI]
+DatasetDependencyType = tuple[str, int]
 
 logger = logging.getLogger("datachain")
 
@@ -185,38 +180,6 @@ class QueryStep(StartingStep):
         )
 
 
-@frozen
-class IndexingStep(StartingStep):
-    path: str
-    catalog: "Catalog"
-    kwargs: dict[str, Any]
-    recursive: Optional[bool] = True
-
-    def apply(self):
-        self.catalog.index([self.path], **self.kwargs)
-        uri, path = Client.parse_url(self.path)
-        _partial_id, partial_path = self.catalog.metastore.get_valid_partial_id(
-            uri, path
-        )
-        dataset = self.catalog.get_dataset(Storage.dataset_name(uri, partial_path))
-        dataset_rows = self.catalog.warehouse.dataset_rows(
-            dataset, dataset.latest_version
-        )
-
-        def q(*columns):
-            col_names = [c.name for c in columns]
-            return self.catalog.warehouse.nodes_dataset_query(
-                dataset_rows,
-                column_names=col_names,
-                path=path,
-                recursive=self.recursive,
-            )
-
-        storage = self.catalog.metastore.get_storage(uri)
-
-        return step_result(q, dataset_rows.c, dependencies=[storage.uri])
-
-
 def generator_then_call(generator, func: Callable):
     """
     Yield items from generator then execute a function and yield
@@ -230,7 +193,7 @@ def generator_then_call(generator, func: Callable):
 class DatasetDiffOperation(Step):
     """
     Abstract class for operations that are calculation some kind of diff between
-    datasets queries like subtract, changed etc.
+    datasets queries like subtract etc.
     """
 
     dq: "DatasetQuery"
@@ -304,28 +267,6 @@ class Subtract(DatasetDiffOperation):
         return sq.select().except_(sq.select().where(where_clause))
 
 
-@frozen
-class Changed(DatasetDiffOperation):
-    """
-    Calculates rows that are changed in a source query compared to target query
-    Changed means it has same source + path but different last_modified
-    Example:
-        >>> ds = DatasetQuery(name="dogs_cats") # some older dataset with embeddings
-        >>> ds_updated = (
-                DatasetQuery("gs://dvcx-datalakes/dogs-and-cats")
-                .filter(C.size > 1000) # we can also filter out source query
-                .changed(ds)
-                .add_signals(calc_embeddings) # calculae embeddings only on changed rows
-                .union(ds) # union with old dataset that's missing updated rows
-                .save("dogs_cats_updated")
-            )
-
-    """
-
-    def query(self, source_query: Select, target_query: Select) -> Select:
-        return self.catalog.warehouse.changed_query(source_query, target_query)
-
-
 def adjust_outputs(
     warehouse: "AbstractWarehouse", row: dict[str, Any], udf_col_types: list[tuple]
 ) -> dict[str, Any]:
@@ -423,7 +364,7 @@ def get_generated_callback(is_generator: bool = False) -> Callback:
 
 @frozen
 class UDFStep(Step, ABC):
-    udf: UDFType
+    udf: UDFBase
     catalog: "Catalog"
     partition_by: Optional[PartitionByType] = None
     parallel: Optional[int] = None
@@ -529,12 +470,6 @@ class UDFStep(Step, ABC):
 
             else:
                 # Otherwise process single-threaded (faster for smaller UDFs)
-                # Optionally instantiate the UDF instance if a class is provided.
-                if isinstance(self.udf, UDFFactory):
-                    udf: UDFBase = self.udf()
-                else:
-                    udf = self.udf
-
                 warehouse = self.catalog.warehouse
 
                 with contextlib.closing(
@@ -544,7 +479,7 @@ class UDFStep(Step, ABC):
                     processed_cb = get_processed_callback()
                     generated_cb = get_generated_callback(self.is_generator)
                     try:
-                        udf_results = udf.run(
+                        udf_results = self.udf.run(
                             udf_fields,
                             udf_inputs,
                             self.catalog,
@@ -557,7 +492,7 @@ class UDFStep(Step, ABC):
                             warehouse,
                             udf_table,
                             udf_results,
-                            udf,
+                            self.udf,
                             cb=generated_cb,
                         )
                     finally:
@@ -1096,28 +1031,14 @@ class ResultIter:
 class DatasetQuery:
     def __init__(
         self,
-        path: str = "",
-        name: str = "",
+        name: str,
         version: Optional[int] = None,
         catalog: Optional["Catalog"] = None,
-        client_config=None,
-        recursive: Optional[bool] = True,
         session: Optional[Session] = None,
-        anon: bool = False,
-        indexing_feature_schema: Optional[dict] = None,
         indexing_column_types: Optional[dict[str, Any]] = None,
-        update: Optional[bool] = False,
         in_memory: bool = False,
     ):
-        if client_config is None:
-            client_config = {}
-
-        if anon:
-            client_config["anon"] = True
-
-        self.session = Session.get(
-            session, catalog=catalog, client_config=client_config, in_memory=in_memory
-        )
+        self.session = Session.get(session, catalog=catalog, in_memory=in_memory)
         self.catalog = catalog or self.session.catalog
         self.steps: list[Step] = []
         self._chunk_index: Optional[int] = None
@@ -1131,26 +1052,14 @@ class DatasetQuery:
         self.feature_schema: Optional[dict] = None
         self.column_types: Optional[dict[str, Any]] = None
 
-        if path:
-            kwargs = {"update": True} if update else {}
-            self.starting_step = IndexingStep(path, self.catalog, kwargs, recursive)
-            self.feature_schema = indexing_feature_schema
-            self.column_types = indexing_column_types
-        elif name:
-            self.name = name
-            ds = self.catalog.get_dataset(name)
-            self.version = version or ds.latest_version
-            self.feature_schema = ds.get_version(self.version).feature_schema
-            self.column_types = copy(ds.schema)
-            if "sys__id" in self.column_types:
-                self.column_types.pop("sys__id")
-            self.starting_step = QueryStep(self.catalog, name, self.version)
-        else:
-            raise ValueError("must provide path or name")
-
-    @staticmethod
-    def is_storage_path(path):
-        return bool(re.compile(r"^[a-zA-Z0-9]+://").match(path))
+        self.name = name
+        ds = self.catalog.get_dataset(name)
+        self.version = version or ds.latest_version
+        self.feature_schema = ds.get_version(self.version).feature_schema
+        self.column_types = copy(ds.schema)
+        if "sys__id" in self.column_types:
+            self.column_types.pop("sys__id")
+        self.starting_step = QueryStep(self.catalog, name, self.version)
 
     def __iter__(self):
         return iter(self.db_results())
@@ -1556,7 +1465,7 @@ class DatasetQuery:
     @detach
     def add_signals(
         self,
-        udf: UDFType,
+        udf: UDFBase,
         parallel: Optional[int] = None,
         workers: Union[bool, int] = False,
         min_task_size: Optional[int] = None,
@@ -1577,9 +1486,6 @@ class DatasetQuery:
         at least that minimum number of rows to each distributed worker, mostly useful
         if there are a very large number of small tasks to process.
         """
-        if isinstance(udf, UDFClassWrapper):  # type: ignore[unreachable]
-            # This is a bare decorated class, "instantiate" it now.
-            udf = udf()  # type: ignore[unreachable]
         query = self.clone()
         query.steps.append(
             UDFSignal(
@@ -1595,34 +1501,21 @@ class DatasetQuery:
         return query
 
     @detach
-    def subtract(self, dq: "DatasetQuery") -> "Self":
-        return self._subtract(dq, on=[("source", "source"), ("path", "path")])
-
-    @detach
-    def _subtract(self, dq: "DatasetQuery", on: Sequence[tuple[str, str]]) -> "Self":
+    def subtract(self, dq: "DatasetQuery", on: Sequence[tuple[str, str]]) -> "Self":
         query = self.clone()
         query.steps.append(Subtract(dq, self.catalog, on=on))
         return query
 
-    @detach
-    def changed(self, dq: "DatasetQuery") -> "Self":
-        query = self.clone()
-        query.steps.append(Changed(dq, self.catalog))
-        return query
-
     @detach
     def generate(
         self,
-        udf: UDFType,
+        udf: UDFBase,
         parallel: Optional[int] = None,
         workers: Union[bool, int] = False,
         min_task_size: Optional[int] = None,
         partition_by: Optional[PartitionByType] = None,
         cache: bool = False,
     ) -> "Self":
-        if isinstance(udf, UDFClassWrapper):  # type: ignore[unreachable]
-            # This is a bare decorated class, "instantiate" it now.
-            udf = udf()  # type: ignore[unreachable]
         query = self.clone()
         steps = query.steps
         steps.append(
@@ -1640,24 +1533,13 @@ class DatasetQuery:
 
     def _add_dependencies(self, dataset: "DatasetRecord", version: int):
         for dependency in self.dependencies:
-            if isinstance(dependency, tuple):
-                # dataset dependency
-                ds_dependency_name, ds_dependency_version = dependency
-                self.catalog.metastore.add_dataset_dependency(
-                    dataset.name,
-                    version,
-                    ds_dependency_name,
-                    ds_dependency_version,
-                )
-            else:
-                # storage dependency - its name is a valid StorageURI
-                storage = self.catalog.metastore.get_storage(dependency)
-                self.catalog.metastore.add_storage_dependency(
-                    StorageURI(dataset.name),
-                    version,
-                    storage.uri,
-                    storage.timestamp_str,
-                )
+            ds_dependency_name, ds_dependency_version = dependency
+            self.catalog.metastore.add_dataset_dependency(
+                dataset.name,
+                version,
+                ds_dependency_name,
+                ds_dependency_version,
+            )
 
     def exec(self) -> "Self":
         """Execute the query."""

datachain/query/dispatch.py

@@ -27,7 +27,7 @@ from datachain.query.queue import (
     put_into_queue,
     unmarshal,
 )
-from datachain.query.udf import UDFBase, UDFFactory, UDFResult
+from datachain.query.udf import UDFBase, UDFResult
 from datachain.utils import batched_it
 
 DEFAULT_BATCH_SIZE = 10000
@@ -156,8 +156,6 @@ class UDFDispatcher:
 
     @property
     def batch_size(self):
-        if not self.udf:
-            self.udf = self.udf_factory()
         if self._batch_size is None:
             if hasattr(self.udf, "properties") and hasattr(
                 self.udf.properties, "batch"
@@ -181,18 +179,7 @@ class UDFDispatcher:
             self.catalog = Catalog(
                 id_generator, metastore, warehouse, **self.catalog_init_params
             )
-        udf = loads(self.udf_data)
-        # isinstance cannot be used here, as cloudpickle packages the entire class
-        # definition, and so these two types are not considered exactly equal,
-        # even if they have the same import path.
-        if full_module_type_path(type(udf)) != full_module_type_path(UDFFactory):
-            self.udf = udf
-        else:
-            self.udf = None
-            self.udf_factory = udf
-        if not self.udf:
-            self.udf = self.udf_factory()
-
+        self.udf = loads(self.udf_data)
         return UDFWorker(
             self.catalog,
             self.udf,

datachain/query/schema.py

@@ -9,6 +9,7 @@ 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:
@@ -19,6 +20,17 @@ if TYPE_CHECKING:
 DEFAULT_DELIMITER = "__"
 
 
+def file_signals(row, signal_name="file"):
+    # TODO this is workaround until we decide what to do with these classes
+    prefix = f"{signal_name}{DEFAULT_DELIMITER}"
+    return {
+        c_name.removeprefix(prefix): c_value
+        for c_name, c_value in row.items()
+        if c_name.startswith(prefix)
+        and DEFAULT_DELIMITER not in c_name.removeprefix(prefix)
+    }
+
+
 class ColumnMeta(type):
     @staticmethod
     def to_db_name(name: str) -> str:
@@ -86,11 +98,11 @@ class Object(UDFParameter):
         cb: Callback = DEFAULT_CALLBACK,
         **kwargs,
     ) -> Any:
-        client = catalog.get_client(row["source"])
-        uid = catalog._get_row_uid(row)
+        file = File._from_row(file_signals(row))
+        client = catalog.get_client(file.source)
         if cache:
-            client.download(uid, callback=cb)
-        with client.open_object(uid, use_cache=cache, cb=cb) as f:
+            client.download(file, callback=cb)
+        with client.open_object(file, use_cache=cache, cb=cb) as f:
             return self.reader(f)
 
     async def get_value_async(
@@ -103,12 +115,12 @@ class Object(UDFParameter):
         cb: Callback = DEFAULT_CALLBACK,
         **kwargs,
     ) -> Any:
-        client = catalog.get_client(row["source"])
-        uid = catalog._get_row_uid(row)
+        file = File._from_row(file_signals(row))
+        client = catalog.get_client(file.source)
         if cache:
-            await client._download(uid, callback=cb)
+            await client._download(file, callback=cb)
         obj = await mapper.to_thread(
-            functools.partial(client.open_object, uid, use_cache=cache, cb=cb)
+            functools.partial(client.open_object, file, use_cache=cache, cb=cb)
         )
         with obj:
             return await mapper.to_thread(self.reader, obj)
@@ -129,11 +141,11 @@ class Stream(UDFParameter):
         cb: Callback = DEFAULT_CALLBACK,
         **kwargs,
     ) -> Any:
-        client = catalog.get_client(row["source"])
-        uid = catalog._get_row_uid(row)
+        file = File._from_row(file_signals(row))
+        client = catalog.get_client(file.source)
         if cache:
-            client.download(uid, callback=cb)
-        return client.open_object(uid, use_cache=cache, cb=cb)
+            client.download(file, callback=cb)
+        return client.open_object(file, use_cache=cache, cb=cb)
 
     async def get_value_async(
         self,
@@ -145,12 +157,12 @@ class Stream(UDFParameter):
         cb: Callback = DEFAULT_CALLBACK,
         **kwargs,
     ) -> Any:
-        client = catalog.get_client(row["source"])
-        uid = catalog._get_row_uid(row)
+        file = File._from_row(file_signals(row))
+        client = catalog.get_client(file.source)
         if cache:
-            await client._download(uid, callback=cb)
+            await client._download(file, callback=cb)
         return await mapper.to_thread(
-            functools.partial(client.open_object, uid, use_cache=cache, cb=cb)
+            functools.partial(client.open_object, file, use_cache=cache, cb=cb)
         )
 
 
@@ -178,10 +190,10 @@ class LocalFilename(UDFParameter):
             # If the glob pattern is specified and the row filename
             # does not match it, then return None
             return None
-        client = catalog.get_client(row["source"])
-        uid = catalog._get_row_uid(row)
-        client.download(uid, callback=cb)
-        return client.cache.get_path(uid)
+        file = File._from_row(file_signals(row))
+        client = catalog.get_client(file.source)
+        client.download(file, callback=cb)
+        return client.cache.get_path(file)
 
     async def get_value_async(
         self,
@@ -197,10 +209,10 @@ class LocalFilename(UDFParameter):
             # If the glob pattern is specified and the row filename
             # does not match it, then return None
             return None
-        client = catalog.get_client(row["source"])
-        uid = catalog._get_row_uid(row)
-        await client._download(uid, callback=cb)
-        return client.cache.get_path(uid)
+        file = File._from_row(file_signals(row))
+        client = catalog.get_client(file.source)
+        await client._download(file, callback=cb)
+        return client.cache.get_path(file)
 
 
 UDFParamSpec = Union[str, Column, UDFParameter]

datachain/query/udf.py

@@ -1,14 +1,9 @@
 import typing
-from collections.abc import Iterable, Iterator, Mapping, Sequence
+from collections.abc import Iterable, Iterator, Sequence
 from dataclasses import dataclass
-from functools import WRAPPER_ASSIGNMENTS
-from inspect import isclass
 from typing import (
     TYPE_CHECKING,
     Any,
-    Callable,
-    Optional,
-    Union,
 )
 
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
@@ -23,11 +18,7 @@ from .batch import (
     RowsOutputBatch,
     UDFInputBatch,
 )
-from .schema import (
-    UDFParameter,
-    UDFParamSpec,
-    normalize_param,
-)
+from .schema import UDFParameter
 
 if TYPE_CHECKING:
     from datachain.catalog import Catalog
@@ -66,41 +57,6 @@ class UDFProperties:
         return self.output.keys()
 
 
-def udf(
-    params: Sequence[UDFParamSpec],
-    output: UDFOutputSpec,
-    *,
-    method: Optional[str] = None,  # only used for class-based UDFs
-    batch: int = 1,
-):
-    """
-    Decorate a function or a class to be used as a UDF.
-
-    The decorator expects both the outputs and inputs of the UDF to be specified.
-    The outputs are defined as a collection of tuples containing the signal name
-    and type.
-    Parameters are defined as a list of column objects (e.g. C.name).
-    Optionally, UDFs can be run on batches of rows to improve performance, this
-    is determined by the 'batch' parameter. When operating on batches of inputs,
-    the UDF function will be called with a single argument - a list
-    of tuples containing inputs (e.g. ((input1_a, input1_b), (input2_a, input2b))).
-    """
-    if isinstance(params, str):
-        params = (params,)
-    if not isinstance(output, Mapping):
-        raise TypeError(f"'output' must be a mapping, got {type(output).__name__}")
-
-    properties = UDFProperties([normalize_param(p) for p in params], output, batch)
-
-    def decorator(udf_base: Union[Callable, type]):
-        if isclass(udf_base):
-            return UDFClassWrapper(udf_base, properties, method=method)
-        if callable(udf_base):
-            return UDFWrapper(udf_base, properties)
-
-    return decorator
-
-
 class UDFBase:
     """A base class for implementing stateful UDFs."""
 
@@ -168,105 +124,3 @@ class UDFBase:
             for row_id, signals in zip(row_ids, results)
             if signals is not None  # skip rows with no output
         ]
-
-
-class UDFClassWrapper:
-    """
-    A wrapper for class-based (stateful) UDFs.
-    """
-
-    def __init__(
-        self,
-        udf_class: type,
-        properties: UDFProperties,
-        method: Optional[str] = None,
-    ):
-        self.udf_class = udf_class
-        self.udf_method = method
-        self.properties = properties
-        self.output = properties.output
-
-    def __call__(self, *args, **kwargs) -> "UDFFactory":
-        return UDFFactory(
-            self.udf_class,
-            args,
-            kwargs,
-            self.properties,
-            self.udf_method,
-        )
-
-
-class UDFWrapper(UDFBase):
-    """A wrapper class for function UDFs to be used in custom signal generation."""
-
-    def __init__(
-        self,
-        func: Callable,
-        properties: UDFProperties,
-    ):
-        self.func = func
-        super().__init__(properties)
-        # This emulates the behavior of functools.wraps for a class decorator
-        for attr in WRAPPER_ASSIGNMENTS:
-            if hasattr(func, attr):
-                setattr(self, attr, getattr(func, attr))
-
-    def run_once(
-        self,
-        catalog: "Catalog",
-        arg: "UDFInput",
-        is_generator: bool = False,
-        cache: bool = False,
-        cb: Callback = DEFAULT_CALLBACK,
-    ) -> Iterable[UDFResult]:
-        if isinstance(arg, UDFInputBatch):
-            udf_inputs = [
-                self.bind_parameters(catalog, row, cache=cache, cb=cb)
-                for row in arg.rows
-            ]
-            udf_outputs = self.func(udf_inputs)
-            return self._process_results(arg.rows, udf_outputs, is_generator)
-        if isinstance(arg, RowDict):
-            udf_inputs = self.bind_parameters(catalog, arg, cache=cache, cb=cb)
-            udf_outputs = self.func(*udf_inputs)
-            if not is_generator:
-                # udf_outputs is generator already if is_generator=True
-                udf_outputs = [udf_outputs]
-            return self._process_results([arg], udf_outputs, is_generator)
-        raise ValueError(f"Unexpected UDF argument: {arg}")
-
-    # This emulates the behavior of functools.wraps for a class decorator
-    def __repr__(self):
-        return repr(self.func)
-
-
-class UDFFactory:
-    """
-    A wrapper for late instantiation of UDF classes, primarily for use in parallelized
-    execution.
-    """
-
-    def __init__(
-        self,
-        udf_class: type,
-        args,
-        kwargs,
-        properties: UDFProperties,
-        method: Optional[str] = None,
-    ):
-        self.udf_class = udf_class
-        self.udf_method = method
-        self.args = args
-        self.kwargs = kwargs
-        self.properties = properties
-        self.output = properties.output
-
-    def __call__(self) -> UDFWrapper:
-        udf_func = self.udf_class(*self.args, **self.kwargs)
-        if self.udf_method:
-            udf_func = getattr(udf_func, self.udf_method)
-
-        return UDFWrapper(udf_func, self.properties)
-
-
-UDFType = Union[UDFBase, UDFFactory]

datachain/sql/types.py

@@ -12,11 +12,11 @@ for sqlite we can use `sqlite.register_converter`
 ( https://docs.python.org/3/library/sqlite3.html#sqlite3.register_converter )
 """
 
-import json
 from datetime import datetime
 from types import MappingProxyType
 from typing import Any, Union
 
+import orjson
 import sqlalchemy as sa
 from sqlalchemy import TypeDecorator, types
 
@@ -312,7 +312,7 @@ class Array(SQLType):
     def on_read_convert(self, value, dialect):
         r = read_converter(dialect).array(value, self.item_type, dialect)
         if isinstance(self.item_type, JSON):
-            r = [json.loads(item) if isinstance(item, str) else item for item in r]
+            r = [orjson.loads(item) if isinstance(item, str) else item for item in r]
         return r
 
 
@@ -420,6 +420,8 @@ class TypeReadConverter:
         return [item_type.on_read_convert(x, dialect) for x in value]
 
     def json(self, value):
+        if isinstance(value, str):
+            return orjson.loads(value)
         return value
 
     def datetime(self, value):