datachain 0.34.0__py3-none-any.whl → 0.34.2__py3-none-any.whl

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/catalog/loader.py

@@ -3,6 +3,7 @@ import sys
 from importlib import import_module
 from typing import TYPE_CHECKING, Any, Optional
 
+from datachain.plugins import ensure_plugins_loaded
 from datachain.utils import get_envs_by_prefix
 
 if TYPE_CHECKING:
@@ -24,6 +25,8 @@ IN_MEMORY_ERROR_MESSAGE = "In-memory is only supported on SQLite"
 
 
 def get_metastore(in_memory: bool = False) -> "AbstractMetastore":
+    ensure_plugins_loaded()
+
     from datachain.data_storage import AbstractMetastore
     from datachain.data_storage.serializer import deserialize
 
@@ -64,6 +67,8 @@ def get_metastore(in_memory: bool = False) -> "AbstractMetastore":
 
 
 def get_warehouse(in_memory: bool = False) -> "AbstractWarehouse":
+    ensure_plugins_loaded()
+
     from datachain.data_storage import AbstractWarehouse
     from datachain.data_storage.serializer import deserialize
 

datachain/data_storage/serializer.py

@@ -1,29 +1,119 @@
 import base64
-import pickle
+import json
 from abc import abstractmethod
 from collections.abc import Callable
-from typing import Any
+from typing import Any, ClassVar
+
+from datachain.plugins import ensure_plugins_loaded
+
+
+class CallableRegistry:
+    _registry: ClassVar[dict[str, Callable]] = {}
+
+    @classmethod
+    def register(cls, callable_obj: Callable, name: str) -> str:
+        cls._registry[name] = callable_obj
+        return name
+
+    @classmethod
+    def get(cls, name: str) -> Callable:
+        return cls._registry[name]
 
 
 class Serializable:
+    @classmethod
+    @abstractmethod
+    def serialize_callable_name(cls) -> str:
+        """Return the registered name used for this class' factory callable."""
+
     @abstractmethod
     def clone_params(self) -> tuple[Callable[..., Any], list[Any], dict[str, Any]]:
-        """
-        Returns the class, args, and kwargs needed to instantiate a cloned copy
-        of this instance for use in separate processes or machines.
-        """
+        """Return (callable, args, kwargs) necessary to recreate this object."""
+
+    def _prepare(self, params: tuple) -> dict:
+        callable, args, kwargs = params
+        callable_name = callable.__self__.serialize_callable_name()
+        return {
+            "callable": callable_name,
+            "args": args,
+            "kwargs": {
+                k: self._prepare(v) if isinstance(v, tuple) else v
+                for k, v in kwargs.items()
+            },
+        }
 
     def serialize(self) -> str:
-        """
-        Returns a string representation of clone params.
-        This is useful for storing the state of an object in environment variable.
-        """
-        return base64.b64encode(pickle.dumps(self.clone_params())).decode()
+        """Return a base64-encoded JSON string with registered callable + params."""
+        _ensure_default_callables_registered()
+        data = self.clone_params()
+        return base64.b64encode(json.dumps(self._prepare(data)).encode()).decode()
 
 
 def deserialize(s: str) -> Serializable:
+    """Deserialize from base64-encoded JSON using only registered callables.
+
+    Nested serialized objects are instantiated automatically except for those
+    passed via clone parameter tuples (keys ending with ``_clone_params``),
+    which must remain as (callable, args, kwargs) for later factory usage.
     """
-    Returns a new instance of the class represented by the string.
-    """
-    (f, args, kwargs) = pickle.loads(base64.b64decode(s.encode()))  # noqa: S301
-    return f(*args, **kwargs)
+    ensure_plugins_loaded()
+    _ensure_default_callables_registered()
+    decoded = base64.b64decode(s.encode())
+    data = json.loads(decoded.decode())
+
+    def _is_serialized(obj: Any) -> bool:
+        return isinstance(obj, dict) and {"callable", "args", "kwargs"}.issubset(
+            obj.keys()
+        )
+
+    def _reconstruct(obj: Any, nested: bool = False) -> Any:
+        if not _is_serialized(obj):
+            return obj
+        callable_name: str = obj["callable"]
+        args: list[Any] = obj["args"]
+        kwargs: dict[str, Any] = obj["kwargs"]
+        # Recurse only inside kwargs because serialize() only nests through kwargs
+        for k, v in list(kwargs.items()):
+            if _is_serialized(v):
+                kwargs[k] = _reconstruct(v, True)
+        callable_obj = CallableRegistry.get(callable_name)
+        if nested:
+            return (callable_obj, args, kwargs)
+        # Otherwise instantiate
+        return callable_obj(*args, **kwargs)
+
+    if not _is_serialized(data):
+        raise ValueError("Invalid serialized data format")
+    return _reconstruct(data, False)
+
+
+class _DefaultsState:
+    registered = False
+
+
+def _ensure_default_callables_registered() -> None:
+    if _DefaultsState.registered:
+        return
+
+    from datachain.data_storage.sqlite import (
+        SQLiteDatabaseEngine,
+        SQLiteMetastore,
+        SQLiteWarehouse,
+    )
+
+    # Register (idempotent by name overwrite is fine) using class-level
+    # serialization names to avoid hard-coded literals here.
+    CallableRegistry.register(
+        SQLiteDatabaseEngine.from_db_file,
+        SQLiteDatabaseEngine.serialize_callable_name(),
+    )
+    CallableRegistry.register(
+        SQLiteMetastore.init_after_clone,
+        SQLiteMetastore.serialize_callable_name(),
+    )
+    CallableRegistry.register(
+        SQLiteWarehouse.init_after_clone,
+        SQLiteWarehouse.serialize_callable_name(),
+    )
+
+    _DefaultsState.registered = True

datachain/data_storage/sqlite.py

@@ -201,10 +201,14 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         """
         return (
             SQLiteDatabaseEngine.from_db_file,
-            [self.db_file],
+            [str(self.db_file)],
             {},
         )
 
+    @classmethod
+    def serialize_callable_name(cls) -> str:
+        return "sqlite.from_db_file"
+
     def _reconnect(self) -> None:
         if not self.is_closed:
             raise RuntimeError("Cannot reconnect on still-open DB!")
@@ -403,6 +407,10 @@ class SQLiteMetastore(AbstractDBMetastore):
             },
         )
 
+    @classmethod
+    def serialize_callable_name(cls) -> str:
+        return "sqlite.metastore.init_after_clone"
+
     @classmethod
     def init_after_clone(
         cls,
@@ -610,6 +618,10 @@ class SQLiteWarehouse(AbstractWarehouse):
             {"db_clone_params": self.db.clone_params()},
         )
 
+    @classmethod
+    def serialize_callable_name(cls) -> str:
+        return "sqlite.warehouse.init_after_clone"
+
     @classmethod
     def init_after_clone(
         cls,

datachain/lib/dc/datachain.py

@@ -1963,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.
@@ -1984,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(
@@ -2006,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/plugins.py

@@ -0,0 +1,30 @@
+"""Plugin loader for DataChain callables.
+
+Discovers and invokes entry points in the group "datachain.callables" once
+per process. This enables external packages (e.g., Studio) to register
+their callables with the serializer registry without explicit imports.
+"""
+
+from importlib import metadata as importlib_metadata
+
+_plugins_loaded = False
+
+
+def ensure_plugins_loaded() -> None:
+    global _plugins_loaded  # noqa: PLW0603
+    if _plugins_loaded:
+        return
+
+    # Compatible across importlib.metadata versions
+    eps_obj = importlib_metadata.entry_points()
+    if hasattr(eps_obj, "select"):
+        eps_list = eps_obj.select(group="datachain.callables")
+    else:
+        # Compatibility for older versions of importlib_metadata, Python 3.9
+        eps_list = eps_obj.get("datachain.callables", [])  # type: ignore[attr-defined]
+
+    for ep in eps_list:
+        func = ep.load()
+        func()
+
+    _plugins_loaded = True

{datachain-0.34.0.dist-info → datachain-0.34.2.dist-info}/METADATA

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

{datachain-0.34.0.dist-info → datachain-0.34.2.dist-info}/RECORD

@@ -14,6 +14,7 @@ datachain/namespace.py,sha256=sgIF90KEaC_VlMFivDIJiFz8RUsTftMxW4kOUTyxo3A,2356
 datachain/node.py,sha256=KWDT0ClYXB7FYI-QOvzAa-UDkLJErUI2eWm5FBteYuU,5577
 datachain/nodes_fetcher.py,sha256=_wgaKyqEjkqdwJ_Hj6D8vUYz7hnU7g6xhm0H6ZnYxmE,1095
 datachain/nodes_thread_pool.py,sha256=mdo0s-VybuSZkRUARcUO4Tjh8KFfZr9foHqmupx2SmM,3989
+datachain/plugins.py,sha256=eWOeKg1uBZYEJND8s4D8eZ9b2oEHGN9hahkQDMR45Jc,932
 datachain/progress.py,sha256=lRzxoYP4Qv2XBwD78sOkmYRzHFpZ2ExVNJF8wAeICtY,770
 datachain/project.py,sha256=90D4GpJSA3t0fayYZbzrL3sk4U7EJhQo8psnWvdI7_o,2280
 datachain/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -23,9 +24,9 @@ datachain/studio.py,sha256=IS8o4BZnhUo73Bd8m4CJxFc5utdmh2miIs25WswkFBA,15283
 datachain/telemetry.py,sha256=0A4IOPPp9VlP5pyW9eBfaTK3YhHGzHl7dQudQjUAx9A,994
 datachain/utils.py,sha256=yW-Df5R6npqcqlNZMlBRBwyhUFmXpl9sQipPmy9HfQU,15797
 datachain/catalog/__init__.py,sha256=9NBaywvAOaXdkyqiHjbBEiXs7JImR1OJsY9r8D5Q16g,403
-datachain/catalog/catalog.py,sha256=oI4YBuuOJGVx_Fp1cDoFb56lPV7Or27ZquzR8oM1m3Y,69133
+datachain/catalog/catalog.py,sha256=a1AN6eDHWWzII1wi46T_1JvTsW1AeMudwR_6sVQ4f7I,67588
 datachain/catalog/datasource.py,sha256=IkGMh0Ttg6Q-9DWfU_H05WUnZepbGa28HYleECi6K7I,1353
-datachain/catalog/loader.py,sha256=53VnuSRkt_CO9RdlHWkzQsPF55qMxcXvEm3ecsZREw8,6150
+datachain/catalog/loader.py,sha256=H25cESk72rzs_oAP22jIbthVHunslQbdr63CvV54Pko,6260
 datachain/cli/__init__.py,sha256=so3WxEQF03KdGvjav15Sw7a6-lriiE24uDSGbBDBp8o,8298
 datachain/cli/utils.py,sha256=wrLnAh7Wx8O_ojZE8AE4Lxn5WoxHbOj7as8NWlLAA74,3036
 datachain/cli/commands/__init__.py,sha256=zp3bYIioO60x_X04A4-IpZqSYVnpwOa1AdERQaRlIhI,493
@@ -54,8 +55,8 @@ datachain/data_storage/db_engine.py,sha256=n8ojCbvVMPY2e3SG8fUaaD0b9GkVfpl_Naa_6
 datachain/data_storage/job.py,sha256=NGFhXg0C0zRFTaF6ccjXZJT4xI4_gUr1WcxTLK6WYDE,448
 datachain/data_storage/metastore.py,sha256=9Wd0MfdVrdpgvFXOddUvyz61MnoRDipv0-A38aRsqzw,61021
 datachain/data_storage/schema.py,sha256=DmxxXjNIsXib9gj5jcrb1CVjGzHf7HZLOehs1RmuiMA,9891
-datachain/data_storage/serializer.py,sha256=6G2YtOFqqDzJf1KbvZraKGXl2XHZyVml2krunWUum5o,927
-datachain/data_storage/sqlite.py,sha256=Z6KlFk7hWoXBbjzxfk2NuIBecqP86AJzp5iEE2W4yw0,30603
+datachain/data_storage/serializer.py,sha256=oL8i8smyAeVUyDepk8Xhf3lFOGOEHMoZjA5GdFzvfGI,3862
+datachain/data_storage/sqlite.py,sha256=YNHXPdJeTEoWfhZYb5fsLf1CIjiEhB7VG4OgQzDrWVU,30936
 datachain/data_storage/warehouse.py,sha256=7jc69CtWdfQlc_9WbJ5l6yQooarpLFBrDk4fY-svi_0,32783
 datachain/diff/__init__.py,sha256=v03JfMxH1VvwFl3rniedS4YWs6EXSfaLCULJTKNECE4,9603
 datachain/fs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -107,7 +108,7 @@ datachain/lib/convert/values_to_tuples.py,sha256=j5yZMrVUH6W7b-7yUvdCTGI7JCUAYUO
 datachain/lib/dc/__init__.py,sha256=UrUzmDH6YyVl8fxM5iXTSFtl5DZTUzEYm1MaazK4vdQ,900
 datachain/lib/dc/csv.py,sha256=wUsDPpLD4lts92yn0gejZHqTv8qQBbv8JYRwiIepj0o,4471
 datachain/lib/dc/database.py,sha256=sTpos1rE4BS5BTzzixykhWIO2JxVYKH1GTRncdpu4dU,14716
-datachain/lib/dc/datachain.py,sha256=uUAPchtNXyJo1tzFd3z1MLWhVC2dzO2ZjhTS0naqXiE,104032
+datachain/lib/dc/datachain.py,sha256=Xh7Hwpvow_3QHPhsPSpP99HDKlwcJOpZEZJUNa_Ex9c,104396
 datachain/lib/dc/datasets.py,sha256=pVRcrVEPVPHMf8sLqqhjXbilB3QuUqKE-byvZ-XlJNE,15347
 datachain/lib/dc/hf.py,sha256=B7pubDQTDmth9uILXyhpQNtOAT3UOLjR-peU__tpypk,2884
 datachain/lib/dc/json.py,sha256=-vJ-pUpp2JxK4_vOfznE09FIoEOrvCwoIZSLxM6pjmY,2742
@@ -164,9 +165,9 @@ datachain/sql/sqlite/vector.py,sha256=ncW4eu2FlJhrP_CIpsvtkUabZlQdl2D5Lgwy_cbfqR
 datachain/toolkit/__init__.py,sha256=eQ58Q5Yf_Fgv1ZG0IO5dpB4jmP90rk8YxUWmPc1M2Bo,68
 datachain/toolkit/split.py,sha256=ktGWzY4kyzjWyR86dhvzw-Zhl0lVk_LOX3NciTac6qo,2914
 datachain/torch/__init__.py,sha256=gIS74PoEPy4TB3X6vx9nLO0Y3sLJzsA8ckn8pRWihJM,579
-datachain-0.34.0.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.34.0.dist-info/METADATA,sha256=YBmM_daqadosEKHBY-QLxSRxYn55XuhB0S0tfeEfzts,13655
-datachain-0.34.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-datachain-0.34.0.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.34.0.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.34.0.dist-info/RECORD,,
+datachain-0.34.2.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.34.2.dist-info/METADATA,sha256=p-mulDC4TJ2QOJr2peiHCygfiVP1bwwdubi-fyfLQkg,13655
+datachain-0.34.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+datachain-0.34.2.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.34.2.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.34.2.dist-info/RECORD,,