You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@sedona.apache.org by "pwrliang (via GitHub)" <gi...@apache.org> on 2026/02/10 02:26:54 UTC
[PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
pwrliang opened a new pull request, #586:
URL: https://github.com/apache/sedona-db/pull/586
This PR adds Rust wrappers and tests for c/sedona-libgpuspatial. `libgpuspatial.rs` contains the Rust wrappers for `GpuSpatialIndexFloat2D` and `GpuSpatialRefiner`; `lib.rs` contains a high-level struct `GpuSpatial` that supports all functions of libgpuspatial.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #586:
URL: https://github.com/apache/sedona-db/pull/586#discussion_r2817584548
##########
c/sedona-libgpuspatial/src/predicate.rs:
##########
@@ -0,0 +1,61 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::os::raw::c_uint;
+
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum GpuSpatialRelationPredicate {
+ Equals,
+ Disjoint,
+ Touches,
+ Contains,
+ Covers,
+ Intersects,
+ Within,
+ CoveredBy,
+}
+
+#[allow(dead_code)] // not used if the GPU feature is disabled
Review Comment:
```suggestion
#[cfg(not(gpu_available))] // not used if the GPU feature is disabled
```
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,305 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo_types::Rect;
+
+mod error;
#[cfg(gpu_available)]
+mod libgpuspatial;
mod libgpuspatial_glue_bindgen;
+mod options;
+mod predicate;
+
+pub use error::GpuSpatialError;
+pub use options::GpuSpatialOptions;
+pub use predicate::GpuSpatialRelationPredicate;
+pub use sys::{GpuSpatialIndex, GpuSpatialRefiner};
+
+#[cfg(gpu_available)]
+mod sys {
+ use super::libgpuspatial;
+ use super::*;
+ use libgpuspatial::GpuSpatialRuntimeWrapper;
+ use std::sync::{Arc, Mutex};
+
+ pub type Result<T> = std::result::Result<T, GpuSpatialError>;
+
+ // Global Runtime State
+ unsafe impl Send for GpuSpatialRuntimeWrapper {}
+ unsafe impl Sync for GpuSpatialRuntimeWrapper {}
Review Comment:
These should go right below the struct definition for `GpuSpatialRuntimeWrapper` (makes it easier to spot members or functions of the struct that might not be Send or Sync)
##########
c/sedona-libgpuspatial/libgpuspatial/src/gpuspatial_c.cc:
##########
@@ -194,12 +192,25 @@ struct GpuSpatialIndexFloat2DExporter {
}
static int CProbe(self_t* self, SedonaSpatialIndexContext* context, const float* buf,
- uint32_t n_rects) {
- return SafeExecute(static_cast<context_t*>(context->private_data), [&] {
+ uint32_t n_rects,
+ void (*callback)(const uint32_t* build_indices,
+ const uint32_t* probe_indices, uint32_t length,
+ void* user_data),
+ void* user_data) {
+ // Do not use SafeExecute because this method is thread-safe and we don't want to set
+ // last_error for the whole index if one thread encounters an error
+ try {
auto* rects = reinterpret_cast<const spatial_index_t::box_t*>(buf);
auto& buff = static_cast<context_t*>(context->private_data)->payload;
use_index(self).Probe(rects, n_rects, &buff.build_indices, &buff.probe_indices);
- });
+ callback(buff.build_indices.data(), buff.probe_indices.data(),
+ buff.build_indices.size(), user_data);
+ return 0;
+ } catch (const std::exception& e) { // user should call context_get_last_error
+ return EINVAL;
+ } catch (...) {
+ return EINVAL;
+ }
Review Comment:
You should be able to set the context error string here safely, correct? If `use_index(self).Probe(...)` already does this, you can add a comment here. If the process of probing never returns a meaningful error message, you can set the context error message to "Unknown error" or something.
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,305 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo_types::Rect;
+
+mod error;
#[cfg(gpu_available)]
+mod libgpuspatial;
mod libgpuspatial_glue_bindgen;
+mod options;
+mod predicate;
+
+pub use error::GpuSpatialError;
+pub use options::GpuSpatialOptions;
+pub use predicate::GpuSpatialRelationPredicate;
+pub use sys::{GpuSpatialIndex, GpuSpatialRefiner};
+
+#[cfg(gpu_available)]
+mod sys {
+ use super::libgpuspatial;
+ use super::*;
+ use libgpuspatial::GpuSpatialRuntimeWrapper;
+ use std::sync::{Arc, Mutex};
+
+ pub type Result<T> = std::result::Result<T, GpuSpatialError>;
+
+ // Global Runtime State
+ unsafe impl Send for GpuSpatialRuntimeWrapper {}
+ unsafe impl Sync for GpuSpatialRuntimeWrapper {}
+
+ static GLOBAL_GPUSPATIAL_RUNTIME: Mutex<Option<Arc<GpuSpatialRuntimeWrapper>>> =
+ Mutex::new(None);
+ /// Handles initialization of the GPU runtime.
+ pub struct SpatialContext {
+ runtime: Arc<GpuSpatialRuntimeWrapper>,
+ }
+
+ impl SpatialContext {
+ pub fn try_new(options: &GpuSpatialOptions) -> Result<Self> {
+ // Lock the mutex globally
+ let mut guard = GLOBAL_GPUSPATIAL_RUNTIME
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Global mutex poisoned".into()))?;
+
+ // Check if it already exists
+ if let Some(existing_runtime) = guard.as_ref() {
+ if existing_runtime.device_id != options.device_id {
+ return Err(GpuSpatialError::Init(format!(
+ "Runtime conflict: Initialized on Device {}, requested Device {}.",
+ existing_runtime.device_id, options.device_id
+ )));
+ }
+ // Return the existing one
+ return Ok(Self {
+ runtime: existing_runtime.clone(),
+ });
+ }
Review Comment:
I think a `OnceLock<Arc<GpuSpatialRuntimeWrapper>>` is the Rusty way to do this. In action:
https://github.com/apache/sedona-db/blob/7eac903e08c84ae9aa6396d2ff1e8101d1e6b8df/rust/sedona-functions/src/st_setsrid.rs#L434-L448
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,305 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo_types::Rect;
+
+mod error;
#[cfg(gpu_available)]
+mod libgpuspatial;
mod libgpuspatial_glue_bindgen;
+mod options;
+mod predicate;
+
+pub use error::GpuSpatialError;
+pub use options::GpuSpatialOptions;
+pub use predicate::GpuSpatialRelationPredicate;
+pub use sys::{GpuSpatialIndex, GpuSpatialRefiner};
+
+#[cfg(gpu_available)]
+mod sys {
+ use super::libgpuspatial;
+ use super::*;
+ use libgpuspatial::GpuSpatialRuntimeWrapper;
+ use std::sync::{Arc, Mutex};
+
+ pub type Result<T> = std::result::Result<T, GpuSpatialError>;
+
+ // Global Runtime State
+ unsafe impl Send for GpuSpatialRuntimeWrapper {}
+ unsafe impl Sync for GpuSpatialRuntimeWrapper {}
+
+ static GLOBAL_GPUSPATIAL_RUNTIME: Mutex<Option<Arc<GpuSpatialRuntimeWrapper>>> =
+ Mutex::new(None);
+ /// Handles initialization of the GPU runtime.
+ pub struct SpatialContext {
+ runtime: Arc<GpuSpatialRuntimeWrapper>,
+ }
+
+ impl SpatialContext {
+ pub fn try_new(options: &GpuSpatialOptions) -> Result<Self> {
+ // Lock the mutex globally
+ let mut guard = GLOBAL_GPUSPATIAL_RUNTIME
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Global mutex poisoned".into()))?;
+
+ // Check if it already exists
+ if let Some(existing_runtime) = guard.as_ref() {
+ if existing_runtime.device_id != options.device_id {
+ return Err(GpuSpatialError::Init(format!(
+ "Runtime conflict: Initialized on Device {}, requested Device {}.",
+ existing_runtime.device_id, options.device_id
+ )));
+ }
+ // Return the existing one
+ return Ok(Self {
+ runtime: existing_runtime.clone(),
+ });
+ }
+
+ let out_path = std::path::PathBuf::from(env!("OUT_DIR"));
+ let ptx_root = out_path.join("share/gpuspatial/shaders");
+ let ptx_root_str = ptx_root
+ .to_str()
+ .ok_or_else(|| GpuSpatialError::Init("Invalid PTX path".to_string()))?;
+
+ let wrapper = libgpuspatial::GpuSpatialRuntimeWrapper::try_new(
+ options.device_id,
+ ptx_root_str,
+ options.cuda_use_memory_pool,
+ options.cuda_memory_pool_init_percent,
+ )?;
+
+ let arc_wrapper = Arc::new(wrapper);
+
+ *guard = Some(arc_wrapper.clone());
+
+ Ok(Self {
+ runtime: arc_wrapper,
+ })
+ }
+ }
+
+ pub struct GpuSpatialIndex {
+ inner: libgpuspatial::FloatIndex2D,
+ }
+
+ impl GpuSpatialIndex {
Review Comment:
I think eventually we will want to collapse all of this wrapper code into the `FloatIndex2D` but no need to do that now. One good reason to do this is that it means you only have to document it once. If we keep this wrapper for this PR, we still need to document both (although you can link to the inner documentation from the documentation for this struct).
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,502 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::cell::UnsafeCell;
+use std::convert::TryFrom;
+use std::ffi::{c_void, CStr, CString};
+use std::os::raw::c_char;
+use std::sync::Arc;
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: UnsafeCell<GpuSpatialRuntime>,
+ /// Store which device the runtime is created on
+ pub device_id: i32,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::Init("init function is None".to_string()));
+ }
+
+ Ok(GpuSpatialRuntimeWrapper {
+ runtime: UnsafeCell::new(runtime),
+ device_id,
+ })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ let runtime = self.runtime.get_mut();
+ let release_fn = runtime.release.expect("release function is None");
+ unsafe {
+ release_fn(runtime as *mut _);
+ }
+ }
+}
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D` struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<GpuSpatialRuntimeWrapper>,
+}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ let release_fn = self.index.release.expect("release function is None");
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+}
+
+pub struct FloatIndex2D {
+ inner: FloatIndex2DWrapper,
+}
+
Review Comment:
```suggestion
unsafe impl Send for FloatIndex2D;
unsafe impl Sync for FloatIndex2D;
```
##########
c/sedona-libgpuspatial/libgpuspatial/include/gpuspatial/gpuspatial_c.h:
##########
@@ -84,30 +84,23 @@ struct SedonaFloatIndex2D {
* @return 0 on success, non-zero on failure
*/
int (*finish_building)(struct SedonaFloatIndex2D* self);
+
/**
* Probe the spatial index with the given rectangles, each rectangle is represented by 4
* floats: [min_x, min_y, max_x, max_y] Points can also be probed by providing [x, y, x,
* y] but points and rectangles cannot be mixed in one Probe call. The results of the
- * probe will be stored in the context.
+ * probe will be stored in the context. The callback function will be called for each
+ * batch of results, with the build and probe indices of the candidate pairs in the
+ * batch. The user_data pointer will be passed to the callback
*
* @return 0 on success, non-zero on failure
*/
int (*probe)(struct SedonaFloatIndex2D* self, struct SedonaSpatialIndexContext* context,
- const float* buf, uint32_t n_rects);
- /** Get the build indices buffer from the context
- *
- * @return A pointer to the buffer and its length
- */
- void (*get_build_indices_buffer)(struct SedonaSpatialIndexContext* context,
- uint32_t** build_indices,
- uint32_t* build_indices_length);
- /** Get the probe indices buffer from the context
- *
- * @return A pointer to the buffer and its length
- */
- void (*get_probe_indices_buffer)(struct SedonaSpatialIndexContext* context,
- uint32_t** probe_indices,
- uint32_t* probe_indices_length);
+ const float* buf, uint32_t n_rects,
+ void (*callback)(const uint32_t* build_indices,
+ const uint32_t* probe_indices, uint32_t length,
+ void* user_data),
Review Comment:
We should let the callback return an error code. Right now you only visit exactly one pre-allocated buffer but hopefully we can split that up into multiple calls for huge probe results and returning an error code will allow that process to be cancelled.
```suggestion
int (*callback)(const uint32_t* build_indices,
const uint32_t* probe_indices, uint32_t length,
void* user_data),
```
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,502 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::cell::UnsafeCell;
+use std::convert::TryFrom;
+use std::ffi::{c_void, CStr, CString};
+use std::os::raw::c_char;
+use std::sync::Arc;
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: UnsafeCell<GpuSpatialRuntime>,
+ /// Store which device the runtime is created on
+ pub device_id: i32,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::Init("init function is None".to_string()));
+ }
+
+ Ok(GpuSpatialRuntimeWrapper {
+ runtime: UnsafeCell::new(runtime),
+ device_id,
+ })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ let runtime = self.runtime.get_mut();
+ let release_fn = runtime.release.expect("release function is None");
+ unsafe {
+ release_fn(runtime as *mut _);
+ }
+ }
+}
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D` struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<GpuSpatialRuntimeWrapper>,
+}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ let release_fn = self.index.release.expect("release function is None");
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+}
+
+pub struct FloatIndex2D {
+ inner: FloatIndex2DWrapper,
+}
+
+impl FloatIndex2D {
+ pub fn try_new(
+ runtime: Arc<GpuSpatialRuntimeWrapper>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let config = GpuSpatialIndexConfig {
+ runtime: runtime.runtime.get(),
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ return Err(GpuSpatialError::Init("Index Create failed".into()));
+ }
+ }
+
+ Ok(Self {
+ inner: FloatIndex2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ },
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.inner.index.clear {
+ unsafe {
+ clear_fn(&mut self.inner.index as *mut _);
+ }
+ }
+ }
+
+ pub fn push_build(&mut self, buf: *const f32, n_rects: u32) -> Result<(), GpuSpatialError> {
+ let push_fn =
+ self.inner.index.push_build.ok_or_else(|| {
+ GpuSpatialError::PushBuild("push_build function is None".to_string())
+ })?;
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut SedonaFloatIndex2D;
+
+ unsafe {
+ check_ffi_call(
+ move || push_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )
+ }
+ }
+
+ pub fn finish_building(&mut self) -> Result<(), GpuSpatialError> {
+ let finish_fn = self
+ .inner
+ .index
+ .finish_building
+ .ok_or_else(|| GpuSpatialError::FinishBuild("finish_building missing".into()))?;
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut SedonaFloatIndex2D;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_fn(&mut self.inner.index),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )
+ }
+ }
+
+ pub fn probe(
+ &self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(Vec<u32>, Vec<u32>), GpuSpatialError> {
+ let probe_fn = self
+ .inner
+ .index
+ .probe
+ .ok_or_else(|| GpuSpatialError::Probe("probe function is None".into()))?;
+ let create_context_fn = self.inner.index.create_context;
+ let destroy_context_fn = self.inner.index.destroy_context;
+ let context_err_fn = self.inner.index.context_get_last_error;
+ let index_ptr = &self.inner.index as *const _ as *mut SedonaFloatIndex2D;
+
+ let mut ctx = SedonaSpatialIndexContext {
+ private_data: std::ptr::null_mut(),
+ };
+ let mut state = ProbeState {
+ results: (Vec::new(), Vec::new()),
+ error: None,
+ };
+
+ unsafe {
+ if let Some(create_ctx) = create_context_fn {
+ create_ctx(&mut ctx);
+ }
+
+ let status = probe_fn(
+ index_ptr,
+ &mut ctx,
+ buf,
+ n_rects,
+ Some(probe_callback_wrapper),
+ &mut state as *mut _ as *mut c_void,
+ );
+
+ if status != 0 {
+ let error_string = if let Some(get_ctx_err) = context_err_fn {
+ CStr::from_ptr(get_ctx_err(&mut ctx))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error during probe".to_string()
+ };
+
+ if let Some(destroy_ctx) = destroy_context_fn {
+ destroy_ctx(&mut ctx);
+ }
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+
+ if let Some(callback_error) = state.error {
+ if let Some(destroy_ctx) = destroy_context_fn {
+ destroy_ctx(&mut ctx);
+ }
+ return Err(callback_error);
+ }
+
+ if let Some(destroy_ctx) = destroy_context_fn {
+ destroy_ctx(&mut ctx);
+ }
+ }
+
+ Ok(state.results)
+ }
+}
+
+struct RefinerWrapper {
+ refiner: SedonaSpatialRefiner,
+ _runtime: Arc<GpuSpatialRuntimeWrapper>,
+}
+
+impl Drop for RefinerWrapper {
+ fn drop(&mut self) {
+ let release_fn = self.refiner.release.expect("release function is None");
+ unsafe {
+ release_fn(&mut self.refiner as *mut _);
+ }
+ }
+}
+pub struct Refiner {
+ inner: RefinerWrapper,
+}
+
Review Comment:
```suggestion
unsafe impl Send for Refiner;
unsafe impl Sync for Refiner;
```
##########
c/sedona-libgpuspatial/libgpuspatial/src/gpuspatial_c.cc:
##########
@@ -194,12 +192,25 @@ struct GpuSpatialIndexFloat2DExporter {
}
static int CProbe(self_t* self, SedonaSpatialIndexContext* context, const float* buf,
- uint32_t n_rects) {
- return SafeExecute(static_cast<context_t*>(context->private_data), [&] {
+ uint32_t n_rects,
+ void (*callback)(const uint32_t* build_indices,
+ const uint32_t* probe_indices, uint32_t length,
+ void* user_data),
Review Comment:
This callback should return an error code (this will let long-running probe operations be cancelled if or when you split up the output of huge probe results).
```suggestion
int (*callback)(const uint32_t* build_indices,
const uint32_t* probe_indices, uint32_t length,
void* user_data),
```
##########
c/sedona-libgpuspatial/libgpuspatial/src/gpuspatial_c.cc:
##########
@@ -194,12 +192,25 @@ struct GpuSpatialIndexFloat2DExporter {
}
static int CProbe(self_t* self, SedonaSpatialIndexContext* context, const float* buf,
- uint32_t n_rects) {
- return SafeExecute(static_cast<context_t*>(context->private_data), [&] {
+ uint32_t n_rects,
+ void (*callback)(const uint32_t* build_indices,
+ const uint32_t* probe_indices, uint32_t length,
+ void* user_data),
+ void* user_data) {
+ // Do not use SafeExecute because this method is thread-safe and we don't want to set
+ // last_error for the whole index if one thread encounters an error
+ try {
auto* rects = reinterpret_cast<const spatial_index_t::box_t*>(buf);
auto& buff = static_cast<context_t*>(context->private_data)->payload;
use_index(self).Probe(rects, n_rects, &buff.build_indices, &buff.probe_indices);
- });
+ callback(buff.build_indices.data(), buff.probe_indices.data(),
+ buff.build_indices.size(), user_data);
+ return 0;
Review Comment:
```suggestion
return callback(buff.build_indices.data(), buff.probe_indices.data(),
buff.build_indices.size(), user_data);
```
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,502 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::cell::UnsafeCell;
+use std::convert::TryFrom;
+use std::ffi::{c_void, CStr, CString};
+use std::os::raw::c_char;
+use std::sync::Arc;
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: UnsafeCell<GpuSpatialRuntime>,
+ /// Store which device the runtime is created on
+ pub device_id: i32,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::Init("init function is None".to_string()));
+ }
+
+ Ok(GpuSpatialRuntimeWrapper {
+ runtime: UnsafeCell::new(runtime),
+ device_id,
+ })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ let runtime = self.runtime.get_mut();
+ let release_fn = runtime.release.expect("release function is None");
+ unsafe {
+ release_fn(runtime as *mut _);
+ }
+ }
+}
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D` struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<GpuSpatialRuntimeWrapper>,
+}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ let release_fn = self.index.release.expect("release function is None");
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+}
+
+pub struct FloatIndex2D {
+ inner: FloatIndex2DWrapper,
+}
+
+impl FloatIndex2D {
+ pub fn try_new(
+ runtime: Arc<GpuSpatialRuntimeWrapper>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let config = GpuSpatialIndexConfig {
+ runtime: runtime.runtime.get(),
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ return Err(GpuSpatialError::Init("Index Create failed".into()));
+ }
+ }
+
+ Ok(Self {
+ inner: FloatIndex2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ },
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.inner.index.clear {
+ unsafe {
+ clear_fn(&mut self.inner.index as *mut _);
+ }
+ }
+ }
+
+ pub fn push_build(&mut self, buf: *const f32, n_rects: u32) -> Result<(), GpuSpatialError> {
+ let push_fn =
+ self.inner.index.push_build.ok_or_else(|| {
+ GpuSpatialError::PushBuild("push_build function is None".to_string())
+ })?;
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut SedonaFloatIndex2D;
+
+ unsafe {
+ check_ffi_call(
+ move || push_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )
+ }
+ }
+
+ pub fn finish_building(&mut self) -> Result<(), GpuSpatialError> {
+ let finish_fn = self
+ .inner
+ .index
+ .finish_building
+ .ok_or_else(|| GpuSpatialError::FinishBuild("finish_building missing".into()))?;
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut SedonaFloatIndex2D;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_fn(&mut self.inner.index),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )
+ }
+ }
+
+ pub fn probe(
+ &self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(Vec<u32>, Vec<u32>), GpuSpatialError> {
+ let probe_fn = self
+ .inner
+ .index
+ .probe
+ .ok_or_else(|| GpuSpatialError::Probe("probe function is None".into()))?;
+ let create_context_fn = self.inner.index.create_context;
+ let destroy_context_fn = self.inner.index.destroy_context;
+ let context_err_fn = self.inner.index.context_get_last_error;
+ let index_ptr = &self.inner.index as *const _ as *mut SedonaFloatIndex2D;
+
+ let mut ctx = SedonaSpatialIndexContext {
+ private_data: std::ptr::null_mut(),
+ };
+ let mut state = ProbeState {
+ results: (Vec::new(), Vec::new()),
+ error: None,
+ };
+
+ unsafe {
+ if let Some(create_ctx) = create_context_fn {
+ create_ctx(&mut ctx);
+ }
+
+ let status = probe_fn(
+ index_ptr,
+ &mut ctx,
+ buf,
+ n_rects,
+ Some(probe_callback_wrapper),
+ &mut state as *mut _ as *mut c_void,
+ );
+
+ if status != 0 {
+ let error_string = if let Some(get_ctx_err) = context_err_fn {
+ CStr::from_ptr(get_ctx_err(&mut ctx))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error during probe".to_string()
+ };
+
+ if let Some(destroy_ctx) = destroy_context_fn {
+ destroy_ctx(&mut ctx);
+ }
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+
+ if let Some(callback_error) = state.error {
+ if let Some(destroy_ctx) = destroy_context_fn {
+ destroy_ctx(&mut ctx);
+ }
+ return Err(callback_error);
+ }
+
+ if let Some(destroy_ctx) = destroy_context_fn {
+ destroy_ctx(&mut ctx);
+ }
+ }
+
+ Ok(state.results)
+ }
+}
+
+struct RefinerWrapper {
+ refiner: SedonaSpatialRefiner,
+ _runtime: Arc<GpuSpatialRuntimeWrapper>,
+}
+
+impl Drop for RefinerWrapper {
+ fn drop(&mut self) {
+ let release_fn = self.refiner.release.expect("release function is None");
+ unsafe {
+ release_fn(&mut self.refiner as *mut _);
+ }
+ }
+}
+pub struct Refiner {
+ inner: RefinerWrapper,
+}
+
+impl Refiner {
+ pub fn try_new(
+ runtime: Arc<GpuSpatialRuntimeWrapper>,
+ concurrency: u32,
+ compress_bvh: bool,
+ pipeline_batches: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut refiner = SedonaSpatialRefiner {
+ clear: None,
+ init_schema: None,
+ push_build: None,
+ finish_building: None,
+ refine: None,
+ get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let config = GpuSpatialRefinerConfig {
+ runtime: runtime.runtime.get(),
+ concurrency,
+ compress_bvh,
+ pipeline_batches,
+ };
+
+ unsafe {
+ GpuSpatialRefinerCreate(&mut refiner, &config);
+ }
+
+ Ok(Self {
+ inner: RefinerWrapper {
+ refiner,
+ _runtime: runtime.clone(),
+ },
+ })
+ }
+
+ pub fn init_schema(
+ &mut self,
+ build_dt: &DataType,
+ probe_dt: &DataType,
+ ) -> Result<(), GpuSpatialError> {
+ let build_ffi = FFI_ArrowSchema::try_from(build_dt)?;
+ let probe_ffi = FFI_ArrowSchema::try_from(probe_dt)?;
+ let init_fn = self.inner.refiner.init_schema.unwrap();
+ let get_last_error = self.inner.refiner.get_last_error;
+ let refiner_ptr = &mut self.inner.refiner as *mut SedonaSpatialRefiner;
+
+ unsafe {
+ check_ffi_call(
+ || {
+ init_fn(
+ &mut self.inner.refiner,
+ &build_ffi as *const _ as *const _,
+ &probe_ffi as *const _ as *const _,
+ )
+ },
+ get_last_error,
+ refiner_ptr,
+ GpuSpatialError::Init,
+ )
+ }
+ }
+
+ pub fn push_build(&mut self, array: &ArrayRef) -> Result<(), GpuSpatialError> {
+ let (ffi_array, _) = arrow_array::ffi::to_ffi(&array.to_data())?;
+ let push_fn = self.inner.refiner.push_build.unwrap();
+ let get_last_error = self.inner.refiner.get_last_error;
+ let refiner_ptr = &mut self.inner.refiner as *mut SedonaSpatialRefiner;
+
+ unsafe {
+ check_ffi_call(
+ || push_fn(&mut self.inner.refiner, &ffi_array as *const _ as *const _),
+ get_last_error,
+ refiner_ptr,
+ GpuSpatialError::PushBuild,
+ )
+ }
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.inner.refiner.clear {
+ unsafe {
+ clear_fn(&mut self.inner.refiner as *mut _);
+ }
+ }
+ }
+
+ pub fn finish_building(&mut self) -> Result<(), GpuSpatialError> {
+ let finish_fn = self.inner.refiner.finish_building.unwrap();
+ let get_last_error = self.inner.refiner.get_last_error;
+ let refiner_ptr = &mut self.inner.refiner as *mut SedonaSpatialRefiner;
+
+ unsafe {
+ check_ffi_call(
+ || finish_fn(&mut self.inner.refiner),
+ get_last_error,
+ refiner_ptr,
+ GpuSpatialError::FinishBuild,
+ )
+ }
+ }
+
+ pub fn refine(
+ &self,
+ array: &ArrayRef,
+ predicate: GpuSpatialRelationPredicate,
+ build_indices: &mut Vec<u32>,
+ probe_indices: &mut Vec<u32>,
+ ) -> Result<(), GpuSpatialError> {
+ let (ffi_array, _) = arrow_array::ffi::to_ffi(&array.to_data())?;
+ let refine_fn = self.inner.refiner.refine.unwrap();
+ let mut new_len: u32 = 0;
+
+ unsafe {
+ check_ffi_call(
+ || {
+ refine_fn(
+ &self.inner.refiner as *const _ as *mut _,
+ &ffi_array as *const _ as *mut _,
+ predicate.as_c_uint(),
+ build_indices.as_mut_ptr(),
+ probe_indices.as_mut_ptr(),
+ build_indices.len() as u32,
+ &mut new_len,
+ )
+ },
+ self.inner.refiner.get_last_error,
+ &self.inner.refiner as *const _ as *mut _,
+ GpuSpatialError::Refine,
+ )?;
+ }
+ build_indices.truncate(new_len as usize);
+ probe_indices.truncate(new_len as usize);
+ Ok(())
+ }
+}
+
+// ----------------------------------------------------------------------
+// Helper Functions
+// ----------------------------------------------------------------------
+
+// Define the exact signature of the C error-getting function
+type ErrorFn<T> = unsafe extern "C" fn(*mut T) -> *const c_char;
+struct ProbeState {
+ results: (Vec<u32>, Vec<u32>),
+ error: Option<GpuSpatialError>,
+}
+/// Helper to handle the common pattern of calling a C function returning an int status,
+/// checking if it failed, and retrieving the error message if so.
+unsafe fn check_ffi_call<T, F, ErrMap>(
+ call_fn: F,
+ get_error_fn: Option<ErrorFn<T>>,
+ obj_ptr: *mut T,
+ err_mapper: ErrMap,
+) -> Result<(), GpuSpatialError>
+where
+ F: FnOnce() -> i32,
+ ErrMap: FnOnce(String) -> GpuSpatialError,
+{
+ if call_fn() != 0 {
+ let error_string = if let Some(get_err) = get_error_fn {
+ let err_ptr = get_err(obj_ptr);
+ if !err_ptr.is_null() {
+ CStr::from_ptr(err_ptr).to_string_lossy().into_owned()
+ } else {
+ "Unknown error (null error message)".to_string()
+ }
+ } else {
+ "Unknown error (get_last_error not available)".to_string()
+ };
+
+ return Err(err_mapper(error_string));
+ }
+ Ok(())
+}
+
+unsafe extern "C" fn probe_callback_wrapper(
+ build_indices: *const u32,
+ probe_indices: *const u32,
+ length: u32,
+ user_data: *mut c_void,
+) {
Review Comment:
Not for this PR, but we should expose the visitor interface via Rust because it is so much more memory efficient if there are any polygons whose bbox is not selective.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #586:
URL: https://github.com/apache/sedona-db/pull/586#discussion_r2796791609
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,537 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo::Rect;
+
+mod error;
#[cfg(gpu_available)]
+mod libgpuspatial;
mod libgpuspatial_glue_bindgen;
+mod options;
+mod predicate;
+// Re-export to the users
+pub use error::GpuSpatialError;
+pub use options::GpuSpatialOptions;
+pub use predicate::GpuSpatialRelationPredicate;
+
+#[cfg(gpu_available)]
+mod sys {
+ use super::libgpuspatial;
+ use super::libgpuspatial_glue_bindgen;
+ use super::*;
+
+ pub type Result<T> = std::result::Result<T, GpuSpatialError>;
+
+ use std::sync::{Arc, Mutex};
+
+ use libgpuspatial::{
+ FloatIndex2DBuilder, GpuSpatialRefinerWrapper, GpuSpatialRuntimeWrapper, SharedFloatIndex2D,
+ };
+ // To be used in the global state. We need to ensure these are Send + Sync since they will be shared across threads.
+ unsafe impl Send for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
+ unsafe impl Sync for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
+
+ // -- Global State --
+ static GLOBAL_GPUSPATIAL_RUNTIME: Mutex<Option<Arc<Mutex<GpuSpatialRuntimeWrapper>>>> =
+ Mutex::new(None);
+
+ // -- The Actual Implementation Struct --
+ pub struct SpatialImpl {
+ runtime: Option<Arc<Mutex<GpuSpatialRuntimeWrapper>>>,
+ // Index state is either building (Builder) or built (Shared)
+ index_builder: Option<FloatIndex2DBuilder>,
+ index: Option<SharedFloatIndex2D>,
+ refiner: Option<GpuSpatialRefinerWrapper>,
+ }
Review Comment:
Just discovered that Dewey has already gave a more comprehensive suggestion. https://github.com/apache/sedona-db/pull/586#pullrequestreview-3781255055
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "pwrliang (via GitHub)" <gi...@apache.org>.
pwrliang commented on code in PR #586:
URL: https://github.com/apache/sedona-db/pull/586#discussion_r2819992207
##########
c/sedona-libgpuspatial/libgpuspatial/src/gpuspatial_c.cc:
##########
@@ -194,12 +192,25 @@ struct GpuSpatialIndexFloat2DExporter {
}
static int CProbe(self_t* self, SedonaSpatialIndexContext* context, const float* buf,
- uint32_t n_rects) {
- return SafeExecute(static_cast<context_t*>(context->private_data), [&] {
+ uint32_t n_rects,
+ void (*callback)(const uint32_t* build_indices,
+ const uint32_t* probe_indices, uint32_t length,
+ void* user_data),
+ void* user_data) {
+ // Do not use SafeExecute because this method is thread-safe and we don't want to set
+ // last_error for the whole index if one thread encounters an error
+ try {
auto* rects = reinterpret_cast<const spatial_index_t::box_t*>(buf);
auto& buff = static_cast<context_t*>(context->private_data)->payload;
use_index(self).Probe(rects, n_rects, &buff.build_indices, &buff.probe_indices);
- });
+ callback(buff.build_indices.data(), buff.probe_indices.data(),
+ buff.build_indices.size(), user_data);
+ return 0;
+ } catch (const std::exception& e) { // user should call context_get_last_error
+ return EINVAL;
+ } catch (...) {
+ return EINVAL;
+ }
Review Comment:
Thanks, forgot to set the error message.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "pwrliang (via GitHub)" <gi...@apache.org>.
pwrliang commented on code in PR #586:
URL: https://github.com/apache/sedona-db/pull/586#discussion_r2796356271
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,628 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo::Rect;
+
+#[cfg(gpu_available)]
+pub mod error;
+#[cfg(gpu_available)]
+mod libgpuspatial;
#[cfg(gpu_available)]
mod libgpuspatial_glue_bindgen;
+
+pub struct GpuSpatialOptions {
+ pub cuda_use_memory_pool: bool,
+ pub cuda_memory_pool_init_percent: i32,
+ pub concurrency: u32,
+ pub device_id: i32,
+ pub compress_bvh: bool,
+ pub pipeline_batches: u32,
+}
+
+/// Spatial predicates for GPU operations
+#[repr(u32)]
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum GpuSpatialRelationPredicate {
+ Equals = 0,
+ Disjoint = 1,
+ Touches = 2,
+ Contains = 3,
+ Covers = 4,
+ Intersects = 5,
+ Within = 6,
+ CoveredBy = 7,
+}
+
+impl std::fmt::Display for GpuSpatialRelationPredicate {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ GpuSpatialRelationPredicate::Equals => write!(f, "equals"),
+ GpuSpatialRelationPredicate::Disjoint => write!(f, "disjoint"),
+ GpuSpatialRelationPredicate::Touches => write!(f, "touches"),
+ GpuSpatialRelationPredicate::Contains => write!(f, "contains"),
+ GpuSpatialRelationPredicate::Covers => write!(f, "covers"),
+ GpuSpatialRelationPredicate::Intersects => write!(f, "intersects"),
+ GpuSpatialRelationPredicate::Within => write!(f, "within"),
+ GpuSpatialRelationPredicate::CoveredBy => write!(f, "coveredby"),
+ }
+ }
+}
+
+// Re-export Error type from the sys module abstraction
+pub use sys::GpuSpatialError;
+pub type Result<T> = std::result::Result<T, GpuSpatialError>;
+
+// 1. GPU Available Implementation
+#[cfg(gpu_available)]
+mod sys {
+ use super::libgpuspatial;
+ use super::libgpuspatial_glue_bindgen;
+ use super::*;
+
+ use nvml_wrapper::Nvml;
+ use std::sync::{Arc, Mutex};
+
+ // Re-export the error from the parent module so the public API works
+ pub use super::error::GpuSpatialError;
+
+ use libgpuspatial::{
+ GpuSpatialIndexFloat2DWrapper, GpuSpatialRefinerWrapper,
+ GpuSpatialRelationPredicateWrapper, GpuSpatialRuntimeWrapper,
+ };
+ use libgpuspatial_glue_bindgen::SedonaSpatialIndexContext;
+
+ // -- Thread Safety Setup --
+ unsafe impl Send for SedonaSpatialIndexContext {}
+ unsafe impl Send for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
+ unsafe impl Sync for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
+ unsafe impl Send for libgpuspatial_glue_bindgen::SedonaFloatIndex2D {}
+ unsafe impl Send for libgpuspatial_glue_bindgen::SedonaSpatialRefiner {}
+ unsafe impl Sync for libgpuspatial_glue_bindgen::SedonaFloatIndex2D {}
+ unsafe impl Sync for libgpuspatial_glue_bindgen::SedonaSpatialRefiner {}
Review Comment:
I need two of these as the runtime is a singleton and it is declared as `static GLOBAL_GPUSPATIAL_RUNTIME: Mutex<Option<Arc<Mutex<GpuSpatialRuntimeWrapper>>>> =
Mutex::new(None);`
```rust
unsafe impl Send for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
unsafe impl Sync for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #586:
URL: https://github.com/apache/sedona-db/pull/586#discussion_r2790148191
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,628 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo::Rect;
+
+#[cfg(gpu_available)]
+pub mod error;
+#[cfg(gpu_available)]
+mod libgpuspatial;
#[cfg(gpu_available)]
mod libgpuspatial_glue_bindgen;
+
+pub struct GpuSpatialOptions {
+ pub cuda_use_memory_pool: bool,
+ pub cuda_memory_pool_init_percent: i32,
+ pub concurrency: u32,
+ pub device_id: i32,
+ pub compress_bvh: bool,
+ pub pipeline_batches: u32,
+}
+
+/// Spatial predicates for GPU operations
+#[repr(u32)]
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum GpuSpatialRelationPredicate {
+ Equals = 0,
+ Disjoint = 1,
+ Touches = 2,
+ Contains = 3,
+ Covers = 4,
+ Intersects = 5,
+ Within = 6,
+ CoveredBy = 7,
+}
+
+impl std::fmt::Display for GpuSpatialRelationPredicate {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ GpuSpatialRelationPredicate::Equals => write!(f, "equals"),
+ GpuSpatialRelationPredicate::Disjoint => write!(f, "disjoint"),
+ GpuSpatialRelationPredicate::Touches => write!(f, "touches"),
+ GpuSpatialRelationPredicate::Contains => write!(f, "contains"),
+ GpuSpatialRelationPredicate::Covers => write!(f, "covers"),
+ GpuSpatialRelationPredicate::Intersects => write!(f, "intersects"),
+ GpuSpatialRelationPredicate::Within => write!(f, "within"),
+ GpuSpatialRelationPredicate::CoveredBy => write!(f, "coveredby"),
+ }
+ }
+}
Review Comment:
This is duplicated in libgpuspatial.rs (perhaps you can move this and `GpuSpatialOptions` to a new module `options.rs`).
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,628 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo::Rect;
+
+#[cfg(gpu_available)]
+pub mod error;
+#[cfg(gpu_available)]
+mod libgpuspatial;
#[cfg(gpu_available)]
mod libgpuspatial_glue_bindgen;
+
+pub struct GpuSpatialOptions {
+ pub cuda_use_memory_pool: bool,
+ pub cuda_memory_pool_init_percent: i32,
+ pub concurrency: u32,
+ pub device_id: i32,
+ pub compress_bvh: bool,
+ pub pipeline_batches: u32,
+}
+
+/// Spatial predicates for GPU operations
+#[repr(u32)]
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum GpuSpatialRelationPredicate {
+ Equals = 0,
+ Disjoint = 1,
+ Touches = 2,
+ Contains = 3,
+ Covers = 4,
+ Intersects = 5,
+ Within = 6,
+ CoveredBy = 7,
+}
+
+impl std::fmt::Display for GpuSpatialRelationPredicate {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ GpuSpatialRelationPredicate::Equals => write!(f, "equals"),
+ GpuSpatialRelationPredicate::Disjoint => write!(f, "disjoint"),
+ GpuSpatialRelationPredicate::Touches => write!(f, "touches"),
+ GpuSpatialRelationPredicate::Contains => write!(f, "contains"),
+ GpuSpatialRelationPredicate::Covers => write!(f, "covers"),
+ GpuSpatialRelationPredicate::Intersects => write!(f, "intersects"),
+ GpuSpatialRelationPredicate::Within => write!(f, "within"),
+ GpuSpatialRelationPredicate::CoveredBy => write!(f, "coveredby"),
+ }
+ }
+}
+
+// Re-export Error type from the sys module abstraction
+pub use sys::GpuSpatialError;
+pub type Result<T> = std::result::Result<T, GpuSpatialError>;
+
+// 1. GPU Available Implementation
+#[cfg(gpu_available)]
+mod sys {
+ use super::libgpuspatial;
+ use super::libgpuspatial_glue_bindgen;
+ use super::*;
+
+ use nvml_wrapper::Nvml;
+ use std::sync::{Arc, Mutex};
+
+ // Re-export the error from the parent module so the public API works
+ pub use super::error::GpuSpatialError;
Review Comment:
I feel like I'm missing something with the errors but it's probably best to have them not depend on the gpu availability if at all possible.
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,628 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo::Rect;
+
+#[cfg(gpu_available)]
+pub mod error;
+#[cfg(gpu_available)]
+mod libgpuspatial;
#[cfg(gpu_available)]
mod libgpuspatial_glue_bindgen;
+
+pub struct GpuSpatialOptions {
+ pub cuda_use_memory_pool: bool,
+ pub cuda_memory_pool_init_percent: i32,
+ pub concurrency: u32,
+ pub device_id: i32,
+ pub compress_bvh: bool,
+ pub pipeline_batches: u32,
+}
+
+/// Spatial predicates for GPU operations
+#[repr(u32)]
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum GpuSpatialRelationPredicate {
+ Equals = 0,
+ Disjoint = 1,
+ Touches = 2,
+ Contains = 3,
+ Covers = 4,
+ Intersects = 5,
+ Within = 6,
+ CoveredBy = 7,
+}
+
+impl std::fmt::Display for GpuSpatialRelationPredicate {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ GpuSpatialRelationPredicate::Equals => write!(f, "equals"),
+ GpuSpatialRelationPredicate::Disjoint => write!(f, "disjoint"),
+ GpuSpatialRelationPredicate::Touches => write!(f, "touches"),
+ GpuSpatialRelationPredicate::Contains => write!(f, "contains"),
+ GpuSpatialRelationPredicate::Covers => write!(f, "covers"),
+ GpuSpatialRelationPredicate::Intersects => write!(f, "intersects"),
+ GpuSpatialRelationPredicate::Within => write!(f, "within"),
+ GpuSpatialRelationPredicate::CoveredBy => write!(f, "coveredby"),
+ }
+ }
+}
+
+// Re-export Error type from the sys module abstraction
+pub use sys::GpuSpatialError;
+pub type Result<T> = std::result::Result<T, GpuSpatialError>;
Review Comment:
```suggestion
```
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,628 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo::Rect;
+
+#[cfg(gpu_available)]
+pub mod error;
Review Comment:
```suggestion
pub mod error;
```
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,628 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo::Rect;
+
+#[cfg(gpu_available)]
+pub mod error;
+#[cfg(gpu_available)]
+mod libgpuspatial;
#[cfg(gpu_available)]
mod libgpuspatial_glue_bindgen;
+
+pub struct GpuSpatialOptions {
+ pub cuda_use_memory_pool: bool,
+ pub cuda_memory_pool_init_percent: i32,
+ pub concurrency: u32,
+ pub device_id: i32,
+ pub compress_bvh: bool,
+ pub pipeline_batches: u32,
+}
+
+/// Spatial predicates for GPU operations
+#[repr(u32)]
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum GpuSpatialRelationPredicate {
+ Equals = 0,
+ Disjoint = 1,
+ Touches = 2,
+ Contains = 3,
+ Covers = 4,
+ Intersects = 5,
+ Within = 6,
+ CoveredBy = 7,
+}
+
+impl std::fmt::Display for GpuSpatialRelationPredicate {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ GpuSpatialRelationPredicate::Equals => write!(f, "equals"),
+ GpuSpatialRelationPredicate::Disjoint => write!(f, "disjoint"),
+ GpuSpatialRelationPredicate::Touches => write!(f, "touches"),
+ GpuSpatialRelationPredicate::Contains => write!(f, "contains"),
+ GpuSpatialRelationPredicate::Covers => write!(f, "covers"),
+ GpuSpatialRelationPredicate::Intersects => write!(f, "intersects"),
+ GpuSpatialRelationPredicate::Within => write!(f, "within"),
+ GpuSpatialRelationPredicate::CoveredBy => write!(f, "coveredby"),
+ }
+ }
+}
+
+// Re-export Error type from the sys module abstraction
+pub use sys::GpuSpatialError;
+pub type Result<T> = std::result::Result<T, GpuSpatialError>;
+
+// 1. GPU Available Implementation
+#[cfg(gpu_available)]
+mod sys {
+ use super::libgpuspatial;
+ use super::libgpuspatial_glue_bindgen;
+ use super::*;
+
+ use nvml_wrapper::Nvml;
+ use std::sync::{Arc, Mutex};
+
+ // Re-export the error from the parent module so the public API works
+ pub use super::error::GpuSpatialError;
+
+ use libgpuspatial::{
+ GpuSpatialIndexFloat2DWrapper, GpuSpatialRefinerWrapper,
+ GpuSpatialRelationPredicateWrapper, GpuSpatialRuntimeWrapper,
+ };
+ use libgpuspatial_glue_bindgen::SedonaSpatialIndexContext;
+
+ // -- Thread Safety Setup --
+ unsafe impl Send for SedonaSpatialIndexContext {}
+ unsafe impl Send for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
+ unsafe impl Sync for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
+ unsafe impl Send for libgpuspatial_glue_bindgen::SedonaFloatIndex2D {}
+ unsafe impl Send for libgpuspatial_glue_bindgen::SedonaSpatialRefiner {}
+ unsafe impl Sync for libgpuspatial_glue_bindgen::SedonaFloatIndex2D {}
+ unsafe impl Sync for libgpuspatial_glue_bindgen::SedonaSpatialRefiner {}
Review Comment:
I don't think these are correct `unsafe impl`s. The wrappers are the appropriate place to do the `unsafe impl Send`, and I am not sure you have implemented any that are `unsafe impl Sync`. It is probably better to use a `Mutex` or `RwLock` to ensure something is `impl Sync`.
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,523 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+use crate::libgpuspatial_glue_bindgen::*;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::{c_char, c_uint};
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Helper Functions
+// ----------------------------------------------------------------------
+/// Helper to handle the common pattern of calling a C function returning an int status,
+/// checking if it failed, and retrieving the error message if so.
+///
+/// T: The type of the object (Runtime, Index, Refiner) being operated on.
+unsafe fn check_ffi_call<T, F, ErrMap>(
Review Comment:
Rust will let you put these at the bottom of the file before they are used (which usually helps the readability of the file).
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,523 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+use crate::libgpuspatial_glue_bindgen::*;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::{c_char, c_uint};
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Helper Functions
+// ----------------------------------------------------------------------
+/// Helper to handle the common pattern of calling a C function returning an int status,
+/// checking if it failed, and retrieving the error message if so.
+///
+/// T: The type of the object (Runtime, Index, Refiner) being operated on.
+unsafe fn check_ffi_call<T, F, ErrMap>(
+ call_fn: F,
+ get_error_fn: Option<unsafe extern "C" fn(*mut T) -> *const c_char>,
+ obj_ptr: *mut T,
+ err_mapper: ErrMap,
+) -> Result<(), GpuSpatialError>
+where
+ F: FnOnce() -> i32,
+ ErrMap: FnOnce(String) -> GpuSpatialError,
+{
+ if call_fn() != 0 {
+ let error_string = if let Some(get_err) = get_error_fn {
+ let err_ptr = get_err(obj_ptr);
+ if !err_ptr.is_null() {
+ CStr::from_ptr(err_ptr).to_string_lossy().into_owned()
+ } else {
+ "Unknown error (null error message)".to_string()
+ }
+ } else {
+ "Unknown error (get_last_error not available)".to_string()
+ };
+
+ log::error!("GpuSpatial FFI Error: {}", error_string);
+ return Err(err_mapper(error_string));
+ }
+ Ok(())
+}
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ /// Initializes the GpuSpatialRuntime.
+ /// This function should only be called once per engine instance.
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ check_ffi_call(
+ || init_fn(&runtime as *const _ as *mut _, &mut config),
+ runtime.get_last_error,
+ &runtime as *const _ as *mut _,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialIndexFloat2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+impl GpuSpatialIndexFloat2DWrapper {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ // Can't use check_ffi_call helper here easily because 'index' isn't fully initialized/wrapped yet
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(GpuSpatialIndexFloat2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.index.clear {
+ unsafe {
+ clear_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
+ &mut self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.index.push_build {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ Ok(())
+ }
+
+ pub fn finish_building(&mut self) -> Result<(), GpuSpatialError> {
+ if let Some(finish_building_fn) = self.index.finish_building {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn create_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(create_context_fn) = self.index.create_context {
+ unsafe {
+ create_context_fn(ctx as *mut _);
+ }
+ }
+ }
Review Comment:
It would be nice for these contexts to get managed as part of the wrapper (e.g., how these structs keep a reference to an `Arc` of the runtime engine to ensure it is kept alive while required).
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,523 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+use crate::libgpuspatial_glue_bindgen::*;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::{c_char, c_uint};
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Helper Functions
+// ----------------------------------------------------------------------
+/// Helper to handle the common pattern of calling a C function returning an int status,
+/// checking if it failed, and retrieving the error message if so.
+///
+/// T: The type of the object (Runtime, Index, Refiner) being operated on.
+unsafe fn check_ffi_call<T, F, ErrMap>(
+ call_fn: F,
+ get_error_fn: Option<unsafe extern "C" fn(*mut T) -> *const c_char>,
+ obj_ptr: *mut T,
+ err_mapper: ErrMap,
+) -> Result<(), GpuSpatialError>
+where
+ F: FnOnce() -> i32,
+ ErrMap: FnOnce(String) -> GpuSpatialError,
+{
+ if call_fn() != 0 {
+ let error_string = if let Some(get_err) = get_error_fn {
+ let err_ptr = get_err(obj_ptr);
+ if !err_ptr.is_null() {
+ CStr::from_ptr(err_ptr).to_string_lossy().into_owned()
+ } else {
+ "Unknown error (null error message)".to_string()
+ }
+ } else {
+ "Unknown error (get_last_error not available)".to_string()
+ };
+
+ log::error!("GpuSpatial FFI Error: {}", error_string);
+ return Err(err_mapper(error_string));
+ }
+ Ok(())
+}
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ /// Initializes the GpuSpatialRuntime.
+ /// This function should only be called once per engine instance.
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ check_ffi_call(
+ || init_fn(&runtime as *const _ as *mut _, &mut config),
+ runtime.get_last_error,
+ &runtime as *const _ as *mut _,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialIndexFloat2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+impl GpuSpatialIndexFloat2DWrapper {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ // Can't use check_ffi_call helper here easily because 'index' isn't fully initialized/wrapped yet
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(GpuSpatialIndexFloat2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.index.clear {
+ unsafe {
+ clear_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
+ &mut self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.index.push_build {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ Ok(())
+ }
+
+ pub fn finish_building(&mut self) -> Result<(), GpuSpatialError> {
+ if let Some(finish_building_fn) = self.index.finish_building {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn create_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(create_context_fn) = self.index.create_context {
+ unsafe {
+ create_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub fn destroy_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(destroy_context_fn) = self.index.destroy_context {
+ unsafe {
+ destroy_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn probe(
+ &self,
+ ctx: &mut SedonaSpatialIndexContext,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
Review Comment:
I think this comment got lost in the last review, but unless there's a good reason to avoid this pattern, this should use the same signature as `refine()` (i.e., the caller provides the buffers into which the prove callback can populate the index pairs).
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,523 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+use crate::libgpuspatial_glue_bindgen::*;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::{c_char, c_uint};
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Helper Functions
+// ----------------------------------------------------------------------
+/// Helper to handle the common pattern of calling a C function returning an int status,
+/// checking if it failed, and retrieving the error message if so.
+///
+/// T: The type of the object (Runtime, Index, Refiner) being operated on.
+unsafe fn check_ffi_call<T, F, ErrMap>(
+ call_fn: F,
+ get_error_fn: Option<unsafe extern "C" fn(*mut T) -> *const c_char>,
+ obj_ptr: *mut T,
+ err_mapper: ErrMap,
+) -> Result<(), GpuSpatialError>
+where
+ F: FnOnce() -> i32,
+ ErrMap: FnOnce(String) -> GpuSpatialError,
+{
+ if call_fn() != 0 {
+ let error_string = if let Some(get_err) = get_error_fn {
+ let err_ptr = get_err(obj_ptr);
+ if !err_ptr.is_null() {
+ CStr::from_ptr(err_ptr).to_string_lossy().into_owned()
+ } else {
+ "Unknown error (null error message)".to_string()
+ }
+ } else {
+ "Unknown error (get_last_error not available)".to_string()
+ };
+
+ log::error!("GpuSpatial FFI Error: {}", error_string);
+ return Err(err_mapper(error_string));
+ }
+ Ok(())
+}
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ /// Initializes the GpuSpatialRuntime.
+ /// This function should only be called once per engine instance.
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ check_ffi_call(
+ || init_fn(&runtime as *const _ as *mut _, &mut config),
+ runtime.get_last_error,
+ &runtime as *const _ as *mut _,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialIndexFloat2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+impl GpuSpatialIndexFloat2DWrapper {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ // Can't use check_ffi_call helper here easily because 'index' isn't fully initialized/wrapped yet
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(GpuSpatialIndexFloat2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.index.clear {
+ unsafe {
+ clear_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
+ &mut self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.index.push_build {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ Ok(())
+ }
+
+ pub fn finish_building(&mut self) -> Result<(), GpuSpatialError> {
+ if let Some(finish_building_fn) = self.index.finish_building {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn create_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(create_context_fn) = self.index.create_context {
+ unsafe {
+ create_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub fn destroy_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(destroy_context_fn) = self.index.destroy_context {
+ unsafe {
+ destroy_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn probe(
+ &self,
+ ctx: &mut SedonaSpatialIndexContext,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(probe_fn) = self.index.probe {
+ // Note: probe uses context_get_last_error, not the index one
+ if probe_fn(
+ &self.index as *const _ as *mut _,
+ ctx as *mut _,
+ buf,
+ n_rects,
+ ) != 0
+ {
+ let error_string = if let Some(get_ctx_err) = self.index.context_get_last_error {
+ CStr::from_ptr(get_ctx_err(ctx))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error".to_string()
+ };
+ log::error!("DEBUG FFI: probe failed: {}", error_string);
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+ }
+ Ok(())
+ }
+
+ fn get_indices_buffer_helper(
+ &self,
+ ctx: &mut SedonaSpatialIndexContext,
+ func: Option<unsafe extern "C" fn(*mut SedonaSpatialIndexContext, *mut *mut u32, *mut u32)>,
+ ) -> &[u32] {
+ if let Some(f) = func {
+ let mut ptr: *mut u32 = std::ptr::null_mut();
+ let mut len: u32 = 0;
+ unsafe {
+ f(ctx, &mut ptr, &mut len);
+ if len > 0 && !ptr.is_null() {
+ return std::slice::from_raw_parts(ptr, len as usize);
+ }
+ }
+ }
+ &[]
+ }
+
+ pub fn get_build_indices_buffer(&self, ctx: &mut SedonaSpatialIndexContext) -> &[u32] {
+ self.get_indices_buffer_helper(ctx, self.index.get_build_indices_buffer)
+ }
+
+ pub fn get_probe_indices_buffer(&self, ctx: &mut SedonaSpatialIndexContext) -> &[u32] {
+ self.get_indices_buffer_helper(ctx, self.index.get_probe_indices_buffer)
+ }
+}
+
+impl Drop for GpuSpatialIndexFloat2DWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.index.release {
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Predicate Wrapper
+// ----------------------------------------------------------------------
+
+#[repr(u32)]
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum GpuSpatialRelationPredicateWrapper {
+ Equals = 0,
+ Disjoint = 1,
+ Touches = 2,
+ Contains = 3,
+ Covers = 4,
+ Intersects = 5,
+ Within = 6,
+ CoveredBy = 7,
+}
+
+impl TryFrom<c_uint> for GpuSpatialRelationPredicateWrapper {
+ type Error = &'static str;
+
+ fn try_from(v: c_uint) -> Result<Self, Self::Error> {
+ match v {
+ 0 => Ok(Self::Equals),
+ 1 => Ok(Self::Disjoint),
+ 2 => Ok(Self::Touches),
+ 3 => Ok(Self::Contains),
+ 4 => Ok(Self::Covers),
+ 5 => Ok(Self::Intersects),
+ 6 => Ok(Self::Within),
+ 7 => Ok(Self::CoveredBy),
+ _ => Err("Invalid GpuSpatialPredicate value"),
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Refiner Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRefinerWrapper {
+ refiner: SedonaSpatialRefiner,
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+impl GpuSpatialRefinerWrapper {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ compress_bvh: bool,
+ pipeline_batches: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut refiner = SedonaSpatialRefiner {
+ clear: None,
+ init_schema: None,
+ push_build: None,
+ finish_building: None,
+ refine: None,
+ get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialRefinerConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ compress_bvh,
+ pipeline_batches,
+ };
+
+ unsafe {
+ if GpuSpatialRefinerCreate(&mut refiner, &config) != 0 {
+ let msg = if let Some(get_err) = refiner.get_last_error {
+ CStr::from_ptr(get_err(&refiner as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Refiner Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(GpuSpatialRefinerWrapper {
+ refiner,
+ _runtime: runtime.clone(),
+ })
+ }
+
+ pub fn init_schema(
+ &self,
+ build_data_type: &DataType,
+ probe_data_type: &DataType,
+ ) -> Result<(), GpuSpatialError> {
+ let ffi_build_schema = FFI_ArrowSchema::try_from(build_data_type)?;
+ let ffi_probe_schema = FFI_ArrowSchema::try_from(probe_data_type)?;
+
+ if let Some(init_schema_fn) = self.refiner.init_schema {
+ // Use pointer casting instead of transmute for safer FFI
+ let ffi_build_ptr = &ffi_build_schema as *const _ as *const ArrowSchema;
+ let ffi_probe_ptr = &ffi_probe_schema as *const _ as *const ArrowSchema;
Review Comment:
These functions should take `&mut self`, because they must not be called concurrently (i.e., the caller must own a single mutable reference). In general you should not have to do `as *const _ as *mut _` (use `&mut self.refiner as *mut _`).
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,628 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo::Rect;
+
+#[cfg(gpu_available)]
+pub mod error;
+#[cfg(gpu_available)]
+mod libgpuspatial;
#[cfg(gpu_available)]
mod libgpuspatial_glue_bindgen;
+
+pub struct GpuSpatialOptions {
+ pub cuda_use_memory_pool: bool,
+ pub cuda_memory_pool_init_percent: i32,
+ pub concurrency: u32,
+ pub device_id: i32,
+ pub compress_bvh: bool,
+ pub pipeline_batches: u32,
+}
+
+/// Spatial predicates for GPU operations
+#[repr(u32)]
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum GpuSpatialRelationPredicate {
+ Equals = 0,
+ Disjoint = 1,
+ Touches = 2,
+ Contains = 3,
+ Covers = 4,
+ Intersects = 5,
+ Within = 6,
+ CoveredBy = 7,
+}
+
+impl std::fmt::Display for GpuSpatialRelationPredicate {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ GpuSpatialRelationPredicate::Equals => write!(f, "equals"),
+ GpuSpatialRelationPredicate::Disjoint => write!(f, "disjoint"),
+ GpuSpatialRelationPredicate::Touches => write!(f, "touches"),
+ GpuSpatialRelationPredicate::Contains => write!(f, "contains"),
+ GpuSpatialRelationPredicate::Covers => write!(f, "covers"),
+ GpuSpatialRelationPredicate::Intersects => write!(f, "intersects"),
+ GpuSpatialRelationPredicate::Within => write!(f, "within"),
+ GpuSpatialRelationPredicate::CoveredBy => write!(f, "coveredby"),
+ }
+ }
+}
+
+// Re-export Error type from the sys module abstraction
+pub use sys::GpuSpatialError;
+pub type Result<T> = std::result::Result<T, GpuSpatialError>;
+
+// 1. GPU Available Implementation
+#[cfg(gpu_available)]
+mod sys {
+ use super::libgpuspatial;
+ use super::libgpuspatial_glue_bindgen;
+ use super::*;
+
+ use nvml_wrapper::Nvml;
+ use std::sync::{Arc, Mutex};
+
+ // Re-export the error from the parent module so the public API works
+ pub use super::error::GpuSpatialError;
+
+ use libgpuspatial::{
+ GpuSpatialIndexFloat2DWrapper, GpuSpatialRefinerWrapper,
+ GpuSpatialRelationPredicateWrapper, GpuSpatialRuntimeWrapper,
+ };
+ use libgpuspatial_glue_bindgen::SedonaSpatialIndexContext;
+
+ // -- Thread Safety Setup --
+ unsafe impl Send for SedonaSpatialIndexContext {}
+ unsafe impl Send for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
+ unsafe impl Sync for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
+ unsafe impl Send for libgpuspatial_glue_bindgen::SedonaFloatIndex2D {}
+ unsafe impl Send for libgpuspatial_glue_bindgen::SedonaSpatialRefiner {}
+ unsafe impl Sync for libgpuspatial_glue_bindgen::SedonaFloatIndex2D {}
+ unsafe impl Sync for libgpuspatial_glue_bindgen::SedonaSpatialRefiner {}
+
+ // -- Global State --
+ static GLOBAL_GPUSPATIAL_RUNTIME: Mutex<Option<Arc<Mutex<GpuSpatialRuntimeWrapper>>>> =
+ Mutex::new(None);
+
+ // -- Conversion Trait --
+ impl From<GpuSpatialRelationPredicate> for GpuSpatialRelationPredicateWrapper {
+ fn from(pred: GpuSpatialRelationPredicate) -> Self {
+ match pred {
+ GpuSpatialRelationPredicate::Equals => GpuSpatialRelationPredicateWrapper::Equals,
+ GpuSpatialRelationPredicate::Disjoint => {
+ GpuSpatialRelationPredicateWrapper::Disjoint
+ }
+ GpuSpatialRelationPredicate::Touches => GpuSpatialRelationPredicateWrapper::Touches,
+ GpuSpatialRelationPredicate::Contains => {
+ GpuSpatialRelationPredicateWrapper::Contains
+ }
+ GpuSpatialRelationPredicate::Covers => GpuSpatialRelationPredicateWrapper::Covers,
+ GpuSpatialRelationPredicate::Intersects => {
+ GpuSpatialRelationPredicateWrapper::Intersects
+ }
+ GpuSpatialRelationPredicate::Within => GpuSpatialRelationPredicateWrapper::Within,
+ GpuSpatialRelationPredicate::CoveredBy => {
+ GpuSpatialRelationPredicateWrapper::CoveredBy
+ }
+ }
+ }
+ }
+
+ // -- The Actual Implementation Struct --
+ pub struct SpatialImpl {
+ runtime: Option<Arc<Mutex<GpuSpatialRuntimeWrapper>>>,
+ index: Option<GpuSpatialIndexFloat2DWrapper>,
+ refiner: Option<GpuSpatialRefinerWrapper>,
+ }
+
+ impl SpatialImpl {
+ pub fn new() -> Result<Self> {
+ Ok(Self {
+ runtime: None,
+ index: None,
+ refiner: None,
+ })
+ }
+
+ pub fn init(&mut self, options: GpuSpatialOptions) -> Result<()> {
+ let mut global_runtime_guard = GLOBAL_GPUSPATIAL_RUNTIME.lock().unwrap();
+
+ if global_runtime_guard.is_none() {
+ let out_path = std::path::PathBuf::from(env!("OUT_DIR"));
+ let ptx_root = out_path.join("share/gpuspatial/shaders");
+ let ptx_root_str = ptx_root
+ .to_str()
+ .ok_or_else(|| GpuSpatialError::Init("Invalid PTX path".to_string()))?;
+
+ let runtime = GpuSpatialRuntimeWrapper::try_new(
+ options.device_id,
+ ptx_root_str,
+ options.cuda_use_memory_pool,
+ options.cuda_memory_pool_init_percent,
+ )?;
+ *global_runtime_guard = Some(Arc::new(Mutex::new(runtime)));
+ }
+
+ let runtime_ref = global_runtime_guard.as_ref().unwrap().clone();
+ self.runtime = Some(runtime_ref);
+
+ // FIX: Clone the Arc here
+ let index = GpuSpatialIndexFloat2DWrapper::try_new(
+ self.runtime.as_ref().unwrap().clone(),
+ options.concurrency,
+ )?;
+ self.index = Some(index);
+
+ // FIX: Clone the Arc here as well
+ let refiner = GpuSpatialRefinerWrapper::try_new(
+ self.runtime.as_ref().unwrap().clone(),
+ options.concurrency,
+ options.compress_bvh,
+ options.pipeline_batches,
+ )?;
+ self.refiner = Some(refiner);
+
+ Ok(())
+ }
+
+ pub fn index_clear(&mut self) -> Result<()> {
+ let index = self
+ .index
+ .as_mut()
+ .ok_or_else(|| GpuSpatialError::Init("GPU index is not available".into()))?;
+ index.clear();
+ Ok(())
+ }
+
+ pub fn index_push_build(&mut self, rects: &[Rect<f32>]) -> Result<()> {
+ let index = self
+ .index
+ .as_mut()
+ .ok_or_else(|| GpuSpatialError::Init("GPU index not available".into()))?;
+ unsafe { index.push_build(rects.as_ptr() as *const f32, rects.len() as u32) }
+ }
+
+ pub fn index_finish_building(&mut self) -> Result<()> {
+ self.index
+ .as_mut()
+ .ok_or_else(|| GpuSpatialError::Init("GPU index not available".into()))?
+ .finish_building()
+ }
+
+ pub fn probe(&self, rects: &[Rect<f32>]) -> Result<(Vec<u32>, Vec<u32>)> {
+ let index = self
+ .index
+ .as_ref()
+ .ok_or_else(|| GpuSpatialError::Init("GPU index not available".into()))?;
+
+ let mut ctx = SedonaSpatialIndexContext {
+ private_data: std::ptr::null_mut(),
+ };
+ index.create_context(&mut ctx);
+
+ let result = (|| -> Result<(Vec<u32>, Vec<u32>)> {
+ unsafe {
+ index.probe(&mut ctx, rects.as_ptr() as *const f32, rects.len() as u32)?;
+ }
+ let build_indices = index.get_build_indices_buffer(&mut ctx).to_vec();
+ let probe_indices = index.get_probe_indices_buffer(&mut ctx).to_vec();
+ Ok((build_indices, probe_indices))
+ })();
+
+ index.destroy_context(&mut ctx);
+ result
+ }
+
+ pub fn refiner_clear(&mut self) -> Result<()> {
+ let refiner = self
+ .refiner
+ .as_mut()
+ .ok_or_else(|| GpuSpatialError::Init("GPU refiner is not available".into()))?;
+ refiner.clear();
+ Ok(())
+ }
+
+ pub fn refiner_init_schema(
+ &mut self,
+ build_data_type: &DataType,
+ probe_data_type: &DataType,
+ ) -> Result<()> {
+ let refiner = self
+ .refiner
+ .as_mut()
+ .ok_or_else(|| GpuSpatialError::Init("GPU refiner not available".into()))?;
+ refiner.init_schema(build_data_type, probe_data_type)
+ }
+
+ pub fn refiner_push_build(&mut self, array: &arrow_array::ArrayRef) -> Result<()> {
+ let refiner = self
+ .refiner
+ .as_ref()
+ .ok_or_else(|| GpuSpatialError::Init("GPU refiner not available".into()))?;
+ refiner.push_build(array)
+ }
+
+ pub fn refiner_finish_building(&mut self) -> Result<()> {
+ let refiner = self
+ .refiner
+ .as_mut()
+ .ok_or_else(|| GpuSpatialError::Init("GPU refiner not available".into()))?;
+ refiner.finish_building()
+ }
+
+ pub fn refine(
+ &self,
+ probe_array: &arrow_array::ArrayRef,
+ predicate: GpuSpatialRelationPredicate,
+ build_indices: &mut Vec<u32>,
+ probe_indices: &mut Vec<u32>,
+ ) -> Result<()> {
+ let refiner = self
+ .refiner
+ .as_ref()
+ .ok_or_else(|| GpuSpatialError::Init("GPU refiner not available".into()))?;
+
+ refiner.refine(
+ probe_array,
+ GpuSpatialRelationPredicateWrapper::from(predicate),
+ build_indices,
+ probe_indices,
+ )
+ }
+ }
+
+ pub fn is_gpu_available() -> bool {
+ match Nvml::init() {
+ Ok(instance) => instance.device_count().map(|c| c > 0).unwrap_or(false),
+ Err(_) => false,
+ }
+ }
Review Comment:
If at all possible I would like to avoid `Nvml` here. I think having this be a function that is exposed from the C API will simplify the build/link portion of this (i.e., I'd like to constrain the place where we get CUDA linking right in exactly one place that is not Rust!)
##########
c/sedona-libgpuspatial/src/error.rs:
##########
@@ -0,0 +1,60 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+use arrow_schema::ArrowError;
+use std::fmt;
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum GpuSpatialError {
+ Arrow(ArrowError),
+ Init(String),
+ PushBuild(String),
+ FinishBuild(String),
+ Probe(String),
+ Refine(String),
Review Comment:
```suggestion
Refine(String),
GpuNotAvailable,
```
(Maybe this can solve the multiple errors issue?)
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,523 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+use crate::libgpuspatial_glue_bindgen::*;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::{c_char, c_uint};
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Helper Functions
+// ----------------------------------------------------------------------
+/// Helper to handle the common pattern of calling a C function returning an int status,
+/// checking if it failed, and retrieving the error message if so.
+///
+/// T: The type of the object (Runtime, Index, Refiner) being operated on.
+unsafe fn check_ffi_call<T, F, ErrMap>(
+ call_fn: F,
+ get_error_fn: Option<unsafe extern "C" fn(*mut T) -> *const c_char>,
+ obj_ptr: *mut T,
+ err_mapper: ErrMap,
+) -> Result<(), GpuSpatialError>
+where
+ F: FnOnce() -> i32,
+ ErrMap: FnOnce(String) -> GpuSpatialError,
+{
+ if call_fn() != 0 {
+ let error_string = if let Some(get_err) = get_error_fn {
+ let err_ptr = get_err(obj_ptr);
+ if !err_ptr.is_null() {
+ CStr::from_ptr(err_ptr).to_string_lossy().into_owned()
+ } else {
+ "Unknown error (null error message)".to_string()
+ }
+ } else {
+ "Unknown error (get_last_error not available)".to_string()
+ };
+
+ log::error!("GpuSpatial FFI Error: {}", error_string);
+ return Err(err_mapper(error_string));
+ }
+ Ok(())
+}
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ /// Initializes the GpuSpatialRuntime.
+ /// This function should only be called once per engine instance.
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ check_ffi_call(
+ || init_fn(&runtime as *const _ as *mut _, &mut config),
+ runtime.get_last_error,
+ &runtime as *const _ as *mut _,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialIndexFloat2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+impl GpuSpatialIndexFloat2DWrapper {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ // Can't use check_ffi_call helper here easily because 'index' isn't fully initialized/wrapped yet
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(GpuSpatialIndexFloat2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.index.clear {
+ unsafe {
+ clear_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
+ &mut self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.index.push_build {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ Ok(())
+ }
+
+ pub fn finish_building(&mut self) -> Result<(), GpuSpatialError> {
+ if let Some(finish_building_fn) = self.index.finish_building {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn create_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(create_context_fn) = self.index.create_context {
+ unsafe {
+ create_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub fn destroy_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(destroy_context_fn) = self.index.destroy_context {
+ unsafe {
+ destroy_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn probe(
+ &self,
+ ctx: &mut SedonaSpatialIndexContext,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(probe_fn) = self.index.probe {
+ // Note: probe uses context_get_last_error, not the index one
+ if probe_fn(
+ &self.index as *const _ as *mut _,
+ ctx as *mut _,
+ buf,
+ n_rects,
+ ) != 0
+ {
+ let error_string = if let Some(get_ctx_err) = self.index.context_get_last_error {
+ CStr::from_ptr(get_ctx_err(ctx))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error".to_string()
+ };
+ log::error!("DEBUG FFI: probe failed: {}", error_string);
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+ }
+ Ok(())
+ }
+
+ fn get_indices_buffer_helper(
+ &self,
+ ctx: &mut SedonaSpatialIndexContext,
+ func: Option<unsafe extern "C" fn(*mut SedonaSpatialIndexContext, *mut *mut u32, *mut u32)>,
+ ) -> &[u32] {
+ if let Some(f) = func {
+ let mut ptr: *mut u32 = std::ptr::null_mut();
+ let mut len: u32 = 0;
+ unsafe {
+ f(ctx, &mut ptr, &mut len);
+ if len > 0 && !ptr.is_null() {
+ return std::slice::from_raw_parts(ptr, len as usize);
+ }
+ }
+ }
+ &[]
+ }
+
+ pub fn get_build_indices_buffer(&self, ctx: &mut SedonaSpatialIndexContext) -> &[u32] {
+ self.get_indices_buffer_helper(ctx, self.index.get_build_indices_buffer)
+ }
+
+ pub fn get_probe_indices_buffer(&self, ctx: &mut SedonaSpatialIndexContext) -> &[u32] {
+ self.get_indices_buffer_helper(ctx, self.index.get_probe_indices_buffer)
+ }
+}
+
+impl Drop for GpuSpatialIndexFloat2DWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.index.release {
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Predicate Wrapper
+// ----------------------------------------------------------------------
+
+#[repr(u32)]
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum GpuSpatialRelationPredicateWrapper {
+ Equals = 0,
+ Disjoint = 1,
+ Touches = 2,
+ Contains = 3,
+ Covers = 4,
+ Intersects = 5,
+ Within = 6,
+ CoveredBy = 7,
+}
+
+impl TryFrom<c_uint> for GpuSpatialRelationPredicateWrapper {
+ type Error = &'static str;
+
+ fn try_from(v: c_uint) -> Result<Self, Self::Error> {
+ match v {
+ 0 => Ok(Self::Equals),
+ 1 => Ok(Self::Disjoint),
+ 2 => Ok(Self::Touches),
+ 3 => Ok(Self::Contains),
+ 4 => Ok(Self::Covers),
+ 5 => Ok(Self::Intersects),
+ 6 => Ok(Self::Within),
+ 7 => Ok(Self::CoveredBy),
+ _ => Err("Invalid GpuSpatialPredicate value"),
+ }
+ }
+}
Review Comment:
The details of how this is passed to C shouldn't affect the public API of this crate or be used to construct these...if you need an integer from this you can use a private method in `impl GpuSpatialRelationPredicateWrapper`.
```suggestion
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum GpuSpatialRelationPredicateWrapper {
Equals,
Disjoint,
Touches,
Contains,
Covers,
Intersects,
Within,
CoveredBy,
}
```
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,523 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+use crate::libgpuspatial_glue_bindgen::*;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::{c_char, c_uint};
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Helper Functions
+// ----------------------------------------------------------------------
+/// Helper to handle the common pattern of calling a C function returning an int status,
+/// checking if it failed, and retrieving the error message if so.
+///
+/// T: The type of the object (Runtime, Index, Refiner) being operated on.
+unsafe fn check_ffi_call<T, F, ErrMap>(
+ call_fn: F,
+ get_error_fn: Option<unsafe extern "C" fn(*mut T) -> *const c_char>,
+ obj_ptr: *mut T,
+ err_mapper: ErrMap,
+) -> Result<(), GpuSpatialError>
+where
+ F: FnOnce() -> i32,
+ ErrMap: FnOnce(String) -> GpuSpatialError,
+{
+ if call_fn() != 0 {
+ let error_string = if let Some(get_err) = get_error_fn {
+ let err_ptr = get_err(obj_ptr);
+ if !err_ptr.is_null() {
+ CStr::from_ptr(err_ptr).to_string_lossy().into_owned()
+ } else {
+ "Unknown error (null error message)".to_string()
+ }
+ } else {
+ "Unknown error (get_last_error not available)".to_string()
+ };
+
+ log::error!("GpuSpatial FFI Error: {}", error_string);
Review Comment:
The returned error should be sufficient? (The caller can always log this if they would like to?)
```suggestion
```
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,523 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+use crate::libgpuspatial_glue_bindgen::*;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::{c_char, c_uint};
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Helper Functions
+// ----------------------------------------------------------------------
+/// Helper to handle the common pattern of calling a C function returning an int status,
+/// checking if it failed, and retrieving the error message if so.
+///
+/// T: The type of the object (Runtime, Index, Refiner) being operated on.
+unsafe fn check_ffi_call<T, F, ErrMap>(
+ call_fn: F,
+ get_error_fn: Option<unsafe extern "C" fn(*mut T) -> *const c_char>,
+ obj_ptr: *mut T,
+ err_mapper: ErrMap,
+) -> Result<(), GpuSpatialError>
+where
+ F: FnOnce() -> i32,
+ ErrMap: FnOnce(String) -> GpuSpatialError,
+{
+ if call_fn() != 0 {
+ let error_string = if let Some(get_err) = get_error_fn {
+ let err_ptr = get_err(obj_ptr);
+ if !err_ptr.is_null() {
+ CStr::from_ptr(err_ptr).to_string_lossy().into_owned()
+ } else {
+ "Unknown error (null error message)".to_string()
+ }
+ } else {
+ "Unknown error (get_last_error not available)".to_string()
+ };
+
+ log::error!("GpuSpatial FFI Error: {}", error_string);
+ return Err(err_mapper(error_string));
+ }
+ Ok(())
+}
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ /// Initializes the GpuSpatialRuntime.
+ /// This function should only be called once per engine instance.
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ check_ffi_call(
+ || init_fn(&runtime as *const _ as *mut _, &mut config),
+ runtime.get_last_error,
+ &runtime as *const _ as *mut _,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialIndexFloat2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+impl GpuSpatialIndexFloat2DWrapper {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ // Can't use check_ffi_call helper here easily because 'index' isn't fully initialized/wrapped yet
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(GpuSpatialIndexFloat2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.index.clear {
+ unsafe {
+ clear_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
+ &mut self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.index.push_build {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ Ok(())
+ }
+
+ pub fn finish_building(&mut self) -> Result<(), GpuSpatialError> {
+ if let Some(finish_building_fn) = self.index.finish_building {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn create_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(create_context_fn) = self.index.create_context {
+ unsafe {
+ create_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub fn destroy_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(destroy_context_fn) = self.index.destroy_context {
+ unsafe {
+ destroy_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn probe(
+ &self,
+ ctx: &mut SedonaSpatialIndexContext,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(probe_fn) = self.index.probe {
+ // Note: probe uses context_get_last_error, not the index one
+ if probe_fn(
+ &self.index as *const _ as *mut _,
+ ctx as *mut _,
+ buf,
+ n_rects,
+ ) != 0
+ {
+ let error_string = if let Some(get_ctx_err) = self.index.context_get_last_error {
+ CStr::from_ptr(get_ctx_err(ctx))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error".to_string()
+ };
+ log::error!("DEBUG FFI: probe failed: {}", error_string);
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+ }
+ Ok(())
+ }
+
+ fn get_indices_buffer_helper(
+ &self,
+ ctx: &mut SedonaSpatialIndexContext,
+ func: Option<unsafe extern "C" fn(*mut SedonaSpatialIndexContext, *mut *mut u32, *mut u32)>,
+ ) -> &[u32] {
+ if let Some(f) = func {
+ let mut ptr: *mut u32 = std::ptr::null_mut();
+ let mut len: u32 = 0;
+ unsafe {
+ f(ctx, &mut ptr, &mut len);
+ if len > 0 && !ptr.is_null() {
+ return std::slice::from_raw_parts(ptr, len as usize);
+ }
+ }
+ }
+ &[]
+ }
+
+ pub fn get_build_indices_buffer(&self, ctx: &mut SedonaSpatialIndexContext) -> &[u32] {
+ self.get_indices_buffer_helper(ctx, self.index.get_build_indices_buffer)
+ }
+
+ pub fn get_probe_indices_buffer(&self, ctx: &mut SedonaSpatialIndexContext) -> &[u32] {
+ self.get_indices_buffer_helper(ctx, self.index.get_probe_indices_buffer)
+ }
+}
+
+impl Drop for GpuSpatialIndexFloat2DWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.index.release {
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Predicate Wrapper
+// ----------------------------------------------------------------------
+
+#[repr(u32)]
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum GpuSpatialRelationPredicateWrapper {
+ Equals = 0,
+ Disjoint = 1,
+ Touches = 2,
+ Contains = 3,
+ Covers = 4,
+ Intersects = 5,
+ Within = 6,
+ CoveredBy = 7,
+}
+
+impl TryFrom<c_uint> for GpuSpatialRelationPredicateWrapper {
+ type Error = &'static str;
+
+ fn try_from(v: c_uint) -> Result<Self, Self::Error> {
+ match v {
+ 0 => Ok(Self::Equals),
+ 1 => Ok(Self::Disjoint),
+ 2 => Ok(Self::Touches),
+ 3 => Ok(Self::Contains),
+ 4 => Ok(Self::Covers),
+ 5 => Ok(Self::Intersects),
+ 6 => Ok(Self::Within),
+ 7 => Ok(Self::CoveredBy),
+ _ => Err("Invalid GpuSpatialPredicate value"),
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Refiner Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRefinerWrapper {
+ refiner: SedonaSpatialRefiner,
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+impl GpuSpatialRefinerWrapper {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ compress_bvh: bool,
+ pipeline_batches: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut refiner = SedonaSpatialRefiner {
+ clear: None,
+ init_schema: None,
+ push_build: None,
+ finish_building: None,
+ refine: None,
+ get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialRefinerConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ compress_bvh,
+ pipeline_batches,
+ };
+
+ unsafe {
+ if GpuSpatialRefinerCreate(&mut refiner, &config) != 0 {
+ let msg = if let Some(get_err) = refiner.get_last_error {
+ CStr::from_ptr(get_err(&refiner as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Refiner Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(GpuSpatialRefinerWrapper {
+ refiner,
+ _runtime: runtime.clone(),
+ })
+ }
+
+ pub fn init_schema(
+ &self,
+ build_data_type: &DataType,
+ probe_data_type: &DataType,
+ ) -> Result<(), GpuSpatialError> {
+ let ffi_build_schema = FFI_ArrowSchema::try_from(build_data_type)?;
+ let ffi_probe_schema = FFI_ArrowSchema::try_from(probe_data_type)?;
+
+ if let Some(init_schema_fn) = self.refiner.init_schema {
+ // Use pointer casting instead of transmute for safer FFI
+ let ffi_build_ptr = &ffi_build_schema as *const _ as *const ArrowSchema;
+ let ffi_probe_ptr = &ffi_probe_schema as *const _ as *const ArrowSchema;
+
+ unsafe {
+ check_ffi_call(
+ || {
+ init_schema_fn(
+ &self.refiner as *const _ as *mut _,
+ ffi_build_ptr,
+ ffi_probe_ptr,
+ )
+ },
+ self.refiner.get_last_error,
+ &self.refiner as *const _ as *mut _,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn clear(&self) {
+ if let Some(clear_fn) = self.refiner.clear {
+ unsafe {
+ clear_fn(&self.refiner as *const _ as *mut _);
+ }
+ }
+ }
+
+ pub fn push_build(&self, array: &ArrayRef) -> Result<(), GpuSpatialError> {
+ // Keep ffi_array alive until the C function returns
+ let (ffi_array, _ffi_schema) = arrow_array::ffi::to_ffi(&array.to_data())?;
+
+ if let Some(push_build_fn) = self.refiner.push_build {
+ let ffi_array_ptr = &ffi_array as *const _ as *const ArrowArray;
+ unsafe {
+ check_ffi_call(
+ || push_build_fn(&self.refiner as *const _ as *mut _, ffi_array_ptr as *mut _),
+ self.refiner.get_last_error,
+ &self.refiner as *const _ as *mut _,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn finish_building(&self) -> Result<(), GpuSpatialError> {
+ if let Some(finish_building_fn) = self.refiner.finish_building {
+ unsafe {
+ check_ffi_call(
+ || finish_building_fn(&self.refiner as *const _ as *mut _),
+ self.refiner.get_last_error,
+ &self.refiner as *const _ as *mut _,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn refine(
+ &self,
+ probe_array: &ArrayRef,
+ predicate: GpuSpatialRelationPredicateWrapper,
+ build_indices: &mut Vec<u32>,
+ probe_indices: &mut Vec<u32>,
+ ) -> Result<(), GpuSpatialError> {
+ let (ffi_array, _ffi_schema) = arrow_array::ffi::to_ffi(&probe_array.to_data())?;
+
+ if let Some(refine_fn) = self.refiner.refine {
Review Comment:
Should this code be responsible for reserving some capacity on `build_indices` and/or `probe_indices`?
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,523 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+use crate::libgpuspatial_glue_bindgen::*;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::{c_char, c_uint};
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Helper Functions
+// ----------------------------------------------------------------------
+/// Helper to handle the common pattern of calling a C function returning an int status,
+/// checking if it failed, and retrieving the error message if so.
+///
+/// T: The type of the object (Runtime, Index, Refiner) being operated on.
+unsafe fn check_ffi_call<T, F, ErrMap>(
+ call_fn: F,
+ get_error_fn: Option<unsafe extern "C" fn(*mut T) -> *const c_char>,
+ obj_ptr: *mut T,
+ err_mapper: ErrMap,
+) -> Result<(), GpuSpatialError>
+where
+ F: FnOnce() -> i32,
+ ErrMap: FnOnce(String) -> GpuSpatialError,
+{
+ if call_fn() != 0 {
+ let error_string = if let Some(get_err) = get_error_fn {
+ let err_ptr = get_err(obj_ptr);
+ if !err_ptr.is_null() {
+ CStr::from_ptr(err_ptr).to_string_lossy().into_owned()
+ } else {
+ "Unknown error (null error message)".to_string()
+ }
+ } else {
+ "Unknown error (get_last_error not available)".to_string()
+ };
+
+ log::error!("GpuSpatial FFI Error: {}", error_string);
+ return Err(err_mapper(error_string));
+ }
+ Ok(())
+}
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ /// Initializes the GpuSpatialRuntime.
+ /// This function should only be called once per engine instance.
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ check_ffi_call(
+ || init_fn(&runtime as *const _ as *mut _, &mut config),
+ runtime.get_last_error,
+ &runtime as *const _ as *mut _,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialIndexFloat2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+impl GpuSpatialIndexFloat2DWrapper {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ // Can't use check_ffi_call helper here easily because 'index' isn't fully initialized/wrapped yet
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(GpuSpatialIndexFloat2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.index.clear {
+ unsafe {
+ clear_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
+ &mut self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.index.push_build {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ Ok(())
+ }
+
+ pub fn finish_building(&mut self) -> Result<(), GpuSpatialError> {
+ if let Some(finish_building_fn) = self.index.finish_building {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn create_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(create_context_fn) = self.index.create_context {
+ unsafe {
+ create_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub fn destroy_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(destroy_context_fn) = self.index.destroy_context {
+ unsafe {
+ destroy_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn probe(
+ &self,
+ ctx: &mut SedonaSpatialIndexContext,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(probe_fn) = self.index.probe {
+ // Note: probe uses context_get_last_error, not the index one
+ if probe_fn(
+ &self.index as *const _ as *mut _,
+ ctx as *mut _,
+ buf,
+ n_rects,
+ ) != 0
+ {
+ let error_string = if let Some(get_ctx_err) = self.index.context_get_last_error {
+ CStr::from_ptr(get_ctx_err(ctx))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error".to_string()
+ };
+ log::error!("DEBUG FFI: probe failed: {}", error_string);
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+ }
+ Ok(())
+ }
+
+ fn get_indices_buffer_helper(
+ &self,
+ ctx: &mut SedonaSpatialIndexContext,
+ func: Option<unsafe extern "C" fn(*mut SedonaSpatialIndexContext, *mut *mut u32, *mut u32)>,
+ ) -> &[u32] {
+ if let Some(f) = func {
+ let mut ptr: *mut u32 = std::ptr::null_mut();
+ let mut len: u32 = 0;
+ unsafe {
+ f(ctx, &mut ptr, &mut len);
+ if len > 0 && !ptr.is_null() {
+ return std::slice::from_raw_parts(ptr, len as usize);
+ }
+ }
+ }
+ &[]
+ }
+
+ pub fn get_build_indices_buffer(&self, ctx: &mut SedonaSpatialIndexContext) -> &[u32] {
+ self.get_indices_buffer_helper(ctx, self.index.get_build_indices_buffer)
+ }
+
+ pub fn get_probe_indices_buffer(&self, ctx: &mut SedonaSpatialIndexContext) -> &[u32] {
+ self.get_indices_buffer_helper(ctx, self.index.get_probe_indices_buffer)
+ }
+}
+
+impl Drop for GpuSpatialIndexFloat2DWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.index.release {
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Predicate Wrapper
+// ----------------------------------------------------------------------
+
+#[repr(u32)]
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum GpuSpatialRelationPredicateWrapper {
+ Equals = 0,
+ Disjoint = 1,
+ Touches = 2,
+ Contains = 3,
+ Covers = 4,
+ Intersects = 5,
+ Within = 6,
+ CoveredBy = 7,
+}
+
+impl TryFrom<c_uint> for GpuSpatialRelationPredicateWrapper {
+ type Error = &'static str;
+
+ fn try_from(v: c_uint) -> Result<Self, Self::Error> {
+ match v {
+ 0 => Ok(Self::Equals),
+ 1 => Ok(Self::Disjoint),
+ 2 => Ok(Self::Touches),
+ 3 => Ok(Self::Contains),
+ 4 => Ok(Self::Covers),
+ 5 => Ok(Self::Intersects),
+ 6 => Ok(Self::Within),
+ 7 => Ok(Self::CoveredBy),
+ _ => Err("Invalid GpuSpatialPredicate value"),
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Refiner Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRefinerWrapper {
+ refiner: SedonaSpatialRefiner,
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+impl GpuSpatialRefinerWrapper {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ compress_bvh: bool,
+ pipeline_batches: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut refiner = SedonaSpatialRefiner {
+ clear: None,
+ init_schema: None,
+ push_build: None,
+ finish_building: None,
+ refine: None,
+ get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialRefinerConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ compress_bvh,
+ pipeline_batches,
+ };
+
+ unsafe {
+ if GpuSpatialRefinerCreate(&mut refiner, &config) != 0 {
+ let msg = if let Some(get_err) = refiner.get_last_error {
+ CStr::from_ptr(get_err(&refiner as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Refiner Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(GpuSpatialRefinerWrapper {
+ refiner,
+ _runtime: runtime.clone(),
+ })
+ }
+
+ pub fn init_schema(
+ &self,
+ build_data_type: &DataType,
+ probe_data_type: &DataType,
+ ) -> Result<(), GpuSpatialError> {
+ let ffi_build_schema = FFI_ArrowSchema::try_from(build_data_type)?;
+ let ffi_probe_schema = FFI_ArrowSchema::try_from(probe_data_type)?;
+
+ if let Some(init_schema_fn) = self.refiner.init_schema {
+ // Use pointer casting instead of transmute for safer FFI
+ let ffi_build_ptr = &ffi_build_schema as *const _ as *const ArrowSchema;
+ let ffi_probe_ptr = &ffi_probe_schema as *const _ as *const ArrowSchema;
+
+ unsafe {
+ check_ffi_call(
+ || {
+ init_schema_fn(
+ &self.refiner as *const _ as *mut _,
+ ffi_build_ptr,
+ ffi_probe_ptr,
+ )
+ },
+ self.refiner.get_last_error,
+ &self.refiner as *const _ as *mut _,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn clear(&self) {
+ if let Some(clear_fn) = self.refiner.clear {
+ unsafe {
+ clear_fn(&self.refiner as *const _ as *mut _);
+ }
+ }
+ }
+
+ pub fn push_build(&self, array: &ArrayRef) -> Result<(), GpuSpatialError> {
+ // Keep ffi_array alive until the C function returns
+ let (ffi_array, _ffi_schema) = arrow_array::ffi::to_ffi(&array.to_data())?;
+
+ if let Some(push_build_fn) = self.refiner.push_build {
+ let ffi_array_ptr = &ffi_array as *const _ as *const ArrowArray;
+ unsafe {
+ check_ffi_call(
+ || push_build_fn(&self.refiner as *const _ as *mut _, ffi_array_ptr as *mut _),
Review Comment:
On the interface side, the refiner callback should accept a `const ArrowArray*` if it doesn't intend on releasing it. These is the same pattern used by `ArrowArrayViewSetArray()`, which only reads from its array input and will never release it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "pwrliang (via GitHub)" <gi...@apache.org>.
pwrliang commented on code in PR #586:
URL: https://github.com/apache/sedona-db/pull/586#discussion_r2791089980
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,523 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+use crate::libgpuspatial_glue_bindgen::*;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::{c_char, c_uint};
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Helper Functions
+// ----------------------------------------------------------------------
+/// Helper to handle the common pattern of calling a C function returning an int status,
+/// checking if it failed, and retrieving the error message if so.
+///
+/// T: The type of the object (Runtime, Index, Refiner) being operated on.
+unsafe fn check_ffi_call<T, F, ErrMap>(
+ call_fn: F,
+ get_error_fn: Option<unsafe extern "C" fn(*mut T) -> *const c_char>,
+ obj_ptr: *mut T,
+ err_mapper: ErrMap,
+) -> Result<(), GpuSpatialError>
+where
+ F: FnOnce() -> i32,
+ ErrMap: FnOnce(String) -> GpuSpatialError,
+{
+ if call_fn() != 0 {
+ let error_string = if let Some(get_err) = get_error_fn {
+ let err_ptr = get_err(obj_ptr);
+ if !err_ptr.is_null() {
+ CStr::from_ptr(err_ptr).to_string_lossy().into_owned()
+ } else {
+ "Unknown error (null error message)".to_string()
+ }
+ } else {
+ "Unknown error (get_last_error not available)".to_string()
+ };
+
+ log::error!("GpuSpatial FFI Error: {}", error_string);
+ return Err(err_mapper(error_string));
+ }
+ Ok(())
+}
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ /// Initializes the GpuSpatialRuntime.
+ /// This function should only be called once per engine instance.
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ check_ffi_call(
+ || init_fn(&runtime as *const _ as *mut _, &mut config),
+ runtime.get_last_error,
+ &runtime as *const _ as *mut _,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialIndexFloat2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+impl GpuSpatialIndexFloat2DWrapper {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ // Can't use check_ffi_call helper here easily because 'index' isn't fully initialized/wrapped yet
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(GpuSpatialIndexFloat2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.index.clear {
+ unsafe {
+ clear_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
+ &mut self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.index.push_build {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ Ok(())
+ }
+
+ pub fn finish_building(&mut self) -> Result<(), GpuSpatialError> {
+ if let Some(finish_building_fn) = self.index.finish_building {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn create_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(create_context_fn) = self.index.create_context {
+ unsafe {
+ create_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub fn destroy_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(destroy_context_fn) = self.index.destroy_context {
+ unsafe {
+ destroy_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn probe(
+ &self,
+ ctx: &mut SedonaSpatialIndexContext,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(probe_fn) = self.index.probe {
+ // Note: probe uses context_get_last_error, not the index one
+ if probe_fn(
+ &self.index as *const _ as *mut _,
+ ctx as *mut _,
+ buf,
+ n_rects,
+ ) != 0
+ {
+ let error_string = if let Some(get_ctx_err) = self.index.context_get_last_error {
+ CStr::from_ptr(get_ctx_err(ctx))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error".to_string()
+ };
+ log::error!("DEBUG FFI: probe failed: {}", error_string);
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+ }
+ Ok(())
+ }
+
+ fn get_indices_buffer_helper(
+ &self,
+ ctx: &mut SedonaSpatialIndexContext,
+ func: Option<unsafe extern "C" fn(*mut SedonaSpatialIndexContext, *mut *mut u32, *mut u32)>,
+ ) -> &[u32] {
+ if let Some(f) = func {
+ let mut ptr: *mut u32 = std::ptr::null_mut();
+ let mut len: u32 = 0;
+ unsafe {
+ f(ctx, &mut ptr, &mut len);
+ if len > 0 && !ptr.is_null() {
+ return std::slice::from_raw_parts(ptr, len as usize);
+ }
+ }
+ }
+ &[]
+ }
+
+ pub fn get_build_indices_buffer(&self, ctx: &mut SedonaSpatialIndexContext) -> &[u32] {
+ self.get_indices_buffer_helper(ctx, self.index.get_build_indices_buffer)
+ }
+
+ pub fn get_probe_indices_buffer(&self, ctx: &mut SedonaSpatialIndexContext) -> &[u32] {
+ self.get_indices_buffer_helper(ctx, self.index.get_probe_indices_buffer)
+ }
+}
+
+impl Drop for GpuSpatialIndexFloat2DWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.index.release {
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Predicate Wrapper
+// ----------------------------------------------------------------------
+
+#[repr(u32)]
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum GpuSpatialRelationPredicateWrapper {
+ Equals = 0,
+ Disjoint = 1,
+ Touches = 2,
+ Contains = 3,
+ Covers = 4,
+ Intersects = 5,
+ Within = 6,
+ CoveredBy = 7,
+}
+
+impl TryFrom<c_uint> for GpuSpatialRelationPredicateWrapper {
+ type Error = &'static str;
+
+ fn try_from(v: c_uint) -> Result<Self, Self::Error> {
+ match v {
+ 0 => Ok(Self::Equals),
+ 1 => Ok(Self::Disjoint),
+ 2 => Ok(Self::Touches),
+ 3 => Ok(Self::Contains),
+ 4 => Ok(Self::Covers),
+ 5 => Ok(Self::Intersects),
+ 6 => Ok(Self::Within),
+ 7 => Ok(Self::CoveredBy),
+ _ => Err("Invalid GpuSpatialPredicate value"),
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Refiner Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRefinerWrapper {
+ refiner: SedonaSpatialRefiner,
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+impl GpuSpatialRefinerWrapper {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ compress_bvh: bool,
+ pipeline_batches: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut refiner = SedonaSpatialRefiner {
+ clear: None,
+ init_schema: None,
+ push_build: None,
+ finish_building: None,
+ refine: None,
+ get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialRefinerConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ compress_bvh,
+ pipeline_batches,
+ };
+
+ unsafe {
+ if GpuSpatialRefinerCreate(&mut refiner, &config) != 0 {
+ let msg = if let Some(get_err) = refiner.get_last_error {
+ CStr::from_ptr(get_err(&refiner as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Refiner Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(GpuSpatialRefinerWrapper {
+ refiner,
+ _runtime: runtime.clone(),
+ })
+ }
+
+ pub fn init_schema(
+ &self,
+ build_data_type: &DataType,
+ probe_data_type: &DataType,
+ ) -> Result<(), GpuSpatialError> {
+ let ffi_build_schema = FFI_ArrowSchema::try_from(build_data_type)?;
+ let ffi_probe_schema = FFI_ArrowSchema::try_from(probe_data_type)?;
+
+ if let Some(init_schema_fn) = self.refiner.init_schema {
+ // Use pointer casting instead of transmute for safer FFI
+ let ffi_build_ptr = &ffi_build_schema as *const _ as *const ArrowSchema;
+ let ffi_probe_ptr = &ffi_probe_schema as *const _ as *const ArrowSchema;
+
+ unsafe {
+ check_ffi_call(
+ || {
+ init_schema_fn(
+ &self.refiner as *const _ as *mut _,
+ ffi_build_ptr,
+ ffi_probe_ptr,
+ )
+ },
+ self.refiner.get_last_error,
+ &self.refiner as *const _ as *mut _,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn clear(&self) {
+ if let Some(clear_fn) = self.refiner.clear {
+ unsafe {
+ clear_fn(&self.refiner as *const _ as *mut _);
+ }
+ }
+ }
+
+ pub fn push_build(&self, array: &ArrayRef) -> Result<(), GpuSpatialError> {
+ // Keep ffi_array alive until the C function returns
+ let (ffi_array, _ffi_schema) = arrow_array::ffi::to_ffi(&array.to_data())?;
+
+ if let Some(push_build_fn) = self.refiner.push_build {
+ let ffi_array_ptr = &ffi_array as *const _ as *const ArrowArray;
+ unsafe {
+ check_ffi_call(
+ || push_build_fn(&self.refiner as *const _ as *mut _, ffi_array_ptr as *mut _),
+ self.refiner.get_last_error,
+ &self.refiner as *const _ as *mut _,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn finish_building(&self) -> Result<(), GpuSpatialError> {
+ if let Some(finish_building_fn) = self.refiner.finish_building {
+ unsafe {
+ check_ffi_call(
+ || finish_building_fn(&self.refiner as *const _ as *mut _),
+ self.refiner.get_last_error,
+ &self.refiner as *const _ as *mut _,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn refine(
+ &self,
+ probe_array: &ArrayRef,
+ predicate: GpuSpatialRelationPredicateWrapper,
+ build_indices: &mut Vec<u32>,
+ probe_indices: &mut Vec<u32>,
+ ) -> Result<(), GpuSpatialError> {
+ let (ffi_array, _ffi_schema) = arrow_array::ffi::to_ffi(&probe_array.to_data())?;
+
+ if let Some(refine_fn) = self.refiner.refine {
Review Comment:
Since this is a refinement operation, the sizes of build_indices and probe_indices will only decrease, so no reservation is needed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "pwrliang (via GitHub)" <gi...@apache.org>.
pwrliang commented on PR #586:
URL: https://github.com/apache/sedona-db/pull/586#issuecomment-3918217305
Thank you for the review. I think the above issues have been addressed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #586:
URL: https://github.com/apache/sedona-db/pull/586#discussion_r2796777787
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,537 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo::Rect;
+
+mod error;
#[cfg(gpu_available)]
+mod libgpuspatial;
mod libgpuspatial_glue_bindgen;
+mod options;
+mod predicate;
+// Re-export to the users
+pub use error::GpuSpatialError;
+pub use options::GpuSpatialOptions;
+pub use predicate::GpuSpatialRelationPredicate;
+
+#[cfg(gpu_available)]
+mod sys {
+ use super::libgpuspatial;
+ use super::libgpuspatial_glue_bindgen;
+ use super::*;
+
+ pub type Result<T> = std::result::Result<T, GpuSpatialError>;
+
+ use std::sync::{Arc, Mutex};
+
+ use libgpuspatial::{
+ FloatIndex2DBuilder, GpuSpatialRefinerWrapper, GpuSpatialRuntimeWrapper, SharedFloatIndex2D,
+ };
+ // To be used in the global state. We need to ensure these are Send + Sync since they will be shared across threads.
+ unsafe impl Send for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
+ unsafe impl Sync for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
+
+ // -- Global State --
+ static GLOBAL_GPUSPATIAL_RUNTIME: Mutex<Option<Arc<Mutex<GpuSpatialRuntimeWrapper>>>> =
+ Mutex::new(None);
+
+ // -- The Actual Implementation Struct --
+ pub struct SpatialImpl {
+ runtime: Option<Arc<Mutex<GpuSpatialRuntimeWrapper>>>,
+ // Index state is either building (Builder) or built (Shared)
+ index_builder: Option<FloatIndex2DBuilder>,
+ index: Option<SharedFloatIndex2D>,
+ refiner: Option<GpuSpatialRefinerWrapper>,
+ }
Review Comment:
I have some suggestions about the overall design of SpatialImpl. I am not familiar with the special complexities brought by GPU programming, so there might be some misunderstanding.
1. `SpatialImpl` needs to call `init()` to work properly. I think a more Rusty idiom here is RAII: we implement a `try_new` that accepts `options: GpuSpatialOptions`, and once the `SpatialImpl` is created it is ready to work. No two-step construction.
2. `index` does not need to be a member of `SpatialImpl`. `index_finish_building` can just consume the builder and return a `SharedFloatIndex2D` value to the caller. `probe` doesn't actually use any other field of `SpatialImpl` — it only needs the index — so it could take the index as a parameter, or even become a method on `SharedFloatIndex2D` directly.
After doing this, `SpatialImpl` mostly won't need `Option` members anymore. `runtime` and `refiner` get set in `try_new`, `index` is gone (returned to the caller). The only one that still needs to be `Option` is `index_builder`, since it's consumed by `finish` — but that's a clear lifecycle, not a "was init called?" check. We get rid of all the `.ok_or_else(|| GpuSpatialError::Init(...))` ceremony in the refiner and probe methods.
If someone wants to defer initialization, they can use `Option<SpatialImpl>`. The requirement of "only use it after initializing" just becomes transparent.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "pwrliang (via GitHub)" <gi...@apache.org>.
pwrliang commented on code in PR #586:
URL: https://github.com/apache/sedona-db/pull/586#discussion_r2820055407
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,305 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo_types::Rect;
+
+mod error;
#[cfg(gpu_available)]
+mod libgpuspatial;
mod libgpuspatial_glue_bindgen;
+mod options;
+mod predicate;
+
+pub use error::GpuSpatialError;
+pub use options::GpuSpatialOptions;
+pub use predicate::GpuSpatialRelationPredicate;
+pub use sys::{GpuSpatialIndex, GpuSpatialRefiner};
+
+#[cfg(gpu_available)]
+mod sys {
+ use super::libgpuspatial;
+ use super::*;
+ use libgpuspatial::GpuSpatialRuntimeWrapper;
+ use std::sync::{Arc, Mutex};
+
+ pub type Result<T> = std::result::Result<T, GpuSpatialError>;
+
+ // Global Runtime State
+ unsafe impl Send for GpuSpatialRuntimeWrapper {}
+ unsafe impl Sync for GpuSpatialRuntimeWrapper {}
+
+ static GLOBAL_GPUSPATIAL_RUNTIME: Mutex<Option<Arc<GpuSpatialRuntimeWrapper>>> =
+ Mutex::new(None);
+ /// Handles initialization of the GPU runtime.
+ pub struct SpatialContext {
+ runtime: Arc<GpuSpatialRuntimeWrapper>,
+ }
+
+ impl SpatialContext {
+ pub fn try_new(options: &GpuSpatialOptions) -> Result<Self> {
+ // Lock the mutex globally
+ let mut guard = GLOBAL_GPUSPATIAL_RUNTIME
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Global mutex poisoned".into()))?;
+
+ // Check if it already exists
+ if let Some(existing_runtime) = guard.as_ref() {
+ if existing_runtime.device_id != options.device_id {
+ return Err(GpuSpatialError::Init(format!(
+ "Runtime conflict: Initialized on Device {}, requested Device {}.",
+ existing_runtime.device_id, options.device_id
+ )));
+ }
+ // Return the existing one
+ return Ok(Self {
+ runtime: existing_runtime.clone(),
+ });
+ }
Review Comment:
I need `get_or_try_init`, which is unstable in the latest Rust stable release. Can we replace this later?
```rust
let runtime_ref = GLOBAL_GPUSPATIAL_RUNTIME.get_or_try_init(|| {
let out_path = PathBuf::from(env!("OUT_DIR"));
let ptx_root = out_path.join("share/gpuspatial/shaders");
let ptx_root_str = ptx_root
.to_str()
.ok_or_else(|| GpuSpatialError::Init("Invalid PTX path".to_string()))?;
// Return the raw Wrapper here (not an Arc)
libgpuspatial::GpuSpatialRuntimeWrapper::try_new(
options.device_id,
ptx_root_str,
options.cuda_use_memory_pool,
options.cuda_memory_pool_init_percent,
)
})?;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "pwrliang (via GitHub)" <gi...@apache.org>.
pwrliang commented on code in PR #586:
URL: https://github.com/apache/sedona-db/pull/586#discussion_r2805824007
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,583 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::c_char;
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Internal Wrapper
+// ----------------------------------------------------------------------
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D` struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+// The C library is designed for thread safety when used correctly (separate contexts per thread)
+unsafe impl Send for FloatIndex2DWrapper {}
+unsafe impl Sync for FloatIndex2DWrapper {}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.index.release {
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Builder
+// ----------------------------------------------------------------------
+
+/// Builder for the Spatial Index. This struct has exclusive ownership
+/// and is not thread-safe (Send but not Sync) because building is a
+/// single-threaded operation.
+pub struct FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper,
+}
+
+impl FloatIndex2DBuilder {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ },
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.inner.index.clear {
+ unsafe {
+ clear_fn(&mut self.inner.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
+ &mut self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.inner.index.push_build {
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut _;
+
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ Ok(())
+ }
+
+ /// Consumes the builder and returns a shared, thread-safe index wrapper.
+ pub fn finish(mut self) -> Result<SharedFloatIndex2D, GpuSpatialError> {
+ if let Some(finish_building_fn) = self.inner.index.finish_building {
+ // Extract to local vars
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+
+ Ok(SharedFloatIndex2D {
+ inner: Arc::new(self.inner),
+ })
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Shared Read-Only Index
+// ----------------------------------------------------------------------
+
+/// Thread-safe wrapper around the built index.
+/// Used to spawn thread-local contexts for probing.
+#[derive(Clone)]
+pub struct SharedFloatIndex2D {
+ inner: Arc<FloatIndex2DWrapper>,
+}
+
+unsafe impl Send for SharedFloatIndex2D {}
+unsafe impl Sync for SharedFloatIndex2D {}
+
+impl SharedFloatIndex2D {
+ pub fn create_context(&self) -> Result<FloatIndex2DContext, GpuSpatialError> {
+ let mut ctx = SedonaSpatialIndexContext {
+ private_data: std::ptr::null_mut(),
+ };
+
+ if let Some(create_context_fn) = self.inner.index.create_context {
+ unsafe {
+ create_context_fn(&mut ctx);
+ }
+ }
+
+ Ok(FloatIndex2DContext {
+ inner: self.inner.clone(),
+ context: ctx,
+ })
+ }
+ /// Probes the index using the provided thread-local context.
+ /// The context is modified to contain the result buffers.
+ pub unsafe fn probe(
+ &self,
+ ctx: &mut FloatIndex2DContext,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(probe_fn) = self.inner.index.probe {
+ // Get mutable pointer to the index (C API requirement, safe due to internal locking/context usage)
+ let index_ptr = &self.inner.index as *const _ as *mut SedonaFloatIndex2D;
+
+ // Pass the context from the argument
+ if probe_fn(index_ptr, &mut ctx.context, buf, n_rects) != 0 {
+ let error_string =
+ if let Some(get_ctx_err) = self.inner.index.context_get_last_error {
+ CStr::from_ptr(get_ctx_err(&mut ctx.context))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error".to_string()
+ };
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+ }
+ Ok(())
+ }
+}
+
+/// Thread-local context for probing the index.
+/// This struct is Send (can be moved between threads) but NOT Sync.
+pub struct FloatIndex2DContext {
+ inner: Arc<FloatIndex2DWrapper>, // Shared reference to the index wrapper to ensure it lives as long as the context
+ context: SedonaSpatialIndexContext, // The actual C context struct that holds thread-local state and result buffers
+}
+
+unsafe impl Send for FloatIndex2DContext {}
+
+impl FloatIndex2DContext {
+ fn get_indices_buffer_helper(
+ &mut self,
+ func: Option<unsafe extern "C" fn(*mut SedonaSpatialIndexContext, *mut *mut u32, *mut u32)>,
+ ) -> &[u32] {
+ if let Some(f) = func {
+ let mut ptr: *mut u32 = std::ptr::null_mut();
+ let mut len: u32 = 0;
+ unsafe {
+ f(&mut self.context, &mut ptr, &mut len);
+ if len > 0 && !ptr.is_null() {
+ return std::slice::from_raw_parts(ptr, len as usize);
+ }
+ }
+ }
+ &[]
+ }
+
+ pub fn get_build_indices_buffer(&mut self) -> &[u32] {
+ self.get_indices_buffer_helper(self.inner.index.get_build_indices_buffer)
+ }
+
+ pub fn get_probe_indices_buffer(&mut self) -> &[u32] {
+ self.get_indices_buffer_helper(self.inner.index.get_probe_indices_buffer)
+ }
+}
+
+impl Drop for FloatIndex2DContext {
+ fn drop(&mut self) {
+ if let Some(destroy_context_fn) = self.inner.index.destroy_context {
+ unsafe {
+ destroy_context_fn(&mut self.context);
+ }
+ }
+ }
+}
+
+struct RefinerWrapper {
+ refiner: SedonaSpatialRefiner,
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+unsafe impl Send for RefinerWrapper {}
+unsafe impl Sync for RefinerWrapper {}
+
+impl Drop for RefinerWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.refiner.release {
+ unsafe {
+ release_fn(&mut self.refiner as *mut _);
+ }
+ }
+ }
+}
+
+pub struct GpuSpatialRefinerBuilder {
+ inner: RefinerWrapper,
+}
Review Comment:
Refiner can be used in parallel, just like the index. The geometries from the build side are pushed to the refiner, and will be shared by all partitions on the stream side.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #586:
URL: https://github.com/apache/sedona-db/pull/586#discussion_r2805059025
##########
c/sedona-libgpuspatial/Cargo.toml:
##########
@@ -35,3 +35,17 @@ gpu = []
bindgen = "0.72.1"
cmake = "0.1"
which = "8.0"
+
+[dependencies]
+arrow-array = { workspace = true, features = ["ffi"] }
+arrow-schema = { workspace = true }
+thiserror = { workspace = true }
+log = { workspace = true }
+geo = { workspace = true }
Review Comment:
If we can get away with not requiring `geo` here except in the development dependencies. You can use `geo-types` if all you need is `Rect` (much lighter).
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,437 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo::Rect;
+
+mod error;
#[cfg(gpu_available)]
+mod libgpuspatial;
mod libgpuspatial_glue_bindgen;
+mod options;
+mod predicate;
+
+pub use error::GpuSpatialError;
+pub use options::GpuSpatialOptions;
+pub use predicate::GpuSpatialRelationPredicate;
+
+#[cfg(gpu_available)]
+mod sys {
+ use super::libgpuspatial;
+ use super::libgpuspatial_glue_bindgen;
+ use super::*;
+ use std::sync::{Arc, Mutex};
+
+ pub type Result<T> = std::result::Result<T, GpuSpatialError>;
+
+ // Direct type aliases to C++ wrappers
+ pub type IndexBuilderInner = libgpuspatial::FloatIndex2DBuilder;
+ pub type IndexInner = libgpuspatial::SharedFloatIndex2D;
+
+ // Refiner aliases
+ pub type RefinerBuilderInner = libgpuspatial::GpuSpatialRefinerBuilder;
+ pub type RefinerInner = libgpuspatial::GpuSpatialRefinerWrapper;
+
+ use libgpuspatial::GpuSpatialRuntimeWrapper;
+
+ // Global Runtime State
+ unsafe impl Send for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
+ unsafe impl Sync for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
+
+ static GLOBAL_GPUSPATIAL_RUNTIME: Mutex<Option<Arc<Mutex<GpuSpatialRuntimeWrapper>>>> =
+ Mutex::new(None);
+
+ /// Handles initialization of the GPU runtime.
+ pub struct SpatialContext {
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ }
+
+ impl SpatialContext {
+ pub fn try_new(options: &GpuSpatialOptions) -> Result<Self> {
+ let mut global_runtime_guard = GLOBAL_GPUSPATIAL_RUNTIME.lock().unwrap();
+
+ if global_runtime_guard.is_none() {
+ let out_path = std::path::PathBuf::from(env!("OUT_DIR"));
+ let ptx_root = out_path.join("share/gpuspatial/shaders");
+ let ptx_root_str = ptx_root
+ .to_str()
+ .ok_or_else(|| GpuSpatialError::Init("Invalid PTX path".to_string()))?;
+
+ let runtime = GpuSpatialRuntimeWrapper::try_new(
+ options.device_id,
+ ptx_root_str,
+ options.cuda_use_memory_pool,
+ options.cuda_memory_pool_init_percent,
+ )?;
+ *global_runtime_guard = Some(Arc::new(Mutex::new(runtime)));
+ }
+
+ Ok(Self {
+ runtime: global_runtime_guard.as_ref().unwrap().clone(),
+ })
+ }
+
+ pub fn runtime(&self) -> Arc<Mutex<GpuSpatialRuntimeWrapper>> {
+ self.runtime.clone()
+ }
+ }
+
+ pub struct IndexBuilderImpl {
+ inner: IndexBuilderInner,
+ }
+
+ impl IndexBuilderImpl {
+ pub fn try_new(options: &GpuSpatialOptions) -> Result<Self> {
+ let ctx = SpatialContext::try_new(options)?;
+ let inner = IndexBuilderInner::try_new(ctx.runtime(), options.concurrency)?;
+ Ok(Self { inner })
+ }
+
+ pub fn clear(&mut self) {
+ self.inner.clear()
+ }
+ pub fn push_build(&mut self, rects: &[Rect<f32>]) -> Result<()> {
+ unsafe {
+ self.inner
+ .push_build(rects.as_ptr() as *const f32, rects.len() as u32)
+ }
+ }
+
+ pub fn finish_building(self) -> Result<IndexImpl> {
+ let index = self.inner.finish()?;
+ Ok(IndexImpl { inner: index })
+ }
+ }
+
+ pub struct IndexImpl {
+ inner: IndexInner,
+ }
+
+ impl IndexImpl {
+ pub fn probe(&self, rects: &[Rect<f32>]) -> Result<(Vec<u32>, Vec<u32>)> {
+ // 1. Create a thread-local context
+ let mut ctx = self.inner.create_context()?;
+
+ // 2. Perform the probe using the context
+ unsafe {
+ self.inner
+ .probe(&mut ctx, rects.as_ptr() as *const f32, rects.len() as u32)?;
+ }
+
+ // 3. Extract results from the context
+ Ok((
+ ctx.get_build_indices_buffer().to_vec(),
+ ctx.get_probe_indices_buffer().to_vec(),
+ ))
+ }
+ }
+
+ pub struct RefinerBuilderImpl {
+ inner: RefinerBuilderInner,
+ }
+
+ impl RefinerBuilderImpl {
+ pub fn try_new(options: &GpuSpatialOptions) -> Result<Self> {
+ let ctx = SpatialContext::try_new(options)?;
+ let inner = RefinerBuilderInner::try_new(
+ ctx.runtime(),
+ options.concurrency,
+ options.compress_bvh,
+ options.pipeline_batches,
+ )?;
+ Ok(Self { inner })
+ }
+
+ pub fn init_schema(&mut self, build: &DataType, probe: &DataType) -> Result<()> {
+ self.inner.init_schema(build, probe)
+ }
+
+ pub fn clear(&mut self) {
+ self.inner.clear()
+ }
+ pub fn push_build(&mut self, array: &arrow_array::ArrayRef) -> Result<()> {
+ self.inner.push_build(array)
+ }
+
+ pub fn finish_building(self) -> Result<RefinerImpl> {
+ let refiner_wrapper = self.inner.finish_building()?;
+ Ok(RefinerImpl {
+ inner: refiner_wrapper,
+ })
+ }
+ }
+
+ pub struct RefinerImpl {
+ inner: RefinerInner,
+ }
+
+ impl RefinerImpl {
+ pub fn refine(
+ &self,
+ probe: &arrow_array::ArrayRef,
+ pred: GpuSpatialRelationPredicate,
+ bi: &mut Vec<u32>,
+ pi: &mut Vec<u32>,
+ ) -> Result<()> {
+ self.inner.refine(probe, pred, bi, pi)
+ }
+ }
+}
+
+#[cfg(not(gpu_available))]
+mod sys {
+ use super::*;
+ pub type Result<T> = std::result::Result<T, crate::error::GpuSpatialError>;
+
+ pub struct IndexBuilderImpl;
+ pub struct IndexImpl;
+ pub struct RefinerBuilderImpl;
+ pub struct RefinerImpl;
+
+ impl IndexBuilderImpl {
+ pub fn try_new(_opts: &GpuSpatialOptions) -> Result<Self> {
+ Err(GpuSpatialError::GpuNotAvailable)
+ }
+ pub fn clear(&mut self) {}
+ pub fn push_build(&mut self, _r: &[Rect<f32>]) -> Result<()> {
+ Err(GpuSpatialError::GpuNotAvailable)
+ }
+ pub fn finish_building(self) -> Result<IndexImpl> {
+ Err(GpuSpatialError::GpuNotAvailable)
+ }
+ }
+ impl IndexImpl {
+ pub fn probe(&self, _r: &[Rect<f32>]) -> Result<(Vec<u32>, Vec<u32>)> {
+ Err(GpuSpatialError::GpuNotAvailable)
+ }
+ }
+ impl RefinerBuilderImpl {
+ pub fn try_new(_opts: &GpuSpatialOptions) -> Result<Self> {
+ Err(GpuSpatialError::GpuNotAvailable)
+ }
+ pub fn init_schema(&mut self, _b: &DataType, _p: &DataType) -> Result<()> {
+ Err(GpuSpatialError::GpuNotAvailable)
+ }
+ pub fn clear(&mut self) {}
+ pub fn push_build(&mut self, _arr: &arrow_array::ArrayRef) -> Result<()> {
+ Err(GpuSpatialError::GpuNotAvailable)
+ }
+ pub fn finish_building(self) -> Result<RefinerImpl> {
+ Err(GpuSpatialError::GpuNotAvailable)
+ }
+ }
+ impl RefinerImpl {
+ pub fn refine(
+ &self,
+ _p: &arrow_array::ArrayRef,
+ _pr: GpuSpatialRelationPredicate,
+ _bi: &mut Vec<u32>,
+ _pi: &mut Vec<u32>,
+ ) -> Result<()> {
+ Err(GpuSpatialError::GpuNotAvailable)
+ }
+ }
+}
+
+/// Builder for creating a GPU Spatial Index.
+pub struct GpuSpatialIndexBuilder {
+ inner: sys::IndexBuilderImpl,
+}
Review Comment:
I think you can avoid this level of indirection with `pub use sys::{Float2DIndexBuilder, Refiner}`.
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,583 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::c_char;
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Internal Wrapper
+// ----------------------------------------------------------------------
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D` struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+// The C library is designed for thread safety when used correctly (separate contexts per thread)
+unsafe impl Send for FloatIndex2DWrapper {}
+unsafe impl Sync for FloatIndex2DWrapper {}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.index.release {
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Builder
+// ----------------------------------------------------------------------
+
+/// Builder for the Spatial Index. This struct has exclusive ownership
+/// and is not thread-safe (Send but not Sync) because building is a
+/// single-threaded operation.
+pub struct FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper,
+}
+
+impl FloatIndex2DBuilder {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ },
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.inner.index.clear {
+ unsafe {
+ clear_fn(&mut self.inner.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
+ &mut self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.inner.index.push_build {
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut _;
+
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ Ok(())
+ }
+
+ /// Consumes the builder and returns a shared, thread-safe index wrapper.
+ pub fn finish(mut self) -> Result<SharedFloatIndex2D, GpuSpatialError> {
+ if let Some(finish_building_fn) = self.inner.index.finish_building {
+ // Extract to local vars
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+
+ Ok(SharedFloatIndex2D {
+ inner: Arc::new(self.inner),
+ })
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Shared Read-Only Index
+// ----------------------------------------------------------------------
+
+/// Thread-safe wrapper around the built index.
+/// Used to spawn thread-local contexts for probing.
+#[derive(Clone)]
+pub struct SharedFloatIndex2D {
+ inner: Arc<FloatIndex2DWrapper>,
+}
+
+unsafe impl Send for SharedFloatIndex2D {}
+unsafe impl Sync for SharedFloatIndex2D {}
+
+impl SharedFloatIndex2D {
+ pub fn create_context(&self) -> Result<FloatIndex2DContext, GpuSpatialError> {
+ let mut ctx = SedonaSpatialIndexContext {
+ private_data: std::ptr::null_mut(),
+ };
+
+ if let Some(create_context_fn) = self.inner.index.create_context {
+ unsafe {
+ create_context_fn(&mut ctx);
+ }
+ }
+
+ Ok(FloatIndex2DContext {
+ inner: self.inner.clone(),
+ context: ctx,
+ })
+ }
+ /// Probes the index using the provided thread-local context.
+ /// The context is modified to contain the result buffers.
+ pub unsafe fn probe(
+ &self,
+ ctx: &mut FloatIndex2DContext,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(probe_fn) = self.inner.index.probe {
+ // Get mutable pointer to the index (C API requirement, safe due to internal locking/context usage)
+ let index_ptr = &self.inner.index as *const _ as *mut SedonaFloatIndex2D;
+
+ // Pass the context from the argument
+ if probe_fn(index_ptr, &mut ctx.context, buf, n_rects) != 0 {
+ let error_string =
+ if let Some(get_ctx_err) = self.inner.index.context_get_last_error {
+ CStr::from_ptr(get_ctx_err(&mut ctx.context))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error".to_string()
+ };
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+ }
+ Ok(())
+ }
+}
+
+/// Thread-local context for probing the index.
+/// This struct is Send (can be moved between threads) but NOT Sync.
+pub struct FloatIndex2DContext {
+ inner: Arc<FloatIndex2DWrapper>, // Shared reference to the index wrapper to ensure it lives as long as the context
+ context: SedonaSpatialIndexContext, // The actual C context struct that holds thread-local state and result buffers
+}
+
+unsafe impl Send for FloatIndex2DContext {}
+
+impl FloatIndex2DContext {
+ fn get_indices_buffer_helper(
+ &mut self,
+ func: Option<unsafe extern "C" fn(*mut SedonaSpatialIndexContext, *mut *mut u32, *mut u32)>,
+ ) -> &[u32] {
+ if let Some(f) = func {
+ let mut ptr: *mut u32 = std::ptr::null_mut();
+ let mut len: u32 = 0;
+ unsafe {
+ f(&mut self.context, &mut ptr, &mut len);
+ if len > 0 && !ptr.is_null() {
+ return std::slice::from_raw_parts(ptr, len as usize);
+ }
+ }
+ }
+ &[]
+ }
+
+ pub fn get_build_indices_buffer(&mut self) -> &[u32] {
+ self.get_indices_buffer_helper(self.inner.index.get_build_indices_buffer)
+ }
+
+ pub fn get_probe_indices_buffer(&mut self) -> &[u32] {
+ self.get_indices_buffer_helper(self.inner.index.get_probe_indices_buffer)
+ }
+}
+
+impl Drop for FloatIndex2DContext {
+ fn drop(&mut self) {
+ if let Some(destroy_context_fn) = self.inner.index.destroy_context {
+ unsafe {
+ destroy_context_fn(&mut self.context);
+ }
+ }
Review Comment:
Make sure all of these functions do something loud and error-y for the case where `self.inner.index.destroy_context` is `None`. It's unlikely but if it happens we want to know about it! Using `.expect()` is probably best (I think some of our other wrappers return an internal error when they can but either is fine).
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,583 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::c_char;
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Internal Wrapper
+// ----------------------------------------------------------------------
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D` struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+// The C library is designed for thread safety when used correctly (separate contexts per thread)
+unsafe impl Send for FloatIndex2DWrapper {}
+unsafe impl Sync for FloatIndex2DWrapper {}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.index.release {
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Builder
+// ----------------------------------------------------------------------
+
+/// Builder for the Spatial Index. This struct has exclusive ownership
+/// and is not thread-safe (Send but not Sync) because building is a
+/// single-threaded operation.
+pub struct FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper,
+}
+
+impl FloatIndex2DBuilder {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ },
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.inner.index.clear {
+ unsafe {
+ clear_fn(&mut self.inner.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
Review Comment:
The other functions here use an unsafe block rather than mark their functions as unsafe, which I think is better than `unsafe fn` here (i.e., the caller of this function shouldn't need an unsafe block)
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,523 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+use crate::libgpuspatial_glue_bindgen::*;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::{c_char, c_uint};
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Helper Functions
+// ----------------------------------------------------------------------
+/// Helper to handle the common pattern of calling a C function returning an int status,
+/// checking if it failed, and retrieving the error message if so.
+///
+/// T: The type of the object (Runtime, Index, Refiner) being operated on.
+unsafe fn check_ffi_call<T, F, ErrMap>(
+ call_fn: F,
+ get_error_fn: Option<unsafe extern "C" fn(*mut T) -> *const c_char>,
+ obj_ptr: *mut T,
+ err_mapper: ErrMap,
+) -> Result<(), GpuSpatialError>
+where
+ F: FnOnce() -> i32,
+ ErrMap: FnOnce(String) -> GpuSpatialError,
+{
+ if call_fn() != 0 {
+ let error_string = if let Some(get_err) = get_error_fn {
+ let err_ptr = get_err(obj_ptr);
+ if !err_ptr.is_null() {
+ CStr::from_ptr(err_ptr).to_string_lossy().into_owned()
+ } else {
+ "Unknown error (null error message)".to_string()
+ }
+ } else {
+ "Unknown error (get_last_error not available)".to_string()
+ };
+
+ log::error!("GpuSpatial FFI Error: {}", error_string);
+ return Err(err_mapper(error_string));
+ }
+ Ok(())
+}
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ /// Initializes the GpuSpatialRuntime.
+ /// This function should only be called once per engine instance.
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ check_ffi_call(
+ || init_fn(&runtime as *const _ as *mut _, &mut config),
+ runtime.get_last_error,
+ &runtime as *const _ as *mut _,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialIndexFloat2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+impl GpuSpatialIndexFloat2DWrapper {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ // Can't use check_ffi_call helper here easily because 'index' isn't fully initialized/wrapped yet
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(GpuSpatialIndexFloat2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.index.clear {
+ unsafe {
+ clear_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
+ &mut self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.index.push_build {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ Ok(())
+ }
+
+ pub fn finish_building(&mut self) -> Result<(), GpuSpatialError> {
+ if let Some(finish_building_fn) = self.index.finish_building {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn create_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(create_context_fn) = self.index.create_context {
+ unsafe {
+ create_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub fn destroy_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(destroy_context_fn) = self.index.destroy_context {
+ unsafe {
+ destroy_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn probe(
+ &self,
+ ctx: &mut SedonaSpatialIndexContext,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
Review Comment:
Got it! That's a good reason for the context to manage the results. Perhaps add that to the documentation to the C API if it's not there already since it wasn't clear to me.
##########
c/sedona-libgpuspatial/src/lib.rs:
##########
@@ -15,6 +15,437 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo::Rect;
+
+mod error;
#[cfg(gpu_available)]
+mod libgpuspatial;
mod libgpuspatial_glue_bindgen;
+mod options;
+mod predicate;
+
+pub use error::GpuSpatialError;
+pub use options::GpuSpatialOptions;
+pub use predicate::GpuSpatialRelationPredicate;
+
+#[cfg(gpu_available)]
+mod sys {
+ use super::libgpuspatial;
+ use super::libgpuspatial_glue_bindgen;
+ use super::*;
+ use std::sync::{Arc, Mutex};
+
+ pub type Result<T> = std::result::Result<T, GpuSpatialError>;
+
+ // Direct type aliases to C++ wrappers
+ pub type IndexBuilderInner = libgpuspatial::FloatIndex2DBuilder;
+ pub type IndexInner = libgpuspatial::SharedFloatIndex2D;
+
+ // Refiner aliases
+ pub type RefinerBuilderInner = libgpuspatial::GpuSpatialRefinerBuilder;
+ pub type RefinerInner = libgpuspatial::GpuSpatialRefinerWrapper;
+
+ use libgpuspatial::GpuSpatialRuntimeWrapper;
+
+ // Global Runtime State
+ unsafe impl Send for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
+ unsafe impl Sync for libgpuspatial_glue_bindgen::GpuSpatialRuntime {}
+
+ static GLOBAL_GPUSPATIAL_RUNTIME: Mutex<Option<Arc<Mutex<GpuSpatialRuntimeWrapper>>>> =
+ Mutex::new(None);
+
+ /// Handles initialization of the GPU runtime.
+ pub struct SpatialContext {
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ }
+
+ impl SpatialContext {
+ pub fn try_new(options: &GpuSpatialOptions) -> Result<Self> {
+ let mut global_runtime_guard = GLOBAL_GPUSPATIAL_RUNTIME.lock().unwrap();
+
+ if global_runtime_guard.is_none() {
+ let out_path = std::path::PathBuf::from(env!("OUT_DIR"));
+ let ptx_root = out_path.join("share/gpuspatial/shaders");
+ let ptx_root_str = ptx_root
+ .to_str()
+ .ok_or_else(|| GpuSpatialError::Init("Invalid PTX path".to_string()))?;
+
+ let runtime = GpuSpatialRuntimeWrapper::try_new(
+ options.device_id,
+ ptx_root_str,
+ options.cuda_use_memory_pool,
+ options.cuda_memory_pool_init_percent,
+ )?;
+ *global_runtime_guard = Some(Arc::new(Mutex::new(runtime)));
+ }
+
+ Ok(Self {
+ runtime: global_runtime_guard.as_ref().unwrap().clone(),
+ })
+ }
+
+ pub fn runtime(&self) -> Arc<Mutex<GpuSpatialRuntimeWrapper>> {
+ self.runtime.clone()
+ }
+ }
+
+ pub struct IndexBuilderImpl {
+ inner: IndexBuilderInner,
+ }
+
+ impl IndexBuilderImpl {
Review Comment:
Do you need this extra level of wrapper? I think the implementations you have in libgpuspatial.rs are high-level enough that this extra level isn't needed.
In other words, I think `pub use FloatIndex2D` and `pub use Refiner` would let you avoid this level of indirection.
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,583 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::c_char;
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Internal Wrapper
+// ----------------------------------------------------------------------
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D` struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+// The C library is designed for thread safety when used correctly (separate contexts per thread)
+unsafe impl Send for FloatIndex2DWrapper {}
+unsafe impl Sync for FloatIndex2DWrapper {}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.index.release {
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Builder
+// ----------------------------------------------------------------------
+
+/// Builder for the Spatial Index. This struct has exclusive ownership
+/// and is not thread-safe (Send but not Sync) because building is a
+/// single-threaded operation.
+pub struct FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper,
+}
+
+impl FloatIndex2DBuilder {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ },
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.inner.index.clear {
+ unsafe {
+ clear_fn(&mut self.inner.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
+ &mut self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.inner.index.push_build {
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut _;
+
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ Ok(())
+ }
+
+ /// Consumes the builder and returns a shared, thread-safe index wrapper.
+ pub fn finish(mut self) -> Result<SharedFloatIndex2D, GpuSpatialError> {
+ if let Some(finish_building_fn) = self.inner.index.finish_building {
+ // Extract to local vars
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+
+ Ok(SharedFloatIndex2D {
+ inner: Arc::new(self.inner),
+ })
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Shared Read-Only Index
+// ----------------------------------------------------------------------
+
+/// Thread-safe wrapper around the built index.
+/// Used to spawn thread-local contexts for probing.
+#[derive(Clone)]
+pub struct SharedFloatIndex2D {
+ inner: Arc<FloatIndex2DWrapper>,
+}
+
+unsafe impl Send for SharedFloatIndex2D {}
+unsafe impl Sync for SharedFloatIndex2D {}
+
+impl SharedFloatIndex2D {
+ pub fn create_context(&self) -> Result<FloatIndex2DContext, GpuSpatialError> {
+ let mut ctx = SedonaSpatialIndexContext {
+ private_data: std::ptr::null_mut(),
+ };
+
+ if let Some(create_context_fn) = self.inner.index.create_context {
+ unsafe {
+ create_context_fn(&mut ctx);
+ }
+ }
+
+ Ok(FloatIndex2DContext {
+ inner: self.inner.clone(),
+ context: ctx,
+ })
+ }
+ /// Probes the index using the provided thread-local context.
+ /// The context is modified to contain the result buffers.
+ pub unsafe fn probe(
+ &self,
+ ctx: &mut FloatIndex2DContext,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(probe_fn) = self.inner.index.probe {
+ // Get mutable pointer to the index (C API requirement, safe due to internal locking/context usage)
+ let index_ptr = &self.inner.index as *const _ as *mut SedonaFloatIndex2D;
+
+ // Pass the context from the argument
+ if probe_fn(index_ptr, &mut ctx.context, buf, n_rects) != 0 {
+ let error_string =
+ if let Some(get_ctx_err) = self.inner.index.context_get_last_error {
+ CStr::from_ptr(get_ctx_err(&mut ctx.context))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error".to_string()
+ };
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+ }
+ Ok(())
+ }
+}
+
+/// Thread-local context for probing the index.
+/// This struct is Send (can be moved between threads) but NOT Sync.
+pub struct FloatIndex2DContext {
+ inner: Arc<FloatIndex2DWrapper>, // Shared reference to the index wrapper to ensure it lives as long as the context
+ context: SedonaSpatialIndexContext, // The actual C context struct that holds thread-local state and result buffers
+}
+
+unsafe impl Send for FloatIndex2DContext {}
+
+impl FloatIndex2DContext {
+ fn get_indices_buffer_helper(
Review Comment:
I think this is the appropriate place for `probe()` so that the caller doesn't need to keep track of the shared index AND the context.
I think you can also have the input and output here be more Rust-y...this file is the right level of abstraction to handle creating slices from pointers and vice versa.
```suggestion
impl FloatIndex2DContext {
pub fn probe(&mut self, &[(f64, f64, f64, f64)]) -> Result<(&[u32], &[u32]), GpuSpatialError> {
todo!()
}
fn get_indices_buffer_helper(
```
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,583 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::c_char;
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Internal Wrapper
+// ----------------------------------------------------------------------
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D` struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+// The C library is designed for thread safety when used correctly (separate contexts per thread)
+unsafe impl Send for FloatIndex2DWrapper {}
+unsafe impl Sync for FloatIndex2DWrapper {}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.index.release {
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Builder
+// ----------------------------------------------------------------------
+
+/// Builder for the Spatial Index. This struct has exclusive ownership
+/// and is not thread-safe (Send but not Sync) because building is a
+/// single-threaded operation.
+pub struct FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper,
+}
+
Review Comment:
(Apologies if I missed this somewhere else)
```suggestion
unsafe impl Send for FloatIndex2DBuilder {}
```
##########
c/sedona-libgpuspatial/src/predicate.rs:
##########
@@ -0,0 +1,61 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::os::raw::c_uint;
+
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum GpuSpatialRelationPredicate {
+ Equals,
+ Disjoint,
+ Touches,
+ Contains,
+ Covers,
+ Intersects,
+ Within,
+ CoveredBy,
+}
+
+impl GpuSpatialRelationPredicate {
+ /// Internal helper to convert the Rust enum to the C-compatible integer.
+ #[allow(dead_code)] // not used if the GPU feature is disabled
Review Comment:
I think this will work instead of allowing dead code:
```suggestion
#[cfg(gpu_available)]
impl GpuSpatialRelationPredicate {
/// Internal helper to convert the Rust enum to the C-compatible integer.
```
##########
c/sedona-libgpuspatial/src/options.rs:
##########
@@ -0,0 +1,25 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+pub struct GpuSpatialOptions {
+ pub cuda_use_memory_pool: bool,
+ pub cuda_memory_pool_init_percent: i32,
+ pub concurrency: u32,
+ pub device_id: i32,
+ pub compress_bvh: bool,
+ pub pipeline_batches: u32,
+}
Review Comment:
Can you add documentation to these public items? I know this is repetitive but this is the place that other SedonaDB developers will interact with this and it is worth a brief note about what these are and where to look for more details would be helpful.
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,583 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::c_char;
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Internal Wrapper
+// ----------------------------------------------------------------------
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D` struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+// The C library is designed for thread safety when used correctly (separate contexts per thread)
+unsafe impl Send for FloatIndex2DWrapper {}
+unsafe impl Sync for FloatIndex2DWrapper {}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.index.release {
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Builder
+// ----------------------------------------------------------------------
+
+/// Builder for the Spatial Index. This struct has exclusive ownership
+/// and is not thread-safe (Send but not Sync) because building is a
+/// single-threaded operation.
+pub struct FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper,
+}
+
+impl FloatIndex2DBuilder {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ },
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.inner.index.clear {
+ unsafe {
+ clear_fn(&mut self.inner.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
+ &mut self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.inner.index.push_build {
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut _;
+
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ Ok(())
+ }
+
+ /// Consumes the builder and returns a shared, thread-safe index wrapper.
+ pub fn finish(mut self) -> Result<SharedFloatIndex2D, GpuSpatialError> {
+ if let Some(finish_building_fn) = self.inner.index.finish_building {
+ // Extract to local vars
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+
+ Ok(SharedFloatIndex2D {
+ inner: Arc::new(self.inner),
+ })
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Shared Read-Only Index
+// ----------------------------------------------------------------------
+
+/// Thread-safe wrapper around the built index.
+/// Used to spawn thread-local contexts for probing.
+#[derive(Clone)]
+pub struct SharedFloatIndex2D {
+ inner: Arc<FloatIndex2DWrapper>,
+}
+
+unsafe impl Send for SharedFloatIndex2D {}
+unsafe impl Sync for SharedFloatIndex2D {}
+
+impl SharedFloatIndex2D {
+ pub fn create_context(&self) -> Result<FloatIndex2DContext, GpuSpatialError> {
+ let mut ctx = SedonaSpatialIndexContext {
+ private_data: std::ptr::null_mut(),
+ };
+
+ if let Some(create_context_fn) = self.inner.index.create_context {
+ unsafe {
+ create_context_fn(&mut ctx);
+ }
+ }
+
+ Ok(FloatIndex2DContext {
+ inner: self.inner.clone(),
+ context: ctx,
+ })
+ }
+ /// Probes the index using the provided thread-local context.
+ /// The context is modified to contain the result buffers.
+ pub unsafe fn probe(
+ &self,
+ ctx: &mut FloatIndex2DContext,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(probe_fn) = self.inner.index.probe {
+ // Get mutable pointer to the index (C API requirement, safe due to internal locking/context usage)
+ let index_ptr = &self.inner.index as *const _ as *mut SedonaFloatIndex2D;
+
+ // Pass the context from the argument
+ if probe_fn(index_ptr, &mut ctx.context, buf, n_rects) != 0 {
+ let error_string =
+ if let Some(get_ctx_err) = self.inner.index.context_get_last_error {
+ CStr::from_ptr(get_ctx_err(&mut ctx.context))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error".to_string()
+ };
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+ }
+ Ok(())
+ }
+}
+
+/// Thread-local context for probing the index.
+/// This struct is Send (can be moved between threads) but NOT Sync.
+pub struct FloatIndex2DContext {
+ inner: Arc<FloatIndex2DWrapper>, // Shared reference to the index wrapper to ensure it lives as long as the context
+ context: SedonaSpatialIndexContext, // The actual C context struct that holds thread-local state and result buffers
+}
+
+unsafe impl Send for FloatIndex2DContext {}
+
+impl FloatIndex2DContext {
+ fn get_indices_buffer_helper(
+ &mut self,
+ func: Option<unsafe extern "C" fn(*mut SedonaSpatialIndexContext, *mut *mut u32, *mut u32)>,
+ ) -> &[u32] {
+ if let Some(f) = func {
+ let mut ptr: *mut u32 = std::ptr::null_mut();
+ let mut len: u32 = 0;
+ unsafe {
+ f(&mut self.context, &mut ptr, &mut len);
+ if len > 0 && !ptr.is_null() {
+ return std::slice::from_raw_parts(ptr, len as usize);
+ }
+ }
+ }
+ &[]
+ }
+
+ pub fn get_build_indices_buffer(&mut self) -> &[u32] {
+ self.get_indices_buffer_helper(self.inner.index.get_build_indices_buffer)
+ }
+
+ pub fn get_probe_indices_buffer(&mut self) -> &[u32] {
+ self.get_indices_buffer_helper(self.inner.index.get_probe_indices_buffer)
+ }
+}
+
+impl Drop for FloatIndex2DContext {
+ fn drop(&mut self) {
+ if let Some(destroy_context_fn) = self.inner.index.destroy_context {
+ unsafe {
+ destroy_context_fn(&mut self.context);
+ }
+ }
+ }
+}
+
+struct RefinerWrapper {
+ refiner: SedonaSpatialRefiner,
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+unsafe impl Send for RefinerWrapper {}
+unsafe impl Sync for RefinerWrapper {}
+
+impl Drop for RefinerWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.refiner.release {
+ unsafe {
+ release_fn(&mut self.refiner as *mut _);
+ }
+ }
+ }
+}
+
+pub struct GpuSpatialRefinerBuilder {
+ inner: RefinerWrapper,
+}
Review Comment:
I am not sure you need both a Builder and a Wrapper here because to my reading a refiner can never be queried in parallel.
In other words, I think you can just call this `Refiner` and add `pub fn refine(&mut self)`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "pwrliang (via GitHub)" <gi...@apache.org>.
pwrliang commented on code in PR #586:
URL: https://github.com/apache/sedona-db/pull/586#discussion_r2796242363
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,523 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+use crate::libgpuspatial_glue_bindgen::*;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::{c_char, c_uint};
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Helper Functions
+// ----------------------------------------------------------------------
+/// Helper to handle the common pattern of calling a C function returning an int status,
+/// checking if it failed, and retrieving the error message if so.
+///
+/// T: The type of the object (Runtime, Index, Refiner) being operated on.
+unsafe fn check_ffi_call<T, F, ErrMap>(
+ call_fn: F,
+ get_error_fn: Option<unsafe extern "C" fn(*mut T) -> *const c_char>,
+ obj_ptr: *mut T,
+ err_mapper: ErrMap,
+) -> Result<(), GpuSpatialError>
+where
+ F: FnOnce() -> i32,
+ ErrMap: FnOnce(String) -> GpuSpatialError,
+{
+ if call_fn() != 0 {
+ let error_string = if let Some(get_err) = get_error_fn {
+ let err_ptr = get_err(obj_ptr);
+ if !err_ptr.is_null() {
+ CStr::from_ptr(err_ptr).to_string_lossy().into_owned()
+ } else {
+ "Unknown error (null error message)".to_string()
+ }
+ } else {
+ "Unknown error (get_last_error not available)".to_string()
+ };
+
+ log::error!("GpuSpatial FFI Error: {}", error_string);
+ return Err(err_mapper(error_string));
+ }
+ Ok(())
+}
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ /// Initializes the GpuSpatialRuntime.
+ /// This function should only be called once per engine instance.
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ check_ffi_call(
+ || init_fn(&runtime as *const _ as *mut _, &mut config),
+ runtime.get_last_error,
+ &runtime as *const _ as *mut _,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialIndexFloat2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+impl GpuSpatialIndexFloat2DWrapper {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ // Can't use check_ffi_call helper here easily because 'index' isn't fully initialized/wrapped yet
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(GpuSpatialIndexFloat2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.index.clear {
+ unsafe {
+ clear_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
+ &mut self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.index.push_build {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ Ok(())
+ }
+
+ pub fn finish_building(&mut self) -> Result<(), GpuSpatialError> {
+ if let Some(finish_building_fn) = self.index.finish_building {
+ let get_last_error = self.index.get_last_error;
+ let index_ptr = &mut self.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+ Ok(())
+ }
+
+ pub fn create_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(create_context_fn) = self.index.create_context {
+ unsafe {
+ create_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub fn destroy_context(&self, ctx: &mut SedonaSpatialIndexContext) {
+ if let Some(destroy_context_fn) = self.index.destroy_context {
+ unsafe {
+ destroy_context_fn(ctx as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn probe(
+ &self,
+ ctx: &mut SedonaSpatialIndexContext,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
Review Comment:
Do you suggest that the caller should manage the buffer that stores the probe results (build indices and probe indices)? This may not work for the probe interface because the buffer size is unknown to the caller.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #586:
URL: https://github.com/apache/sedona-db/pull/586#discussion_r2806164268
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,583 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::c_char;
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.runtime.release {
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Internal Wrapper
+// ----------------------------------------------------------------------
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D` struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+// The C library is designed for thread safety when used correctly (separate contexts per thread)
+unsafe impl Send for FloatIndex2DWrapper {}
+unsafe impl Sync for FloatIndex2DWrapper {}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.index.release {
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Builder
+// ----------------------------------------------------------------------
+
+/// Builder for the Spatial Index. This struct has exclusive ownership
+/// and is not thread-safe (Send but not Sync) because building is a
+/// single-threaded operation.
+pub struct FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper,
+}
+
+impl FloatIndex2DBuilder {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ },
+ })
+ }
+
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.inner.index.clear {
+ unsafe {
+ clear_fn(&mut self.inner.index as *mut _);
+ }
+ }
+ }
+
+ pub unsafe fn push_build(
+ &mut self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.inner.index.push_build {
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut _;
+
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ Ok(())
+ }
+
+ /// Consumes the builder and returns a shared, thread-safe index wrapper.
+ pub fn finish(mut self) -> Result<SharedFloatIndex2D, GpuSpatialError> {
+ if let Some(finish_building_fn) = self.inner.index.finish_building {
+ // Extract to local vars
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ }
+
+ Ok(SharedFloatIndex2D {
+ inner: Arc::new(self.inner),
+ })
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Shared Read-Only Index
+// ----------------------------------------------------------------------
+
+/// Thread-safe wrapper around the built index.
+/// Used to spawn thread-local contexts for probing.
+#[derive(Clone)]
+pub struct SharedFloatIndex2D {
+ inner: Arc<FloatIndex2DWrapper>,
+}
+
+unsafe impl Send for SharedFloatIndex2D {}
+unsafe impl Sync for SharedFloatIndex2D {}
+
+impl SharedFloatIndex2D {
+ pub fn create_context(&self) -> Result<FloatIndex2DContext, GpuSpatialError> {
+ let mut ctx = SedonaSpatialIndexContext {
+ private_data: std::ptr::null_mut(),
+ };
+
+ if let Some(create_context_fn) = self.inner.index.create_context {
+ unsafe {
+ create_context_fn(&mut ctx);
+ }
+ }
+
+ Ok(FloatIndex2DContext {
+ inner: self.inner.clone(),
+ context: ctx,
+ })
+ }
+ /// Probes the index using the provided thread-local context.
+ /// The context is modified to contain the result buffers.
+ pub unsafe fn probe(
+ &self,
+ ctx: &mut FloatIndex2DContext,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(), GpuSpatialError> {
+ if let Some(probe_fn) = self.inner.index.probe {
+ // Get mutable pointer to the index (C API requirement, safe due to internal locking/context usage)
+ let index_ptr = &self.inner.index as *const _ as *mut SedonaFloatIndex2D;
+
+ // Pass the context from the argument
+ if probe_fn(index_ptr, &mut ctx.context, buf, n_rects) != 0 {
+ let error_string =
+ if let Some(get_ctx_err) = self.inner.index.context_get_last_error {
+ CStr::from_ptr(get_ctx_err(&mut ctx.context))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error".to_string()
+ };
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+ }
+ Ok(())
+ }
+}
+
+/// Thread-local context for probing the index.
+/// This struct is Send (can be moved between threads) but NOT Sync.
+pub struct FloatIndex2DContext {
+ inner: Arc<FloatIndex2DWrapper>, // Shared reference to the index wrapper to ensure it lives as long as the context
+ context: SedonaSpatialIndexContext, // The actual C context struct that holds thread-local state and result buffers
+}
+
+unsafe impl Send for FloatIndex2DContext {}
+
+impl FloatIndex2DContext {
+ fn get_indices_buffer_helper(
+ &mut self,
+ func: Option<unsafe extern "C" fn(*mut SedonaSpatialIndexContext, *mut *mut u32, *mut u32)>,
+ ) -> &[u32] {
+ if let Some(f) = func {
+ let mut ptr: *mut u32 = std::ptr::null_mut();
+ let mut len: u32 = 0;
+ unsafe {
+ f(&mut self.context, &mut ptr, &mut len);
+ if len > 0 && !ptr.is_null() {
+ return std::slice::from_raw_parts(ptr, len as usize);
+ }
+ }
+ }
+ &[]
+ }
+
+ pub fn get_build_indices_buffer(&mut self) -> &[u32] {
+ self.get_indices_buffer_helper(self.inner.index.get_build_indices_buffer)
+ }
+
+ pub fn get_probe_indices_buffer(&mut self) -> &[u32] {
+ self.get_indices_buffer_helper(self.inner.index.get_probe_indices_buffer)
+ }
+}
+
+impl Drop for FloatIndex2DContext {
+ fn drop(&mut self) {
+ if let Some(destroy_context_fn) = self.inner.index.destroy_context {
+ unsafe {
+ destroy_context_fn(&mut self.context);
+ }
+ }
+ }
+}
+
+struct RefinerWrapper {
+ refiner: SedonaSpatialRefiner,
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+unsafe impl Send for RefinerWrapper {}
+unsafe impl Sync for RefinerWrapper {}
+
+impl Drop for RefinerWrapper {
+ fn drop(&mut self) {
+ if let Some(release_fn) = self.refiner.release {
+ unsafe {
+ release_fn(&mut self.refiner as *mut _);
+ }
+ }
+ }
+}
+
+pub struct GpuSpatialRefinerBuilder {
+ inner: RefinerWrapper,
+}
Review Comment:
Ah, I see...but we don't need the Context because there because the indices are owned by the caller for the `refine()` call. I had thought that the purpose of the index was to reduce the number of geometries whose full definitions needed to be loaded into the GPU.
If the only purpose of the context is to accumulate indices, it would be nice to eliminate it because it complicates the interface considerably. Would a callback work?
```c
static int CProbe(self_t* self, SedonaSpatialIndexContext* context, const float* buf, (*callback)(const uint32_t* build_indices, uint32_t* probe_indices, void* private_data), void* private_data);
```
...this is roughly how GEOS' index is probed as well and eliminates the need to accumulate all of the output pairs in memory at once (potentially n_build times n_probe / 2 ish). Then you don't need to separate the builder for the index.
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,607 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::c_char;
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::Init("init function is None".to_string()));
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ let release_fn = self.runtime.release.expect("release function is None");
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Internal Wrapper
+// ----------------------------------------------------------------------
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D` struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+// The C library is designed for thread safety when used correctly (separate contexts per thread)
+unsafe impl Send for FloatIndex2DWrapper {}
+unsafe impl Sync for FloatIndex2DWrapper {}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ let release_fn = self.index.release.expect("release function is None");
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Builder
+// ----------------------------------------------------------------------
+
+/// Builder for the Spatial Index. This struct has exclusive ownership
+/// and is not thread-safe (Send but not Sync) because building is a
+/// single-threaded operation.
+pub struct FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper,
+}
+
+unsafe impl Send for FloatIndex2DBuilder {}
+
+impl FloatIndex2DBuilder {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ },
+ })
+ }
+
+ pub fn clear(&mut self) {
+ let clear_fn = self.inner.index.clear.expect("clear function is None");
+ unsafe {
+ clear_fn(&mut self.inner.index as *mut _);
+ }
+ }
+
+ pub fn push_build(&mut self, buf: *const f32, n_rects: u32) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.inner.index.push_build {
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut _;
+ unsafe {
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::PushBuild(
+ "push_build function is None".to_string(),
+ ));
+ }
+ Ok(())
+ }
+
+ /// Consumes the builder and returns a shared, thread-safe index wrapper.
+ pub fn finish_building(mut self) -> Result<SharedFloatIndex2D, GpuSpatialError> {
+ if let Some(finish_building_fn) = self.inner.index.finish_building {
+ // Extract to local vars
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::FinishBuild(
+ "finish_building function is None".to_string(),
+ ));
+ }
+
+ Ok(SharedFloatIndex2D {
+ inner: Arc::new(self.inner),
+ })
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Shared Read-Only Index
+// ----------------------------------------------------------------------
+
+/// Thread-safe wrapper around the built index.
+/// Used to spawn thread-local contexts for probing.
+#[derive(Clone)]
+pub struct SharedFloatIndex2D {
+ inner: Arc<FloatIndex2DWrapper>,
+}
+
+unsafe impl Send for SharedFloatIndex2D {}
+unsafe impl Sync for SharedFloatIndex2D {}
+
+impl SharedFloatIndex2D {
+ pub fn create_context(&self) -> Result<FloatIndex2DContext, GpuSpatialError> {
+ let mut ctx = SedonaSpatialIndexContext {
+ private_data: std::ptr::null_mut(),
+ };
+
+ if let Some(create_context_fn) = self.inner.index.create_context {
+ unsafe {
+ create_context_fn(&mut ctx);
+ }
+ } else {
+ return Err(GpuSpatialError::Init(
+ "create_index function is None".to_string(),
+ ));
+ }
+
+ Ok(FloatIndex2DContext {
+ inner: self.inner.clone(),
+ context: ctx,
+ })
+ }
+}
+
+/// Thread-local context for probing the index.
+/// This struct is Send (can be moved between threads) but NOT Sync.
+pub struct FloatIndex2DContext {
+ inner: Arc<FloatIndex2DWrapper>, // Shared reference to the index wrapper to ensure it lives as long as the context
+ context: SedonaSpatialIndexContext, // The actual C context struct that holds thread-local state and result buffers
+}
+
+unsafe impl Send for FloatIndex2DContext {}
+
+impl FloatIndex2DContext {
+ /// Probes the index using the provided thread-local context.
+ /// The context is modified to contain the result buffers.
+ pub fn probe(&mut self, buf: *const f32, n_rects: u32) -> Result<(), GpuSpatialError> {
+ if let Some(probe_fn) = self.inner.index.probe {
+ // Get mutable pointer to the index (C API requirement, safe due to internal locking/context usage)
+ let index_ptr = &self.inner.index as *const _ as *mut SedonaFloatIndex2D;
+ unsafe {
+ // Pass the context from the argument
+ if probe_fn(index_ptr, &mut self.context, buf, n_rects) != 0 {
+ let error_string =
+ if let Some(get_ctx_err) = self.inner.index.context_get_last_error {
+ CStr::from_ptr(get_ctx_err(&mut self.context))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error".to_string()
+ };
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+ }
+ } else {
+ return Err(GpuSpatialError::Probe("probe function is None".to_string()));
+ }
+ Ok(())
+ }
+ fn get_indices_buffer_helper(
+ &mut self,
+ func: Option<unsafe extern "C" fn(*mut SedonaSpatialIndexContext, *mut *mut u32, *mut u32)>,
+ ) -> &[u32] {
+ let f = func.expect("get_indices_buffer function is None");
+ let mut ptr: *mut u32 = std::ptr::null_mut();
+ let mut len: u32 = 0;
+ unsafe {
+ f(&mut self.context, &mut ptr, &mut len);
+ if len > 0 && !ptr.is_null() {
+ return std::slice::from_raw_parts(ptr, len as usize);
+ }
+ }
+ &[]
+ }
+
+ pub fn get_build_indices_buffer(&mut self) -> &[u32] {
+ self.get_indices_buffer_helper(self.inner.index.get_build_indices_buffer)
+ }
+
+ pub fn get_probe_indices_buffer(&mut self) -> &[u32] {
+ self.get_indices_buffer_helper(self.inner.index.get_probe_indices_buffer)
+ }
+}
+
+impl Drop for FloatIndex2DContext {
+ fn drop(&mut self) {
+ let destroy_context_fn = self
+ .inner
+ .index
+ .destroy_context
+ .expect("destroy_context function is None");
+ unsafe {
+ destroy_context_fn(&mut self.context);
+ }
+ }
+}
+
+struct RefinerWrapper {
+ refiner: SedonaSpatialRefiner,
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+unsafe impl Send for RefinerWrapper {}
+unsafe impl Sync for RefinerWrapper {}
Review Comment:
```suggestion
unsafe impl Send for RefinerWrapper {}
```
I think the `unsafe impl Sync` would go on the `GpuSpatialRefinerBuilder`.
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,607 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::convert::TryFrom;
+use std::ffi::{CStr, CString};
+use std::os::raw::c_char;
+use std::sync::{Arc, Mutex};
+
+// ----------------------------------------------------------------------
+// Runtime Wrapper
+// ----------------------------------------------------------------------
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: GpuSpatialRuntime,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::Init("init function is None".to_string()));
+ }
+
+ Ok(GpuSpatialRuntimeWrapper { runtime })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ let release_fn = self.runtime.release.expect("release function is None");
+ unsafe {
+ release_fn(&mut self.runtime as *mut _);
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Internal Wrapper
+// ----------------------------------------------------------------------
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D` struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+// The C library is designed for thread safety when used correctly (separate contexts per thread)
+unsafe impl Send for FloatIndex2DWrapper {}
+unsafe impl Sync for FloatIndex2DWrapper {}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ let release_fn = self.index.release.expect("release function is None");
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Builder
+// ----------------------------------------------------------------------
+
+/// Builder for the Spatial Index. This struct has exclusive ownership
+/// and is not thread-safe (Send but not Sync) because building is a
+/// single-threaded operation.
+pub struct FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper,
+}
+
+unsafe impl Send for FloatIndex2DBuilder {}
+
+impl FloatIndex2DBuilder {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_build_indices_buffer: None,
+ get_probe_indices_buffer: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialIndexConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ let msg = if let Some(get_err) = index.get_last_error {
+ CStr::from_ptr(get_err(&index as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Index Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(FloatIndex2DBuilder {
+ inner: FloatIndex2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ },
+ })
+ }
+
+ pub fn clear(&mut self) {
+ let clear_fn = self.inner.index.clear.expect("clear function is None");
+ unsafe {
+ clear_fn(&mut self.inner.index as *mut _);
+ }
+ }
+
+ pub fn push_build(&mut self, buf: *const f32, n_rects: u32) -> Result<(), GpuSpatialError> {
+ if let Some(push_build_fn) = self.inner.index.push_build {
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut _;
+ unsafe {
+ check_ffi_call(
+ move || push_build_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::PushBuild(
+ "push_build function is None".to_string(),
+ ));
+ }
+ Ok(())
+ }
+
+ /// Consumes the builder and returns a shared, thread-safe index wrapper.
+ pub fn finish_building(mut self) -> Result<SharedFloatIndex2D, GpuSpatialError> {
+ if let Some(finish_building_fn) = self.inner.index.finish_building {
+ // Extract to local vars
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(index_ptr),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::FinishBuild(
+ "finish_building function is None".to_string(),
+ ));
+ }
+
+ Ok(SharedFloatIndex2D {
+ inner: Arc::new(self.inner),
+ })
+ }
+}
+
+// ----------------------------------------------------------------------
+// Spatial Index - Shared Read-Only Index
+// ----------------------------------------------------------------------
+
+/// Thread-safe wrapper around the built index.
+/// Used to spawn thread-local contexts for probing.
+#[derive(Clone)]
+pub struct SharedFloatIndex2D {
+ inner: Arc<FloatIndex2DWrapper>,
+}
+
+unsafe impl Send for SharedFloatIndex2D {}
+unsafe impl Sync for SharedFloatIndex2D {}
+
+impl SharedFloatIndex2D {
+ pub fn create_context(&self) -> Result<FloatIndex2DContext, GpuSpatialError> {
+ let mut ctx = SedonaSpatialIndexContext {
+ private_data: std::ptr::null_mut(),
+ };
+
+ if let Some(create_context_fn) = self.inner.index.create_context {
+ unsafe {
+ create_context_fn(&mut ctx);
+ }
+ } else {
+ return Err(GpuSpatialError::Init(
+ "create_index function is None".to_string(),
+ ));
+ }
+
+ Ok(FloatIndex2DContext {
+ inner: self.inner.clone(),
+ context: ctx,
+ })
+ }
+}
+
+/// Thread-local context for probing the index.
+/// This struct is Send (can be moved between threads) but NOT Sync.
+pub struct FloatIndex2DContext {
+ inner: Arc<FloatIndex2DWrapper>, // Shared reference to the index wrapper to ensure it lives as long as the context
+ context: SedonaSpatialIndexContext, // The actual C context struct that holds thread-local state and result buffers
+}
+
+unsafe impl Send for FloatIndex2DContext {}
+
+impl FloatIndex2DContext {
+ /// Probes the index using the provided thread-local context.
+ /// The context is modified to contain the result buffers.
+ pub fn probe(&mut self, buf: *const f32, n_rects: u32) -> Result<(), GpuSpatialError> {
+ if let Some(probe_fn) = self.inner.index.probe {
+ // Get mutable pointer to the index (C API requirement, safe due to internal locking/context usage)
+ let index_ptr = &self.inner.index as *const _ as *mut SedonaFloatIndex2D;
+ unsafe {
+ // Pass the context from the argument
+ if probe_fn(index_ptr, &mut self.context, buf, n_rects) != 0 {
+ let error_string =
+ if let Some(get_ctx_err) = self.inner.index.context_get_last_error {
+ CStr::from_ptr(get_ctx_err(&mut self.context))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error".to_string()
+ };
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+ }
+ } else {
+ return Err(GpuSpatialError::Probe("probe function is None".to_string()));
+ }
+ Ok(())
+ }
+ fn get_indices_buffer_helper(
+ &mut self,
+ func: Option<unsafe extern "C" fn(*mut SedonaSpatialIndexContext, *mut *mut u32, *mut u32)>,
+ ) -> &[u32] {
+ let f = func.expect("get_indices_buffer function is None");
+ let mut ptr: *mut u32 = std::ptr::null_mut();
+ let mut len: u32 = 0;
+ unsafe {
+ f(&mut self.context, &mut ptr, &mut len);
+ if len > 0 && !ptr.is_null() {
+ return std::slice::from_raw_parts(ptr, len as usize);
+ }
+ }
+ &[]
+ }
+
+ pub fn get_build_indices_buffer(&mut self) -> &[u32] {
+ self.get_indices_buffer_helper(self.inner.index.get_build_indices_buffer)
+ }
+
+ pub fn get_probe_indices_buffer(&mut self) -> &[u32] {
+ self.get_indices_buffer_helper(self.inner.index.get_probe_indices_buffer)
+ }
+}
+
+impl Drop for FloatIndex2DContext {
+ fn drop(&mut self) {
+ let destroy_context_fn = self
+ .inner
+ .index
+ .destroy_context
+ .expect("destroy_context function is None");
+ unsafe {
+ destroy_context_fn(&mut self.context);
+ }
+ }
+}
+
+struct RefinerWrapper {
+ refiner: SedonaSpatialRefiner,
+ _runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+}
+
+unsafe impl Send for RefinerWrapper {}
+unsafe impl Sync for RefinerWrapper {}
+
+impl Drop for RefinerWrapper {
+ fn drop(&mut self) {
+ let release_fn = self.refiner.release.expect("release function is None");
+ unsafe {
+ release_fn(&mut self.refiner as *mut _);
+ }
+ }
+}
+
+pub struct GpuSpatialRefinerBuilder {
+ inner: RefinerWrapper,
+}
+
+impl GpuSpatialRefinerBuilder {
+ pub fn try_new(
+ runtime: Arc<Mutex<GpuSpatialRuntimeWrapper>>,
+ concurrency: u32,
+ compress_bvh: bool,
+ pipeline_batches: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut refiner = SedonaSpatialRefiner {
+ clear: None,
+ init_schema: None,
+ push_build: None,
+ finish_building: None,
+ refine: None,
+ get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let mut engine_guard = runtime
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Failed to acquire mutex lock".to_string()))?;
+
+ let config = GpuSpatialRefinerConfig {
+ runtime: &mut engine_guard.runtime,
+ concurrency,
+ compress_bvh,
+ pipeline_batches,
+ };
+
+ unsafe {
+ if GpuSpatialRefinerCreate(&mut refiner, &config) != 0 {
+ let msg = if let Some(get_err) = refiner.get_last_error {
+ CStr::from_ptr(get_err(&refiner as *const _ as *mut _))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown error during Refiner Create".into()
+ };
+ return Err(GpuSpatialError::Init(msg));
+ }
+ }
+
+ Ok(GpuSpatialRefinerBuilder {
+ inner: RefinerWrapper {
+ refiner,
+ _runtime: runtime.clone(),
+ },
+ })
+ }
+
+ pub fn clear(&mut self) {
+ let clear_fn = self.inner.refiner.clear.expect("clear function is None");
+ unsafe {
+ clear_fn(&mut self.inner.refiner as *mut _);
+ }
+ }
+
+ pub fn init_schema(
+ &mut self,
+ build_data_type: &DataType,
+ probe_data_type: &DataType,
+ ) -> Result<(), GpuSpatialError> {
+ let ffi_build_schema = FFI_ArrowSchema::try_from(build_data_type)?;
+ let ffi_probe_schema = FFI_ArrowSchema::try_from(probe_data_type)?;
+
+ if let Some(init_schema_fn) = self.inner.refiner.init_schema {
+ let ffi_build_ptr = &ffi_build_schema as *const _ as *const ArrowSchema;
+ let ffi_probe_ptr = &ffi_probe_schema as *const _ as *const ArrowSchema;
+
+ let get_last_error = self.inner.refiner.get_last_error;
+ let refiner_ptr = &mut self.inner.refiner as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || init_schema_fn(refiner_ptr, ffi_build_ptr, ffi_probe_ptr),
+ get_last_error,
+ refiner_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::Init(
+ "init_schema function is None".to_string(),
+ ));
+ }
+ Ok(())
+ }
+
+ pub fn push_build(&mut self, array: &ArrayRef) -> Result<(), GpuSpatialError> {
+ let (ffi_array, _ffi_schema) = arrow_array::ffi::to_ffi(&array.to_data())?;
+
+ if let Some(push_build_fn) = self.inner.refiner.push_build {
+ let ffi_array_ptr = &ffi_array as *const _ as *const ArrowArray;
+
+ let get_last_error = self.inner.refiner.get_last_error;
+ let refiner_ptr = &mut self.inner.refiner as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || push_build_fn(refiner_ptr, ffi_array_ptr),
+ get_last_error,
+ refiner_ptr,
+ GpuSpatialError::PushBuild,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::PushBuild(
+ "push_build function is None".to_string(),
+ ));
+ }
+ Ok(())
+ }
+
+ pub fn finish_building(mut self) -> Result<GpuSpatialRefinerWrapper, GpuSpatialError> {
+ if let Some(finish_building_fn) = self.inner.refiner.finish_building {
+ let get_last_error = self.inner.refiner.get_last_error;
+ let refiner_ptr = &mut self.inner.refiner as *mut _;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_building_fn(refiner_ptr),
+ get_last_error,
+ refiner_ptr,
+ GpuSpatialError::FinishBuild,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::FinishBuild(
+ "finish_building function is None".to_string(),
+ ));
+ }
+
+ Ok(GpuSpatialRefinerWrapper {
+ inner: Arc::new(self.inner),
+ })
+ }
Review Comment:
I think that you can safely add `fn refine(&self, ...)` here without separating the builder (and you can impl unsafe Sync for this if that can be done safely from multiple threads).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "pwrliang (via GitHub)" <gi...@apache.org>.
pwrliang commented on PR #586:
URL: https://github.com/apache/sedona-db/pull/586#issuecomment-3894395917
@Kontinuation @paleolimbot Thank you for your suggestions. I have refactored both index and refiner with the builder pattern. There are no Option fields anymore. The resource allocation also follows the RAII idiom. Originally, I thought we could reuse the index and refiner after the execution of the spatial join, so there would be less memory allocation overhead. It may be an anti-pattern in Rust, but I'm happy with the changes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "pwrliang (via GitHub)" <gi...@apache.org>.
pwrliang commented on PR #586:
URL: https://github.com/apache/sedona-db/pull/586#issuecomment-3906159250
Thank you! I have changed the signature of the probe, making it take a callback function to return results. `libgpuspatial.rs` has also been refactored to remove the builders.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #586:
URL: https://github.com/apache/sedona-db/pull/586#discussion_r2824795442
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,502 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::cell::UnsafeCell;
+use std::convert::TryFrom;
+use std::ffi::{c_void, CStr, CString};
+use std::os::raw::c_char;
+use std::sync::Arc;
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: UnsafeCell<GpuSpatialRuntime>,
+ /// Store which device the runtime is created on
+ pub device_id: i32,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::Init("init function is None".to_string()));
+ }
+
+ Ok(GpuSpatialRuntimeWrapper {
+ runtime: UnsafeCell::new(runtime),
+ device_id,
+ })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ let runtime = self.runtime.get_mut();
+ let release_fn = runtime.release.expect("release function is None");
+ unsafe {
+ release_fn(runtime as *mut _);
+ }
+ }
+}
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D` struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<GpuSpatialRuntimeWrapper>,
+}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ let release_fn = self.index.release.expect("release function is None");
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+}
+
+pub struct FloatIndex2D {
+ inner: FloatIndex2DWrapper,
+}
+
Review Comment:
If a use case pops up where it is needed you can...I believe that Rust will automatically know that a struct whose members are all `Send`/`Sync` is also `Send`/`Sync` but I'm not clear on the details. You will get an error implementing the next level up if you need to.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot merged PR #586:
URL: https://github.com/apache/sedona-db/pull/586
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "pwrliang (via GitHub)" <gi...@apache.org>.
pwrliang commented on code in PR #586:
URL: https://github.com/apache/sedona-db/pull/586#discussion_r2820029223
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -0,0 +1,502 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::cell::UnsafeCell;
+use std::convert::TryFrom;
+use std::ffi::{c_void, CStr, CString};
+use std::os::raw::c_char;
+use std::sync::Arc;
+
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: UnsafeCell<GpuSpatialRuntime>,
+ /// Store which device the runtime is created on
+ pub device_id: i32,
+}
+
+impl GpuSpatialRuntimeWrapper {
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::Init("init function is None".to_string()));
+ }
+
+ Ok(GpuSpatialRuntimeWrapper {
+ runtime: UnsafeCell::new(runtime),
+ device_id,
+ })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ let runtime = self.runtime.get_mut();
+ let release_fn = runtime.release.expect("release function is None");
+ unsafe {
+ release_fn(runtime as *mut _);
+ }
+ }
+}
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D` struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the index
+ _runtime: Arc<GpuSpatialRuntimeWrapper>,
+}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ let release_fn = self.index.release.expect("release function is None");
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+}
+
+pub struct FloatIndex2D {
+ inner: FloatIndex2DWrapper,
+}
+
Review Comment:
Should we also annotate the public interfaces (GpuSpatialIndex and GpuSpatialRefiner) with these traits?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-libgpuspatial): Add Rust Wrapper [sedona-db]
Posted by "pwrliang (via GitHub)" <gi...@apache.org>.
pwrliang commented on PR #586:
URL: https://github.com/apache/sedona-db/pull/586#issuecomment-3888498724
Thanks for your suggestions @paleolimbot I think I have fixed most of the issues except the one related to the design of `probe` interface. Could you elaborate on the current design flaw?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org