paleolimbot commented on code in PR #615: URL: https://github.com/apache/sedona-db/pull/615#discussion_r2890664875
########## docs/reference/sql/rs_contains.qmd: ########## @@ -0,0 +1,57 @@ +--- +# 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. + +title: RS_Contains +description: Returns true if the first argument's extent contains the second. +kernels: + - returns: boolean + args: [raster, geometry] + - returns: boolean + args: [geometry, raster] + - returns: boolean + args: + - name: rastA + type: raster + description: Input raster + - name: rastB + type: raster + description: Input raster +--- + +## Description + +Returns `true` if the convex hull of the first argument completely contains +the second argument. Both rasters and geometries are accepted in either +argument position. When two rasters are provided, their convex hulls are +compared. + +If the arguments have different CRSes, the geometry is transformed into the +raster's CRS before evaluating the predicate. For two rasters, the second +raster is transformed into the first raster's CRS. If the preferred +transformation fails, both sides are transformed to WGS 84 as a fallback. Review Comment: ```suggestion transformation fails, the extent of both sides are transformed to WGS 84 as a fallback. ``` ########## rust/sedona-raster-functions/src/crs_utils.rs: ########## @@ -0,0 +1,428 @@ +// 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. + +use std::sync::{Arc, LazyLock}; + +use datafusion_common::{exec_datafusion_err, DataFusionError, Result}; +use sedona_geometry::transform::{transform, CrsEngine}; +use sedona_proj::transform::with_global_proj_engine; +use sedona_schema::crs::{deserialize_crs, lnglat, CoordinateReferenceSystem}; +use wkb::reader::read_wkb; + +/// Cached default CRS (WGS84 longitude/latitude). Initialized once on first access. +static DEFAULT_CRS: LazyLock<Arc<dyn CoordinateReferenceSystem + Send + Sync>> = + LazyLock::new(|| lnglat().expect("lnglat() should always succeed")); Review Comment: I commented above too, but I don't think this is a good place to apply a default CRS. In SedonaDB it is practically quite difficult to get something without a CRS (when there's a raster reader, the reader would be a good place to apply that). ########## 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)?; + + let result = evaluate_predicate_with_crs::<Op>( + &raster_wkb, + raster.crs(), + geom_wkb, + maybe_geom_crs, + false, + )?; + builder.append_value(result); + } + _ => builder.append_null(), + } + Ok(()) + })?; + + executor.finish(Arc::new(builder.finish())) + } + + /// Invoke RS_<Predicate>(geometry, raster) + fn invoke_geom_raster( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result<ColumnarValue> { + // Reorder so executor always sees (raster, geom) + let exec_arg_types = vec![arg_types[1].clone(), arg_types[0].clone()]; + let exec_args = vec![args[1].clone(), args[0].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)?; + + // Note: order is geometry, raster for the predicate + let result = evaluate_predicate_with_crs::<Op>( + geom_wkb, + maybe_geom_crs, + &raster_wkb, + raster.crs(), + true, + )?; + builder.append_value(result); + } + _ => builder.append_null(), + } + Ok(()) + })?; + + executor.finish(Arc::new(builder.finish())) + } + + /// Invoke RS_<Predicate>(raster1, raster2) + fn invoke_raster_raster( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result<ColumnarValue> { + // Ensure executor always sees (raster, raster) + 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 wkb0 = Vec::with_capacity(CONVEXHULL_WKB_SIZE); + let mut wkb1 = Vec::with_capacity(CONVEXHULL_WKB_SIZE); + + executor.execute_raster_raster_void(|_i, r0_opt, r1_opt| { + match (r0_opt, r1_opt) { + (Some(r0), Some(r1)) => { + wkb0.clear(); + wkb1.clear(); + write_convexhull_wkb(r0, &mut wkb0)?; + write_convexhull_wkb(r1, &mut wkb1)?; + + let result = + evaluate_predicate_with_crs::<Op>(&wkb0, r0.crs(), &wkb1, r1.crs(), false)?; + builder.append_value(result); + } + _ => builder.append_null(), + } + Ok(()) + })?; + + executor.finish(Arc::new(builder.finish())) + } +} + +/// Evaluate a spatial predicate with CRS handling +/// +/// Rules: +/// - If no CRS defined, assume WGS84 +/// - If both same CRS, compare directly +/// - Otherwise, try transforming one side to the other's CRS for comparison. +/// If that fails, transform both to WGS84 and compare. +fn evaluate_predicate_with_crs<Op: tg::BinaryPredicate>( + wkb_a: &[u8], + crs_a: Option<&str>, + wkb_b: &[u8], + crs_b: Option<&str>, + from_a_to_b: bool, +) -> Result<bool> { Review Comment: I expected to see `with_global_proj_engine()` outside the loop with a reference to the engine passed in here. I think it would be better to keep that in the UDF execute because it will be easier to grab the session engine from the ConfigOptions when we get there. ########## rust/sedona-raster-functions/src/executor.rs: ########## @@ -35,6 +40,177 @@ pub struct RasterExecutor<'a, 'b> { num_iterations: usize, } +#[derive(Clone)] +enum ItemWkbAccessor { + Binary(BinaryArray), + BinaryView(BinaryViewArray), +} + +impl ItemWkbAccessor { + #[inline] + fn get(&self, i: usize) -> Option<&[u8]> { + match self { + Self::Binary(arr) => { + if arr.is_null(i) { + None + } else { + Some(arr.value(i)) + } + } + Self::BinaryView(arr) => { + if arr.is_null(i) { + None + } else { + Some(arr.value(i)) + } + } + } + } +} + +enum GeomWkbCrsAccessor { + WkbArray { + wkb: ItemWkbAccessor, + static_crs: Option<String>, + }, + WkbScalar { + wkb: Option<Vec<u8>>, + static_crs: Option<String>, + }, + ItemCrsArray { + struct_array: StructArray, + item: ItemWkbAccessor, + crs: StringViewArray, + item_static_crs: Option<String>, + }, + ItemCrsScalar { + struct_array: StructArray, + item: ItemWkbAccessor, + crs: StringViewArray, + item_static_crs: Option<String>, + }, + Null, +} + +impl GeomWkbCrsAccessor { + #[inline] + fn get(&self, i: usize) -> Result<(Option<&[u8]>, Option<&str>)> { + match self { + Self::Null => Ok((None, None)), + Self::WkbArray { wkb, static_crs } => { + let maybe_wkb = wkb.get(i); + if maybe_wkb.is_none() { + return Ok((None, None)); + } + Ok((maybe_wkb, static_crs.as_deref())) + } + Self::WkbScalar { wkb, static_crs } => { + if wkb.is_none() { + return Ok((None, None)); + } + let _ = i; + Ok((wkb.as_deref(), static_crs.as_deref())) + } + Self::ItemCrsArray { + struct_array, + item, + crs, + item_static_crs, + } => { + if struct_array.is_null(i) { + return Ok((None, None)); + } + + let maybe_wkb = item.get(i); + if maybe_wkb.is_none() { + return Ok((None, None)); + } + + let item_crs_str = if crs.is_null(i) { + None + } else { + Some(crs.value(i)) + }; + let static_crs_str = item_static_crs.as_deref(); + let crs_out = resolve_item_crs(item_crs_str, static_crs_str)?; + Ok((maybe_wkb, crs_out)) + } + Self::ItemCrsScalar { + struct_array, + item, + crs, + item_static_crs, + } => { + if struct_array.is_null(0) { + return Ok((None, None)); + } + + let maybe_wkb = item.get(0); + if maybe_wkb.is_none() { + return Ok((None, None)); + } + + let item_crs_str = if crs.is_null(0) { + None + } else { + Some(crs.value(0)) + }; + let static_crs_str = item_static_crs.as_deref(); + let crs_out = resolve_item_crs(item_crs_str, static_crs_str)?; + let _ = i; + Ok((maybe_wkb, crs_out)) + } + } + } +} + +fn resolve_item_crs<'a>( + item_crs: Option<&'a str>, + static_crs: Option<&'a str>, +) -> Result<Option<&'a str>> { + match (item_crs, static_crs) { + (None, None) => Ok(None), + (Some(item), None) => Ok(Some(item)), + (None, Some(st)) => Ok(Some(st)), + (Some(item), Some(st)) => { + if item == st { + return Ok(Some(item)); + } + + let item_crs = deserialize_crs(item)?; + let static_crs = deserialize_crs(st)?; + if item_crs == static_crs { + Ok(Some(item)) + } else { + exec_err!("CRS values not equal: {item_crs:?} vs {static_crs:?}") + } + } + } +} + +fn crs_string_from_sedona_type(sedona_type: &SedonaType) -> Result<Option<String>> { + match sedona_type { + SedonaType::Wkb(_, crs) | SedonaType::WkbView(_, crs) => { + if let Some(crs) = crs { + Ok(Some( + crs.to_authority_code()?.unwrap_or_else(|| crs.to_json()), + )) Review Comment: I believe we have a CRS abbreviator already (and the executor should probably not automatically abbreviate so that functions can use the full definitions if they are provided). ########## rust/sedona-raster-functions/src/executor.rs: ########## @@ -35,6 +40,177 @@ pub struct RasterExecutor<'a, 'b> { num_iterations: usize, } +#[derive(Clone)] +enum ItemWkbAccessor { + Binary(BinaryArray), + BinaryView(BinaryViewArray), +} + +impl ItemWkbAccessor { + #[inline] + fn get(&self, i: usize) -> Option<&[u8]> { Review Comment: The pattern of switching on type for every element is probably not a great one. Possibly a `Box<dyn Iterator<Item = ...>>`? In the context of raster iteration where it is less likely there will be billions of them this is probably OK. You could also use that for raster/CRS iteration and use the existing executor to handle the arrays (it's almost always free to cast the CRS array to string view and there is less variation in the raster type than the geometry type). ########## rust/sedona-raster-functions/src/crs_utils.rs: ########## @@ -0,0 +1,428 @@ +// 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. + +use std::sync::{Arc, LazyLock}; + +use datafusion_common::{exec_datafusion_err, DataFusionError, Result}; +use sedona_geometry::transform::{transform, CrsEngine}; +use sedona_proj::transform::with_global_proj_engine; +use sedona_schema::crs::{deserialize_crs, lnglat, CoordinateReferenceSystem}; +use wkb::reader::read_wkb; + +/// Cached default CRS (WGS84 longitude/latitude). Initialized once on first access. +static DEFAULT_CRS: LazyLock<Arc<dyn CoordinateReferenceSystem + Send + Sync>> = + LazyLock::new(|| lnglat().expect("lnglat() should always succeed")); + +/// Run a closure with the active CRS engine, abstracting away the engine choice. +/// +/// We keep this API engine-agnostic to allow future engines beyond PROJ; the +/// current implementation uses PROJ via `with_global_proj_engine`. +pub fn with_crs_engine<R, F: FnMut(&dyn CrsEngine) -> Result<R>>(mut func: F) -> Result<R> { + with_global_proj_engine(|engine| func(engine)) +} + +/// Resolve an optional CRS string to a concrete CRS object. +/// +/// - If `crs_str` is `Some` and deserializes to a known CRS, that CRS is returned. +/// - Otherwise (None, empty, "0", etc.), the default WGS84 lnglat CRS is returned. +pub fn resolve_crs( + crs_str: Option<&str>, +) -> Result<Arc<dyn CoordinateReferenceSystem + Send + Sync>> { + if let Some(crs_str) = crs_str { + let crs = deserialize_crs(crs_str)?; + Ok(crs.unwrap_or_else(|| DEFAULT_CRS.clone())) + } else { + Ok(DEFAULT_CRS.clone()) + } +} + +/// Return a reference to the default CRS (WGS84 longitude/latitude). +/// +/// This is a zero-cost accessor backed by a `LazyLock` static — no allocation +/// or atomic ref-count increment on each call. +pub fn default_crs() -> &'static (dyn CoordinateReferenceSystem + Send + Sync) { + DEFAULT_CRS.as_ref() +} + +/// Transform a geometry encoded as WKB from one CRS to another. +/// +/// This is a utility used by raster/spatial functions to reproject a geometry +/// without leaking PROJ engine details into call sites. +/// +/// **Behavior** +/// - If `from_crs` and `to_crs` are equal, returns the original WKB (clone) without decoding. +/// - Otherwise, builds a PROJ pipeline and transforms all coordinates. +/// +/// **Errors** +/// - Returns an error if WKB parsing fails, PROJ cannot build the CRS-to-CRS transform, +/// or if the coordinate transformation itself fails. +pub fn crs_transform_wkb( + wkb: &[u8], + from_crs: &dyn CoordinateReferenceSystem, + to_crs: &dyn CoordinateReferenceSystem, +) -> Result<Vec<u8>> { Review Comment: The identical CRS check should probably happen before this to avoid the clone. For a bbox the clone is trivial but this could also be a complex geometry (e.g., division boundary) on the other side of the predicate. ########## 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: I imagine this is not a huge performance issue, but technically this is unnecessarily recomputing the raster convex hull for the raster-is-a-scalar/geometry-is-a-column case. ########## 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 Review Comment: In SedonaDB I've tried hard to avoid this assumption (I am guessing this is a hangover from Sedona not propagating CRSes from GeoParquet). We can make it easy to apply this assumption explicitly if it causes issues (e.g., `ST_ApplyDefaultCRS()` or a Python dataframe level function that does it for all spatial columns). ########## rust/sedona-raster-functions/src/crs_utils.rs: ########## @@ -0,0 +1,428 @@ +// 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. + +use std::sync::{Arc, LazyLock}; + +use datafusion_common::{exec_datafusion_err, DataFusionError, Result}; +use sedona_geometry::transform::{transform, CrsEngine}; +use sedona_proj::transform::with_global_proj_engine; +use sedona_schema::crs::{deserialize_crs, lnglat, CoordinateReferenceSystem}; +use wkb::reader::read_wkb; + +/// Cached default CRS (WGS84 longitude/latitude). Initialized once on first access. +static DEFAULT_CRS: LazyLock<Arc<dyn CoordinateReferenceSystem + Send + Sync>> = + LazyLock::new(|| lnglat().expect("lnglat() should always succeed")); + +/// Run a closure with the active CRS engine, abstracting away the engine choice. +/// +/// We keep this API engine-agnostic to allow future engines beyond PROJ; the +/// current implementation uses PROJ via `with_global_proj_engine`. +pub fn with_crs_engine<R, F: FnMut(&dyn CrsEngine) -> Result<R>>(mut func: F) -> Result<R> { + with_global_proj_engine(|engine| func(engine)) +} Review Comment: I think we can apply this at the UDF level where it's a bit more explicit (and amortized over a batch). -- 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]
