boti-dask 1.0.0__py3-none-any.whl

boti_dask/__init__.py

@@ -0,0 +1,62 @@
+"""Dask runtime and resilience utilities for the Boti ecosystem."""
+
+from boti_dask.diagnostics import (
+    UniqueValuesExtractor,
+    describe_frame,
+    diagnostics_logger,
+    inspect_graph,
+)
+from boti_dask.resilience import (
+    RECOVERABLE_DASK_ERRORS,
+    async_safe_compute,
+    async_safe_gather,
+    async_safe_head,
+    async_safe_persist,
+    async_safe_wait,
+    dask_is_empty,
+    dask_is_probably_empty,
+    safe_compute,
+    safe_gather,
+    safe_head,
+    safe_persist,
+    safe_wait,
+)
+from boti_dask.session import (
+    DaskSession,
+    DaskSessionSettings,
+    apply_recommended_dask_config,
+    current_client_summary,
+    dask_session,
+    dask_session_from_env_prefix,
+    describe_client,
+    recommended_dask_config,
+)
+
+__all__ = [
+    "DaskSession",
+    "DaskSessionSettings",
+    "RECOVERABLE_DASK_ERRORS",
+    "UniqueValuesExtractor",
+    "apply_recommended_dask_config",
+    "async_safe_compute",
+    "async_safe_gather",
+    "async_safe_head",
+    "async_safe_persist",
+    "async_safe_wait",
+    "current_client_summary",
+    "dask_is_empty",
+    "dask_is_probably_empty",
+    "dask_session",
+    "dask_session_from_env_prefix",
+    "describe_frame",
+    "describe_client",
+    "diagnostics_logger",
+    "inspect_graph",
+    "recommended_dask_config",
+    "safe_compute",
+    "safe_gather",
+    "safe_head",
+    "safe_persist",
+    "safe_wait",
+]
+

boti_dask/_internal.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from typing import Any
+
+import dask.dataframe as dd
+
+
+def _log(logger: Any | None, level: str, message: str) -> None:
+    if logger is None:
+        return
+    log_fn = getattr(logger, level, None)
+    if callable(log_fn):
+        log_fn(message)
+
+
+def _is_dask_dataframe_like(obj: Any) -> bool:
+    return isinstance(obj, (dd.DataFrame, dd.Series)) or hasattr(obj, "_meta")
+
+
+def _is_running_client(client: Any | None) -> bool:
+    if client is None:
+        return False
+    return getattr(client, "status", "running") in {"running", "started"}

boti_dask/diagnostics.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+import dask.dataframe as dd
+import pandas as pd
+
+try:
+    import pyarrow as pa
+except ImportError:
+    pa = None
+
+try:
+    import polars as pl
+except ImportError:
+    pl = None
+
+from boti.core.logger import Logger
+
+from ._internal import _log
+
+
+def _has_dask_graph(obj: Any) -> bool:
+    return hasattr(obj, "__dask_graph__")
+
+
+def inspect_graph(obj: Any, *, logger: Any | None = None) -> dict[str, Any]:
+    """Return compact Dask graph metrics for diagnostics and dry-run usage."""
+    if not _has_dask_graph(obj):
+        metrics = {"type": type(obj).__name__, "is_dask": False}
+        _log(logger, "info", f"Dask graph inspection metrics={metrics}")
+        return metrics
+
+    graph = obj.__dask_graph__()
+    graph_layers: int | None = None
+    if hasattr(getattr(obj, "dask", None), "layers"):
+        graph_layers = len(obj.dask.layers)
+    elif hasattr(graph, "layers"):
+        graph_layers = len(graph.layers)
+    metrics = {
+        "type": type(obj).__name__,
+        "is_dask": True,
+        "task_count": len(graph) if hasattr(graph, "__len__") else None,
+        "npartitions": getattr(obj, "npartitions", None),
+        "graph_layers": graph_layers,
+    }
+    _log(logger, "info", f"Dask graph inspection metrics={metrics}")
+    return metrics
+
+
+def describe_frame(frame: Any) -> dict[str, Any]:
+    """Return compact frame metrics suitable for runtime diagnostics logs."""
+    if isinstance(frame, dd.DataFrame):
+        graph = frame.__dask_graph__()
+        return {
+            "engine": "dask",
+            "columns": len(frame.columns),
+            "npartitions": frame.npartitions,
+            "graph_tasks": len(graph) if hasattr(graph, "__len__") else None,
+            "graph_layers": len(frame.dask.layers) if hasattr(frame.dask, "layers") else None,
+            "known_divisions": frame.known_divisions,
+        }
+    if isinstance(frame, pd.DataFrame):
+        return {
+            "engine": "pandas",
+            "rows": len(frame.index),
+            "columns": len(frame.columns),
+        }
+    if pa is not None and isinstance(frame, pa.Table):
+        return {
+            "engine": "arrow",
+            "rows": frame.num_rows,
+            "columns": len(frame.column_names),
+        }
+    if pl is not None and isinstance(frame, pl.DataFrame):
+        return {
+            "engine": "polars",
+            "rows": frame.height,
+            "columns": frame.width,
+        }
+    return {"engine": type(frame).__name__}
+
+
+def diagnostics_logger(logger: Any | None, *, name: str) -> Any:
+    """Return provided logger or a :class:`boti.Logger` fallback scoped by *name*."""
+    if logger is not None:
+        return logger
+    return Logger.default_logger(logger_name=name)
+
+
+class UniqueValuesExtractor:
+    """Best-effort unique value extraction for Dask-backed columns."""
+
+    def __init__(self, *, dask_client: Any | None = None, logger: Any | None = None) -> None:
+        self.dask_client = dask_client
+        self.logger = logger
+
+    def _extract_one(self, frame: dd.DataFrame, column: str, limit: int) -> tuple[str, list[Any]]:
+        from .resilience import safe_compute
+
+        if column not in frame.columns:
+            _log(self.logger, "warning", f"Column '{column}' not found for unique extraction.")
+            return column, []
+
+        expression = frame[column].dropna().drop_duplicates()
+        if self.dask_client is not None:
+            values = safe_compute(expression, dask_client=self.dask_client, logger=self.logger)
+        else:
+            values = expression.compute(scheduler="threads")
+        as_series = pd.Series(values)
+        unique_values = as_series.dropna().unique().tolist()
+        if len(unique_values) > limit:
+            _log(
+                self.logger,
+                "warning",
+                f"Unique value extraction for column '{column}' truncated at {limit} items.",
+            )
+            unique_values = unique_values[:limit]
+        return column, unique_values
+
+    async def extract_unique_values(
+        self,
+        frame: dd.DataFrame,
+        *columns: str,
+        limit: int = 100_000,
+    ) -> dict[str, list[Any]]:
+        from .resilience import safe_persist
+
+        if self.dask_client is not None:
+            frame = safe_persist(frame, dask_client=self.dask_client, logger=self.logger)
+        else:
+            frame = frame.persist()
+        _log(self.logger, "info", f"Persisted {frame.npartitions}-partition frame for unique extraction.")
+        pairs = await asyncio.gather(
+            *(asyncio.to_thread(self._extract_one, frame, column, limit) for column in columns)
+        )
+        return dict(pairs)
+
+
+__all__ = [
+    "UniqueValuesExtractor",
+    "describe_frame",
+    "diagnostics_logger",
+    "inspect_graph",
+]