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 82c1bb05 feat(c/sedona-libgpuspatial): Change the interface of
SpatialRefiner (#717)
82c1bb05 is described below
commit 82c1bb05a105c428c1831fb7226b72a40896029c
Author: Liang Geng <[email protected]>
AuthorDate: Thu Mar 19 05:57:32 2026 +0800
feat(c/sedona-libgpuspatial): Change the interface of SpatialRefiner (#717)
---
.../include/gpuspatial/gpuspatial_c.h | 14 +++++---
.../libgpuspatial/src/gpuspatial_c.cc | 38 ++++++++++------------
.../libgpuspatial/test/c_wrapper_test.cc | 10 +++---
c/sedona-libgpuspatial/src/lib.rs | 16 ++++-----
c/sedona-libgpuspatial/src/libgpuspatial.rs | 26 ++++++---------
5 files changed, 51 insertions(+), 53 deletions(-)
diff --git
a/c/sedona-libgpuspatial/libgpuspatial/include/gpuspatial/gpuspatial_c.h
b/c/sedona-libgpuspatial/libgpuspatial/include/gpuspatial/gpuspatial_c.h
index 5842bbd5..0b949f31 100644
--- a/c/sedona-libgpuspatial/libgpuspatial/include/gpuspatial/gpuspatial_c.h
+++ b/c/sedona-libgpuspatial/libgpuspatial/include/gpuspatial/gpuspatial_c.h
@@ -155,9 +155,13 @@ struct SedonaSpatialRefiner {
/** Clear all built geometries from the refiner */
int (*clear)(struct SedonaSpatialRefiner* self);
- int (*init_schema)(struct SedonaSpatialRefiner* self,
- const struct ArrowSchema* build_schema,
- const struct ArrowSchema* probe_schema);
+ /** Initialize the spatial refiner with the schema of the build geometries
+ *
+ * @param build_schema The Arrow schema of the build geometries;
+ * @return 0 on success, non-zero on failure
+ */
+ int (*init_build_schema)(struct SedonaSpatialRefiner* self,
+ const struct ArrowSchema* build_schema);
/** Push geometries for building the spatial refiner
*
@@ -177,6 +181,7 @@ struct SedonaSpatialRefiner {
* Refine candidate pairs of geometries
*
* @param probe_array The Arrow array of the probe geometries
+ * @param probe_schema The Arrow schema of the probe geometries
* @param predicate The spatial relation predicate to evaluate
* @param build_indices An array of build-side indices corresponding to
candidate pairs.
* This is a global index from 0 to N-1, where N is the total number of
build geometries
@@ -188,7 +193,8 @@ struct SedonaSpatialRefiner {
* @param new_indices_size Output parameter to store the number of refined
pairs
* @return 0 on success, non-zero on failure
*/
- int (*refine)(struct SedonaSpatialRefiner* self, const struct ArrowArray*
probe_array,
+ int (*refine)(struct SedonaSpatialRefiner* self, const struct ArrowSchema*
probe_schema,
+ const struct ArrowArray* probe_array,
enum SedonaSpatialRelationPredicate predicate, uint32_t*
build_indices,
uint32_t* probe_indices, uint32_t indices_size,
uint32_t* new_indices_size);
diff --git a/c/sedona-libgpuspatial/libgpuspatial/src/gpuspatial_c.cc
b/c/sedona-libgpuspatial/libgpuspatial/src/gpuspatial_c.cc
index 0e0c5627..e0fae7a4 100644
--- a/c/sedona-libgpuspatial/libgpuspatial/src/gpuspatial_c.cc
+++ b/c/sedona-libgpuspatial/libgpuspatial/src/gpuspatial_c.cc
@@ -271,7 +271,6 @@ struct GpuSpatialRefinerExporter {
struct Payload {
std::unique_ptr<gpuspatial::SpatialRefiner> refiner;
nanoarrow::UniqueArrayView build_array_view;
- nanoarrow::UniqueArrayView probe_array_view;
runtime_data_t* rdata;
};
using private_data_t = GpuSpatialWrapper<Payload>;
@@ -293,16 +292,14 @@ struct GpuSpatialRefinerExporter {
auto refiner = gpuspatial::CreateRTSpatialRefiner(refiner_config);
out->clear = &CClear;
- out->init_schema = &CInitSchema;
+ out->init_build_schema = &CInitBuildSchema;
out->push_build = &CPushBuild;
out->finish_building = &CFinishBuilding;
out->refine = &CRefine;
out->get_last_error = &CGetLastError;
out->release = &CRelease;
- out->private_data =
- new private_data_t{Payload{std::move(refiner),
nanoarrow::UniqueArrayView(),
- nanoarrow::UniqueArrayView(), rdata},
- ""};
+ out->private_data = new private_data_t{
+ Payload{std::move(refiner), nanoarrow::UniqueArrayView(), rdata}, ""};
}
static int CClear(SedonaSpatialRefiner* self) {
@@ -310,8 +307,8 @@ struct GpuSpatialRefinerExporter {
[&] { use_refiner(self).Clear(); });
}
- static int CInitSchema(SedonaSpatialRefiner* self, const ArrowSchema*
build_schema,
- const ArrowSchema* probe_schema) {
+ static int CInitBuildSchema(SedonaSpatialRefiner* self,
+ const ArrowSchema* build_schema) {
return SafeExecute(static_cast<private_data_t*>(self->private_data), [&] {
auto* private_data = static_cast<private_data_t*>(self->private_data);
ArrowError arrow_error;
@@ -320,11 +317,6 @@ struct GpuSpatialRefinerExporter {
throw std::runtime_error("ArrowArrayViewInitFromSchema error " +
std::string(arrow_error.message));
}
- if
(ArrowArrayViewInitFromSchema(private_data->payload.probe_array_view.get(),
- probe_schema, &arrow_error) !=
NANOARROW_OK) {
- throw std::runtime_error("ArrowArrayViewInitFromSchema error " +
- std::string(arrow_error.message));
- }
});
}
@@ -348,23 +340,29 @@ struct GpuSpatialRefinerExporter {
[&] { use_refiner(self).FinishBuilding(); });
}
- static int CRefine(SedonaSpatialRefiner* self, const ArrowArray* probe_array,
+ static int CRefine(SedonaSpatialRefiner* self, const ArrowSchema*
probe_schema,
+ const ArrowArray* probe_array,
SedonaSpatialRelationPredicate predicate, uint32_t*
build_indices,
uint32_t* probe_indices, uint32_t indices_size,
uint32_t* new_indices_size) {
return SafeExecute(static_cast<private_data_t*>(self->private_data), [&] {
- auto* private_data = static_cast<private_data_t*>(self->private_data);
- auto* array_view = private_data->payload.build_array_view.get();
+ // We need to create a local ArrayView to make sure this method is
thread-safe
ArrowError arrow_error;
-
- if (ArrowArrayViewSetArray(array_view, probe_array, &arrow_error) !=
NANOARROW_OK) {
+ nanoarrow::UniqueArrayView probe_array_view;
+ if (ArrowArrayViewInitFromSchema(probe_array_view.get(), probe_schema,
+ &arrow_error) != NANOARROW_OK) {
+ throw std::runtime_error("ArrowArrayViewInitFromSchema error " +
+ std::string(arrow_error.message));
+ }
+ if (ArrowArrayViewSetArray(probe_array_view.get(), probe_array,
&arrow_error) !=
+ NANOARROW_OK) {
throw std::runtime_error("ArrowArrayViewSetArray error " +
std::string(arrow_error.message));
}
*new_indices_size = use_refiner(self).Refine(
- array_view, static_cast<gpuspatial::Predicate>(predicate),
build_indices,
- probe_indices, indices_size);
+ probe_array_view.get(),
static_cast<gpuspatial::Predicate>(predicate),
+ build_indices, probe_indices, indices_size);
});
}
diff --git a/c/sedona-libgpuspatial/libgpuspatial/test/c_wrapper_test.cc
b/c/sedona-libgpuspatial/libgpuspatial/test/c_wrapper_test.cc
index 35f850f0..481fe9f2 100644
--- a/c/sedona-libgpuspatial/libgpuspatial/test/c_wrapper_test.cc
+++ b/c/sedona-libgpuspatial/libgpuspatial/test/c_wrapper_test.cc
@@ -122,7 +122,7 @@ class CWrapperTest : public ::testing::Test {
runtime_config.ptx_root = ptx_root.c_str();
runtime_config.device_id = 0;
- runtime_config.use_cuda_memory_pool = true;
+ runtime_config.use_cuda_memory_pool = false;
runtime_config.cuda_memory_pool_init_precent = 10;
ASSERT_EQ(runtime_.init(&runtime_, &runtime_config), 0);
@@ -242,8 +242,6 @@ TEST_F(CWrapperTest, InitializeJoiner) {
NANOARROW_OK)
<< error.message;
- refiner_.init_schema(&refiner_, build_schema.get(), probe_schema.get());
-
for (int64_t j = 0; j < probe_array->length; j++) {
ArrowBufferView wkb = ArrowArrayViewGetBytesUnsafe(probe_view.get(), j);
auto geom = wkb_reader.read(wkb.data.as_uint8, wkb.size_bytes);
@@ -282,13 +280,17 @@ TEST_F(CWrapperTest, InitializeJoiner) {
},
&intersection_ids);
+ if (i == 0) {
+ ASSERT_EQ(refiner_.init_build_schema(&refiner_, build_schema.get()), 0);
+ }
+
refiner_.clear(&refiner_);
ASSERT_EQ(refiner_.push_build(&refiner_, build_array.get()), 0);
ASSERT_EQ(refiner_.finish_building(&refiner_), 0);
uint32_t new_len;
ASSERT_EQ(refiner_.refine(
- &refiner_, probe_array.get(),
+ &refiner_, probe_schema.get(), probe_array.get(),
SedonaSpatialRelationPredicate::SedonaSpatialPredicateContains,
intersection_ids.build_indices_ptr,
intersection_ids.probe_indices_ptr,
intersection_ids.length, &new_len),
diff --git a/c/sedona-libgpuspatial/src/lib.rs
b/c/sedona-libgpuspatial/src/lib.rs
index 5dc05a60..437a35e8 100644
--- a/c/sedona-libgpuspatial/src/lib.rs
+++ b/c/sedona-libgpuspatial/src/lib.rs
@@ -151,10 +151,9 @@ mod sys {
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)
+ /// Initializes the schema for the refiner based on the data types of
the build geometries.
+ pub fn init_build_schema(&mut self, build: &DataType) -> Result<()> {
+ self.inner.init_build_schema(build)
}
/// Clears any previously inserted data from the refiner, allowing it
to be reused for building a new set of geometries.
@@ -207,7 +206,7 @@ mod sys {
pub fn push_build(&mut self, _r: &[Rect<f32>]) -> Result<()> {
Err(GpuSpatialError::GpuNotAvailable)
}
- pub fn finish_building(self) -> Result<GpuSpatialIndex> {
+ pub fn finish_building(&mut self) -> Result<GpuSpatialIndex> {
Err(GpuSpatialError::GpuNotAvailable)
}
pub fn probe(&self, _r: &[Rect<f32>]) -> Result<(Vec<u32>, Vec<u32>)> {
@@ -219,9 +218,10 @@ mod sys {
pub fn try_new(_opts: &GpuSpatialOptions) -> Result<Self> {
Err(GpuSpatialError::GpuNotAvailable)
}
- pub fn init_schema(&mut self, _b: &DataType, _p: &DataType) ->
Result<()> {
+ pub fn init_build_schema(&mut self, _b: &DataType) -> Result<()> {
Err(GpuSpatialError::GpuNotAvailable)
}
+
pub fn clear(&mut self) {}
pub fn push_build(&mut self, _arr: &arrow_array::ArrayRef) ->
Result<()> {
Err(GpuSpatialError::GpuNotAvailable)
@@ -319,9 +319,7 @@ mod tests {
let points = create_array_storage(point_values, &WKB_GEOMETRY);
// 2. Build Refiner
- refiner
- .init_schema(polygons.data_type(), points.data_type())
- .unwrap();
+ refiner.init_build_schema(polygons.data_type()).unwrap();
refiner.push_build(&polygons).unwrap();
diff --git a/c/sedona-libgpuspatial/src/libgpuspatial.rs
b/c/sedona-libgpuspatial/src/libgpuspatial.rs
index 8a53b61d..f1cb2e92 100644
--- a/c/sedona-libgpuspatial/src/libgpuspatial.rs
+++ b/c/sedona-libgpuspatial/src/libgpuspatial.rs
@@ -346,7 +346,7 @@ impl Refiner {
) -> Result<Self, GpuSpatialError> {
let mut refiner = SedonaSpatialRefiner {
clear: None,
- init_schema: None,
+ init_build_schema: None,
push_build: None,
finish_building: None,
refine: None,
@@ -377,26 +377,19 @@ impl Refiner {
/// 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> {
+ pub fn init_build_schema(&mut self, build_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 init_fn = self
+ .inner
+ .refiner
+ .init_build_schema
+ .expect("init_build_schema function is None");
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 _,
- )
- },
+ || init_fn(&mut self.inner.refiner, &build_ffi as *const _ as
*const _),
get_last_error,
refiner_ptr,
GpuSpatialError::Init,
@@ -464,7 +457,7 @@ impl Refiner {
build_indices: &mut Vec<u32>,
probe_indices: &mut Vec<u32>,
) -> Result<(), GpuSpatialError> {
- let (ffi_array, _) = arrow_array::ffi::to_ffi(&array.to_data())?;
+ let (ffi_array, ffi_schema) =
arrow_array::ffi::to_ffi(&array.to_data())?;
let refine_fn = self.inner.refiner.refine.unwrap();
let mut new_len: u32 = 0;
@@ -473,6 +466,7 @@ impl Refiner {
|| {
refine_fn(
&self.inner.refiner as *const _ as *mut _,
+ &ffi_schema as *const _ as *mut _,
&ffi_array as *const _ as *mut _,
predicate.as_c_uint(),
build_indices.as_mut_ptr(),
