james-willis commented on code in PR #1073: URL: https://github.com/apache/sedona-db/pull/1073#discussion_r3643054272
########## rust/sedona-spatial-join-raster/src/join_provider.rs: ########## @@ -0,0 +1,474 @@ +// 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; + +use arrow_array::{builder::BinaryBuilder, Array, ArrayRef, StructArray}; +use datafusion_common::{exec_datafusion_err, JoinType, Result}; +use datafusion_expr::ColumnarValue; +use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err, SpatialJoinOptions}; +use sedona_expr::statistics::GeoStatistics; +use sedona_geometry::{transform::CrsEngine, wkb_factory::write_wkb_polygon}; +use sedona_proj::transform::with_global_proj_engine; +use sedona_raster::array::RasterStructArray; +use sedona_raster::traits::RasterRef; +use sedona_raster_functions::crs_utils::resolve_crs; +use sedona_raster_functions::footprint::{raster_footprint_corners, write_convexhull_wkb}; +use sedona_schema::{ + crs::{lnglat, CoordinateReferenceSystem, Crs}, + datatypes::SedonaType, + datatypes::WKB_GEOMETRY, +}; +use sedona_spatial_join::{ + index::{spatial_index_builder::SpatialJoinBuildMetrics, SpatialIndexBuilder}, + join_provider::{DefaultSpatialJoinProvider, SpatialJoinProvider}, + operand_evaluator::{EvaluatedGeometryArray, EvaluatedGeometryArrayFactory}, + utils::bounds::Bounds2D, + SpatialPredicate, +}; + +/// [`SpatialJoinProvider`] for raster/geometry spatial joins. +/// +/// The R-tree index builder and the memory estimate delegate to the default +/// provider; only the operand evaluator is raster-aware. The factory it produces +/// reprojects raster footprints into `target_crs` (the geometry operand's CRS). +#[derive(Debug)] +pub(crate) struct RasterJoinProvider { + default: DefaultSpatialJoinProvider, + target_crs: Crs, +} + +impl RasterJoinProvider { + pub(crate) fn new(target_crs: Crs) -> Self { + Self { + default: DefaultSpatialJoinProvider, + target_crs, + } + } +} + +impl SpatialJoinProvider for RasterJoinProvider { + fn try_new_spatial_index_builder( + &self, + schema: arrow_schema::SchemaRef, + spatial_predicate: SpatialPredicate, + options: SpatialJoinOptions, + join_type: JoinType, + probe_threads_count: usize, + metrics: SpatialJoinBuildMetrics, + ) -> Result<Box<dyn SpatialIndexBuilder>> { + // Footprints are ordinary planar WKB polygons, so the default R-tree + // builder and WKB refiner apply unchanged. + self.default.try_new_spatial_index_builder( + schema, + spatial_predicate, + options, + join_type, + probe_threads_count, + metrics, + ) + } + + fn estimate_extra_memory_usage( + &self, + geo_stats: &GeoStatistics, + spatial_predicate: &SpatialPredicate, + options: &SpatialJoinOptions, + ) -> usize { + self.default + .estimate_extra_memory_usage(geo_stats, spatial_predicate, options) + } + + fn evaluated_array_factory(&self) -> Arc<dyn EvaluatedGeometryArrayFactory> { + Arc::new(RasterGeometryArrayFactory { + target_crs: self.target_crs.clone(), + }) + } +} + +/// Evaluates the operands of a raster/geometry spatial predicate. +/// +/// The same factory sees both operands of the join. A raster operand is turned +/// into its footprint — the convex hull of the raster's four corners +/// (see [`RasterGeometryArrayFactory::evaluate_raster`]) — and a geometry operand +/// is evaluated with the default planar behavior, since the geometry is already +/// in the target CRS. +/// +/// Cross-CRS raster footprints are indexed and refined as the convex hull of the +/// raster's four reprojected corners; projection curvature along the edges is not +/// modeled, so for large-extent rasters reprojected between very different CRSs +/// the hull can slightly under-cover the true footprint (rare missed matches at +/// the extreme edges). Same-CRS joins are exact. +#[derive(Debug)] +struct RasterGeometryArrayFactory { + /// The geometry operand's CRS: the common CRS this join compares in. Raster + /// footprints are reprojected into it. WGS84 when the geometry operand carries + /// no CRS — a missing CRS is assumed WGS84 (matching the `RS_*` kernel and + /// Sedona Spark), so the CRS-less geometry is compared as-is in a WGS84 frame. + target_crs: Crs, +} + +impl EvaluatedGeometryArrayFactory for RasterGeometryArrayFactory { + fn try_new_evaluated_array( + &self, + geometry_array: ArrayRef, + sedona_type: &SedonaType, + distance_columnar_value: Option<&ColumnarValue>, + ) -> Result<EvaluatedGeometryArray> { + // Relation predicates (Intersects/Contains/Within) carry no distance; + // this factory is only wired up for those by the raster planner. + if distance_columnar_value.is_some() { + return sedona_internal_err!( + "Raster spatial joins do not support a distance predicate" + ); + } + + match sedona_type { + SedonaType::Raster => self.evaluate_raster(geometry_array), + // The geometry operand is already in the target CRS, so its default + // planar evaluation (Cartesian WKB bounds + WKB) is exactly right. + _ => EvaluatedGeometryArray::try_new(geometry_array, sedona_type), + } + } +} + +impl RasterGeometryArrayFactory { + /// Evaluate a raster struct array into footprint polygons plus bounding + /// rectangles, reprojecting each footprint into the target CRS. + fn evaluate_raster(&self, raster_array: ArrayRef) -> Result<EvaluatedGeometryArray> { + let struct_array = raster_array + .as_any() + .downcast_ref::<StructArray>() + .ok_or_else(|| { + sedona_internal_datafusion_err!("Expected StructArray for raster operand") + })?; + let rasters = RasterStructArray::try_new(struct_array) + .map_err(|e| exec_datafusion_err!("Failed to read raster array: {e}"))?; + + let num_rows = rasters.len(); + let mut builder = BinaryBuilder::with_capacity(num_rows, num_rows * 96); + let mut rects = Vec::with_capacity(num_rows); + + // A missing CRS on either side is assumed to be WGS84 (matching the `RS_*` + // kernel and Sedona Spark); resolve it once rather than per row. + let lnglat_crs = lnglat(); + let wgs84 = lnglat_crs.as_deref().expect("lnglat() is always Some"); + + with_global_proj_engine(|engine| { + for i in 0..num_rows { + // A null raster produces a null footprint that never matches. + if rasters.is_null(i) { + builder.append_null(); + rects.push(Bounds2D::empty()); + continue; + } + + let raster = rasters + .get(i) + .map_err(|e| exec_datafusion_err!("Failed to read raster row {i}: {e}"))?; + let rect = self.append_footprint(&raster, wgs84, engine, &mut builder)?; + rects.push(rect); + } + Ok(()) + })?; + + let footprint_array: ArrayRef = Arc::new(builder.finish()); + EvaluatedGeometryArray::try_new_with_rects(footprint_array, rects, &WKB_GEOMETRY) + } + + /// Append one raster's footprint WKB to `builder` and return its bounding + /// rectangle, reconciling the raster's CRS against the target CRS. `wgs84` is + /// the WGS84 CRS a missing side is assumed to be. + /// + /// CRS rules mirror the `RS_*` kernel: a missing CRS on either side is assumed + /// to be WGS84 (matching Sedona Spark), after which an equal CRS compares + /// directly and a genuine CRS difference reprojects the footprint's four + /// corners into the target CRS. The footprint is always the convex hull of the + /// four corners — projection curvature along the edges is not modeled (see the + /// type-level note on [`RasterGeometryArrayFactory`]). + fn append_footprint( + &self, + raster: &dyn RasterRef, + wgs84: &(dyn CoordinateReferenceSystem + Send + Sync), + engine: &dyn CrsEngine, + builder: &mut BinaryBuilder, + ) -> Result<Bounds2D> { + let raster_crs = resolve_crs(raster.crs())?; + let corners = raster_footprint_corners(raster); + + // Assume WGS84 for whichever side lacks a CRS, then compare in the target + // frame. Both-absent and same-CRS collapse to a direct comparison. + let raster_crs = raster_crs.as_deref().unwrap_or(wgs84); + let target_crs = self.target_crs.as_deref().unwrap_or(wgs84); Review Comment: https://github.com/apache/sedona/issues/3153 https://github.com/apache/sedona-db/issues/1075 IMO we should keep the behavior in sync. I filed matching bugs in each project -- 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]
