feat: add dedicated download functions

CHANGELOG.md

@@ -7,6 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Automatic total time wrapper decorator to aggregate nested function calls
+- Parameter `columns_to_download` for selecting columns to download from the dataset [#23](https://github.com/kraina-ai/overturemaestro/issues/23)
+
+### Changed
+
+- Refactored available release versions caching [#24](https://github.com/kraina-ai/overturemaestro/issues/24)
+- Removed hive partitioned parquet schema columns from GeoDataFrame loading
+
 ## [0.1.1] - 2024-11-24
 
 ### Changed

README.md

@@ -80,6 +80,8 @@ Required:
 
 - `geoarrow-rust-core (>=0.3.0)`: For transforming Arrow data to Shapely objects
 
+- `duckdb (>=1.1.0)`: For transforming downloaded data to the wide format
+
 - `pooch (>=1.6.0)`: For downloading precalculated dataset indexes
 
 - `rich (>=12.0.0)`: For showing progress bars

overturemaestro/_duckdb.py

@@ -0,0 +1,27 @@
+"""Helper functions for DuckDB."""
+
+from pathlib import Path
+from typing import Union
+
+import duckdb
+
+
+def _sql_escape(value: str) -> str:
+    """Escape value for SQL query."""
+    return value.replace("'", "''")
+
+
+def _set_up_duckdb_connection(tmp_dir_path: Union[str, Path]) -> "duckdb.DuckDBPyConnection":
+    """Create DuckDB connection in a given directory."""
+    local_db_file = "db.duckdb"
+    connection = duckdb.connect(
+        database=str(Path(tmp_dir_path) / local_db_file),
+        config=dict(preserve_insertion_order=False),
+    )
+    connection.sql("SET enable_progress_bar = false;")
+    connection.sql("SET enable_progress_bar_print = false;")
+
+    connection.install_extension("spatial")
+    connection.load_extension("spatial")
+
+    return connection

overturemaestro/_exceptions.py

@@ -1 +1,7 @@
 class QueryNotGeocodedError(ValueError): ...
+
+
+class MissingColumnError(ValueError): ...
+
+
+class HierarchyDepthOutOfBoundsError(ValueError): ...

overturemaestro/_parquet_multiprocessing.py

@@ -1,8 +1,9 @@
 import multiprocessing
+from multiprocessing.managers import SyncManager
 from pathlib import Path
 from queue import Empty, Queue
 from time import sleep, time
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast
 
 from overturemaestro._rich_progress import VERBOSITY_MODE, TrackProgressSpinner
 
@@ -31,13 +32,14 @@ def _job(
     columns: Optional[list[str]],
     filesystem: "fs.FileSystem",
 ) -> None:  # pragma: no cover
+    import hashlib
+
     import pyarrow.dataset as ds
     import pyarrow.parquet as pq
 
     current_pid = multiprocessing.current_process().pid
 
-    filepath = save_path / f"{current_pid}.parquet"
-    writer = None
+    writers = {}
     while not queue.empty():
         try:
             file_name, row_group_index = None, None
@@ -61,10 +63,16 @@ def _job(
                     tracker.value += 1
                 continue
 
-            if not writer:
-                writer = pq.ParquetWriter(filepath, result_table.schema)
+            h = hashlib.new("sha256")
+            h.update(result_table.schema.to_string().encode())
+            schema_hash = h.hexdigest()
+
+            if schema_hash not in writers:
+                filepath = save_path / str(current_pid) / f"{schema_hash}.parquet"
+                filepath.parent.mkdir(exist_ok=True, parents=True)
+                writers[schema_hash] = pq.ParquetWriter(filepath, result_table.schema)
 
-            writer.write_table(result_table)
+            writers[schema_hash].write_table(result_table)
 
             with tracker_lock:
                 tracker.value += 1
@@ -80,7 +88,7 @@ def _job(
             )
             raise MultiprocessingRuntimeError(msg) from ex
 
-    if writer:
+    for writer in writers.values():
         writer.close()
 
 
@@ -107,6 +115,13 @@ def exception(self) -> Optional[tuple[Exception, str]]:
         return self._exception
 
 
+class SingletonContextManager(SyncManager):
+    def __new__(cls, ctx: multiprocessing.context.SpawnContext) -> "SingletonContextManager":
+        if not hasattr(cls, "instance"):
+            cls.instance = ctx.Manager()
+        return cast(SingletonContextManager, cls.instance)
+
+
 def _read_row_group_number(path: str, filesystem: "fs.FileSystem") -> int:
     import pyarrow.parquet as pq
 
@@ -154,7 +169,7 @@ def map_parquet_dataset(
 
         from overturemaestro._rich_progress import TrackProgressBar
 
-        manager = ctx.Manager()
+        manager = SingletonContextManager(ctx=ctx)
 
         queue: Queue[tuple[str, int]] = manager.Queue()
         tracker: ValueProxy[int] = manager.Value("i", 0)

overturemaestro/advanced_functions/__init__.py

@@ -0,0 +1,49 @@
+# """
+# Advanced functions.
+
+# This module contains dedicated functions for specific use cases.
+# """
+
+# from overturemaestro.advanced_functions.buildings import (
+#     convert_bounding_box_to_buildings_geodataframe,
+#     convert_bounding_box_to_buildings_parquet,
+#     convert_geometry_to_buildings_geodataframe,
+#     convert_geometry_to_buildings_parquet,
+# )
+# from overturemaestro.advanced_functions.poi import (
+#     convert_bounding_box_to_pois_geodataframe,
+#     convert_bounding_box_to_pois_parquet,
+#     convert_geometry_to_pois_geodataframe,
+#     convert_geometry_to_pois_parquet,
+# )
+# from overturemaestro.advanced_functions.transportation import (
+#     convert_bounding_box_to_roads_geodataframe,
+#     convert_bounding_box_to_roads_parquet,
+#     convert_geometry_to_roads_geodataframe,
+#     convert_geometry_to_roads_parquet,
+# )
+# from overturemaestro.advanced_functions.wide_form import (
+#     convert_bounding_box_to_wide_form_geodataframe,
+#     convert_bounding_box_to_wide_form_parquet,
+#     convert_geometry_to_wide_form_geodataframe,
+#     convert_geometry_to_wide_form_parquet,
+# )
+
+# __all__ = [
+#     "convert_bounding_box_to_buildings_geodataframe",
+#     "convert_bounding_box_to_buildings_parquet",
+#     "convert_bounding_box_to_pois_geodataframe",
+#     "convert_bounding_box_to_pois_parquet",
+#     "convert_bounding_box_to_roads_geodataframe",
+#     "convert_bounding_box_to_roads_parquet",
+#     "convert_bounding_box_to_wide_form_geodataframe",
+#     "convert_bounding_box_to_wide_form_parquet",
+#     "convert_geometry_to_buildings_geodataframe",
+#     "convert_geometry_to_buildings_parquet",
+#     "convert_geometry_to_pois_geodataframe",
+#     "convert_geometry_to_pois_parquet",
+#     "convert_geometry_to_roads_geodataframe",
+#     "convert_geometry_to_roads_parquet",
+#     "convert_geometry_to_wide_form_geodataframe",
+#     "convert_geometry_to_wide_form_parquet",
+# ]