datachain 0.34.0__tar.gz → 0.34.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.34.0
+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.34.0 → 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.34.0 → datachain-0.34.1}/src/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-0.34.0 → datachain-0.34.1}/src/datachain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.34.0
+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.34.0 → 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

{datachain-0.34.0 → datachain-0.34.1}/tests/unit/lib/test_datachain.py

@@ -3,6 +3,7 @@ import json
 import math
 import os
 import re
+from collections import Counter
 from collections.abc import Generator, Iterator
 from unittest.mock import ANY, patch
 
@@ -564,6 +565,32 @@ def test_from_features_simple_types_in_memory():
     assert df["odds"].tolist() == values
 
 
+@skip_if_not_sqlite
+def test_to_pandas_as_object_preserves_none(test_session):
+    timestamp = datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc)
+    chain = dc.read_values(
+        id=[1, None],
+        value=[3.14, None],
+        ts=[timestamp, None],
+        session=test_session,
+    )
+
+    df_default = chain.to_pandas()
+    assert df_default["id"].dtype != object
+    assert df_default["value"].dtype != object
+    assert df_default["id"].isna().sum() == 1
+    assert df_default["value"].isna().sum() == 1
+    assert pd.isna(df_default.loc[df_default["id"].isna(), "ts"]).all()
+
+    df_object = chain.to_pandas(as_object=True)
+    assert df_object["id"].dtype == object
+    assert df_object["value"].dtype == object
+    assert df_object["ts"].dtype == object
+    assert Counter(df_object["id"].tolist()) == Counter([1, None])
+    assert Counter(df_object["value"].tolist()) == Counter([3.14, None])
+    assert Counter(df_object["ts"].tolist()) == Counter([timestamp, None])
+
+
 def test_from_features_more_simple_types(test_session):
     ds = dc.read_values(
         t1=features,

{datachain-0.34.0 → datachain-0.34.1}/tests/unit/test_query.py

@@ -42,31 +42,12 @@ def test_args(catalog, mock_popen):
     mock_popen.assert_called_once_with(["mypython", "-c", "pass"], env=expected_env)
 
 
-def test_capture_stdout(catalog, mock_popen):
-    mock_popen.stdout = io.BytesIO(b"Hello, World!\rLorem Ipsum\nDolor Sit Amet\nconse")
-    stdout = []
-
-    catalog.query("pass", stdout_callback=stdout.append)
-    assert stdout == ["Hello, World!\r", "Lorem Ipsum\n", "Dolor Sit Amet\n", "conse"]
-
-
-def test_capture_stderr(catalog, mock_popen):
-    mock_popen.stderr = io.BytesIO(b"Hello, World!\rLorem Ipsum\nDolor Sit Amet\nconse")
-    stderr = []
-
-    catalog.query("pass", stderr_callback=stderr.append)
-    assert stderr == ["Hello, World!\r", "Lorem Ipsum\n", "Dolor Sit Amet\n", "conse"]
-
-
 def test_capture_output(catalog, mock_popen):
     mock_popen.stdout = io.BytesIO(b"Hello, World!\rLorem Ipsum\nDolor Sit Amet\nconse")
-    mock_popen.stderr = io.BytesIO(b"foo\nbar")
-    stdout = []
-    stderr = []
+    lines = []
 
-    catalog.query("pass", stdout_callback=stdout.append, stderr_callback=stderr.append)
-    assert stdout == ["Hello, World!\r", "Lorem Ipsum\n", "Dolor Sit Amet\n", "conse"]
-    assert stderr == ["foo\n", "bar"]
+    catalog.query("pass", capture_output=True, output_hook=lines.append)
+    assert lines == ["Hello, World!\r", "Lorem Ipsum\n", "Dolor Sit Amet\n", "conse"]
 
 
 def test_canceled_by_user(catalog, mock_popen):