benches: bench sh script and fixed python bench

Signed-off-by: nuts_rice <[email protected]>

Author: nuts-rice

benches/generate_data.py

@@ -2,9 +2,9 @@
 import shapely
 
 
-gdf = gpd.read_file("Utah.geojson.zip", engine="pyogrio")
+gdf = gpd.read_file("./bench_data/8dc58605f9dd484295c7d065694cdc0f_0.geojson", engine="pyogrio")
 bounds = shapely.bounds(gdf.geometry)
 print(bounds.shape)
 buf = bounds.tobytes("C")
-with open("bounds.raw", "wb") as f:
+with open("./bench_data/bounds.raw", "wb") as f:
     f.write(buf)

benches/requirements.txt

@@ -0,0 +1,4 @@
+geoindex-rs
+geopandas 
+pyogrio
+pyarrow

benches/rtree.py

@@ -0,0 +1,40 @@
+import geopandas as gpd
+import numpy as np
+import shapely
+from geoindex_rs import rtree as rt
+import timeit
+import pyogrio
+import requests
+def load_data():
+    path = "./bench_data/taxi_zones_4326.parquet"
+    gdf = gpd.read_parquet(path)
+
+    bounds = gdf.bounds
+    print(bounds)
+    return bounds
+
+
+def construct_wsg84_tree(bounds):
+    builder = rt.RTreeBuilder(bounds.shape[0])
+    min_x= np.array(bounds["minx"].values)
+    min_y=np.array(bounds["miny"].values)
+    max_x=np.array(bounds["maxx"].values)
+    max_y=np.array(bounds["maxy"].values)
+    builder.add(min_x, min_y, max_x, max_y)
+    return builder.finish()
+
+def construct_shapely_tree():
+    path = "./bench_data/taxi_zones_4326.parquet"
+    gdf = gpd.read_parquet(path)
+    tree = shapely.SRTree(gdf.geometry)
+    return tree
+
+
+if __name__ == "__main__":
+    bounds = load_data()
+    print(timeit.timeit(stmt='construct_wsg84_tree(bounds)', number=100, globals=globals()))
+
+
+
+
+

benches/rtree.rs

@@ -9,7 +9,7 @@ use rstar::AABB;
 use std::fs::read;
 
 fn load_data() -> Vec<f64> {
-    let buf = read("benches/bounds.raw").unwrap();
+    let buf = read("benches/bench_data/bounds.raw").unwrap();
     cast_slice(&buf).to_vec()
 }
 

scripts/bench.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+
+cd ./benches/bench_data 
+if [ ! -f "8dc58605f9dd484295c7d065694cdc0f_0.geojson" ]
+  then 
+    echo "Downloading geojson benchmark data..."
+    wget https://opendata.arcgis.com/datasets/8dc58605f9dd484295c7d065694cdc0f_0.geojson 
+  else
+    echo "Benchmark data already downloaded"
+fi    
+if [ ! -f "taxi_zones_4326.parquet" ]
+  then 
+    echo "Downloading parquet benchmark data..."
+    wget https://data.source.coop/cholmes/nyc-taxi-zones/taxi_zones_4326.parquet 
+  else
+    echo "Parquet Benchmark data already downloaded"
+fi    
+
+
+cd ../
+pip install -r requirements.txt
+python generate_data.py
+cd ../
+echo "Running base benchmarks..."
+cargo bench --bench rtree  
+echo "Running benchmarks with rayon feature..."
+cargo bench --bench rtree --features rayon
+cd ./benches
+python rtree.py