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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
