Kontinuation commented on code in PR #615: URL: https://github.com/apache/sedona-db/pull/615#discussion_r2891150613
########## rust/sedona-raster-functions/src/rs_spatial_predicates.rs: ########## @@ -0,0 +1,615 @@ +// 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. + +//! RS_Intersects, RS_Contains, RS_Within functions +//! +//! These functions test spatial relationships between rasters and geometries. +//! Each function supports three overloads: (raster, geometry), (geometry, raster), +//! and (raster, raster). Rasters are compared via their convex hulls. +//! +//! CRS transformation rules: +//! - If a raster or geometry does not have a defined SRID, it is assumed to be in WGS84 +//! - If both sides are in the same CRS, perform the relationship test directly +//! - For raster/geometry pairs, the geometry is transformed into the raster's CRS +//! - For raster/raster pairs, the second raster is transformed into the first's CRS +//! - If the preferred transformation fails, both sides are transformed to WGS84 as a fallback + +use std::sync::Arc; + +use crate::crs_utils::crs_transform_wkb; +use crate::crs_utils::default_crs; +use crate::crs_utils::resolve_crs; +use crate::executor::RasterExecutor; +use arrow_array::builder::BooleanBuilder; +use arrow_schema::DataType; +use datafusion_common::exec_datafusion_err; +use datafusion_common::DataFusionError; +use datafusion_common::Result; +use datafusion_expr::{ColumnarValue, Volatility}; +use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; +use sedona_geometry::wkb_factory::write_wkb_polygon; +use sedona_raster::affine_transformation::to_world_coordinate; +use sedona_raster::traits::RasterRef; +use sedona_schema::{datatypes::SedonaType, matchers::ArgMatcher}; +use sedona_tg::tg; + +/// RS_Intersects() scalar UDF +/// +/// Returns true if the extents of the two arguments intersect. Supports +/// (raster, geometry), (geometry, raster), and (raster, raster) overloads. +/// Rasters are compared via their convex hulls. +pub fn rs_intersects_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_intersects", + vec![ + Arc::new(RsSpatialPredicate::<tg::Intersects>::raster_geom()), + Arc::new(RsSpatialPredicate::<tg::Intersects>::geom_raster()), + Arc::new(RsSpatialPredicate::<tg::Intersects>::raster_raster()), + ], + Volatility::Immutable, + ) +} + +/// RS_Contains() scalar UDF +/// +/// Returns true if the first argument's extent completely contains the second. +/// Supports (raster, geometry), (geometry, raster), and (raster, raster) overloads. +/// Rasters are compared via their convex hulls. +pub fn rs_contains_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_contains", + vec![ + Arc::new(RsSpatialPredicate::<tg::Contains>::raster_geom()), + Arc::new(RsSpatialPredicate::<tg::Contains>::geom_raster()), + Arc::new(RsSpatialPredicate::<tg::Contains>::raster_raster()), + ], + Volatility::Immutable, + ) +} + +/// RS_Within() scalar UDF +/// +/// Returns true if the first argument's extent is completely within the second. +/// Supports (raster, geometry), (geometry, raster), and (raster, raster) overloads. +/// Rasters are compared via their convex hulls. +pub fn rs_within_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_within", + vec![ + Arc::new(RsSpatialPredicate::<tg::Within>::raster_geom()), + Arc::new(RsSpatialPredicate::<tg::Within>::geom_raster()), + Arc::new(RsSpatialPredicate::<tg::Within>::raster_raster()), + ], + Volatility::Immutable, + ) +} + +/// Argument order for the spatial predicate +#[derive(Debug, Clone, Copy)] +enum ArgOrder { + /// First arg is raster, second is geometry + RasterGeom, + /// First arg is geometry, second is raster + GeomRaster, + /// Both args are rasters + RasterRaster, +} + +#[derive(Debug)] +struct RsSpatialPredicate<Op: tg::BinaryPredicate> { + arg_order: ArgOrder, + _op: std::marker::PhantomData<Op>, +} + +impl<Op: tg::BinaryPredicate> RsSpatialPredicate<Op> { + fn raster_geom() -> Self { + Self { + arg_order: ArgOrder::RasterGeom, + _op: std::marker::PhantomData, + } + } + + fn geom_raster() -> Self { + Self { + arg_order: ArgOrder::GeomRaster, + _op: std::marker::PhantomData, + } + } + + fn raster_raster() -> Self { + Self { + arg_order: ArgOrder::RasterRaster, + _op: std::marker::PhantomData, + } + } +} + +impl<Op: tg::BinaryPredicate + Send + Sync> SedonaScalarKernel for RsSpatialPredicate<Op> { + fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> { + let matcher = match self.arg_order { + ArgOrder::RasterGeom => ArgMatcher::new( + vec![ArgMatcher::is_raster(), ArgMatcher::is_geometry()], + SedonaType::Arrow(DataType::Boolean), + ), + ArgOrder::GeomRaster => ArgMatcher::new( + vec![ArgMatcher::is_geometry(), ArgMatcher::is_raster()], + SedonaType::Arrow(DataType::Boolean), + ), + ArgOrder::RasterRaster => ArgMatcher::new( + vec![ArgMatcher::is_raster(), ArgMatcher::is_raster()], + SedonaType::Arrow(DataType::Boolean), + ), + }; + + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result<ColumnarValue> { + match self.arg_order { + ArgOrder::RasterGeom => self.invoke_raster_geom(arg_types, args), + ArgOrder::GeomRaster => self.invoke_geom_raster(arg_types, args), + ArgOrder::RasterRaster => self.invoke_raster_raster(arg_types, args), + } + } +} + +impl<Op: tg::BinaryPredicate + Send + Sync> RsSpatialPredicate<Op> { + /// Invoke RS_<Predicate>(raster, geometry) + fn invoke_raster_geom( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result<ColumnarValue> { + // Ensure executor always sees (raster, geom) + let exec_arg_types = vec![arg_types[0].clone(), arg_types[1].clone()]; + let exec_args = vec![args[0].clone(), args[1].clone()]; + let executor = RasterExecutor::new(&exec_arg_types, &exec_args); + let mut builder = BooleanBuilder::with_capacity(executor.num_iterations()); + let mut raster_wkb = Vec::with_capacity(CONVEXHULL_WKB_SIZE); + + executor.execute_raster_wkb_crs_void(|raster_opt, maybe_wkb, maybe_geom_crs| { + match (raster_opt, maybe_wkb) { + (Some(raster), Some(geom_wkb)) => { + raster_wkb.clear(); + write_convexhull_wkb(raster, &mut raster_wkb)?; Review Comment: The raster-is-a-scalar case should be rare in practice, leaving this case unoptimized to simplify the code should be a sane trade off. -- 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]
