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