paleolimbot commented on code in PR #1068: URL: https://github.com/apache/sedona-db/pull/1068#discussion_r3640820615
########## python/sedonadb/tests/functions/test_rs_reprojectmatch.py: ########## @@ -0,0 +1,299 @@ +# 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_ReprojectMatch cross-checked against a rasterio reference implementation. + +Each test writes an input raster and a reference raster (numpy array + GDAL +geotransform + CRS) and reprojects the input onto the reference's grid through +both raster engines from `sedonadb.raster_testing`: + +- **Same-CRS regrid** onto a finer/coarser reference shares GDAL's warp with + rasterio's `reproject`. Nearest-neighbour is integer selection, so pixels are + compared exactly; bilinear accumulates over neighbours so it uses a small + tolerance (rasterio may bundle a different GDAL build). +- **Cross-CRS** (EPSG:4326 -> EPSG:3857) reprojects onto the reference grid; + both engines wrap the same GDAL warp, so nearest-neighbour pixels match + exactly. +- Cells the reprojected input does not cover fill with the input band nodata. + +Both rasters travel as table columns (not literals) so the kernel runs its real +array path. +""" + +import numpy as np +import pyarrow as pa +import pytest + +from sedonadb.raster_testing import ( + Rasterio, + SedonaDB, + decode_raster, + random_raster_data, + write_geotiff, +) + +DTYPES = ["uint8", "uint16", "int16", "int32", "float32", "float64"] Review Comment: Are `int64` and `uint64` supported for warping? ########## python/sedonadb/tests/functions/test_rs_reprojectmatch.py: ########## @@ -0,0 +1,299 @@ +# 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_ReprojectMatch cross-checked against a rasterio reference implementation. + +Each test writes an input raster and a reference raster (numpy array + GDAL +geotransform + CRS) and reprojects the input onto the reference's grid through +both raster engines from `sedonadb.raster_testing`: + +- **Same-CRS regrid** onto a finer/coarser reference shares GDAL's warp with + rasterio's `reproject`. Nearest-neighbour is integer selection, so pixels are + compared exactly; bilinear accumulates over neighbours so it uses a small + tolerance (rasterio may bundle a different GDAL build). +- **Cross-CRS** (EPSG:4326 -> EPSG:3857) reprojects onto the reference grid; + both engines wrap the same GDAL warp, so nearest-neighbour pixels match + exactly. +- Cells the reprojected input does not cover fill with the input band nodata. + +Both rasters travel as table columns (not literals) so the kernel runs its real +array path. +""" + +import numpy as np +import pyarrow as pa +import pytest + +from sedonadb.raster_testing import ( + Rasterio, + SedonaDB, + decode_raster, + random_raster_data, + write_geotiff, +) + +DTYPES = ["uint8", "uint16", "int16", "int32", "float32", "float64"] + + [email protected]() +def sedona(con): + return SedonaDB(con) + + [email protected]() +def reference(): + return Rasterio.create_or_skip() Review Comment: Can we inline these instead of use magic fixtures? ########## python/sedonadb/tests/functions/test_rs_reprojectmatch.py: ########## @@ -0,0 +1,299 @@ +# 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_ReprojectMatch cross-checked against a rasterio reference implementation. + +Each test writes an input raster and a reference raster (numpy array + GDAL +geotransform + CRS) and reprojects the input onto the reference's grid through +both raster engines from `sedonadb.raster_testing`: + +- **Same-CRS regrid** onto a finer/coarser reference shares GDAL's warp with + rasterio's `reproject`. Nearest-neighbour is integer selection, so pixels are + compared exactly; bilinear accumulates over neighbours so it uses a small + tolerance (rasterio may bundle a different GDAL build). +- **Cross-CRS** (EPSG:4326 -> EPSG:3857) reprojects onto the reference grid; + both engines wrap the same GDAL warp, so nearest-neighbour pixels match + exactly. +- Cells the reprojected input does not cover fill with the input band nodata. + +Both rasters travel as table columns (not literals) so the kernel runs its real +array path. +""" + +import numpy as np +import pyarrow as pa +import pytest + +from sedonadb.raster_testing import ( + Rasterio, + SedonaDB, + decode_raster, + random_raster_data, + write_geotiff, +) + +DTYPES = ["uint8", "uint16", "int16", "int32", "float32", "float64"] + + [email protected]() +def sedona(con): + return SedonaDB(con) + + [email protected]() +def reference(): + return Rasterio.create_or_skip() + + +def _write( + tmp_path, + dtype, + *, + bands, + height, + width, + nodata=None, + transform, + crs, + name, +): + path = tmp_path / f"reprojectmatch_{name}_{dtype}_{width}x{height}.tif" + data = random_raster_data(dtype, bands=bands, height=height, width=width) + write_geotiff(path, data, gdal_transform=transform, nodata=nodata, crs=crs) + return path + + +def _write_grid(tmp_path, *, transform, width, height, crs, name): + """Write a zeroed reference raster whose only role is to define a grid.""" + data = np.zeros((1, height, width), dtype="uint8") + path = tmp_path / f"reprojectmatch_{name}_{width}x{height}.tif" + write_geotiff(path, data, gdal_transform=transform, crs=crs) + return path + + +def _assert_transform_and_nodata(got, expected): + assert got.gdal_transform == pytest.approx( + expected.gdal_transform, rel=1e-12, abs=1e-12 + ) + assert got.nodata == expected.nodata Review Comment: Are these utilities that already exist or should be added to raster_testing so they can be reused? ########## c/sedona-gdal/src/warp.rs: ########## @@ -0,0 +1,241 @@ +// 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. + +//! GDAL warp / reproject API wrappers. +//! +//! Sits alongside the RasterIO resampling in [`crate::raster::rasterband`]: this +//! is the (re)gridding path that can move the output origin off the source grid, +//! grow the output extent beyond the source footprint, and change the CRS — +//! things RasterIO's read-into-a-buffer cannot do. + +use std::ffi::CString; +use std::os::raw::c_void; +use std::ptr::{null, null_mut}; + +use crate::dataset::Dataset; +use crate::errors::{GdalError, Result}; +use crate::gdal_api::{call_gdal_api, GdalApi}; +use crate::gdal_dyn_bindgen::{CE_Failure, CE_None}; +use crate::geo_transform::GeoTransform; +use crate::raster::types::ResampleAlg; +use crate::spatial_ref::SpatialRef; + +/// Reproject/warp `src` into the already-created `dst` dataset. +/// +/// Both datasets carry their own spatial reference (set via `set_projection` / +/// `set_spatial_ref`), so the warp reprojects from the source SRS to the +/// destination SRS, resampling with `alg`. The output grid — origin, pixel size, +/// dimensions — is whatever `dst` was created with. +/// +/// Destination cells that the (reprojected) source footprint does not cover are +/// left **untouched**: `GDALReprojectImage` with no warp options writes only +/// covered pixels. Callers that grow the extent must therefore pre-fill `dst`'s +/// band buffers with the desired background/nodata value before warping. +pub fn reproject_image( + api: &'static GdalApi, + src: &Dataset, + dst: &Dataset, + alg: ResampleAlg, +) -> Result<()> { + let gra = alg.to_gdal_warp().ok_or_else(|| { + GdalError::BadArgument(format!( + "resample algorithm {alg:?} is not supported by the warp API" + )) + })?; + + let rv = unsafe { + call_gdal_api!( + api, + GDALReprojectImage, + src.c_dataset(), + null(), // src WKT: NULL means "use the source dataset's own SRS" + dst.c_dataset(), + null(), // dst WKT: NULL means "use the destination dataset's own SRS" + gra, + 0.0, // warp memory limit (0 = GDAL default) Review Comment: I believe there are `GdalOptions` as a member of the ConfigOptions (`sedona.gdal.xxx`)...exposing this value there seems like a good idea for now. Setting it automatically would be tricky (DataFusion does not have great memory accounting, and especially does not have any accounting for the memory used by scalar functions). ########## rust/sedona-raster-gdal/src/rs_reproject_match.rs: ########## @@ -0,0 +1,538 @@ +// 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_ReprojectMatch UDF - Reproject a raster onto a reference raster's grid. +//! +//! Reprojects the input raster onto the reference raster's CRS, pixel grid, and +//! envelope: the output always has the *same* extent, resolution, dimensions, +//! and CRS as the reference (in the spirit of `rioxarray`'s `reproject_match`). +//! The input's band count/order, per-band data type, and nodata are preserved. +//! Pixel values are recomputed by GDAL's warp (`GDALReprojectImage`); output +//! cells the reprojected input footprint does not cover are filled with nodata. +//! +//! The reference raster contributes only its grid — transform, dimensions, and +//! CRS — never its pixels. + +use std::sync::Arc; + +use arrow_array::ArrayRef; +use arrow_schema::DataType; +use datafusion_common::cast::as_string_array; +use datafusion_common::config::ConfigOptions; +use datafusion_common::error::Result; +use datafusion_common::{exec_err, ScalarValue}; +use datafusion_expr::{ColumnarValue, Volatility}; + +use sedona_common::sedona_internal_err; +use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; +use sedona_gdal::geo_transform::GeoTransform; +use sedona_gdal::raster::types::ResampleAlg; +use sedona_raster::array::RasterRefImpl; +use sedona_raster::builder::RasterBuilder; +use sedona_raster::traits::RasterRef; +use sedona_raster_functions::rs_ensure_loaded::{ + NEEDS_PIXELS_METADATA_KEY, RETURNS_BYTES_METADATA_KEY, +}; +use sedona_raster_functions::RasterExecutor; +use sedona_schema::datatypes::{SedonaType, RASTER}; +use sedona_schema::matchers::ArgMatcher; + +use crate::gdal_common::{raster_ref_to_gdal_mem, with_gdal, GdalBandLayout}; +use crate::gdal_dataset_provider::configure_thread_local_options; +use crate::utils::{append_warped_nd_from_dataset, WarpGrid}; + +/// RS_ReprojectMatch() scalar UDF implementation. +/// +/// Reprojects `raster` onto `reference`'s CRS + grid + envelope. +/// +/// Signatures (matching Apache Sedona (Spark)): +/// - `RS_ReprojectMatch(raster, reference)` — 2 args (algorithm defaults to +/// `NearestNeighbor`) +/// - `RS_ReprojectMatch(raster, reference, algorithm)` — 3 args +pub fn rs_reproject_match_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_reprojectmatch", + vec![ + Arc::new(RsReprojectMatch { arg_count: 2 }), + Arc::new(RsReprojectMatch { arg_count: 3 }), + ], + Volatility::Immutable, + ) + // Reads band pixels (so the planner materializes OutDb rasters via + // RS_EnsureLoaded first) and emits a fresh InDb raster (so its output is + // already loaded and isn't wrapped again). + .with_metadata(NEEDS_PIXELS_METADATA_KEY, "true") + .with_metadata(RETURNS_BYTES_METADATA_KEY, "true") +} + +/// Kernel implementation for RS_ReprojectMatch. +#[derive(Debug)] +struct RsReprojectMatch { + /// Number of arguments in the matched signature (2 or 3). + arg_count: usize, +} + +impl SedonaScalarKernel for RsReprojectMatch { + fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> { + let matchers = match self.arg_count { + 2 => vec![ArgMatcher::is_raster(), ArgMatcher::is_raster()], + 3 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_raster(), + ArgMatcher::is_string(), + ], + _ => { + return sedona_internal_err!( + "RS_ReprojectMatch: unexpected arg_count {}", + self.arg_count + ); + } + }; + ArgMatcher::new(matchers, RASTER).match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result<ColumnarValue> { + self.invoke_batch_from_args(arg_types, args, &SedonaType::Arrow(DataType::Null), 0, None) + } + + fn invoke_batch_from_args( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + _return_type: &SedonaType, + _num_rows: usize, + config_options: Option<&ConfigOptions>, + ) -> Result<ColumnarValue> { + let num_iterations = RasterExecutor::num_iterations_over(args); + + // Algorithm string at index 2 (when arg_count == 3); otherwise the + // Spark default `NearestNeighbor`. Expand to an array so a per-row + // column and a scalar are handled identically. + let algorithm_array = if self.arg_count >= 3 { + args[2] + .clone() + .cast_to(&DataType::Utf8, None)? Review Comment: Using a utf8view here should be slightly better for the scalar-to-array case -- 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]
