You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@sedona.apache.org by "Kontinuation (via GitHub)" <gi...@apache.org> on 2026/06/14 08:18:35 UTC
[PR] feat(rust/sedona-raster-gdal): add rs_asraster [sedona-db]
Kontinuation opened a new pull request, #956:
URL: https://github.com/apache/sedona-db/pull/956
## Summary
- add the GDAL-backed raster function
- normalize the SQL name to
- add SQL reference docs and a dedicated Criterion benchmark
## Dependency
- Independent branch based on
## Testing
-
-
-
running 80 tests
test gdal_common::tests::test_band_data_type_to_gdal ... ok
test gdal_common::tests::test_bytes_to_f64 ... ok
test gdal_common::tests::test_normalize_outdb_source_path_http_and_https ... ok
test gdal_common::tests::test_normalize_outdb_source_path_gcs_schemes ... ok
test gdal_common::tests::test_normalize_outdb_source_path_s3_and_s3a ... ok
test gdal_common::tests::test_normalize_outdb_source_path_azure_schemes ... ok
test gdal_common::tests::test_normalize_outdb_source_path_local_and_vsi_paths ... ok
test gdal_common::tests::test_raster_ref_to_gdal_mem_multi_nonspatial_dims_round_trip ... ok
test gdal_common::tests::test_raster_ref_to_gdal_mem_mixed_2d_and_nd_round_trips ... ok
test gdal_common::tests::test_to_raster_metadata ... ok
test gdal_common::tests::test_raster_ref_to_gdal_mem_rejects_outdb_bands ... ok
test gdal_common::tests::test_to_gdal_geotransform ... ok
test gdal_common::tests::test_raster_ref_to_gdal_mem_nd_band_stacks_and_round_trips ... ok
test gdal_dataset_provider::tests::test_outdb_source_key_normalizes_gcs_variants ... ok
test gdal_dataset_provider::tests::test_outdb_source_key_normalizes_azure_variants ... ok
test gdal_dataset_provider::tests::test_outdb_source_key_normalizes_s3_variants ... ok
test gdal_common::tests::test_raster_ref_to_gdal_mem_rejects_non_trailing_spatial_pair ... ok
test gdal_dataset_provider::tests::test_provider_rejects_nd_band_in_vrt_path ... ok
test gdal_dataset_provider::tests::test_provider_selects_outdb_band_via_band_fragment ... ok
test gdal_dataset_provider::tests::test_vrt_key_distinguishes_lossless_uint64_nodata ... ok
test gdal_dataset_provider::tests::test_vrt_key_normalizes_equivalent_outdb_urls ... ok
test raster_loader::tests::gdal_loader_accepts_latlon_dim_names ... ok
test raster_loader::tests::gdal_loader_defaults_to_band_1_when_fragment_missing ... ok
test gdal_dataset_provider::tests::test_outdb_vrt_cache_reuse ... ok
test gdal_dataset_provider::tests::test_vrt_key_normalizes_gcs_and_azure_equivalent_outdb_urls ... ok
test raster_loader::tests::gdal_loader_errors_on_band_index_out_of_range ... ok
test raster_loader::tests::gdal_loader_rejects_non_2d_source_shape ... ok
test raster_loader::tests::gdal_loader_rejects_request_over_byte_cap ... ok
test raster_loader::tests::gdal_loader_rejects_unrecognized_dim_names ... ok
test raster_loader::tests::read_band_blockwise_honours_pre_cancelled_flag ... ok
test rs_as_raster::tests::test_parse_pixel_type ... ok
test raster_loader::tests::gdal_loader_reads_2d_uint8_geotiff ... ok
test rs_as_raster::tests::test_rs_as_raster_udf_name ... ok
test raster_loader::tests::gdal_loader_errors_when_dtype_disagrees_with_file ... ok
test raster_loader::tests::gdal_loader_errors_on_missing_file ... ok
test rs_frompath::tests::test_invoke_rs_from_path ... FAILED
test raster_loader::tests::gdal_loader_multi_strip_read_matches_single_call ... ok
test rs_frompath::tests::test_invoke_rs_from_path_propagates_nulls ... FAILED
test rs_frompath::tests::test_invoke_rs_from_path_invalid_path_errors ... ok
test rs_frompath::tests::test_invoke_rs_from_path_scalar_ignores_num_rows_for_shape ... FAILED
test rs_frompath::tests::test_rs_from_path_udf_name ... ok
test source_uri::tests::band_fragment_extracts_index ... ok
test gdal_common::tests::test_raster_ref_to_gdal_mem_preserves_band_order_data_and_nodata ... ok
test gdal_dataset_provider::tests::test_provider_returns_mem_dataset_for_indb_raster ... ok
test source_uri::tests::band_one_fragment_round_trips ... ok
test source_uri::tests::band_zero_errors ... ok
test source_uri::tests::band_overflow_errors ... ok
test source_uri::tests::empty_band_value_errors ... ok
test source_uri::tests::empty_fragment_passes_through ... ok
test source_uri::tests::gdal_subdataset_gtiff_dir_passthrough ... ok
test source_uri::tests::empty_uri ... ok
test source_uri::tests::gdal_subdataset_hdf5_passthrough ... ok
test source_uri::tests::band_max_u32_accepted ... ok
test source_uri::tests::gdal_subdataset_netcdf_passthrough ... ok
test source_uri::tests::gdal_subdataset_with_band_fragment_extracts_band ... ok
test source_uri::tests::no_fragment_defaults_to_band_one ... ok
test source_uri::tests::local_path_with_band_fragment ... ok
test source_uri::tests::non_band_fragment_passes_through ... ok
test source_uri::tests::non_numeric_band_errors ... ok
test source_uri::tests::trailing_band_wins_over_earlier_anchor ... ok
test source_uri::tests::url_query_string_preserved_with_band_fragment ... ok
test source_uri::tests::local_path_without_fragment ... ok
test source_uri::tests::negative_band_errors ... ok
test source_uri::tests::url_query_string_with_non_band_fragment_passes_through ... ok
test utils::tests::append_as_outdb_raster_preserves_int64_nodata ... ok
test utils::tests::append_as_outdb_raster_reads_single_band_geotiff ... FAILED
test utils::tests::dataset_to_indb_raster_preserves_int64_nodata_and_data ... ok
test utils::tests::dataset_to_indb_raster_preserves_multi_band_data_and_nodata ... ok
test utils::tests::append_as_indb_raster_appends_multiple_rasters ... ok
test gdal_dataset_provider::tests::test_provider_preserves_band_order_for_mixed_raster ... ok
test rs_as_raster::tests::test_rs_as_raster_use_geometry_extent ... ok
test utils::tests::dataset_to_indb_raster_preserves_uint16_nodata_and_data ... ok
test rs_as_raster::tests::test_rs_as_raster_udf_batch ... ok
test rs_as_raster::tests::test_rs_as_raster_use_reference_extent ... ok
test gdal_dataset_provider::tests::test_provider_returns_empty_dataset_for_zero_band_raster ... ok
test utils::tests::dataset_to_indb_raster_preserves_uint64_nodata_and_data ... ok
test utils::tests::dataset_to_indb_raster_reads_single_band_geotiff ... ok
test gdal_common::tests::test_raster_ref_to_gdal_empty_preserves_metadata_and_crs ... ok
test utils::tests::append_as_outdb_raster_preserves_uint64_nodata ... ok
test utils::tests::dataset_to_indb_raster_preserves_per_band_nodata_for_mem_dataset ... ok
failures:
---- rs_frompath::tests::test_invoke_rs_from_path stdout ----
thread 'rs_frompath::tests::test_invoke_rs_from_path' (436084) panicked at rust/sedona-raster-gdal/src/rs_frompath.rs:147:46:
test4.tiff should exist: External("SedonaDB internal error: sedona-testing raster file 'test4.tiff' not found at '../../submodules/sedona-testing/data/raster/test4.tiff', run submodules/download-assets.py or check the file name.\nThis issue was likely caused by a bug in SedonaDB's code. Please help us to resolve this by filing a bug report in our issue tracker: https://github.com/apache/sedona-db/issues")
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
---- rs_frompath::tests::test_invoke_rs_from_path_propagates_nulls stdout ----
thread 'rs_frompath::tests::test_invoke_rs_from_path_propagates_nulls' (436086) panicked at rust/sedona-raster-gdal/src/rs_frompath.rs:207:46:
test4.tiff should exist: External("SedonaDB internal error: sedona-testing raster file 'test4.tiff' not found at '../../submodules/sedona-testing/data/raster/test4.tiff', run submodules/download-assets.py or check the file name.\nThis issue was likely caused by a bug in SedonaDB's code. Please help us to resolve this by filing a bug report in our issue tracker: https://github.com/apache/sedona-db/issues")
---- rs_frompath::tests::test_invoke_rs_from_path_scalar_ignores_num_rows_for_shape stdout ----
thread 'rs_frompath::tests::test_invoke_rs_from_path_scalar_ignores_num_rows_for_shape' (436087) panicked at rust/sedona-raster-gdal/src/rs_frompath.rs:250:46:
test4.tiff should exist: External("SedonaDB internal error: sedona-testing raster file 'test4.tiff' not found at '../../submodules/sedona-testing/data/raster/test4.tiff', run submodules/download-assets.py or check the file name.\nThis issue was likely caused by a bug in SedonaDB's code. Please help us to resolve this by filing a bug report in our issue tracker: https://github.com/apache/sedona-db/issues")
---- utils::tests::append_as_outdb_raster_reads_single_band_geotiff stdout ----
thread 'utils::tests::append_as_outdb_raster_reads_single_band_geotiff' (436113) panicked at rust/sedona-raster-gdal/src/utils.rs:453:46:
test4.tiff should exist: External("SedonaDB internal error: sedona-testing raster file 'test4.tiff' not found at '../../submodules/sedona-testing/data/raster/test4.tiff', run submodules/download-assets.py or check the file name.\nThis issue was likely caused by a bug in SedonaDB's code. Please help us to resolve this by filing a bug report in our issue tracker: https://github.com/apache/sedona-db/issues")
failures:
rs_frompath::tests::test_invoke_rs_from_path
rs_frompath::tests::test_invoke_rs_from_path_propagates_nulls
rs_frompath::tests::test_invoke_rs_from_path_scalar_ignores_num_rows_for_shape
utils::tests::append_as_outdb_raster_reads_single_band_geotiff
test result: FAILED. 76 passed; 4 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.05s
- blocked in this environment by existing missing raster fixture files in other tests
-
- front matter validated; example execution is currently blocked by an existing Python import mismatch in this environment
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(rust/sedona-raster-gdal): add RS_AsRaster for rasterizing geometries [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #956:
URL: https://github.com/apache/sedona-db/pull/956#discussion_r3560549874
##########
python/sedonadb/tests/functions/test_raster_functions.py:
##########
@@ -268,3 +268,191 @@ def test_rs_setbandnodatavalue_two_arg_requires_single_band():
SedonaDB().assert_query_result(
"SELECT RS_SetBandNoDataValue(RS_Example(), 0)", None
)
+
+
+def _rs_as_raster_sql(
+ pixel_type, all_touched, burn_value, nodata_value, use_geometry_extent
+):
+ return f"""
+ WITH src AS (SELECT RS_FromPath($1) AS raster)
+ SELECT RS_AsRaster(
+ ST_GeomFromText('POLYGON((2 8, 2 5, 5 5, 5 8, 2 8))'),
+ raster,
+ '{pixel_type}',
+ {str(all_touched).upper()},
+ {burn_value},
+ {nodata_value},
+ {str(use_geometry_extent).upper()}
+ ) AS raster
+ FROM src
+ """
+
+
+@pytest.mark.parametrize(
+ ("all_touched", "use_geometry_extent"),
+ [(False, False), (False, True), (True, False)],
+)
+def test_rs_as_raster_matches_rasterio(
+ con, sedona_testing, all_touched, use_geometry_extent
+):
+ import numpy as np
+
+ pytest.importorskip("rasterio")
+ from rasterio.features import rasterize
+ from rasterio.transform import Affine
+ from shapely import wkt
+
+ path = sedona_testing / "data/raster/test4.tiff"
+ transform = (0.0, 1.0, 0.0, 10.0, 0.0, -1.0)
+
+ result = (
+ con.sql(
+ _rs_as_raster_sql("float64", all_touched, 7.0, 0.0, use_geometry_extent),
+ params=(str(path),),
+ )
+ .to_arrow_table()["raster"][0]
+ .as_py()
Review Comment:
It doesn't matter for the test, but because of some pyarrow behaviour the only way to get a zero copy numpy array out is `.to_pandas().iloc[...]` or `Raster(arrow_table["raster"], 0)`. It might be good to update the syntax so that future contributors pick up a better pattern.
##########
python/sedonadb/tests/functions/test_raster_functions.py:
##########
@@ -268,3 +268,191 @@ def test_rs_setbandnodatavalue_two_arg_requires_single_band():
SedonaDB().assert_query_result(
"SELECT RS_SetBandNoDataValue(RS_Example(), 0)", None
)
+
+
+def _rs_as_raster_sql(
+ pixel_type, all_touched, burn_value, nodata_value, use_geometry_extent
+):
+ return f"""
+ WITH src AS (SELECT RS_FromPath($1) AS raster)
+ SELECT RS_AsRaster(
+ ST_GeomFromText('POLYGON((2 8, 2 5, 5 5, 5 8, 2 8))'),
+ raster,
+ '{pixel_type}',
+ {str(all_touched).upper()},
+ {burn_value},
+ {nodata_value},
+ {str(use_geometry_extent).upper()}
+ ) AS raster
+ FROM src
+ """
+
+
+@pytest.mark.parametrize(
+ ("all_touched", "use_geometry_extent"),
+ [(False, False), (False, True), (True, False)],
+)
+def test_rs_as_raster_matches_rasterio(
+ con, sedona_testing, all_touched, use_geometry_extent
+):
+ import numpy as np
Review Comment:
I think we can move this to the top level (it's in test dependencies and should always be around)
##########
docs/reference/sql/rs_asraster.qmd:
##########
@@ -0,0 +1,125 @@
+---
+# 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_AsRaster
+description: Rasterizes a geometry using the grid of a reference raster.
+kernels:
+ - returns: raster
+ args:
+ - name: geometry
+ type: geometry
+ - name: reference_raster
+ type: raster
+ - name: pixel_type
+ type: string
Review Comment:
What we settled on in other PRs is to mark the raster SQL signatures as experimental before we release them to give us the freedom to update the signatures and break compatibility. It might be that with the generated R and Python APIs, users can pass these as keyword args (maybe we can make a list of them and force that behaviour in a future change).
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(rust/sedona-raster-gdal): add RS_AsRaster for rasterizing geometries [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #956:
URL: https://github.com/apache/sedona-db/pull/956#discussion_r3566070758
##########
python/sedonadb/tests/functions/test_raster_functions.py:
##########
@@ -268,3 +268,191 @@ def test_rs_setbandnodatavalue_two_arg_requires_single_band():
SedonaDB().assert_query_result(
"SELECT RS_SetBandNoDataValue(RS_Example(), 0)", None
)
+
+
+def _rs_as_raster_sql(
+ pixel_type, all_touched, burn_value, nodata_value, use_geometry_extent
+):
+ return f"""
+ WITH src AS (SELECT RS_FromPath($1) AS raster)
+ SELECT RS_AsRaster(
+ ST_GeomFromText('POLYGON((2 8, 2 5, 5 5, 5 8, 2 8))'),
+ raster,
+ '{pixel_type}',
+ {str(all_touched).upper()},
+ {burn_value},
+ {nodata_value},
+ {str(use_geometry_extent).upper()}
+ ) AS raster
+ FROM src
+ """
+
+
+@pytest.mark.parametrize(
+ ("all_touched", "use_geometry_extent"),
+ [(False, False), (False, True), (True, False)],
+)
+def test_rs_as_raster_matches_rasterio(
+ con, sedona_testing, all_touched, use_geometry_extent
+):
+ import numpy as np
+
+ pytest.importorskip("rasterio")
+ from rasterio.features import rasterize
+ from rasterio.transform import Affine
+ from shapely import wkt
+
+ path = sedona_testing / "data/raster/test4.tiff"
+ transform = (0.0, 1.0, 0.0, 10.0, 0.0, -1.0)
+
+ result = (
+ con.sql(
+ _rs_as_raster_sql("float64", all_touched, 7.0, 0.0, use_geometry_extent),
+ params=(str(path),),
+ )
+ .to_arrow_table()["raster"][0]
+ .as_py()
Review Comment:
Switched to `Raster(table["raster"], 0)`
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(rust/sedona-raster-gdal): add RS_AsRaster for rasterizing geometries [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on PR #956:
URL: https://github.com/apache/sedona-db/pull/956#issuecomment-4950745733
Updated RS_AsRaster to handle ItemCrs geometries and follow the CRS auto-reprojection rule mentioned in https://github.com/apache/sedona-db/pull/974#discussion_r3455680920.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(rust/sedona-raster-gdal): add RS_AsRaster for rasterizing geometries [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #956:
URL: https://github.com/apache/sedona-db/pull/956#discussion_r3576202545
##########
rust/sedona-raster-gdal/Cargo.toml:
##########
@@ -45,6 +45,7 @@ sedona-gdal = { workspace = true }
sedona-raster = { workspace = true }
sedona-raster-functions = { workspace = true }
sedona-schema = { workspace = true }
+sedona-proj = { workspace = true, features = ["proj-sys"] }
Review Comment:
I think this is causing the R build failures...we can't use proj-sys the feature (we need lazily defined proj). The CrsEngine is now available in the SedonaOptions, so hopefully that removes the need for this!
##########
rust/sedona-raster-gdal/src/rs_as_raster.rs:
##########
@@ -0,0 +1,1014 @@
+// 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_AsRaster UDF - rasterize a vector geometry onto a raster grid.
+
+use std::{convert::TryFrom, sync::Arc};
+
+use arrow_array::ArrayRef;
+use arrow_schema::DataType;
+use datafusion_common::cast::{as_boolean_array, as_float64_array, as_string_array};
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::error::Result;
+use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue};
+use datafusion_expr::{ColumnarValue, Volatility};
+use sedona_expr::{
+ item_crs::parse_item_crs_arg_type,
+ scalar_udf::{SedonaScalarKernel, SedonaScalarUDF},
+};
+use sedona_gdal::dataset::Dataset;
+use sedona_gdal::gdal::Gdal;
+use sedona_gdal::mem::MemDatasetBuilder;
+use sedona_gdal::raster::{rasterband::RasterBand, types::Buffer};
+use sedona_proj::transform::with_global_proj_engine;
+use sedona_raster::array::RasterRefImpl;
+use sedona_raster::builder::RasterBuilder;
+use sedona_raster::traits::{BandMetadata, RasterMetadata, RasterRef};
+use sedona_raster_functions::{
+ crs_utils::{align_wkb_to_crs, resolve_crs},
+ RasterExecutor,
+};
+use sedona_schema::datatypes::{SedonaType, RASTER};
+use sedona_schema::matchers::ArgMatcher;
+use sedona_schema::raster::{BandDataType, StorageType};
+
+use crate::gdal_common::{band_data_type_to_gdal, with_gdal};
+use crate::gdal_dataset_provider::{configure_thread_local_options, thread_local_provider};
+
+pub fn rs_as_raster_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "rs_asraster",
+ vec![
+ Arc::new(RsAsRaster { arg_count: 3 }),
+ Arc::new(RsAsRaster { arg_count: 4 }),
+ Arc::new(RsAsRaster { arg_count: 5 }),
+ Arc::new(RsAsRaster { arg_count: 6 }),
+ Arc::new(RsAsRaster { arg_count: 7 }),
+ ],
+ Volatility::Immutable,
+ )
+}
+
+#[derive(Debug)]
+struct RsAsRaster {
+ arg_count: usize,
+}
+
+impl SedonaScalarKernel for RsAsRaster {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ let Some(first_arg) = args.first() else {
+ return Ok(None);
+ };
+ let (geometry_type, _) = parse_item_crs_arg_type(first_arg)?;
+ let mut arg_types = args.to_vec();
+ arg_types[0] = geometry_type;
+
+ let mut matchers = vec![
+ ArgMatcher::is_geometry_or_geography(),
+ ArgMatcher::is_raster(),
+ ArgMatcher::is_string(),
+ ];
+
+ if self.arg_count >= 4 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+ if self.arg_count >= 5 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 6 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 7 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+
+ ArgMatcher::new(matchers, RASTER).match_args(&arg_types)
+ }
+
+ 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 = args
+ .iter()
+ .find_map(|arg| match arg {
+ ColumnarValue::Array(array) => Some(array.len()),
+ ColumnarValue::Scalar(_) => None,
+ })
+ .unwrap_or(1);
+
+ 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_with_num_iterations(&exec_arg_types, &exec_args, num_iterations);
+
+ let pixel_type_array = args[2]
+ .clone()
+ .cast_to(&DataType::Utf8, None)?
+ .into_array(num_iterations)?;
+ let pixel_type_array = as_string_array(&pixel_type_array)?;
+
+ let all_touched_array = if self.arg_count >= 4 {
+ args[3]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(false)).to_array_of_size(num_iterations)?
+ };
+ let all_touched_array = as_boolean_array(&all_touched_array)?;
+
+ let burn_value_array = if self.arg_count >= 5 {
+ args[4]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(Some(1.0)).to_array_of_size(num_iterations)?
+ };
+ let burn_value_array = as_float64_array(&burn_value_array)?;
+
+ let nodata_value_array = if self.arg_count >= 6 {
+ args[5]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(None).to_array_of_size(num_iterations)?
+ };
+ let nodata_value_array = as_float64_array(&nodata_value_array)?;
+
+ let use_geom_extent_array = if self.arg_count >= 7 {
+ args[6]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(true)).to_array_of_size(num_iterations)?
+ };
+ let use_geom_extent_array = as_boolean_array(&use_geom_extent_array)?;
+
+ let mut builder = RasterBuilder::new(num_iterations);
+
+ let mut pixel_type_iter = pixel_type_array.iter();
+ let mut all_touched_iter = all_touched_array.iter();
+ let mut burn_value_iter = burn_value_array.iter();
+ let mut nodata_value_iter = nodata_value_array.iter();
+ let mut use_geom_extent_iter = use_geom_extent_array.iter();
+
+ with_global_proj_engine(|engine| {
Review Comment:
No need to do it here, but the CrsEngine is now available from the options. I can also take a pass through all the raster UDFs sometime before the next release.
##########
rust/sedona-raster-gdal/src/rs_as_raster.rs:
##########
@@ -0,0 +1,1014 @@
+// 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_AsRaster UDF - rasterize a vector geometry onto a raster grid.
+
+use std::{convert::TryFrom, sync::Arc};
+
+use arrow_array::ArrayRef;
+use arrow_schema::DataType;
+use datafusion_common::cast::{as_boolean_array, as_float64_array, as_string_array};
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::error::Result;
+use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue};
+use datafusion_expr::{ColumnarValue, Volatility};
+use sedona_expr::{
+ item_crs::parse_item_crs_arg_type,
+ scalar_udf::{SedonaScalarKernel, SedonaScalarUDF},
+};
+use sedona_gdal::dataset::Dataset;
+use sedona_gdal::gdal::Gdal;
+use sedona_gdal::mem::MemDatasetBuilder;
+use sedona_gdal::raster::{rasterband::RasterBand, types::Buffer};
+use sedona_proj::transform::with_global_proj_engine;
+use sedona_raster::array::RasterRefImpl;
+use sedona_raster::builder::RasterBuilder;
+use sedona_raster::traits::{BandMetadata, RasterMetadata, RasterRef};
+use sedona_raster_functions::{
+ crs_utils::{align_wkb_to_crs, resolve_crs},
+ RasterExecutor,
+};
+use sedona_schema::datatypes::{SedonaType, RASTER};
+use sedona_schema::matchers::ArgMatcher;
+use sedona_schema::raster::{BandDataType, StorageType};
+
+use crate::gdal_common::{band_data_type_to_gdal, with_gdal};
+use crate::gdal_dataset_provider::{configure_thread_local_options, thread_local_provider};
+
+pub fn rs_as_raster_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "rs_asraster",
+ vec![
+ Arc::new(RsAsRaster { arg_count: 3 }),
+ Arc::new(RsAsRaster { arg_count: 4 }),
+ Arc::new(RsAsRaster { arg_count: 5 }),
+ Arc::new(RsAsRaster { arg_count: 6 }),
+ Arc::new(RsAsRaster { arg_count: 7 }),
+ ],
+ Volatility::Immutable,
+ )
+}
+
+#[derive(Debug)]
+struct RsAsRaster {
+ arg_count: usize,
+}
+
+impl SedonaScalarKernel for RsAsRaster {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ let Some(first_arg) = args.first() else {
+ return Ok(None);
+ };
+ let (geometry_type, _) = parse_item_crs_arg_type(first_arg)?;
+ let mut arg_types = args.to_vec();
+ arg_types[0] = geometry_type;
+
+ let mut matchers = vec![
+ ArgMatcher::is_geometry_or_geography(),
+ ArgMatcher::is_raster(),
+ ArgMatcher::is_string(),
+ ];
+
+ if self.arg_count >= 4 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+ if self.arg_count >= 5 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 6 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 7 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+
+ ArgMatcher::new(matchers, RASTER).match_args(&arg_types)
+ }
+
+ 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 = args
+ .iter()
+ .find_map(|arg| match arg {
+ ColumnarValue::Array(array) => Some(array.len()),
+ ColumnarValue::Scalar(_) => None,
+ })
+ .unwrap_or(1);
+
+ 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_with_num_iterations(&exec_arg_types, &exec_args, num_iterations);
+
+ let pixel_type_array = args[2]
+ .clone()
+ .cast_to(&DataType::Utf8, None)?
+ .into_array(num_iterations)?;
+ let pixel_type_array = as_string_array(&pixel_type_array)?;
+
+ let all_touched_array = if self.arg_count >= 4 {
+ args[3]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(false)).to_array_of_size(num_iterations)?
+ };
+ let all_touched_array = as_boolean_array(&all_touched_array)?;
+
+ let burn_value_array = if self.arg_count >= 5 {
+ args[4]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(Some(1.0)).to_array_of_size(num_iterations)?
+ };
+ let burn_value_array = as_float64_array(&burn_value_array)?;
+
+ let nodata_value_array = if self.arg_count >= 6 {
+ args[5]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(None).to_array_of_size(num_iterations)?
+ };
+ let nodata_value_array = as_float64_array(&nodata_value_array)?;
+
+ let use_geom_extent_array = if self.arg_count >= 7 {
+ args[6]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(true)).to_array_of_size(num_iterations)?
+ };
+ let use_geom_extent_array = as_boolean_array(&use_geom_extent_array)?;
+
+ let mut builder = RasterBuilder::new(num_iterations);
+
+ let mut pixel_type_iter = pixel_type_array.iter();
+ let mut all_touched_iter = all_touched_array.iter();
+ let mut burn_value_iter = burn_value_array.iter();
+ let mut nodata_value_iter = nodata_value_array.iter();
+ let mut use_geom_extent_iter = use_geom_extent_array.iter();
+
+ with_global_proj_engine(|engine| {
+ with_gdal(|gdal| {
+ configure_thread_local_options(gdal, config_options)?;
+
+ executor.execute_raster_wkb_crs_void(|raster_opt, geom_opt, geom_crs| {
+ let pixel_type_opt = pixel_type_iter.next().unwrap();
+ let all_touched_opt = all_touched_iter.next().unwrap().unwrap_or(false);
+ let burn_value_opt = burn_value_iter.next().unwrap().unwrap_or(1.0);
+ let nodata_value_opt = nodata_value_iter.next().unwrap();
+ let use_geometry_extent_opt =
+ use_geom_extent_iter.next().unwrap().unwrap_or(true);
+
+ let raster = match raster_opt {
+ Some(raster) => raster,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+ let geom_wkb = match geom_opt {
+ Some(geom_wkb) => geom_wkb,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+ let pixel_type = match pixel_type_opt {
+ Some(pixel_type) => pixel_type,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+
+ let band_type = parse_pixel_type(pixel_type)?;
+ let raster_crs = resolve_crs(raster.crs())?;
+ let geom_wkb = align_wkb_to_crs(
+ geom_wkb,
+ geom_crs,
+ raster_crs.as_deref(),
+ "geometry",
+ "reference raster",
+ engine,
+ )?;
+
+ let (out_metadata, out_band_metadata, out_band_bytes) = as_raster(
+ gdal,
+ &geom_wkb,
+ raster,
+ band_type,
+ all_touched_opt,
+ burn_value_opt,
+ nodata_value_opt,
+ use_geometry_extent_opt,
+ )
+ .map_err(|e| exec_datafusion_err!("RS_AsRaster failed: {}", e))?;
+
+ builder
+ .start_raster(&out_metadata, raster.crs())
+ .map_err(|e| {
+ exec_datafusion_err!("Failed to start output raster: {}", e)
+ })?;
+ builder.start_band(out_band_metadata).map_err(|e| {
+ exec_datafusion_err!("Failed to start output raster band: {}", e)
+ })?;
+ builder.band_data_writer().append_value(out_band_bytes);
+ builder.finish_band().map_err(|e| {
+ exec_datafusion_err!("Failed to finish output raster band: {}", e)
+ })?;
+ builder.finish_raster().map_err(|e| {
+ exec_datafusion_err!("Failed to finish output raster: {}", e)
+ })?;
+
+ Ok(())
+ })?;
+
+ Ok(())
+ })
+ })?;
+
+ let result: ArrayRef = Arc::new(builder.finish()?);
+ executor.finish(result)
+ }
+}
+
+fn parse_pixel_type(value: &str) -> Result<BandDataType> {
+ match value.trim().to_ascii_lowercase().as_str() {
+ "d" | "float64" => Ok(BandDataType::Float64),
+ "f" | "float32" => Ok(BandDataType::Float32),
+ "i" | "int32" => Ok(BandDataType::Int32),
+ "ui" | "uint32" => Ok(BandDataType::UInt32),
+ "s" | "int16" => Ok(BandDataType::Int16),
+ "us" | "uint16" => Ok(BandDataType::UInt16),
+ "b" | "uint8" => Ok(BandDataType::UInt8),
+ "i8" | "int8" => Ok(BandDataType::Int8),
+ "u64" | "uint64" => Ok(BandDataType::UInt64),
+ "i64" | "int64" => Ok(BandDataType::Int64),
+ other => exec_err!(
+ "Unsupported pixelType: {} (expected one of D/F/I/UI/S/US/B/I8/U64/I64 or int8/uint8/int16/uint16/int32/uint32/int64/uint64/float32/float64)",
+ other
+ ),
+ }
+}
+
+#[derive(Clone, Copy, Debug)]
+enum TypedBandValue {
+ UInt8(u8),
+ Int8(i8),
+ UInt16(u16),
+ Int16(i16),
+ UInt32(u32),
+ Int32(i32),
+ UInt64(u64),
+ Int64(i64),
+ Float32(f32),
+ Float64(f64),
+}
Review Comment:
After this merges perhaps @james-willis can move these to a reusable location (this concept also shows up in one of the functions I just reviewed and this does a better job handling i64 and u64s)
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(rust/sedona-raster-gdal): add RS_AsRaster for rasterizing geometries [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation merged PR #956:
URL: https://github.com/apache/sedona-db/pull/956
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(rust/sedona-raster-gdal): add RS_AsRaster for rasterizing geometries [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #956:
URL: https://github.com/apache/sedona-db/pull/956#discussion_r3453628531
##########
docs/reference/sql/rs_asraster.qmd:
##########
@@ -0,0 +1,125 @@
+---
+# 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_AsRaster
+description: Rasterizes a geometry using the grid of a reference raster.
+kernels:
+ - returns: raster
+ args:
+ - name: geometry
+ type: geometry
+ - name: reference_raster
+ type: raster
+ - name: pixel_type
+ type: string
+ - returns: raster
+ args:
+ - name: geometry
+ type: geometry
+ - name: reference_raster
+ type: raster
+ - name: pixel_type
+ type: string
+ - name: all_touched
+ type: boolean
+ - returns: raster
+ args:
+ - name: geometry
+ type: geometry
+ - name: reference_raster
+ type: raster
+ - name: pixel_type
+ type: string
+ - name: all_touched
+ type: boolean
+ - name: burn_value
+ type: double
+ - returns: raster
+ args:
+ - name: geometry
+ type: geometry
+ - name: reference_raster
+ type: raster
+ - name: pixel_type
+ type: string
+ - name: all_touched
+ type: boolean
+ - name: burn_value
+ type: double
+ - name: nodata_value
+ type: double
+ - returns: raster
+ args:
+ - name: geometry
+ type: geometry
+ - name: reference_raster
+ type: raster
+ - name: pixel_type
+ type: string
+ - name: all_touched
+ type: boolean
+ - name: burn_value
+ type: double
+ - name: nodata_value
+ type: double
+ - name: use_geometry_extent
+ type: boolean
+ - returns: raster
+ args:
+ - name: geometry
+ type: geometry
+ - name: reference_raster
+ type: raster
+ - name: pixel_type
+ type: string
+ - name: all_touched
+ type: boolean
+ - name: burn_value
+ type: double
+ - name: nodata_value
+ type: double
+ - name: use_geometry_extent
+ type: boolean
+---
+
+## Description
+
+Creates a single-band in-db raster by rasterizing `geometry` onto the grid
+defined by `reference_raster`.
+
+The output pixel type is controlled by `pixel_type`. Optional arguments control
+whether pixels touched by the geometry are burned, the value burned into covered
+pixels, the nodata fill value, and whether the output extent is cropped to the
+geometry extent or kept equal to the full reference raster extent.
+
+## Examples
+
+```sql
+SELECT
+ RS_Width(
+ RS_AsRaster(
+ ST_GeomFromWKT('POLYGON((1 1, 1 2, 2 2, 2 1, 1 1))'),
+ RS_FromPath('../../../submodules/sedona-testing/data/raster/test4.tiff'),
Review Comment:
Can this example be user-runnable using a URL (can be URL to the sedona-testing repo tiff)
##########
docs/reference/sql/rs_asraster.qmd:
##########
@@ -0,0 +1,125 @@
+---
+# 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_AsRaster
+description: Rasterizes a geometry using the grid of a reference raster.
+kernels:
+ - returns: raster
+ args:
+ - name: geometry
+ type: geometry
+ - name: reference_raster
+ type: raster
+ - name: pixel_type
+ type: string
Review Comment:
I don't think we should solve this now but these signatures result in pretty much unreadable SQL (we can probably improve this with some unified "options" handling in the future using struct arrays or JSON).
##########
rust/sedona-raster-gdal/src/rs_as_raster.rs:
##########
@@ -0,0 +1,702 @@
+// 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_AsRaster UDF - rasterize a vector geometry onto a raster grid.
+
+use std::sync::Arc;
+
+use arrow_array::{Array, StructArray};
+use arrow_schema::DataType;
+use datafusion_common::cast::{
+ as_binary_array, as_boolean_array, as_float64_array, as_string_array,
+};
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::error::Result;
+use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue};
+use datafusion_expr::{ColumnarValue, Volatility};
+use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err};
+use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
+use sedona_gdal::dataset::Dataset;
+use sedona_gdal::gdal::Gdal;
+use sedona_gdal::mem::MemDatasetBuilder;
+use sedona_gdal::raster::types::Buffer;
+use sedona_raster::array::{RasterRefImpl, RasterStructArray};
+use sedona_raster::builder::RasterBuilder;
+use sedona_raster::traits::{BandMetadata, RasterMetadata, RasterRef};
+use sedona_schema::datatypes::{SedonaType, RASTER};
+use sedona_schema::matchers::ArgMatcher;
+use sedona_schema::raster::{BandDataType, StorageType};
+
+use crate::gdal_common::{band_data_type_to_gdal, nodata_f64_to_bytes, with_gdal};
+use crate::gdal_dataset_provider::{configure_thread_local_options, thread_local_provider};
+
+pub fn rs_as_raster_udf() -> SedonaScalarUDF {
Review Comment:
Not for now, but this seems like it would be most effective as an aggregate function (ST_Collect + AsRaster is a workaround but probably much more intensive)
##########
rust/sedona-raster-gdal/src/rs_as_raster.rs:
##########
@@ -0,0 +1,702 @@
+// 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_AsRaster UDF - rasterize a vector geometry onto a raster grid.
+
+use std::sync::Arc;
+
+use arrow_array::{Array, StructArray};
+use arrow_schema::DataType;
+use datafusion_common::cast::{
+ as_binary_array, as_boolean_array, as_float64_array, as_string_array,
+};
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::error::Result;
+use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue};
+use datafusion_expr::{ColumnarValue, Volatility};
+use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err};
+use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
+use sedona_gdal::dataset::Dataset;
+use sedona_gdal::gdal::Gdal;
+use sedona_gdal::mem::MemDatasetBuilder;
+use sedona_gdal::raster::types::Buffer;
+use sedona_raster::array::{RasterRefImpl, RasterStructArray};
+use sedona_raster::builder::RasterBuilder;
+use sedona_raster::traits::{BandMetadata, RasterMetadata, RasterRef};
+use sedona_schema::datatypes::{SedonaType, RASTER};
+use sedona_schema::matchers::ArgMatcher;
+use sedona_schema::raster::{BandDataType, StorageType};
+
+use crate::gdal_common::{band_data_type_to_gdal, nodata_f64_to_bytes, with_gdal};
+use crate::gdal_dataset_provider::{configure_thread_local_options, thread_local_provider};
+
+pub fn rs_as_raster_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "rs_asraster",
+ vec![
+ Arc::new(RsAsRaster { arg_count: 3 }),
+ Arc::new(RsAsRaster { arg_count: 4 }),
+ Arc::new(RsAsRaster { arg_count: 5 }),
+ Arc::new(RsAsRaster { arg_count: 6 }),
+ Arc::new(RsAsRaster { arg_count: 7 }),
+ ],
+ Volatility::Immutable,
+ )
+}
+
+#[derive(Debug)]
+struct RsAsRaster {
+ arg_count: usize,
+}
+
+impl SedonaScalarKernel for RsAsRaster {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ let mut matchers = vec![
+ ArgMatcher::is_geometry_or_geography(),
+ ArgMatcher::is_raster(),
+ ArgMatcher::is_string(),
+ ];
+
+ if self.arg_count >= 4 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+ if self.arg_count >= 5 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 6 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 7 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+
+ 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 = calc_num_iterations(args);
+
+ let geom_array = args[0].clone().into_array(num_iterations)?;
+ let geom_array = as_binary_array(&geom_array)?;
+
+ let pixel_type_array = args[2]
+ .clone()
+ .cast_to(&DataType::Utf8, None)?
+ .into_array(num_iterations)?;
+ let pixel_type_array = as_string_array(&pixel_type_array)?;
+
+ let all_touched_array = if self.arg_count >= 4 {
+ args[3]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(false)).to_array_of_size(num_iterations)?
+ };
+ let all_touched_array = as_boolean_array(&all_touched_array)?;
+
+ let burn_value_array = if self.arg_count >= 5 {
+ args[4]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(Some(1.0)).to_array_of_size(num_iterations)?
+ };
+ let burn_value_array = as_float64_array(&burn_value_array)?;
+
+ let nodata_value_array = if self.arg_count >= 6 {
+ args[5]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(None).to_array_of_size(num_iterations)?
+ };
+ let nodata_value_array = as_float64_array(&nodata_value_array)?;
+
+ let use_geom_extent_array = if self.arg_count >= 7 {
+ args[6]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(true)).to_array_of_size(num_iterations)?
+ };
+ let use_geom_extent_array = as_boolean_array(&use_geom_extent_array)?;
+
+ let mut builder = RasterBuilder::new(num_iterations);
+
+ let mut geom_iter = geom_array.iter();
+ let mut pixel_type_iter = pixel_type_array.iter();
+ let mut all_touched_iter = all_touched_array.iter();
+ let mut burn_value_iter = burn_value_array.iter();
+ let mut nodata_value_iter = nodata_value_array.iter();
+ let mut use_geom_extent_iter = use_geom_extent_array.iter();
+
+ with_gdal(|gdal| {
+ configure_thread_local_options(gdal, config_options)?;
+
+ execute_raster_arg(arg_types, args, 1, num_iterations, |_, raster_opt| {
+ let geom_opt = geom_iter.next().unwrap();
+ let pixel_type_opt = pixel_type_iter.next().unwrap();
+ let all_touched_opt = all_touched_iter.next().unwrap().unwrap_or(false);
+ let burn_value_opt = burn_value_iter.next().unwrap().unwrap_or(1.0);
+ let nodata_value_opt = nodata_value_iter.next().unwrap();
+ let use_geometry_extent_opt = use_geom_extent_iter.next().unwrap().unwrap_or(true);
+
+ let raster = match raster_opt {
+ Some(raster) => raster,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+ let geom_wkb = match geom_opt {
+ Some(geom_wkb) => geom_wkb,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+ let pixel_type = match pixel_type_opt {
+ Some(pixel_type) => pixel_type,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+
+ let band_type = parse_pixel_type(pixel_type)?;
+
+ match as_raster(
+ gdal,
+ geom_wkb,
+ raster,
+ band_type,
+ all_touched_opt,
+ burn_value_opt,
+ nodata_value_opt,
+ use_geometry_extent_opt,
+ ) {
+ Ok((out_metadata, out_band_metadata, out_band_bytes)) => {
+ builder
+ .start_raster(&out_metadata, raster.crs())
+ .map_err(|e| {
+ exec_datafusion_err!("Failed to start output raster: {}", e)
+ })?;
+ builder.start_band(out_band_metadata).map_err(|e| {
+ exec_datafusion_err!("Failed to start output raster band: {}", e)
+ })?;
+ builder.band_data_writer().append_value(out_band_bytes);
+ builder.finish_band().map_err(|e| {
+ exec_datafusion_err!("Failed to finish output raster band: {}", e)
+ })?;
+ builder.finish_raster().map_err(|e| {
+ exec_datafusion_err!("Failed to finish output raster: {}", e)
+ })?;
+ }
+ Err(_) => builder.append_null()?,
+ }
+
+ Ok(())
+ })?;
+
+ finish_result(args, Arc::new(builder.finish()?))
+ })
+ }
+}
+
+fn parse_pixel_type(value: &str) -> Result<BandDataType> {
+ match value.trim().to_ascii_uppercase().as_str() {
+ "D" => Ok(BandDataType::Float64),
+ "F" => Ok(BandDataType::Float32),
+ "I" => Ok(BandDataType::Int32),
+ "S" => Ok(BandDataType::Int16),
+ "US" => Ok(BandDataType::UInt16),
+ "B" => Ok(BandDataType::UInt8),
+ "I8" | "INT8" => Ok(BandDataType::Int8),
+ "U64" | "UINT64" => Ok(BandDataType::UInt64),
+ "I64" | "INT64" => Ok(BandDataType::Int64),
+ other => exec_err!(
+ "Unsupported pixelType: {} (expected one of D, F, I, S, US, B, I8, U64, I64)",
+ other
+ ),
+ }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn as_raster(
+ gdal: &Gdal,
+ geom_wkb: &[u8],
+ reference_raster: &RasterRefImpl<'_>,
+ band_type: BandDataType,
+ all_touched: bool,
+ burn_value: f64,
+ nodata_value: Option<f64>,
+ use_geometry_extent: bool,
+) -> Result<(RasterMetadata, BandMetadata, Vec<u8>)> {
+ let ref_md = reference_raster.metadata();
+
+ if ref_md.skew_x() != 0.0 || ref_md.skew_y() != 0.0 {
+ return exec_err!(
+ "RS_AsRaster currently requires skew_x=0 and skew_y=0 in the reference raster"
+ );
+ }
+
+ let geometry = gdal
+ .geometry_from_wkb(geom_wkb)
+ .map_err(|e| exec_datafusion_err!("Failed to parse geometry from WKB: {}", e))?;
+
+ let (out_width, out_height, out_ulx, out_uly) = if use_geometry_extent {
+ let env = geometry.envelope();
+ let ulx = ref_md.upper_left_x();
+ let uly = ref_md.upper_left_y();
+ let scale_x = ref_md.scale_x();
+ let scale_y = ref_md.scale_y();
+
+ if scale_x == 0.0 || scale_y == 0.0 {
+ return exec_err!("Reference raster has zero scale");
+ }
+
+ let start_col = ((env.MinX - ulx) / scale_x).floor() as isize;
+ let end_col_excl = ((env.MaxX - ulx) / scale_x).ceil() as isize;
+ let start_row = ((env.MaxY - uly) / scale_y).floor() as isize;
+ let end_row_excl = ((env.MinY - uly) / scale_y).ceil() as isize;
+
+ let width = (end_col_excl - start_col).max(0) as usize;
+ let height = (end_row_excl - start_row).max(0) as usize;
Review Comment:
Should these use a checked cast? I may be reading this incorrectly but it is probably possible to accidentally pass a huge geometry and a very small reference raster and overflow an `isize`.
##########
rust/sedona-raster-gdal/src/rs_as_raster.rs:
##########
@@ -0,0 +1,702 @@
+// 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_AsRaster UDF - rasterize a vector geometry onto a raster grid.
+
+use std::sync::Arc;
+
+use arrow_array::{Array, StructArray};
+use arrow_schema::DataType;
+use datafusion_common::cast::{
+ as_binary_array, as_boolean_array, as_float64_array, as_string_array,
+};
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::error::Result;
+use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue};
+use datafusion_expr::{ColumnarValue, Volatility};
+use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err};
+use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
+use sedona_gdal::dataset::Dataset;
+use sedona_gdal::gdal::Gdal;
+use sedona_gdal::mem::MemDatasetBuilder;
+use sedona_gdal::raster::types::Buffer;
+use sedona_raster::array::{RasterRefImpl, RasterStructArray};
+use sedona_raster::builder::RasterBuilder;
+use sedona_raster::traits::{BandMetadata, RasterMetadata, RasterRef};
+use sedona_schema::datatypes::{SedonaType, RASTER};
+use sedona_schema::matchers::ArgMatcher;
+use sedona_schema::raster::{BandDataType, StorageType};
+
+use crate::gdal_common::{band_data_type_to_gdal, nodata_f64_to_bytes, with_gdal};
+use crate::gdal_dataset_provider::{configure_thread_local_options, thread_local_provider};
+
+pub fn rs_as_raster_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "rs_asraster",
+ vec![
+ Arc::new(RsAsRaster { arg_count: 3 }),
+ Arc::new(RsAsRaster { arg_count: 4 }),
+ Arc::new(RsAsRaster { arg_count: 5 }),
+ Arc::new(RsAsRaster { arg_count: 6 }),
+ Arc::new(RsAsRaster { arg_count: 7 }),
+ ],
+ Volatility::Immutable,
+ )
+}
+
+#[derive(Debug)]
+struct RsAsRaster {
+ arg_count: usize,
+}
+
+impl SedonaScalarKernel for RsAsRaster {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ let mut matchers = vec![
+ ArgMatcher::is_geometry_or_geography(),
+ ArgMatcher::is_raster(),
+ ArgMatcher::is_string(),
+ ];
+
+ if self.arg_count >= 4 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+ if self.arg_count >= 5 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 6 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 7 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+
+ 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 = calc_num_iterations(args);
+
+ let geom_array = args[0].clone().into_array(num_iterations)?;
+ let geom_array = as_binary_array(&geom_array)?;
+
+ let pixel_type_array = args[2]
+ .clone()
+ .cast_to(&DataType::Utf8, None)?
+ .into_array(num_iterations)?;
+ let pixel_type_array = as_string_array(&pixel_type_array)?;
+
+ let all_touched_array = if self.arg_count >= 4 {
+ args[3]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(false)).to_array_of_size(num_iterations)?
+ };
+ let all_touched_array = as_boolean_array(&all_touched_array)?;
+
+ let burn_value_array = if self.arg_count >= 5 {
+ args[4]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(Some(1.0)).to_array_of_size(num_iterations)?
+ };
+ let burn_value_array = as_float64_array(&burn_value_array)?;
+
+ let nodata_value_array = if self.arg_count >= 6 {
+ args[5]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(None).to_array_of_size(num_iterations)?
+ };
+ let nodata_value_array = as_float64_array(&nodata_value_array)?;
+
+ let use_geom_extent_array = if self.arg_count >= 7 {
+ args[6]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(true)).to_array_of_size(num_iterations)?
+ };
+ let use_geom_extent_array = as_boolean_array(&use_geom_extent_array)?;
+
+ let mut builder = RasterBuilder::new(num_iterations);
+
+ let mut geom_iter = geom_array.iter();
+ let mut pixel_type_iter = pixel_type_array.iter();
+ let mut all_touched_iter = all_touched_array.iter();
+ let mut burn_value_iter = burn_value_array.iter();
+ let mut nodata_value_iter = nodata_value_array.iter();
+ let mut use_geom_extent_iter = use_geom_extent_array.iter();
+
+ with_gdal(|gdal| {
+ configure_thread_local_options(gdal, config_options)?;
+
+ execute_raster_arg(arg_types, args, 1, num_iterations, |_, raster_opt| {
+ let geom_opt = geom_iter.next().unwrap();
+ let pixel_type_opt = pixel_type_iter.next().unwrap();
+ let all_touched_opt = all_touched_iter.next().unwrap().unwrap_or(false);
+ let burn_value_opt = burn_value_iter.next().unwrap().unwrap_or(1.0);
+ let nodata_value_opt = nodata_value_iter.next().unwrap();
+ let use_geometry_extent_opt = use_geom_extent_iter.next().unwrap().unwrap_or(true);
+
+ let raster = match raster_opt {
+ Some(raster) => raster,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+ let geom_wkb = match geom_opt {
+ Some(geom_wkb) => geom_wkb,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+ let pixel_type = match pixel_type_opt {
+ Some(pixel_type) => pixel_type,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+
+ let band_type = parse_pixel_type(pixel_type)?;
+
+ match as_raster(
+ gdal,
+ geom_wkb,
+ raster,
+ band_type,
+ all_touched_opt,
+ burn_value_opt,
+ nodata_value_opt,
+ use_geometry_extent_opt,
+ ) {
+ Ok((out_metadata, out_band_metadata, out_band_bytes)) => {
+ builder
+ .start_raster(&out_metadata, raster.crs())
+ .map_err(|e| {
+ exec_datafusion_err!("Failed to start output raster: {}", e)
+ })?;
+ builder.start_band(out_band_metadata).map_err(|e| {
+ exec_datafusion_err!("Failed to start output raster band: {}", e)
+ })?;
+ builder.band_data_writer().append_value(out_band_bytes);
+ builder.finish_band().map_err(|e| {
+ exec_datafusion_err!("Failed to finish output raster band: {}", e)
+ })?;
+ builder.finish_raster().map_err(|e| {
+ exec_datafusion_err!("Failed to finish output raster: {}", e)
+ })?;
+ }
+ Err(_) => builder.append_null()?,
+ }
+
+ Ok(())
+ })?;
+
+ finish_result(args, Arc::new(builder.finish()?))
+ })
+ }
+}
+
+fn parse_pixel_type(value: &str) -> Result<BandDataType> {
+ match value.trim().to_ascii_uppercase().as_str() {
+ "D" => Ok(BandDataType::Float64),
+ "F" => Ok(BandDataType::Float32),
+ "I" => Ok(BandDataType::Int32),
+ "S" => Ok(BandDataType::Int16),
+ "US" => Ok(BandDataType::UInt16),
+ "B" => Ok(BandDataType::UInt8),
+ "I8" | "INT8" => Ok(BandDataType::Int8),
+ "U64" | "UINT64" => Ok(BandDataType::UInt64),
+ "I64" | "INT64" => Ok(BandDataType::Int64),
+ other => exec_err!(
+ "Unsupported pixelType: {} (expected one of D, F, I, S, US, B, I8, U64, I64)",
+ other
+ ),
+ }
+}
Review Comment:
These are pretty cryptic...can we also accept the strings we use elsewhere (`int8`, `uint8`, etc.)?
##########
rust/sedona-raster-gdal/src/rs_as_raster.rs:
##########
@@ -0,0 +1,702 @@
+// 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_AsRaster UDF - rasterize a vector geometry onto a raster grid.
+
+use std::sync::Arc;
+
+use arrow_array::{Array, StructArray};
+use arrow_schema::DataType;
+use datafusion_common::cast::{
+ as_binary_array, as_boolean_array, as_float64_array, as_string_array,
+};
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::error::Result;
+use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue};
+use datafusion_expr::{ColumnarValue, Volatility};
+use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err};
+use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
+use sedona_gdal::dataset::Dataset;
+use sedona_gdal::gdal::Gdal;
+use sedona_gdal::mem::MemDatasetBuilder;
+use sedona_gdal::raster::types::Buffer;
+use sedona_raster::array::{RasterRefImpl, RasterStructArray};
+use sedona_raster::builder::RasterBuilder;
+use sedona_raster::traits::{BandMetadata, RasterMetadata, RasterRef};
+use sedona_schema::datatypes::{SedonaType, RASTER};
+use sedona_schema::matchers::ArgMatcher;
+use sedona_schema::raster::{BandDataType, StorageType};
+
+use crate::gdal_common::{band_data_type_to_gdal, nodata_f64_to_bytes, with_gdal};
+use crate::gdal_dataset_provider::{configure_thread_local_options, thread_local_provider};
+
+pub fn rs_as_raster_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "rs_asraster",
+ vec![
+ Arc::new(RsAsRaster { arg_count: 3 }),
+ Arc::new(RsAsRaster { arg_count: 4 }),
+ Arc::new(RsAsRaster { arg_count: 5 }),
+ Arc::new(RsAsRaster { arg_count: 6 }),
+ Arc::new(RsAsRaster { arg_count: 7 }),
+ ],
+ Volatility::Immutable,
+ )
+}
+
+#[derive(Debug)]
+struct RsAsRaster {
+ arg_count: usize,
+}
+
+impl SedonaScalarKernel for RsAsRaster {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ let mut matchers = vec![
+ ArgMatcher::is_geometry_or_geography(),
+ ArgMatcher::is_raster(),
+ ArgMatcher::is_string(),
+ ];
+
+ if self.arg_count >= 4 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+ if self.arg_count >= 5 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 6 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 7 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+
+ 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 = calc_num_iterations(args);
+
+ let geom_array = args[0].clone().into_array(num_iterations)?;
+ let geom_array = as_binary_array(&geom_array)?;
+
+ let pixel_type_array = args[2]
+ .clone()
+ .cast_to(&DataType::Utf8, None)?
+ .into_array(num_iterations)?;
+ let pixel_type_array = as_string_array(&pixel_type_array)?;
+
+ let all_touched_array = if self.arg_count >= 4 {
+ args[3]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(false)).to_array_of_size(num_iterations)?
+ };
+ let all_touched_array = as_boolean_array(&all_touched_array)?;
+
+ let burn_value_array = if self.arg_count >= 5 {
+ args[4]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(Some(1.0)).to_array_of_size(num_iterations)?
+ };
+ let burn_value_array = as_float64_array(&burn_value_array)?;
+
+ let nodata_value_array = if self.arg_count >= 6 {
+ args[5]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(None).to_array_of_size(num_iterations)?
+ };
+ let nodata_value_array = as_float64_array(&nodata_value_array)?;
+
+ let use_geom_extent_array = if self.arg_count >= 7 {
+ args[6]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(true)).to_array_of_size(num_iterations)?
+ };
+ let use_geom_extent_array = as_boolean_array(&use_geom_extent_array)?;
+
+ let mut builder = RasterBuilder::new(num_iterations);
+
+ let mut geom_iter = geom_array.iter();
+ let mut pixel_type_iter = pixel_type_array.iter();
+ let mut all_touched_iter = all_touched_array.iter();
+ let mut burn_value_iter = burn_value_array.iter();
+ let mut nodata_value_iter = nodata_value_array.iter();
+ let mut use_geom_extent_iter = use_geom_extent_array.iter();
+
+ with_gdal(|gdal| {
+ configure_thread_local_options(gdal, config_options)?;
+
+ execute_raster_arg(arg_types, args, 1, num_iterations, |_, raster_opt| {
+ let geom_opt = geom_iter.next().unwrap();
+ let pixel_type_opt = pixel_type_iter.next().unwrap();
+ let all_touched_opt = all_touched_iter.next().unwrap().unwrap_or(false);
+ let burn_value_opt = burn_value_iter.next().unwrap().unwrap_or(1.0);
+ let nodata_value_opt = nodata_value_iter.next().unwrap();
+ let use_geometry_extent_opt = use_geom_extent_iter.next().unwrap().unwrap_or(true);
+
+ let raster = match raster_opt {
+ Some(raster) => raster,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+ let geom_wkb = match geom_opt {
+ Some(geom_wkb) => geom_wkb,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+ let pixel_type = match pixel_type_opt {
+ Some(pixel_type) => pixel_type,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+
+ let band_type = parse_pixel_type(pixel_type)?;
+
+ match as_raster(
+ gdal,
+ geom_wkb,
+ raster,
+ band_type,
+ all_touched_opt,
+ burn_value_opt,
+ nodata_value_opt,
+ use_geometry_extent_opt,
+ ) {
+ Ok((out_metadata, out_band_metadata, out_band_bytes)) => {
+ builder
+ .start_raster(&out_metadata, raster.crs())
+ .map_err(|e| {
+ exec_datafusion_err!("Failed to start output raster: {}", e)
+ })?;
+ builder.start_band(out_band_metadata).map_err(|e| {
+ exec_datafusion_err!("Failed to start output raster band: {}", e)
+ })?;
+ builder.band_data_writer().append_value(out_band_bytes);
+ builder.finish_band().map_err(|e| {
+ exec_datafusion_err!("Failed to finish output raster band: {}", e)
+ })?;
+ builder.finish_raster().map_err(|e| {
+ exec_datafusion_err!("Failed to finish output raster: {}", e)
+ })?;
+ }
+ Err(_) => builder.append_null()?,
+ }
+
+ Ok(())
+ })?;
+
+ finish_result(args, Arc::new(builder.finish()?))
+ })
+ }
+}
+
+fn parse_pixel_type(value: &str) -> Result<BandDataType> {
+ match value.trim().to_ascii_uppercase().as_str() {
+ "D" => Ok(BandDataType::Float64),
+ "F" => Ok(BandDataType::Float32),
+ "I" => Ok(BandDataType::Int32),
+ "S" => Ok(BandDataType::Int16),
+ "US" => Ok(BandDataType::UInt16),
+ "B" => Ok(BandDataType::UInt8),
+ "I8" | "INT8" => Ok(BandDataType::Int8),
+ "U64" | "UINT64" => Ok(BandDataType::UInt64),
+ "I64" | "INT64" => Ok(BandDataType::Int64),
+ other => exec_err!(
+ "Unsupported pixelType: {} (expected one of D, F, I, S, US, B, I8, U64, I64)",
+ other
+ ),
+ }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn as_raster(
+ gdal: &Gdal,
+ geom_wkb: &[u8],
+ reference_raster: &RasterRefImpl<'_>,
+ band_type: BandDataType,
+ all_touched: bool,
+ burn_value: f64,
+ nodata_value: Option<f64>,
+ use_geometry_extent: bool,
+) -> Result<(RasterMetadata, BandMetadata, Vec<u8>)> {
+ let ref_md = reference_raster.metadata();
+
+ if ref_md.skew_x() != 0.0 || ref_md.skew_y() != 0.0 {
+ return exec_err!(
+ "RS_AsRaster currently requires skew_x=0 and skew_y=0 in the reference raster"
+ );
+ }
+
+ let geometry = gdal
+ .geometry_from_wkb(geom_wkb)
+ .map_err(|e| exec_datafusion_err!("Failed to parse geometry from WKB: {}", e))?;
+
+ let (out_width, out_height, out_ulx, out_uly) = if use_geometry_extent {
+ let env = geometry.envelope();
+ let ulx = ref_md.upper_left_x();
+ let uly = ref_md.upper_left_y();
+ let scale_x = ref_md.scale_x();
+ let scale_y = ref_md.scale_y();
+
+ if scale_x == 0.0 || scale_y == 0.0 {
+ return exec_err!("Reference raster has zero scale");
+ }
+
+ let start_col = ((env.MinX - ulx) / scale_x).floor() as isize;
+ let end_col_excl = ((env.MaxX - ulx) / scale_x).ceil() as isize;
+ let start_row = ((env.MaxY - uly) / scale_y).floor() as isize;
+ let end_row_excl = ((env.MinY - uly) / scale_y).ceil() as isize;
+
+ let width = (end_col_excl - start_col).max(0) as usize;
+ let height = (end_row_excl - start_row).max(0) as usize;
+
+ if width == 0 || height == 0 {
+ return exec_err!("Geometry extent produced an empty raster");
+ }
+
+ (
+ width,
+ height,
+ ulx + (start_col as f64) * scale_x,
+ uly + (start_row as f64) * scale_y,
+ )
+ } else {
+ (
+ ref_md.width() as usize,
+ ref_md.height() as usize,
+ ref_md.upper_left_x(),
+ ref_md.upper_left_y(),
+ )
+ };
+
+ let out_dataset = MemDatasetBuilder::create(
+ gdal,
+ out_width,
+ out_height,
+ 1,
+ band_data_type_to_gdal(&band_type),
+ )
+ .map_err(|e| exec_datafusion_err!("Failed to create dataset: {}", e))?;
+
+ out_dataset
+ .set_geo_transform(&[
+ out_ulx,
+ ref_md.scale_x(),
+ ref_md.skew_x(),
+ out_uly,
+ ref_md.skew_y(),
+ ref_md.scale_y(),
+ ])
+ .map_err(|e| exec_datafusion_err!("Failed to set geotransform: {}", e))?;
+
+ let provider = thread_local_provider(gdal)
+ .map_err(|e| exec_datafusion_err!("Failed to init GDAL provider: {}", e))?;
+ let ref_raster_ds = provider
+ .raster_ref_to_gdal(reference_raster)
+ .map_err(|e| exec_datafusion_err!("Failed to create GDAL dataset: {}", e))?;
+ if let Ok(srs) = ref_raster_ds.as_dataset().spatial_ref() {
+ out_dataset
+ .set_spatial_ref(&srs)
+ .map_err(|e| exec_datafusion_err!("Failed to set spatial reference: {}", e))?;
+ }
+
+ let init_value = nodata_value.unwrap_or(0.0);
+ initialize_band(&out_dataset, &band_type, out_width, out_height, init_value)?;
+
+ if let Some(nodata) = nodata_value {
+ let band = out_dataset
+ .rasterband(1)
+ .map_err(|e| exec_datafusion_err!("Failed to get output band: {}", e))?;
+ match band_type {
+ BandDataType::UInt64 => band
+ .set_no_data_value_u64(Some(nodata as u64))
+ .map_err(|e| exec_datafusion_err!("Failed to set nodata value: {}", e))?,
+ BandDataType::Int64 => band
+ .set_no_data_value_i64(Some(nodata as i64))
+ .map_err(|e| exec_datafusion_err!("Failed to set nodata value: {}", e))?,
+ _ => band
+ .set_no_data_value(Some(nodata))
+ .map_err(|e| exec_datafusion_err!("Failed to set nodata value: {}", e))?,
+ }
+ }
+
+ gdal.rasterize_affine(&out_dataset, &[1], &[geometry], &[burn_value], all_touched)
+ .map_err(|e| exec_datafusion_err!("Failed to rasterize geometry: {}", e))?;
+
+ let band = out_dataset
+ .rasterband(1)
+ .map_err(|e| exec_datafusion_err!("Failed to get output band: {}", e))?;
+ let band_bytes = band
+ .read_as_bytes(
+ (0, 0),
+ (out_width, out_height),
+ (out_width, out_height),
+ None,
+ )
+ .map_err(|e| exec_datafusion_err!("Failed to read band data: {}", e))?;
+
+ Ok((
+ RasterMetadata {
+ width: out_width as i64,
+ height: out_height as i64,
+ upperleft_x: out_ulx,
+ upperleft_y: out_uly,
+ scale_x: ref_md.scale_x(),
+ scale_y: ref_md.scale_y(),
+ skew_x: ref_md.skew_x(),
+ skew_y: ref_md.skew_y(),
+ },
+ BandMetadata {
+ nodata_value: nodata_value.map(|value| nodata_f64_to_bytes(value, &band_type)),
+ storage_type: StorageType::InDb,
+ datatype: band_type,
+ outdb_url: None,
+ outdb_band_id: None,
+ },
+ band_bytes,
+ ))
+}
+
+fn initialize_band(
+ dataset: &Dataset,
+ band_type: &BandDataType,
+ width: usize,
+ height: usize,
+ init_value: f64,
+) -> Result<()> {
+ match band_type {
+ BandDataType::UInt8 => initialize_band_t::<u8>(dataset, width, height, init_value as u8),
+ BandDataType::Int8 => initialize_band_t::<i8>(dataset, width, height, init_value as i8),
+ BandDataType::UInt16 => initialize_band_t::<u16>(dataset, width, height, init_value as u16),
+ BandDataType::Int16 => initialize_band_t::<i16>(dataset, width, height, init_value as i16),
+ BandDataType::UInt32 => initialize_band_t::<u32>(dataset, width, height, init_value as u32),
+ BandDataType::Int32 => initialize_band_t::<i32>(dataset, width, height, init_value as i32),
+ BandDataType::UInt64 => initialize_band_t::<u64>(dataset, width, height, init_value as u64),
+ BandDataType::Int64 => initialize_band_t::<i64>(dataset, width, height, init_value as i64),
+ BandDataType::Float32 => {
+ initialize_band_t::<f32>(dataset, width, height, init_value as f32)
+ }
+ BandDataType::Float64 => initialize_band_t::<f64>(dataset, width, height, init_value),
+ }
+}
+
+fn initialize_band_t<T: sedona_gdal::raster::types::GdalType + Copy>(
+ dataset: &Dataset,
+ width: usize,
+ height: usize,
+ init_value: T,
+) -> Result<()> {
+ let band = dataset
+ .rasterband(1)
+ .map_err(|e| exec_datafusion_err!("Failed to get output band: {}", e))?;
+ let mut buffer = Buffer::new((width, height), vec![init_value; width * height]);
+ band.write((0, 0), (width, height), &mut buffer)
+ .map_err(|e| exec_datafusion_err!("Failed to initialize band: {}", e))?;
+ Ok(())
+}
+
+fn calc_num_iterations(args: &[ColumnarValue]) -> usize {
+ args.iter()
+ .find_map(|arg| match arg {
+ ColumnarValue::Array(array) => Some(array.len()),
+ ColumnarValue::Scalar(_) => None,
+ })
+ .unwrap_or(1)
+}
+
+fn finish_result(args: &[ColumnarValue], out: Arc<dyn Array>) -> Result<ColumnarValue> {
+ if args
+ .iter()
+ .any(|arg| matches!(arg, ColumnarValue::Array(_)))
+ {
+ Ok(ColumnarValue::Array(out))
+ } else {
+ Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(&out, 0)?))
+ }
+}
+
+fn execute_raster_arg<F>(
+ arg_types: &[SedonaType],
+ args: &[ColumnarValue],
+ index: usize,
+ num_iterations: usize,
+ mut func: F,
+) -> Result<()>
+where
+ F: FnMut(usize, Option<&RasterRefImpl<'_>>) -> Result<()>,
+{
+ if arg_types.get(index) != Some(&RASTER) {
+ return sedona_internal_err!("Argument {index} must be a raster type");
+ }
+
+ match &args[index] {
+ ColumnarValue::Array(array) => {
+ let raster_struct = array
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .ok_or_else(|| {
+ sedona_internal_datafusion_err!("Expected StructArray for raster data")
+ })?;
+ let raster_array = RasterStructArray::try_new(raster_struct)?;
+
+ for i in 0..num_iterations {
+ if raster_array.is_null(i) {
+ func(i, None)?;
+ } else {
+ let raster = raster_array.get(i)?;
+ func(i, Some(&raster))?;
+ }
+ }
+
+ Ok(())
+ }
+ ColumnarValue::Scalar(ScalarValue::Struct(raster_struct)) => {
+ let raster_array = RasterStructArray::try_new(raster_struct.as_ref())?;
+ let raster = if raster_array.is_null(0) {
+ None
+ } else {
+ Some(raster_array.get(0)?)
+ };
+
+ for i in 0..num_iterations {
+ func(i, raster.as_ref())?;
+ }
+
+ Ok(())
+ }
+ ColumnarValue::Scalar(ScalarValue::Null) => {
+ for i in 0..num_iterations {
+ func(i, None)?;
+ }
+ Ok(())
+ }
+ _ => sedona_internal_err!("Expected Struct scalar for raster"),
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow_array::{BooleanArray, Float64Array, StringArray};
+ use datafusion_expr::{ScalarUDF, ScalarUDFImpl};
+ use sedona_raster::array::RasterStructArray;
+ use sedona_schema::datatypes::RASTER;
+ use sedona_schema::datatypes::WKB_GEOMETRY;
+ use sedona_schema::raster::BandDataType;
+ use sedona_testing::raster_spec::{assert_rasters_equal, RasterSpec};
+ use sedona_testing::{create::create_array, testers::ScalarUdfTester};
+
+ use super::*;
+
+ fn reference_raster_spec() -> RasterSpec {
+ RasterSpec::d2(4, 3)
+ .transform([10.0, 2.0, 0.0, 20.0, 0.0, -2.0])
+ .crs(Some("EPSG:4326"))
+ .band_values(&[0u8; 12])
+ .nodata(0u8)
+ }
+
+ fn wkb_from_wkt(gdal: &Gdal, wkt: &str) -> Result<Vec<u8>> {
+ let geom = gdal.geometry_from_wkt(wkt).unwrap();
+ geom.wkb().map_err(|e| exec_datafusion_err!("{e}"))
+ }
+
+ fn bytes_to_f64_vec(bytes: &[u8]) -> Vec<f64> {
+ bytes
+ .chunks_exact(8)
+ .map(|chunk| f64::from_le_bytes(chunk.try_into().unwrap()))
+ .collect()
+ }
Review Comment:
I have seen both of these functions before. Can we use the previous versions of them?
##########
rust/sedona-raster-gdal/benches/rs_as_raster.rs:
##########
@@ -0,0 +1,169 @@
+// 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.
+
+//! Benchmarks for RS_AsRaster UDF.
+
+use std::{hint::black_box, sync::Arc};
+
+use arrow_array::{ArrayRef, BooleanArray, Float64Array, StringArray};
+use arrow_schema::DataType;
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
+use datafusion_expr::ScalarUDF;
+use sedona_gdal::global::with_global_gdal;
+use sedona_gdal::raster::types::Buffer;
+use sedona_raster::array::RasterStructArray;
+use sedona_raster::builder::RasterBuilder;
+use sedona_raster::traits::RasterRef;
+use sedona_schema::datatypes::{SedonaType, RASTER, WKB_GEOMETRY};
+use sedona_testing::{create::create_array, testers::ScalarUdfTester};
+
+fn base_raster() -> arrow_array::StructArray {
+ with_global_gdal(|gdal| {
+ let driver = gdal.get_driver_by_name("MEM").unwrap();
+ let dataset = driver.create_with_band_type::<u8>("", 4, 3, 1).unwrap();
+ dataset
+ .set_geo_transform(&[10.0, 2.0, 0.0, 20.0, 0.0, -2.0])
+ .unwrap();
+ dataset.set_projection("EPSG:4326").unwrap();
+ let band = dataset.rasterband(1).unwrap();
+ band.set_no_data_value(Some(0.0)).unwrap();
+ let mut buffer = Buffer::new((4, 3), vec![0u8; 12]);
+ band.write((0, 0), (4, 3), &mut buffer).unwrap();
+ sedona_raster_gdal::dataset_to_indb_raster(&dataset).unwrap()
+ })
+ .unwrap()
+}
Review Comment:
Can we use the `RasterSpec` to generate this a little more readably?
##########
rust/sedona-raster-gdal/src/rs_as_raster.rs:
##########
@@ -0,0 +1,702 @@
+// 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_AsRaster UDF - rasterize a vector geometry onto a raster grid.
+
+use std::sync::Arc;
+
+use arrow_array::{Array, StructArray};
+use arrow_schema::DataType;
+use datafusion_common::cast::{
+ as_binary_array, as_boolean_array, as_float64_array, as_string_array,
+};
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::error::Result;
+use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue};
+use datafusion_expr::{ColumnarValue, Volatility};
+use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err};
+use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
+use sedona_gdal::dataset::Dataset;
+use sedona_gdal::gdal::Gdal;
+use sedona_gdal::mem::MemDatasetBuilder;
+use sedona_gdal::raster::types::Buffer;
+use sedona_raster::array::{RasterRefImpl, RasterStructArray};
+use sedona_raster::builder::RasterBuilder;
+use sedona_raster::traits::{BandMetadata, RasterMetadata, RasterRef};
+use sedona_schema::datatypes::{SedonaType, RASTER};
+use sedona_schema::matchers::ArgMatcher;
+use sedona_schema::raster::{BandDataType, StorageType};
+
+use crate::gdal_common::{band_data_type_to_gdal, nodata_f64_to_bytes, with_gdal};
+use crate::gdal_dataset_provider::{configure_thread_local_options, thread_local_provider};
+
+pub fn rs_as_raster_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "rs_asraster",
+ vec![
+ Arc::new(RsAsRaster { arg_count: 3 }),
+ Arc::new(RsAsRaster { arg_count: 4 }),
+ Arc::new(RsAsRaster { arg_count: 5 }),
+ Arc::new(RsAsRaster { arg_count: 6 }),
+ Arc::new(RsAsRaster { arg_count: 7 }),
+ ],
+ Volatility::Immutable,
+ )
+}
+
+#[derive(Debug)]
+struct RsAsRaster {
+ arg_count: usize,
+}
+
+impl SedonaScalarKernel for RsAsRaster {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ let mut matchers = vec![
+ ArgMatcher::is_geometry_or_geography(),
+ ArgMatcher::is_raster(),
+ ArgMatcher::is_string(),
+ ];
+
+ if self.arg_count >= 4 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+ if self.arg_count >= 5 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 6 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 7 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+
+ 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 = calc_num_iterations(args);
+
+ let geom_array = args[0].clone().into_array(num_iterations)?;
+ let geom_array = as_binary_array(&geom_array)?;
+
+ let pixel_type_array = args[2]
+ .clone()
+ .cast_to(&DataType::Utf8, None)?
+ .into_array(num_iterations)?;
+ let pixel_type_array = as_string_array(&pixel_type_array)?;
+
+ let all_touched_array = if self.arg_count >= 4 {
+ args[3]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(false)).to_array_of_size(num_iterations)?
+ };
+ let all_touched_array = as_boolean_array(&all_touched_array)?;
+
+ let burn_value_array = if self.arg_count >= 5 {
+ args[4]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(Some(1.0)).to_array_of_size(num_iterations)?
+ };
+ let burn_value_array = as_float64_array(&burn_value_array)?;
+
+ let nodata_value_array = if self.arg_count >= 6 {
+ args[5]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(None).to_array_of_size(num_iterations)?
+ };
+ let nodata_value_array = as_float64_array(&nodata_value_array)?;
+
+ let use_geom_extent_array = if self.arg_count >= 7 {
+ args[6]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(true)).to_array_of_size(num_iterations)?
+ };
+ let use_geom_extent_array = as_boolean_array(&use_geom_extent_array)?;
+
+ let mut builder = RasterBuilder::new(num_iterations);
+
+ let mut geom_iter = geom_array.iter();
+ let mut pixel_type_iter = pixel_type_array.iter();
+ let mut all_touched_iter = all_touched_array.iter();
+ let mut burn_value_iter = burn_value_array.iter();
+ let mut nodata_value_iter = nodata_value_array.iter();
+ let mut use_geom_extent_iter = use_geom_extent_array.iter();
+
+ with_gdal(|gdal| {
+ configure_thread_local_options(gdal, config_options)?;
+
+ execute_raster_arg(arg_types, args, 1, num_iterations, |_, raster_opt| {
+ let geom_opt = geom_iter.next().unwrap();
+ let pixel_type_opt = pixel_type_iter.next().unwrap();
+ let all_touched_opt = all_touched_iter.next().unwrap().unwrap_or(false);
+ let burn_value_opt = burn_value_iter.next().unwrap().unwrap_or(1.0);
+ let nodata_value_opt = nodata_value_iter.next().unwrap();
+ let use_geometry_extent_opt = use_geom_extent_iter.next().unwrap().unwrap_or(true);
+
+ let raster = match raster_opt {
+ Some(raster) => raster,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+ let geom_wkb = match geom_opt {
+ Some(geom_wkb) => geom_wkb,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+ let pixel_type = match pixel_type_opt {
+ Some(pixel_type) => pixel_type,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+
+ let band_type = parse_pixel_type(pixel_type)?;
+
+ match as_raster(
+ gdal,
+ geom_wkb,
+ raster,
+ band_type,
+ all_touched_opt,
+ burn_value_opt,
+ nodata_value_opt,
+ use_geometry_extent_opt,
+ ) {
+ Ok((out_metadata, out_band_metadata, out_band_bytes)) => {
+ builder
+ .start_raster(&out_metadata, raster.crs())
+ .map_err(|e| {
+ exec_datafusion_err!("Failed to start output raster: {}", e)
+ })?;
+ builder.start_band(out_band_metadata).map_err(|e| {
+ exec_datafusion_err!("Failed to start output raster band: {}", e)
+ })?;
+ builder.band_data_writer().append_value(out_band_bytes);
+ builder.finish_band().map_err(|e| {
+ exec_datafusion_err!("Failed to finish output raster band: {}", e)
+ })?;
+ builder.finish_raster().map_err(|e| {
+ exec_datafusion_err!("Failed to finish output raster: {}", e)
+ })?;
+ }
+ Err(_) => builder.append_null()?,
+ }
+
+ Ok(())
+ })?;
+
+ finish_result(args, Arc::new(builder.finish()?))
+ })
+ }
+}
+
+fn parse_pixel_type(value: &str) -> Result<BandDataType> {
+ match value.trim().to_ascii_uppercase().as_str() {
+ "D" => Ok(BandDataType::Float64),
+ "F" => Ok(BandDataType::Float32),
+ "I" => Ok(BandDataType::Int32),
+ "S" => Ok(BandDataType::Int16),
+ "US" => Ok(BandDataType::UInt16),
+ "B" => Ok(BandDataType::UInt8),
+ "I8" | "INT8" => Ok(BandDataType::Int8),
+ "U64" | "UINT64" => Ok(BandDataType::UInt64),
+ "I64" | "INT64" => Ok(BandDataType::Int64),
+ other => exec_err!(
+ "Unsupported pixelType: {} (expected one of D, F, I, S, US, B, I8, U64, I64)",
+ other
+ ),
+ }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn as_raster(
+ gdal: &Gdal,
+ geom_wkb: &[u8],
+ reference_raster: &RasterRefImpl<'_>,
+ band_type: BandDataType,
+ all_touched: bool,
+ burn_value: f64,
+ nodata_value: Option<f64>,
+ use_geometry_extent: bool,
+) -> Result<(RasterMetadata, BandMetadata, Vec<u8>)> {
+ let ref_md = reference_raster.metadata();
+
+ if ref_md.skew_x() != 0.0 || ref_md.skew_y() != 0.0 {
+ return exec_err!(
+ "RS_AsRaster currently requires skew_x=0 and skew_y=0 in the reference raster"
+ );
+ }
+
+ let geometry = gdal
+ .geometry_from_wkb(geom_wkb)
+ .map_err(|e| exec_datafusion_err!("Failed to parse geometry from WKB: {}", e))?;
+
+ let (out_width, out_height, out_ulx, out_uly) = if use_geometry_extent {
+ let env = geometry.envelope();
+ let ulx = ref_md.upper_left_x();
+ let uly = ref_md.upper_left_y();
+ let scale_x = ref_md.scale_x();
+ let scale_y = ref_md.scale_y();
+
+ if scale_x == 0.0 || scale_y == 0.0 {
+ return exec_err!("Reference raster has zero scale");
+ }
+
+ let start_col = ((env.MinX - ulx) / scale_x).floor() as isize;
+ let end_col_excl = ((env.MaxX - ulx) / scale_x).ceil() as isize;
+ let start_row = ((env.MaxY - uly) / scale_y).floor() as isize;
+ let end_row_excl = ((env.MinY - uly) / scale_y).ceil() as isize;
+
+ let width = (end_col_excl - start_col).max(0) as usize;
+ let height = (end_row_excl - start_row).max(0) as usize;
+
+ if width == 0 || height == 0 {
+ return exec_err!("Geometry extent produced an empty raster");
+ }
+
+ (
+ width,
+ height,
+ ulx + (start_col as f64) * scale_x,
+ uly + (start_row as f64) * scale_y,
+ )
+ } else {
+ (
+ ref_md.width() as usize,
+ ref_md.height() as usize,
+ ref_md.upper_left_x(),
+ ref_md.upper_left_y(),
+ )
+ };
+
+ let out_dataset = MemDatasetBuilder::create(
+ gdal,
+ out_width,
+ out_height,
+ 1,
+ band_data_type_to_gdal(&band_type),
+ )
+ .map_err(|e| exec_datafusion_err!("Failed to create dataset: {}", e))?;
+
+ out_dataset
+ .set_geo_transform(&[
+ out_ulx,
+ ref_md.scale_x(),
+ ref_md.skew_x(),
+ out_uly,
+ ref_md.skew_y(),
+ ref_md.scale_y(),
+ ])
+ .map_err(|e| exec_datafusion_err!("Failed to set geotransform: {}", e))?;
+
+ let provider = thread_local_provider(gdal)
+ .map_err(|e| exec_datafusion_err!("Failed to init GDAL provider: {}", e))?;
+ let ref_raster_ds = provider
+ .raster_ref_to_gdal(reference_raster)
+ .map_err(|e| exec_datafusion_err!("Failed to create GDAL dataset: {}", e))?;
+ if let Ok(srs) = ref_raster_ds.as_dataset().spatial_ref() {
+ out_dataset
+ .set_spatial_ref(&srs)
+ .map_err(|e| exec_datafusion_err!("Failed to set spatial reference: {}", e))?;
+ }
+
+ let init_value = nodata_value.unwrap_or(0.0);
+ initialize_band(&out_dataset, &band_type, out_width, out_height, init_value)?;
+
+ if let Some(nodata) = nodata_value {
+ let band = out_dataset
+ .rasterband(1)
+ .map_err(|e| exec_datafusion_err!("Failed to get output band: {}", e))?;
+ match band_type {
+ BandDataType::UInt64 => band
+ .set_no_data_value_u64(Some(nodata as u64))
+ .map_err(|e| exec_datafusion_err!("Failed to set nodata value: {}", e))?,
+ BandDataType::Int64 => band
+ .set_no_data_value_i64(Some(nodata as i64))
+ .map_err(|e| exec_datafusion_err!("Failed to set nodata value: {}", e))?,
+ _ => band
+ .set_no_data_value(Some(nodata))
+ .map_err(|e| exec_datafusion_err!("Failed to set nodata value: {}", e))?,
+ }
+ }
+
+ gdal.rasterize_affine(&out_dataset, &[1], &[geometry], &[burn_value], all_touched)
+ .map_err(|e| exec_datafusion_err!("Failed to rasterize geometry: {}", e))?;
+
+ let band = out_dataset
+ .rasterband(1)
+ .map_err(|e| exec_datafusion_err!("Failed to get output band: {}", e))?;
+ let band_bytes = band
+ .read_as_bytes(
+ (0, 0),
+ (out_width, out_height),
+ (out_width, out_height),
+ None,
+ )
+ .map_err(|e| exec_datafusion_err!("Failed to read band data: {}", e))?;
+
+ Ok((
+ RasterMetadata {
+ width: out_width as i64,
+ height: out_height as i64,
+ upperleft_x: out_ulx,
+ upperleft_y: out_uly,
+ scale_x: ref_md.scale_x(),
+ scale_y: ref_md.scale_y(),
+ skew_x: ref_md.skew_x(),
+ skew_y: ref_md.skew_y(),
+ },
+ BandMetadata {
+ nodata_value: nodata_value.map(|value| nodata_f64_to_bytes(value, &band_type)),
+ storage_type: StorageType::InDb,
+ datatype: band_type,
+ outdb_url: None,
+ outdb_band_id: None,
+ },
+ band_bytes,
+ ))
+}
+
+fn initialize_band(
+ dataset: &Dataset,
+ band_type: &BandDataType,
+ width: usize,
+ height: usize,
+ init_value: f64,
+) -> Result<()> {
+ match band_type {
+ BandDataType::UInt8 => initialize_band_t::<u8>(dataset, width, height, init_value as u8),
+ BandDataType::Int8 => initialize_band_t::<i8>(dataset, width, height, init_value as i8),
+ BandDataType::UInt16 => initialize_band_t::<u16>(dataset, width, height, init_value as u16),
+ BandDataType::Int16 => initialize_band_t::<i16>(dataset, width, height, init_value as i16),
+ BandDataType::UInt32 => initialize_band_t::<u32>(dataset, width, height, init_value as u32),
+ BandDataType::Int32 => initialize_band_t::<i32>(dataset, width, height, init_value as i32),
+ BandDataType::UInt64 => initialize_band_t::<u64>(dataset, width, height, init_value as u64),
+ BandDataType::Int64 => initialize_band_t::<i64>(dataset, width, height, init_value as i64),
+ BandDataType::Float32 => {
+ initialize_band_t::<f32>(dataset, width, height, init_value as f32)
+ }
+ BandDataType::Float64 => initialize_band_t::<f64>(dataset, width, height, init_value),
+ }
+}
+
+fn initialize_band_t<T: sedona_gdal::raster::types::GdalType + Copy>(
+ dataset: &Dataset,
+ width: usize,
+ height: usize,
+ init_value: T,
+) -> Result<()> {
+ let band = dataset
+ .rasterband(1)
+ .map_err(|e| exec_datafusion_err!("Failed to get output band: {}", e))?;
+ let mut buffer = Buffer::new((width, height), vec![init_value; width * height]);
+ band.write((0, 0), (width, height), &mut buffer)
+ .map_err(|e| exec_datafusion_err!("Failed to initialize band: {}", e))?;
+ Ok(())
+}
+
+fn calc_num_iterations(args: &[ColumnarValue]) -> usize {
Review Comment:
This is fine for now but we should improve our pattern for iterating over pairs of raster(s)/geometry(ies) without expanding scalars.
##########
rust/sedona-raster-gdal/src/rs_as_raster.rs:
##########
@@ -0,0 +1,702 @@
+// 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_AsRaster UDF - rasterize a vector geometry onto a raster grid.
+
+use std::sync::Arc;
+
+use arrow_array::{Array, StructArray};
+use arrow_schema::DataType;
+use datafusion_common::cast::{
+ as_binary_array, as_boolean_array, as_float64_array, as_string_array,
+};
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::error::Result;
+use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue};
+use datafusion_expr::{ColumnarValue, Volatility};
+use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err};
+use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
+use sedona_gdal::dataset::Dataset;
+use sedona_gdal::gdal::Gdal;
+use sedona_gdal::mem::MemDatasetBuilder;
+use sedona_gdal::raster::types::Buffer;
+use sedona_raster::array::{RasterRefImpl, RasterStructArray};
+use sedona_raster::builder::RasterBuilder;
+use sedona_raster::traits::{BandMetadata, RasterMetadata, RasterRef};
+use sedona_schema::datatypes::{SedonaType, RASTER};
+use sedona_schema::matchers::ArgMatcher;
+use sedona_schema::raster::{BandDataType, StorageType};
+
+use crate::gdal_common::{band_data_type_to_gdal, nodata_f64_to_bytes, with_gdal};
+use crate::gdal_dataset_provider::{configure_thread_local_options, thread_local_provider};
+
+pub fn rs_as_raster_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "rs_asraster",
+ vec![
+ Arc::new(RsAsRaster { arg_count: 3 }),
+ Arc::new(RsAsRaster { arg_count: 4 }),
+ Arc::new(RsAsRaster { arg_count: 5 }),
+ Arc::new(RsAsRaster { arg_count: 6 }),
+ Arc::new(RsAsRaster { arg_count: 7 }),
+ ],
+ Volatility::Immutable,
+ )
+}
+
+#[derive(Debug)]
+struct RsAsRaster {
+ arg_count: usize,
+}
+
+impl SedonaScalarKernel for RsAsRaster {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ let mut matchers = vec![
+ ArgMatcher::is_geometry_or_geography(),
+ ArgMatcher::is_raster(),
+ ArgMatcher::is_string(),
+ ];
+
+ if self.arg_count >= 4 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+ if self.arg_count >= 5 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 6 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 7 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+
+ 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 = calc_num_iterations(args);
+
+ let geom_array = args[0].clone().into_array(num_iterations)?;
+ let geom_array = as_binary_array(&geom_array)?;
+
+ let pixel_type_array = args[2]
+ .clone()
+ .cast_to(&DataType::Utf8, None)?
+ .into_array(num_iterations)?;
+ let pixel_type_array = as_string_array(&pixel_type_array)?;
+
+ let all_touched_array = if self.arg_count >= 4 {
+ args[3]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(false)).to_array_of_size(num_iterations)?
+ };
+ let all_touched_array = as_boolean_array(&all_touched_array)?;
+
+ let burn_value_array = if self.arg_count >= 5 {
+ args[4]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(Some(1.0)).to_array_of_size(num_iterations)?
+ };
+ let burn_value_array = as_float64_array(&burn_value_array)?;
+
+ let nodata_value_array = if self.arg_count >= 6 {
+ args[5]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(None).to_array_of_size(num_iterations)?
+ };
+ let nodata_value_array = as_float64_array(&nodata_value_array)?;
+
+ let use_geom_extent_array = if self.arg_count >= 7 {
+ args[6]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(true)).to_array_of_size(num_iterations)?
+ };
+ let use_geom_extent_array = as_boolean_array(&use_geom_extent_array)?;
+
+ let mut builder = RasterBuilder::new(num_iterations);
+
+ let mut geom_iter = geom_array.iter();
+ let mut pixel_type_iter = pixel_type_array.iter();
+ let mut all_touched_iter = all_touched_array.iter();
+ let mut burn_value_iter = burn_value_array.iter();
+ let mut nodata_value_iter = nodata_value_array.iter();
+ let mut use_geom_extent_iter = use_geom_extent_array.iter();
+
+ with_gdal(|gdal| {
+ configure_thread_local_options(gdal, config_options)?;
+
+ execute_raster_arg(arg_types, args, 1, num_iterations, |_, raster_opt| {
+ let geom_opt = geom_iter.next().unwrap();
+ let pixel_type_opt = pixel_type_iter.next().unwrap();
+ let all_touched_opt = all_touched_iter.next().unwrap().unwrap_or(false);
+ let burn_value_opt = burn_value_iter.next().unwrap().unwrap_or(1.0);
+ let nodata_value_opt = nodata_value_iter.next().unwrap();
+ let use_geometry_extent_opt = use_geom_extent_iter.next().unwrap().unwrap_or(true);
+
+ let raster = match raster_opt {
+ Some(raster) => raster,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+ let geom_wkb = match geom_opt {
+ Some(geom_wkb) => geom_wkb,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+ let pixel_type = match pixel_type_opt {
+ Some(pixel_type) => pixel_type,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+
+ let band_type = parse_pixel_type(pixel_type)?;
+
+ match as_raster(
+ gdal,
+ geom_wkb,
+ raster,
+ band_type,
+ all_touched_opt,
+ burn_value_opt,
+ nodata_value_opt,
+ use_geometry_extent_opt,
+ ) {
+ Ok((out_metadata, out_band_metadata, out_band_bytes)) => {
+ builder
+ .start_raster(&out_metadata, raster.crs())
+ .map_err(|e| {
+ exec_datafusion_err!("Failed to start output raster: {}", e)
+ })?;
+ builder.start_band(out_band_metadata).map_err(|e| {
+ exec_datafusion_err!("Failed to start output raster band: {}", e)
+ })?;
+ builder.band_data_writer().append_value(out_band_bytes);
+ builder.finish_band().map_err(|e| {
+ exec_datafusion_err!("Failed to finish output raster band: {}", e)
+ })?;
+ builder.finish_raster().map_err(|e| {
+ exec_datafusion_err!("Failed to finish output raster: {}", e)
+ })?;
+ }
+ Err(_) => builder.append_null()?,
+ }
+
+ Ok(())
+ })?;
+
+ finish_result(args, Arc::new(builder.finish()?))
+ })
+ }
+}
+
+fn parse_pixel_type(value: &str) -> Result<BandDataType> {
+ match value.trim().to_ascii_uppercase().as_str() {
+ "D" => Ok(BandDataType::Float64),
+ "F" => Ok(BandDataType::Float32),
+ "I" => Ok(BandDataType::Int32),
+ "S" => Ok(BandDataType::Int16),
+ "US" => Ok(BandDataType::UInt16),
+ "B" => Ok(BandDataType::UInt8),
+ "I8" | "INT8" => Ok(BandDataType::Int8),
+ "U64" | "UINT64" => Ok(BandDataType::UInt64),
+ "I64" | "INT64" => Ok(BandDataType::Int64),
+ other => exec_err!(
+ "Unsupported pixelType: {} (expected one of D, F, I, S, US, B, I8, U64, I64)",
+ other
+ ),
+ }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn as_raster(
+ gdal: &Gdal,
+ geom_wkb: &[u8],
+ reference_raster: &RasterRefImpl<'_>,
+ band_type: BandDataType,
+ all_touched: bool,
+ burn_value: f64,
+ nodata_value: Option<f64>,
+ use_geometry_extent: bool,
+) -> Result<(RasterMetadata, BandMetadata, Vec<u8>)> {
+ let ref_md = reference_raster.metadata();
+
+ if ref_md.skew_x() != 0.0 || ref_md.skew_y() != 0.0 {
+ return exec_err!(
+ "RS_AsRaster currently requires skew_x=0 and skew_y=0 in the reference raster"
+ );
+ }
+
+ let geometry = gdal
+ .geometry_from_wkb(geom_wkb)
+ .map_err(|e| exec_datafusion_err!("Failed to parse geometry from WKB: {}", e))?;
+
+ let (out_width, out_height, out_ulx, out_uly) = if use_geometry_extent {
+ let env = geometry.envelope();
+ let ulx = ref_md.upper_left_x();
+ let uly = ref_md.upper_left_y();
+ let scale_x = ref_md.scale_x();
+ let scale_y = ref_md.scale_y();
+
+ if scale_x == 0.0 || scale_y == 0.0 {
+ return exec_err!("Reference raster has zero scale");
+ }
+
+ let start_col = ((env.MinX - ulx) / scale_x).floor() as isize;
+ let end_col_excl = ((env.MaxX - ulx) / scale_x).ceil() as isize;
+ let start_row = ((env.MaxY - uly) / scale_y).floor() as isize;
+ let end_row_excl = ((env.MinY - uly) / scale_y).ceil() as isize;
+
+ let width = (end_col_excl - start_col).max(0) as usize;
+ let height = (end_row_excl - start_row).max(0) as usize;
+
+ if width == 0 || height == 0 {
+ return exec_err!("Geometry extent produced an empty raster");
+ }
+
+ (
+ width,
+ height,
+ ulx + (start_col as f64) * scale_x,
+ uly + (start_row as f64) * scale_y,
+ )
+ } else {
+ (
+ ref_md.width() as usize,
+ ref_md.height() as usize,
+ ref_md.upper_left_x(),
+ ref_md.upper_left_y(),
+ )
+ };
+
+ let out_dataset = MemDatasetBuilder::create(
+ gdal,
+ out_width,
+ out_height,
+ 1,
+ band_data_type_to_gdal(&band_type),
+ )
+ .map_err(|e| exec_datafusion_err!("Failed to create dataset: {}", e))?;
+
+ out_dataset
+ .set_geo_transform(&[
+ out_ulx,
+ ref_md.scale_x(),
+ ref_md.skew_x(),
+ out_uly,
+ ref_md.skew_y(),
+ ref_md.scale_y(),
+ ])
+ .map_err(|e| exec_datafusion_err!("Failed to set geotransform: {}", e))?;
+
+ let provider = thread_local_provider(gdal)
+ .map_err(|e| exec_datafusion_err!("Failed to init GDAL provider: {}", e))?;
+ let ref_raster_ds = provider
+ .raster_ref_to_gdal(reference_raster)
+ .map_err(|e| exec_datafusion_err!("Failed to create GDAL dataset: {}", e))?;
+ if let Ok(srs) = ref_raster_ds.as_dataset().spatial_ref() {
+ out_dataset
+ .set_spatial_ref(&srs)
+ .map_err(|e| exec_datafusion_err!("Failed to set spatial reference: {}", e))?;
+ }
+
+ let init_value = nodata_value.unwrap_or(0.0);
+ initialize_band(&out_dataset, &band_type, out_width, out_height, init_value)?;
+
+ if let Some(nodata) = nodata_value {
+ let band = out_dataset
+ .rasterband(1)
+ .map_err(|e| exec_datafusion_err!("Failed to get output band: {}", e))?;
+ match band_type {
+ BandDataType::UInt64 => band
+ .set_no_data_value_u64(Some(nodata as u64))
+ .map_err(|e| exec_datafusion_err!("Failed to set nodata value: {}", e))?,
+ BandDataType::Int64 => band
+ .set_no_data_value_i64(Some(nodata as i64))
+ .map_err(|e| exec_datafusion_err!("Failed to set nodata value: {}", e))?,
+ _ => band
+ .set_no_data_value(Some(nodata))
+ .map_err(|e| exec_datafusion_err!("Failed to set nodata value: {}", e))?,
+ }
+ }
+
+ gdal.rasterize_affine(&out_dataset, &[1], &[geometry], &[burn_value], all_touched)
+ .map_err(|e| exec_datafusion_err!("Failed to rasterize geometry: {}", e))?;
+
+ let band = out_dataset
+ .rasterband(1)
+ .map_err(|e| exec_datafusion_err!("Failed to get output band: {}", e))?;
+ let band_bytes = band
+ .read_as_bytes(
+ (0, 0),
+ (out_width, out_height),
+ (out_width, out_height),
+ None,
+ )
+ .map_err(|e| exec_datafusion_err!("Failed to read band data: {}", e))?;
+
+ Ok((
+ RasterMetadata {
+ width: out_width as i64,
+ height: out_height as i64,
+ upperleft_x: out_ulx,
+ upperleft_y: out_uly,
+ scale_x: ref_md.scale_x(),
+ scale_y: ref_md.scale_y(),
+ skew_x: ref_md.skew_x(),
+ skew_y: ref_md.skew_y(),
+ },
+ BandMetadata {
+ nodata_value: nodata_value.map(|value| nodata_f64_to_bytes(value, &band_type)),
+ storage_type: StorageType::InDb,
+ datatype: band_type,
+ outdb_url: None,
+ outdb_band_id: None,
+ },
+ band_bytes,
+ ))
+}
+
+fn initialize_band(
+ dataset: &Dataset,
+ band_type: &BandDataType,
+ width: usize,
+ height: usize,
+ init_value: f64,
+) -> Result<()> {
+ match band_type {
+ BandDataType::UInt8 => initialize_band_t::<u8>(dataset, width, height, init_value as u8),
+ BandDataType::Int8 => initialize_band_t::<i8>(dataset, width, height, init_value as i8),
+ BandDataType::UInt16 => initialize_band_t::<u16>(dataset, width, height, init_value as u16),
+ BandDataType::Int16 => initialize_band_t::<i16>(dataset, width, height, init_value as i16),
+ BandDataType::UInt32 => initialize_band_t::<u32>(dataset, width, height, init_value as u32),
+ BandDataType::Int32 => initialize_band_t::<i32>(dataset, width, height, init_value as i32),
+ BandDataType::UInt64 => initialize_band_t::<u64>(dataset, width, height, init_value as u64),
+ BandDataType::Int64 => initialize_band_t::<i64>(dataset, width, height, init_value as i64),
Review Comment:
These two could generate potentially lossy values. It may be worth casting the nodata/init value based on the requested data type (or erroring for now until that can be supported).
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(rust/sedona-raster-gdal): add RS_AsRaster for rasterizing geometries [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #956:
URL: https://github.com/apache/sedona-db/pull/956#discussion_r3576190869
##########
rust/sedona-raster-gdal/src/rs_as_raster.rs:
##########
@@ -0,0 +1,702 @@
+// 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_AsRaster UDF - rasterize a vector geometry onto a raster grid.
+
+use std::sync::Arc;
+
+use arrow_array::{Array, StructArray};
+use arrow_schema::DataType;
+use datafusion_common::cast::{
+ as_binary_array, as_boolean_array, as_float64_array, as_string_array,
+};
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::error::Result;
+use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue};
+use datafusion_expr::{ColumnarValue, Volatility};
+use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err};
+use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
+use sedona_gdal::dataset::Dataset;
+use sedona_gdal::gdal::Gdal;
+use sedona_gdal::mem::MemDatasetBuilder;
+use sedona_gdal::raster::types::Buffer;
+use sedona_raster::array::{RasterRefImpl, RasterStructArray};
+use sedona_raster::builder::RasterBuilder;
+use sedona_raster::traits::{BandMetadata, RasterMetadata, RasterRef};
+use sedona_schema::datatypes::{SedonaType, RASTER};
+use sedona_schema::matchers::ArgMatcher;
+use sedona_schema::raster::{BandDataType, StorageType};
+
+use crate::gdal_common::{band_data_type_to_gdal, nodata_f64_to_bytes, with_gdal};
+use crate::gdal_dataset_provider::{configure_thread_local_options, thread_local_provider};
+
+pub fn rs_as_raster_udf() -> SedonaScalarUDF {
Review Comment:
I'm not sure I prefer it, but it seems like it might be a useful addition on the day that we can allocate the rasterizer and burn in values lazily as they come.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(rust/sedona-raster-gdal): add RS_AsRaster for rasterizing geometries [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #956:
URL: https://github.com/apache/sedona-db/pull/956#discussion_r3565846432
##########
rust/sedona-raster-gdal/src/rs_as_raster.rs:
##########
@@ -0,0 +1,702 @@
+// 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_AsRaster UDF - rasterize a vector geometry onto a raster grid.
+
+use std::sync::Arc;
+
+use arrow_array::{Array, StructArray};
+use arrow_schema::DataType;
+use datafusion_common::cast::{
+ as_binary_array, as_boolean_array, as_float64_array, as_string_array,
+};
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::error::Result;
+use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue};
+use datafusion_expr::{ColumnarValue, Volatility};
+use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err};
+use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
+use sedona_gdal::dataset::Dataset;
+use sedona_gdal::gdal::Gdal;
+use sedona_gdal::mem::MemDatasetBuilder;
+use sedona_gdal::raster::types::Buffer;
+use sedona_raster::array::{RasterRefImpl, RasterStructArray};
+use sedona_raster::builder::RasterBuilder;
+use sedona_raster::traits::{BandMetadata, RasterMetadata, RasterRef};
+use sedona_schema::datatypes::{SedonaType, RASTER};
+use sedona_schema::matchers::ArgMatcher;
+use sedona_schema::raster::{BandDataType, StorageType};
+
+use crate::gdal_common::{band_data_type_to_gdal, nodata_f64_to_bytes, with_gdal};
+use crate::gdal_dataset_provider::{configure_thread_local_options, thread_local_provider};
+
+pub fn rs_as_raster_udf() -> SedonaScalarUDF {
Review Comment:
Why do you prefer defining RS_AsRaster as an aggregate function? Is it because it is usually used for dealing with mosaics where we want to stitch polygons extracted on multiple tiles together?
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(rust/sedona-raster-gdal): add RS_AsRaster for rasterizing geometries [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #956:
URL: https://github.com/apache/sedona-db/pull/956#discussion_r3565880521
##########
rust/sedona-raster-gdal/src/rs_as_raster.rs:
##########
@@ -0,0 +1,702 @@
+// 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_AsRaster UDF - rasterize a vector geometry onto a raster grid.
+
+use std::sync::Arc;
+
+use arrow_array::{Array, StructArray};
+use arrow_schema::DataType;
+use datafusion_common::cast::{
+ as_binary_array, as_boolean_array, as_float64_array, as_string_array,
+};
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::error::Result;
+use datafusion_common::{exec_datafusion_err, exec_err, ScalarValue};
+use datafusion_expr::{ColumnarValue, Volatility};
+use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err};
+use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
+use sedona_gdal::dataset::Dataset;
+use sedona_gdal::gdal::Gdal;
+use sedona_gdal::mem::MemDatasetBuilder;
+use sedona_gdal::raster::types::Buffer;
+use sedona_raster::array::{RasterRefImpl, RasterStructArray};
+use sedona_raster::builder::RasterBuilder;
+use sedona_raster::traits::{BandMetadata, RasterMetadata, RasterRef};
+use sedona_schema::datatypes::{SedonaType, RASTER};
+use sedona_schema::matchers::ArgMatcher;
+use sedona_schema::raster::{BandDataType, StorageType};
+
+use crate::gdal_common::{band_data_type_to_gdal, nodata_f64_to_bytes, with_gdal};
+use crate::gdal_dataset_provider::{configure_thread_local_options, thread_local_provider};
+
+pub fn rs_as_raster_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "rs_asraster",
+ vec![
+ Arc::new(RsAsRaster { arg_count: 3 }),
+ Arc::new(RsAsRaster { arg_count: 4 }),
+ Arc::new(RsAsRaster { arg_count: 5 }),
+ Arc::new(RsAsRaster { arg_count: 6 }),
+ Arc::new(RsAsRaster { arg_count: 7 }),
+ ],
+ Volatility::Immutable,
+ )
+}
+
+#[derive(Debug)]
+struct RsAsRaster {
+ arg_count: usize,
+}
+
+impl SedonaScalarKernel for RsAsRaster {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ let mut matchers = vec![
+ ArgMatcher::is_geometry_or_geography(),
+ ArgMatcher::is_raster(),
+ ArgMatcher::is_string(),
+ ];
+
+ if self.arg_count >= 4 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+ if self.arg_count >= 5 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 6 {
+ matchers.push(ArgMatcher::is_numeric());
+ }
+ if self.arg_count >= 7 {
+ matchers.push(ArgMatcher::is_boolean());
+ }
+
+ 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 = calc_num_iterations(args);
+
+ let geom_array = args[0].clone().into_array(num_iterations)?;
+ let geom_array = as_binary_array(&geom_array)?;
+
+ let pixel_type_array = args[2]
+ .clone()
+ .cast_to(&DataType::Utf8, None)?
+ .into_array(num_iterations)?;
+ let pixel_type_array = as_string_array(&pixel_type_array)?;
+
+ let all_touched_array = if self.arg_count >= 4 {
+ args[3]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(false)).to_array_of_size(num_iterations)?
+ };
+ let all_touched_array = as_boolean_array(&all_touched_array)?;
+
+ let burn_value_array = if self.arg_count >= 5 {
+ args[4]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(Some(1.0)).to_array_of_size(num_iterations)?
+ };
+ let burn_value_array = as_float64_array(&burn_value_array)?;
+
+ let nodata_value_array = if self.arg_count >= 6 {
+ args[5]
+ .clone()
+ .cast_to(&DataType::Float64, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Float64(None).to_array_of_size(num_iterations)?
+ };
+ let nodata_value_array = as_float64_array(&nodata_value_array)?;
+
+ let use_geom_extent_array = if self.arg_count >= 7 {
+ args[6]
+ .clone()
+ .cast_to(&DataType::Boolean, None)?
+ .into_array(num_iterations)?
+ } else {
+ ScalarValue::Boolean(Some(true)).to_array_of_size(num_iterations)?
+ };
+ let use_geom_extent_array = as_boolean_array(&use_geom_extent_array)?;
+
+ let mut builder = RasterBuilder::new(num_iterations);
+
+ let mut geom_iter = geom_array.iter();
+ let mut pixel_type_iter = pixel_type_array.iter();
+ let mut all_touched_iter = all_touched_array.iter();
+ let mut burn_value_iter = burn_value_array.iter();
+ let mut nodata_value_iter = nodata_value_array.iter();
+ let mut use_geom_extent_iter = use_geom_extent_array.iter();
+
+ with_gdal(|gdal| {
+ configure_thread_local_options(gdal, config_options)?;
+
+ execute_raster_arg(arg_types, args, 1, num_iterations, |_, raster_opt| {
+ let geom_opt = geom_iter.next().unwrap();
+ let pixel_type_opt = pixel_type_iter.next().unwrap();
+ let all_touched_opt = all_touched_iter.next().unwrap().unwrap_or(false);
+ let burn_value_opt = burn_value_iter.next().unwrap().unwrap_or(1.0);
+ let nodata_value_opt = nodata_value_iter.next().unwrap();
+ let use_geometry_extent_opt = use_geom_extent_iter.next().unwrap().unwrap_or(true);
+
+ let raster = match raster_opt {
+ Some(raster) => raster,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+ let geom_wkb = match geom_opt {
+ Some(geom_wkb) => geom_wkb,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+ let pixel_type = match pixel_type_opt {
+ Some(pixel_type) => pixel_type,
+ None => {
+ builder.append_null()?;
+ return Ok(());
+ }
+ };
+
+ let band_type = parse_pixel_type(pixel_type)?;
+
+ match as_raster(
+ gdal,
+ geom_wkb,
+ raster,
+ band_type,
+ all_touched_opt,
+ burn_value_opt,
+ nodata_value_opt,
+ use_geometry_extent_opt,
+ ) {
+ Ok((out_metadata, out_band_metadata, out_band_bytes)) => {
+ builder
+ .start_raster(&out_metadata, raster.crs())
+ .map_err(|e| {
+ exec_datafusion_err!("Failed to start output raster: {}", e)
+ })?;
+ builder.start_band(out_band_metadata).map_err(|e| {
+ exec_datafusion_err!("Failed to start output raster band: {}", e)
+ })?;
+ builder.band_data_writer().append_value(out_band_bytes);
+ builder.finish_band().map_err(|e| {
+ exec_datafusion_err!("Failed to finish output raster band: {}", e)
+ })?;
+ builder.finish_raster().map_err(|e| {
+ exec_datafusion_err!("Failed to finish output raster: {}", e)
+ })?;
+ }
+ Err(_) => builder.append_null()?,
+ }
+
+ Ok(())
+ })?;
+
+ finish_result(args, Arc::new(builder.finish()?))
+ })
+ }
+}
+
+fn parse_pixel_type(value: &str) -> Result<BandDataType> {
+ match value.trim().to_ascii_uppercase().as_str() {
+ "D" => Ok(BandDataType::Float64),
+ "F" => Ok(BandDataType::Float32),
+ "I" => Ok(BandDataType::Int32),
+ "S" => Ok(BandDataType::Int16),
+ "US" => Ok(BandDataType::UInt16),
+ "B" => Ok(BandDataType::UInt8),
+ "I8" | "INT8" => Ok(BandDataType::Int8),
+ "U64" | "UINT64" => Ok(BandDataType::UInt64),
+ "I64" | "INT64" => Ok(BandDataType::Int64),
+ other => exec_err!(
+ "Unsupported pixelType: {} (expected one of D, F, I, S, US, B, I8, U64, I64)",
+ other
+ ),
+ }
+}
+
+#[allow(clippy::too_many_arguments)]
+fn as_raster(
+ gdal: &Gdal,
+ geom_wkb: &[u8],
+ reference_raster: &RasterRefImpl<'_>,
+ band_type: BandDataType,
+ all_touched: bool,
+ burn_value: f64,
+ nodata_value: Option<f64>,
+ use_geometry_extent: bool,
+) -> Result<(RasterMetadata, BandMetadata, Vec<u8>)> {
+ let ref_md = reference_raster.metadata();
+
+ if ref_md.skew_x() != 0.0 || ref_md.skew_y() != 0.0 {
+ return exec_err!(
+ "RS_AsRaster currently requires skew_x=0 and skew_y=0 in the reference raster"
+ );
+ }
+
+ let geometry = gdal
+ .geometry_from_wkb(geom_wkb)
+ .map_err(|e| exec_datafusion_err!("Failed to parse geometry from WKB: {}", e))?;
+
+ let (out_width, out_height, out_ulx, out_uly) = if use_geometry_extent {
+ let env = geometry.envelope();
+ let ulx = ref_md.upper_left_x();
+ let uly = ref_md.upper_left_y();
+ let scale_x = ref_md.scale_x();
+ let scale_y = ref_md.scale_y();
+
+ if scale_x == 0.0 || scale_y == 0.0 {
+ return exec_err!("Reference raster has zero scale");
+ }
+
+ let start_col = ((env.MinX - ulx) / scale_x).floor() as isize;
+ let end_col_excl = ((env.MaxX - ulx) / scale_x).ceil() as isize;
+ let start_row = ((env.MaxY - uly) / scale_y).floor() as isize;
+ let end_row_excl = ((env.MinY - uly) / scale_y).ceil() as isize;
+
+ let width = (end_col_excl - start_col).max(0) as usize;
+ let height = (end_row_excl - start_row).max(0) as usize;
+
+ if width == 0 || height == 0 {
+ return exec_err!("Geometry extent produced an empty raster");
+ }
+
+ (
+ width,
+ height,
+ ulx + (start_col as f64) * scale_x,
+ uly + (start_row as f64) * scale_y,
+ )
+ } else {
+ (
+ ref_md.width() as usize,
+ ref_md.height() as usize,
+ ref_md.upper_left_x(),
+ ref_md.upper_left_y(),
+ )
+ };
+
+ let out_dataset = MemDatasetBuilder::create(
+ gdal,
+ out_width,
+ out_height,
+ 1,
+ band_data_type_to_gdal(&band_type),
+ )
+ .map_err(|e| exec_datafusion_err!("Failed to create dataset: {}", e))?;
+
+ out_dataset
+ .set_geo_transform(&[
+ out_ulx,
+ ref_md.scale_x(),
+ ref_md.skew_x(),
+ out_uly,
+ ref_md.skew_y(),
+ ref_md.scale_y(),
+ ])
+ .map_err(|e| exec_datafusion_err!("Failed to set geotransform: {}", e))?;
+
+ let provider = thread_local_provider(gdal)
+ .map_err(|e| exec_datafusion_err!("Failed to init GDAL provider: {}", e))?;
+ let ref_raster_ds = provider
+ .raster_ref_to_gdal(reference_raster)
+ .map_err(|e| exec_datafusion_err!("Failed to create GDAL dataset: {}", e))?;
+ if let Ok(srs) = ref_raster_ds.as_dataset().spatial_ref() {
+ out_dataset
+ .set_spatial_ref(&srs)
+ .map_err(|e| exec_datafusion_err!("Failed to set spatial reference: {}", e))?;
+ }
+
+ let init_value = nodata_value.unwrap_or(0.0);
+ initialize_band(&out_dataset, &band_type, out_width, out_height, init_value)?;
+
+ if let Some(nodata) = nodata_value {
+ let band = out_dataset
+ .rasterband(1)
+ .map_err(|e| exec_datafusion_err!("Failed to get output band: {}", e))?;
+ match band_type {
+ BandDataType::UInt64 => band
+ .set_no_data_value_u64(Some(nodata as u64))
+ .map_err(|e| exec_datafusion_err!("Failed to set nodata value: {}", e))?,
+ BandDataType::Int64 => band
+ .set_no_data_value_i64(Some(nodata as i64))
+ .map_err(|e| exec_datafusion_err!("Failed to set nodata value: {}", e))?,
+ _ => band
+ .set_no_data_value(Some(nodata))
+ .map_err(|e| exec_datafusion_err!("Failed to set nodata value: {}", e))?,
+ }
+ }
+
+ gdal.rasterize_affine(&out_dataset, &[1], &[geometry], &[burn_value], all_touched)
+ .map_err(|e| exec_datafusion_err!("Failed to rasterize geometry: {}", e))?;
+
+ let band = out_dataset
+ .rasterband(1)
+ .map_err(|e| exec_datafusion_err!("Failed to get output band: {}", e))?;
+ let band_bytes = band
+ .read_as_bytes(
+ (0, 0),
+ (out_width, out_height),
+ (out_width, out_height),
+ None,
+ )
+ .map_err(|e| exec_datafusion_err!("Failed to read band data: {}", e))?;
+
+ Ok((
+ RasterMetadata {
+ width: out_width as i64,
+ height: out_height as i64,
+ upperleft_x: out_ulx,
+ upperleft_y: out_uly,
+ scale_x: ref_md.scale_x(),
+ scale_y: ref_md.scale_y(),
+ skew_x: ref_md.skew_x(),
+ skew_y: ref_md.skew_y(),
+ },
+ BandMetadata {
+ nodata_value: nodata_value.map(|value| nodata_f64_to_bytes(value, &band_type)),
+ storage_type: StorageType::InDb,
+ datatype: band_type,
+ outdb_url: None,
+ outdb_band_id: None,
+ },
+ band_bytes,
+ ))
+}
+
+fn initialize_band(
+ dataset: &Dataset,
+ band_type: &BandDataType,
+ width: usize,
+ height: usize,
+ init_value: f64,
+) -> Result<()> {
+ match band_type {
+ BandDataType::UInt8 => initialize_band_t::<u8>(dataset, width, height, init_value as u8),
+ BandDataType::Int8 => initialize_band_t::<i8>(dataset, width, height, init_value as i8),
+ BandDataType::UInt16 => initialize_band_t::<u16>(dataset, width, height, init_value as u16),
+ BandDataType::Int16 => initialize_band_t::<i16>(dataset, width, height, init_value as i16),
+ BandDataType::UInt32 => initialize_band_t::<u32>(dataset, width, height, init_value as u32),
+ BandDataType::Int32 => initialize_band_t::<i32>(dataset, width, height, init_value as i32),
+ BandDataType::UInt64 => initialize_band_t::<u64>(dataset, width, height, init_value as u64),
+ BandDataType::Int64 => initialize_band_t::<i64>(dataset, width, height, init_value as i64),
Review Comment:
The code added for casting in rs_as_raster.rs is massive. The newly added type `TypedBandValue` and the related functions seems to be useful for implementing other RS features. We can create one follow-up PR to extract these functions to a common module such as gdal_common.rs
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(rust/sedona-raster-gdal): add RS_AsRaster for rasterizing geometries [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #956:
URL: https://github.com/apache/sedona-db/pull/956#discussion_r3565875721
##########
docs/reference/sql/rs_asraster.qmd:
##########
@@ -0,0 +1,125 @@
+---
+# 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_AsRaster
+description: Rasterizes a geometry using the grid of a reference raster.
+kernels:
+ - returns: raster
+ args:
+ - name: geometry
+ type: geometry
+ - name: reference_raster
+ type: raster
+ - name: pixel_type
+ type: string
Review Comment:
I have marked this RS_AsRaster function as experimental as well.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(rust/sedona-raster-gdal): add RS_AsRaster for rasterizing geometries [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #956:
URL: https://github.com/apache/sedona-db/pull/956#discussion_r3566071240
##########
python/sedonadb/tests/functions/test_raster_functions.py:
##########
@@ -268,3 +268,191 @@ def test_rs_setbandnodatavalue_two_arg_requires_single_band():
SedonaDB().assert_query_result(
"SELECT RS_SetBandNoDataValue(RS_Example(), 0)", None
)
+
+
+def _rs_as_raster_sql(
+ pixel_type, all_touched, burn_value, nodata_value, use_geometry_extent
+):
+ return f"""
+ WITH src AS (SELECT RS_FromPath($1) AS raster)
+ SELECT RS_AsRaster(
+ ST_GeomFromText('POLYGON((2 8, 2 5, 5 5, 5 8, 2 8))'),
+ raster,
+ '{pixel_type}',
+ {str(all_touched).upper()},
+ {burn_value},
+ {nodata_value},
+ {str(use_geometry_extent).upper()}
+ ) AS raster
+ FROM src
+ """
+
+
+@pytest.mark.parametrize(
+ ("all_touched", "use_geometry_extent"),
+ [(False, False), (False, True), (True, False)],
+)
+def test_rs_as_raster_matches_rasterio(
+ con, sedona_testing, all_touched, use_geometry_extent
+):
+ import numpy as np
Review Comment:
Hoisted the local imports generated by LLM
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org