feat(c/sedona-libgpuspatial): Change the interface of SpatialRefiner (#717)

pwrliang · web-flow · commit 82c1bb05a105 · 2026-03-18T16:57:32.000-05:00
diff --git 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
@@ -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,25 +292,23 @@ 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) {
     return SafeExecute(static_cast<private_data_t*>(self->private_data),
                        [&] { 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
@@ -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
@@ -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
@@ -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(),
