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

datachain/query/session.py

@@ -1,5 +1,8 @@
 import atexit
+import logging
+import os
 import re
+import sys
 from typing import TYPE_CHECKING, Optional
 from uuid import uuid4
 
@@ -9,6 +12,8 @@ from datachain.error import TableMissingError
 if TYPE_CHECKING:
     from datachain.catalog import Catalog
 
+logger = logging.getLogger("datachain")
+
 
 class Session:
     """
@@ -35,6 +40,7 @@ class Session:
 
     GLOBAL_SESSION_CTX: Optional["Session"] = None
     GLOBAL_SESSION: Optional["Session"] = None
+    ORIGINAL_EXCEPT_HOOK = None
 
     DATASET_PREFIX = "session_"
     GLOBAL_SESSION_NAME = "global"
@@ -58,6 +64,7 @@ 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
@@ -67,6 +74,9 @@ class Session:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type:
+            self._cleanup_created_versions(self.name)
+
         self._cleanup_temp_datasets()
         if self.is_new_catalog:
             self.catalog.metastore.close_on_exit()
@@ -88,6 +98,21 @@ 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:
+            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
+            )
+            self.catalog.remove_dataset_version(dataset, version)
+
     @classmethod
     def get(
         cls,
@@ -114,9 +139,23 @@ class Session:
                 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
 
+    @staticmethod
+    def except_hook(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)
+
     @classmethod
     def cleanup_for_tests(cls):
         if cls.GLOBAL_SESSION_CTX is not None:
@@ -125,6 +164,9 @@ class Session:
             cls.GLOBAL_SESSION_CTX = None
             atexit.unregister(cls._global_cleanup)
 
+        if cls.ORIGINAL_EXCEPT_HOOK:
+            sys.excepthook = cls.ORIGINAL_EXCEPT_HOOK
+
     @staticmethod
     def _global_cleanup():
         if Session.GLOBAL_SESSION_CTX is not None:

datachain/sql/functions/string.py

@@ -37,6 +37,18 @@ class regexp_replace(GenericFunction):  # noqa: N801
     inherit_cache = True
 
 
+class replace(GenericFunction):  # noqa: N801
+    """
+    Replaces substring with another string.
+    """
+
+    type = String()
+    package = "string"
+    name = "replace"
+    inherit_cache = True
+
+
 compiler_not_implemented(length)
 compiler_not_implemented(split)
 compiler_not_implemented(regexp_replace)
+compiler_not_implemented(replace)

datachain/sql/sqlite/base.py

@@ -78,7 +78,8 @@ def setup():
     compiles(array.length, "sqlite")(compile_array_length)
     compiles(string.length, "sqlite")(compile_string_length)
     compiles(string.split, "sqlite")(compile_string_split)
-    compiles(string.regexp_replace, "sqlite")(compile_regexp_replace)
+    compiles(string.regexp_replace, "sqlite")(compile_string_regexp_replace)
+    compiles(string.replace, "sqlite")(compile_string_replace)
     compiles(conditional.greatest, "sqlite")(compile_greatest)
     compiles(conditional.least, "sqlite")(compile_least)
     compiles(Values, "sqlite")(compile_values)
@@ -273,10 +274,6 @@ def path_file_ext(path):
     return func.substr(path, func.length(path) - path_file_ext_length(path) + 1)
 
 
-def compile_regexp_replace(element, compiler, **kwargs):
-    return f"regexp_replace({compiler.process(element.clauses, **kwargs)})"
-
-
 def compile_path_parent(element, compiler, **kwargs):
     return compiler.process(path_parent(*element.clauses.clauses), **kwargs)
 
@@ -331,6 +328,14 @@ def compile_string_split(element, compiler, **kwargs):
     return compiler.process(func.split(*element.clauses.clauses), **kwargs)
 
 
+def compile_string_regexp_replace(element, compiler, **kwargs):
+    return f"regexp_replace({compiler.process(element.clauses, **kwargs)})"
+
+
+def compile_string_replace(element, compiler, **kwargs):
+    return compiler.process(func.replace(*element.clauses.clauses), **kwargs)
+
+
 def compile_greatest(element, compiler, **kwargs):
     """
     Compiles a sql function for `greatest(*args)` taking 1 or more args

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

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

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

@@ -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=FuKuIiCwPgN5Ea25hnFe_ZFZH9YEUZ2ma9k_Lczk-JU,63867
+datachain/catalog/catalog.py,sha256=BsMyk2RQibQYHgrmovFZeSEpPVMTwgb_7ntVYdc7t-E,64090
 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
@@ -33,17 +33,17 @@ datachain/data_storage/__init__.py,sha256=cEOJpyu1JDZtfUupYucCDNFI6e5Wmp_Oyzq6rZ
 datachain/data_storage/db_engine.py,sha256=81Ol1of9TTTzD97ORajCnP366Xz2mEJt6C-kTUCaru4,3406
 datachain/data_storage/id_generator.py,sha256=lCEoU0BM37Ai2aRpSbwo5oQT0GqZnSpYwwvizathRMQ,4292
 datachain/data_storage/job.py,sha256=w-7spowjkOa1P5fUVtJou3OltT0L48P0RYWZ9rSJ9-s,383
-datachain/data_storage/metastore.py,sha256=NV4FJ_W16Q19Sx70i5Qtre-n4DC2kMD0qw0vBz3j7Ks,52228
+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=EBKJncuzcyQfcKFm2mUjvHjHRTODsteM-k_zndunBrw,28834
+datachain/data_storage/sqlite.py,sha256=fW08P7AbJ0cDbTbcTKuAGpvMXvBjg-QkGsKT_Dslyws,28383
 datachain/data_storage/warehouse.py,sha256=fXhVfao3NfWFGbbG5uJ-Ga4bX1FiKVfcbDyQgECYfk8,32122
 datachain/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datachain/lib/arrow.py,sha256=aUsoQmxDmuSnB8Ik9p57Y66gc_dgx6NBqkDDIfLsvno,7630
+datachain/lib/arrow.py,sha256=0R2CYsN82nNa5_03iS6jVix9EKeeqNZNAMgpSQP2hfo,9482
 datachain/lib/clip.py,sha256=lm5CzVi4Cj1jVLEKvERKArb-egb9j1Ls-fwTItT6vlI,6150
-datachain/lib/data_model.py,sha256=gHIjlow84GMRDa78yLL1Ud-N18or21fnTyPEwsatpXY,2045
+datachain/lib/data_model.py,sha256=ECTbvlnzM98hp2mZ4fo82Yi0-MuoqTIQasQKGIyd89I,2040
 datachain/lib/dataset_info.py,sha256=srPPhI2UHf6hFPBecyFEVw2SS5aPisIIMsvGgKqi7ss,2366
-datachain/lib/dc.py,sha256=yTyHrKIswCzdlvl2n-wdEVZEEF5VQpkLJPzPfUL9CTU,72054
+datachain/lib/dc.py,sha256=XmAFU9k79wUHIh0gYab8j-wF4vIlyW6opJcOy8fmoVc,76666
 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,11 +53,11 @@ 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=iqgubjCBRiUJB30miv05qFX4uU04dA_Pzi3DCUsHZGs,24177
+datachain/lib/signal_schema.py,sha256=gj45dRQuOsKDmaKaJxb5j63HYVGw-Ks1fyAS1FpyOWA,24145
 datachain/lib/tar.py,sha256=3WIzao6yD5fbLqXLTt9GhPGNonbFIs_fDRu-9vgLgsA,1038
 datachain/lib/text.py,sha256=UNHm8fhidk7wdrWqacEWaA6I9ykfYqarQ2URby7jc7M,1261
-datachain/lib/udf.py,sha256=nG7DDuPgZ5ZuijwvDoCq-OZMxlDM8vFNzyxMmik0Y1c,11716
-datachain/lib/udf_signature.py,sha256=gMStcEeYJka5M6cg50Z9orC6y6HzCAJ3MkFqqn1fjZg,7137
+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/vfile.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/lib/webdataset.py,sha256=o7SHk5HOUWsZ5Ln04xOM04eQqiBHiJNO7xLgyVBrwo8,6924
@@ -67,17 +67,16 @@ datachain/lib/convert/flatten.py,sha256=Uebc5CeqCsacp-nr6IG9i6OGuUavXqdqnoGctZBk
 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/unflatten.py,sha256=Ogvh_5wg2f38_At_1lN0D_e2uZOOpYEvwvB2xdq56Tw,2012
-datachain/lib/convert/values_to_tuples.py,sha256=YOdbjzHq-uj6-cV2Qq43G72eN2avMNDGl4x5t6yQMl8,3931
+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=-vlpINJiertlnaoUVv1C95RatU0F6zuhpIYRufJRo1M,3660
-datachain/query/dataset.py,sha256=tLCTaj4K93BY93GgOPv9PknZByEF89zpHc7y9s8ZF_w,53610
-datachain/query/dispatch.py,sha256=CFAc09O6UllcyUSSEY1GUlEMPzeO8RYhXinNN4HBl9M,12405
+datachain/query/batch.py,sha256=3QlwshhpUc1amZRtXWVXEEuq47hEQgQlY0Ji48DR6hg,3508
+datachain/query/dataset.py,sha256=MF_E7yjbFQV6NcP4gKbJFXiWuoQkpQ7-Jmxa59FxenE,53630
+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=UPH5Z4fzCDsvj81ji0e8GA6Mgra3bOAEpVq4htqOtis,4317
-datachain/query/udf.py,sha256=HB2hbEuiGA4ch9P2mh9iLA5Jj9mRj-4JFy9VfjTLJ8U,3622
+datachain/query/session.py,sha256=kpFFJMfWBnxaMPojMGhJRbk-BOsSYI8Ckl6vvqnx7d0,5787
 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
@@ -91,15 +90,15 @@ datachain/sql/functions/array.py,sha256=EB7nJSncUc1PuxlHyzU2gVhF8DuXaxpGlxb5e8X2
 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=NSQIpmtQgm68hz3TFJsgHMBuo4MjBNhDSyEIC3pWkT8,916
+datachain/sql/functions/string.py,sha256=DYgiw8XSk7ge7GXvyRI1zbaMruIizNeI-puOjriQGZQ,1148
 datachain/sql/sqlite/__init__.py,sha256=TAdJX0Bg28XdqPO-QwUVKy8rg78cgMileHvMNot7d04,166
-datachain/sql/sqlite/base.py,sha256=WLPHBhZbXbiqPoRV1VgDrXJqku4UuvJpBhYeQ0k5rI8,13364
+datachain/sql/sqlite/base.py,sha256=3gDMLKSWkxnbiZ1dykYa5VuHSSlg5sLY9ihMqcH_o1M,13578
 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.5.0.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.5.0.dist-info/METADATA,sha256=tKSZNiHZY0WJ_w6irkpSF7qDfuOTfiYNEQ6St3eBs-M,17156
-datachain-0.5.0.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-datachain-0.5.0.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.5.0.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.5.0.dist-info/RECORD,,
+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/query/udf.py

@@ -1,126 +0,0 @@
-import typing
-from collections.abc import Iterable, Iterator, Sequence
-from dataclasses import dataclass
-from typing import (
-    TYPE_CHECKING,
-    Any,
-)
-
-from fsspec.callbacks import DEFAULT_CALLBACK, Callback
-
-from datachain.dataset import RowDict
-
-from .batch import (
-    Batch,
-    BatchingStrategy,
-    NoBatching,
-    Partition,
-    RowsOutputBatch,
-    UDFInputBatch,
-)
-from .schema import UDFParameter
-
-if TYPE_CHECKING:
-    from datachain.catalog import Catalog
-
-    from .batch import RowsOutput, UDFInput
-
-ColumnType = Any
-
-
-# Specification for the output of a UDF
-UDFOutputSpec = typing.Mapping[str, ColumnType]
-
-# Result type when calling the UDF wrapper around the actual
-# Python function / class implementing it.
-UDFResult = dict[str, Any]
-
-
-@dataclass
-class UDFProperties:
-    """Container for basic UDF properties."""
-
-    params: list[UDFParameter]
-    output: UDFOutputSpec
-    batch: int = 1
-
-    def get_batching(self, use_partitioning: bool = False) -> BatchingStrategy:
-        if use_partitioning:
-            return Partition()
-        if self.batch == 1:
-            return NoBatching()
-        if self.batch > 1:
-            return Batch(self.batch)
-        raise ValueError(f"invalid batch size {self.batch}")
-
-    def signal_names(self) -> Iterable[str]:
-        return self.output.keys()
-
-
-class UDFBase:
-    """A base class for implementing stateful UDFs."""
-
-    def __init__(
-        self,
-        properties: UDFProperties,
-    ):
-        self.properties = properties
-        self.signal_names = properties.signal_names()
-        self.output = properties.output
-
-    def run(
-        self,
-        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"]]:
-        for batch in udf_inputs:
-            if isinstance(batch, RowsOutputBatch):
-                n_rows = len(batch.rows)
-                inputs: UDFInput = UDFInputBatch(
-                    [RowDict(zip(udf_fields, row)) for row in batch.rows]
-                )
-            else:
-                n_rows = 1
-                inputs = RowDict(zip(udf_fields, batch))
-            output = self.run_once(catalog, inputs, is_generator, cache, cb=download_cb)
-            processed_cb.relative_update(n_rows)
-            yield output
-
-    def run_once(
-        self,
-        catalog: "Catalog",
-        arg: "UDFInput",
-        is_generator: bool = False,
-        cache: bool = False,
-        cb: Callback = DEFAULT_CALLBACK,
-    ) -> Iterable[UDFResult]:
-        raise NotImplementedError
-
-    def bind_parameters(self, catalog: "Catalog", row: "RowDict", **kwargs) -> list:
-        return [p.get_value(catalog, row, **kwargs) for p in self.properties.params]
-
-    def _process_results(
-        self,
-        rows: Sequence["RowDict"],
-        results: Sequence[Sequence[Any]],
-        is_generator=False,
-    ) -> Iterable[UDFResult]:
-        """Create a list of dictionaries representing UDF results."""
-
-        # outputting rows
-        if is_generator:
-            # each row in results is a tuple of column values
-            return (dict(zip(self.signal_names, row)) for row in results)
-
-        # outputting signals
-        row_ids = [row["sys__id"] for row in rows]
-        return [
-            {"sys__id": row_id} | dict(zip(self.signal_names, signals))
-            for row_id, signals in zip(row_ids, results)
-            if signals is not None  # skip rows with no output
-        ]