Merge pull request #67 from NREL/rjf/omf-revisions

Rjf/omf revisions

Author: robfitzgerald

rust/bambam-omf/src/collection/collector.rs

@@ -1,13 +1,13 @@
 use arrow::array::RecordBatch;
 use chrono::NaiveDate;
 use futures::stream::{self, StreamExt};
+use itertools::Itertools;
 use object_store::{path::Path, ListResult, ObjectMeta, ObjectStore};
 use parquet::arrow::arrow_reader::ArrowPredicate;
 use parquet::arrow::arrow_reader::ArrowReaderOptions;
 use parquet::arrow::async_reader::ParquetObjectReader;
 use parquet::arrow::async_reader::ParquetRecordBatchStreamBuilder;
 use rayon::prelude::*;
-use serde::de::DeserializeOwned;
 use std::sync::Arc;
 use std::time::Instant;
 
@@ -184,47 +184,29 @@ impl OvertureMapsCollector {
             .collect::<Result<Vec<RecordBatch>, _>>()
             .map_err(|e| OvertureMapsCollectionError::RecordBatchRetrievalError { source: e })?;
 
-        // Deserialize batches into record types
-        let records: Vec<Vec<OvertureRecord>> = match record_type {
-            OvertureRecordType::Places => record_batches
-                .par_iter()
-                .map(deserialize_batch::<PlacesRecord>)
-                .map(|records_result| {
-                    records_result
-                        .map(|records| records.into_iter().map(OvertureRecord::Places).collect())
-                })
-                .collect::<Result<Vec<_>, OvertureMapsCollectionError>>()?,
-            OvertureRecordType::Buildings => record_batches
-                .par_iter()
-                .map(deserialize_batch::<BuildingsRecord>)
-                .map(|records_result| {
-                    records_result
-                        .map(|records| records.into_iter().map(OvertureRecord::Buildings).collect())
-                })
-                .collect::<Result<Vec<_>, OvertureMapsCollectionError>>()?,
-            OvertureRecordType::Segment => record_batches
-                .par_iter()
-                .map(deserialize_batch::<TransportationSegmentRecord>)
-                .map(|records_result| {
-                    records_result
-                        .map(|records| records.into_iter().map(OvertureRecord::Segment).collect())
-                })
-                .collect::<Result<Vec<_>, OvertureMapsCollectionError>>()?,
-            OvertureRecordType::Connector => record_batches
-                .par_iter()
-                .map(deserialize_batch::<TransportationConnectorRecord>)
-                .map(|records_result| {
-                    records_result
-                        .map(|records| records.into_iter().map(OvertureRecord::Connector).collect())
-                })
-                .collect::<Result<Vec<_>, OvertureMapsCollectionError>>()?,
-        };
-        log::info!("Deserialization time {:?}", start_collection.elapsed());
+        let start_deserialization = Instant::now();
+        let records: Vec<OvertureRecord> = record_batches
+            .par_iter()
+            .map(|batch| match record_type {
+                OvertureRecordType::Places => record_type.process_batch::<PlacesRecord>(batch),
+                OvertureRecordType::Buildings => {
+                    record_type.process_batch::<BuildingsRecord>(batch)
+                }
+                OvertureRecordType::Segment => {
+                    record_type.process_batch::<TransportationSegmentRecord>(batch)
+                }
+                OvertureRecordType::Connector => {
+                    record_type.process_batch::<TransportationConnectorRecord>(batch)
+                }
+            })
+            .collect::<Result<Vec<_>, _>>()?
+            .into_iter()
+            .flatten()
+            .collect_vec();
 
-        // Flatten the collection
-        let flatten_records = records.into_iter().flatten().collect();
+        log::info!("Deserialization time {:?}", start_deserialization.elapsed());
         log::info!("Total time {:?}", start_collection.elapsed());
-        Ok(flatten_records)
+        Ok(records)
     }
 
     pub fn collect_from_release(
@@ -237,21 +219,12 @@ impl OvertureMapsCollector {
             ReleaseVersion::Latest => self.get_latest_release()?,
             other => String::from(other),
         };
-        log::info!("Collecting OvertureMaps records from release {release_str}");
+        log::info!("Collecting OvertureMaps {record_type} records from release {release_str}");
         let path = Path::from(record_type.format_url(release_str));
         self.collect_from_path(path, record_type, row_filter_config)
     }
 }
 
-/// Deserialize recordBatch into type T
-fn deserialize_batch<T>(record_batch: &RecordBatch) -> Result<Vec<T>, OvertureMapsCollectionError>
-where
-    T: DeserializeOwned,
-{
-    serde_arrow::from_record_batch(record_batch)
-        .map_err(|e| OvertureMapsCollectionError::DeserializeError(format!("Serde error: {e}")))
-}
-
 #[cfg(test)]
 mod test {
     use crate::collection::{

rust/bambam-omf/src/collection/filter/bbox_row_predicate.rs

@@ -5,6 +5,7 @@ use arrow::{
 };
 use parquet::arrow::arrow_reader::ArrowPredicate;
 
+/// tests if a row is contained within a bounding box.
 pub struct BboxRowPredicate {
     bbox: Bbox,
     projection_mask: parquet::arrow::ProjectionMask,
@@ -24,11 +25,13 @@ impl ArrowPredicate for BboxRowPredicate {
         &self.projection_mask
     }
 
+    /// tests the bounding box of each row in the record batch, filtering entries
+    /// that are not fully-contained.
     fn evaluate(
         &mut self,
         batch: arrow::array::RecordBatch,
     ) -> Result<arrow::array::BooleanArray, arrow::error::ArrowError> {
-        let struct_array = batch
+        let bbox_struct = batch
             .column_by_name("bbox")
             .ok_or(ArrowError::ParquetError(String::from(
                 "`bbox` column not found",
@@ -39,36 +42,47 @@ impl ArrowPredicate for BboxRowPredicate {
                 "Cannot cast column `bbox` to StructArray type",
             )))?;
 
-        let x_min_col = struct_array
-            .column_by_name("xmin")
-            .ok_or(ArrowError::ParquetError(String::from(
-                "`bbox.xmin` column not found",
-            )))?
-            .as_any()
-            .downcast_ref::<Float32Array>()
-            .ok_or(ArrowError::ParquetError(String::from(
-                "Cannot cast column `bbox.xmin` to Float32Array type",
-            )))?;
+        let xmins = get_column::<Float32Array>("xmin", bbox_struct)?;
+        let ymins = get_column::<Float32Array>("ymin", bbox_struct)?;
+        let xmaxs = get_column::<Float32Array>("xmax", bbox_struct)?;
+        let ymaxs = get_column::<Float32Array>("ymax", bbox_struct)?;
 
-        let y_min_col = struct_array
-            .column_by_name("ymin")
-            .ok_or(ArrowError::ParquetError(String::from(
-                "`bbox.ymin` column not found",
-            )))?
-            .as_any()
-            .downcast_ref::<Float32Array>()
-            .ok_or(ArrowError::ParquetError(String::from(
-                "Cannot cast column `bbox.ymin` to Float32Array type",
-            )))?;
-
-        let boolean_values: Vec<bool> = (0..struct_array.len())
-            .map(|i| {
-                self.bbox.xmin < x_min_col.value(i)
-                    && x_min_col.value(i) < self.bbox.xmax
-                    && self.bbox.ymin < y_min_col.value(i)
-                    && y_min_col.value(i) < self.bbox.ymax
-            })
+        let boolean_values: Vec<bool> = (0..bbox_struct.len())
+            .map(|i| within_box(i, xmins, ymins, xmaxs, ymaxs, &self.bbox))
             .collect();
         Ok(BooleanArray::from(boolean_values))
     }
 }
+
+/// helper function to get a column by name from a struct array and return it as
+/// the expected type.
+fn get_column<'b, T>(col: &str, struct_array: &'b StructArray) -> Result<&'b T, ArrowError>
+where
+    T: 'static,
+{
+    struct_array
+        .column_by_name(col)
+        .ok_or(ArrowError::ParquetError(format!(
+            "'bbox.{col}' column not found"
+        )))?
+        .as_any()
+        .downcast_ref::<T>()
+        .ok_or(ArrowError::ParquetError(format!(
+            "Cannot cast column 'bbox.{col}' to expected type"
+        )))
+}
+
+/// helper function to test whether a given row's values are contained within the bounding box.
+fn within_box(
+    index: usize,
+    xmins: &Float32Array,
+    ymins: &Float32Array,
+    xmaxs: &Float32Array,
+    ymaxs: &Float32Array,
+    bbox: &Bbox,
+) -> bool {
+    bbox.xmin <= xmins.value(index)
+        && xmaxs.value(index) <= bbox.xmax
+        && bbox.ymin <= ymins.value(index)
+        && ymaxs.value(index) <= bbox.ymax
+}

rust/bambam-omf/src/collection/record/overture_record.rs

@@ -9,3 +9,27 @@ pub enum OvertureRecord {
     Segment(TransportationSegmentRecord),
     Connector(TransportationConnectorRecord),
 }
+
+impl From<PlacesRecord> for OvertureRecord {
+    fn from(value: PlacesRecord) -> Self {
+        Self::Places(value)
+    }
+}
+
+impl From<BuildingsRecord> for OvertureRecord {
+    fn from(value: BuildingsRecord) -> Self {
+        Self::Buildings(value)
+    }
+}
+
+impl From<TransportationSegmentRecord> for OvertureRecord {
+    fn from(value: TransportationSegmentRecord) -> Self {
+        Self::Segment(value)
+    }
+}
+
+impl From<TransportationConnectorRecord> for OvertureRecord {
+    fn from(value: TransportationConnectorRecord) -> Self {
+        Self::Connector(value)
+    }
+}

rust/bambam-omf/src/collection/record/record_type.rs

@@ -1,3 +1,8 @@
+use arrow::array::RecordBatch;
+use serde::de::DeserializeOwned;
+
+use crate::collection::{OvertureMapsCollectionError, OvertureRecord};
+
 pub enum OvertureRecordType {
     Places,
     Buildings,
@@ -22,4 +27,31 @@ impl OvertureRecordType {
             }
         }
     }
+
+    /// processes an arrow [RecordBatch] into an [OvertureRecord] collection,
+    /// deserializing into the underlying row type struct along the way.
+    pub fn process_batch<R>(
+        &self,
+        record_batch: &RecordBatch,
+    ) -> Result<Vec<OvertureRecord>, OvertureMapsCollectionError>
+    where
+        R: DeserializeOwned + Into<OvertureRecord>,
+    {
+        let as_rows: Vec<R> = serde_arrow::from_record_batch(record_batch).map_err(|e| {
+            OvertureMapsCollectionError::DeserializeError(format!("Serde error: {e}"))
+        })?;
+        let as_result: Vec<OvertureRecord> = as_rows.into_iter().map(Into::into).collect();
+        Ok(as_result)
+    }
+}
+
+impl std::fmt::Display for OvertureRecordType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Places => write!(f, "Places"),
+            Self::Buildings => write!(f, "Buildings"),
+            Self::Segment => write!(f, "Segment"),
+            Self::Connector => write!(f, "Connector"),
+        }
+    }
 }

script/setup_test_bambam.py

@@ -1,15 +1,21 @@
 print("importing osmnx, compass")
 import osmnx as ox
 from nrel.routee.compass.io import generate_compass_dataset
-
+from nrel.routee.compass.io.generate_dataset import GeneratePipelinePhase
 
 if __name__ == "__main__":
+
+    phases = [
+        GeneratePipelinePhase.CONFIG,
+        GeneratePipelinePhase.GRAPH
+    ]
+
     print("downloading graph")
     g = ox.graph_from_place("Denver, Colorado, USA", network_type="drive")
     print("processing graph into compass dataset")
-    generate_compass_dataset(g, output_directory="denver_co")
+    generate_compass_dataset(g, output_directory="denver_co", phases=phases)
 
     # Boulder graph for GTFS
     g = ox.graph_from_place("Boulder, Colorado, USA", network_type="drive")
     print("processing graph into compass dataset")
-    generate_compass_dataset(g, output_directory="boulder_co")
+    generate_compass_dataset(g, output_directory="boulder_co", phases=phases)