Feature : CI option release

overturemaps/cli.py

@@ -5,20 +5,21 @@
 specified bounding box in a few different file formats.
 
 """
+
 import json
 import os
 import sys
 from typing import Optional
 
 import click
 import pyarrow as pa
-import pyarrow.dataset as ds
 import pyarrow.compute as pc
+import pyarrow.dataset as ds
 import pyarrow.fs as fs
 import pyarrow.parquet as pq
 import shapely.wkb
 
-from . core import record_batch_reader, get_all_overture_types
+from .core import get_all_overture_types, record_batch_reader
 
 
 def get_writer(output_format, path, schema):
@@ -74,18 +75,24 @@ def cli():
     required=True,
 )
 @click.option("-o", "--output", required=False, type=click.Path())
+@click.option("-r", "--release", required=False, type=str)
 @click.option(
     "-t",
     "--type",
     "type_",
     type=click.Choice(get_all_overture_types()),
     required=True,
 )
-def download(bbox, output_format, output, type_):
+def download(bbox, output_format, output, type_, release):
     if output is None:
         output = sys.stdout
 
-    reader = record_batch_reader(type_, bbox)
+    reader = record_batch_reader(
+        type_,
+        bbox,
+        **({"release": release} if release is not None else {}),
+    )
+
     if reader is None:
         return
 

overturemaps/core.py

@@ -6,12 +6,15 @@
 import pyarrow.fs as fs
 
 
-def record_batch_reader(overture_type, bbox=None) -> Optional[pa.RecordBatchReader]:
+def record_batch_reader(
+    overture_type,
+    bbox=None,
+    release="2024-04-16-beta.0",
+) -> Optional[pa.RecordBatchReader]:
     """
     Return a pyarrow RecordBatchReader for the desired bounding box and s3 path
     """
-    path = _dataset_path(overture_type)
-
+    path = _dataset_path(overture_type, release)
     if bbox:
         xmin, ymin, xmax, ymax = bbox
         filter = (
@@ -26,6 +29,7 @@ def record_batch_reader(overture_type, bbox=None) -> Optional[pa.RecordBatchRead
     dataset = ds.dataset(
         path, filesystem=fs.S3FileSystem(anonymous=True, region="us-west-2")
     )
+
     batches = dataset.to_batches(filter=filter)
 
     # to_batches() can yield many batches with no rows. I've seen
@@ -34,6 +38,7 @@ def record_batch_reader(overture_type, bbox=None) -> Optional[pa.RecordBatchRead
     # each one bloating the size of a parquet file. Just omit
     # them so the RecordBatchReader only has non-empty ones. Use
     # the generator syntax so the batches are streamed out
+
     non_empty_batches = (b for b in batches if b.num_rows > 0)
 
     geoarrow_schema = geoarrow_schema_adapter(dataset.schema)
@@ -85,7 +90,10 @@ def geoarrow_schema_adapter(schema: pa.Schema) -> pa.Schema:
 }
 
 
-def _dataset_path(overture_type: str) -> str:
+def _dataset_path(
+    overture_type: str,
+    release: str,
+) -> str:
     """
     Returns the s3 path of the Overture dataset to use. This assumes overture_type has
     been validated, e.g. by the CLI
@@ -95,8 +103,10 @@ def _dataset_path(overture_type: str) -> str:
     # complete s3 path. Could be discovered by reading from the top-level s3
     # location but this allows to only read the files in the necessary partition.
     theme = type_theme_map[overture_type]
-    return f"overturemaps-us-west-2/release/2024-04-16-beta.0/theme={theme}/type={overture_type}/"
+    return (
+        f"overturemaps-us-west-2/release/{release}/theme={theme}/type={overture_type}/"
+    )
 
 
 def get_all_overture_types() -> List[str]:
-    return list(type_theme_map.keys())
+    return list(type_theme_map.keys())