This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/sedona-db.git
The following commit(s) were added to refs/heads/main by this push:
new a780bb88 feat(c/sedona-libgpuspatial): Add Rust Wrapper (#586)
a780bb88 is described below
commit a780bb88e29948c1d6fc163e6cf7b3aaa7bbd025
Author: Liang Geng <[email protected]>
AuthorDate: Wed Feb 18 17:47:22 2026 -0500
feat(c/sedona-libgpuspatial): Add Rust Wrapper (#586)
Co-authored-by: Dewey Dunnington <[email protected]>
---
Cargo.lock | 10 +
c/sedona-libgpuspatial/Cargo.toml | 14 +
.../include/gpuspatial/gpuspatial_c.h | 46 +-
.../libgpuspatial/src/gpuspatial_c.cc | 26 +-
.../libgpuspatial/test/c_wrapper_test.cc | 40 +-
c/sedona-libgpuspatial/src/error.rs | 65 +++
c/sedona-libgpuspatial/src/lib.rs | 326 +++++++++++-
c/sedona-libgpuspatial/src/libgpuspatial.rs | 578 +++++++++++++++++++++
.../src/libgpuspatial_glue_bindgen.rs | 8 +-
c/sedona-libgpuspatial/src/{lib.rs => options.rs} | 18 +-
c/sedona-libgpuspatial/src/predicate.rs | 61 +++
11 files changed, 1141 insertions(+), 51 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index f518de7f..a13a037a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5440,9 +5440,19 @@ dependencies = [
name = "sedona-libgpuspatial"
version = "0.3.0"
dependencies = [
+ "arrow-array",
+ "arrow-schema",
"bindgen",
"cmake",
+ "geo",
+ "geo-types",
+ "sedona-expr",
+ "sedona-geos",
+ "sedona-schema",
+ "sedona-testing",
+ "thiserror 2.0.17",
"which",
+ "wkt 0.14.0",
]
[[package]]
diff --git a/c/sedona-libgpuspatial/Cargo.toml
b/c/sedona-libgpuspatial/Cargo.toml
index 01840813..9b872f7b 100644
--- a/c/sedona-libgpuspatial/Cargo.toml
+++ b/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 }
+geo-types = { workspace = true }
+sedona-schema = { workspace = true }
+
+[dev-dependencies]
+wkt = { workspace = true }
+geo = { workspace = true }
+sedona-expr = { path = "../../rust/sedona-expr" }
+sedona-geos = { path = "../sedona-geos" }
+sedona-testing = { path = "../../rust/sedona-testing" }
diff --git
a/c/sedona-libgpuspatial/libgpuspatial/include/gpuspatial/gpuspatial_c.h
b/c/sedona-libgpuspatial/libgpuspatial/include/gpuspatial/gpuspatial_c.h
index 01821ac0..5842bbd5 100644
--- a/c/sedona-libgpuspatial/libgpuspatial/include/gpuspatial/gpuspatial_c.h
+++ b/c/sedona-libgpuspatial/libgpuspatial/include/gpuspatial/gpuspatial_c.h
@@ -54,7 +54,7 @@ void GpuSpatialRuntimeCreate(struct GpuSpatialRuntime*
runtime);
struct GpuSpatialIndexConfig {
/** Pointer to an initialized GpuSpatialRuntime struct */
- struct GpuSpatialRuntime* runtime;
+ const struct GpuSpatialRuntime* runtime;
/** How many threads will concurrently call Probe method */
uint32_t concurrency;
};
@@ -71,10 +71,10 @@ struct SedonaFloatIndex2D {
void (*create_context)(struct SedonaSpatialIndexContext* context);
/** Destroy a previously created context */
void (*destroy_context)(struct SedonaSpatialIndexContext* context);
- /** Push rectangles for building the spatial index, each rectangle is
represented by 4
- * floats: [min_x, min_y, max_x, max_y].
+ /** Push rectangles for building the spatial index.
+ * @param buf each rectangle is represented by 4 floats: [min_x, min_y,
max_x, max_y].
* Points can also be indexed by providing degenerated rectangles [x, y, x,
y].
- *
+ * @param n_rects The number of rectangles in the buffer
* @return 0 on success, non-zero on failure
*/
int (*push_build)(struct SedonaFloatIndex2D* self, const float* buf,
uint32_t n_rects);
@@ -85,29 +85,27 @@ struct SedonaFloatIndex2D {
*/
int (*finish_building)(struct SedonaFloatIndex2D* self);
/**
- * Probe the spatial index with the given rectangles, each rectangle is
represented by 4
- * floats: [min_x, min_y, max_x, max_y] Points can also be probed by
providing [x, y, x,
- * y] but points and rectangles cannot be mixed in one Probe call. The
results of the
- * probe will be stored in the context.
+ * Probe the spatial index with the given rectangles.
+ * @param buf The buffer of rectangles to probe, stored in the same format
as the build
+ * rectangles.
+ * @param n_rects The number of rectangles in the probe buffer.
+ * @param callback The callback function to call for each batch of results.
+ * The callback should return 0 to continue receiving results, or non-zero
to stop the
+ * probe early. The callback will be called with arrays of build and probe
indices
+ * corresponding to candidate pairs of rectangles that intersect.
+ * The user-provided callback is required to return a value that will be
further passed
+ * to the probe function to indicate whether there's an error during the
callback
+ * execution.
+ * @param user_data The user_data pointer will be passed to the callback
*
* @return 0 on success, non-zero on failure
*/
int (*probe)(struct SedonaFloatIndex2D* self, struct
SedonaSpatialIndexContext* context,
- const float* buf, uint32_t n_rects);
- /** Get the build indices buffer from the context
- *
- * @return A pointer to the buffer and its length
- */
- void (*get_build_indices_buffer)(struct SedonaSpatialIndexContext* context,
- uint32_t** build_indices,
- uint32_t* build_indices_length);
- /** Get the probe indices buffer from the context
- *
- * @return A pointer to the buffer and its length
- */
- void (*get_probe_indices_buffer)(struct SedonaSpatialIndexContext* context,
- uint32_t** probe_indices,
- uint32_t* probe_indices_length);
+ const float* buf, uint32_t n_rects,
+ int (*callback)(const uint32_t* build_indices,
+ const uint32_t* probe_indices, uint32_t length,
+ void* user_data),
+ void* user_data);
/** Get the last error message from either the index
*
* @return A pointer to the error message string
@@ -131,7 +129,7 @@ int GpuSpatialIndexFloat2DCreate(struct SedonaFloatIndex2D*
index,
struct GpuSpatialRefinerConfig {
/** Pointer to an initialized GpuSpatialRuntime struct */
- struct GpuSpatialRuntime* runtime;
+ const struct GpuSpatialRuntime* runtime;
/** How many threads will concurrently call Probe method */
uint32_t concurrency;
/** Whether to compress the BVH structures to save memory */
diff --git a/c/sedona-libgpuspatial/libgpuspatial/src/gpuspatial_c.cc
b/c/sedona-libgpuspatial/libgpuspatial/src/gpuspatial_c.cc
index 5c6b530a..0e0c5627 100644
--- a/c/sedona-libgpuspatial/libgpuspatial/src/gpuspatial_c.cc
+++ b/c/sedona-libgpuspatial/libgpuspatial/src/gpuspatial_c.cc
@@ -159,8 +159,6 @@ struct GpuSpatialIndexFloat2DExporter {
out->push_build = &CPushBuild;
out->finish_building = &CFinishBuilding;
out->probe = &CProbe;
- out->get_build_indices_buffer = &CGetBuildIndicesBuffer;
- out->get_probe_indices_buffer = &CGetProbeIndicesBuffer;
out->get_last_error = &CGetLastError;
out->context_get_last_error = &CContextGetLastError;
out->release = &CRelease;
@@ -194,12 +192,28 @@ struct GpuSpatialIndexFloat2DExporter {
}
static int CProbe(self_t* self, SedonaSpatialIndexContext* context, const
float* buf,
- uint32_t n_rects) {
- return SafeExecute(static_cast<context_t*>(context->private_data), [&] {
+ uint32_t n_rects,
+ int (*callback)(const uint32_t* build_indices,
+ const uint32_t* probe_indices, uint32_t
length,
+ void* user_data),
+ void* user_data) {
+ auto* p_ctx = static_cast<context_t*>(context->private_data);
+ // Do not use SafeExecute because this method is thread-safe and we don't
want to set
+ // last_error for the whole index if one thread encounters an error
+ try {
auto* rects = reinterpret_cast<const spatial_index_t::box_t*>(buf);
- auto& buff = static_cast<context_t*>(context->private_data)->payload;
+ auto& buff = p_ctx->payload;
use_index(self).Probe(rects, n_rects, &buff.build_indices,
&buff.probe_indices);
- });
+
+ return callback(buff.build_indices.data(), buff.probe_indices.data(),
+ buff.build_indices.size(), user_data);
+ } catch (const std::exception& e) { // user should call
context_get_last_error
+ p_ctx->last_error = std::string(e.what());
+ return EINVAL;
+ } catch (...) {
+ p_ctx->last_error = "Unknown internal error";
+ return EINVAL;
+ }
}
static void CGetBuildIndicesBuffer(struct SedonaSpatialIndexContext* context,
diff --git a/c/sedona-libgpuspatial/libgpuspatial/test/c_wrapper_test.cc
b/c/sedona-libgpuspatial/libgpuspatial/test/c_wrapper_test.cc
index 3de7fadc..35f850f0 100644
--- a/c/sedona-libgpuspatial/libgpuspatial/test/c_wrapper_test.cc
+++ b/c/sedona-libgpuspatial/libgpuspatial/test/c_wrapper_test.cc
@@ -261,19 +261,26 @@ TEST_F(CWrapperTest, InitializeJoiner) {
fpoint_t((float)xmax, (float)ymax));
}
- // Run SUT Probe
+ struct IntersectionIDs {
+ uint32_t* build_indices_ptr;
+ uint32_t* probe_indices_ptr;
+ uint32_t length;
+ };
+ IntersectionIDs intersection_ids{nullptr, nullptr, 0};
+
SedonaSpatialIndexContext idx_ctx;
index_.create_context(&idx_ctx);
- index_.probe(&index_, &idx_ctx, (float*)queries.data(), queries.size());
-
- // Retrieve SUT Results
- uint32_t* build_indices_ptr;
- uint32_t* probe_indices_ptr;
- uint32_t build_indices_length;
- uint32_t probe_indices_length;
-
- index_.get_build_indices_buffer(&idx_ctx, &build_indices_ptr,
&build_indices_length);
- index_.get_probe_indices_buffer(&idx_ctx, &probe_indices_ptr,
&probe_indices_length);
+ index_.probe(
+ &index_, &idx_ctx, (float*)queries.data(), queries.size(),
+ [](const uint32_t* build_indices, const uint32_t* probe_indices,
uint32_t length,
+ void* user_data) {
+ IntersectionIDs* ids = reinterpret_cast<IntersectionIDs*>(user_data);
+ ids->build_indices_ptr = (uint32_t*)build_indices;
+ ids->probe_indices_ptr = (uint32_t*)probe_indices;
+ ids->length = length;
+ return 0;
+ },
+ &intersection_ids);
refiner_.clear(&refiner_);
ASSERT_EQ(refiner_.push_build(&refiner_, build_array.get()), 0);
@@ -283,13 +290,14 @@ TEST_F(CWrapperTest, InitializeJoiner) {
ASSERT_EQ(refiner_.refine(
&refiner_, probe_array.get(),
SedonaSpatialRelationPredicate::SedonaSpatialPredicateContains,
- build_indices_ptr, probe_indices_ptr, build_indices_length,
&new_len),
+ intersection_ids.build_indices_ptr,
intersection_ids.probe_indices_ptr,
+ intersection_ids.length, &new_len),
0);
- std::vector<uint32_t> sut_build_indices(build_indices_ptr,
- build_indices_ptr + new_len);
- std::vector<uint32_t> sut_stream_indices(probe_indices_ptr,
- probe_indices_ptr + new_len);
+ std::vector<uint32_t> sut_build_indices(intersection_ids.build_indices_ptr,
+ intersection_ids.build_indices_ptr
+ new_len);
+ std::vector<uint32_t> sut_stream_indices(
+ intersection_ids.probe_indices_ptr, intersection_ids.probe_indices_ptr
+ new_len);
index_.destroy_context(&idx_ctx);
diff --git a/c/sedona-libgpuspatial/src/error.rs
b/c/sedona-libgpuspatial/src/error.rs
new file mode 100644
index 00000000..4ed701e3
--- /dev/null
+++ b/c/sedona-libgpuspatial/src/error.rs
@@ -0,0 +1,65 @@
+// 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;
+
+/// Errors that can occur during GPU spatial operations.
+#[derive(Error, Debug)]
+pub enum GpuSpatialError {
+ Arrow(ArrowError),
+ Init(String),
+ PushBuild(String),
+ FinishBuild(String),
+ Probe(String),
+ Refine(String),
+ GpuNotAvailable,
+}
+
+impl From<ArrowError> for GpuSpatialError {
+ fn from(value: ArrowError) -> Self {
+ GpuSpatialError::Arrow(value)
+ }
+}
+
+impl fmt::Display for GpuSpatialError {
+ fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+ match self {
+ GpuSpatialError::Arrow(error) => {
+ write!(f, "{error}")
+ }
+ GpuSpatialError::Init(errmsg) => {
+ write!(f, "Initialization failed: {}", errmsg)
+ }
+ GpuSpatialError::PushBuild(errmsg) => {
+ write!(f, "Push build failed: {}", errmsg)
+ }
+ GpuSpatialError::FinishBuild(errmsg) => {
+ write!(f, "Finish building failed: {}", errmsg)
+ }
+ GpuSpatialError::Probe(errmsg) => {
+ write!(f, "Probe failed: {}", errmsg)
+ }
+ GpuSpatialError::Refine(errmsg) => {
+ write!(f, "Refine failed: {}", errmsg)
+ }
+ GpuSpatialError::GpuNotAvailable => {
+ write!(f, "GPU is not available")
+ }
+ }
+ }
+}
diff --git a/c/sedona-libgpuspatial/src/lib.rs
b/c/sedona-libgpuspatial/src/lib.rs
index 8d9fd61c..5dc05a60 100644
--- a/c/sedona-libgpuspatial/src/lib.rs
+++ b/c/sedona-libgpuspatial/src/lib.rs
@@ -15,6 +15,330 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
+use arrow_schema::DataType;
+use geo_types::Rect;
+
+mod error;
#[cfg(gpu_available)]
+mod libgpuspatial;
mod libgpuspatial_glue_bindgen;
+mod options;
+mod predicate;
+
+pub use error::GpuSpatialError;
+pub use options::GpuSpatialOptions;
+pub use predicate::GpuSpatialRelationPredicate;
+pub use sys::{GpuSpatialIndex, GpuSpatialRefiner};
+
+#[cfg(gpu_available)]
+mod sys {
+ use super::libgpuspatial;
+ use super::*;
+ use libgpuspatial::GpuSpatialRuntimeWrapper;
+ use std::sync::{Arc, Mutex};
+
+ pub type Result<T> = std::result::Result<T, GpuSpatialError>;
+
+ static GLOBAL_GPUSPATIAL_RUNTIME:
Mutex<Option<Arc<GpuSpatialRuntimeWrapper>>> =
+ Mutex::new(None);
+ /// Handles initialization of the GPU runtime.
+ pub struct SpatialContext {
+ runtime: Arc<GpuSpatialRuntimeWrapper>,
+ }
+
+ impl SpatialContext {
+ pub fn try_new(options: &GpuSpatialOptions) -> Result<Self> {
+ // Lock the mutex globally
+ let mut guard = GLOBAL_GPUSPATIAL_RUNTIME
+ .lock()
+ .map_err(|_| GpuSpatialError::Init("Global mutex
poisoned".into()))?;
+
+ // Check if it already exists
+ if let Some(existing_runtime) = guard.as_ref() {
+ if existing_runtime.device_id != options.device_id {
+ return Err(GpuSpatialError::Init(format!(
+ "Runtime conflict: Initialized on Device {}, requested
Device {}.",
+ existing_runtime.device_id, options.device_id
+ )));
+ }
+ // Return the existing one
+ return Ok(Self {
+ runtime: existing_runtime.clone(),
+ });
+ }
+
+ let out_path = std::path::PathBuf::from(env!("OUT_DIR"));
+ let ptx_root = out_path.join("share/gpuspatial/shaders");
+ let ptx_root_str = ptx_root
+ .to_str()
+ .ok_or_else(|| GpuSpatialError::Init("Invalid PTX
path".to_string()))?;
+
+ let wrapper = libgpuspatial::GpuSpatialRuntimeWrapper::try_new(
+ options.device_id,
+ ptx_root_str,
+ options.cuda_use_memory_pool,
+ options.cuda_memory_pool_init_percent,
+ )?;
+
+ let arc_wrapper = Arc::new(wrapper);
+
+ *guard = Some(arc_wrapper.clone());
+
+ Ok(Self {
+ runtime: arc_wrapper,
+ })
+ }
+ }
+
+ /// A GPU-accelerated spatial index for 2D rectangles in FP32.
+ /// Once built, the index is immutable and can be safely shared across
threads for read-only probe operations.
+ pub struct GpuSpatialIndex {
+ inner: libgpuspatial::FloatIndex2D,
+ }
+
+ impl GpuSpatialIndex {
+ /// Creates a new GPU spatial index with the specified options.
+ /// This initializes the GPU runtime if it hasn't been initialized yet.
+ pub fn try_new(options: &GpuSpatialOptions) -> Result<Self> {
+ let ctx = SpatialContext::try_new(options)?;
+ let inner = libgpuspatial::FloatIndex2D::try_new(ctx.runtime,
options.concurrency)?;
+ Ok(Self { inner })
+ }
+
+ /// Clears any previously inserted data from the builder, allowing it
to be reused for building a new index.
+ pub fn clear(&mut self) {
+ self.inner.clear()
+ }
+
+ /// Inserts a batch of bounding boxes into the index.
+ /// Each rectangle is represented as a `Rect<f32>` with minimum and
maximum x and y coordinates.
+ /// This method accumulates these rectangles until `finish_building`
is called to finalize the index.
+ /// The method can be called multiple times to insert data in batches
before finalizing.
+ pub fn push_build(&mut self, rects: &[Rect<f32>]) -> Result<()> {
+ // Re-interpreting Rect<f32> as flat f32 array (xmin, ymin, xmax,
ymax)
+ let raw_ptr = rects.as_ptr() as *const f32;
+ self.inner.push_build(raw_ptr, rects.len() as u32)
+ }
+
+ /// Finalizes the building process and returns an immutable spatial
index that can be probed.
+ pub fn finish_building(&mut self) -> Result<()> {
+ self.inner.finish_building()
+ }
+
+ /// Probes the spatial index with a batch of rectangles and returns
pairs of matching indices from the build and probe sets.
+ pub fn probe(&self, rects: &[Rect<f32>]) -> Result<(Vec<u32>,
Vec<u32>)> {
+ let raw_ptr = rects.as_ptr() as *const f32;
+ self.inner.probe(raw_ptr, rects.len() as u32)
+ }
+ }
+
+ /// A GPU-accelerated spatial refiner that can perform various spatial
relation tests (e.g., Intersects, Contains) between two sets of geometries.
+ pub struct GpuSpatialRefiner {
+ inner: libgpuspatial::Refiner,
+ }
+
+ impl GpuSpatialRefiner {
+ /// Creates a new GPU spatial refiner with the specified options.
+ /// This initializes the GPU runtime if it hasn't been initialized yet.
+ pub fn try_new(options: &GpuSpatialOptions) -> Result<Self> {
+ let ctx = SpatialContext::try_new(options)?;
+ let inner = libgpuspatial::Refiner::try_new(
+ ctx.runtime,
+ options.concurrency,
+ options.compress_bvh,
+ options.pipeline_batches,
+ )?;
+ Ok(Self { inner })
+ }
+
+ /// Initializes the schema for the refiner based on the data types of
the build and probe geometries.
+ /// This allows the refiner to understand how to interpret the
geometry data for both sets.
+ pub fn init_schema(&mut self, build: &DataType, probe: &DataType) ->
Result<()> {
+ self.inner.init_schema(build, probe)
+ }
+
+ /// Clears any previously inserted data from the refiner, allowing it
to be reused for building a new set of geometries.
+ pub fn clear(&mut self) {
+ self.inner.clear()
+ }
+
+ /// Inserts a batch of geometries into the refiner for the build side.
+ /// The geometries are provided as an Arrow array reference, and the
refiner will process them according to the initialized schema.
+ /// This method accumulates these geometries until `finish_building`
is called to finalize the refiner.
+ pub fn push_build(&mut self, array: &arrow_array::ArrayRef) ->
Result<()> {
+ self.inner.push_build(array)
+ }
+
+ /// Finalizes the building process and prepares the refiner for
refinement operations.
+ /// After this call, the refiner is ready to perform spatial relation
tests against probe geometries.
+ pub fn finish_building(&mut self) -> Result<()> {
+ self.inner.finish_building()
+ }
+
+ /// Refines the candidate pairs of geometries based on the specified
spatial relation predicate.
+ /// The probe geometries are provided as an Arrow array reference,
+ /// and the method updates the provided vectors of build and probe
indices to
+ /// include only those pairs that satisfy the spatial relation
predicate.
+ pub fn refine(
+ &self,
+ probe: &arrow_array::ArrayRef,
+ pred: GpuSpatialRelationPredicate,
+ build_indices: &mut Vec<u32>,
+ probe_indices: &mut Vec<u32>,
+ ) -> Result<()> {
+ self.inner.refine(probe, pred, build_indices, probe_indices)
+ }
+ }
+}
+
+#[cfg(not(gpu_available))]
+mod sys {
+ use super::*;
+ pub type Result<T> = std::result::Result<T, crate::error::GpuSpatialError>;
+
+ pub struct GpuSpatialIndex;
+ pub struct GpuSpatialRefiner;
+
+ impl GpuSpatialIndex {
+ 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<GpuSpatialIndex> {
+ Err(GpuSpatialError::GpuNotAvailable)
+ }
+ pub fn probe(&self, _r: &[Rect<f32>]) -> Result<(Vec<u32>, Vec<u32>)> {
+ Err(GpuSpatialError::GpuNotAvailable)
+ }
+ }
+
+ impl GpuSpatialRefiner {
+ 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(&mut self) -> Result<()> {
+ Err(GpuSpatialError::GpuNotAvailable)
+ }
+ pub fn refine(
+ &self,
+ _p: &arrow_array::ArrayRef,
+ _pr: GpuSpatialRelationPredicate,
+ _build_indices: &mut Vec<u32>,
+ _probe_indices: &mut Vec<u32>,
+ ) -> Result<()> {
+ Err(GpuSpatialError::GpuNotAvailable)
+ }
+ }
+}
+
+#[cfg(gpu_available)]
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use geo::{BoundingRect, Point, Polygon};
+ use sedona_schema::datatypes::WKB_GEOMETRY;
+ use sedona_testing::create::create_array_storage;
+ use wkt::TryFromWkt;
+
+ #[test]
+ fn test_spatial_index() {
+ let options = GpuSpatialOptions {
+ concurrency: 1,
+ device_id: 0,
+ compress_bvh: false,
+ pipeline_batches: 1,
+ cuda_use_memory_pool: true,
+ cuda_memory_pool_init_percent: 10,
+ };
+
+ // 1. Create Builder
+ let mut index = GpuSpatialIndex::try_new(&options).expect("Failed to
create builder");
+
+ let polygon_values = &[
+ Some("POLYGON ((30 10, 40 40, 20 40, 10 20, 30 10))"),
+ Some("POLYGON ((35 10, 45 45, 15 40, 10 20, 35 10), (20 30, 35 35,
30 20, 20 30))"),
+ ];
+ let rects: Vec<Rect<f32>> = polygon_values
+ .iter()
+ .map(|w| {
+ Polygon::try_from_wkt_str(w.unwrap())
+ .unwrap()
+ .bounding_rect()
+ .unwrap()
+ })
+ .collect();
+
+ // 2. Insert Data
+ index.push_build(&rects).expect("Failed to insert");
+
+ // 3. Finish (Consumes Builder -> Returns Index)
+ index.finish_building().expect("Failed to finish building");
+
+ // 4. Probe (Index is immutable and safe)
+ let point_values = &[Some("POINT (30 20)")];
+ let points: Vec<Rect<f32>> = point_values
+ .iter()
+ .map(|w|
Point::try_from_wkt_str(w.unwrap()).unwrap().bounding_rect())
+ .collect();
+
+ let (build_idx, probe_idx) = index.probe(&points).unwrap();
+
+ assert!(!build_idx.is_empty());
+ assert_eq!(build_idx.len(), probe_idx.len());
+ }
+
+ #[test]
+ fn test_spatial_refiner() {
+ let options = GpuSpatialOptions {
+ concurrency: 1,
+ device_id: 0,
+ compress_bvh: false,
+ pipeline_batches: 1,
+ cuda_use_memory_pool: true,
+ cuda_memory_pool_init_percent: 10,
+ };
+
+ // 1. Create Refiner Builder
+ let mut refiner =
+ GpuSpatialRefiner::try_new(&options).expect("Failed to create
refiner builder");
+
+ let polygon_values = &[Some("POLYGON ((30 10, 40 40, 20 40, 10 20, 30
10))")];
+ let polygons = create_array_storage(polygon_values, &WKB_GEOMETRY);
+
+ let point_values = &[Some("POINT (30 20)")];
+ let points = create_array_storage(point_values, &WKB_GEOMETRY);
+
+ // 2. Build Refiner
+ refiner
+ .init_schema(polygons.data_type(), points.data_type())
+ .unwrap();
+
+ refiner.push_build(&polygons).unwrap();
+
+ // 3. Finish (Consumes Builder -> Returns Refiner)
+ refiner.finish_building().expect("Failed to finish refiner");
+
+ // 4. Use Refiner
+ let mut build_idx = vec![0];
+ let mut probe_idx = vec![0];
+
+ refiner
+ .refine(
+ &points,
+ GpuSpatialRelationPredicate::Intersects,
+ &mut build_idx,
+ &mut probe_idx,
+ )
+ .unwrap();
+ }
+}
diff --git a/c/sedona-libgpuspatial/src/libgpuspatial.rs
b/c/sedona-libgpuspatial/src/libgpuspatial.rs
new file mode 100644
index 00000000..8a53b61d
--- /dev/null
+++ b/c/sedona-libgpuspatial/src/libgpuspatial.rs
@@ -0,0 +1,578 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::error::GpuSpatialError;
+#[cfg(gpu_available)]
+use crate::libgpuspatial_glue_bindgen::*;
+use crate::predicate::GpuSpatialRelationPredicate;
+use arrow_array::{Array, ArrayRef};
+use arrow_schema::ffi::FFI_ArrowSchema;
+use arrow_schema::DataType;
+use std::cell::UnsafeCell;
+use std::convert::TryFrom;
+use std::ffi::{c_void, CStr, CString};
+use std::os::raw::c_char;
+use std::sync::Arc;
+
+/// Public wrapper around the C `GpuSpatialRuntime` struct that manages its
lifecycle and provides safe Rust methods to interact with it.
+pub struct GpuSpatialRuntimeWrapper {
+ runtime: UnsafeCell<GpuSpatialRuntime>,
+ /// Store which device the runtime is created on
+ pub device_id: i32,
+}
+
+unsafe impl Send for GpuSpatialRuntimeWrapper {}
+unsafe impl Sync for GpuSpatialRuntimeWrapper {}
+
+impl GpuSpatialRuntimeWrapper {
+ /// Creates a new `GpuSpatialRuntimeWrapper` instance by initializing the
underlying C struct with the provided configuration. It returns a
`GpuSpatialError` if initialization fails.
+ /// # Arguments
+ /// * `device_id` - The ID of the GPU device to use for the runtime
+ /// * `ptx_root` - The root directory where the PTX files are located
+ /// * `use_cuda_memory_pool` - Whether to use the CUDA memory pool for
allocations
+ /// * `cuda_memory_pool_init_precent` - The initial percentage of the CUDA
memory pool to use (0-100)
+ pub fn try_new(
+ device_id: i32,
+ ptx_root: &str,
+ use_cuda_memory_pool: bool,
+ cuda_memory_pool_init_precent: i32,
+ ) -> Result<GpuSpatialRuntimeWrapper, GpuSpatialError> {
+ let mut runtime = GpuSpatialRuntime {
+ init: None,
+ release: None,
+ get_last_error: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ unsafe {
+ GpuSpatialRuntimeCreate(&mut runtime);
+ }
+
+ if let Some(init_fn) = runtime.init {
+ let c_ptx_root = CString::new(ptx_root).map_err(|_| {
+ GpuSpatialError::Init("Failed to convert ptx_root to
CString".into())
+ })?;
+
+ let mut config = GpuSpatialRuntimeConfig {
+ device_id,
+ ptx_root: c_ptx_root.as_ptr(),
+ use_cuda_memory_pool,
+ cuda_memory_pool_init_precent,
+ };
+
+ unsafe {
+ let get_last_error = runtime.get_last_error;
+ let runtime_ptr = &mut runtime as *mut GpuSpatialRuntime;
+
+ check_ffi_call(
+ move || init_fn(runtime_ptr as *mut _, &mut config),
+ get_last_error,
+ runtime_ptr,
+ GpuSpatialError::Init,
+ )?;
+ }
+ } else {
+ return Err(GpuSpatialError::Init("init function is
None".to_string()));
+ }
+
+ Ok(GpuSpatialRuntimeWrapper {
+ runtime: UnsafeCell::new(runtime),
+ device_id,
+ })
+ }
+}
+
+impl Drop for GpuSpatialRuntimeWrapper {
+ fn drop(&mut self) {
+ let runtime = self.runtime.get_mut();
+ let release_fn = runtime.release.expect("release function is None");
+ unsafe {
+ release_fn(runtime as *mut _);
+ }
+ }
+}
+
+/// Internal wrapper that manages the lifecycle of the C `SedonaFloatIndex2D`
struct.
+/// It is wrapped in an `Arc` by the public structs to ensure thread safety.
+struct FloatIndex2DWrapper {
+ index: SedonaFloatIndex2D,
+ // Keep a reference to the RT engine to ensure it lives as long as the
index
+ _runtime: Arc<GpuSpatialRuntimeWrapper>,
+}
+
+impl Drop for FloatIndex2DWrapper {
+ fn drop(&mut self) {
+ let release_fn = self.index.release.expect("release function is None");
+ unsafe {
+ release_fn(&mut self.index as *mut _);
+ }
+ }
+}
+
+/// Public struct representing the 2D float spatial index. It provides safe
Rust methods to interact with the underlying C implementation.
+pub struct FloatIndex2D {
+ inner: FloatIndex2DWrapper,
+}
+
+unsafe impl Send for FloatIndex2D {}
+unsafe impl Sync for FloatIndex2D {}
+
+impl FloatIndex2D {
+ /// Creates a new `FloatIndex2D` instance by initializing the underlying C
struct with the provided configuration. It returns a `GpuSpatialError` if
initialization fails.
+ /// # Arguments
+ /// * `runtime` - An `Arc` to the `GpuSpatialRuntimeWrapper` that the
index will use for GPU operations
+ /// * `concurrency` - The maximum level of concurrency allowed to use for
probing the index
+ pub fn try_new(
+ runtime: Arc<GpuSpatialRuntimeWrapper>,
+ concurrency: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut index = SedonaFloatIndex2D {
+ clear: None,
+ create_context: None,
+ destroy_context: None,
+ push_build: None,
+ finish_building: None,
+ probe: None,
+ get_last_error: None,
+ context_get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let config = GpuSpatialIndexConfig {
+ runtime: runtime.runtime.get(),
+ concurrency,
+ };
+
+ unsafe {
+ if GpuSpatialIndexFloat2DCreate(&mut index, &config) != 0 {
+ return Err(GpuSpatialError::Init("Index Create
failed".into()));
+ }
+ }
+
+ Ok(Self {
+ inner: FloatIndex2DWrapper {
+ index,
+ _runtime: runtime.clone(),
+ },
+ })
+ }
+
+ /// Clears the internal state of the index.
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.inner.index.clear {
+ unsafe {
+ clear_fn(&mut self.inner.index as *mut _);
+ }
+ }
+ }
+
+ /// Pushes a batch of rectangles to the index for building.
+ /// # Arguments
+ /// * `buf` - A pointer to a buffer containing the rectangle data in the
format [x_min, y_min, x_max, y_max] for each rectangle
+ /// * `n_rects` - The number of rectangles in the buffer
+ /// # Returns
+ /// * `Ok(())` if the push operation is successful
+ /// * `Err(GpuSpatialError)` if an error occurs during the push operation,
with the error message retrieved from the C struct
+ pub fn push_build(&mut self, buf: *const f32, n_rects: u32) -> Result<(),
GpuSpatialError> {
+ let push_fn =
+ self.inner.index.push_build.ok_or_else(|| {
+ GpuSpatialError::PushBuild("push_build function is
None".to_string())
+ })?;
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut SedonaFloatIndex2D;
+
+ unsafe {
+ check_ffi_call(
+ move || push_fn(index_ptr, buf, n_rects),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::PushBuild,
+ )
+ }
+ }
+
+ /// Finalizes the building process of the index.
+ /// # Returns
+ /// * `Ok(())` if the finish building operation is successful
+ /// * `Err(GpuSpatialError)` if an error occurs during the finish building
operation, with the error message retrieved from the C struct
+ pub fn finish_building(&mut self) -> Result<(), GpuSpatialError> {
+ let finish_fn = self
+ .inner
+ .index
+ .finish_building
+ .ok_or_else(|| GpuSpatialError::FinishBuild("finish_building
missing".into()))?;
+ let get_last_error = self.inner.index.get_last_error;
+ let index_ptr = &mut self.inner.index as *mut SedonaFloatIndex2D;
+
+ unsafe {
+ check_ffi_call(
+ move || finish_fn(&mut self.inner.index),
+ get_last_error,
+ index_ptr,
+ GpuSpatialError::FinishBuild,
+ )
+ }
+ }
+
+ /// Probes the index with a batch of rectangles and retrieves the
candidate pairs.
+ /// # Arguments
+ /// * `buf` - A pointer to a buffer containing the rectangle data in the
format [x_min, y_min, x_max, y_max] for each rectangle
+ /// * `n_rects` - The number of rectangles in the buffer
+ /// # Returns
+ /// * `Ok((Vec<u32>, Vec<u32>))` containing the build and probe indices of
the candidate pairs if the probe operation is successful
+ /// * `Err(GpuSpatialError)` if an error occurs during the probe
operation, with the error message retrieved from the C struct or from the
callback wrapper
+ pub fn probe(
+ &self,
+ buf: *const f32,
+ n_rects: u32,
+ ) -> Result<(Vec<u32>, Vec<u32>), GpuSpatialError> {
+ let probe_fn = self
+ .inner
+ .index
+ .probe
+ .ok_or_else(|| GpuSpatialError::Probe("probe function is
None".into()))?;
+ let create_context_fn = self.inner.index.create_context;
+ let destroy_context_fn = self.inner.index.destroy_context;
+ let context_err_fn = self.inner.index.context_get_last_error;
+ let index_ptr = &self.inner.index as *const _ as *mut
SedonaFloatIndex2D;
+
+ let mut ctx = SedonaSpatialIndexContext {
+ private_data: std::ptr::null_mut(),
+ };
+ let mut state = ProbeState {
+ results: (Vec::new(), Vec::new()),
+ error: None,
+ };
+
+ unsafe {
+ if let Some(create_ctx) = create_context_fn {
+ create_ctx(&mut ctx);
+ }
+
+ let status = probe_fn(
+ index_ptr,
+ &mut ctx,
+ buf,
+ n_rects,
+ Some(probe_callback_wrapper),
+ &mut state as *mut _ as *mut c_void,
+ );
+
+ if status != 0 {
+ // IMPROVEMENT: Check Rust error first!
+ // If the callback returned -1, 'state.error' has the real
reason (the panic).
+ if let Some(callback_error) = state.error {
+ if let Some(destroy_ctx) = destroy_context_fn {
+ destroy_ctx(&mut ctx);
+ }
+ return Err(callback_error);
+ }
+
+ // If no Rust error, it was a genuine C-side error
+ let error_string = if let Some(get_ctx_err) = context_err_fn {
+ CStr::from_ptr(get_ctx_err(&mut ctx))
+ .to_string_lossy()
+ .into_owned()
+ } else {
+ "Unknown context error during probe".to_string()
+ };
+
+ if let Some(destroy_ctx) = destroy_context_fn {
+ destroy_ctx(&mut ctx);
+ }
+ return Err(GpuSpatialError::Probe(error_string));
+ }
+
+ // Cleanup on success
+ if let Some(destroy_ctx) = destroy_context_fn {
+ destroy_ctx(&mut ctx);
+ }
+ }
+
+ Ok(state.results)
+ }
+}
+
+/// Internal wrapper that manages the lifecycle of the C
`SedonaSpatialRefiner` struct.
+struct RefinerWrapper {
+ refiner: SedonaSpatialRefiner,
+ _runtime: Arc<GpuSpatialRuntimeWrapper>,
+}
+
+impl Drop for RefinerWrapper {
+ fn drop(&mut self) {
+ let release_fn = self.refiner.release.expect("release function is
None");
+ unsafe {
+ release_fn(&mut self.refiner as *mut _);
+ }
+ }
+}
+
+/// Public struct representing the spatial refiner. It provides safe Rust
methods to interact with the underlying C implementation.
+pub struct Refiner {
+ inner: RefinerWrapper,
+}
+
+unsafe impl Send for Refiner {}
+unsafe impl Sync for Refiner {}
+
+impl Refiner {
+ /// Creates a new `Refiner` instance by initializing the underlying C
struct with the provided configuration. It returns a `GpuSpatialError` if
initialization fails.
+ /// # Arguments
+ /// * `runtime` - An `Arc` to the `GpuSpatialRuntimeWrapper` that the
refiner will use for GPU operations
+ /// * `concurrency` - The maximum level of concurrency allowed to use for
refining candidate pairs
+ /// * `compress_bvh` - Whether to compress the BVH used internally by the
refiner to save memory at the cost of potentially slower refinement times
+ /// * `pipeline_batches` - The number of batches to use for pipelining the
refinement process, which can improve performance by overlapping GPU
computation with WKB parsing. A value of 1 means no pipelining.
+ pub fn try_new(
+ runtime: Arc<GpuSpatialRuntimeWrapper>,
+ concurrency: u32,
+ compress_bvh: bool,
+ pipeline_batches: u32,
+ ) -> Result<Self, GpuSpatialError> {
+ let mut refiner = SedonaSpatialRefiner {
+ clear: None,
+ init_schema: None,
+ push_build: None,
+ finish_building: None,
+ refine: None,
+ get_last_error: None,
+ release: None,
+ private_data: std::ptr::null_mut(),
+ };
+
+ let config = GpuSpatialRefinerConfig {
+ runtime: runtime.runtime.get(),
+ concurrency,
+ compress_bvh,
+ pipeline_batches,
+ };
+
+ unsafe {
+ GpuSpatialRefinerCreate(&mut refiner, &config);
+ }
+
+ Ok(Self {
+ inner: RefinerWrapper {
+ refiner,
+ _runtime: runtime.clone(),
+ },
+ })
+ }
+
+ /// Initializes the schema for the refiner using the provided build and
probe data types.
+ /// It converts the Arrow `DataType` to the C-compatible `FFI_ArrowSchema`
and calls the underlying C function.
+ /// If initialization fails, it retrieves the error message from the C
struct and returns a `GpuSpatialError`.
+ pub fn init_schema(
+ &mut self,
+ build_dt: &DataType,
+ probe_dt: &DataType,
+ ) -> Result<(), GpuSpatialError> {
+ let build_ffi = FFI_ArrowSchema::try_from(build_dt)?;
+ let probe_ffi = FFI_ArrowSchema::try_from(probe_dt)?;
+ let init_fn = self.inner.refiner.init_schema.unwrap();
+ let get_last_error = self.inner.refiner.get_last_error;
+ let refiner_ptr = &mut self.inner.refiner as *mut SedonaSpatialRefiner;
+
+ unsafe {
+ check_ffi_call(
+ || {
+ init_fn(
+ &mut self.inner.refiner,
+ &build_ffi as *const _ as *const _,
+ &probe_ffi as *const _ as *const _,
+ )
+ },
+ get_last_error,
+ refiner_ptr,
+ GpuSpatialError::Init,
+ )
+ }
+ }
+
+ /// Pushes a batch of data to the refiner for building.
+ /// It converts the provided Arrow array to its FFI representation and
calls the underlying C function.
+ /// If the push operation fails, it retrieves the error message from the C
struct and returns a `GpuSpatialError`.
+ pub fn push_build(&mut self, array: &ArrayRef) -> Result<(),
GpuSpatialError> {
+ let (ffi_array, _) = arrow_array::ffi::to_ffi(&array.to_data())?;
+ let push_fn = self.inner.refiner.push_build.unwrap();
+ let get_last_error = self.inner.refiner.get_last_error;
+ let refiner_ptr = &mut self.inner.refiner as *mut SedonaSpatialRefiner;
+
+ unsafe {
+ check_ffi_call(
+ || push_fn(&mut self.inner.refiner, &ffi_array as *const _ as
*const _),
+ get_last_error,
+ refiner_ptr,
+ GpuSpatialError::PushBuild,
+ )
+ }
+ }
+
+ /// Clears the internal state of the refiner.
+ pub fn clear(&mut self) {
+ if let Some(clear_fn) = self.inner.refiner.clear {
+ unsafe {
+ clear_fn(&mut self.inner.refiner as *mut _);
+ }
+ }
+ }
+
+ /// Finalizes the building process of the refiner.
+ pub fn finish_building(&mut self) -> Result<(), GpuSpatialError> {
+ let finish_fn = self.inner.refiner.finish_building.unwrap();
+ let get_last_error = self.inner.refiner.get_last_error;
+ let refiner_ptr = &mut self.inner.refiner as *mut SedonaSpatialRefiner;
+
+ unsafe {
+ check_ffi_call(
+ || finish_fn(&mut self.inner.refiner),
+ get_last_error,
+ refiner_ptr,
+ GpuSpatialError::FinishBuild,
+ )
+ }
+ }
+
+ /// Refines the candidate pairs based on the provided predicate.
+ /// # Arguments
+ /// * `array` - The probe array containing the geometries to be refined.
+ /// * `predicate` - The spatial relation predicate to apply for refinement.
+ /// * `build_indices` - A mutable vector of build indices corresponding to
the candidate pairs
+ /// * `probe_indices` - A mutable vector of probe indices corresponding to
the candidate pairs
+ /// # Returns
+ /// * `Ok(())` if the refinement is successful, with `build_indices` and
`probe_indices` updated to contain only the pairs that satisfy the predicate.
+ /// * `Err(GpuSpatialError)` if an error occurs during refinement, with
the error message retrieved from the C struct.
+ pub fn refine(
+ &self,
+ array: &ArrayRef,
+ predicate: GpuSpatialRelationPredicate,
+ build_indices: &mut Vec<u32>,
+ probe_indices: &mut Vec<u32>,
+ ) -> Result<(), GpuSpatialError> {
+ let (ffi_array, _) = arrow_array::ffi::to_ffi(&array.to_data())?;
+ let refine_fn = self.inner.refiner.refine.unwrap();
+ let mut new_len: u32 = 0;
+
+ unsafe {
+ check_ffi_call(
+ || {
+ refine_fn(
+ &self.inner.refiner as *const _ as *mut _,
+ &ffi_array as *const _ as *mut _,
+ predicate.as_c_uint(),
+ build_indices.as_mut_ptr(),
+ probe_indices.as_mut_ptr(),
+ build_indices.len() as u32,
+ &mut new_len,
+ )
+ },
+ self.inner.refiner.get_last_error,
+ &self.inner.refiner as *const _ as *mut _,
+ GpuSpatialError::Refine,
+ )?;
+ }
+ build_indices.truncate(new_len as usize);
+ probe_indices.truncate(new_len as usize);
+ Ok(())
+ }
+}
+
+// ----------------------------------------------------------------------
+// Helper Functions
+// ----------------------------------------------------------------------
+
+// Define the exact signature of the C error-getting function
+type ErrorFn<T> = unsafe extern "C" fn(*mut T) -> *const c_char;
+struct ProbeState {
+ results: (Vec<u32>, Vec<u32>),
+ error: Option<GpuSpatialError>,
+}
+/// Helper to handle the common pattern of calling a C function returning an
int status,
+/// checking if it failed, and retrieving the error message if so.
+unsafe fn check_ffi_call<T, F, ErrMap>(
+ call_fn: F,
+ get_error_fn: Option<ErrorFn<T>>,
+ obj_ptr: *mut T,
+ err_mapper: ErrMap,
+) -> Result<(), GpuSpatialError>
+where
+ F: FnOnce() -> i32,
+ ErrMap: FnOnce(String) -> GpuSpatialError,
+{
+ if call_fn() != 0 {
+ let error_string = if let Some(get_err) = get_error_fn {
+ let err_ptr = get_err(obj_ptr);
+ if !err_ptr.is_null() {
+ CStr::from_ptr(err_ptr).to_string_lossy().into_owned()
+ } else {
+ "Unknown error (null error message)".to_string()
+ }
+ } else {
+ "Unknown error (get_last_error not available)".to_string()
+ };
+
+ return Err(err_mapper(error_string));
+ }
+ Ok(())
+}
+
+/// Wrapper for the probe callback that C will call.
+/// It safely converts the raw pointers to Rust slices and updates the
`ProbeState` with the results.
+/// It also catches any panics that occur within the callback and stores the
error message in the `ProbeState`,
+/// returning an error code to C to indicate that the callback failed.
+unsafe extern "C" fn probe_callback_wrapper(
+ build_indices: *const u32,
+ probe_indices: *const u32,
+ length: u32,
+ user_data: *mut c_void,
+) -> i32 {
+ let state = &mut *(user_data as *mut ProbeState);
+
+ // 1. Short-circuit: If previous error exists, tell C to stop immediately.
+ if state.error.is_some() {
+ return -1;
+ }
+
+ let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+ if length > 0 {
+ let build_slice = std::slice::from_raw_parts(build_indices, length
as usize);
+ let probe_slice = std::slice::from_raw_parts(probe_indices, length
as usize);
+
+ state.results.0.extend_from_slice(build_slice);
+ state.results.1.extend_from_slice(probe_slice);
+ }
+ }));
+
+ match result {
+ Ok(_) => 0, // Success code
+ Err(payload) => {
+ // Extract panic message
+ let msg = if let Some(s) = payload.downcast_ref::<&str>() {
+ format!("Panic in callback: {}", s)
+ } else if let Some(s) = payload.downcast_ref::<String>() {
+ format!("Panic in callback: {}", s)
+ } else {
+ "Unknown panic in callback".to_string()
+ };
+
+ // Set state and return error code to C
+ state.error = Some(GpuSpatialError::Probe(msg));
+ -1
+ }
+ }
+}
diff --git a/c/sedona-libgpuspatial/src/libgpuspatial_glue_bindgen.rs
b/c/sedona-libgpuspatial/src/libgpuspatial_glue_bindgen.rs
index ce5f4aad..8eac3563 100644
--- a/c/sedona-libgpuspatial/src/libgpuspatial_glue_bindgen.rs
+++ b/c/sedona-libgpuspatial/src/libgpuspatial_glue_bindgen.rs
@@ -19,5 +19,11 @@
#![allow(non_camel_case_types)]
#![allow(non_snake_case)]
#![allow(dead_code)]
+#![allow(clippy::type_complexity)]
-include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
+mod sys {
+ include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
+}
+
+#[cfg(gpu_available)]
+pub use sys::*;
diff --git a/c/sedona-libgpuspatial/src/lib.rs
b/c/sedona-libgpuspatial/src/options.rs
similarity index 51%
copy from c/sedona-libgpuspatial/src/lib.rs
copy to c/sedona-libgpuspatial/src/options.rs
index 8d9fd61c..f8c5d715 100644
--- a/c/sedona-libgpuspatial/src/lib.rs
+++ b/c/sedona-libgpuspatial/src/options.rs
@@ -15,6 +15,18 @@
// specific language governing permissions and limitations
// under the License.
-// Module declarations
-#[cfg(gpu_available)]
-mod libgpuspatial_glue_bindgen;
+/// Options for GPU-accelerated index and refiner.
+pub struct GpuSpatialOptions {
+ /// Whether to use CUDA memory pool for allocations
+ pub cuda_use_memory_pool: bool,
+ /// Ratio of initial memory pool size to total GPU memory, between 0 and
100
+ pub cuda_memory_pool_init_percent: i32,
+ /// How many threads will concurrently use the library
+ pub concurrency: u32,
+ /// The device id to use
+ pub device_id: i32,
+ /// Whether to build a compressed BVH, which can reduce memory usage, but
may increase build time
+ pub compress_bvh: bool,
+ /// The number of batches for pipelined refinement that overlaps the WKB
loading and refinement. Setting 1 effectively disables pipelining.
+ pub pipeline_batches: u32,
+}
diff --git a/c/sedona-libgpuspatial/src/predicate.rs
b/c/sedona-libgpuspatial/src/predicate.rs
new file mode 100644
index 00000000..1287b9b2
--- /dev/null
+++ b/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.
+#[cfg(gpu_available)]
+use std::os::raw::c_uint;
+
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum GpuSpatialRelationPredicate {
+ Equals,
+ Disjoint,
+ Touches,
+ Contains,
+ Covers,
+ Intersects,
+ Within,
+ CoveredBy,
+}
+
+#[cfg(gpu_available)] // not used if the GPU feature is disabled
+impl GpuSpatialRelationPredicate {
+ /// Internal helper to convert the Rust enum to the C-compatible integer.
+ pub(crate) fn as_c_uint(self) -> c_uint {
+ match self {
+ Self::Equals => 0,
+ Self::Disjoint => 1,
+ Self::Touches => 2,
+ Self::Contains => 3,
+ Self::Covers => 4,
+ Self::Intersects => 5,
+ Self::Within => 6,
+ Self::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"),
+ }
+ }
+}
