sibi-dst 2025.8.9__tar.gz → 2025.9.2__tar.gz

{sibi_dst-2025.8.9 → sibi_dst-2025.9.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sibi-dst
-Version: 2025.8.9
+Version: 2025.9.2
 Summary: Data Science Toolkit
 Author: Luis Valverde
 Author-email: lvalverdeb@gmail.com
@@ -21,6 +21,7 @@ Requires-Dist: pyarrow (>=20.0.0,<21.0.0)
 Requires-Dist: pydantic (>=2.11.7,<3.0.0)
 Requires-Dist: pyiceberg[hive,s3fs] (>=0.9.1,<0.10.0)
 Requires-Dist: pymysql (>=1.1.1,<2.0.0)
+Requires-Dist: pyrosm (>=0.6.2,<0.7.0)
 Requires-Dist: s3fs (>=2025.5.1,<2026.0.0)
 Requires-Dist: sqlalchemy (>=2.0.41,<3.0.0)
 Requires-Dist: sse-starlette (>=3.0.2,<4.0.0)

{sibi_dst-2025.8.9 → sibi_dst-2025.9.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sibi-dst"
-version = "2025.8.9"
+version = "2025.9.2"
 description = "Data Science Toolkit"
 authors = ["Luis Valverde <lvalverdeb@gmail.com>"]
 readme = "README.md"
@@ -25,6 +25,7 @@ opentelemetry-exporter-otlp = "^1.35.0"
 opentelemetry-sdk = "^1.35.0"
 pyiceberg = {extras = ["hive", "s3fs"], version = "^0.9.1"}
 sse-starlette = "^3.0.2"
+pyrosm = "^0.6.2"
 
 [tool.poetry.group.dev]
 optional = true

{sibi_dst-2025.8.9 → sibi_dst-2025.9.2}/sibi_dst/df_helper/_df_helper.py

@@ -137,6 +137,7 @@ class DfHelper(ManagedResource):
     def __init__(self, backend="sqlalchemy", **kwargs):
         self.default_config = self.default_config or {}
         kwargs = {**self.default_config.copy(), **kwargs}
+        kwargs.setdefault("auto_sse", True)
         super().__init__(**kwargs)
         self.backend = backend
 
@@ -166,6 +167,18 @@ class DfHelper(ManagedResource):
         self.backend_strategy = strategy_cls(self)
 
     # ---------- ManagedResource hooks ----------
+    def get_sse(self):
+        return self._ensure_sse()
+
+    def _emit_bg(self, event: str, **data: Any) -> None:
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            # no running loop: run to completion
+            asyncio.run(self.emit(event, **data))
+        else:
+            loop.create_task(self.emit(event, **data))
+
     def _cleanup(self):
         attr_name = self._BACKEND_ATTR_MAP.get(self.backend)
         if not attr_name:
@@ -213,6 +226,7 @@ class DfHelper(ManagedResource):
         timeout: Optional[float] = None,
         **options
     ) -> Union[pd.DataFrame, dd.DataFrame]:
+        await self.emit(f"{self.__class__.__name__} load:start", message=f"Pulling data from {self.backend} backend")
         # 1) Async load if available, else run sync load in a thread.
         if hasattr(self.backend_strategy, "aload"):
             load_awaitable = self.backend_strategy.aload(**options)
@@ -224,18 +238,20 @@ class DfHelper(ManagedResource):
         self.total_records = total
 
         # 2) Post-processing steps are sync; offload to threads.
+        await self.emit(event=f"{self.__class__.__name__} load:progress", message=f"Post-processing {len(df)} records")
         df = await asyncio.to_thread(self._process_loaded_data, df)
         df = await asyncio.to_thread(self._post_process_df, df)
 
         # 3) Persist and compute can block; offload when needed.
         if persist and _is_dask_df(df):
             df = await asyncio.to_thread(df.persist)
-
         if as_pandas and _is_dask_df(df):
             # Allow separate timeout for compute if desired; reuse same timeout here.
             compute_awaitable = asyncio.to_thread(df.compute)
             return await (asyncio.wait_for(compute_awaitable, timeout) if timeout else compute_awaitable)
 
+        await self.emit(event=f"{self.__class__.__name__} load:progress", message=f"Returning {len(df)} records")
+
         return df
 
     # ---------- dataframe post-processing ----------
@@ -307,6 +323,11 @@ class DfHelper(ManagedResource):
 
         self.logger.debug(f"Successfully saved '{parquet_filename}' to '{path}'.", extra=self.logger_extra)
 
+    async def asave_to_parquet(self, df: dd.DataFrame, **kwargs):
+        await self.emit(event=f"{self.__class__.__name__} save:start", message=f"Saving {len(df)} records to parquet")
+        await asyncio.to_thread(self.save_to_parquet, df, **kwargs)
+        await self.emit(event=f"{self.__class__.__name__} save:end", message=f"Saved {len(df)} records to parquet")
+
     def save_to_clickhouse(self, df: dd.DataFrame, **credentials):
         if not self._has_any_rows(df):
             self.logger.warning("Skipping save to ClickHouse: The provided DataFrame is empty.", extra=self.logger_extra)
@@ -315,6 +336,11 @@ class DfHelper(ManagedResource):
             writer.save_to_clickhouse(df)
             self.logger.debug("Save to ClickHouse completed.", extra=self.logger_extra)
 
+    async def asave_to_clickhouse(self, df: dd.DataFrame, **credentials):
+        await self.emit(event=f"{self.__class__.__name__} save:start", message=f"Saving {len(df)} records to ClickHouse")
+        await asyncio.to_thread(self.save_to_clickhouse, df, **credentials)
+        await self.emit(event=f"{self.__class__.__name__} save:end", message=f"Saved {len(df)} records to ClickHouse")
+
     # ---------- period loaders ----------
     def load_period(self, dt_field: str, start: str, end: str, **kwargs):
         final_kwargs = self._prepare_period_filters(dt_field, start, end, **kwargs)
@@ -346,5 +372,3 @@ class DfHelper(ManagedResource):
             return bool(ddf.head(1, npartitions=-1).shape[0])
         except Exception:
             return False
-
-

{sibi_dst-2025.8.9 → sibi_dst-2025.9.2}/sibi_dst/df_helper/backends/parquet/_parquet_options.py

@@ -205,6 +205,8 @@ class ParquetConfig(BaseModel):
                 filesystem=self.fs,
                 filters=pq_filters,
                 # Toggle based on file count; False is safer for many tiny files.
+                aggregate_files=True,
+                split_row_groups=True,
                 gather_statistics=False,
                 ignore_metadata_file=True,
             )

{sibi_dst-2025.8.9 → sibi_dst-2025.9.2}/sibi_dst/df_helper/backends/sqlalchemy/_load_from_db.py

@@ -30,7 +30,7 @@ class SqlAlchemyLoadFromDb(ManagedResource):
         self.engine = self.db_connection.engine
         self.query_config = plugin_query
         self.params_config = plugin_params
-        self.chunk_size = kwargs.get("chunk_size", self.params_config.df_params.get("chunk_size", 1000) if self.params_config else 1000)
+        self.chunk_size = kwargs.get("chunk_size", self.params_config.df_params.get("chunk_size", 10000) if self.params_config else 10000)
         self.total_records = -1
 
     def build_and_load(self) -> Tuple[int, dd.DataFrame]:

{sibi_dst-2025.8.9 → sibi_dst-2025.9.2}/sibi_dst/osmnx_helper/utils.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import math
 import os
 import pickle
@@ -9,236 +11,102 @@ import numpy as np
 import osmnx as ox
 from geopy.distance import geodesic
 
-
-#
-# options = {
-#    'ox_files_save_path': ox_files_save_path,
-#    'network_type': 'drive',
-#    'place': 'Costa Rica',
-#    'files_prefix': 'costa-rica-',
-# }
-# Usage example
-# handler = PBFHandler(**options)
-# handler.load()
-
+from typing import Optional
+from fsspec.core import url_to_fs
 
 class PBFHandler:
     """
-    Handles the creation, management, and visualization of graph data derived
-    from .pbf (Protocolbuffer Binary Format) files. This class enables the
-    loading, processing, saving, and reutilization of graph, node, and edge
-    data for geographical regions, supporting verbose mode for detailed outputs.
-
-    :ivar graph: The generated graph object representing the spatial network; can be None if not yet loaded or processed.
-    :type graph: Optional[NetworkX.Graph]
-    :ivar nodes: GeoDataFrame representing the nodes of the graph; can be None if not yet loaded or processed.
-    :type nodes: Optional[geopandas.GeoDataFrame]
-    :ivar edges: GeoDataFrame representing the edges of the graph; can be None if not yet loaded or processed.
-    :type edges: Optional[geopandas.GeoDataFrame]
-    :ivar rebuild: Indicates whether to rebuild the graph data, ignoring any existing cached files. Default is ``False``.
-    :type rebuild: bool
-    :ivar verbose: Enables verbose mode to provide detailed status messages during operations. Default is ``False``.
-    :type verbose: bool
-    :ivar place: The name of the geographical region to process with OpenStreetMap. Default is ``Costa Rica``.
-    :type place: str
-    :ivar filepath: The path to the directory where the graph, nodes, and edges pickle files are saved. Default is ``gis_data/``.
-    :type filepath: str
-    :ivar file_prefix: The prefix for the filenames of the saved graph, node, and edge pickle files. Default is ``costa-rica-``.
-    :type file_prefix: str
-    :ivar network_type: The type of network to extract from OpenStreetMap, such as "all" or other specific network types. Default is ``all``.
-    :type network_type: str
-    :ivar graph_file: Full path of the file to save or load the graph data as a pickle file.
-    :type graph_file: str
-    :ivar node_file: Full path of the file to save or load the graph's node data as a pickle file.
-    :type node_file: str
-    :ivar edge_file: Full path of the file to save or load the graph's edge data as a pickle file.
-    :type edge_file: str
+    Build/load OSMnx graph + nodes/edges; persist as pickle via fsspec.
     """
+
     def __init__(self, **kwargs):
         self.graph = None
-        self.nodes = None
-        self.edges = None
-        self.rebuild = kwargs.setdefault("rebuild", False)
-        self.verbose = kwargs.setdefault("verbose", False)
-        self.place = kwargs.setdefault('place', 'Costa Rica')
-        self.filepath = kwargs.setdefault('ox_files_save_path', "gis_data/")
-        self.file_prefix = kwargs.setdefault('file_prefix', 'costa-rica-')
-        self.network_type = kwargs.setdefault('network_type', 'all')
-        self.graph_file = f"{self.filepath}{self.file_prefix}graph.pkl"
-        self.node_file = f"{self.filepath}{self.file_prefix}nodes.pkl"
-        self.edge_file = f"{self.filepath}{self.file_prefix}edges.pkl"
-
-    def load(self):
-        """
-        Loads the required data files for processing. If the files do not exist or
-        if the `rebuild` flag is set to True, it will process and recreate the
-        necessary data from the source. Otherwise, it will load the data from
-        existing pickle files. This function ensures the target directory exists,
-        and processes files conditionally based on their presence.
-
-        :param verbose: Flag to control the verbosity of the function's output.
-        :param rebuild: Indicates whether the data should be rebuilt from the raw
-            source files.
-        :param graph_file: Path to the graph file to be loaded or rebuilt.
-        :param node_file: Path to the node file to be loaded or rebuilt.
-        :param edge_file: Path to the edge file to be loaded or rebuilt.
-        :param filepath: Path to the directory where files are processed and saved.
-
-        :return: None
-        """
+        self.nodes: Optional[gpd.GeoDataFrame] = None
+        self.edges: Optional[gpd.GeoDataFrame] = None
+
+        self.rebuild: bool = kwargs.setdefault("rebuild", False)
+        self.verbose: bool = kwargs.setdefault("verbose", False)
+        self.place: str = kwargs.setdefault("place", "Costa Rica")
+        self.network_type: str = kwargs.setdefault("network_type", "all")
+        base_url: str = kwargs.setdefault("data_path", "osmnx_data/pbf_files")
+        prefix: str = kwargs.setdefault("files_prefix", "costa-rica-").rstrip("-") + "-"
+
+        # Allow passing an fsspec instance directly
+        fs = kwargs.get("fs")
+        if fs is not None:
+            self.fs = fs
+            self.base = base_url.rstrip("/")
+        else:
+            self.fs, self.base = url_to_fs(base_url)
+
+        self.fs.mkdirs(self.base, exist_ok=True)
+
+        self.graph_file = f"{self.base.rstrip('/')}/{prefix}graph.pkl"
+        self.node_file = f"{self.base.rstrip('/')}/{prefix}nodes.pkl"
+        self.edge_file = f"{self.base.rstrip('/')}/{prefix}edges.pkl"
+
         if self.verbose:
-            print("Loading data...")
+            print(f"[PBFHandler] base={self.base}")
+            print(f"  graph={self.graph_file}")
+            print(f"  nodes={self.node_file}")
+            print(f"  edges={self.edge_file}")
 
-        files_to_check = [self.graph_file, self.node_file, self.edge_file]
+    # ---------- public API ----------
+    def load(self) -> None:
+        if self.verbose:
+            print("[PBFHandler] load()")
 
         if self.rebuild:
-            for file in files_to_check:
-                if os.path.exists(file):
-                    os.remove(file)
-        if not os.path.exists(self.filepath):
-            os.makedirs(self.filepath, exist_ok=True)
-            # self.process_pbf()
-            # self.save_to_pickle()
-        if not all(os.path.exists(f) for f in files_to_check):
+            self._delete_artifacts()
+
+        if not self._artifacts_exist():
             self.process_pbf()
             self.save_to_pickle()
         else:
             self.load_from_pickle()
 
+    def process_pbf(self) -> None:
+        if self.verbose:
+            print(f"[PBFHandler] processing: {self.place}")
+        self.graph = ox.graph_from_place(self.place, network_type=self.network_type)
+        self.nodes, self.edges = ox.graph_to_gdfs(self.graph)
+
+    def save_to_pickle(self) -> None:
+        if self.verbose:
+            print("[PBFHandler] saving via fsspec")
+        for path, obj in {
+            self.graph_file: self.graph,
+            self.node_file: self.nodes,
+            self.edge_file: self.edges,
+        }.items():
+            if obj is not None:
+                with self.fs.open(path, "wb") as f:
+                    pickle.dump(obj, f, protocol=pickle.HIGHEST_PROTOCOL)
+
+    def load_from_pickle(self) -> None:
+        if self.verbose:
+            print("[PBFHandler] loading via fsspec")
+        self.graph = self._load_pickle(self.graph_file)
+        self.nodes = self._load_pickle(self.node_file)
+        self.edges = self._load_pickle(self.edge_file)
+
+    # ---------- helpers ----------
+    def _artifacts_exist(self) -> bool:
+        return all(self.fs.exists(p) for p in (self.graph_file, self.node_file, self.edge_file))
+
+    def _delete_artifacts(self) -> None:
         if self.verbose:
-            print("Data loaded successfully.")
-
-    def process_pbf(self):
-        """
-        Processes the Protocolbuffer Binary Format (PBF) data specified for a given place by
-        utilizing the OSMnx library to create a graph representation and extracts nodes and
-        edges into GeoDataFrames. The function provides verbose output if enabled.
-
-        :param self: Refers to the current instance of the class containing this method.
-
-        :param self.verbose: bool
-            A flag to control verbose output. If True, detailed processing status messages are
-            logged to the console.
-
-        :param self.place: str
-            The name or description of the geographic place for which PBF data is processed. It
-            is used to construct a graph representation of the place.
-
-        :param self.network_type: str
-            The type of network graph to be created, typically one of 'all', 'walk', 'drive',
-            etc., reflecting the type of paths or streets included in the graph.
-
-        :return: None
-            This function does not return a value, but updates class attributes ``graph``,
-            ``nodes``, and ``edges``.
-
-        :raises Exception:
-            Raises a general exception when there is an error in processing the PBF data. Error
-            details are printed when verbose output is enabled.
-        """
-        try:
-            if self.verbose:
-                print(f"Processing PBF for {self.place}...")
-
-            self.graph = ox.graph_from_place(self.place, network_type=self.network_type)
-            self.nodes, self.edges = ox.graph_to_gdfs(self.graph)
-
-            if self.verbose:
-                print("PBF processed successfully.")
-        except Exception as e:
-            print(f"Error processing PBF: {e}")
-            raise
-
-    def save_to_pickle(self):
-        """
-        Saves data, including graph, nodes, and edges, to pickle files. Each data object is
-        saved to its corresponding file if available. If verbose mode is enabled, prints
-        messages indicating the saving progress and success.
-
-        :param self:
-            Represents the instance of the class that contains attributes `graph_file`,
-            `graph`, `node_file`, `nodes`, `edge_file`, `edges`, and `verbose`. These
-            attributes determine the files to save to and the data to save.
-
-        :raises Exception:
-            Raises an exception if an error occurs during the saving process.
-
-        :return:
-            None
-        """
-        try:
-            if self.verbose:
-                print("Saving data to pickle files...")
-
-            data_to_save = {
-                self.graph_file: self.graph,
-                self.node_file: self.nodes,
-                self.edge_file: self.edges
-            }
-
-            for file, data in data_to_save.items():
-                if data is not None:
-                    with open(file, 'wb') as f:
-                        pickle.dump(data, f, pickle.HIGHEST_PROTOCOL)
-
-            if self.verbose:
-                print("Data saved to pickle files successfully.")
-        except Exception as e:
-            print(f"Error saving to pickle: {e}")
-            raise
-
-    def load_from_pickle(self):
-        """
-        Loads data from pickle files specified by the attributes `graph_file`, `node_file`,
-        and `edge_file` and assigns them to the corresponding attributes `graph`,
-        `nodes`, and `edges`, respectively. Displays verbose messages during the load
-        process if the `verbose` attribute is set to True.
-
-        :raises Exception: If an error occurs during reading or deserialization of the
-                           pickle files.
-        """
-        try:
-            if self.verbose:
-                print("Loading data from pickle files...")
-
-            files_to_load = {
-                self.graph_file: 'graph',
-                self.node_file: 'nodes',
-                self.edge_file: 'edges'
-            }
-
-            for file, attr in files_to_load.items():
-                with open(file, 'rb') as f:
-                    setattr(self, attr, pickle.load(f))
-
-            if self.verbose:
-                print("Data loaded from pickle files successfully.")
-        except Exception as e:
-            print(f"Error loading from pickle: {e}")
-            raise
-
-    def plot_graph(self):
-        """
-        Plots the loaded graph using the OSMnx library.
-
-        This method checks if a graph is loaded and, if available, plots it. Outputs
-        verbose messages during the process if verbosity is enabled.
-
-        :raises Exception: Raises if an error occurs during the plotting process.
-        :return: None
-        """
-        try:
-            if self.graph is not None:
-                if self.verbose:
-                    print("Plotting the graph...")
-                ox.plot_graph(self.graph)
-                if self.verbose:
-                    print("Graph plotted successfully.")
-            else:
-                print("Graph is not loaded. Please load a PBF file first.")
-        except Exception as e:
-            print(f"Error plotting the graph: {e}")
-            raise
+            print("[PBFHandler] deleting artifacts (rebuild=True)")
+        for p in (self.graph_file, self.node_file, self.edge_file):
+            if self.fs.exists(p):
+                try:
+                    self.fs.rm_file(p)
+                except Exception:
+                    self.fs.rm(p)
+
+    def _load_pickle(self, path: str):
+        with self.fs.open(path, "rb") as f:
+            return pickle.load(f)
 
 
 def get_bounding_box_from_points(gps_points, margin=0.001):