add py bench

Change-Id: I82756252fbb0af27941d3af0eee92af1da935410

Author: ddupg

python/python/benchmarks/test_search.py

@@ -505,3 +505,64 @@ def test_late_materialization(test_dataset, benchmark, use_index):
         filter=f"{column} = 0",
         batch_size=32,
     )
+
+
+@pytest.fixture(scope="module")
+def test_geo_dataset(tmpdir_factory):
+    from geoarrow.rust.core import (
+        point,
+        points,
+    )
+
+    num_rows = 1_000_000
+    points_2d = points(
+        [np.random.randn(num_rows) * 100, np.random.randn(num_rows) * 100]
+    )
+
+    schema = pa.schema(
+        [
+            pa.field(point("xy")).with_name("points"),
+        ]
+    )
+    table = pa.Table.from_arrays([points_2d], schema=schema)
+    uri = str(tmpdir_factory.mktemp("test_geo_dataset"))
+    lance.write_dataset(table, uri)
+    ds = lance.dataset(uri)
+    return ds
+
+
+@pytest.mark.benchmark(group="geo")
+@pytest.mark.parametrize(
+    "use_index",
+    (False, True),
+    ids=["no_index", "with_index"],
+)
+def test_geo_rtree(test_geo_dataset, benchmark, use_index):
+    if use_index:
+        test_geo_dataset.create_scalar_index(
+            column="points",
+            index_type="RTREE",
+            replace=True,
+        )
+
+    print(
+        test_geo_dataset.scanner(
+            columns=["points"],
+            filter="""
+                St_Contains(points,
+                ST_GeomFromText('POLYGON (( 0 0, 2 0, 0 2, 2 2, 0 0 ))'))
+            """,
+            batch_size=32,
+            use_scalar_index=use_index,
+        ).explain_plan(True)
+    )
+    benchmark(
+        test_geo_dataset.to_table,
+        columns=["points"],
+        filter="""
+            St_Contains(points,
+            ST_GeomFromText('POLYGON (( 0 0, 2 0, 0 2, 2 2, 0 0 ))'))
+        """,
+        batch_size=32,
+        use_scalar_index=use_index,
+    )

rust/lance-index/benches/geo.rs

@@ -111,7 +111,7 @@ async fn rect_search_rtree(
 
 fn bench_rtree(c: &mut Criterion) {
     let rt = tokio::runtime::Builder::new_multi_thread().build().unwrap();
-    let num_points = 1_000_000;
+    let num_rows = 1_000_000;
 
     let tempdir = tempfile::tempdir().unwrap();
     let index_dir = Path::from_filesystem_path(tempdir.path()).unwrap();
@@ -123,7 +123,7 @@ fn bench_rtree(c: &mut Criterion) {
         ))
     });
 
-    let geo_data = rt.block_on(async { black_box(generate_geo_data(num_points, 42)) });
+    let geo_data = rt.block_on(async { black_box(generate_geo_data(num_rows, 42)) });
 
     let mut group = c.benchmark_group("RTree");
     group.sample_size(10);