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

datachain/__init__.py

@@ -1,21 +1,23 @@
 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 (
+    ArrowRow,
     File,
     FileError,
     ImageFile,
-    IndexedFile,
     TarVFile,
     TextFile,
 )
 from datachain.lib.model_store import ModelStore
 from datachain.lib.udf import Aggregator, Generator, Mapper
 from datachain.lib.utils import AbstractUDF, DataChainError
+from datachain.query import metrics, param
 from datachain.query.session import Session
 
 __all__ = [
     "AbstractUDF",
     "Aggregator",
+    "ArrowRow",
     "C",
     "Column",
     "DataChain",
@@ -26,7 +28,6 @@ __all__ = [
     "FileError",
     "Generator",
     "ImageFile",
-    "IndexedFile",
     "Mapper",
     "ModelStore",
     "Session",
@@ -34,4 +35,6 @@ __all__ = [
     "TarVFile",
     "TextFile",
     "is_chain_type",
+    "metrics",
+    "param",
 ]

datachain/catalog/catalog.py

@@ -1,4 +1,3 @@
-import ast
 import glob
 import io
 import json
@@ -53,9 +52,9 @@ from datachain.error import (
     DataChainError,
     DatasetInvalidVersionError,
     DatasetNotFoundError,
+    DatasetVersionNotFoundError,
     PendingIndexingError,
     QueryScriptCancelError,
-    QueryScriptCompileError,
     QueryScriptRunError,
 )
 from datachain.listing import Listing
@@ -588,37 +587,6 @@ class Catalog:
     def generate_query_dataset_name(cls) -> str:
         return f"{QUERY_DATASET_PREFIX}_{uuid4().hex}"
 
-    def attach_query_wrapper(self, code_ast):
-        if code_ast.body:
-            last_expr = code_ast.body[-1]
-            if isinstance(last_expr, ast.Expr):
-                new_expressions = [
-                    ast.Import(
-                        names=[ast.alias(name="datachain.query.dataset", asname=None)]
-                    ),
-                    ast.Expr(
-                        value=ast.Call(
-                            func=ast.Attribute(
-                                value=ast.Attribute(
-                                    value=ast.Attribute(
-                                        value=ast.Name(id="datachain", ctx=ast.Load()),
-                                        attr="query",
-                                        ctx=ast.Load(),
-                                    ),
-                                    attr="dataset",
-                                    ctx=ast.Load(),
-                                ),
-                                attr="query_wrapper",
-                                ctx=ast.Load(),
-                            ),
-                            args=[last_expr],
-                            keywords=[],
-                        )
-                    ),
-                ]
-                code_ast.body[-1:] = new_expressions
-        return code_ast
-
     def get_client(self, uri: str, **config: Any) -> Client:
         """
         Return the client corresponding to the given source `uri`.
@@ -1218,7 +1186,9 @@ class Catalog:
 
         dataset_version = dataset.get_version(version)
         if not dataset_version:
-            raise ValueError(f"Dataset {dataset.name} does not have version {version}")
+            raise DatasetVersionNotFoundError(
+                f"Dataset {dataset.name} does not have version {version}"
+            )
 
         if not dataset_version.is_final_status():
             raise ValueError("Cannot register dataset version in non final status")
@@ -1581,7 +1551,7 @@ class Catalog:
 
         try:
             remote_dataset_version = remote_dataset.get_version(version)
-        except (ValueError, StopIteration) as exc:
+        except (DatasetVersionNotFoundError, StopIteration) as exc:
             raise DataChainError(
                 f"Dataset {remote_dataset_name} doesn't have version {version}"
                 " on server"
@@ -1722,64 +1692,24 @@ class Catalog:
         query_script: str,
         env: Optional[Mapping[str, str]] = None,
         python_executable: str = sys.executable,
-        save: bool = False,
-        capture_output: bool = True,
+        capture_output: bool = False,
         output_hook: Callable[[str], None] = noop,
         params: Optional[dict[str, str]] = None,
         job_id: Optional[str] = None,
-        _execute_last_expression: bool = False,
     ) -> None:
-        """
-        Method to run custom user Python script to run a query and, as result,
-        creates new dataset from the results of a query.
-        Returns tuple of result dataset and script output.
-
-        Constraints on query script:
-            1. datachain.query.DatasetQuery should be used in order to create query
-            for a dataset
-            2. There should not be any .save() call on DatasetQuery since the idea
-            is to create only one dataset as the outcome of the script
-            3. Last statement must be an instance of DatasetQuery
-
-        If save is set to True, we are creating new dataset with results
-        from dataset query. If it's set to False, we will just print results
-        without saving anything
-
-        Example of query script:
-            from datachain.query import DatasetQuery, C
-            DatasetQuery('s3://ldb-public/remote/datasets/mnist-tiny/').filter(
-                C.size > 1000
-            )
-        """
-        if _execute_last_expression:
-            try:
-                code_ast = ast.parse(query_script)
-                code_ast = self.attach_query_wrapper(code_ast)
-                query_script_compiled = ast.unparse(code_ast)
-            except Exception as exc:
-                raise QueryScriptCompileError(
-                    f"Query script failed to compile, reason: {exc}"
-                ) from exc
-        else:
-            query_script_compiled = query_script
-            assert not save
-
+        cmd = [python_executable, "-c", query_script]
         env = dict(env or os.environ)
         env.update(
             {
                 "DATACHAIN_QUERY_PARAMS": json.dumps(params or {}),
-                "PYTHONPATH": os.getcwd(),  # For local imports
-                "DATACHAIN_QUERY_SAVE": "1" if save else "",
-                "PYTHONUNBUFFERED": "1",
                 "DATACHAIN_JOB_ID": job_id or "",
             },
         )
-        popen_kwargs = {}
+        popen_kwargs: dict[str, Any] = {}
         if capture_output:
             popen_kwargs = {"stdout": subprocess.PIPE, "stderr": subprocess.STDOUT}
 
-        cmd = [python_executable, "-c", query_script_compiled]
-        with subprocess.Popen(cmd, env=env, **popen_kwargs) as proc:  # type: ignore[call-overload]  # noqa: S603
+        with subprocess.Popen(cmd, env=env, **popen_kwargs) as proc:  # noqa: S603
             if capture_output:
                 args = (proc.stdout, output_hook)
                 thread = Thread(target=_process_stream, args=args, daemon=True)

datachain/cli.py

@@ -804,7 +804,6 @@ def query(
         catalog.query(
             script_content,
             python_executable=python_executable,
-            capture_output=False,
             params=params,
             job_id=job_id,
         )

datachain/dataset.py

@@ -12,6 +12,7 @@ from typing import (
 from urllib.parse import urlparse
 
 from datachain.client import Client
+from datachain.error import DatasetVersionNotFoundError
 from datachain.sql.types import NAME_TYPES_MAPPING, SQLType
 
 if TYPE_CHECKING:
@@ -417,7 +418,9 @@ class DatasetRecord:
 
     def get_version(self, version: int) -> DatasetVersion:
         if not self.has_version(version):
-            raise ValueError(f"Dataset {self.name} does not have version {version}")
+            raise DatasetVersionNotFoundError(
+                f"Dataset {self.name} does not have version {version}"
+            )
         return next(
             v
             for v in self.versions  # type: ignore [union-attr]
@@ -435,7 +438,9 @@ class DatasetRecord:
         Get identifier in the form my-dataset@v3
         """
         if not self.has_version(version):
-            raise ValueError(f"Dataset {self.name} doesn't have a version {version}")
+            raise DatasetVersionNotFoundError(
+                f"Dataset {self.name} doesn't have a version {version}"
+            )
         return f"{self.name}@v{version}"
 
     def uri(self, version: int) -> str:

datachain/error.py

@@ -10,6 +10,10 @@ class DatasetNotFoundError(NotFoundError):
     pass
 
 
+class DatasetVersionNotFoundError(NotFoundError):
+    pass
+
+
 class DatasetInvalidVersionError(Exception):
     pass
 
@@ -32,14 +36,12 @@ class QueryScriptRunError(Exception):
     Attributes:
         message      Explanation of the error
         return_code  Code returned by the subprocess
-        output       STDOUT + STDERR output of the subprocess
     """
 
-    def __init__(self, message: str, return_code: int = 0, output: str = ""):
+    def __init__(self, message: str, return_code: int = 0):
         self.message = message
         self.return_code = return_code
-        self.output = output
-        super().__init__(self.message)
+        super().__init__(message)
 
 
 class QueryScriptCancelError(QueryScriptRunError):

datachain/lib/arrow.py

@@ -4,11 +4,11 @@ from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING, Optional
 
 import pyarrow as pa
-from pyarrow.dataset import dataset
+from pyarrow.dataset import CsvFileFormat, dataset
 from tqdm import tqdm
 
 from datachain.lib.data_model import dict_to_data_model
-from datachain.lib.file import File, IndexedFile
+from datachain.lib.file import ArrowRow, File
 from datachain.lib.model_store import ModelStore
 from datachain.lib.udf import Generator
 
@@ -84,7 +84,12 @@ class ArrowGenerator(Generator):
                                 vals_dict[field] = val
                         vals = [self.output_schema(**vals_dict)]
                     if self.source:
-                        yield [IndexedFile(file=file, index=index), *vals]
+                        kwargs: dict = self.kwargs
+                        # Can't serialize CsvFileFormat; may lose formatting options.
+                        if isinstance(kwargs.get("format"), CsvFileFormat):
+                            kwargs["format"] = "csv"
+                        arrow_file = ArrowRow(file=file, index=index, kwargs=kwargs)
+                        yield [arrow_file, *vals]
                     else:
                         yield vals
                     index += 1

datachain/lib/dc.py

@@ -26,8 +26,8 @@ from datachain.lib.convert.python_to_sql import python_to_sql
 from datachain.lib.convert.values_to_tuples import values_to_tuples
 from datachain.lib.data_model import DataModel, DataType, dict_to_data_model
 from datachain.lib.dataset_info import DatasetInfo
+from datachain.lib.file import ArrowRow, File, get_file_type
 from datachain.lib.file import ExportPlacement as FileExportPlacement
-from datachain.lib.file import File, IndexedFile, get_file_type
 from datachain.lib.listing import (
     is_listing_dataset,
     is_listing_expired,
@@ -1614,7 +1614,7 @@ class DataChain(DatasetQuery):
                 for name, info in output.model_fields.items()
             }
         if source:
-            output = {"source": IndexedFile} | output  # type: ignore[assignment,operator]
+            output = {"source": ArrowRow} | output  # type: ignore[assignment,operator]
         return self.gen(
             ArrowGenerator(schema, model, source, nrows, **kwargs), output=output
         )

datachain/lib/file.py

@@ -17,6 +17,7 @@ from urllib.request import url2pathname
 
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 from PIL import Image
+from pyarrow.dataset import dataset
 from pydantic import Field, field_validator
 
 if TYPE_CHECKING:
@@ -439,14 +440,31 @@ class ImageFile(File):
         self.read().save(destination)
 
 
-class IndexedFile(DataModel):
-    """Metadata indexed from tabular files.
-
-    Includes `file` and `index` signals.
-    """
+class ArrowRow(DataModel):
+    """`DataModel` for reading row from Arrow-supported file."""
 
     file: File
     index: int
+    kwargs: dict
+
+    @contextmanager
+    def open(self):
+        """Stream row contents from indexed file."""
+        if self.file._caching_enabled:
+            self.file.ensure_cached()
+            path = self.file.get_local_path()
+            ds = dataset(path, **self.kwargs)
+
+        else:
+            path = self.file.get_path()
+            ds = dataset(path, filesystem=self.file.get_fs(), **self.kwargs)
+
+        return ds.take([self.index]).to_reader()
+
+    def read(self):
+        """Returns row contents as dict."""
+        with self.open() as record_batch:
+            return record_batch.to_pylist()[0]
 
 
 def get_file_type(type_: Literal["binary", "text", "image"] = "binary") -> type[File]:

datachain/query/dataset.py

@@ -1604,25 +1604,3 @@ class DatasetQuery:
         finally:
             self.cleanup()
         return self.__class__(name=name, version=version, catalog=self.catalog)
-
-
-def query_wrapper(dataset_query: Any) -> Any:
-    """
-    Wrapper function that wraps the last statement of user query script.
-    Last statement MUST be instance of DatasetQuery, otherwise script exits with
-    error code 10
-    """
-    if not isinstance(dataset_query, DatasetQuery):
-        return dataset_query
-
-    catalog = dataset_query.catalog
-    save = bool(os.getenv("DATACHAIN_QUERY_SAVE"))
-
-    is_session_temp_dataset = dataset_query.name and dataset_query.name.startswith(
-        dataset_query.session.get_temp_prefix()
-    )
-
-    if save and (is_session_temp_dataset or not dataset_query.attached):
-        name = catalog.generate_query_dataset_name()
-        dataset_query = dataset_query.save(name)
-    return dataset_query

{datachain-0.3.18.dist-info → datachain-0.3.19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.3.18
+Version: 0.3.19
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -33,7 +33,6 @@ 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: dill ==0.3.8
 Requires-Dist: cloudpickle
 Requires-Dist: orjson >=3.10.5
 Requires-Dist: pydantic <3,>=2

{datachain-0.3.18.dist-info → datachain-0.3.19.dist-info}/RECORD

@@ -1,12 +1,12 @@
-datachain/__init__.py,sha256=GeyhE-5LgfJav2OKYGaieP2lBvf2Gm-ihj7thnK9zjI,800
+datachain/__init__.py,sha256=ofPJ6B-d-ybSDRrE7J6wqF_ZRAB2W9U8l-eeuBtqPLg,865
 datachain/__main__.py,sha256=hG3Y4ARGEqe1AWwNMd259rBlqtphx1Wk39YbueQ0yV8,91
 datachain/asyn.py,sha256=Lg3Ck1PQLjQziMx9KU4atzbEnJXTE0924WMYkhgWtGU,8247
 datachain/cache.py,sha256=s0YHN7qurmQv-eC265TjeureK84TebWWAnL07cxchZQ,2997
-datachain/cli.py,sha256=tRuUvlFey5zYE0UVkGylqGiG5t89gUBo2SJ_yPsvu1I,30129
+datachain/cli.py,sha256=TQ1OKMulAcsJndKLCyxJpfNqbMWQgOa4Aeihnu36cR8,30095
 datachain/cli_utils.py,sha256=jrn9ejGXjybeO1ur3fjdSiAyCHZrX0qsLLbJzN9ErPM,2418
 datachain/config.py,sha256=PfC7W5yO6HFO6-iMB4YB-0RR88LPiGmD6sS_SfVbGso,1979
-datachain/dataset.py,sha256=sHnsmKfMg2bK88gZH1izk8jlbmJDEhQpyOemdaPQVFo,14761
-datachain/error.py,sha256=OnZ8OaBtDdTZPy8XQiy29SAjqdQArQeorYbP5ju7ldc,1199
+datachain/dataset.py,sha256=2NCQU9ZSgNGhA01SP5ON18VhMohXif-btOB4Lz-Uvds,14911
+datachain/error.py,sha256=vbIbamnFMIojh1UpmxWoA6Omup7WFAFNJnf8xAkGWwI,1146
 datachain/job.py,sha256=Jt4sNutMHJReaGsj3r3scueN5aESLGfhimAa8pUP7Is,1271
 datachain/listing.py,sha256=TkMmBzCiru26x4RaZiagWJTmTGbiy6yGrAsSJMr8cFE,8213
 datachain/node.py,sha256=ThE6Ue4BqpaBvrkFFJW_ljLxchixUX2aWz3l_nbwY54,5195
@@ -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=tICInYEeCRJow9hNSFnlA50hCOjFPN7fyGgoN5shcf8,67985
+datachain/catalog/catalog.py,sha256=poTu_B5va35MTCV60ntsn4jvAFXepqa2peCjYCXWeU0,64982
 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
@@ -39,12 +39,12 @@ datachain/data_storage/serializer.py,sha256=6G2YtOFqqDzJf1KbvZraKGXl2XHZyVml2kru
 datachain/data_storage/sqlite.py,sha256=EBKJncuzcyQfcKFm2mUjvHjHRTODsteM-k_zndunBrw,28834
 datachain/data_storage/warehouse.py,sha256=Vwhu_OfcNAoTtg1BHui80VCzlPeTUjZQL0QWziu8awY,32186
 datachain/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datachain/lib/arrow.py,sha256=uYn9RQwJy4MsMkhu18_6cgtVO3HkniBcB1NdFmkwtvo,7292
+datachain/lib/arrow.py,sha256=aUsoQmxDmuSnB8Ik9p57Y66gc_dgx6NBqkDDIfLsvno,7630
 datachain/lib/clip.py,sha256=lm5CzVi4Cj1jVLEKvERKArb-egb9j1Ls-fwTItT6vlI,6150
 datachain/lib/data_model.py,sha256=gHIjlow84GMRDa78yLL1Ud-N18or21fnTyPEwsatpXY,2045
 datachain/lib/dataset_info.py,sha256=srPPhI2UHf6hFPBecyFEVw2SS5aPisIIMsvGgKqi7ss,2366
-datachain/lib/dc.py,sha256=oc9tPf5G9X3DmVCPaTuuwp8LlLogoYuEHpOkq_W7h6Y,68984
-datachain/lib/file.py,sha256=flKGvmrotXWZqQQafaZQEeSQlqUVTSVWB7JIkEsr0MM,14255
+datachain/lib/dc.py,sha256=kabEHnqbcoat7gd-yl0PvmuC6SyKbRa8r7NWKcN6GEQ,68978
+datachain/lib/file.py,sha256=LjTW_-PDAnoUhvyB4bJ8Y8n__XGqrxvmd9mDOF0Gir8,14875
 datachain/lib/hf.py,sha256=cPnmLuprr0pYABH7KqA5FARQ1JGlywdDwD3yDzVAm4k,5920
 datachain/lib/image.py,sha256=AMXYwQsmarZjRbPCZY3M1jDsM2WAB_b3cTY4uOIuXNU,2675
 datachain/lib/listing.py,sha256=cHPN5-Fq8yb0gP6DARImhmZWxykDDNqhhJujDxEp53A,4104
@@ -70,7 +70,7 @@ datachain/lib/convert/unflatten.py,sha256=Ogvh_5wg2f38_At_1lN0D_e2uZOOpYEvwvB2xd
 datachain/lib/convert/values_to_tuples.py,sha256=YOdbjzHq-uj6-cV2Qq43G72eN2avMNDGl4x5t6yQMl8,3931
 datachain/query/__init__.py,sha256=0NBOZVgIDpCcj1Ci883dQ9A0iiwe03xzmotkOCFbxYc,293
 datachain/query/batch.py,sha256=-vlpINJiertlnaoUVv1C95RatU0F6zuhpIYRufJRo1M,3660
-datachain/query/dataset.py,sha256=k2jU0uZ86i9vr3On-o7GzHrubK5bCJjZEvz9P8extmw,54347
+datachain/query/dataset.py,sha256=F9WEVhDuFm6NQT6l-Vi3PMU-mQVpqwKHMgZIA4eWB18,53602
 datachain/query/dispatch.py,sha256=CFAc09O6UllcyUSSEY1GUlEMPzeO8RYhXinNN4HBl9M,12405
 datachain/query/metrics.py,sha256=r5b0ygYhokbXp8Mg3kCH8iFSRw0jxzyeBe-C-J_bKFc,938
 datachain/query/params.py,sha256=O_j89mjYRLOwWNhYZl-z7mi-rkdP7WyFmaDufsdTryE,863
@@ -97,9 +97,9 @@ datachain/sql/sqlite/base.py,sha256=WLPHBhZbXbiqPoRV1VgDrXJqku4UuvJpBhYeQ0k5rI8,
 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.3.18.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.3.18.dist-info/METADATA,sha256=_LpwSHtaSTA-rz4rG9nHIbO2mLlrlI4mCnlxKx8vePo,17185
-datachain-0.3.18.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-datachain-0.3.18.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.3.18.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.3.18.dist-info/RECORD,,
+datachain-0.3.19.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.3.19.dist-info/METADATA,sha256=yMBpXwOmeoWOmpS0m_hp8GFiMs3Zu_ixMzkG6GF_Z2U,17157
+datachain-0.3.19.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+datachain-0.3.19.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.3.19.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.3.19.dist-info/RECORD,,