sibi-dst 2025.8.7__py3-none-any.whl → 2025.8.9__py3-none-any.whl

sibi_dst/utils/parquet_saver.py

@@ -1,5 +1,7 @@
 import warnings
 
+from pandas.api.types import is_period_dtype, is_bool_dtype, is_string_dtype
+import pandas as pd
 import dask.dataframe as dd
 import pyarrow as pa
 
@@ -15,6 +17,7 @@ class ParquetSaver(ManagedResource):
 
     Assumes `df_result` is a Dask DataFrame.
     """
+    logger_extra = {"sibi_dst_component": __name__}
 
     def __init__(
         self,
@@ -32,6 +35,10 @@ class ParquetSaver(ManagedResource):
         if "://" in self.parquet_storage_path:
             self.protocol = self.parquet_storage_path.split(":", 1)[0]
 
+        self.persist = kwargs.get("persist",True)
+        self.write_index = kwargs.get("write_index", False)
+        self.write_metadata_file = kwargs.get("write_metadata_file", True)
+
     def save_to_parquet(self, output_directory_name: str = "default_output", overwrite: bool = True):
         """
         Saves the Dask DataFrame to a Parquet dataset.
@@ -42,17 +49,19 @@ class ParquetSaver(ManagedResource):
         full_path = f"{self.parquet_storage_path}/{output_directory_name}"
 
         if overwrite and self.fs and self.fs.exists(full_path):
-            self.logger.info(f"Overwrite is True, clearing destination path: {full_path}")
+            self.logger.info(f"Overwrite is True, clearing destination path: {full_path}", extra=self.logger_extra)
             self._clear_directory_safely(full_path)
 
         # Ensure the base directory exists after clearing
         self.fs.mkdirs(full_path, exist_ok=True)
 
         schema = self._define_schema()
-        self.logger.info(f"Saving DataFrame to Parquet dataset at: {full_path}")
+        self.logger.info(f"Saving DataFrame to Parquet dataset at: {full_path}", extra=self.logger_extra)
+        # 1) Normalize to declared schema (fixes bool→string, Period→string, etc.)
+        ddf = self._coerce_ddf_to_schema(self.df_result, schema)
 
-        # persist then write (lets the graph be shared if the caller reuses it)
-        ddf = self.df_result.persist()
+        # 2) Persist after coercion so all partitions share the coerced dtypes
+        ddf = ddf.persist() if self.persist else ddf
 
         try:
             ddf.to_parquet(
@@ -61,11 +70,12 @@ class ParquetSaver(ManagedResource):
                 schema=schema,
                 overwrite=False,         # we've handled deletion already
                 filesystem=self.fs,
-                write_index=False,
+                write_index=self.write_index,  # whether to write the index
+                write_metadata_file=self.write_metadata_file,  # write _metadata for easier reading later
             )
-            self.logger.info(f"Successfully saved Parquet dataset to: {full_path}")
+            self.logger.info(f"Successfully saved Parquet dataset to: {full_path}", extra=self.logger_extra)
         except Exception as e:
-            self.logger.error(f"Failed to save Parquet dataset to {full_path}: {e}")
+            self.logger.error(f"Failed to save Parquet dataset to {full_path}: {e}", extra=self.logger_extra)
             raise
 
     def _clear_directory_safely(self, directory: str):
@@ -91,7 +101,7 @@ class ParquetSaver(ManagedResource):
                     else:
                         self.fs.rm(path, recursive=False)
                 except Exception as e:
-                    self.logger.warning(f"Failed to delete '{path}': {e}")
+                    self.logger.warning(f"Failed to delete '{path}': {e}", extra=self.logger_extra)
             # remove the (now empty) directory if present
             try:
                 self.fs.rm(directory, recursive=False)
@@ -120,4 +130,95 @@ class ParquetSaver(ManagedResource):
             pa.field(c, pandas_dtype_to_pa.get(str(d), pa.string()))
             for c, d in self.df_result.dtypes.items()
         ]
-        return pa.schema(fields)
+        return pa.schema(fields)
+
+
+    def _coerce_ddf_to_schema(self, ddf: dd.DataFrame, schema: pa.Schema) -> dd.DataFrame:
+        """
+        Coerce Dask DataFrame columns to match the provided PyArrow schema.
+        - Ensures cross-partition consistency.
+        - Converts troublesome dtypes (Period, mixed object/bool) to the declared type.
+        """
+        # Build a map: name -> target kind
+        target = {field.name: field.type for field in schema}
+
+        def _coerce_partition(pdf: pd.DataFrame) -> pd.DataFrame:
+            for col, typ in target.items():
+                if col not in pdf.columns:
+                    continue
+
+                pa_type = typ
+
+                # String targets
+                if pa.types.is_string(pa_type) or pa.types.is_large_string(pa_type):
+                    # Convert Period or any dtype to string with NA-preservation
+                    s = pdf[col]
+                    if is_period_dtype(s):
+                        pdf[col] = s.astype(str)
+                    elif not is_string_dtype(s):
+                        # astype("string") keeps NA; str(s) can produce "NaT" strings
+                        try:
+                            pdf[col] = s.astype("string")
+                        except Exception:
+                            pdf[col] = s.astype(str).astype("string")
+                    continue
+
+                # Boolean targets
+                if pa.types.is_boolean(pa_type):
+                    s = pdf[col]
+                    # Allow object/bool mixtures; coerce via pandas nullable boolean then to bool
+                    try:
+                        pdf[col] = s.astype("boolean").astype(bool)
+                    except Exception:
+                        pdf[col] = s.astype(bool)
+                    continue
+
+                # Integer targets
+                if pa.types.is_integer(pa_type):
+                    s = pdf[col]
+                    # Go through pandas nullable Int64 to preserve NA, then to int64 if clean
+                    s2 = pd.to_numeric(s, errors="coerce").astype("Int64")
+                    # If there are no nulls, downcast to numpy int64 for speed
+                    if not s2.isna().any():
+                        s2 = s2.astype("int64")
+                    pdf[col] = s2
+                    continue
+
+                # Floating targets
+                if pa.types.is_floating(pa_type):
+                    pdf[col] = pd.to_numeric(pdf[col], errors="coerce").astype("float64")
+                    continue
+
+                # Timestamp[ns] (optionally with tz)
+                if pa.types.is_timestamp(pa_type):
+                    # If tz in Arrow type, you may want to localize; here we just ensure ns
+                    pdf[col] = pd.to_datetime(pdf[col], errors="coerce")
+                    continue
+
+                # Fallback: leave as-is
+            return pdf
+
+        # Provide a meta with target dtypes to avoid meta mismatch warnings
+        meta = {}
+        for name, typ in target.items():
+            # Rough meta mapping; Arrow large_string vs string both → 'string'
+            if pa.types.is_string(typ) or pa.types.is_large_string(typ):
+                meta[name] = pd.Series([], dtype="string")
+            elif pa.types.is_boolean(typ):
+                meta[name] = pd.Series([], dtype="bool")
+            elif pa.types.is_integer(typ):
+                meta[name] = pd.Series([], dtype="Int64")  # nullable int
+            elif pa.types.is_floating(typ):
+                meta[name] = pd.Series([], dtype="float64")
+            elif pa.types.is_timestamp(typ):
+                meta[name] = pd.Series([], dtype="datetime64[ns]")
+            else:
+                meta[name] = pd.Series([], dtype="object")
+
+        # Start from current meta and update known columns
+        new_meta = ddf._meta.copy()
+        for k, v in meta.items():
+            if k in new_meta.columns:
+                new_meta[k] = v
+
+        return ddf.map_partitions(_coerce_partition, meta=new_meta)

sibi_dst/utils/progress/__init__.py

@@ -0,0 +1,5 @@
+from .sse_runner import SSERunner, _as_sse_msg
+
+__all__ = [
+    "SSERunner", "_as_sse_msg"
+]

sibi_dst/utils/progress/jobs.py

@@ -0,0 +1,82 @@
+# jobs.py
+import asyncio, json, uuid
+import contextlib
+import os
+
+import redis.asyncio as redis
+from fastapi import APIRouter
+from sse_starlette.sse import EventSourceResponse
+host = os.getenv("REDIS_HOST", "0.0.0.0")
+port = int(os.getenv("REDIS_PORT", 6379))
+db = int(os.getenv("REDIS_DB", 0))
+router = APIRouter(prefix="/jobs", tags=["Jobs"])
+r = redis.Redis(host=host, port=port, db=db, decode_responses=True)  # strings for pubsub
+
+CHANNEL = lambda job_id: f"job:{job_id}:events"
+KEY_STATUS = lambda job_id: f"job:{job_id}:status"   # JSON blob with state/progress
+KEY_RESULT = lambda job_id: f"job:{job_id}:result"   # final payload
+
+async def publish(job_id: str, event: str, data: dict):
+    msg = json.dumps({"event": event, "data": data})
+    await r.publish(CHANNEL(job_id), msg)
+    # store last status
+    await r.set(KEY_STATUS(job_id), json.dumps({"event": event, "data": data}))
+
+# ---- Worker entry (can live in a separate process) ----
+async def run_job(job_id: str):
+    try:
+        await publish(job_id, "progress", {"message": "Initializing..."})
+        # ... do actual work, emit more progress
+        await asyncio.sleep(0.2)
+        # compute result
+        result = [{"id": 1, "ok": True}]
+        await r.set(KEY_RESULT(job_id), json.dumps(result), ex=3600)
+        await publish(job_id, "complete", {"records": len(result)})
+    except Exception as e:
+        await publish(job_id, "error", {"detail": str(e)})
+
+# ---- API ----
+@router.post("/start")
+async def start_job():
+    job_id = str(uuid.uuid4())
+    # enqueue: prefer Celery/RQ/etc. For demo we detach a task.
+    asyncio.create_task(run_job(job_id))
+    return {"job_id": job_id}
+
+@router.get("/{job_id}/stream")
+async def stream(job_id: str):
+    pubsub = r.pubsub()
+    await pubsub.subscribe(CHANNEL(job_id))
+
+    async def gen():
+        try:
+            # emit latest known status immediately, if any
+            if (s := await r.get(KEY_STATUS(job_id))):
+                yield {"event": "progress", "data": s}
+            while True:
+                msg = await pubsub.get_message(ignore_subscribe_messages=True, timeout=30.0)
+                if msg and msg["type"] == "message":
+                    payload = msg["data"]  # already a JSON string
+                    yield {"event": "message", "data": payload}
+                await asyncio.sleep(0.01)
+        finally:
+            with contextlib.suppress(Exception):
+                await pubsub.unsubscribe(CHANNEL(job_id))
+                await pubsub.close()
+
+    return EventSourceResponse(
+        gen(),
+        ping=15,
+        headers={"Cache-Control": "no-cache, no-transform", "X-Accel-Buffering": "no"},
+    )
+
+@router.get("/{job_id}/status")
+async def status(job_id: str):
+    s = await r.get(KEY_STATUS(job_id))
+    done = await r.exists(KEY_RESULT(job_id))
+    return {"job_id": job_id, "status": json.loads(s) if s else None, "done": bool(done)}
+
+@router.get("/{job_id}/result")
+async def result(job_id: str):
+    data = await r.get(KEY_RESULT(job_id))
+    return {"job_id": job_id, "result": json.loads(data) if data else None}

sibi_dst/utils/progress/sse_runner.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+import asyncio, contextlib, inspect, json
+from typing import Any, Awaitable, Callable, Dict, Optional, Union
+from fastapi import Request
+from sse_starlette.sse import EventSourceResponse
+from sibi_dst.utils import Logger
+
+Payload = Union[str, bytes, dict, list, None]
+Task2 = Callable[[asyncio.Queue, str], Awaitable[Payload]]
+Task3 = Callable[[asyncio.Queue, str, Dict[str, Any]], Awaitable[Payload]]
+TaskFn = Union[Task2, Task3]
+
+def _as_sse_msg(event: str, data: Any) -> dict:
+    return {"event": event, "data": json.dumps(data) if not isinstance(data, (str, bytes)) else data}
+
+class SSERunner:
+    def __init__(self, *, task: TaskFn, logger: Logger, ping: int = 15,
+                 headers: Optional[dict] = None, auto_complete: bool = True) -> None:
+        self.task = task
+        self.logger = logger
+        self.ping = ping
+        self.headers = headers or {"Cache-Control": "no-cache, no-transform", "X-Accel-Buffering": "no"}
+        self.auto_complete = auto_complete
+        self._expects_ctx = len(inspect.signature(task).parameters) >= 3
+
+    async def _call_task(self, queue: asyncio.Queue, task_id: str, ctx: Dict[str, Any]) -> Payload:
+        if self._expects_ctx:
+            return await self.task(queue, task_id, ctx)  # type: ignore[misc]
+        return await self.task(queue, task_id)           # type: ignore[misc]
+
+    async def _worker(self, queue: asyncio.Queue, task_id: str, ctx: Dict[str, Any]) -> None:
+        self.logger.info(f"SSE {task_id}: start")
+        try:
+            await queue.put(_as_sse_msg("progress", {"message": "Task started"}))
+            payload = await self._call_task(queue, task_id, ctx)
+            if self.auto_complete:
+                final = payload if payload is not None else {"status": "complete"}
+                await queue.put(_as_sse_msg("complete", final))
+            self.logger.info(f"SSE {task_id}: complete")
+        except asyncio.CancelledError:
+            raise
+        except Exception as e:
+            self.logger.error(f"SSE {task_id} failed: {e}", exc_info=True)
+            await queue.put(_as_sse_msg("error", {"detail": str(e)}))
+        finally:
+            await queue.put(None)
+
+    def endpoint(self):
+        async def handler(request: Request):  # <-- only Request
+            queue: asyncio.Queue = asyncio.Queue()
+            task_id = str(asyncio.get_running_loop().time()).replace(".", "")
+
+            ctx: Dict[str, Any] = {
+                "path": dict(request.path_params),          # <-- pull path params here
+                "query": dict(request.query_params),
+                "method": request.method,
+            }
+            if request.headers.get("content-type", "").startswith("application/json"):
+                try:
+                    ctx["body"] = await request.json()
+                except Exception:
+                    ctx["body"] = None
+
+            worker = asyncio.create_task(self._worker(queue, task_id, ctx))
+
+            async def gen():
+                try:
+                    while True:
+                        msg = await queue.get()
+                        if msg is None:
+                            break
+                        yield msg
+                finally:
+                    if not worker.done():
+                        worker.cancel()
+                        with contextlib.suppress(Exception):
+                            await worker
+
+            return EventSourceResponse(gen(), ping=self.ping, headers=self.headers)
+        return handler
+
+__all__ = ["SSERunner", "_as_sse_msg"]

sibi_dst/utils/storage_hive.py

@@ -1,4 +1,7 @@
 from __future__ import annotations
+
+import asyncio
+
 import pandas as pd
 import dask.dataframe as dd
 from typing import Iterable, Optional, List, Tuple, Union
@@ -192,4 +195,38 @@ class HiveDatePartitionedStore:
             clauses.append([("yyyy","==",y)])
         for m in range(1, eM):
             clauses.append([("yyyy","==",eY),("mm","==",m)])
-        return clauses
+        return clauses
+
+    async def write_async(
+            self,
+            df: dd.DataFrame,
+            *,
+            repartition: int | None = None,
+            overwrite: bool = False,
+            timeout: float | None = None,
+    ) -> None:
+        async def _run():
+            return await asyncio.to_thread(self.write, df, repartition=repartition, overwrite=overwrite)
+
+        return await (asyncio.wait_for(_run(), timeout) if timeout else _run())
+
+    async def read_range_async(
+            self,
+            start, end, *, columns: Iterable[str] | None = None, timeout: float | None = None
+    ) -> dd.DataFrame:
+        async def _run():
+            return await asyncio.to_thread(self.read_range, start, end, columns=columns)
+
+        return await (asyncio.wait_for(_run(), timeout) if timeout else _run())
+
+    async def read_month_async(self, year: int, month: int, *, columns=None, timeout: float | None = None):
+        async def _run():
+            return await asyncio.to_thread(self.read_month, year, month, columns=columns)
+
+        return await (asyncio.wait_for(_run(), timeout) if timeout else _run())
+
+    async def read_day_async(self, year: int, month: int, day: int, *, columns=None, timeout: float | None = None):
+        async def _run():
+            return await asyncio.to_thread(self.read_day, year, month, day, columns=columns)
+
+        return await (asyncio.wait_for(_run(), timeout) if timeout else _run())