pwrliang commented on code in PR #556:
URL: https://github.com/apache/sedona-db/pull/556#discussion_r2764313596
##########
c/sedona-libgpuspatial/src/libgpuspatial.rs:
##########
@@ -331,179 +327,387 @@ impl GpuSpatialJoinerWrapper {
&[]
}
- pub fn get_stream_indices_buffer(&self, ctx: &mut GpuSpatialJoinerContext)
-> &[u32] {
- if let Some(get_stream_indices_buffer_fn) =
self.joiner.get_stream_indices_buffer {
- let mut stream_indices_ptr: *mut c_void = std::ptr::null_mut();
- let mut stream_indices_len: u32 = 0;
+ pub fn get_probe_indices_buffer(&self, ctx: &mut
SedonaSpatialIndexContext) -> &[u32] {
+ if let Some(get_probe_indices_buffer_fn) =
self.index.get_probe_indices_buffer {
+ let mut probe_indices_ptr: *mut u32 = std::ptr::null_mut();
+ let mut probe_indices_len: u32 = 0;
unsafe {
- get_stream_indices_buffer_fn(
+ get_probe_indices_buffer_fn(
ctx as *mut _,
- &mut stream_indices_ptr as *mut *mut c_void,
- &mut stream_indices_len as *mut u32,
+ &mut probe_indices_ptr as *mut *mut u32,
+ &mut probe_indices_len as *mut u32,
);
// Check length first - empty vectors return empty slice
- if stream_indices_len == 0 {
+ if probe_indices_len == 0 {
return &[];
}
// Validate pointer (should not be null if length > 0)
- if stream_indices_ptr.is_null() {
+ if probe_indices_ptr.is_null() {
return &[];
}
// Convert the raw pointer to a slice. This is safe to do
because
// we've validated the pointer is non-null and length is valid.
- let typed_ptr = stream_indices_ptr as *const u32;
+ let typed_ptr = probe_indices_ptr as *const u32;
// Safety: We've checked ptr is non-null and len > 0
- return std::slice::from_raw_parts(typed_ptr,
stream_indices_len as usize);
+ return std::slice::from_raw_parts(typed_ptr, probe_indices_len
as usize);
}
}
&[]
}
+}
- pub fn release(&mut self) {
- // Call the release function if it exists
- if let Some(release_fn) = self.joiner.release {
- unsafe {
- release_fn(&mut self.joiner as *mut _);
- }
+impl Default for GpuSpatialIndexFloat2DWrapper {
+ fn default() -> Self {
+ GpuSpatialIndexFloat2DWrapper {
+ 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(),
+ },
+ _runtime:
Arc::new(Mutex::new(GpuSpatialRuntimeWrapper::default())),
}
}
}
-impl Drop for GpuSpatialJoinerWrapper {
+impl Drop for GpuSpatialIndexFloat2DWrapper {
fn drop(&mut self) {
// Call the release function if it exists
- if let Some(release_fn) = self.joiner.release {
+ if let Some(release_fn) = self.index.release {
unsafe {
- release_fn(&mut self.joiner as *mut _);
+ release_fn(&mut self.index as *mut _);
}
}
}
}
-#[cfg(test)]
-mod test {
- use super::*;
- use sedona_expr::scalar_udf::SedonaScalarUDF;
- use sedona_geos::register::scalar_kernels;
- use sedona_schema::crs::lnglat;
- use sedona_schema::datatypes::{Edges, SedonaType, WKB_GEOMETRY};
- use sedona_testing::create::create_array_storage;
- use sedona_testing::testers::ScalarUdfTester;
- use std::env;
- use std::path::PathBuf;
-
- #[test]
- fn test_gpu_joiner_end2end() {
- let mut joiner = GpuSpatialJoinerWrapper::new();
-
- let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());
- let ptx_root = out_path.join("share/gpuspatial/shaders");
-
- joiner
- .init(
- 1,
- ptx_root.to_str().expect("Failed to convert path to string"),
- )
- .expect("Failed to init GpuSpatialJoiner");
-
- 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))"),
- Some("POLYGON ((0 0, 10 0, 10 10, 0 10, 0 0), (2 2, 3 2, 3 3, 2 3,
2 2), (6 6, 8 6, 8 8, 6 8, 6 6))"),
- Some("POLYGON ((30 0, 60 20, 50 50, 10 50, 0 20, 30 0), (20 30, 25
40, 15 40, 20 30), (30 30, 35 40, 25 40, 30 30), (40 30, 45 40, 35 40, 40
30))"),
- Some("POLYGON ((40 0, 50 30, 80 20, 90 70, 60 90, 30 80, 20 40, 40
0), (50 20, 65 30, 60 50, 45 40, 50 20), (30 60, 50 70, 45 80, 30 60))"),
- ];
- let polygons = create_array_storage(polygon_values, &WKB_GEOMETRY);
-
- // Let the gpusaptial joiner to parse WKBs and get building boxes
- joiner
- .push_build(&polygons, 0, polygons.len().try_into().unwrap())
- .expect("Failed to push building");
- // Build a spatial index for Build internally on GPU
- joiner.finish_building().expect("Failed to finish building");
-
- // Each thread that performs spatial joins should have its own context.
- // The context is passed to PushStream calls to perform spatial joins.
- let mut ctx = GpuSpatialJoinerContext {
- last_error: std::ptr::null(),
+#[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,
+}
Review Comment:
This one maps the C library enum to a Rust enum. Where's the second one?
--
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]
