Add support for geoindex rtrees. (#42)

* Add geoindex rtree generation to benchmarks.

* Remove randomization from geometry generation.

* Initial pass at spatial rtree utilization.

* Remove unneeded column skipping and ensure projection columns are kept.

* Refactor process_chunk_with_row_filter_async to reduce aray scan.

* Refactor geometry expression parsing.

* Add geoindex option to test store.

* Use explict members of the ArraysToGenerate enum rather than all.

* Include low level wkb benchmark that demonstrates slow vlen decoding.

* Formatting fixes.

* Bbox S3 benchmark.

* Use geoindex utility to downcast bbox during rtree construction.

* Make all imports top level.

* Use MEGA_BYTE constant.

* Use geoindex bbox casting utility for tests.

* Move all imports to top level.

* Refactor bespoke wkt parsing logic to use wkt crate.

* Include test for invalid rtree index error handling.

Author: sharkinsspatial

Cargo.lock

@@ -10,7 +10,10 @@ arrow-schema = "57.0.0"
 async-trait = "0.1.89"
 datafusion = "52.0"
 futures = "0.3.31"
+geo-index = "0.2"
+geo-traits = "0.3"
 geodatafusion = "0.3"
+wkt = { version = "0.14", default-features = false }
 geoarrow-schema = "0.7.0"
 icechunk = "0.3.16"
 object_store = "0.12.4"
@@ -29,6 +32,7 @@ chrono = "0.4.44"
 criterion = { version = "0.5.1", features = ["async_tokio"] }
 dhat = "0.3"
 geo = "0.31"
+geo-index = "0.2"
 geoarrow-array = "0.7.0"
 rand = "0.8"
 tempfile = "3"
@@ -59,3 +63,10 @@ test = false
 name = "bbox_columns_s3"
 harness = false
 test = false
+
+[[bench]]
+name = "bbox_zarrs_raw"
+harness = false
+test = false
+
+

Cargo.toml

@@ -13,7 +13,7 @@ use zarr_datafusion_search::table_provider::ZarrTableProvider;
 fn bbox_columns_bench_local(c: &mut Criterion) {
     let rt = Runtime::new().unwrap();
     let (session, _temp_dir) =
-        generate_icechunk_store_local(&rt, ArraysToGenerate::BboxColumns).unwrap();
+        generate_icechunk_store_local(&rt, &[ArraysToGenerate::BboxColumns]).unwrap();
     let table_provider = Arc::new(
         rt.block_on(ZarrTableProvider::new_icechunk(session, "/meta"))
             .unwrap(),

benches/bbox_columns_local.rs

@@ -11,7 +11,15 @@ use zarr_datafusion_search::table_provider::ZarrTableProvider;
 
 fn bbox_bench_local(c: &mut Criterion) {
     let rt = Runtime::new().unwrap();
-    let (session, _temp_dir) = generate_icechunk_store_local(&rt, ArraysToGenerate::Bbox).unwrap();
+    let (session, _temp_dir) = generate_icechunk_store_local(
+        &rt,
+        &[
+            ArraysToGenerate::Datetime,
+            ArraysToGenerate::Bbox,
+            ArraysToGenerate::RtreeIndex,
+        ],
+    )
+    .unwrap();
     let table_provider = Arc::new(
         rt.block_on(ZarrTableProvider::new_icechunk(session, "/meta"))
             .unwrap(),

benches/bbox_local.rs

@@ -0,0 +1,47 @@
+mod common;
+
+use common::{BBOX_SQL, run_bench};
+use criterion::{Criterion, criterion_group, criterion_main};
+use datafusion::prelude::SessionContext;
+use icechunk::{ObjectStorage, Repository, repository::VersionInfo};
+use std::collections::HashMap;
+use std::sync::Arc;
+use tokio::runtime::Runtime;
+use zarr_datafusion_search::table_provider::ZarrTableProvider;
+
+fn bbox_bench_s3(c: &mut Criterion) {
+    let bucket = "zarr-datafusion-search".to_string();
+    let prefix = "".to_string();
+
+    let rt = Runtime::new().unwrap();
+    let storage = rt
+        .block_on(ObjectStorage::new_s3(
+            bucket,
+            Some(prefix),
+            None, // credentials - uses default AWS credential chain
+            None, // config - uses default S3 options
+        ))
+        .unwrap();
+    let repo = rt
+        .block_on(Repository::open_or_create(
+            None,
+            Arc::new(storage),
+            HashMap::new(),
+        ))
+        .unwrap();
+    let session = rt
+        .block_on(repo.readonly_session(&VersionInfo::BranchTipRef("main".to_string())))
+        .unwrap();
+
+    let table_provider = Arc::new(
+        rt.block_on(ZarrTableProvider::new_icechunk(session, "/meta"))
+            .unwrap(),
+    );
+
+    let ctx = SessionContext::new();
+    ctx.register_table("zarr_data", table_provider).unwrap();
+    run_bench(c, &rt, &ctx, "bbox_bench_s3", "bbox_bench_s3", BBOX_SQL);
+}
+
+criterion_group!(benches_s3, bbox_bench_s3);
+criterion_main!(benches_s3);

benches/bbox_s3.rs

@@ -0,0 +1,59 @@
+mod common;
+
+use common::{ArraysToGenerate, generate_icechunk_store_local};
+use criterion::{Criterion, SamplingMode, criterion_group, criterion_main};
+use std::sync::Arc;
+use tokio::runtime::Runtime;
+use zarrs::array::Array;
+use zarrs_icechunk::AsyncIcechunkStore;
+
+fn bbox_zarrs_raw_bench(c: &mut Criterion) {
+    let rt = Runtime::new().unwrap();
+    let (session, _temp_dir) =
+        generate_icechunk_store_local(&rt, &[ArraysToGenerate::Bbox]).unwrap();
+    let store = Arc::new(AsyncIcechunkStore::new(session));
+
+    // Criterion benchmark
+    let mut group = c.benchmark_group("bbox_zarrs_raw");
+    group.sample_size(10);
+    group.sampling_mode(SamplingMode::Flat);
+    group.warm_up_time(std::time::Duration::from_secs(1));
+    group.measurement_time(std::time::Duration::from_secs(2));
+
+    group.bench_function("fetch_all_bbox_chunks", |b| {
+        b.to_async(&rt).iter(|| async {
+            let bbox_array = Array::async_open(store.clone(), "/meta/bbox")
+                .await
+                .unwrap();
+
+            let chunk_grid_shape = bbox_array.chunk_grid_shape();
+            let num_chunks: u64 = chunk_grid_shape.iter().product();
+
+            let mut total_rows = 0usize;
+            let max_rows = 3_000_000;
+
+            for chunk_idx in 0..num_chunks {
+                if total_rows >= max_rows {
+                    break;
+                }
+
+                let chunk_indices = vec![chunk_idx];
+                let subset = bbox_array.chunk_subset_bounded(&chunk_indices).unwrap();
+
+                let data: Vec<Vec<u8>> = bbox_array
+                    .async_retrieve_array_subset_elements(&subset)
+                    .await
+                    .unwrap();
+
+                total_rows += data.len();
+            }
+
+            total_rows.min(max_rows)
+        });
+    });
+
+    group.finish();
+}
+
+criterion_group!(benches, bbox_zarrs_raw_bench);
+criterion_main!(benches);