RTree Buffer protocol, python binding tests (#55)

* Draft buffer protocol (memoryview-mediated, apparently zero-copy), explicit from_buffer, to_buffer methods (latter causes a copy to occur). TODO: add tests equiv to geoarrow-rs

* Add poetry infra, tests, try to avoid segfault on shared memory (unsuccessfully)

* Swap test names

* Rip out unnecessary py03 feature, avoid breaking out build.rs

* Rip out unnecessary py03 annotations, clippy silencing, switch to opendal-esque (basically verbatim) __getbuffer__

* Remove copy-prone to_buffer (superseded in all python versions since ~3.0 by the buffer protocol), from_buffer

* Trim down tests to just buffer imports

* Remove from_buffer, to_buffer in pyi

* Draft pytest (non-matrix) + python-wheels CI conf

* Swap back to fixed python version in publish step (suspect setup-python might not actually be necessary), try out matrixed pytest

* updates

---------

Co-authored-by: Kyle Barron <kyle@developmentseed.org>

.github/workflows/python-wheels.yml

@@ -19,6 +19,7 @@ jobs:
     runs-on: ${{ matrix.platform.runner }}
     strategy:
       matrix:
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
         platform:
           - runner: ubuntu-latest
             target: x86_64
@@ -36,7 +37,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.10"
+          python-version: ${{ matrix.python-version }}
       - name: Build wheels
         uses: PyO3/maturin-action@v1
         with:
@@ -54,6 +55,7 @@ jobs:
     runs-on: ${{ matrix.platform.runner }}
     strategy:
       matrix:
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
         platform:
           - runner: windows-latest
             target: x64
@@ -63,7 +65,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.10"
+          python-version: ${{ matrix.python-version }}
           architecture: ${{ matrix.platform.target }}
       - name: Build wheels
         uses: PyO3/maturin-action@v1
@@ -81,6 +83,7 @@ jobs:
     runs-on: ${{ matrix.platform.runner }}
     strategy:
       matrix:
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
         platform:
           - runner: macos-latest
             target: x86_64
@@ -90,7 +93,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.10"
+          python-version: ${{ matrix.python-version }}
       - name: Build wheels
         uses: PyO3/maturin-action@v1
         with:

.github/workflows/test.yml

@@ -33,3 +33,36 @@ jobs:
         run: |
           cargo test --all
           cargo test --all --all-features
+  pytest:
+    name: Pytest
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install Rust
+        uses: dtolnay/rust-toolchain@stable
+        with:
+          components: rustfmt
+
+      - uses: Swatinem/rust-cache@v2
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install and configure Poetry
+        uses: snok/install-poetry@v1
+        with:
+          version: 1.7.0
+          virtualenvs-create: true
+          virtualenvs-in-project: true
+          installer-parallel: true
+      - name: Install poetry venv
+        working-directory: python
+        run: poetry install
+      - name: maturin venv Build
+        working-directory: python
+        run: poetry run maturin develop
+      - name: Run pytest
+        working-directory: python
+        run: poetry run test

python/Cargo.toml

@@ -25,7 +25,7 @@ geo-index = { path = "../", features = ["rayon"] }
 # This is the fork used by polars
 # https://github.com/pola-rs/polars/blob/fac700d9670feb57f1df32beaeee38377725fccf/py-polars/Cargo.toml#L33-L35
 numpy = { git = "https://github.com/stinodego/rust-numpy.git", rev = "9ba9962ae57ba26e35babdce6f179edf5fe5b9c8", default-features = false }
-pyo3 = { version = "0.21", features = ["abi3-py38"] }
+pyo3 = { version = "0.21.0", features = ["macros"] }
 thiserror = "1"
 
 [profile.release]

python/pyproject.toml

@@ -6,6 +6,7 @@ build-backend = "maturin"
 name = "geoindex-rs"
 requires-python = ">=3.8"
 dependencies = []
+dynamic = ["version"]
 classifiers = [
     "Programming Language :: Rust",
     "Programming Language :: Python :: Implementation :: CPython",
@@ -16,3 +17,22 @@ classifiers = [
 features = ["pyo3/extension-module"]
 module-name = "geoindex_rs._rust"
 python-source = "python"
+
+[tool.poetry]
+name = "geoindex-rs"
+version = "0.0.0"
+description = ""
+authors = []
+packages = [{ include = "python/geoindex_rs" }]
+
+[tool.poetry.dependencies]
+# Note: this is only for dev dependencies
+python = "^3.9"
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.0.0"
+maturin = "^1.4.0"
+numpy = "*"
+
+[tool.poetry.scripts]
+test = "pytest:main"

python/python/geoindex_rs/tests/test_buffers.py

@@ -0,0 +1,17 @@
+import numpy as np
+
+from .. import RTree
+
+
+def generate_random_boxes():
+    random_points = np.random.random_sample((100, 2))
+    boxes = np.concatenate([random_points, random_points], axis=1)
+    return boxes
+
+
+def test_buffer_protocol():
+    boxes = generate_random_boxes()
+    initial = RTree.from_interleaved(boxes)
+    # construct a memoryview transparently
+    view = memoryview(initial)
+    assert initial.num_bytes == view.nbytes

python/src/kdtree.rs

@@ -87,7 +87,8 @@ impl KDTree {
         max_x: f64,
         max_y: f64,
     ) -> Bound<'py, PyArray1<usize>> {
-        let result = py.allow_threads(move || self.0.as_ref().range(min_x, min_y, max_x, max_y));
+        let result =
+            py.allow_threads(move || self.0.as_kdtree_ref().range(min_x, min_y, max_x, max_y));
         PyArray1::from_vec_bound(py, result)
     }
 
@@ -105,7 +106,7 @@ impl KDTree {
         qy: f64,
         r: f64,
     ) -> Bound<'py, PyArray1<usize>> {
-        let result = py.allow_threads(move || self.0.as_ref().within(qx, qy, r));
+        let result = py.allow_threads(move || self.0.as_kdtree_ref().within(qx, qy, r));
         PyArray1::from_vec_bound(py, result)
     }
 }