From 1f92806bd4e948d3fa2db54dd32082187db405c7 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 21 Aug 2024 11:15:37 -0400 Subject: [PATCH 01/12] Import via buffer protocol --- python/Cargo.toml | 2 +- python/src/rtree.rs | 198 +++++++++++++++++++++++++++++++++++++++++++- src/rtree/index.rs | 15 +++- src/rtree/mod.rs | 2 +- 4 files changed, 212 insertions(+), 5 deletions(-) diff --git a/python/Cargo.toml b/python/Cargo.toml index 3f67a71..9349016 100644 --- a/python/Cargo.toml +++ b/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", features = [] } thiserror = "1" [profile.release] diff --git a/python/src/rtree.rs b/python/src/rtree.rs index 339565d..4e6adb7 100644 --- a/python/src/rtree.rs +++ b/python/src/rtree.rs @@ -1,9 +1,11 @@ +use geo_index::indices::Indices; use geo_index::rtree::sort::{HilbertSort, STRSort}; use geo_index::rtree::util::f64_box_to_f32; -use geo_index::rtree::{OwnedRTree, RTreeBuilder, RTreeIndex}; -use geo_index::IndexableNum; +use geo_index::rtree::{OwnedRTree, RTreeBuilder, RTreeIndex, TreeMetadata}; +use geo_index::{CoordType, IndexableNum}; use numpy::ndarray::{ArrayView1, ArrayView2}; use numpy::{PyArray1, PyArrayMethods, PyReadonlyArray1, PyReadonlyArray2}; +use pyo3::buffer::PyBuffer; use pyo3::exceptions::{PyIndexError, PyTypeError, PyValueError}; use pyo3::intern; use pyo3::prelude::*; @@ -27,6 +29,198 @@ impl<'a> FromPyObject<'a> for RTreeMethod { } } +struct PyU8Buffer(PyBuffer); + +impl<'py> FromPyObject<'py> for PyU8Buffer { + fn extract_bound(obj: &Bound<'py, PyAny>) -> PyResult { + let buffer = PyBuffer::::get_bound(obj)?; + if !buffer.readonly() { + return Err(PyValueError::new_err("Must be read-only byte buffer.")); + } + if buffer.dimensions() != 1 { + return Err(PyValueError::new_err("Expected 1-dimensional array.")); + } + // Note: this is probably superfluous for 1D array + if !buffer.is_c_contiguous() { + return Err(PyValueError::new_err("Expected c-contiguous array.")); + } + if buffer.len_bytes() == 0 { + return Err(PyValueError::new_err("Buffer has no data.")); + } + + Ok(Self(buffer)) + } +} + +impl AsRef<[u8]> for PyU8Buffer { + fn as_ref(&self) -> &[u8] { + let len = self.0.item_count(); + let data = self.0.buf_ptr() as *const u8; + unsafe { std::slice::from_raw_parts(data, len) } + } +} + +struct Pyf64RTreeRef { + buffer: PyU8Buffer, + metadata: TreeMetadata, +} + +impl Pyf64RTreeRef { + fn try_new(buffer: PyU8Buffer) -> PyResult { + let metadata = TreeMetadata::try_new(buffer.as_ref()).unwrap(); + Ok(Self { buffer, metadata }) + } +} + +impl AsRef<[u8]> for Pyf64RTreeRef { + fn as_ref(&self) -> &[u8] { + self.buffer.as_ref() + } +} + +impl RTreeIndex for Pyf64RTreeRef { + fn boxes(&self) -> &[f64] { + self.metadata.boxes_slice(self.as_ref()) + } + + fn indices(&self) -> Indices { + self.metadata.indices_slice(self.as_ref()) + } + + fn level_bounds(&self) -> &[usize] { + self.metadata.level_bounds() + } + + fn node_size(&self) -> usize { + self.metadata.node_size() + } + + fn num_items(&self) -> usize { + self.metadata.num_items() + } + + fn num_nodes(&self) -> usize { + self.metadata.num_nodes() + } +} + +struct Pyf32RTreeRef { + buffer: PyU8Buffer, + metadata: TreeMetadata, +} + +impl Pyf32RTreeRef { + fn try_new(buffer: PyU8Buffer) -> PyResult { + let metadata = TreeMetadata::try_new(buffer.as_ref()).unwrap(); + Ok(Self { buffer, metadata }) + } +} + +impl AsRef<[u8]> for Pyf32RTreeRef { + fn as_ref(&self) -> &[u8] { + self.buffer.as_ref() + } +} + +impl RTreeIndex for Pyf32RTreeRef { + fn boxes(&self) -> &[f32] { + self.metadata.boxes_slice(self.as_ref()) + } + + fn indices(&self) -> Indices { + self.metadata.indices_slice(self.as_ref()) + } + + fn level_bounds(&self) -> &[usize] { + self.metadata.level_bounds() + } + + fn node_size(&self) -> usize { + self.metadata.node_size() + } + + fn num_items(&self) -> usize { + self.metadata.num_items() + } + + fn num_nodes(&self) -> usize { + self.metadata.num_nodes() + } +} + +enum PyRTreeRef { + Float32(Pyf32RTreeRef), + Float64(Pyf64RTreeRef), +} + +impl<'py> FromPyObject<'py> for PyRTreeRef { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + let buffer = PyU8Buffer::extract_bound(ob)?; + let ct = CoordType::from_buffer(&buffer.as_ref()).unwrap(); + match ct { + CoordType::Float32 => Ok(Self::Float32(Pyf32RTreeRef::try_new(buffer)?)), + CoordType::Float64 => Ok(Self::Float64(Pyf64RTreeRef::try_new(buffer)?)), + _ => todo!(), + } + } +} + +impl PyRTreeRef { + fn num_items(&self) -> usize { + match self { + Self::Float32(index) => index.num_items(), + Self::Float64(index) => index.num_items(), + } + } + + fn num_nodes(&self) -> usize { + match self { + Self::Float32(index) => index.num_nodes(), + Self::Float64(index) => index.num_nodes(), + } + } + + fn node_size(&self) -> usize { + match self { + Self::Float32(index) => index.node_size(), + Self::Float64(index) => index.node_size(), + } + } + + fn num_levels(&self) -> usize { + match self { + Self::Float32(index) => index.num_levels(), + Self::Float64(index) => index.num_levels(), + } + } + + fn num_bytes(&self) -> usize { + match self { + Self::Float32(index) => AsRef::as_ref(index).len(), + Self::Float64(index) => AsRef::as_ref(index).len(), + } + } + + fn boxes_at_level<'py>(&'py self, py: Python<'py>, level: usize) -> PyResult { + match self { + Self::Float32(index) => { + let boxes = index + .boxes_at_level(level) + .map_err(|err| PyIndexError::new_err(err.to_string()))?; + let array = PyArray1::from_slice_bound(py, boxes); + Ok(array.reshape([boxes.len() / 4, 4])?.into_py(py)) + } + Self::Float64(index) => { + let boxes = index + .boxes_at_level(level) + .map_err(|err| PyIndexError::new_err(err.to_string()))?; + let array = PyArray1::from_slice_bound(py, boxes); + Ok(array.reshape([boxes.len() / 4, 4])?.into_py(py)) + } + } + } +} + enum RTreeInner { Float32(OwnedRTree), Float64(OwnedRTree), diff --git a/src/rtree/index.rs b/src/rtree/index.rs index b7b00ad..d2a2369 100644 --- a/src/rtree/index.rs +++ b/src/rtree/index.rs @@ -12,7 +12,7 @@ use crate::rtree::util::compute_num_nodes; /// Common metadata to describe a tree #[derive(Debug, Clone, PartialEq)] -pub(crate) struct TreeMetadata { +pub struct TreeMetadata { pub(crate) node_size: usize, pub(crate) num_items: usize, pub(crate) num_nodes: usize, @@ -112,6 +112,19 @@ impl TreeMetadata { [8 + self.nodes_byte_length..8 + self.nodes_byte_length + self.indices_byte_length]; Indices::new(indices_buf, self.num_nodes) } + + pub fn node_size(&self) -> usize { + self.node_size + } + pub fn num_items(&self) -> usize { + self.num_items + } + pub fn num_nodes(&self) -> usize { + self.num_nodes + } + pub fn level_bounds(&self) -> &[usize] { + self.level_bounds.as_slice() + } } /// An owned RTree buffer. diff --git a/src/rtree/mod.rs b/src/rtree/mod.rs index 8c973e6..f5c8540 100644 --- a/src/rtree/mod.rs +++ b/src/rtree/mod.rs @@ -9,5 +9,5 @@ pub mod traversal; pub mod util; pub use builder::RTreeBuilder; -pub use index::{OwnedRTree, RTreeRef}; +pub use index::{OwnedRTree, RTreeRef, TreeMetadata}; pub use r#trait::RTreeIndex; From 9db47b9db3634445544a6ec8ef4e0ebe049990b3 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 21 Aug 2024 11:19:16 -0400 Subject: [PATCH 02/12] add search function --- python/src/lib.rs | 2 ++ python/src/rtree.rs | 32 +++++++++++++++++++++++++++++++- 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/python/src/lib.rs b/python/src/lib.rs index 6145ac5..dfb6691 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -14,6 +14,8 @@ fn ___version() -> &'static str { fn _rust(_py: Python, m: &Bound) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(___version))?; + m.add_wrapped(wrap_pyfunction!(rtree::search_rtree))?; + m.add_class::()?; m.add_class::()?; diff --git a/python/src/rtree.rs b/python/src/rtree.rs index 4e6adb7..d71149c 100644 --- a/python/src/rtree.rs +++ b/python/src/rtree.rs @@ -148,7 +148,7 @@ impl RTreeIndex for Pyf32RTreeRef { } } -enum PyRTreeRef { +pub(crate) enum PyRTreeRef { Float32(Pyf32RTreeRef), Float64(Pyf64RTreeRef), } @@ -221,6 +221,36 @@ impl PyRTreeRef { } } +/// Search an RTree given the provided bounding box. +/// +/// Results are the indexes of the inserted objects in insertion order. +/// +/// Args: +/// tree: tree or buffer to search +/// min_x: min x coordinate of bounding box +/// min_y: min y coordinate of bounding box +/// max_x: max x coordinate of bounding box +/// max_y: max y coordinate of bounding box +#[pyfunction] +pub(crate) fn search_rtree( + py: Python, + tree: PyRTreeRef, + min_x: f64, + min_y: f64, + max_x: f64, + max_y: f64, +) -> Bound<'_, PyArray1> { + let result = py.allow_threads(|| match tree { + PyRTreeRef::Float32(tree) => { + let (min_x, min_y, max_x, max_y) = f64_box_to_f32(min_x, min_y, max_x, max_y); + tree.search(min_x, min_y, max_x, max_y) + } + PyRTreeRef::Float64(tree) => tree.search(min_x, min_y, max_x, max_y), + }); + + PyArray1::from_vec_bound(py, result) +} + enum RTreeInner { Float32(OwnedRTree), Float64(OwnedRTree), From a85510883d45fb113f67cc5891c1d164d6228653 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 21 Aug 2024 11:27:10 -0400 Subject: [PATCH 03/12] class --- python/src/rtree.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/src/rtree.rs b/python/src/rtree.rs index d71149c..d133704 100644 --- a/python/src/rtree.rs +++ b/python/src/rtree.rs @@ -221,6 +221,10 @@ impl PyRTreeRef { } } +// This will take the place of RTree +#[pyclass] +struct RTreeRefWrapper(PyRTreeRef); + /// Search an RTree given the provided bounding box. /// /// Results are the indexes of the inserted objects in insertion order. From d2df0859839b7a7cbd7fadb2950ba2cdaec320ea Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 21 Aug 2024 12:14:53 -0400 Subject: [PATCH 04/12] buffer protocol import --- Cargo.lock | 350 ++++++++++++++++++------------------------- python/src/lib.rs | 1 + python/src/rtree.rs | 356 +++++++++++++++++++++++++++++++++++--------- src/rtree/index.rs | 4 + 4 files changed, 443 insertions(+), 268 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7c9c798..4c2c598 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4,9 +4,9 @@ version = 3 [[package]] name = "aho-corasick" -version = "1.1.2" +version = "1.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2969dcb958b36655471fc61f7e416fa76033bdd4bfed0678d8fee1e2d07a1f0" +checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" dependencies = [ "memchr", ] @@ -19,33 +19,27 @@ checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" [[package]] name = "anstyle" -version = "1.0.4" +version = "1.0.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7079075b41f533b8c61d2a4d073c4676e1f8b249ff94a393b0595db304e0dd87" +checksum = "1bec1de6f59aedf83baf9ff929c98f2ad654b97c9510f4e70cf6f661d49fd5b1" [[package]] name = "autocfg" -version = "1.1.0" +version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" - -[[package]] -name = "bitflags" -version = "2.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4682ae6287fcf752ecaabbfcc7b6f9b72aa33933dc23a554d853aea8eea8635" +checksum = "0c4b4d0bd25bd0b74681c0ad21497610ce1b7c91b1022cd21c80c6fbdd9476b0" [[package]] name = "bumpalo" -version = "3.14.0" +version = "3.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f30e7476521f6f8af1a1c4c0b8cc94f0bee37d91763d0ca2665f299b6cd8aec" +checksum = "79296716171880943b8470b5f8d03aa55eb2e645a4874bdbb28adb49162e012c" [[package]] name = "bytemuck" -version = "1.14.0" +version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "374d28ec25809ee0e23827c2ab573d729e293f281dfe393500e7ad618baa61c6" +checksum = "6fd4c6dcc3b0aea2f5c0b4b82c2b15fe39ddbc76041a310848f4706edf76bb31" [[package]] name = "byteorder" @@ -67,9 +61,9 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" [[package]] name = "ciborium" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "effd91f6c78e5a4ace8a5d3c0b6bfaec9e2baaef55f3efc00e45fb2e477ee926" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" dependencies = [ "ciborium-io", "ciborium-ll", @@ -78,15 +72,15 @@ dependencies = [ [[package]] name = "ciborium-io" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cdf919175532b369853f5d5e20b26b43112613fd6fe7aee757e35f7a44642656" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" [[package]] name = "ciborium-ll" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "defaa24ecc093c77630e6c15e17c51f5e187bf35ee514f4e2d67baaa96dae22b" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" dependencies = [ "ciborium-io", "half", @@ -94,18 +88,18 @@ dependencies = [ [[package]] name = "clap" -version = "4.4.6" +version = "4.5.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d04704f56c2cde07f43e8e2c154b43f216dc5c92fc98ada720177362f953b956" +checksum = "ed6719fffa43d0d87e5fd8caeab59be1554fb028cd30edc88fc4369b17971019" dependencies = [ "clap_builder", ] [[package]] name = "clap_builder" -version = "4.4.6" +version = "4.5.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e231faeaca65ebd1ea3c737966bf858971cd38c3849107aa3ea7de90a804e45" +checksum = "216aec2b177652e3846684cbfe25c9964d18ec45234f0f5da5157b207ed1aab6" dependencies = [ "anstyle", "clap_lex", @@ -113,9 +107,9 @@ dependencies = [ [[package]] name = "clap_lex" -version = "0.5.1" +version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd7cc57abe963c6d3b9d8be5b06ba7c8957a930305ca90304f24ef040aa6f961" +checksum = "1462739cb27611015575c0c11df5df7601141071f07518d56fcc1be504cbec97" [[package]] name = "criterion" @@ -155,52 +149,40 @@ dependencies = [ [[package]] name = "crossbeam-deque" -version = "0.8.3" +version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ce6fd6f855243022dcecf8702fef0c297d4338e226845fe067f6341ad9fa0cef" +checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d" dependencies = [ - "cfg-if", "crossbeam-epoch", "crossbeam-utils", ] [[package]] name = "crossbeam-epoch" -version = "0.9.15" +version = "0.9.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae211234986c545741a7dc064309f67ee1e5ad243d0e48335adc0484d960bcc7" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" dependencies = [ - "autocfg", - "cfg-if", "crossbeam-utils", - "memoffset", - "scopeguard", ] [[package]] name = "crossbeam-utils" -version = "0.8.16" +version = "0.8.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a22b2d63d4d1dc0b7f1b6b2747dd0088008a9be28b6ddf0b1e7d335e3037294" -dependencies = [ - "cfg-if", -] +checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" [[package]] -name = "either" -version = "1.9.0" +name = "crunchy" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a26ae43d7bcc3b814de94796a5e736d4029efb0ee900c12e2d54c993ad1a1e07" +checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" [[package]] -name = "errno" -version = "0.3.5" +name = "either" +version = "1.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac3e13f66a2f95e32a39eaa81f6b95d42878ca0e1db0c7543723dfe12557e860" -dependencies = [ - "libc", - "windows-sys", -] +checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" [[package]] name = "float_next_after" @@ -224,9 +206,13 @@ dependencies = [ [[package]] name = "half" -version = "1.8.2" +version = "2.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eabb4a44450da02c90444cf74558da904edde8fb4e9035a9a6a4e15445af0bd7" +checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888" +dependencies = [ + "cfg-if", + "crunchy", +] [[package]] name = "hash32" @@ -249,19 +235,19 @@ dependencies = [ [[package]] name = "hermit-abi" -version = "0.3.3" +version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d77f7ec81a6d05a3abb01ab6eb7590f6083d08449fe5a1c8b1e620283546ccb7" +checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc" [[package]] name = "is-terminal" -version = "0.4.9" +version = "0.4.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb0889898416213fab133e1d33a0e5858a48177452750691bde3666d0fdbaf8b" +checksum = "261f68e344040fbd0edea105bef17c66edf46f984ddb1115b775ce31be948f4b" dependencies = [ "hermit-abi", - "rustix", - "windows-sys", + "libc", + "windows-sys 0.52.0", ] [[package]] @@ -275,24 +261,24 @@ dependencies = [ [[package]] name = "itoa" -version = "1.0.9" +version = "1.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af150ab688ff2122fcef229be89cb50dd66af9e01a4ff320cc137eecc9bacc38" +checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" [[package]] name = "js-sys" -version = "0.3.64" +version = "0.3.70" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c5f195fe497f702db0f318b07fdd68edb16955aed830df8363d837542f8f935a" +checksum = "1868808506b929d7b0cfa8f75951347aa71bb21144b7791bae35d9bccfcfe37a" dependencies = [ "wasm-bindgen", ] [[package]] name = "libc" -version = "0.2.149" +version = "0.2.158" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a08173bc88b7955d1b3145aa561539096c421ac8debde8cbc3612ec635fee29b" +checksum = "d8adc4bb1803a324070e64a98ae98f38934d91957a99cfb3a43dcbc01bc56439" [[package]] name = "libm" @@ -300,38 +286,23 @@ version = "0.2.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058" -[[package]] -name = "linux-raw-sys" -version = "0.4.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da2479e8c062e40bf0066ffa0bc823de0a9368974af99c9f6df941d2c231e03f" - [[package]] name = "log" -version = "0.4.20" +version = "0.4.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f" +checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24" [[package]] name = "memchr" -version = "2.6.4" +version = "2.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f665ee40bc4a3c5590afb1e9677db74a508659dfd71e126420da8274909a0167" - -[[package]] -name = "memoffset" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a634b1c61a95585bd15607c6ab0c4e5b226e695ff2800ba0cdccddf208c406c" -dependencies = [ - "autocfg", -] +checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" [[package]] name = "num-traits" -version = "0.2.17" +version = "0.2.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", "libm", @@ -339,21 +310,21 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.18.0" +version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d" +checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" [[package]] name = "oorandom" -version = "11.1.3" +version = "11.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ab1bc2a289d34bd04a330323ac98a1b4bc82c9d9fcb1e66b63caa84da26b575" +checksum = "b410bbe7e14ab526a0e86877eb47c6996a2bd7746f027ba551028c925390e4e9" [[package]] name = "plotters" -version = "0.3.5" +version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2c224ba00d7cadd4d5c660deaf2098e5e80e07846537c51f9cfa4be50c1fd45" +checksum = "a15b6eccb8484002195a3e44fe65a4ce8e93a625797a063735536fd59cb01cf3" dependencies = [ "num-traits", "plotters-backend", @@ -364,42 +335,42 @@ dependencies = [ [[package]] name = "plotters-backend" -version = "0.3.5" +version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e76628b4d3a7581389a35d5b6e2139607ad7c75b17aed325f210aa91f4a9609" +checksum = "414cec62c6634ae900ea1c56128dfe87cf63e7caece0852ec76aba307cebadb7" [[package]] name = "plotters-svg" -version = "0.3.5" +version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38f6d39893cca0701371e3c27294f09797214b86f1fb951b89ade8ec04e2abab" +checksum = "81b30686a7d9c3e010b84284bdd26a29f2138574f52f5eb6f794fc0ad924e705" dependencies = [ "plotters-backend", ] [[package]] name = "proc-macro2" -version = "1.0.69" +version = "1.0.86" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "134c189feb4956b20f6f547d2cf727d4c0fe06722b20a0eec87ed445a97f92da" +checksum = "5e719e8df665df0d1c8fbfd238015744736151d4445ec0836b8e628aae103b77" dependencies = [ "unicode-ident", ] [[package]] name = "quote" -version = "1.0.33" +version = "1.0.36" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5267fca4496028628a95160fc423a33e8b2e6af8a5302579e322e4b520293cae" +checksum = "0fa76aaf39101c457836aec0ce2316dbdc3ab723cdda1c6bd4e6ad4208acaca7" dependencies = [ "proc-macro2", ] [[package]] name = "rayon" -version = "1.8.0" +version = "1.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c27db03db7734835b3f53954b534c91069375ce6ccaa2e065441e07d9b6cdb1" +checksum = "b418a60154510ca1a002a752ca9714984e21e4241e804d32555251faf8b78ffa" dependencies = [ "either", "rayon-core", @@ -407,9 +378,9 @@ dependencies = [ [[package]] name = "rayon-core" -version = "1.12.0" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ce3fb6ad83f861aac485e76e1985cd109d9a3713802152be56c3b1f0e0658ed" +checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" dependencies = [ "crossbeam-deque", "crossbeam-utils", @@ -417,9 +388,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.10.1" +version = "1.10.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aaac441002f822bc9705a681810a4dd2963094b9ca0ddc41cb963a4c189189ea" +checksum = "4219d74c6b67a3654a9fbebc4b419e22126d13d2f3c4a07ee0cb61ff79a79619" dependencies = [ "aho-corasick", "memchr", @@ -429,9 +400,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.2" +version = "0.4.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5011c7e263a695dc8ca064cddb722af1be54e517a280b12a5356f98366899e5d" +checksum = "38caf58cc5ef2fed281f89292ef23f6365465ed9a41b7a7754eb4e26496c92df" dependencies = [ "aho-corasick", "memchr", @@ -440,9 +411,9 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.8.2" +version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c08c74e62047bb2de4ff487b251e4a92e24f48745648451635cec7d591162d9f" +checksum = "7a66a03ae7c801facd77a29370b4faec201768915ac14a721ba36f20bc9c209b" [[package]] name = "rstar" @@ -455,24 +426,11 @@ dependencies = [ "smallvec", ] -[[package]] -name = "rustix" -version = "0.38.19" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "745ecfa778e66b2b63c88a61cb36e0eea109e803b0b86bf9879fbc77c70e86ed" -dependencies = [ - "bitflags", - "errno", - "libc", - "linux-raw-sys", - "windows-sys", -] - [[package]] name = "ryu" -version = "1.0.15" +version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ad4cc8da4ef723ed60bced201181d83791ad433213d8c24efffda1eec85d741" +checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" [[package]] name = "same-file" @@ -483,26 +441,20 @@ dependencies = [ "winapi-util", ] -[[package]] -name = "scopeguard" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" - [[package]] name = "serde" -version = "1.0.189" +version = "1.0.208" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e422a44e74ad4001bdc8eede9a4570ab52f71190e9c076d14369f38b9200537" +checksum = "cff085d2cb684faa248efb494c39b68e522822ac0de72ccf08109abde717cfb2" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.189" +version = "1.0.208" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e48d1f918009ce3145511378cf68d613e3b3d9137d67272562080d68a2b32d5" +checksum = "24008e81ff7613ed8e5ba0cfaf24e2c2f1e5b8a0495711e44fcd4882fca62bcf" dependencies = [ "proc-macro2", "quote", @@ -511,20 +463,21 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.107" +version = "1.0.125" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b420ce6e3d8bd882e9b243c6eed35dbc9a6110c9769e74b584e0d68d1f20c65" +checksum = "83c8e735a073ccf5be70aa8066aa984eaf2fa000db6c8d0100ae605b366d31ed" dependencies = [ "itoa", + "memchr", "ryu", "serde", ] [[package]] name = "smallvec" -version = "1.11.1" +version = "1.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "942b4a808e05215192e39f4ab80813e599068285906cc91aa64f923db842bd5a" +checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" [[package]] name = "stable_deref_trait" @@ -534,9 +487,9 @@ checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" [[package]] name = "syn" -version = "2.0.38" +version = "2.0.75" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e96b79aaa137db8f61e26363a0c9b47d8b4ec75da28b7d1d614c2303e232408b" +checksum = "f6af063034fc1935ede7be0122941bafa9bacb949334d090b77ca98b5817c7d9" dependencies = [ "proc-macro2", "quote", @@ -545,18 +498,18 @@ dependencies = [ [[package]] name = "thiserror" -version = "1.0.49" +version = "1.0.63" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1177e8c6d7ede7afde3585fd2513e611227efd6481bd78d2e82ba1ce16557ed4" +checksum = "c0342370b38b6a11b6cc11d6a805569958d54cfa061a29969c3b5ce2ea405724" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.49" +version = "1.0.63" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "10712f02019e9288794769fba95cd6847df9874d49d871d062172f9dd41bc4cc" +checksum = "a4558b58466b9ad7ca0f102865eccc95938dca1a74a856f2b57b6629050da261" dependencies = [ "proc-macro2", "quote", @@ -575,9 +528,9 @@ dependencies = [ [[package]] name = "tinyvec" -version = "1.6.0" +version = "1.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "87cc5ceb3875bb20c2890005a4e226a4651264a5c75edb2421b52861a0a0cb50" +checksum = "445e881f4f6d382d5f27c034e25eb92edd7c784ceab92a0937db7f2e9471b938" dependencies = [ "tinyvec_macros", ] @@ -596,9 +549,9 @@ checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" [[package]] name = "walkdir" -version = "2.4.0" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d71d857dc86794ca4c280d616f7da00d2dbfd8cd788846559a6813e6aa4b54ee" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" dependencies = [ "same-file", "winapi-util", @@ -606,19 +559,20 @@ dependencies = [ [[package]] name = "wasm-bindgen" -version = "0.2.87" +version = "0.2.93" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7706a72ab36d8cb1f80ffbf0e071533974a60d0a308d01a5d0375bf60499a342" +checksum = "a82edfc16a6c469f5f44dc7b571814045d60404b55a0ee849f9bcfa2e63dd9b5" dependencies = [ "cfg-if", + "once_cell", "wasm-bindgen-macro", ] [[package]] name = "wasm-bindgen-backend" -version = "0.2.87" +version = "0.2.93" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ef2b6d3c510e9625e5fe6f509ab07d66a760f0885d858736483c32ed7809abd" +checksum = "9de396da306523044d3302746f1208fa71d7532227f15e347e2d93e4145dd77b" dependencies = [ "bumpalo", "log", @@ -631,9 +585,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.87" +version = "0.2.93" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dee495e55982a3bd48105a7b947fd2a9b4a8ae3010041b9e0faab3f9cd028f1d" +checksum = "585c4c91a46b072c92e908d99cb1dcdf95c5218eeb6f3bf1efa991ee7a68cccf" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -641,9 +595,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.87" +version = "0.2.93" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54681b18a46765f095758388f2d0cf16eb8d4169b639ab575a8f5693af210c7b" +checksum = "afc340c74d9005395cf9dd098506f7f44e38f2b4a21c6aaacf9a105ea5e1e836" dependencies = [ "proc-macro2", "quote", @@ -654,69 +608,57 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.87" +version = "0.2.93" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ca6ad05a4870b2bf5fe995117d3728437bd27d7cd5f06f13c17443ef369775a1" +checksum = "c62a0a307cb4a311d3a07867860911ca130c3494e8c2719593806c08bc5d0484" [[package]] name = "web-sys" -version = "0.3.64" +version = "0.3.70" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b85cbef8c220a6abc02aefd892dfc0fc23afb1c6a426316ec33253a3877249b" +checksum = "26fdeaafd9bd129f65e7c031593c24d62186301e0c72c8978fa1678be7d532c0" dependencies = [ "js-sys", "wasm-bindgen", ] [[package]] -name = "winapi" -version = "0.3.9" +name = "winapi-util" +version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "winapi-i686-pc-windows-gnu", - "winapi-x86_64-pc-windows-gnu", + "windows-sys 0.59.0", ] [[package]] -name = "winapi-i686-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" - -[[package]] -name = "winapi-util" -version = "0.1.6" +name = "windows-sys" +version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f29e6f9198ba0d26b4c9f07dbe6f9ed633e1f3d5b8b414090084349e46a52596" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "winapi", + "windows-targets", ] -[[package]] -name = "winapi-x86_64-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" - [[package]] name = "windows-sys" -version = "0.48.0" +version = "0.59.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" dependencies = [ "windows-targets", ] [[package]] name = "windows-targets" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ "windows_aarch64_gnullvm", "windows_aarch64_msvc", "windows_i686_gnu", + "windows_i686_gnullvm", "windows_i686_msvc", "windows_x86_64_gnu", "windows_x86_64_gnullvm", @@ -725,42 +667,48 @@ dependencies = [ [[package]] name = "windows_aarch64_gnullvm" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" [[package]] name = "windows_aarch64_msvc" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" [[package]] name = "windows_i686_gnu" -version = "0.48.5" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" [[package]] name = "windows_i686_msvc" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" [[package]] name = "windows_x86_64_gnu" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" [[package]] name = "windows_x86_64_gnullvm" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" [[package]] name = "windows_x86_64_msvc" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" diff --git a/python/src/lib.rs b/python/src/lib.rs index dfb6691..f0d64eb 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -17,6 +17,7 @@ fn _rust(_py: Python, m: &Bound) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(rtree::search_rtree))?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; Ok(()) diff --git a/python/src/rtree.rs b/python/src/rtree.rs index 0b5285e..1a92cf9 100644 --- a/python/src/rtree.rs +++ b/python/src/rtree.rs @@ -13,6 +13,7 @@ use pyo3::prelude::*; use pyo3::types::PyType; use std::os::raw::c_int; +/// Method for constructing rtree pub enum RTreeMethod { Hilbert, STR, @@ -31,6 +32,40 @@ impl<'a> FromPyObject<'a> for RTreeMethod { } } +/// A Rust buffer that implements the Python buffer protocol +#[pyclass(name = "Buffer")] +struct RustBuffer(Vec); + +#[pymethods] +impl RustBuffer { + // pre PEP 688 buffer protocol + unsafe fn __getbuffer__( + slf: PyRef<'_, Self>, + view: *mut ffi::Py_buffer, + flags: c_int, + ) -> PyResult<()> { + let bytes = slf.0.as_slice(); + let ret = ffi::PyBuffer_FillInfo( + view, + slf.as_ptr() as *mut _, + bytes.as_ptr() as *mut _, + bytes.len().try_into().unwrap(), + 1, // read only + flags, + ); + if ret == -1 { + return Err(PyErr::fetch(slf.py())); + } + Ok(()) + } + + unsafe fn __releasebuffer__(&self, _view: *mut ffi::Py_buffer) { + // is there anything to do here? + } +} + +/// A Rust representation of a Python object that implements the Python buffer protocol, exporting +/// a 1-dimensional `&[u8]` slice. struct PyU8Buffer(PyBuffer); impl<'py> FromPyObject<'py> for PyU8Buffer { @@ -62,70 +97,37 @@ impl AsRef<[u8]> for PyU8Buffer { } } -struct Pyf64RTreeRef { +/// A low-level wrapper around a [PyU8Buffer] that validates that the input is a valid Flatbush +/// buffer. This wrapper implements [RTreeIndex]. +pub(crate) struct PyRTreeRefInner { buffer: PyU8Buffer, - metadata: TreeMetadata, + metadata: TreeMetadata, } -impl Pyf64RTreeRef { +impl PyRTreeRefInner { fn try_new(buffer: PyU8Buffer) -> PyResult { let metadata = TreeMetadata::try_new(buffer.as_ref()).unwrap(); Ok(Self { buffer, metadata }) } -} - -impl AsRef<[u8]> for Pyf64RTreeRef { - fn as_ref(&self) -> &[u8] { - self.buffer.as_ref() - } -} - -impl RTreeIndex for Pyf64RTreeRef { - fn boxes(&self) -> &[f64] { - self.metadata.boxes_slice(self.as_ref()) - } - fn indices(&self) -> Indices { - self.metadata.indices_slice(self.as_ref()) - } - - fn level_bounds(&self) -> &[usize] { - self.metadata.level_bounds() - } - - fn node_size(&self) -> usize { - self.metadata.node_size() - } - - fn num_items(&self) -> usize { - self.metadata.num_items() - } - - fn num_nodes(&self) -> usize { - self.metadata.num_nodes() + fn from_owned_rtree(py: Python, tree: OwnedRTree) -> PyResult { + let metadata = tree.metadata().clone(); + let tree_buf = RustBuffer(tree.into_inner()); + Ok(Self { + buffer: tree_buf.into_py(py).extract(py)?, + metadata, + }) } } -struct Pyf32RTreeRef { - buffer: PyU8Buffer, - metadata: TreeMetadata, -} - -impl Pyf32RTreeRef { - fn try_new(buffer: PyU8Buffer) -> PyResult { - let metadata = TreeMetadata::try_new(buffer.as_ref()).unwrap(); - Ok(Self { buffer, metadata }) - } -} - -impl AsRef<[u8]> for Pyf32RTreeRef { +impl AsRef<[u8]> for PyRTreeRefInner { fn as_ref(&self) -> &[u8] { self.buffer.as_ref() } } -impl RTreeIndex for Pyf32RTreeRef { - fn boxes(&self) -> &[f32] { +impl RTreeIndex for PyRTreeRefInner { + fn boxes(&self) -> &[N] { self.metadata.boxes_slice(self.as_ref()) } @@ -150,9 +152,11 @@ impl RTreeIndex for Pyf32RTreeRef { } } +/// An enum wrapper around [PyRTreeRefInner] that allows use of multiple coordinate types from +/// Python. pub(crate) enum PyRTreeRef { - Float32(Pyf32RTreeRef), - Float64(Pyf64RTreeRef), + Float32(PyRTreeRefInner), + Float64(PyRTreeRefInner), } impl<'py> FromPyObject<'py> for PyRTreeRef { @@ -160,13 +164,34 @@ impl<'py> FromPyObject<'py> for PyRTreeRef { let buffer = PyU8Buffer::extract_bound(ob)?; let ct = CoordType::from_buffer(&buffer.as_ref()).unwrap(); match ct { - CoordType::Float32 => Ok(Self::Float32(Pyf32RTreeRef::try_new(buffer)?)), - CoordType::Float64 => Ok(Self::Float64(Pyf64RTreeRef::try_new(buffer)?)), + CoordType::Float32 => Ok(Self::Float32(PyRTreeRefInner::try_new(buffer)?)), + CoordType::Float64 => Ok(Self::Float64(PyRTreeRefInner::try_new(buffer)?)), _ => todo!(), } } } +impl From> for PyRTreeRef { + fn from(value: PyRTreeRefInner) -> Self { + Self::Float32(value) + } +} + +impl From> for PyRTreeRef { + fn from(value: PyRTreeRefInner) -> Self { + Self::Float64(value) + } +} + +impl AsRef<[u8]> for PyRTreeRef { + fn as_ref(&self) -> &[u8] { + match self { + Self::Float32(inner) => inner.as_ref(), + Self::Float64(inner) => inner.as_ref(), + } + } +} + impl PyRTreeRef { fn num_items(&self) -> usize { match self { @@ -223,9 +248,215 @@ impl PyRTreeRef { } } -// This will take the place of RTree #[pyclass] -struct RTreeRefWrapper(PyRTreeRef); +pub(crate) struct RTree(PyRTreeRef); + +#[pymethods] +impl RTree { + /// Construct an RTree from an existing RTree buffer + #[new] + fn new(py: Python, obj: PyObject) -> PyResult { + Ok(Self(obj.extract(py)?)) + } + + #[classmethod] + #[pyo3( + signature = (boxes, *, method = RTreeMethod::Hilbert, node_size = None), + text_signature = "(boxes, *, method = 'hilbert', node_size = None)") + ] + pub fn from_interleaved( + _cls: &Bound, + py: Python, + boxes: PyObject, + method: RTreeMethod, + node_size: Option, + ) -> PyResult { + // Convert to numpy array (of the same dtype) + let boxes = boxes.call_method0(py, intern!(py, "__array__"))?; + + let result = if let Ok(boxes) = boxes.extract::>(py) { + let boxes = boxes.as_array(); + let tree = py.allow_threads(|| new_interleaved(&boxes, method, node_size)); + Ok(Self(PyRTreeRefInner::from_owned_rtree(py, tree)?.into())) + } else if let Ok(boxes) = boxes.extract::>(py) { + let boxes = boxes.as_array(); + let tree = py.allow_threads(|| new_interleaved(&boxes, method, node_size)); + Ok(Self(PyRTreeRefInner::from_owned_rtree(py, tree)?.into())) + } else { + let dtype = boxes.call_method0(py, intern!(py, "dtype"))?.to_string(); + Err(PyTypeError::new_err(format!( + "Expected a numpy array of dtype float32 or float64, got {}", + dtype + ))) + }; + result + } + + #[classmethod] + #[pyo3( + signature = (min_x, min_y, max_x, max_y, *, method = RTreeMethod::Hilbert, node_size = None), + text_signature = "(min_x, min_y, max_x, max_y, *, method = 'hilbert', node_size = None)") + ] + #[allow(clippy::too_many_arguments)] + pub fn from_separated( + _cls: &Bound, + py: Python, + min_x: PyObject, + min_y: PyObject, + max_x: PyObject, + max_y: PyObject, + method: RTreeMethod, + node_size: Option, + ) -> PyResult { + // Convert to numpy array (of the same dtype) + let min_x = min_x.call_method0(py, intern!(py, "__array__"))?; + let min_y = min_y.call_method0(py, intern!(py, "__array__"))?; + let max_x = max_x.call_method0(py, intern!(py, "__array__"))?; + let max_y = max_y.call_method0(py, intern!(py, "__array__"))?; + + let result = if let Ok(min_x) = min_x.extract::>(py) { + let min_y = min_y.extract::>(py)?; + let max_x = max_x.extract::>(py)?; + let max_y = max_y.extract::>(py)?; + + let min_x_array = min_x.as_array(); + let min_y_array = min_y.as_array(); + let max_x_array = max_x.as_array(); + let max_y_array = max_y.as_array(); + + let tree = py.allow_threads(|| { + new_separated( + min_x_array, + min_y_array, + max_x_array, + max_y_array, + method, + node_size, + ) + }); + Ok(Self(PyRTreeRefInner::from_owned_rtree(py, tree)?.into())) + } else if let Ok(min_x) = min_x.extract::>(py) { + let min_y = min_y.extract::>(py)?; + let max_x = max_x.extract::>(py)?; + let max_y = max_y.extract::>(py)?; + + let min_x_array = min_x.as_array(); + let min_y_array = min_y.as_array(); + let max_x_array = max_x.as_array(); + let max_y_array = max_y.as_array(); + + let tree = py.allow_threads(|| { + new_separated( + min_x_array, + min_y_array, + max_x_array, + max_y_array, + method, + node_size, + ) + }); + Ok(Self(PyRTreeRefInner::from_owned_rtree(py, tree)?.into())) + } else { + let dtype = min_x.call_method0(py, intern!(py, "dtype"))?.to_string(); + Err(PyTypeError::new_err(format!( + "Expected a numpy array of dtype float32 or float64, got {}", + dtype + ))) + }; + result + } + + // pre PEP 688 buffer protocol + unsafe fn __getbuffer__( + slf: PyRef<'_, Self>, + view: *mut ffi::Py_buffer, + flags: c_int, + ) -> PyResult<()> { + let bytes = slf.0.as_ref(); + let ret = ffi::PyBuffer_FillInfo( + view, + slf.as_ptr() as *mut _, + bytes.as_ptr() as *mut _, + bytes.len().try_into().unwrap(), + 1, // read only + flags, + ); + if ret == -1 { + return Err(PyErr::fetch(slf.py())); + } + Ok(()) + } + + unsafe fn __releasebuffer__(&self, _view: *mut ffi::Py_buffer) { + // is there anything to do here? + } + + /// The total number of items contained in this RTree. + #[getter] + pub fn num_items(&self) -> usize { + self.0.num_items() + } + + /// The total number of nodes in this RTree, including both leaf and intermediate nodes. + #[getter] + pub fn num_nodes(&self) -> usize { + self.0.num_nodes() + } + + /// The maximum number of elements in each node. + #[getter] + pub fn node_size(&self) -> usize { + self.0.node_size() + } + + /// The height of the tree + #[getter] + fn num_levels(&self) -> usize { + self.0.num_levels() + } + + /// The number of bytes taken up in memory. + #[getter] + fn num_bytes(&self) -> usize { + self.0.num_bytes() + } + + /// Access the bounding boxes at the given level of the tree. + /// + /// The tree is laid out from bottom to top. Level 0 is the _base_ of the tree. Each integer + /// higher is one level higher of the tree. + fn boxes_at_level<'py>(&'py self, py: Python<'py>, level: usize) -> PyResult { + self.0.boxes_at_level(py, level) + } + + /// Search an RTree given the provided bounding box. + /// + /// Results are the indexes of the inserted objects in insertion order. + /// + /// Args: + /// min_x: min x coordinate of bounding box + /// min_y: min y coordinate of bounding box + /// max_x: max x coordinate of bounding box + /// max_y: max y coordinate of bounding box + pub fn search<'py>( + &'py self, + py: Python<'py>, + min_x: f64, + min_y: f64, + max_x: f64, + max_y: f64, + ) -> Bound<'py, PyArray1> { + let result = py.allow_threads(|| match &self.0 { + PyRTreeRef::Float32(tree) => { + let (min_x, min_y, max_x, max_y) = f64_box_to_f32(min_x, min_y, max_x, max_y); + tree.search(min_x, min_y, max_x, max_y) + } + PyRTreeRef::Float64(tree) => tree.search(min_x, min_y, max_x, max_y), + }); + + PyArray1::from_vec_bound(py, result) + } +} /// Search an RTree given the provided bounding box. /// @@ -394,19 +625,10 @@ impl RTreeInner { } #[pyclass] -pub struct RTree(RTreeInner); - -// TODO: add support for constructing from a buffer. Need to be able to construct (and validate) an -// OwnedRTree -// impl<'a> FromPyObject<'a> for RTree { -// fn extract(ob: &'a PyAny) -> PyResult { -// let s: Vec = ob.extract()?; -// OwnedRTree::from(value) -// } -// } +pub struct RTreeOld(RTreeInner); #[pymethods] -impl RTree { +impl RTreeOld { #[classmethod] #[pyo3( signature = (boxes, *, method = RTreeMethod::Hilbert, node_size = None), @@ -515,7 +737,7 @@ impl RTree { } // pre PEP 688 buffer protocol - pub unsafe fn __getbuffer__( + unsafe fn __getbuffer__( slf: PyRef<'_, Self>, view: *mut ffi::Py_buffer, flags: c_int, @@ -534,7 +756,7 @@ impl RTree { } Ok(()) } - pub unsafe fn __releasebuffer__(&self, _view: *mut ffi::Py_buffer) { + unsafe fn __releasebuffer__(&self, _view: *mut ffi::Py_buffer) { // is there anything to do here? } diff --git a/src/rtree/index.rs b/src/rtree/index.rs index 6f0d4d9..ddd2d1c 100644 --- a/src/rtree/index.rs +++ b/src/rtree/index.rs @@ -153,6 +153,10 @@ impl OwnedRTree { metadata: Cow::Borrowed(&self.metadata), } } + + pub fn metadata(&self) -> &TreeMetadata { + &self.metadata + } } impl AsRef<[u8]> for OwnedRTree { From 2fd5aa45f6f4400dc987186b956f0f70e0a0cd84 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 21 Aug 2024 12:24:43 -0400 Subject: [PATCH 05/12] Update test and typing --- python/python/geoindex_rs/_rust.pyi | 7 ++++++- python/{python/geoindex_rs => }/tests/__init__.py | 0 .../{python/geoindex_rs => }/tests/test_buffers.py | 12 ++++++++---- 3 files changed, 14 insertions(+), 5 deletions(-) rename python/{python/geoindex_rs => }/tests/__init__.py (100%) rename python/{python/geoindex_rs => }/tests/test_buffers.py (57%) diff --git a/python/python/geoindex_rs/_rust.pyi b/python/python/geoindex_rs/_rust.pyi index 835b484..ee0899c 100644 --- a/python/python/geoindex_rs/_rust.pyi +++ b/python/python/geoindex_rs/_rust.pyi @@ -1,4 +1,4 @@ -from typing import Generic, Literal, Optional, Self, TypeVar, Union +from typing import Generic, Literal, Optional, Protocol, Self, TypeVar, Union import numpy as np from numpy.typing import NDArray @@ -32,7 +32,12 @@ class KDTree: # https://stackoverflow.com/a/74634650 T = TypeVar("T", bound=np.generic, covariant=True) +class BufferProtocolExportable(Protocol): + def __buffer__(self, flags: int) -> memoryview: ... + class RTree(Generic[T]): + def __init__(self, obj: BufferProtocolExportable) -> None: ... + def __buffer__(self, flags: int) -> memoryview: ... @classmethod def from_interleaved( cls, diff --git a/python/python/geoindex_rs/tests/__init__.py b/python/tests/__init__.py similarity index 100% rename from python/python/geoindex_rs/tests/__init__.py rename to python/tests/__init__.py diff --git a/python/python/geoindex_rs/tests/test_buffers.py b/python/tests/test_buffers.py similarity index 57% rename from python/python/geoindex_rs/tests/test_buffers.py rename to python/tests/test_buffers.py index b0ee79d..f98ca8d 100644 --- a/python/python/geoindex_rs/tests/test_buffers.py +++ b/python/tests/test_buffers.py @@ -1,6 +1,6 @@ import numpy as np -from .. import RTree +from geoindex_rs import RTree def generate_random_boxes(): @@ -12,6 +12,10 @@ def generate_random_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 + second = RTree(initial) + + assert initial is not second, "Not the same object" + + # Flatbush magic byte + assert memoryview(initial)[0] == 0xFB + assert memoryview(second)[0] == 0xFB From e00dd887ed67874ec4b0d1f350bc3a42d06fcdd1 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 21 Aug 2024 12:26:11 -0400 Subject: [PATCH 06/12] remove old rtree --- python/src/lib.rs | 1 - python/src/rtree.rs | 271 -------------------------------------------- 2 files changed, 272 deletions(-) diff --git a/python/src/lib.rs b/python/src/lib.rs index f0d64eb..dfb6691 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -17,7 +17,6 @@ fn _rust(_py: Python, m: &Bound) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(rtree::search_rtree))?; m.add_class::()?; - m.add_class::()?; m.add_class::()?; Ok(()) diff --git a/python/src/rtree.rs b/python/src/rtree.rs index 1a92cf9..4a9d545 100644 --- a/python/src/rtree.rs +++ b/python/src/rtree.rs @@ -488,11 +488,6 @@ pub(crate) fn search_rtree( PyArray1::from_vec_bound(py, result) } -enum RTreeInner { - Float32(OwnedRTree), - Float64(OwnedRTree), -} - /// These constructors are separated out so they can be generic fn new_interleaved( boxes: &ArrayView2, @@ -560,269 +555,3 @@ fn new_separated( RTreeMethod::STR => builder.finish::(), } } - -impl RTreeInner { - fn num_items(&self) -> usize { - match self { - Self::Float32(index) => index.num_items(), - Self::Float64(index) => index.num_items(), - } - } - - fn num_nodes(&self) -> usize { - match self { - Self::Float32(index) => index.num_nodes(), - Self::Float64(index) => index.num_nodes(), - } - } - - fn node_size(&self) -> usize { - match self { - Self::Float32(index) => index.node_size(), - Self::Float64(index) => index.node_size(), - } - } - - fn num_levels(&self) -> usize { - match self { - Self::Float32(index) => index.num_levels(), - Self::Float64(index) => index.num_levels(), - } - } - - fn num_bytes(&self) -> usize { - match self { - Self::Float32(index) => index.as_ref().len(), - Self::Float64(index) => index.as_ref().len(), - } - } - - fn buffer(&self) -> &[u8] { - match self { - Self::Float32(index) => index.as_ref(), - Self::Float64(index) => index.as_ref(), - } - } - - fn boxes_at_level<'py>(&'py self, py: Python<'py>, level: usize) -> PyResult { - match self { - Self::Float32(index) => { - let boxes = index - .boxes_at_level(level) - .map_err(|err| PyIndexError::new_err(err.to_string()))?; - let array = PyArray1::from_slice_bound(py, boxes); - Ok(array.reshape([boxes.len() / 4, 4])?.into_py(py)) - } - Self::Float64(index) => { - let boxes = index - .boxes_at_level(level) - .map_err(|err| PyIndexError::new_err(err.to_string()))?; - let array = PyArray1::from_slice_bound(py, boxes); - Ok(array.reshape([boxes.len() / 4, 4])?.into_py(py)) - } - } - } -} - -#[pyclass] -pub struct RTreeOld(RTreeInner); - -#[pymethods] -impl RTreeOld { - #[classmethod] - #[pyo3( - signature = (boxes, *, method = RTreeMethod::Hilbert, node_size = None), - text_signature = "(boxes, *, method = 'hilbert', node_size = None)") - ] - pub fn from_interleaved( - _cls: &Bound, - py: Python, - boxes: PyObject, - method: RTreeMethod, - node_size: Option, - ) -> PyResult { - // Convert to numpy array (of the same dtype) - let boxes = boxes.call_method0(py, intern!(py, "__array__"))?; - - let result = if let Ok(boxes) = boxes.extract::>(py) { - let boxes = boxes.as_array(); - let tree = py.allow_threads(|| new_interleaved(&boxes, method, node_size)); - Ok(Self(RTreeInner::Float64(tree))) - } else if let Ok(boxes) = boxes.extract::>(py) { - let boxes = boxes.as_array(); - let tree = py.allow_threads(|| new_interleaved(&boxes, method, node_size)); - Ok(Self(RTreeInner::Float32(tree))) - } else { - let dtype = boxes.call_method0(py, intern!(py, "dtype"))?.to_string(); - Err(PyTypeError::new_err(format!( - "Expected a numpy array of dtype float32 or float64, got {}", - dtype - ))) - }; - result - } - - #[classmethod] - #[pyo3( - signature = (min_x, min_y, max_x, max_y, *, method = RTreeMethod::Hilbert, node_size = None), - text_signature = "(min_x, min_y, max_x, max_y, *, method = 'hilbert', node_size = None)") - ] - #[allow(clippy::too_many_arguments)] - pub fn from_separated( - _cls: &Bound, - py: Python, - min_x: PyObject, - min_y: PyObject, - max_x: PyObject, - max_y: PyObject, - method: RTreeMethod, - node_size: Option, - ) -> PyResult { - // Convert to numpy array (of the same dtype) - let min_x = min_x.call_method0(py, intern!(py, "__array__"))?; - let min_y = min_y.call_method0(py, intern!(py, "__array__"))?; - let max_x = max_x.call_method0(py, intern!(py, "__array__"))?; - let max_y = max_y.call_method0(py, intern!(py, "__array__"))?; - - let result = if let Ok(min_x) = min_x.extract::>(py) { - let min_y = min_y.extract::>(py)?; - let max_x = max_x.extract::>(py)?; - let max_y = max_y.extract::>(py)?; - - let min_x_array = min_x.as_array(); - let min_y_array = min_y.as_array(); - let max_x_array = max_x.as_array(); - let max_y_array = max_y.as_array(); - - let tree = py.allow_threads(|| { - new_separated( - min_x_array, - min_y_array, - max_x_array, - max_y_array, - method, - node_size, - ) - }); - Ok(Self(RTreeInner::Float64(tree))) - } else if let Ok(min_x) = min_x.extract::>(py) { - let min_y = min_y.extract::>(py)?; - let max_x = max_x.extract::>(py)?; - let max_y = max_y.extract::>(py)?; - - let min_x_array = min_x.as_array(); - let min_y_array = min_y.as_array(); - let max_x_array = max_x.as_array(); - let max_y_array = max_y.as_array(); - - let tree = py.allow_threads(|| { - new_separated( - min_x_array, - min_y_array, - max_x_array, - max_y_array, - method, - node_size, - ) - }); - Ok(Self(RTreeInner::Float32(tree))) - } else { - let dtype = min_x.call_method0(py, intern!(py, "dtype"))?.to_string(); - Err(PyTypeError::new_err(format!( - "Expected a numpy array of dtype float32 or float64, got {}", - dtype - ))) - }; - result - } - - // pre PEP 688 buffer protocol - unsafe fn __getbuffer__( - slf: PyRef<'_, Self>, - view: *mut ffi::Py_buffer, - flags: c_int, - ) -> PyResult<()> { - let bytes = slf.0.buffer(); - let ret = ffi::PyBuffer_FillInfo( - view, - slf.as_ptr() as *mut _, - bytes.as_ptr() as *mut _, - bytes.len().try_into().unwrap(), - 1, // read only - flags, - ); - if ret == -1 { - return Err(PyErr::fetch(slf.py())); - } - Ok(()) - } - unsafe fn __releasebuffer__(&self, _view: *mut ffi::Py_buffer) { - // is there anything to do here? - } - - /// The total number of items contained in this RTree. - #[getter] - pub fn num_items(&self) -> usize { - self.0.num_items() - } - - /// The total number of nodes in this RTree, including both leaf and intermediate nodes. - #[getter] - pub fn num_nodes(&self) -> usize { - self.0.num_nodes() - } - - /// The maximum number of elements in each node. - #[getter] - pub fn node_size(&self) -> usize { - self.0.node_size() - } - - /// The height of the tree - #[getter] - fn num_levels(&self) -> usize { - self.0.num_levels() - } - - /// The number of bytes taken up in memory. - #[getter] - fn num_bytes(&self) -> usize { - self.0.num_bytes() - } - - /// Access the bounding boxes at the given level of the tree. - /// - /// The tree is laid out from bottom to top. Level 0 is the _base_ of the tree. Each integer - /// higher is one level higher of the tree. - fn boxes_at_level<'py>(&'py self, py: Python<'py>, level: usize) -> PyResult { - self.0.boxes_at_level(py, level) - } - - /// Search an RTree given the provided bounding box. - /// - /// Results are the indexes of the inserted objects in insertion order. - /// - /// Args: - /// min_x: min x coordinate of bounding box - /// min_y: min y coordinate of bounding box - /// max_x: max x coordinate of bounding box - /// max_y: max y coordinate of bounding box - pub fn search<'py>( - &'py self, - py: Python<'py>, - min_x: f64, - min_y: f64, - max_x: f64, - max_y: f64, - ) -> Bound<'py, PyArray1> { - let result = py.allow_threads(|| match &self.0 { - RTreeInner::Float32(tree) => { - let (min_x, min_y, max_x, max_y) = f64_box_to_f32(min_x, min_y, max_x, max_y); - tree.search(min_x, min_y, max_x, max_y) - } - RTreeInner::Float64(tree) => tree.search(min_x, min_y, max_x, max_y), - }); - - PyArray1::from_vec_bound(py, result) - } -} From 478fdf74132a5711d5416f81e382b8a7bde4df46 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 21 Aug 2024 23:20:10 -0400 Subject: [PATCH 07/12] move buffer code to common --- python/src/common.rs | 76 ++++++++++++++++++++++++++++ python/src/lib.rs | 5 +- python/src/rtree.rs | 118 ++++++++++--------------------------------- 3 files changed, 106 insertions(+), 93 deletions(-) create mode 100644 python/src/common.rs diff --git a/python/src/common.rs b/python/src/common.rs new file mode 100644 index 0000000..3f8d0dc --- /dev/null +++ b/python/src/common.rs @@ -0,0 +1,76 @@ +use pyo3::buffer::PyBuffer; +use pyo3::exceptions::PyValueError; +use pyo3::ffi; +use pyo3::prelude::*; +use std::os::raw::c_int; + +/// A Rust buffer that implements the Python buffer protocol +#[pyclass(name = "Buffer")] +pub(crate) struct RustBuffer(Vec); + +impl RustBuffer { + pub(crate) fn new(buffer: Vec) -> Self { + Self(buffer) + } +} + +#[pymethods] +impl RustBuffer { + // pre PEP 688 buffer protocol + unsafe fn __getbuffer__( + slf: PyRef<'_, Self>, + view: *mut ffi::Py_buffer, + flags: c_int, + ) -> PyResult<()> { + let bytes = slf.0.as_slice(); + let ret = ffi::PyBuffer_FillInfo( + view, + slf.as_ptr() as *mut _, + bytes.as_ptr() as *mut _, + bytes.len().try_into().unwrap(), + 1, // read only + flags, + ); + if ret == -1 { + return Err(PyErr::fetch(slf.py())); + } + Ok(()) + } + + unsafe fn __releasebuffer__(&self, _view: *mut ffi::Py_buffer) { + // is there anything to do here? + } +} + +/// A Rust representation of a Python object that implements the Python buffer protocol, exporting +/// a 1-dimensional `&[u8]` slice. +pub(crate) struct PyU8Buffer(PyBuffer); + +impl<'py> FromPyObject<'py> for PyU8Buffer { + fn extract_bound(obj: &Bound<'py, PyAny>) -> PyResult { + let buffer = PyBuffer::::get_bound(obj)?; + if !buffer.readonly() { + return Err(PyValueError::new_err("Must be read-only byte buffer.")); + } + if buffer.dimensions() != 1 { + return Err(PyValueError::new_err("Expected 1-dimensional array.")); + } + // Note: this is probably superfluous for 1D array + if !buffer.is_c_contiguous() { + return Err(PyValueError::new_err("Expected c-contiguous array.")); + } + if buffer.len_bytes() == 0 { + return Err(PyValueError::new_err("Buffer has no data.")); + } + + Ok(Self(buffer)) + } +} + +impl AsRef<[u8]> for PyU8Buffer { + fn as_ref(&self) -> &[u8] { + let len = self.0.item_count(); + let data = self.0.buf_ptr() as *const u8; + unsafe { std::slice::from_raw_parts(data, len) } + } +} diff --git a/python/src/lib.rs b/python/src/lib.rs index dfb6691..2dade56 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -1,5 +1,6 @@ -pub mod kdtree; -pub mod rtree; +mod common; +mod kdtree; +mod rtree; use pyo3::prelude::*; diff --git a/python/src/rtree.rs b/python/src/rtree.rs index 4a9d545..d8e891b 100644 --- a/python/src/rtree.rs +++ b/python/src/rtree.rs @@ -5,7 +5,6 @@ use geo_index::rtree::{OwnedRTree, RTreeBuilder, RTreeIndex, TreeMetadata}; use geo_index::{CoordType, IndexableNum}; use numpy::ndarray::{ArrayView1, ArrayView2}; use numpy::{PyArray1, PyArrayMethods, PyReadonlyArray1, PyReadonlyArray2}; -use pyo3::buffer::PyBuffer; use pyo3::exceptions::{PyIndexError, PyTypeError, PyValueError}; use pyo3::ffi; use pyo3::intern; @@ -13,10 +12,12 @@ use pyo3::prelude::*; use pyo3::types::PyType; use std::os::raw::c_int; +use crate::common::{PyU8Buffer, RustBuffer}; + /// Method for constructing rtree -pub enum RTreeMethod { +enum RTreeMethod { Hilbert, - STR, + SortTileRecursive, } impl<'a> FromPyObject<'a> for RTreeMethod { @@ -24,7 +25,7 @@ impl<'a> FromPyObject<'a> for RTreeMethod { let s: String = ob.extract()?; match s.to_lowercase().as_str() { "hilbert" => Ok(Self::Hilbert), - "str" => Ok(Self::STR), + "str" => Ok(Self::SortTileRecursive), _ => Err(PyValueError::new_err( "Unexpected method. Should be one of 'hilbert' or 'str'.", )), @@ -32,79 +33,14 @@ impl<'a> FromPyObject<'a> for RTreeMethod { } } -/// A Rust buffer that implements the Python buffer protocol -#[pyclass(name = "Buffer")] -struct RustBuffer(Vec); - -#[pymethods] -impl RustBuffer { - // pre PEP 688 buffer protocol - unsafe fn __getbuffer__( - slf: PyRef<'_, Self>, - view: *mut ffi::Py_buffer, - flags: c_int, - ) -> PyResult<()> { - let bytes = slf.0.as_slice(); - let ret = ffi::PyBuffer_FillInfo( - view, - slf.as_ptr() as *mut _, - bytes.as_ptr() as *mut _, - bytes.len().try_into().unwrap(), - 1, // read only - flags, - ); - if ret == -1 { - return Err(PyErr::fetch(slf.py())); - } - Ok(()) - } - - unsafe fn __releasebuffer__(&self, _view: *mut ffi::Py_buffer) { - // is there anything to do here? - } -} - -/// A Rust representation of a Python object that implements the Python buffer protocol, exporting -/// a 1-dimensional `&[u8]` slice. -struct PyU8Buffer(PyBuffer); - -impl<'py> FromPyObject<'py> for PyU8Buffer { - fn extract_bound(obj: &Bound<'py, PyAny>) -> PyResult { - let buffer = PyBuffer::::get_bound(obj)?; - if !buffer.readonly() { - return Err(PyValueError::new_err("Must be read-only byte buffer.")); - } - if buffer.dimensions() != 1 { - return Err(PyValueError::new_err("Expected 1-dimensional array.")); - } - // Note: this is probably superfluous for 1D array - if !buffer.is_c_contiguous() { - return Err(PyValueError::new_err("Expected c-contiguous array.")); - } - if buffer.len_bytes() == 0 { - return Err(PyValueError::new_err("Buffer has no data.")); - } - - Ok(Self(buffer)) - } -} - -impl AsRef<[u8]> for PyU8Buffer { - fn as_ref(&self) -> &[u8] { - let len = self.0.item_count(); - let data = self.0.buf_ptr() as *const u8; - unsafe { std::slice::from_raw_parts(data, len) } - } -} - /// A low-level wrapper around a [PyU8Buffer] that validates that the input is a valid Flatbush /// buffer. This wrapper implements [RTreeIndex]. -pub(crate) struct PyRTreeRefInner { +pub(crate) struct PyRTreeBuffer { buffer: PyU8Buffer, metadata: TreeMetadata, } -impl PyRTreeRefInner { +impl PyRTreeBuffer { fn try_new(buffer: PyU8Buffer) -> PyResult { let metadata = TreeMetadata::try_new(buffer.as_ref()).unwrap(); Ok(Self { buffer, metadata }) @@ -112,7 +48,7 @@ impl PyRTreeRefInner { fn from_owned_rtree(py: Python, tree: OwnedRTree) -> PyResult { let metadata = tree.metadata().clone(); - let tree_buf = RustBuffer(tree.into_inner()); + let tree_buf = RustBuffer::new(tree.into_inner()); Ok(Self { buffer: tree_buf.into_py(py).extract(py)?, metadata, @@ -120,13 +56,13 @@ impl PyRTreeRefInner { } } -impl AsRef<[u8]> for PyRTreeRefInner { +impl AsRef<[u8]> for PyRTreeBuffer { fn as_ref(&self) -> &[u8] { self.buffer.as_ref() } } -impl RTreeIndex for PyRTreeRefInner { +impl RTreeIndex for PyRTreeBuffer { fn boxes(&self) -> &[N] { self.metadata.boxes_slice(self.as_ref()) } @@ -152,11 +88,11 @@ impl RTreeIndex for PyRTreeRefInner { } } -/// An enum wrapper around [PyRTreeRefInner] that allows use of multiple coordinate types from +/// An enum wrapper around [PyRTreeBuffer] that allows use of multiple coordinate types from /// Python. pub(crate) enum PyRTreeRef { - Float32(PyRTreeRefInner), - Float64(PyRTreeRefInner), + Float32(PyRTreeBuffer), + Float64(PyRTreeBuffer), } impl<'py> FromPyObject<'py> for PyRTreeRef { @@ -164,21 +100,21 @@ impl<'py> FromPyObject<'py> for PyRTreeRef { let buffer = PyU8Buffer::extract_bound(ob)?; let ct = CoordType::from_buffer(&buffer.as_ref()).unwrap(); match ct { - CoordType::Float32 => Ok(Self::Float32(PyRTreeRefInner::try_new(buffer)?)), - CoordType::Float64 => Ok(Self::Float64(PyRTreeRefInner::try_new(buffer)?)), + CoordType::Float32 => Ok(Self::Float32(PyRTreeBuffer::try_new(buffer)?)), + CoordType::Float64 => Ok(Self::Float64(PyRTreeBuffer::try_new(buffer)?)), _ => todo!(), } } } -impl From> for PyRTreeRef { - fn from(value: PyRTreeRefInner) -> Self { +impl From> for PyRTreeRef { + fn from(value: PyRTreeBuffer) -> Self { Self::Float32(value) } } -impl From> for PyRTreeRef { - fn from(value: PyRTreeRefInner) -> Self { +impl From> for PyRTreeRef { + fn from(value: PyRTreeBuffer) -> Self { Self::Float64(value) } } @@ -264,7 +200,7 @@ impl RTree { signature = (boxes, *, method = RTreeMethod::Hilbert, node_size = None), text_signature = "(boxes, *, method = 'hilbert', node_size = None)") ] - pub fn from_interleaved( + fn from_interleaved( _cls: &Bound, py: Python, boxes: PyObject, @@ -277,11 +213,11 @@ impl RTree { let result = if let Ok(boxes) = boxes.extract::>(py) { let boxes = boxes.as_array(); let tree = py.allow_threads(|| new_interleaved(&boxes, method, node_size)); - Ok(Self(PyRTreeRefInner::from_owned_rtree(py, tree)?.into())) + Ok(Self(PyRTreeBuffer::from_owned_rtree(py, tree)?.into())) } else if let Ok(boxes) = boxes.extract::>(py) { let boxes = boxes.as_array(); let tree = py.allow_threads(|| new_interleaved(&boxes, method, node_size)); - Ok(Self(PyRTreeRefInner::from_owned_rtree(py, tree)?.into())) + Ok(Self(PyRTreeBuffer::from_owned_rtree(py, tree)?.into())) } else { let dtype = boxes.call_method0(py, intern!(py, "dtype"))?.to_string(); Err(PyTypeError::new_err(format!( @@ -298,7 +234,7 @@ impl RTree { text_signature = "(min_x, min_y, max_x, max_y, *, method = 'hilbert', node_size = None)") ] #[allow(clippy::too_many_arguments)] - pub fn from_separated( + fn from_separated( _cls: &Bound, py: Python, min_x: PyObject, @@ -334,7 +270,7 @@ impl RTree { node_size, ) }); - Ok(Self(PyRTreeRefInner::from_owned_rtree(py, tree)?.into())) + Ok(Self(PyRTreeBuffer::from_owned_rtree(py, tree)?.into())) } else if let Ok(min_x) = min_x.extract::>(py) { let min_y = min_y.extract::>(py)?; let max_x = max_x.extract::>(py)?; @@ -355,7 +291,7 @@ impl RTree { node_size, ) }); - Ok(Self(PyRTreeRefInner::from_owned_rtree(py, tree)?.into())) + Ok(Self(PyRTreeBuffer::from_owned_rtree(py, tree)?.into())) } else { let dtype = min_x.call_method0(py, intern!(py, "dtype"))?.to_string(); Err(PyTypeError::new_err(format!( @@ -517,7 +453,7 @@ fn new_interleaved( match method { RTreeMethod::Hilbert => builder.finish::(), - RTreeMethod::STR => builder.finish::(), + RTreeMethod::SortTileRecursive => builder.finish::(), } } @@ -552,6 +488,6 @@ fn new_separated( match method { RTreeMethod::Hilbert => builder.finish::(), - RTreeMethod::STR => builder.finish::(), + RTreeMethod::SortTileRecursive => builder.finish::(), } } From 9cd5c2b5203ec6dd0140d435f7c19743a0171cb4 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 21 Aug 2024 23:42:55 -0400 Subject: [PATCH 08/12] Generalize more --- python/src/common.rs | 3 +- python/src/lib.rs | 2 - python/src/rtree.rs | 187 +++++++++++++++++++++++++++++-------------- 3 files changed, 131 insertions(+), 61 deletions(-) diff --git a/python/src/common.rs b/python/src/common.rs index 3f8d0dc..6904b64 100644 --- a/python/src/common.rs +++ b/python/src/common.rs @@ -16,7 +16,7 @@ impl RustBuffer { #[pymethods] impl RustBuffer { - // pre PEP 688 buffer protocol + /// Implements the buffer protocol export unsafe fn __getbuffer__( slf: PyRef<'_, Self>, view: *mut ffi::Py_buffer, @@ -68,6 +68,7 @@ impl<'py> FromPyObject<'py> for PyU8Buffer { } impl AsRef<[u8]> for PyU8Buffer { + /// Extract a slice from a Python object implementing the buffer protocol fn as_ref(&self) -> &[u8] { let len = self.0.item_count(); let data = self.0.buf_ptr() as *const u8; diff --git a/python/src/lib.rs b/python/src/lib.rs index 2dade56..5d9131a 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -15,8 +15,6 @@ fn ___version() -> &'static str { fn _rust(_py: Python, m: &Bound) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(___version))?; - m.add_wrapped(wrap_pyfunction!(rtree::search_rtree))?; - m.add_class::()?; m.add_class::()?; diff --git a/python/src/rtree.rs b/python/src/rtree.rs index d8e891b..58385ca 100644 --- a/python/src/rtree.rs +++ b/python/src/rtree.rs @@ -91,6 +91,12 @@ impl RTreeIndex for PyRTreeBuffer { /// An enum wrapper around [PyRTreeBuffer] that allows use of multiple coordinate types from /// Python. pub(crate) enum PyRTreeRef { + Int8(PyRTreeBuffer), + Int16(PyRTreeBuffer), + Int32(PyRTreeBuffer), + UInt8(PyRTreeBuffer), + UInt16(PyRTreeBuffer), + UInt32(PyRTreeBuffer), Float32(PyRTreeBuffer), Float64(PyRTreeBuffer), } @@ -122,6 +128,12 @@ impl From> for PyRTreeRef { impl AsRef<[u8]> for PyRTreeRef { fn as_ref(&self) -> &[u8] { match self { + Self::Int8(inner) => inner.as_ref(), + Self::Int16(inner) => inner.as_ref(), + Self::Int32(inner) => inner.as_ref(), + Self::UInt8(inner) => inner.as_ref(), + Self::UInt16(inner) => inner.as_ref(), + Self::UInt32(inner) => inner.as_ref(), Self::Float32(inner) => inner.as_ref(), Self::Float64(inner) => inner.as_ref(), } @@ -131,6 +143,12 @@ impl AsRef<[u8]> for PyRTreeRef { impl PyRTreeRef { fn num_items(&self) -> usize { match self { + Self::Int8(index) => index.num_items(), + Self::Int16(index) => index.num_items(), + Self::Int32(index) => index.num_items(), + Self::UInt8(index) => index.num_items(), + Self::UInt16(index) => index.num_items(), + Self::UInt32(index) => index.num_items(), Self::Float32(index) => index.num_items(), Self::Float64(index) => index.num_items(), } @@ -138,6 +156,12 @@ impl PyRTreeRef { fn num_nodes(&self) -> usize { match self { + Self::Int8(index) => index.num_nodes(), + Self::Int16(index) => index.num_nodes(), + Self::Int32(index) => index.num_nodes(), + Self::UInt8(index) => index.num_nodes(), + Self::UInt16(index) => index.num_nodes(), + Self::UInt32(index) => index.num_nodes(), Self::Float32(index) => index.num_nodes(), Self::Float64(index) => index.num_nodes(), } @@ -145,6 +169,12 @@ impl PyRTreeRef { fn node_size(&self) -> usize { match self { + Self::Int8(index) => index.node_size(), + Self::Int16(index) => index.node_size(), + Self::Int32(index) => index.node_size(), + Self::UInt8(index) => index.node_size(), + Self::UInt16(index) => index.node_size(), + Self::UInt32(index) => index.node_size(), Self::Float32(index) => index.node_size(), Self::Float64(index) => index.node_size(), } @@ -152,6 +182,12 @@ impl PyRTreeRef { fn num_levels(&self) -> usize { match self { + Self::Int8(index) => index.num_levels(), + Self::Int16(index) => index.num_levels(), + Self::Int32(index) => index.num_levels(), + Self::UInt8(index) => index.num_levels(), + Self::UInt16(index) => index.num_levels(), + Self::UInt32(index) => index.num_levels(), Self::Float32(index) => index.num_levels(), Self::Float64(index) => index.num_levels(), } @@ -159,39 +195,51 @@ impl PyRTreeRef { fn num_bytes(&self) -> usize { match self { - Self::Float32(index) => AsRef::as_ref(index).len(), - Self::Float64(index) => AsRef::as_ref(index).len(), + Self::Int8(index) => index.as_ref().len(), + Self::Int16(index) => index.as_ref().len(), + Self::Int32(index) => index.as_ref().len(), + Self::UInt8(index) => index.as_ref().len(), + Self::UInt16(index) => index.as_ref().len(), + Self::UInt32(index) => index.as_ref().len(), + Self::Float32(index) => index.as_ref().len(), + Self::Float64(index) => index.as_ref().len(), } } fn boxes_at_level<'py>(&'py self, py: Python<'py>, level: usize) -> PyResult { match self { - Self::Float32(index) => { - let boxes = index - .boxes_at_level(level) - .map_err(|err| PyIndexError::new_err(err.to_string()))?; - let array = PyArray1::from_slice_bound(py, boxes); - Ok(array.reshape([boxes.len() / 4, 4])?.into_py(py)) - } - Self::Float64(index) => { - let boxes = index - .boxes_at_level(level) - .map_err(|err| PyIndexError::new_err(err.to_string()))?; - let array = PyArray1::from_slice_bound(py, boxes); - Ok(array.reshape([boxes.len() / 4, 4])?.into_py(py)) - } + Self::Int8(index) => _boxes_at_level(py, index, level), + Self::Int16(index) => _boxes_at_level(py, index, level), + Self::Int32(index) => _boxes_at_level(py, index, level), + Self::UInt8(index) => _boxes_at_level(py, index, level), + Self::UInt16(index) => _boxes_at_level(py, index, level), + Self::UInt32(index) => _boxes_at_level(py, index, level), + Self::Float32(index) => _boxes_at_level(py, index, level), + Self::Float64(index) => _boxes_at_level(py, index, level), } } } +fn _boxes_at_level( + py: Python, + index: &PyRTreeBuffer, + level: usize, +) -> PyResult { + let boxes = index + .boxes_at_level(level) + .map_err(|err| PyIndexError::new_err(err.to_string()))?; + let array = PyArray1::from_slice_bound(py, boxes); + Ok(array.reshape([boxes.len() / 4, 4])?.into_py(py)) +} + #[pyclass] pub(crate) struct RTree(PyRTreeRef); #[pymethods] impl RTree { /// Construct an RTree from an existing RTree buffer - #[new] - fn new(py: Python, obj: PyObject) -> PyResult { + #[classmethod] + fn from_buffer(_cls: &Bound, py: Python, obj: PyObject) -> PyResult { Ok(Self(obj.extract(py)?)) } @@ -377,53 +425,76 @@ impl RTree { pub fn search<'py>( &'py self, py: Python<'py>, - min_x: f64, - min_y: f64, - max_x: f64, - max_y: f64, - ) -> Bound<'py, PyArray1> { - let result = py.allow_threads(|| match &self.0 { + min_x: PyObject, + min_y: PyObject, + max_x: PyObject, + max_y: PyObject, + ) -> PyResult>> { + let result: Result<_, PyErr> = match &self.0 { + PyRTreeRef::Int8(tree) => { + let min_x = min_x.extract(py)?; + let min_y = min_y.extract(py)?; + let max_x = max_x.extract(py)?; + let max_y = max_y.extract(py)?; + Ok(py.allow_threads(|| tree.search(min_x, min_y, max_x, max_y))) + } + PyRTreeRef::Int16(tree) => { + let min_x = min_x.extract(py)?; + let min_y = min_y.extract(py)?; + let max_x = max_x.extract(py)?; + let max_y = max_y.extract(py)?; + Ok(py.allow_threads(|| tree.search(min_x, min_y, max_x, max_y))) + } + PyRTreeRef::Int32(tree) => { + let min_x = min_x.extract(py)?; + let min_y = min_y.extract(py)?; + let max_x = max_x.extract(py)?; + let max_y = max_y.extract(py)?; + Ok(py.allow_threads(|| tree.search(min_x, min_y, max_x, max_y))) + } + PyRTreeRef::UInt8(tree) => { + let min_x = min_x.extract(py)?; + let min_y = min_y.extract(py)?; + let max_x = max_x.extract(py)?; + let max_y = max_y.extract(py)?; + Ok(py.allow_threads(|| tree.search(min_x, min_y, max_x, max_y))) + } + PyRTreeRef::UInt16(tree) => { + let min_x = min_x.extract(py)?; + let min_y = min_y.extract(py)?; + let max_x = max_x.extract(py)?; + let max_y = max_y.extract(py)?; + Ok(py.allow_threads(|| tree.search(min_x, min_y, max_x, max_y))) + } + PyRTreeRef::UInt32(tree) => { + let min_x = min_x.extract(py)?; + let min_y = min_y.extract(py)?; + let max_x = max_x.extract(py)?; + let max_y = max_y.extract(py)?; + Ok(py.allow_threads(|| tree.search(min_x, min_y, max_x, max_y))) + } PyRTreeRef::Float32(tree) => { + let min_x = min_x.extract::(py)?; + let min_y = min_y.extract::(py)?; + let max_x = max_x.extract::(py)?; + let max_y = max_y.extract::(py)?; + let (min_x, min_y, max_x, max_y) = f64_box_to_f32(min_x, min_y, max_x, max_y); - tree.search(min_x, min_y, max_x, max_y) + Ok(py.allow_threads(|| tree.search(min_x, min_y, max_x, max_y))) + } + PyRTreeRef::Float64(tree) => { + let min_x = min_x.extract(py)?; + let min_y = min_y.extract(py)?; + let max_x = max_x.extract(py)?; + let max_y = max_y.extract(py)?; + Ok(py.allow_threads(|| tree.search(min_x, min_y, max_x, max_y))) } - PyRTreeRef::Float64(tree) => tree.search(min_x, min_y, max_x, max_y), - }); + }; - PyArray1::from_vec_bound(py, result) + Ok(PyArray1::from_vec_bound(py, result?)) } } -/// Search an RTree given the provided bounding box. -/// -/// Results are the indexes of the inserted objects in insertion order. -/// -/// Args: -/// tree: tree or buffer to search -/// min_x: min x coordinate of bounding box -/// min_y: min y coordinate of bounding box -/// max_x: max x coordinate of bounding box -/// max_y: max y coordinate of bounding box -#[pyfunction] -pub(crate) fn search_rtree( - py: Python, - tree: PyRTreeRef, - min_x: f64, - min_y: f64, - max_x: f64, - max_y: f64, -) -> Bound<'_, PyArray1> { - let result = py.allow_threads(|| match tree { - PyRTreeRef::Float32(tree) => { - let (min_x, min_y, max_x, max_y) = f64_box_to_f32(min_x, min_y, max_x, max_y); - tree.search(min_x, min_y, max_x, max_y) - } - PyRTreeRef::Float64(tree) => tree.search(min_x, min_y, max_x, max_y), - }); - - PyArray1::from_vec_bound(py, result) -} - /// These constructors are separated out so they can be generic fn new_interleaved( boxes: &ArrayView2, From 0af56bf8dd1dff593bbdccd38be35af67f7cf023 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Thu, 22 Aug 2024 00:32:43 -0400 Subject: [PATCH 09/12] generalize --- python/Cargo.lock | 8 +- python/Cargo.toml | 2 + python/src/common.rs | 94 +++++++++++++++- python/src/rtree.rs | 254 ++++++++++++++++++++++++++++++++----------- 4 files changed, 287 insertions(+), 71 deletions(-) diff --git a/python/Cargo.lock b/python/Cargo.lock index 2370070..35c9322 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -106,9 +106,9 @@ checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5" [[package]] name = "libc" -version = "0.2.153" +version = "0.2.158" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c198f91728a82281a64e1f4f9eeb25d82cb32a5de251c6bd1b5154d63a8e7bd" +checksum = "d8adc4bb1803a324070e64a98ae98f38934d91957a99cfb3a43dcbc01bc56439" [[package]] name = "lock_api" @@ -172,9 +172,9 @@ dependencies = [ [[package]] name = "num-traits" -version = "0.2.18" +version = "0.2.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da0df0e5185db44f69b44f26786fe401b6c293d1907744beaa7fa62b2e5a517a" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", ] diff --git a/python/Cargo.toml b/python/Cargo.toml index 03d424a..7cb543f 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -17,6 +17,7 @@ name = "_rust" crate-type = ["cdylib"] [dependencies] +# arrow = "52" bytes = "1" geo-index = { path = "../", features = ["rayon"] } # numpy = "0.21" @@ -26,6 +27,7 @@ geo-index = { path = "../", features = ["rayon"] } # 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.0", features = ["macros"] } +# pyo3-arrow = "*" thiserror = "1" [profile.release] diff --git a/python/src/common.rs b/python/src/common.rs index 6904b64..c275b38 100644 --- a/python/src/common.rs +++ b/python/src/common.rs @@ -1,5 +1,7 @@ +use geo_index::IndexableNum; +use numpy::{dtype_bound, PyArray1, PyArrayDescr, PyUntypedArray}; use pyo3::buffer::PyBuffer; -use pyo3::exceptions::PyValueError; +use pyo3::exceptions::{PyTypeError, PyValueError}; use pyo3::ffi; use pyo3::prelude::*; use std::os::raw::c_int; @@ -75,3 +77,93 @@ impl AsRef<[u8]> for PyU8Buffer { unsafe { std::slice::from_raw_parts(data, len) } } } + +pub(crate) enum PyTypedArrayRef<'py, N: IndexableNum + numpy::Element> { + // Arrow((ArrayRef, PhantomData)), + Numpy(&'py PyArray1), +} + +impl<'py, N: IndexableNum + numpy::Element> PyTypedArrayRef<'py, N> { + pub(crate) fn as_slice(&self) -> &[N] { + match self { + Self::Numpy(arr) => unsafe { arr.as_slice() }.unwrap(), + } + } +} + +pub(crate) enum PyArray<'py> { + Int8(PyTypedArrayRef<'py, i8>), + Int16(PyTypedArrayRef<'py, i16>), + Int32(PyTypedArrayRef<'py, i32>), + UInt8(PyTypedArrayRef<'py, u8>), + UInt16(PyTypedArrayRef<'py, u16>), + UInt32(PyTypedArrayRef<'py, u32>), + Float32(PyTypedArrayRef<'py, f32>), + Float64(PyTypedArrayRef<'py, f64>), +} + +impl<'py> FromPyObject<'py> for PyArray<'py> { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + let mut ob = ob.to_owned(); + // call __array__ if it exists + if ob.hasattr("__array__")? { + ob = ob.call_method0("__array__")?; + } + + if let Ok(array) = ob.extract::<&'py PyUntypedArray>() { + if array.ndim() != 1 { + return Err(PyValueError::new_err("Expected 1-dimensional array.")); + } + + let dtype = array.dtype(); + + if is_type::(dtype) { + let arr = array.downcast::>()?; + return Ok(Self::Int8(PyTypedArrayRef::Numpy(arr))); + } + + if is_type::(dtype) { + let arr = array.downcast::>()?; + return Ok(Self::Int16(PyTypedArrayRef::Numpy(arr))); + } + + if is_type::(dtype) { + let arr = array.downcast::>()?; + return Ok(Self::Int32(PyTypedArrayRef::Numpy(arr))); + } + + if is_type::(dtype) { + let arr = array.downcast::>()?; + return Ok(Self::UInt8(PyTypedArrayRef::Numpy(arr))); + } + + if is_type::(dtype) { + let arr = array.downcast::>()?; + return Ok(Self::UInt16(PyTypedArrayRef::Numpy(arr))); + } + + if is_type::(dtype) { + let arr = array.downcast::>()?; + return Ok(Self::UInt32(PyTypedArrayRef::Numpy(arr))); + } + + if is_type::(dtype) { + let arr = array.downcast::>()?; + return Ok(Self::Float32(PyTypedArrayRef::Numpy(arr))); + } + + if is_type::(dtype) { + let arr = array.downcast::>()?; + return Ok(Self::Float64(PyTypedArrayRef::Numpy(arr))); + } + + return Err(PyTypeError::new_err("Unexpected dtype of numpy array.")); + } + + Err(PyTypeError::new_err("Expected numpy array input.")) + } +} + +fn is_type(dtype: &PyArrayDescr) -> bool { + Python::with_gil(|py| dtype.is_equiv_to(dtype_bound::(py).as_gil_ref())) +} diff --git a/python/src/rtree.rs b/python/src/rtree.rs index 58385ca..d407110 100644 --- a/python/src/rtree.rs +++ b/python/src/rtree.rs @@ -3,8 +3,8 @@ use geo_index::rtree::sort::{HilbertSort, STRSort}; use geo_index::rtree::util::f64_box_to_f32; use geo_index::rtree::{OwnedRTree, RTreeBuilder, RTreeIndex, TreeMetadata}; use geo_index::{CoordType, IndexableNum}; -use numpy::ndarray::{ArrayView1, ArrayView2}; -use numpy::{PyArray1, PyArrayMethods, PyReadonlyArray1, PyReadonlyArray2}; +use numpy::ndarray::ArrayView2; +use numpy::{PyArray1, PyArrayMethods, PyReadonlyArray2}; use pyo3::exceptions::{PyIndexError, PyTypeError, PyValueError}; use pyo3::ffi; use pyo3::intern; @@ -12,7 +12,7 @@ use pyo3::prelude::*; use pyo3::types::PyType; use std::os::raw::c_int; -use crate::common::{PyU8Buffer, RustBuffer}; +use crate::common::{PyArray, PyU8Buffer, RustBuffer}; /// Method for constructing rtree enum RTreeMethod { @@ -113,6 +113,42 @@ impl<'py> FromPyObject<'py> for PyRTreeRef { } } +impl From> for PyRTreeRef { + fn from(value: PyRTreeBuffer) -> Self { + Self::Int8(value) + } +} + +impl From> for PyRTreeRef { + fn from(value: PyRTreeBuffer) -> Self { + Self::Int16(value) + } +} + +impl From> for PyRTreeRef { + fn from(value: PyRTreeBuffer) -> Self { + Self::Int32(value) + } +} + +impl From> for PyRTreeRef { + fn from(value: PyRTreeBuffer) -> Self { + Self::UInt8(value) + } +} + +impl From> for PyRTreeRef { + fn from(value: PyRTreeBuffer) -> Self { + Self::UInt16(value) + } +} + +impl From> for PyRTreeRef { + fn from(value: PyRTreeBuffer) -> Self { + Self::UInt32(value) + } +} + impl From> for PyRTreeRef { fn from(value: PyRTreeBuffer) -> Self { Self::Float32(value) @@ -282,72 +318,158 @@ impl RTree { text_signature = "(min_x, min_y, max_x, max_y, *, method = 'hilbert', node_size = None)") ] #[allow(clippy::too_many_arguments)] - fn from_separated( + fn from_separated<'py>( _cls: &Bound, - py: Python, - min_x: PyObject, - min_y: PyObject, - max_x: PyObject, - max_y: PyObject, + py: Python<'py>, + min_x: PyArray<'py>, + min_y: PyArray<'py>, + max_x: PyArray<'py>, + max_y: PyArray<'py>, method: RTreeMethod, node_size: Option, ) -> PyResult { - // Convert to numpy array (of the same dtype) - let min_x = min_x.call_method0(py, intern!(py, "__array__"))?; - let min_y = min_y.call_method0(py, intern!(py, "__array__"))?; - let max_x = max_x.call_method0(py, intern!(py, "__array__"))?; - let max_y = max_y.call_method0(py, intern!(py, "__array__"))?; - - let result = if let Ok(min_x) = min_x.extract::>(py) { - let min_y = min_y.extract::>(py)?; - let max_x = max_x.extract::>(py)?; - let max_y = max_y.extract::>(py)?; - - let min_x_array = min_x.as_array(); - let min_y_array = min_y.as_array(); - let max_x_array = max_x.as_array(); - let max_y_array = max_y.as_array(); - - let tree = py.allow_threads(|| { - new_separated( - min_x_array, - min_y_array, - max_x_array, - max_y_array, + match (min_x, min_y, max_x, max_y) { + ( + PyArray::Int8(min_x), + PyArray::Int8(min_y), + PyArray::Int8(max_x), + PyArray::Int8(max_y), + ) => { + let tree = new_separated_slice( + min_x.as_slice(), + min_y.as_slice(), + max_x.as_slice(), + max_y.as_slice(), method, node_size, - ) - }); - Ok(Self(PyRTreeBuffer::from_owned_rtree(py, tree)?.into())) - } else if let Ok(min_x) = min_x.extract::>(py) { - let min_y = min_y.extract::>(py)?; - let max_x = max_x.extract::>(py)?; - let max_y = max_y.extract::>(py)?; - - let min_x_array = min_x.as_array(); - let min_y_array = min_y.as_array(); - let max_x_array = max_x.as_array(); - let max_y_array = max_y.as_array(); - - let tree = py.allow_threads(|| { - new_separated( - min_x_array, - min_y_array, - max_x_array, - max_y_array, + ); + + Ok(Self(PyRTreeBuffer::from_owned_rtree(py, tree)?.into())) + } + ( + PyArray::Int16(min_x), + PyArray::Int16(min_y), + PyArray::Int16(max_x), + PyArray::Int16(max_y), + ) => { + let tree = new_separated_slice( + min_x.as_slice(), + min_y.as_slice(), + max_x.as_slice(), + max_y.as_slice(), method, node_size, - ) - }); - Ok(Self(PyRTreeBuffer::from_owned_rtree(py, tree)?.into())) - } else { - let dtype = min_x.call_method0(py, intern!(py, "dtype"))?.to_string(); - Err(PyTypeError::new_err(format!( - "Expected a numpy array of dtype float32 or float64, got {}", - dtype - ))) - }; - result + ); + + Ok(Self(PyRTreeBuffer::from_owned_rtree(py, tree)?.into())) + } + ( + PyArray::Int32(min_x), + PyArray::Int32(min_y), + PyArray::Int32(max_x), + PyArray::Int32(max_y), + ) => { + let tree = new_separated_slice( + min_x.as_slice(), + min_y.as_slice(), + max_x.as_slice(), + max_y.as_slice(), + method, + node_size, + ); + + Ok(Self(PyRTreeBuffer::from_owned_rtree(py, tree)?.into())) + } + ( + PyArray::UInt8(min_x), + PyArray::UInt8(min_y), + PyArray::UInt8(max_x), + PyArray::UInt8(max_y), + ) => { + let tree = new_separated_slice( + min_x.as_slice(), + min_y.as_slice(), + max_x.as_slice(), + max_y.as_slice(), + method, + node_size, + ); + + Ok(Self(PyRTreeBuffer::from_owned_rtree(py, tree)?.into())) + } + ( + PyArray::UInt16(min_x), + PyArray::UInt16(min_y), + PyArray::UInt16(max_x), + PyArray::UInt16(max_y), + ) => { + let tree = new_separated_slice( + min_x.as_slice(), + min_y.as_slice(), + max_x.as_slice(), + max_y.as_slice(), + method, + node_size, + ); + + Ok(Self(PyRTreeBuffer::from_owned_rtree(py, tree)?.into())) + } + ( + PyArray::UInt32(min_x), + PyArray::UInt32(min_y), + PyArray::UInt32(max_x), + PyArray::UInt32(max_y), + ) => { + let tree = new_separated_slice( + min_x.as_slice(), + min_y.as_slice(), + max_x.as_slice(), + max_y.as_slice(), + method, + node_size, + ); + + Ok(Self(PyRTreeBuffer::from_owned_rtree(py, tree)?.into())) + } + + ( + PyArray::Float32(min_x), + PyArray::Float32(min_y), + PyArray::Float32(max_x), + PyArray::Float32(max_y), + ) => { + let tree = new_separated_slice( + min_x.as_slice(), + min_y.as_slice(), + max_x.as_slice(), + max_y.as_slice(), + method, + node_size, + ); + + Ok(Self(PyRTreeBuffer::from_owned_rtree(py, tree)?.into())) + } + ( + PyArray::Float64(min_x), + PyArray::Float64(min_y), + PyArray::Float64(max_x), + PyArray::Float64(max_y), + ) => { + let tree = new_separated_slice( + min_x.as_slice(), + min_y.as_slice(), + max_x.as_slice(), + max_y.as_slice(), + method, + node_size, + ); + + Ok(Self(PyRTreeBuffer::from_owned_rtree(py, tree)?.into())) + } + _ => Err(PyTypeError::new_err( + "Expected all input arrays to have the same data type", + )), + } } // pre PEP 688 buffer protocol @@ -528,11 +650,11 @@ fn new_interleaved( } } -fn new_separated( - min_x: ArrayView1, - min_y: ArrayView1, - max_x: ArrayView1, - max_y: ArrayView1, +fn new_separated_slice( + min_x: &[N], + min_y: &[N], + max_x: &[N], + max_y: &[N], method: RTreeMethod, node_size: Option, ) -> OwnedRTree { From bb9433def3c2784a9b9fea3ffbfd2fd12b46acd8 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Thu, 22 Aug 2024 00:34:09 -0400 Subject: [PATCH 10/12] fix typing --- python/python/geoindex_rs/_rust.pyi | 3 ++- python/tests/test_buffers.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/python/python/geoindex_rs/_rust.pyi b/python/python/geoindex_rs/_rust.pyi index ee0899c..9cd5e46 100644 --- a/python/python/geoindex_rs/_rust.pyi +++ b/python/python/geoindex_rs/_rust.pyi @@ -36,9 +36,10 @@ class BufferProtocolExportable(Protocol): def __buffer__(self, flags: int) -> memoryview: ... class RTree(Generic[T]): - def __init__(self, obj: BufferProtocolExportable) -> None: ... def __buffer__(self, flags: int) -> memoryview: ... @classmethod + def from_buffer(cls, obj: BufferProtocolExportable) -> None: ... + @classmethod def from_interleaved( cls, boxes: NDArray[T], diff --git a/python/tests/test_buffers.py b/python/tests/test_buffers.py index f98ca8d..f6a7ade 100644 --- a/python/tests/test_buffers.py +++ b/python/tests/test_buffers.py @@ -12,7 +12,7 @@ def generate_random_boxes(): def test_buffer_protocol(): boxes = generate_random_boxes() initial = RTree.from_interleaved(boxes) - second = RTree(initial) + second = RTree.from_buffer(initial) assert initial is not second, "Not the same object" From 47f85570be64a459de48e32b00469413c0ccc9df Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 15 Oct 2024 21:06:18 -0400 Subject: [PATCH 11/12] Progress --- python/Cargo.lock | 975 ++++++++++++++++++++++++++++++++++++++++--- python/Cargo.toml | 14 +- python/src/common.rs | 169 -------- python/src/lib.rs | 1 - python/src/rtree.rs | 45 +- 5 files changed, 943 insertions(+), 261 deletions(-) delete mode 100644 python/src/common.rs diff --git a/python/Cargo.lock b/python/Cargo.lock index 35c9322..489ab3b 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -2,18 +2,301 @@ # It is not intended for manual editing. version = 3 +[[package]] +name = "ahash" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011" +dependencies = [ + "cfg-if", + "const-random", + "getrandom", + "once_cell", + "version_check", + "zerocopy", +] + +[[package]] +name = "aho-corasick" +version = "1.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" +dependencies = [ + "memchr", +] + +[[package]] +name = "android-tzdata" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0" + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + +[[package]] +name = "arrow" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9ba0d7248932f4e2a12fb37f0a2e3ec82b3bdedbac2a1dce186e036843b8f8c" +dependencies = [ + "arrow-arith", + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-csv", + "arrow-data", + "arrow-ipc", + "arrow-json", + "arrow-ord", + "arrow-row", + "arrow-schema", + "arrow-select", + "arrow-string", +] + +[[package]] +name = "arrow-arith" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d60afcdc004841a5c8d8da4f4fa22d64eb19c0c01ef4bcedd77f175a7cf6e38f" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "half", + "num", +] + +[[package]] +name = "arrow-array" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f16835e8599dbbb1659fd869d865254c4cf32c6c2bb60b6942ac9fc36bfa5da" +dependencies = [ + "ahash", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "chrono-tz", + "half", + "hashbrown 0.14.5", + "num", +] + +[[package]] +name = "arrow-buffer" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a1f34f0faae77da6b142db61deba2cb6d60167592b178be317b341440acba80" +dependencies = [ + "bytes", + "half", + "num", +] + +[[package]] +name = "arrow-cast" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "450e4abb5775bca0740bec0bcf1b1a5ae07eff43bd625661c4436d8e8e4540c4" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "atoi", + "base64", + "chrono", + "half", + "lexical-core", + "num", + "ryu", +] + +[[package]] +name = "arrow-csv" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3a4e4d63830a341713e35d9a42452fbc6241d5f42fa5cf6a4681b8ad91370c4" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-schema", + "chrono", + "csv", + "csv-core", + "lazy_static", + "lexical-core", + "regex", +] + +[[package]] +name = "arrow-data" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b1e618bbf714c7a9e8d97203c806734f012ff71ae3adc8ad1b075689f540634" +dependencies = [ + "arrow-buffer", + "arrow-schema", + "half", + "num", +] + +[[package]] +name = "arrow-ipc" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f98e983549259a2b97049af7edfb8f28b8911682040e99a94e4ceb1196bd65c2" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-schema", + "flatbuffers", +] + +[[package]] +name = "arrow-json" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b198b9c6fcf086501730efbbcb483317b39330a116125af7bb06467d04b352a3" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-schema", + "chrono", + "half", + "indexmap", + "lexical-core", + "num", + "serde", + "serde_json", +] + +[[package]] +name = "arrow-ord" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2427f37b4459a4b9e533045abe87a5183a5e0995a3fc2c2fd45027ae2cc4ef3f" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "half", + "num", +] + +[[package]] +name = "arrow-row" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15959657d92e2261a7a323517640af87f5afd9fd8a6492e424ebee2203c567f6" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "half", +] + +[[package]] +name = "arrow-schema" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbf0388a18fd7f7f3fe3de01852d30f54ed5182f9004db700fbe3ba843ed2794" +dependencies = [ + "bitflags 2.6.0", +] + +[[package]] +name = "arrow-select" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b83e5723d307a38bf00ecd2972cd078d1339c7fd3eb044f609958a9a24463f3a" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "num", +] + +[[package]] +name = "arrow-string" +version = "53.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ab3db7c09dd826e74079661d84ed01ed06547cf75d52c2818ef776d0d852305" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "memchr", + "num", + "regex", + "regex-syntax", +] + +[[package]] +name = "atoi" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" +dependencies = [ + "num-traits", +] + [[package]] name = "autocfg" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f1fdabc7756949593fe60f30ec81974b613357de856987752631dea1e3394c80" +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + [[package]] name = "bitflags" version = "1.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" +[[package]] +name = "bitflags" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de" + +[[package]] +name = "bumpalo" +version = "3.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79296716171880943b8470b5f8d03aa55eb2e645a4874bdbb28adb49162e012c" + [[package]] name = "bytemuck" version = "1.15.0" @@ -26,12 +309,80 @@ version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "514de17de45fdb8dc022b1a7975556c53c86f9f0aa5f534b98977b171857c2c9" +[[package]] +name = "cc" +version = "1.1.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b16803a61b81d9eabb7eae2588776c4c1e584b738ede45fdbb4c972cec1e9945" +dependencies = [ + "shlex", +] + [[package]] name = "cfg-if" version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "chrono" +version = "0.4.38" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a21f936df1771bf62b77f047b726c4625ff2e8aa607c01ec06e5a05bd8463401" +dependencies = [ + "android-tzdata", + "iana-time-zone", + "num-traits", + "windows-targets", +] + +[[package]] +name = "chrono-tz" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd6dd8046d00723a59a2f8c5f295c515b9bb9a331ee4f8f3d4dd49e428acd3b6" +dependencies = [ + "chrono", + "chrono-tz-build", + "phf", +] + +[[package]] +name = "chrono-tz-build" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e94fea34d77a245229e7746bd2beb786cd2a896f306ff491fb8cecb3074b10a7" +dependencies = [ + "parse-zoneinfo", + "phf_codegen", +] + +[[package]] +name = "const-random" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" +dependencies = [ + "const-random-macro", +] + +[[package]] +name = "const-random-macro" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" +dependencies = [ + "getrandom", + "once_cell", + "tiny-keccak", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + [[package]] name = "crossbeam-deque" version = "0.8.5" @@ -57,12 +408,55 @@ version = "0.8.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "248e3bacc7dc6baa3b21e405ee045c3047101a49145e7e9eca583ab4c2ca5345" +[[package]] +name = "crunchy" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7" + +[[package]] +name = "csv" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac574ff4d437a7b5ad237ef331c17ccca63c46479e5b5453eb8e10bb99a759fe" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "csv-core" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5efa2b3d7902f4b634a20cae3c9c4e6209dc4779feb6863329607560143efa70" +dependencies = [ + "memchr", +] + [[package]] name = "either" version = "1.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "11157ac094ffbdde99aa67b23417ebdd801842852b500e395a45a9c0aac03e4a" +[[package]] +name = "equivalent" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" + +[[package]] +name = "flatbuffers" +version = "24.3.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8add37afff2d4ffa83bc748a70b4b1370984f6980768554182424ef71447c35f" +dependencies = [ + "bitflags 1.3.2", + "rustc_version", +] + [[package]] name = "float_next_after" version = "1.0.0" @@ -85,18 +479,87 @@ dependencies = [ name = "geoindex-rs" version = "0.2.0-beta.1" dependencies = [ + "arrow", "bytes", "geo-index", "numpy", "pyo3", + "pyo3-arrow", "thiserror", ] +[[package]] +name = "getrandom" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + +[[package]] +name = "half" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dd08c532ae367adf81c312a4580bc67f1d0fe8bc9c460520283f4c0ff277888" +dependencies = [ + "cfg-if", + "crunchy", + "num-traits", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e087f84d4f86bf4b218b927129862374b72199ae7d8657835f1e89000eea4fb" + [[package]] name = "heck" -version = "0.4.1" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "iana-time-zone" +version = "0.1.61" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "235e081f3925a06703c2d0117ea8b91f042756fd6e7a6e5d901e8ca1a996b220" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "indexmap" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "707907fe3c25f5424cce2cb7e1cbcafee6bdbe735ca90ef77c29e84591e5b9da" +dependencies = [ + "equivalent", + "hashbrown 0.15.0", +] [[package]] name = "indoc" @@ -104,6 +567,91 @@ version = "2.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5" +[[package]] +name = "itoa" +version = "1.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" + +[[package]] +name = "js-sys" +version = "0.3.72" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a88f1bda2bd75b0452a14784937d796722fdebfe50df998aeb3f0b7603019a9" +dependencies = [ + "wasm-bindgen", +] + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "lexical-core" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0431c65b318a590c1de6b8fd6e72798c92291d27762d94c9e6c37ed7a73d8458" +dependencies = [ + "lexical-parse-float", + "lexical-parse-integer", + "lexical-util", + "lexical-write-float", + "lexical-write-integer", +] + +[[package]] +name = "lexical-parse-float" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb17a4bdb9b418051aa59d41d65b1c9be5affab314a872e5ad7f06231fb3b4e0" +dependencies = [ + "lexical-parse-integer", + "lexical-util", + "static_assertions", +] + +[[package]] +name = "lexical-parse-integer" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5df98f4a4ab53bf8b175b363a34c7af608fe31f93cc1fb1bf07130622ca4ef61" +dependencies = [ + "lexical-util", + "static_assertions", +] + +[[package]] +name = "lexical-util" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85314db53332e5c192b6bca611fb10c114a80d1b831ddac0af1e9be1b9232ca0" +dependencies = [ + "static_assertions", +] + +[[package]] +name = "lexical-write-float" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e7c3ad4e37db81c1cbe7cf34610340adc09c322871972f74877a712abc6c809" +dependencies = [ + "lexical-util", + "lexical-write-integer", + "static_assertions", +] + +[[package]] +name = "lexical-write-integer" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb89e9f6958b83258afa3deed90b5de9ef68eef090ad5086c791cd2345610162" +dependencies = [ + "lexical-util", + "static_assertions", +] + [[package]] name = "libc" version = "0.2.158" @@ -111,14 +659,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8adc4bb1803a324070e64a98ae98f38934d91957a99cfb3a43dcbc01bc56439" [[package]] -name = "lock_api" -version = "0.4.11" +name = "libm" +version = "0.2.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c168f8615b12bc01f9c17e2eb0cc07dcae1940121185446edc3744920e8ef45" -dependencies = [ - "autocfg", - "scopeguard", -] +checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058" + +[[package]] +name = "log" +version = "0.4.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24" [[package]] name = "matrixmultiply" @@ -130,6 +680,12 @@ dependencies = [ "rawpointer", ] +[[package]] +name = "memchr" +version = "2.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" + [[package]] name = "memoffset" version = "0.9.1" @@ -152,6 +708,30 @@ dependencies = [ "rawpointer", ] +[[package]] +name = "num" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" +dependencies = [ + "num-bigint", + "num-complex", + "num-integer", + "num-iter", + "num-rational", + "num-traits", +] + +[[package]] +name = "num-bigint" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-complex" version = "0.4.6" @@ -170,6 +750,28 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-iter" +version = "0.1.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1429034a0490724d0075ebb2bc9e875d6503c3cf69e235a8941aa757d83ef5bf" +dependencies = [ + "autocfg", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -177,13 +779,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", + "libm", ] [[package]] name = "numpy" -version = "0.21.0" -source = "git+https://github.com/stinodego/rust-numpy.git?rev=9ba9962ae57ba26e35babdce6f179edf5fe5b9c8#9ba9962ae57ba26e35babdce6f179edf5fe5b9c8" +version = "0.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf314fca279e6e6ac2126a4ff98f26d88aa4ad06bc68fb6ae5cf4bd706758311" dependencies = [ + "half", "libc", "ndarray", "num-complex", @@ -200,26 +805,50 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" [[package]] -name = "parking_lot" -version = "0.12.1" +name = "parse-zoneinfo" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" +checksum = "1f2a05b18d44e2957b88f96ba460715e295bc1d7510468a2f3d3b44535d26c24" dependencies = [ - "lock_api", - "parking_lot_core", + "regex", ] [[package]] -name = "parking_lot_core" -version = "0.9.9" +name = "phf" +version = "0.11.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c42a9226546d68acdd9c0a280d17ce19bfe27a46bf68784e4066115788d008e" +checksum = "ade2d8b8f33c7333b51bcf0428d37e217e9f32192ae4772156f65063b8ce03dc" dependencies = [ - "cfg-if", - "libc", - "redox_syscall", - "smallvec", - "windows-targets", + "phf_shared", +] + +[[package]] +name = "phf_codegen" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8d39688d359e6b34654d328e262234662d16cc0f60ec8dcbe5e718709342a5a" +dependencies = [ + "phf_generator", + "phf_shared", +] + +[[package]] +name = "phf_generator" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48e4cc64c2ad9ebe670cb8fd69dd50ae301650392e81c05f9bfcb2d5bdbc24b0" +dependencies = [ + "phf_shared", + "rand", +] + +[[package]] +name = "phf_shared" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90fcb95eef784c2ac79119d1dd819e162b5da872ce6f3c3abe1e8ca1c082f72b" +dependencies = [ + "siphasher", ] [[package]] @@ -239,15 +868,17 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.21.2" +version = "0.22.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5e00b96a521718e08e03b1a622f01c8a8deb50719335de3f60b3b3950f069d8" +checksum = "3d922163ba1f79c04bc49073ba7b32fd5a8d3b76a87c955921234b8e77333c51" dependencies = [ "cfg-if", + "chrono", + "indexmap", "indoc", "libc", "memoffset", - "parking_lot", + "once_cell", "portable-atomic", "pyo3-build-config", "pyo3-ffi", @@ -255,11 +886,27 @@ dependencies = [ "unindent", ] +[[package]] +name = "pyo3-arrow" +version = "0.5.1" +source = "git+https://github.com/kylebarron/arro3?rev=64a659133d79f0f875cf8e87f30ccd1c34fc8ad3#64a659133d79f0f875cf8e87f30ccd1c34fc8ad3" +dependencies = [ + "arrow", + "arrow-array", + "arrow-buffer", + "arrow-schema", + "half", + "indexmap", + "numpy", + "pyo3", + "thiserror", +] + [[package]] name = "pyo3-build-config" -version = "0.21.2" +version = "0.22.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7883df5835fafdad87c0d888b266c8ec0f4c9ca48a5bed6bbb592e8dedee1b50" +checksum = "bc38c5feeb496c8321091edf3d63e9a6829eab4b863b4a6a65f26f3e9cc6b179" dependencies = [ "once_cell", "target-lexicon", @@ -267,9 +914,9 @@ dependencies = [ [[package]] name = "pyo3-ffi" -version = "0.21.2" +version = "0.22.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "01be5843dc60b916ab4dad1dca6d20b9b4e6ddc8e15f50c47fe6d85f1fb97403" +checksum = "94845622d88ae274d2729fcefc850e63d7a3ddff5e3ce11bd88486db9f1d357d" dependencies = [ "libc", "pyo3-build-config", @@ -277,9 +924,9 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.21.2" +version = "0.22.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77b34069fc0682e11b31dbd10321cbf94808394c56fd996796ce45217dfac53c" +checksum = "e655aad15e09b94ffdb3ce3d217acf652e26bbc37697ef012f5e5e348c716e5e" dependencies = [ "proc-macro2", "pyo3-macros-backend", @@ -289,9 +936,9 @@ dependencies = [ [[package]] name = "pyo3-macros-backend" -version = "0.21.2" +version = "0.22.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08260721f32db5e1a5beae69a55553f56b99bd0e1c3e6e0a5e8851a9d0f5a85c" +checksum = "ae1e3f09eecd94618f60a455a23def79f79eba4dc561a97324bf9ac8c6df30ce" dependencies = [ "heck", "proc-macro2", @@ -309,6 +956,21 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" + [[package]] name = "rawpointer" version = "0.2.1" @@ -336,14 +998,34 @@ dependencies = [ ] [[package]] -name = "redox_syscall" -version = "0.4.1" +name = "regex" +version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa" +checksum = "38200e5ee88914975b69f657f0801b6f6dccafd44fd9326302a4aaeecfacb1d8" dependencies = [ - "bitflags", + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", ] +[[package]] +name = "regex-automata" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "368758f23274712b504848e9d5a6f010445cc8b87a7cdb4d7cbee666c1288da3" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" + [[package]] name = "rustc-hash" version = "1.1.0" @@ -351,16 +1033,75 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" [[package]] -name = "scopeguard" -version = "1.2.0" +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + +[[package]] +name = "ryu" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" + +[[package]] +name = "semver" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61697e0a1c7e512e84a621326239844a24d8207b4669b41bc18b32ea5cbf988b" + +[[package]] +name = "serde" +version = "1.0.210" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8e3592472072e6e22e0a54d5904d9febf8508f65fb8552499a1abc7d1078c3a" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.210" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "243902eda00fad750862fc144cea25caca5e20d615af0a81bee94ca738f1df1f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.128" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ff5456707a1de34e7e37f2a6fd3d3f808c318259cbd01ab6377795054b483d8" +dependencies = [ + "itoa", + "memchr", + "ryu", + "serde", +] + +[[package]] +name = "shlex" +version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" [[package]] -name = "smallvec" -version = "1.13.2" +name = "siphasher" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38b58827f4464d87d377d175e90bf58eb00fd8716ff0a62f80356b5e61555d0d" + +[[package]] +name = "static_assertions" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" [[package]] name = "syn" @@ -399,6 +1140,15 @@ dependencies = [ "syn", ] +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + [[package]] name = "tinyvec" version = "1.6.0" @@ -426,15 +1176,92 @@ version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c7de7d73e1754487cb58364ee906a499937a0dfabd86bcb980fa99ec8c8fa2ce" +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasi" +version = "0.11.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" + +[[package]] +name = "wasm-bindgen" +version = "0.2.95" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "128d1e363af62632b8eb57219c8fd7877144af57558fb2ef0368d0087bddeb2e" +dependencies = [ + "cfg-if", + "once_cell", + "wasm-bindgen-macro", +] + +[[package]] +name = "wasm-bindgen-backend" +version = "0.2.95" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb6dd4d3ca0ddffd1dd1c9c04f94b868c37ff5fac97c30b97cff2d74fce3a358" +dependencies = [ + "bumpalo", + "log", + "once_cell", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.95" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e79384be7f8f5a9dd5d7167216f022090cf1f9ec128e6e6a482a2cb5c5422c56" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.95" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26c6ab57572f7a24a4985830b120de1594465e5d500f24afe89e16b4e833ef68" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-backend", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.95" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65fc09f10666a9f147042251e0dda9c18f166ff7de300607007e96bdebc1068d" + +[[package]] +name = "windows-core" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" +dependencies = [ + "windows-targets", +] + [[package]] name = "windows-targets" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ "windows_aarch64_gnullvm", "windows_aarch64_msvc", "windows_i686_gnu", + "windows_i686_gnullvm", "windows_i686_msvc", "windows_x86_64_gnu", "windows_x86_64_gnullvm", @@ -443,42 +1270,68 @@ dependencies = [ [[package]] name = "windows_aarch64_gnullvm" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" [[package]] name = "windows_aarch64_msvc" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" [[package]] name = "windows_i686_gnu" -version = "0.48.5" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" [[package]] name = "windows_i686_msvc" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" [[package]] name = "windows_x86_64_gnu" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" [[package]] name = "windows_x86_64_gnullvm" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" [[package]] name = "windows_x86_64_msvc" -version = "0.48.5" +version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "zerocopy" +version = "0.7.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.7.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] diff --git a/python/Cargo.toml b/python/Cargo.toml index 7cb543f..079c4e8 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -17,17 +17,13 @@ name = "_rust" crate-type = ["cdylib"] [dependencies] -# arrow = "52" +arrow = "53" bytes = "1" geo-index = { path = "../", features = ["rayon"] } -# numpy = "0.21" -# TODO: Pin to released version once NumPy 2.0 support is merged -# https://github.com/PyO3/rust-numpy/issues/409 -# 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.0", features = ["macros"] } -# pyo3-arrow = "*" +numpy = "0.22" +pyo3 = { version = "0.22.0", features = ["macros"] } +# pyo3-arrow = "0.5.1" +pyo3-arrow = { git = "https://github.com/kylebarron/arro3", rev = "64a659133d79f0f875cf8e87f30ccd1c34fc8ad3" } thiserror = "1" [profile.release] diff --git a/python/src/common.rs b/python/src/common.rs deleted file mode 100644 index c275b38..0000000 --- a/python/src/common.rs +++ /dev/null @@ -1,169 +0,0 @@ -use geo_index::IndexableNum; -use numpy::{dtype_bound, PyArray1, PyArrayDescr, PyUntypedArray}; -use pyo3::buffer::PyBuffer; -use pyo3::exceptions::{PyTypeError, PyValueError}; -use pyo3::ffi; -use pyo3::prelude::*; -use std::os::raw::c_int; - -/// A Rust buffer that implements the Python buffer protocol -#[pyclass(name = "Buffer")] -pub(crate) struct RustBuffer(Vec); - -impl RustBuffer { - pub(crate) fn new(buffer: Vec) -> Self { - Self(buffer) - } -} - -#[pymethods] -impl RustBuffer { - /// Implements the buffer protocol export - unsafe fn __getbuffer__( - slf: PyRef<'_, Self>, - view: *mut ffi::Py_buffer, - flags: c_int, - ) -> PyResult<()> { - let bytes = slf.0.as_slice(); - let ret = ffi::PyBuffer_FillInfo( - view, - slf.as_ptr() as *mut _, - bytes.as_ptr() as *mut _, - bytes.len().try_into().unwrap(), - 1, // read only - flags, - ); - if ret == -1 { - return Err(PyErr::fetch(slf.py())); - } - Ok(()) - } - - unsafe fn __releasebuffer__(&self, _view: *mut ffi::Py_buffer) { - // is there anything to do here? - } -} - -/// A Rust representation of a Python object that implements the Python buffer protocol, exporting -/// a 1-dimensional `&[u8]` slice. -pub(crate) struct PyU8Buffer(PyBuffer); - -impl<'py> FromPyObject<'py> for PyU8Buffer { - fn extract_bound(obj: &Bound<'py, PyAny>) -> PyResult { - let buffer = PyBuffer::::get_bound(obj)?; - if !buffer.readonly() { - return Err(PyValueError::new_err("Must be read-only byte buffer.")); - } - if buffer.dimensions() != 1 { - return Err(PyValueError::new_err("Expected 1-dimensional array.")); - } - // Note: this is probably superfluous for 1D array - if !buffer.is_c_contiguous() { - return Err(PyValueError::new_err("Expected c-contiguous array.")); - } - if buffer.len_bytes() == 0 { - return Err(PyValueError::new_err("Buffer has no data.")); - } - - Ok(Self(buffer)) - } -} - -impl AsRef<[u8]> for PyU8Buffer { - /// Extract a slice from a Python object implementing the buffer protocol - fn as_ref(&self) -> &[u8] { - let len = self.0.item_count(); - let data = self.0.buf_ptr() as *const u8; - unsafe { std::slice::from_raw_parts(data, len) } - } -} - -pub(crate) enum PyTypedArrayRef<'py, N: IndexableNum + numpy::Element> { - // Arrow((ArrayRef, PhantomData)), - Numpy(&'py PyArray1), -} - -impl<'py, N: IndexableNum + numpy::Element> PyTypedArrayRef<'py, N> { - pub(crate) fn as_slice(&self) -> &[N] { - match self { - Self::Numpy(arr) => unsafe { arr.as_slice() }.unwrap(), - } - } -} - -pub(crate) enum PyArray<'py> { - Int8(PyTypedArrayRef<'py, i8>), - Int16(PyTypedArrayRef<'py, i16>), - Int32(PyTypedArrayRef<'py, i32>), - UInt8(PyTypedArrayRef<'py, u8>), - UInt16(PyTypedArrayRef<'py, u16>), - UInt32(PyTypedArrayRef<'py, u32>), - Float32(PyTypedArrayRef<'py, f32>), - Float64(PyTypedArrayRef<'py, f64>), -} - -impl<'py> FromPyObject<'py> for PyArray<'py> { - fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { - let mut ob = ob.to_owned(); - // call __array__ if it exists - if ob.hasattr("__array__")? { - ob = ob.call_method0("__array__")?; - } - - if let Ok(array) = ob.extract::<&'py PyUntypedArray>() { - if array.ndim() != 1 { - return Err(PyValueError::new_err("Expected 1-dimensional array.")); - } - - let dtype = array.dtype(); - - if is_type::(dtype) { - let arr = array.downcast::>()?; - return Ok(Self::Int8(PyTypedArrayRef::Numpy(arr))); - } - - if is_type::(dtype) { - let arr = array.downcast::>()?; - return Ok(Self::Int16(PyTypedArrayRef::Numpy(arr))); - } - - if is_type::(dtype) { - let arr = array.downcast::>()?; - return Ok(Self::Int32(PyTypedArrayRef::Numpy(arr))); - } - - if is_type::(dtype) { - let arr = array.downcast::>()?; - return Ok(Self::UInt8(PyTypedArrayRef::Numpy(arr))); - } - - if is_type::(dtype) { - let arr = array.downcast::>()?; - return Ok(Self::UInt16(PyTypedArrayRef::Numpy(arr))); - } - - if is_type::(dtype) { - let arr = array.downcast::>()?; - return Ok(Self::UInt32(PyTypedArrayRef::Numpy(arr))); - } - - if is_type::(dtype) { - let arr = array.downcast::>()?; - return Ok(Self::Float32(PyTypedArrayRef::Numpy(arr))); - } - - if is_type::(dtype) { - let arr = array.downcast::>()?; - return Ok(Self::Float64(PyTypedArrayRef::Numpy(arr))); - } - - return Err(PyTypeError::new_err("Unexpected dtype of numpy array.")); - } - - Err(PyTypeError::new_err("Expected numpy array input.")) - } -} - -fn is_type(dtype: &PyArrayDescr) -> bool { - Python::with_gil(|py| dtype.is_equiv_to(dtype_bound::(py).as_gil_ref())) -} diff --git a/python/src/lib.rs b/python/src/lib.rs index 5d9131a..c1907e2 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -1,4 +1,3 @@ -mod common; mod kdtree; mod rtree; diff --git a/python/src/rtree.rs b/python/src/rtree.rs index d407110..3f93135 100644 --- a/python/src/rtree.rs +++ b/python/src/rtree.rs @@ -1,3 +1,5 @@ +use arrow::buffer::Buffer; +use arrow::datatypes::DataType; use geo_index::indices::Indices; use geo_index::rtree::sort::{HilbertSort, STRSort}; use geo_index::rtree::util::f64_box_to_f32; @@ -10,10 +12,9 @@ use pyo3::ffi; use pyo3::intern; use pyo3::prelude::*; use pyo3::types::PyType; +use pyo3_arrow::buffer::PyArrowBuffer; use std::os::raw::c_int; -use crate::common::{PyArray, PyU8Buffer, RustBuffer}; - /// Method for constructing rtree enum RTreeMethod { Hilbert, @@ -21,7 +22,7 @@ enum RTreeMethod { } impl<'a> FromPyObject<'a> for RTreeMethod { - fn extract(ob: &'a PyAny) -> PyResult { + fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult { let s: String = ob.extract()?; match s.to_lowercase().as_str() { "hilbert" => Ok(Self::Hilbert), @@ -33,26 +34,24 @@ impl<'a> FromPyObject<'a> for RTreeMethod { } } -/// A low-level wrapper around a [PyU8Buffer] that validates that the input is a valid Flatbush +/// A low-level wrapper around a [PyArrowBuffer] that validates that the input is a valid Flatbush /// buffer. This wrapper implements [RTreeIndex]. pub(crate) struct PyRTreeBuffer { - buffer: PyU8Buffer, + buffer: PyArrowBuffer, metadata: TreeMetadata, } impl PyRTreeBuffer { - fn try_new(buffer: PyU8Buffer) -> PyResult { + fn try_new(buffer: PyArrowBuffer) -> PyResult { let metadata = TreeMetadata::try_new(buffer.as_ref()).unwrap(); Ok(Self { buffer, metadata }) } fn from_owned_rtree(py: Python, tree: OwnedRTree) -> PyResult { let metadata = tree.metadata().clone(); - let tree_buf = RustBuffer::new(tree.into_inner()); - Ok(Self { - buffer: tree_buf.into_py(py).extract(py)?, - metadata, - }) + // let buf = tree.into_inner(); + let buffer = PyArrowBuffer(Buffer::from_vec(tree.into_inner())); + Ok(Self { buffer, metadata }) } } @@ -103,12 +102,17 @@ pub(crate) enum PyRTreeRef { impl<'py> FromPyObject<'py> for PyRTreeRef { fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { - let buffer = PyU8Buffer::extract_bound(ob)?; - let ct = CoordType::from_buffer(&buffer.as_ref()).unwrap(); + let buffer = ob.extract::()?; + let ct = CoordType::from_buffer(&buffer.as_ref().as_ref()).unwrap(); match ct { + CoordType::Int8 => Ok(Self::Int8(PyRTreeBuffer::try_new(buffer)?)), + CoordType::Int16 => Ok(Self::Int16(PyRTreeBuffer::try_new(buffer)?)), + CoordType::Int32 => Ok(Self::Int32(PyRTreeBuffer::try_new(buffer)?)), + CoordType::UInt8 => Ok(Self::UInt8(PyRTreeBuffer::try_new(buffer)?)), + CoordType::UInt16 => Ok(Self::UInt16(PyRTreeBuffer::try_new(buffer)?)), + CoordType::UInt32 => Ok(Self::UInt32(PyRTreeBuffer::try_new(buffer)?)), CoordType::Float32 => Ok(Self::Float32(PyRTreeBuffer::try_new(buffer)?)), CoordType::Float64 => Ok(Self::Float64(PyRTreeBuffer::try_new(buffer)?)), - _ => todo!(), } } } @@ -287,12 +291,11 @@ impl RTree { fn from_interleaved( _cls: &Bound, py: Python, - boxes: PyObject, + boxes: pyo3_arrow::PyArray, method: RTreeMethod, node_size: Option, ) -> PyResult { - // Convert to numpy array (of the same dtype) - let boxes = boxes.call_method0(py, intern!(py, "__array__"))?; + let data_type = boxes.array().data_type(); let result = if let Ok(boxes) = boxes.extract::>(py) { let boxes = boxes.as_array(); @@ -321,10 +324,10 @@ impl RTree { fn from_separated<'py>( _cls: &Bound, py: Python<'py>, - min_x: PyArray<'py>, - min_y: PyArray<'py>, - max_x: PyArray<'py>, - max_y: PyArray<'py>, + min_x: pyo3_arrow::PyArray, + min_y: pyo3_arrow::PyArray, + max_x: pyo3_arrow::PyArray, + max_y: pyo3_arrow::PyArray, method: RTreeMethod, node_size: Option, ) -> PyResult { From 6ca2eb29d8259a135e1beac56c13692ce84a589b Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Tue, 15 Oct 2024 22:42:25 -0400 Subject: [PATCH 12/12] updates --- python/Cargo.lock | 2 +- python/Cargo.toml | 2 +- python/src/rtree.rs | 24 +++++++++++++----------- 3 files changed, 15 insertions(+), 13 deletions(-) diff --git a/python/Cargo.lock b/python/Cargo.lock index 489ab3b..6bfad77 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -889,7 +889,7 @@ dependencies = [ [[package]] name = "pyo3-arrow" version = "0.5.1" -source = "git+https://github.com/kylebarron/arro3?rev=64a659133d79f0f875cf8e87f30ccd1c34fc8ad3#64a659133d79f0f875cf8e87f30ccd1c34fc8ad3" +source = "git+https://github.com/kylebarron/arro3?rev=45be4a12dd62cee025c5d0ecf8c8c081e13643ea#45be4a12dd62cee025c5d0ecf8c8c081e13643ea" dependencies = [ "arrow", "arrow-array", diff --git a/python/Cargo.toml b/python/Cargo.toml index 079c4e8..0d70120 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -23,7 +23,7 @@ geo-index = { path = "../", features = ["rayon"] } numpy = "0.22" pyo3 = { version = "0.22.0", features = ["macros"] } # pyo3-arrow = "0.5.1" -pyo3-arrow = { git = "https://github.com/kylebarron/arro3", rev = "64a659133d79f0f875cf8e87f30ccd1c34fc8ad3" } +pyo3-arrow = { git = "https://github.com/kylebarron/arro3", rev = "45be4a12dd62cee025c5d0ecf8c8c081e13643ea" } thiserror = "1" [profile.release] diff --git a/python/src/rtree.rs b/python/src/rtree.rs index 3f93135..ace5d19 100644 --- a/python/src/rtree.rs +++ b/python/src/rtree.rs @@ -13,6 +13,7 @@ use pyo3::intern; use pyo3::prelude::*; use pyo3::types::PyType; use pyo3_arrow::buffer::PyArrowBuffer; +use pyo3_arrow::PyArray; use std::os::raw::c_int; /// Method for constructing rtree @@ -47,10 +48,9 @@ impl PyRTreeBuffer { Ok(Self { buffer, metadata }) } - fn from_owned_rtree(py: Python, tree: OwnedRTree) -> PyResult { + fn from_owned_rtree(tree: OwnedRTree) -> PyResult { let metadata = tree.metadata().clone(); - // let buf = tree.into_inner(); - let buffer = PyArrowBuffer(Buffer::from_vec(tree.into_inner())); + let buffer = PyArrowBuffer::new(Buffer::from_vec(tree.into_inner())); Ok(Self { buffer, metadata }) } } @@ -103,7 +103,7 @@ pub(crate) enum PyRTreeRef { impl<'py> FromPyObject<'py> for PyRTreeRef { fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { let buffer = ob.extract::()?; - let ct = CoordType::from_buffer(&buffer.as_ref().as_ref()).unwrap(); + let ct = CoordType::from_buffer(&buffer).unwrap(); match ct { CoordType::Int8 => Ok(Self::Int8(PyRTreeBuffer::try_new(buffer)?)), CoordType::Int16 => Ok(Self::Int16(PyRTreeBuffer::try_new(buffer)?)), @@ -278,9 +278,11 @@ pub(crate) struct RTree(PyRTreeRef); #[pymethods] impl RTree { /// Construct an RTree from an existing RTree buffer + /// + /// You can pass any buffer protocol object into this constructor. #[classmethod] - fn from_buffer(_cls: &Bound, py: Python, obj: PyObject) -> PyResult { - Ok(Self(obj.extract(py)?)) + fn from_buffer(_cls: &Bound, obj: PyRTreeRef) -> Self { + Self(obj) } #[classmethod] @@ -291,7 +293,7 @@ impl RTree { fn from_interleaved( _cls: &Bound, py: Python, - boxes: pyo3_arrow::PyArray, + boxes: PyArray, method: RTreeMethod, node_size: Option, ) -> PyResult { @@ -324,10 +326,10 @@ impl RTree { fn from_separated<'py>( _cls: &Bound, py: Python<'py>, - min_x: pyo3_arrow::PyArray, - min_y: pyo3_arrow::PyArray, - max_x: pyo3_arrow::PyArray, - max_y: pyo3_arrow::PyArray, + min_x: PyArray, + min_y: PyArray, + max_x: PyArray, + max_y: PyArray, method: RTreeMethod, node_size: Option, ) -> PyResult {