pwrliang commented on code in PR #310: URL: https://github.com/apache/sedona-db/pull/310#discussion_r2591515949
########## c/sedona-libgpuspatial/src/lib.rs: ########## @@ -0,0 +1,320 @@ +// 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. + +// Module declarations +#[cfg(gpu_available)] +pub mod error; +#[cfg(gpu_available)] +mod libgpuspatial; +#[cfg(gpu_available)] +mod libgpuspatial_glue_bindgen; + +// Import Array trait for len() method (used in gpu_available code) +#[cfg(gpu_available)] +use arrow_array::Array; + +// Re-exports for GPU functionality +#[cfg(gpu_available)] +pub use error::GpuSpatialError; +#[cfg(gpu_available)] +pub use libgpuspatial::{GpuSpatialJoinerWrapper, GpuSpatialPredicateWrapper}; +#[cfg(gpu_available)] +pub use libgpuspatial_glue_bindgen::GpuSpatialJoinerContext; + +// Mark GPU types as Send for thread safety +// SAFETY: The GPU library is designed to be used from multiple threads. +// Each thread gets its own context, and the underlying GPU library handles thread safety. +// The raw pointers inside are managed by the C++ library which ensures proper synchronization. +#[cfg(gpu_available)] +unsafe impl Send for GpuSpatialJoinerContext {} + +#[cfg(gpu_available)] +unsafe impl Send for libgpuspatial_glue_bindgen::GpuSpatialJoiner {} + +#[cfg(gpu_available)] +unsafe impl Send for GpuSpatialJoinerWrapper {} + +// Error type for non-GPU builds +#[cfg(not(gpu_available))] +#[derive(Debug, thiserror::Error)] +pub enum GpuSpatialError { + #[error("GPU not available - CUDA not found during build")] + GpuNotAvailable, +} + +pub type Result<T> = std::result::Result<T, GpuSpatialError>; + +/// High-level wrapper for GPU spatial operations +pub struct GpuSpatialContext { + #[cfg(gpu_available)] + joiner: Option<GpuSpatialJoinerWrapper>, + #[cfg(gpu_available)] + context: Option<GpuSpatialJoinerContext>, + initialized: bool, +} + +impl GpuSpatialContext { + pub fn new() -> Result<Self> { + #[cfg(not(gpu_available))] + { + Err(GpuSpatialError::GpuNotAvailable) + } + + #[cfg(gpu_available)] + { + Ok(Self { + joiner: None, + context: None, + initialized: false, + }) + } + } + + pub fn init(&mut self) -> Result<()> { + #[cfg(not(gpu_available))] + { + Err(GpuSpatialError::GpuNotAvailable) + } + + #[cfg(gpu_available)] + { + let mut joiner = GpuSpatialJoinerWrapper::new(); + + // Get PTX path from OUT_DIR + 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()))?; + + // Initialize with concurrency of 1 for now + joiner.init(1, ptx_root_str)?; + + // Create context + let mut ctx = GpuSpatialJoinerContext { + last_error: std::ptr::null(), + private_data: std::ptr::null_mut(), + build_indices: std::ptr::null_mut(), + stream_indices: std::ptr::null_mut(), + }; + joiner.create_context(&mut ctx); + + self.joiner = Some(joiner); + self.context = Some(ctx); + self.initialized = true; + Ok(()) + } + } + + #[cfg(gpu_available)] + pub fn get_joiner_mut(&mut self) -> Option<&mut GpuSpatialJoinerWrapper> { + self.joiner.as_mut() + } + + #[cfg(gpu_available)] + pub fn get_context_mut(&mut self) -> Option<&mut GpuSpatialJoinerContext> { + self.context.as_mut() + } + + pub fn is_initialized(&self) -> bool { + self.initialized + } + + /// Perform spatial join between two geometry arrays + pub fn spatial_join( + &mut self, + left_geom: arrow_array::ArrayRef, + right_geom: arrow_array::ArrayRef, + predicate: SpatialPredicate, + ) -> Result<(Vec<u32>, Vec<u32>)> { Review Comment: https://github.com/apache/sedona-db/issues/414 -- 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]
