datachain 0.33.1__tar.gz → 0.34.1__tar.gz

{datachain-0.33.1 → datachain-0.34.1}/.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.13.1'
+    rev: 'v0.13.2'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.33.1 → datachain-0.34.1}/PKG-INFO

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

{datachain-0.33.1 → datachain-0.34.1}/src/datachain/catalog/catalog.py

@@ -144,26 +144,19 @@ def shutdown_process(
             return proc.wait()
 
 
-def process_output(stream: IO[bytes], callback: Callable[[str], None]) -> None:
+def _process_stream(stream: "IO[bytes]", callback: Callable[[str], None]) -> None:
     buffer = b""
+    while byt := stream.read(1):  # Read one byte at a time
+        buffer += byt
 
-    try:
-        while byt := stream.read(1):  # Read one byte at a time
-            buffer += byt
-
-            if byt in (b"\n", b"\r"):  # Check for newline or carriage return
-                line = buffer.decode("utf-8", errors="replace")
-                callback(line)
-                buffer = b""  # Clear buffer for the next line
-
-        if buffer:  # Handle any remaining data in the buffer
-            line = buffer.decode("utf-8", errors="replace")
+        if byt in (b"\n", b"\r"):  # Check for newline or carriage return
+            line = buffer.decode("utf-8")
             callback(line)
-    finally:
-        try:
-            stream.close()  # Ensure output is closed
-        except Exception:  # noqa: BLE001, S110
-            pass
+            buffer = b""  # Clear buffer for next line
+
+    if buffer:  # Handle any remaining data in the buffer
+        line = buffer.decode("utf-8")
+        callback(line)
 
 
 class DatasetRowsFetcher(NodesThreadPool):
@@ -1767,13 +1760,13 @@ class Catalog:
             recursive=recursive,
         )
 
-    @staticmethod
     def query(
+        self,
         query_script: str,
         env: Optional[Mapping[str, str]] = None,
         python_executable: str = sys.executable,
-        stdout_callback: Optional[Callable[[str], None]] = None,
-        stderr_callback: Optional[Callable[[str], None]] = None,
+        capture_output: bool = False,
+        output_hook: Callable[[str], None] = noop,
         params: Optional[dict[str, str]] = None,
         job_id: Optional[str] = None,
         interrupt_timeout: Optional[int] = None,
@@ -1788,18 +1781,13 @@ class Catalog:
             },
         )
         popen_kwargs: dict[str, Any] = {}
-
-        if stdout_callback is not None:
-            popen_kwargs = {"stdout": subprocess.PIPE}
-        if stderr_callback is not None:
-            popen_kwargs["stderr"] = subprocess.PIPE
+        if capture_output:
+            popen_kwargs = {"stdout": subprocess.PIPE, "stderr": subprocess.STDOUT}
 
         def raise_termination_signal(sig: int, _: Any) -> NoReturn:
             raise TerminationSignal(sig)
 
-        stdout_thread: Optional[Thread] = None
-        stderr_thread: Optional[Thread] = None
-
+        thread: Optional[Thread] = None
         with subprocess.Popen(cmd, env=env, **popen_kwargs) as proc:  # noqa: S603
             logger.info("Starting process %s", proc.pid)
 
@@ -1813,20 +1801,10 @@ class Catalog:
             orig_sigterm_handler = signal.getsignal(signal.SIGTERM)
             signal.signal(signal.SIGTERM, raise_termination_signal)
             try:
-                if stdout_callback is not None:
-                    stdout_thread = Thread(
-                        target=process_output,
-                        args=(proc.stdout, stdout_callback),
-                        daemon=True,
-                    )
-                    stdout_thread.start()
-                if stderr_callback is not None:
-                    stderr_thread = Thread(
-                        target=process_output,
-                        args=(proc.stderr, stderr_callback),
-                        daemon=True,
-                    )
-                    stderr_thread.start()
+                if capture_output:
+                    args = (proc.stdout, output_hook)
+                    thread = Thread(target=_process_stream, args=args, daemon=True)
+                    thread.start()
 
                 proc.wait()
             except TerminationSignal as exc:
@@ -1844,22 +1822,8 @@ class Catalog:
             finally:
                 signal.signal(signal.SIGTERM, orig_sigterm_handler)
                 signal.signal(signal.SIGINT, orig_sigint_handler)
-                # wait for the reader thread
-                thread_join_timeout_seconds = 30
-                if stdout_thread is not None:
-                    stdout_thread.join(timeout=thread_join_timeout_seconds)
-                    if stdout_thread.is_alive():
-                        logger.warning(
-                            "stdout thread is still alive after %s seconds",
-                            thread_join_timeout_seconds,
-                        )
-                if stderr_thread is not None:
-                    stderr_thread.join(timeout=thread_join_timeout_seconds)
-                    if stderr_thread.is_alive():
-                        logger.warning(
-                            "stderr thread is still alive after %s seconds",
-                            thread_join_timeout_seconds,
-                        )
+                if thread:
+                    thread.join()  # wait for the reader thread
 
         logger.info("Process %s exited with return code %s", proc.pid, proc.returncode)
         if proc.returncode in (

{datachain-0.33.1 → datachain-0.34.1}/src/datachain/data_storage/metastore.py

@@ -21,6 +21,7 @@ from sqlalchemy import (
     Table,
     Text,
     UniqueConstraint,
+    desc,
     select,
 )
 from sqlalchemy.sql import func as f
@@ -399,6 +400,7 @@ class AbstractMetastore(ABC, Serializable):
         workers: int = 1,
         python_version: Optional[str] = None,
         params: Optional[dict[str, str]] = None,
+        parent_job_id: Optional[str] = None,
     ) -> str:
         """
         Creates a new job.
@@ -443,6 +445,10 @@ class AbstractMetastore(ABC, Serializable):
     def list_checkpoints(self, job_id: str, conn=None) -> Iterator["Checkpoint"]:
         """Returns all checkpoints related to some job"""
 
+    @abstractmethod
+    def get_last_checkpoint(self, job_id: str, conn=None) -> Optional[Checkpoint]:
+        """Get last created checkpoint for some job."""
+
     @abstractmethod
     def get_checkpoint_by_id(self, checkpoint_id: str, conn=None) -> Checkpoint:
         """Gets single checkpoint by id"""
@@ -1548,6 +1554,7 @@ class AbstractDBMetastore(AbstractMetastore):
             Column("error_stack", Text, nullable=False, default=""),
             Column("params", JSON, nullable=False),
             Column("metrics", JSON, nullable=False),
+            Column("parent_job_id", Text, nullable=True),
         ]
 
     @cached_property
@@ -1595,6 +1602,7 @@ class AbstractDBMetastore(AbstractMetastore):
         workers: int = 1,
         python_version: Optional[str] = None,
         params: Optional[dict[str, str]] = None,
+        parent_job_id: Optional[str] = None,
         conn: Optional[Any] = None,
     ) -> str:
         """
@@ -1616,6 +1624,7 @@ class AbstractDBMetastore(AbstractMetastore):
                 error_stack="",
                 params=json.dumps(params or {}),
                 metrics=json.dumps({}),
+                parent_job_id=parent_job_id,
             ),
             conn=conn,
         )
@@ -1770,7 +1779,7 @@ class AbstractDBMetastore(AbstractMetastore):
         )
         return self.get_checkpoint_by_id(checkpoint_id)
 
-    def list_checkpoints(self, job_id: str, conn=None) -> Iterator["Checkpoint"]:
+    def list_checkpoints(self, job_id: str, conn=None) -> Iterator[Checkpoint]:
         """List checkpoints by job id."""
         query = self._checkpoints_query().where(self._checkpoints.c.job_id == job_id)
         rows = list(self.db.execute(query, conn=conn))
@@ -1800,3 +1809,15 @@ class AbstractDBMetastore(AbstractMetastore):
         if not rows:
             return None
         return self.checkpoint_class.parse(*rows[0])
+
+    def get_last_checkpoint(self, job_id: str, conn=None) -> Optional[Checkpoint]:
+        query = (
+            self._checkpoints_query()
+            .where(self._checkpoints.c.job_id == job_id)
+            .order_by(desc(self._checkpoints.c.created_at))
+            .limit(1)
+        )
+        rows = list(self.db.execute(query, conn=conn))
+        if not rows:
+            return None
+        return self.checkpoint_class.parse(*rows[0])

{datachain-0.33.1 → datachain-0.34.1}/src/datachain/error.py

@@ -101,3 +101,7 @@ class OutdatedDatabaseSchemaError(DataChainError):
 
 class CheckpointNotFoundError(NotFoundError):
     pass
+
+
+class JobNotFoundError(NotFoundError):
+    pass

{datachain-0.33.1 → datachain-0.34.1}/src/datachain/job.py

@@ -22,6 +22,7 @@ class Job:
     python_version: Optional[str] = None
     error_message: str = ""
     error_stack: str = ""
+    parent_job_id: Optional[str] = None
 
     @classmethod
     def parse(
@@ -39,6 +40,7 @@ class Job:
         error_stack: str,
         params: str,
         metrics: str,
+        parent_job_id: Optional[str],
     ) -> "Job":
         return cls(
             str(id),
@@ -54,4 +56,5 @@ class Job:
             python_version,
             error_message,
             error_stack,
+            parent_job_id,
         )

{datachain-0.33.1 → datachain-0.34.1}/src/datachain/lib/dc/datachain.py

@@ -19,7 +19,6 @@ from typing import (
     cast,
     overload,
 )
-from uuid import uuid4
 
 import sqlalchemy
 import ujson as json
@@ -30,10 +29,15 @@ from tqdm import tqdm
 from datachain import semver
 from datachain.dataset import DatasetRecord
 from datachain.delta import delta_disabled
-from datachain.error import ProjectCreateNotAllowedError, ProjectNotFoundError
+from datachain.error import (
+    JobNotFoundError,
+    ProjectCreateNotAllowedError,
+    ProjectNotFoundError,
+)
 from datachain.func import literal
 from datachain.func.base import Function
 from datachain.func.func import Func
+from datachain.job import Job
 from datachain.lib.convert.python_to_sql import python_to_sql
 from datachain.lib.data_model import (
     DataModel,
@@ -50,11 +54,12 @@ from datachain.lib.signal_schema import SignalResolvingError, SignalSchema
 from datachain.lib.udf import Aggregator, BatchMapper, Generator, Mapper, UDFBase
 from datachain.lib.udf_signature import UdfSignature
 from datachain.lib.utils import DataChainColumnError, DataChainParamsError
+from datachain.project import Project
 from datachain.query import Session
 from datachain.query.dataset import DatasetQuery, PartitionByType
 from datachain.query.schema import DEFAULT_DELIMITER, Column
 from datachain.sql.functions import path as pathfunc
-from datachain.utils import batched_it, inside_notebook, row_to_nested_dict
+from datachain.utils import batched_it, env2bool, inside_notebook, row_to_nested_dict
 
 from .database import DEFAULT_DATABASE_BATCH_SIZE
 from .utils import (
@@ -578,6 +583,19 @@ class DataChain:
             query=self._query.save(project=project, feature_schema=schema)
         )
 
+    def _calculate_job_hash(self, job_id: str) -> str:
+        """
+        Calculates hash of the job at the place of this chain's save method.
+        Hash is calculated using previous job checkpoint hash (if exists) and
+        adding hash of this chain to produce new hash.
+        """
+        last_checkpoint = self.session.catalog.metastore.get_last_checkpoint(job_id)
+
+        return hashlib.sha256(
+            (bytes.fromhex(last_checkpoint.hash) if last_checkpoint else b"")
+            + bytes.fromhex(self.hash())
+        ).hexdigest()
+
     def save(  # type: ignore[override]
         self,
         name: str,
@@ -602,101 +620,171 @@ class DataChain:
             update_version: which part of the dataset version to automatically increase.
                 Available values: `major`, `minor` or `patch`. Default is `patch`.
         """
+
         catalog = self.session.catalog
-        if version is not None:
-            semver.validate(version)
 
-        if update_version is not None and update_version not in [
-            "patch",
-            "major",
-            "minor",
-        ]:
-            raise ValueError(
-                "update_version can have one of the following values: major, minor or"
-                " patch"
-            )
+        result = None  # result chain that will be returned at the end
+
+        # Version validation
+        self._validate_version(version)
+        self._validate_update_version(update_version)
 
         namespace_name, project_name, name = catalog.get_full_dataset_name(
             name,
             namespace_name=self._settings.namespace,
             project_name=self._settings.project,
         )
+        project = self._get_or_create_project(namespace_name, project_name)
+
+        # Checkpoint handling
+        job, _hash, result = self._resolve_checkpoint(name, project, kwargs)
+
+        # Schema preparation
+        schema = self.signals_schema.clone_without_sys_signals().serialize()
 
+        # Handle retry and delta functionality
+        if not result:
+            result = self._handle_delta(name, version, project, schema, kwargs)
+
+        if not result:
+            # calculate chain if we already don't have result from checkpoint or delta
+            result = self._evolve(
+                query=self._query.save(
+                    name=name,
+                    version=version,
+                    project=project,
+                    description=description,
+                    attrs=attrs,
+                    feature_schema=schema,
+                    update_version=update_version,
+                    **kwargs,
+                )
+            )
+
+        if job:
+            catalog.metastore.create_checkpoint(job.id, _hash)  # type: ignore[arg-type]
+
+        return result
+
+    def _validate_version(self, version: Optional[str]) -> None:
+        """Validate dataset version if provided."""
+        if version is not None:
+            semver.validate(version)
+
+    def _validate_update_version(self, update_version: Optional[str]) -> None:
+        """Ensure update_version is one of: major, minor, patch."""
+        allowed = ["major", "minor", "patch"]
+        if update_version not in allowed:
+            raise ValueError(f"update_version must be one of {allowed}")
+
+    def _get_or_create_project(self, namespace: str, project_name: str) -> Project:
+        """Get project or raise if creation not allowed."""
         try:
-            project = self.session.catalog.metastore.get_project(
+            return self.session.catalog.metastore.get_project(
                 project_name,
-                namespace_name,
+                namespace,
                 create=is_studio(),
             )
         except ProjectNotFoundError as e:
-            # not being able to create it as creation is not allowed
             raise ProjectCreateNotAllowedError("Creating project is not allowed") from e
 
-        schema = self.signals_schema.clone_without_sys_signals().serialize()
+    def _resolve_checkpoint(
+        self,
+        name: str,
+        project: Project,
+        kwargs: dict,
+    ) -> tuple[Optional[Job], Optional[str], Optional["DataChain"]]:
+        """Check if checkpoint exists and return cached dataset if possible."""
+        from .datasets import read_dataset
 
-        # Handle retry and delta functionality
-        if self.delta and name:
-            from datachain.delta import delta_retry_update
+        metastore = self.session.catalog.metastore
 
-            # Delta chains must have delta_on defined (ensured by _as_delta method)
-            assert self._delta_on is not None, "Delta chain must have delta_on defined"
+        job_id = os.getenv("DATACHAIN_JOB_ID")
+        checkpoints_reset = env2bool("DATACHAIN_CHECKPOINTS_RESET", undefined=True)
 
-            result_ds, dependencies, has_changes = delta_retry_update(
-                self,
-                namespace_name,
-                project_name,
-                name,
-                on=self._delta_on,
-                right_on=self._delta_result_on,
-                compare=self._delta_compare,
-                delta_retry=self._delta_retry,
+        if not job_id:
+            return None, None, None
+
+        job = metastore.get_job(job_id)
+        if not job:
+            raise JobNotFoundError(f"Job with id {job_id} not found")
+
+        _hash = self._calculate_job_hash(job.id)
+
+        if (
+            job.parent_job_id
+            and not checkpoints_reset
+            and metastore.find_checkpoint(job.parent_job_id, _hash)
+        ):
+            # checkpoint found → reuse dataset
+            chain = read_dataset(
+                name, namespace=project.namespace.name, project=project.name, **kwargs
             )
+            return job, _hash, chain
 
-            if result_ds:
-                return self._evolve(
-                    query=result_ds._query.save(
-                        name=name,
-                        version=version,
-                        project=project,
-                        feature_schema=schema,
-                        dependencies=dependencies,
-                        **kwargs,
-                    )
-                )
+        return job, _hash, None
 
-            if not has_changes:
-                # sources have not been changed so new version of resulting dataset
-                # would be the same as previous one. To avoid duplicating exact
-                # datasets, we won't create new version of it and we will return
-                # current latest version instead.
-                from .datasets import read_dataset
+    def _handle_delta(
+        self,
+        name: str,
+        version: Optional[str],
+        project: Project,
+        schema: dict,
+        kwargs: dict,
+    ) -> Optional["DataChain"]:
+        """Try to save as a delta dataset.
+        Returns:
+            A DataChain if delta logic could handle it, otherwise None to fall back
+            to the regular save path (e.g., on first dataset creation).
+        """
+        from datachain.delta import delta_retry_update
 
-                return read_dataset(
-                    name, namespace=namespace_name, project=project_name, **kwargs
-                )
+        from .datasets import read_dataset
 
-        result = self._evolve(
-            query=self._query.save(
-                name=name,
-                version=version,
-                project=project,
-                description=description,
-                attrs=attrs,
-                feature_schema=schema,
-                update_version=update_version,
-                **kwargs,
-            )
+        if not self.delta or not name:
+            return None
+
+        assert self._delta_on is not None, "Delta chain must have delta_on defined"
+
+        result_ds, dependencies, has_changes = delta_retry_update(
+            self,
+            project.namespace.name,
+            project.name,
+            name,
+            on=self._delta_on,
+            right_on=self._delta_result_on,
+            compare=self._delta_compare,
+            delta_retry=self._delta_retry,
         )
 
-        if job_id := os.getenv("DATACHAIN_JOB_ID"):
-            catalog.metastore.create_checkpoint(
-                job_id,
-                _hash=hashlib.sha256(  # TODO this will be replaced with self.hash()
-                    str(uuid4()).encode()
-                ).hexdigest(),
+        # Case 1: delta produced a new dataset
+        if result_ds:
+            return self._evolve(
+                query=result_ds._query.save(
+                    name=name,
+                    version=version,
+                    project=project,
+                    feature_schema=schema,
+                    dependencies=dependencies,
+                    **kwargs,
+                )
             )
 
-        return result
+        # Case 2: no changes → reuse last version
+        if not has_changes:
+            # sources have not been changed so new version of resulting dataset
+            # would be the same as previous one. To avoid duplicating exact
+            # datasets, we won't create new version of it and we will return
+            # current latest version instead.
+            return read_dataset(
+                name,
+                namespace=project.namespace.name,
+                project=project.name,
+                **kwargs,
+            )
+
+        # Case 3: first creation of dataset
+        return None
 
     def apply(self, func, *args, **kwargs):
         """Apply any function to the chain.
@@ -1875,12 +1963,15 @@ class DataChain:
         self,
         flatten: bool = False,
         include_hidden: bool = True,
+        as_object: bool = False,
     ) -> "pd.DataFrame":
         """Return a pandas DataFrame from the chain.
 
         Parameters:
             flatten: Whether to use a multiindex or flatten column names.
             include_hidden: Whether to include hidden columns.
+            as_object: Whether to emit a dataframe backed by Python objects
+                rather than pandas-inferred dtypes.
 
         Returns:
             pd.DataFrame: A pandas DataFrame representation of the chain.
@@ -1896,6 +1987,9 @@ class DataChain:
             columns = pd.MultiIndex.from_tuples(map(tuple, headers))
 
         results = self.results(include_hidden=include_hidden)
+        if as_object:
+            df = pd.DataFrame(results, columns=columns, dtype=object)
+            return df.where(pd.notna(df), None)
         return pd.DataFrame.from_records(results, columns=columns)
 
     def show(
@@ -1918,7 +2012,11 @@ class DataChain:
         import pandas as pd
 
         dc = self.limit(limit) if limit > 0 else self  # type: ignore[misc]
-        df = dc.to_pandas(flatten, include_hidden=include_hidden)
+        df = dc.to_pandas(
+            flatten,
+            include_hidden=include_hidden,
+            as_object=True,
+        )
 
         if df.empty:
             print("Empty result")

{datachain-0.33.1 → datachain-0.34.1}/src/datachain.egg-info/PKG-INFO

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

{datachain-0.33.1 → datachain-0.34.1}/src/datachain.egg-info/SOURCES.txt

@@ -394,6 +394,7 @@ tests/unit/lib/__init__.py
 tests/unit/lib/conftest.py
 tests/unit/lib/test_arrow.py
 tests/unit/lib/test_audio.py
+tests/unit/lib/test_checkpoints.py
 tests/unit/lib/test_clip.py
 tests/unit/lib/test_datachain.py
 tests/unit/lib/test_datachain_bootstrap.py

{datachain-0.33.1 → datachain-0.34.1}/tests/conftest.py

@@ -456,9 +456,12 @@ def cloud_server(request, tmp_upath_factory, cloud_type, version_aware, tree):
     return make_cloud_server(src_path, cloud_type, tree)
 
 
-@pytest.fixture()
-def datachain_job_id(monkeypatch):
-    job_id = str(uuid.uuid4())
+@pytest.fixture
+def datachain_job_id(test_session, monkeypatch):
+    job_id = test_session.catalog.metastore.create_job(
+        "my-job",
+        'import datachain as dc; dc.read_values(num=[1, 2, 3].save("nums")',
+    )
     monkeypatch.setenv("DATACHAIN_JOB_ID", job_id)
     return job_id
 

{datachain-0.33.1 → datachain-0.34.1}/tests/func/test_datachain.py

@@ -35,6 +35,7 @@ from tests.utils import (
     TARRED_TREE,
     df_equal,
     images_equal,
+    skip_if_not_sqlite,
     sorted_dicts,
     text_embedding,
 )
@@ -559,6 +560,25 @@ def test_show(capsys, test_session):
         assert f"{i} {first_name[i]}" in normalized_output
 
 
+@skip_if_not_sqlite
+def test_show_preserves_none(capsys, test_session):
+    chain = dc.read_values(
+        score=[1, None],
+        ts=[
+            datetime(2020, 1, 1, tzinfo=timezone.utc),
+            None,
+        ],
+        session=test_session,
+    )
+
+    chain.show()
+
+    captured = capsys.readouterr().out
+    assert "NaN" not in captured
+    assert "NaT" not in captured
+    assert captured.count("None") >= 2
+
+
 def test_show_without_temp_datasets(capsys, test_session):
     dc.read_values(
         key=[1, 2, 3, 4], session=test_session