sibi-flux 2025.12.0__py3-none-any.whl

sibi_flux/utils/filepath_generator/_filepath_generator.py

@@ -0,0 +1,185 @@
+import datetime as dt
+import re
+from typing import List, Optional, Iterable
+
+import fsspec
+import pyarrow.dataset as ds
+from fsspec.utils import infer_storage_options
+from sibi_flux.logger import Logger
+
+
+class FilePathGenerator:
+    """
+    Scans Hive-partitioned directories (key=value) and returns paths for a specific date range.
+    Supports single-key date partitions (e.g. partition_date=YYYY-MM-DD) or
+    composite hierarchies (year=YYYY/month=MM/day=DD).
+    """
+
+    def __init__(
+        self,
+        base_path: str = "",
+        *,
+        fs=None,
+        date_partition_key: str = "partition_date",
+        logger: Optional[Logger] = None,
+        debug: bool = False,
+        storage_options: Optional[dict] = None,
+        exclude_patterns: Optional[Iterable[str]] = None,
+        file_extension: str = "parquet",
+    ):
+        """
+        Args:
+            base_path: S3/Local path to the root of the dataset.
+            fs: Optional fsspec filesystem instance.
+            date_partition_key: The hive column name used for filtering.
+                                Defaults to "partition_date".
+                                If the dataset uses standard year/month/day hierarchy,
+                                this is ignored automatically.
+        """
+        self.logger = logger or Logger.default_logger(
+            logger_name=self.__class__.__name__
+        )
+        self.debug = debug
+        self.storage_options = storage_options or {}
+        self.file_extension = file_extension.lstrip(".")
+        self._compiled_exclusions = [re.compile(p) for p in (exclude_patterns or [])]
+        self.date_partition_key = date_partition_key
+
+        # 1. Setup Filesystem (Protocol Stripping)
+        if fs:
+            self.fs = fs
+            self.base_path = infer_storage_options(base_path)["path"]
+        else:
+            self.fs, self.base_path = fsspec.core.url_to_fs(
+                base_path, **self.storage_options
+            )
+
+        # 2. Strict Hive Partitioning
+        # This automatically parses "key=value" directory names.
+        self.partitioning = ds.partitioning(flavor="hive")
+
+        if self.debug:
+            self.logger.debug(
+                f"Init: path={self.base_path}, date_key={self.date_partition_key}"
+            )
+
+    def generate_file_paths(
+        self, start_date, end_date, engine: str = "dask"
+    ) -> List[str]:
+        sd = self._to_date(start_date)
+        ed = self._to_date(end_date)
+        if sd > ed:
+            sd, ed = ed, sd
+
+        # Create Dataset
+        try:
+            dataset = ds.dataset(
+                source=self.base_path,
+                filesystem=self.fs,
+                format=self.file_extension,
+                partitioning=self.partitioning,
+            )
+        except Exception as e:
+            if self.debug:
+                self.logger.error(f"Failed to load dataset: {e}")
+            return []
+
+        # 3. Dynamic Filter Logic
+        schema_names = set(dataset.schema.names)
+        filter_expr = None
+        filter_mode = "unknown"
+
+        # Case A: User-specified single column (e.g., partition_date=2024-01-01)
+        if self.date_partition_key in schema_names:
+            filter_mode = "single_key"
+            # We filter treating the key as a string comparable to YYYY-MM-DD
+            filter_expr = (ds.field(self.date_partition_key) >= str(sd)) & (
+                ds.field(self.date_partition_key) <= str(ed)
+            )
+
+        # Case B: Standard Composite Hierarchy (year=..., month=..., day=...)
+        # We auto-fallback to this if the single key isn't found
+        elif {"year", "month", "day"}.issubset(schema_names):
+            filter_mode = "composite_ymd"
+            # Broad filter on year first for performance
+            filter_expr = (ds.field("year") >= sd.year) & (ds.field("year") <= ed.year)
+
+        else:
+            self.logger.warning(
+                f"Could not find partition key '{self.date_partition_key}' or standard year/month/day keys "
+                f"in dataset schema: {schema_names}. Returning all files (unfiltered)."
+            )
+
+        # 4. Materialize & Refine
+        # We assume PyArrow's filter is "coarse" (partition level), so we double-check checks in Python
+        valid_paths = []
+        fragments = dataset.get_fragments(filter=filter_expr)
+
+        for frag in fragments:
+            if self._is_file_in_range(frag, sd, ed, filter_mode):
+                path = frag.path
+                # Ensure full protocol paths for Pandas/Dask
+                if not path.startswith(self.fs.protocol):
+                    path = self.fs.unstrip_protocol(path)
+
+                if not self._is_excluded(path):
+                    valid_paths.append(path)
+
+        if self.debug:
+            self.logger.debug(
+                f"Found {len(valid_paths)} files via '{filter_mode}' filter."
+            )
+
+        return sorted(valid_paths)
+
+    def _is_file_in_range(self, frag, sd, ed, mode: str) -> bool:
+        """
+        Parses partition keys from fragment metadata and validates exact date range.
+        """
+        try:
+            keys = ds._get_partition_keys(frag.partition_expression)
+
+            if mode == "single_key":
+                val = keys.get(self.date_partition_key)
+                if not val:
+                    return False
+
+                # Handle String vs Date type inference
+                f_date = (
+                    val
+                    if isinstance(val, (dt.date, dt.datetime))
+                    else self._to_date(val)
+                )
+                # Ensure comparison between dates (not datetime vs date)
+                if isinstance(f_date, dt.datetime):
+                    f_date = f_date.date()
+                return sd <= f_date <= ed
+
+            elif mode == "composite_ymd":
+                f_date = dt.date(
+                    int(keys["year"]), int(keys["month"]), int(keys["day"])
+                )
+                return sd <= f_date <= ed
+
+            return True  # If mode unknown, we accept the file (passed by the coarse filter or absence thereof)
+
+        except Exception:
+            # If parsing fails (e.g. malformed folder name), exclude safely
+            return False
+
+    # ------------------------- Helpers -------------------------
+
+    def _is_excluded(self, path: str) -> bool:
+        return any(pat.search(path) for pat in self._compiled_exclusions)
+
+    @staticmethod
+    def _to_date(x) -> dt.date:
+        if isinstance(x, dt.datetime):
+            return x.date()
+        if isinstance(x, dt.date):
+            return x
+        try:
+            return dt.datetime.strptime(str(x), "%Y-%m-%d").date()
+        except ValueError:
+            # Fallback for weird formats if necessary
+            return dt.date.min

sibi_flux/utils/parquet_saver/__init__.py

@@ -0,0 +1,6 @@
+from ._parquet_saver import ParquetSaver, _coerce_partition
+
+__all__ = [
+    "ParquetSaver",
+    "_coerce_partition",
+]

sibi_flux/utils/parquet_saver/_parquet_saver.py

@@ -0,0 +1,436 @@
+from __future__ import annotations
+
+import warnings
+import os
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+from multiprocessing.pool import ThreadPool
+from typing import Any, Dict, Optional, List
+
+import dask.dataframe as dd
+import pandas as pd
+import pyarrow as pa
+
+from sibi_flux.core import ManagedResource
+from sibi_flux.dask_cluster.core import safe_persist
+from ._write_gatekeeper import get_write_sem
+
+warnings.filterwarnings(
+    "ignore", message="Passing 'overwrite=True' to to_parquet is deprecated"
+)
+
+
+def _coerce_partition(
+    pdf: pd.DataFrame,
+    target: Dict[str, pa.DataType],
+    partition_cols: Optional[List[str]] = None,
+) -> pd.DataFrame:
+    """
+    Applies type conversions to a single pandas partition.
+    """
+    partition_cols = partition_cols or []
+
+    for col, pa_type in target.items():
+        if col not in pdf.columns:
+            # Fix: Ensure column exists if it's in the target schema
+            # Fix: Ensure column exists if it's in the target schema
+            if pa.types.is_string(pa_type):
+                pdf[col] = None
+                # We will cast to string[pyarrow] below
+            elif pa.types.is_integer(pa_type):
+                pdf[col] = pd.NA
+            elif pa.types.is_floating(pa_type):
+                pdf[col] = float("nan")
+            elif pa.types.is_boolean(pa_type):
+                pdf[col] = pd.NA
+            elif pa.types.is_timestamp(pa_type):
+                pdf[col] = pd.NaT
+            else:
+                pdf[col] = None
+
+        try:
+            # --- 0. Partition Column Logic ---
+            if col in partition_cols:
+                try:
+                    if not pd.api.types.is_numeric_dtype(pdf[col]):
+                        # SAFE PARSING: Use format='mixed' here too
+                        temp_series = pd.to_datetime(
+                            pdf[col], errors="raise", utc=True, format="mixed"
+                        )
+                        pdf[col] = temp_series.dt.strftime("%Y-%m-%d").astype(
+                            "string[pyarrow]"
+                        )
+                        continue
+                except (ValueError, TypeError):
+                    pass
+
+                pdf[col] = pdf[col].astype("string[pyarrow]")
+                continue
+
+            # --- 1. Timestamp Coercion ---
+            if pa.types.is_timestamp(pa_type):
+                if not pd.api.types.is_datetime64_any_dtype(pdf[col]):
+                    # CRITICAL FIX: format='mixed' handles naive/aware mix in same column
+                    pdf[col] = pd.to_datetime(
+                        pdf[col], errors="coerce", utc=True, format="mixed"
+                    )
+
+                if pdf[col].dt.tz is None:
+                    pdf[col] = pdf[col].dt.tz_localize("UTC")
+                else:
+                    pdf[col] = pdf[col].dt.tz_convert("UTC")
+
+                pdf[col] = pdf[col].astype("timestamp[ns, tz=UTC][pyarrow]")
+                continue
+
+            # --- 2. Other Types ---
+            current_dtype_str = str(pdf[col].dtype)
+
+            if pa.types.is_string(pa_type) and "string" not in current_dtype_str:
+                pdf[col] = pdf[col].astype("string[pyarrow]")
+
+            elif pa.types.is_boolean(pa_type) and "bool" not in current_dtype_str:
+                if pd.api.types.is_object_dtype(
+                    pdf[col]
+                ) or pd.api.types.is_string_dtype(pdf[col]):
+                    pdf[col] = (
+                        pdf[col].astype(str).str.lower().isin(["true", "1", "yes"])
+                    )
+                pdf[col] = pdf[col].astype("boolean[pyarrow]")
+
+            elif pa.types.is_integer(pa_type) and "int" not in current_dtype_str:
+                pdf[col] = pd.to_numeric(pdf[col], errors="coerce").astype(
+                    "int64[pyarrow]"
+                )
+
+            elif pa.types.is_floating(pa_type) and "float" not in current_dtype_str:
+                pdf[col] = pd.to_numeric(pdf[col], errors="coerce").astype(
+                    "float64[pyarrow]"
+                )
+
+        except Exception:
+            pass
+
+    return pdf
+
+
+class ParquetSaver(ManagedResource):
+    """
+    Production-grade Dask -> Parquet writer.
+    """
+
+    logger_extra = {"sibi_flux_component": __name__}
+
+    def __init__(
+        self,
+        df_result: dd.DataFrame,
+        parquet_storage_path: str,
+        *,
+        repartition_size: Optional[str] = "128MB",
+        persist: bool = False,
+        write_index: bool = False,
+        write_metadata_file: bool = True,
+        pyarrow_args: Optional[Dict[str, Any]] = None,
+        writer_threads: int | str = "auto",
+        arrow_cpu: Optional[int] = None,
+        partitions_per_round: int = 24,
+        max_delete_workers: int = 8,
+        write_gate_max: int = 2,
+        write_gate_key: Optional[str] = None,
+        partition_on: Optional[list[str]] = None,
+        dask_client: Optional[Any] = None,
+        schema: Optional[pa.Schema] = None,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+
+        if not isinstance(df_result, dd.DataFrame):
+            raise TypeError("df_result must be a Dask DataFrame")
+        if not self.fs:
+            raise ValueError("File system (fs) must be provided to ParquetSaver.")
+
+        self.df_result = df_result
+        self.dask_client = dask_client
+        self.parquet_storage_path = parquet_storage_path.rstrip("/")
+        self.repartition_size = repartition_size
+        self.persist = persist
+        self.write_index = write_index
+        self.write_metadata_file = write_metadata_file
+        self.pyarrow_args = dict(pyarrow_args or {})
+
+        if writer_threads == "auto" or writer_threads is None:
+            self.writer_threads = self._calculate_optimal_threads()
+            self.logger.debug(f"Auto-tuned writer_threads to: {self.writer_threads}")
+        else:
+            self.writer_threads = max(1, int(writer_threads))
+
+        self.arrow_cpu = None if arrow_cpu is None else max(1, int(arrow_cpu))
+        self.partitions_per_round = max(1, int(partitions_per_round))
+        self.max_delete_workers = max(1, int(max_delete_workers))
+        self.write_gate_max = max(1, int(write_gate_max))
+        self.write_gate_key = (write_gate_key or self.parquet_storage_path).rstrip("/")
+        self.partition_on = partition_on or []
+        self.scheduler = kwargs.get("scheduler", None)
+        self.force_local = kwargs.get("force_local", False)
+        self.schema = schema
+
+        self.pyarrow_args.setdefault("compression", "zstd")
+
+        self.protocol = "file"
+        if "://" in self.parquet_storage_path:
+            self.protocol = self.parquet_storage_path.split(":", 1)[0]
+
+    def save_to_parquet(
+        self, output_directory_name: str = "default_output", overwrite: bool = True
+    ) -> str:
+        import dask.config
+        from contextlib import nullcontext
+
+        # 1. Determine Execution Mode
+        assert self.fs is not None, "FileSystem (fs) must be available"
+        using_distributed = False
+        if self.force_local:
+            using_distributed = False
+        elif self.dask_client is not None:
+            using_distributed = True
+        else:
+            try:
+                from dask.distributed import get_client
+
+                get_client()
+                using_distributed = True
+            except (ImportError, ValueError):
+                pass
+            except Exception:
+                # Fallback if get_client() fails weirdly
+                pass
+
+        # 2. Configure Context Factories
+        # We use functions/lambdas to create FRESH context managers every time 'with' is called.
+        if using_distributed:
+            # Distributed: No-op contexts
+            def get_config_context():
+                return nullcontext()
+
+            def get_pool_context():
+                return nullcontext()
+
+        else:
+            # Local: Use context manager for 'tasks' shuffle, DO NOT SET GLOBAL POOL
+            def get_config_context():
+                return dask.config.set({"dataframe.shuffle.method": "tasks"})
+
+            # For local pool, we pass the pool directly to compute() or allow Dask to handle it.
+            # However, since we are doing ddf.to_parquet(), we can't easily inject the pool argument
+            # without a global config context if not using compute().
+            # Safer approach: Use a context manager for the pool too, but ensure it restores state.
+            def get_pool_context():
+                if self.scheduler == "synchronous":
+                    return dask.config.set(scheduler="synchronous")
+                elif self.scheduler == "threads":
+                    return dask.config.set(
+                        pool=ThreadPool(self.writer_threads), scheduler="threads"
+                    )
+                elif not using_distributed and self.scheduler is None:
+                    # Default to single-threaded for local execution to override any global client
+                    # and avoid graph dependencies mismatch (Missing Dependency error)
+                    return dask.config.set(scheduler="single-threaded")
+                # If scheduler is None and using_distributed is True, rely on global config
+                return nullcontext()
+
+        # 3. Path Resolution
+        if self.partition_on:
+            overwrite = False
+            target_path = self.parquet_storage_path.rstrip("/")
+        else:
+            target_path = f"{self.parquet_storage_path}/{output_directory_name}".rstrip(
+                "/"
+            )
+
+        sem = get_write_sem(self.write_gate_key, self.write_gate_max)
+
+        # 4. Execution
+        # We use the factory to get a fresh context for the config block
+        with get_config_context():
+            with sem:
+                if overwrite and self.fs.exists(target_path):
+                    self._clear_directory_safely(target_path)
+                self.fs.mkdirs(target_path, exist_ok=True)
+
+                schema = self._define_schema()
+                ddf = self._coerce_ddf_to_schema(self.df_result, schema)
+
+                if self.repartition_size:
+                    ddf = ddf.repartition(partition_size=self.repartition_size)
+
+                # Persist: Request a FRESH pool context
+                if self.persist:
+                    with get_pool_context():
+                        if using_distributed:
+                            ddf = safe_persist(ddf)
+                        else:
+                            # Local Persistence: Avoid safe_persist to prevent Client resurrection
+                            ddf = ddf.persist()
+
+                old_arrow_cpu = None
+                if self.arrow_cpu:
+                    old_arrow_cpu = pa.get_cpu_count()
+                    pa.set_cpu_count(self.arrow_cpu)
+
+                try:
+                    params = {
+                        "path": target_path,
+                        "engine": "pyarrow",
+                        "filesystem": self.fs,
+                        "write_index": self.write_index,
+                        "write_metadata_file": self.write_metadata_file,
+                        "schema": schema,
+                        **self.pyarrow_args,
+                    }
+
+                    if self.partition_on:
+                        params["partition_on"] = self.partition_on
+
+                    # Write: Compute with explicit scheduler if determined
+                    # Logic: If we determined a local scheduler (single-threaded, threads, synchronous),
+                    # we should pass it explicitly to compute() to override any global/default client.
+                    # CRITICAL FIX: Use compute=False and explicit dask.compute() to enforce scheduler
+                    write_task = ddf.to_parquet(**params, compute=False)
+
+                    compute_kwargs = {}
+                    if not using_distributed:
+                        if self.scheduler == "synchronous":
+                            compute_kwargs["scheduler"] = "synchronous"
+                        elif self.scheduler == "threads":
+                            compute_kwargs["scheduler"] = "threads"
+                            # Note: Pool logic for threads is handled by get_pool_context if needed,
+                            # but explicit scheduler arg helps.
+                        elif self.scheduler == "single-threaded":
+                            compute_kwargs["scheduler"] = "single-threaded"
+                        elif self.scheduler is None:
+                            # Default fallback for force_local
+                            compute_kwargs["scheduler"] = "single-threaded"
+
+                    with get_pool_context():
+                        dask.compute(write_task, **compute_kwargs)
+                finally:
+                    if old_arrow_cpu is not None:
+                        pa.set_cpu_count(old_arrow_cpu)
+
+        self.logger.info(
+            f"Parquet dataset written: {target_path}", extra=self.logger_extra
+        )
+        return target_path
+
+    # ---------- Internals ----------
+
+    def _calculate_optimal_threads(self) -> int:
+        try:
+            if "OMP_NUM_THREADS" in os.environ:
+                available_cores = int(os.environ["OMP_NUM_THREADS"])
+            elif hasattr(os, "sched_getaffinity"):
+                available_cores = len(os.sched_getaffinity(0))
+            else:
+                available_cores = os.cpu_count() or 4
+            recommended = available_cores * 3
+            return max(4, min(recommended, 32))
+        except Exception:
+            return 8
+
+    def _clear_directory_safely(self, directory: str) -> None:
+        if self.protocol.startswith("s3"):
+            entries = [p for p in self.fs.glob(f"{directory}/**") if p != directory]
+            if not entries:
+                return
+
+            def _rm_one(p: str) -> None:
+                try:
+                    self.fs.rm_file(p)
+                except Exception as e:
+                    self.logger.warning(
+                        f"Delete failed '{p}': {e}", extra=self.logger_extra
+                    )
+
+            with ThreadPoolExecutor(max_workers=self.max_delete_workers) as ex:
+                list(ex.map(_rm_one, entries))
+            try:
+                self.fs.rm(directory, recursive=False)
+            except Exception:
+                pass
+        else:
+            self.fs.rm(directory, recursive=True)
+
+    def _define_schema(self) -> pa.Schema:
+        fields = []
+        # Pre-calculate overrides from self.schema if provided
+        overrides = {f.name: f.type for f in self.schema} if self.schema else {}
+
+        for name, dtype in self.df_result.dtypes.items():
+            if name in overrides:
+                fields.append(pa.field(name, overrides[name]))
+                continue
+
+            if name in self.partition_on:
+                fields.append(pa.field(name, pa.string()))
+                continue
+
+            dtype_str = str(dtype).lower()
+            lower_name = name.lower()
+
+            # --- PRIORITY 1: Heuristics ---
+            if any(x in lower_name for x in ["_dt", "_date", "_at", "time"]):
+                pa_type = pa.timestamp("ns", tz="UTC")
+            elif lower_name.startswith(("is_", "has_", "flag_")):
+                pa_type = pa.bool_()
+            elif (
+                any(x in lower_name for x in ["_id", "id"]) and "guid" not in lower_name
+            ):
+                pa_type = pa.int64()
+            elif any(
+                x in lower_name for x in ["amount", "price", "score", "percent", "rate"]
+            ):
+                pa_type = pa.float64()
+
+            # --- PRIORITY 2: Dtypes ---
+            elif "int" in dtype_str:
+                pa_type = pa.int64()
+            elif "float" in dtype_str or "double" in dtype_str:
+                pa_type = pa.float64()
+            elif "bool" in dtype_str:
+                pa_type = pa.bool_()
+            elif "datetime" in dtype_str or "timestamp" in dtype_str:
+                pa_type = pa.timestamp("ns", tz="UTC")
+            else:
+                pa_type = pa.string()
+
+            fields.append(pa.field(name, pa_type))
+
+        return pa.schema(fields)
+
+    def _coerce_ddf_to_schema(
+        self, ddf: dd.DataFrame, schema: pa.Schema
+    ) -> dd.DataFrame:
+        target = {f.name: f.type for f in schema}
+
+        meta_cols: Dict[str, pd.Series] = {}
+        for name, typ in target.items():
+            if name in self.partition_on:
+                meta_cols[name] = pd.Series([], dtype="string[pyarrow]")
+            elif pa.types.is_timestamp(typ):
+                meta_cols[name] = pd.Series([], dtype="timestamp[ns, tz=UTC][pyarrow]")
+            elif pa.types.is_integer(typ):
+                meta_cols[name] = pd.Series([], dtype="int64[pyarrow]")
+            elif pa.types.is_floating(typ):
+                meta_cols[name] = pd.Series([], dtype="float64[pyarrow]")
+            elif pa.types.is_boolean(typ):
+                meta_cols[name] = pd.Series([], dtype="boolean[pyarrow]")
+            else:
+                meta_cols[name] = pd.Series([], dtype="string[pyarrow]")
+
+        new_meta = pd.DataFrame(meta_cols, index=ddf._meta.index)
+        new_meta = new_meta[list(target.keys())]
+
+        coerce_fn = partial(
+            _coerce_partition, target=target, partition_cols=self.partition_on
+        )
+        return ddf.map_partitions(coerce_fn, meta=new_meta)

sibi_flux/utils/parquet_saver/_write_gatekeeper.py

@@ -0,0 +1,33 @@
+# write_gatekeeper.py
+from __future__ import annotations
+import threading
+
+_LOCK = threading.Lock()
+_SEMS: dict[str, threading.Semaphore] = {}
+
+
+def get_write_sem(key: str, max_concurrency: int) -> threading.Semaphore:
+    """
+    Acquires or creates a semaphore for a given key, ensuring thread-safe access
+    and maximum concurrency.
+
+    This function retrieves an existing semaphore associated with the provided
+    key or creates a new one if it does not exist. The semaphore limits
+    concurrent access based on the provided maximum concurrency value. This is
+    used to manage thread-level synchronization for a specific resource or
+    operation identified by the given key.
+
+    :param key: The unique identifier for which the semaphore is associated.
+    :param max_concurrency: The maximum number of concurrent access allowed. A
+                            value less than 1 will default to a concurrency limit
+                            of 1.
+    :return: A `threading.Semaphore` object for the specified key, initialized
+             with the maximum concurrency.
+    :rtype: threading.Semaphore
+    """
+    with _LOCK:
+        sem = _SEMS.get(key)
+        if sem is None:
+            sem = threading.Semaphore(max(1, int(max_concurrency)))
+            _SEMS[key] = sem
+        return sem