datachain 0.33.0__tar.gz → 0.34.0__tar.gz

{datachain-0.33.0 → datachain-0.34.0}/.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.0 → datachain-0.34.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.33.0
+Version: 0.34.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -103,7 +103,7 @@ Requires-Dist: scipy; extra == "tests"
 Requires-Dist: ultralytics; extra == "tests"
 Provides-Extra: dev
 Requires-Dist: datachain[docs,tests]; extra == "dev"
-Requires-Dist: mypy==1.18.1; extra == "dev"
+Requires-Dist: mypy==1.18.2; extra == "dev"
 Requires-Dist: types-python-dateutil; extra == "dev"
 Requires-Dist: types-dateparser; extra == "dev"
 Requires-Dist: types-pytz; extra == "dev"

{datachain-0.33.0 → datachain-0.34.0}/pyproject.toml

@@ -120,7 +120,7 @@ tests = [
 ]
 dev = [
   "datachain[docs,tests]",
-  "mypy==1.18.1",
+  "mypy==1.18.2",
   "types-python-dateutil",
   "types-dateparser",
   "types-pytz",

{datachain-0.33.0 → datachain-0.34.0}/src/datachain/catalog/catalog.py

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

{datachain-0.33.0 → datachain-0.34.0}/src/datachain/data_storage/job.py

@@ -4,6 +4,7 @@ from enum import Enum
 class JobStatus(int, Enum):
     CREATED = 1
     SCHEDULED = 10
+    PROVISIONING = 12
     QUEUED = 2
     INIT = 3
     RUNNING = 4

{datachain-0.33.0 → datachain-0.34.0}/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.0 → datachain-0.34.0}/src/datachain/diff/__init__.py

@@ -1,5 +1,3 @@
-import random
-import string
 from collections.abc import Sequence
 from enum import Enum
 from typing import TYPE_CHECKING, Optional, Union
@@ -11,16 +9,12 @@ from datachain.query.schema import Column
 if TYPE_CHECKING:
     from datachain.lib.dc import DataChain
 
-
 C = Column
 
 
-def get_status_col_name() -> str:
-    """Returns new unique status col name"""
-    return "diff_" + "".join(
-        random.choice(string.ascii_letters)  # noqa: S311
-        for _ in range(10)
-    )
+STATUS_COL_NAME = "diff_7aeed3aa17ba4d50b8d1c368c76e16a6"
+LEFT_DIFF_COL_NAME = "diff_95f95344064a4b819c8625cd1a5cfc2b"
+RIGHT_DIFF_COL_NAME = "diff_5808838a49b54849aa461d7387376d34"
 
 
 class CompareStatus(str, Enum):
@@ -101,9 +95,9 @@ def _compare(  # noqa: C901, PLR0912
         compare = right_compare = [c for c in cols if c in right_cols and c not in on]  # type: ignore[misc]
 
     # get diff column names
-    diff_col = status_col or get_status_col_name()
-    ldiff_col = get_status_col_name()
-    rdiff_col = get_status_col_name()
+    diff_col = status_col or STATUS_COL_NAME
+    ldiff_col = LEFT_DIFF_COL_NAME
+    rdiff_col = RIGHT_DIFF_COL_NAME
 
     # adding helper diff columns, which will be removed after
     left = left.mutate(**{ldiff_col: 1})
@@ -227,7 +221,7 @@ def compare_and_split(
         )
         ```
     """
-    status_col = get_status_col_name()
+    status_col = STATUS_COL_NAME
 
     res = _compare(
         left,

{datachain-0.33.0 → datachain-0.34.0}/src/datachain/error.py

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

datachain-0.34.0/src/datachain/hash_utils.py

@@ -0,0 +1,147 @@
+import hashlib
+import inspect
+import json
+import textwrap
+from collections.abc import Sequence
+from typing import TypeVar, Union
+
+from sqlalchemy.sql.elements import (
+    BinaryExpression,
+    BindParameter,
+    ColumnElement,
+    Label,
+    Over,
+    UnaryExpression,
+)
+from sqlalchemy.sql.functions import Function
+
+T = TypeVar("T", bound=ColumnElement)
+ColumnLike = Union[str, T]
+
+
+def serialize_column_element(expr: Union[str, ColumnElement]) -> dict:  # noqa: PLR0911
+    """
+    Recursively serialize a SQLAlchemy ColumnElement into a deterministic structure.
+    """
+
+    # Binary operations: col > 5, col1 + col2, etc.
+    if isinstance(expr, BinaryExpression):
+        op = (
+            expr.operator.__name__
+            if hasattr(expr.operator, "__name__")
+            else str(expr.operator)
+        )
+        return {
+            "type": "binary",
+            "op": op,
+            "left": serialize_column_element(expr.left),
+            "right": serialize_column_element(expr.right),
+        }
+
+    # Unary operations: -col, NOT col, etc.
+    if isinstance(expr, UnaryExpression):
+        op = (
+            expr.operator.__name__
+            if expr.operator is not None and hasattr(expr.operator, "__name__")
+            else str(expr.operator)
+        )
+
+        return {
+            "type": "unary",
+            "op": op,
+            "element": serialize_column_element(expr.element),  # type: ignore[arg-type]
+        }
+
+    # Function calls: func.lower(col), func.count(col), etc.
+    if isinstance(expr, Function):
+        return {
+            "type": "function",
+            "name": expr.name,
+            "clauses": [serialize_column_element(c) for c in expr.clauses],
+        }
+
+    # Window functions: func.row_number().over(partition_by=..., order_by=...)
+    if isinstance(expr, Over):
+        return {
+            "type": "window",
+            "function": serialize_column_element(expr.element),
+            "partition_by": [
+                serialize_column_element(p) for p in getattr(expr, "partition_by", [])
+            ],
+            "order_by": [
+                serialize_column_element(o) for o in getattr(expr, "order_by", [])
+            ],
+        }
+
+    # Labeled expressions: col.label("alias")
+    if isinstance(expr, Label):
+        return {
+            "type": "label",
+            "name": expr.name,
+            "element": serialize_column_element(expr.element),
+        }
+
+    # Bound values (constants)
+    if isinstance(expr, BindParameter):
+        return {"type": "bind", "value": expr.value}
+
+    # Plain columns
+    if hasattr(expr, "name"):
+        return {"type": "column", "name": expr.name}
+
+    # Fallback: stringify unknown nodes
+    return {"type": "other", "repr": str(expr)}
+
+
+def hash_column_elements(columns: Sequence[ColumnLike]) -> str:
+    """
+    Hash a list of ColumnElements deterministically, dialect agnostic.
+    Only accepts ordered iterables (like list or tuple).
+    """
+    serialized = [serialize_column_element(c) for c in columns]
+    json_str = json.dumps(serialized, sort_keys=True)  # stable JSON
+    return hashlib.sha256(json_str.encode("utf-8")).hexdigest()
+
+
+def hash_callable(func):
+    """
+    Calculate a hash from a callable.
+    Rules:
+    - Named functions (def) → use source code for stable, cross-version hashing
+    - Lambdas → use bytecode (deterministic in same Python runtime)
+    """
+    if not callable(func):
+        raise TypeError("Expected a callable")
+
+    # Determine if it is a lambda
+    is_lambda = func.__name__ == "<lambda>"
+
+    if not is_lambda:
+        # Try to get exact source of named function
+        try:
+            lines, _ = inspect.getsourcelines(func)
+            payload = textwrap.dedent("".join(lines)).strip()
+        except (OSError, TypeError):
+            # Fallback: bytecode if source not available
+            payload = func.__code__.co_code
+    else:
+        # For lambdas, fall back directly to bytecode
+        payload = func.__code__.co_code
+
+    # Normalize annotations
+    annotations = {
+        k: getattr(v, "__name__", str(v)) for k, v in func.__annotations__.items()
+    }
+
+    # Extras to distinguish functions with same code but different metadata
+    extras = {
+        "name": func.__name__,
+        "defaults": func.__defaults__,
+        "annotations": annotations,
+    }
+
+    # Compute SHA256
+    h = hashlib.sha256()
+    h.update(str(payload).encode() if isinstance(payload, str) else payload)
+    h.update(str(extras).encode())
+    return h.hexdigest()

{datachain-0.33.0 → datachain-0.34.0}/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,
         )