This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/sedona-db.git
The following commit(s) were added to refs/heads/main by this push:
new d5a36f0 feat(rust/sedona-raster-functions): Add a first raster
function - RS_Width (#268)
d5a36f0 is described below
commit d5a36f0f443d6df11e3522e42b625826c48dc3af
Author: jp <[email protected]>
AuthorDate: Wed Nov 5 10:57:45 2025 -0800
feat(rust/sedona-raster-functions): Add a first raster function - RS_Width
(#268)
---
Cargo.lock | 20 ++-
Cargo.toml | 1 +
.../Cargo.toml | 7 +-
rust/sedona-raster-functions/src/executor.rs | 187 +++++++++++++++++++++
.../src/lib.rs | 12 +-
.../src/register.rs} | 36 +++-
rust/sedona-raster-functions/src/rs_size.rs | 128 ++++++++++++++
rust/sedona-raster/Cargo.toml | 3 +
rust/sedona-raster/src/array.rs | 23 ++-
rust/sedona-raster/src/builder.rs | 2 +-
rust/sedona-schema/src/matchers.rs | 11 +-
rust/sedona-testing/Cargo.toml | 1 +
rust/sedona-testing/src/lib.rs | 1 +
rust/sedona-testing/src/rasters.rs | 118 +++++++++++++
14 files changed, 527 insertions(+), 23 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index f34de99..bf0e67e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1,6 +1,6 @@
# This file is automatically @generated by Cargo.
# It is not intended for manual editing.
-version = 3
+version = 4
[[package]]
name = "abi_stable"
@@ -5127,6 +5127,23 @@ dependencies = [
"arrow-schema",
"sedona-common",
"sedona-schema",
+ "sedona-testing",
+]
+
+[[package]]
+name = "sedona-raster-functions"
+version = "0.2.0"
+dependencies = [
+ "arrow-array",
+ "arrow-buffer",
+ "arrow-schema",
+ "datafusion-common",
+ "datafusion-expr",
+ "sedona-common",
+ "sedona-expr",
+ "sedona-raster",
+ "sedona-schema",
+ "sedona-testing",
]
[[package]]
@@ -5224,6 +5241,7 @@ dependencies = [
"sedona-common",
"sedona-expr",
"sedona-geometry",
+ "sedona-raster",
"sedona-schema",
"wkb",
"wkt 0.14.0",
diff --git a/Cargo.toml b/Cargo.toml
index 909676e..5d4bc5c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -31,6 +31,7 @@ members = [
"rust/sedona-geometry",
"rust/sedona-geoparquet",
"rust/sedona-raster",
+ "rust/sedona-raster-functions",
"rust/sedona-schema",
"rust/sedona-spatial-join",
"rust/sedona-testing",
diff --git a/rust/sedona-raster/Cargo.toml
b/rust/sedona-raster-functions/Cargo.toml
similarity index 83%
copy from rust/sedona-raster/Cargo.toml
copy to rust/sedona-raster-functions/Cargo.toml
index 91f35dc..d384e7e 100644
--- a/rust/sedona-raster/Cargo.toml
+++ b/rust/sedona-raster-functions/Cargo.toml
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
[package]
-name = "sedona-raster"
+name = "sedona-raster-functions"
version.workspace = true
homepage.workspace = true
repository.workspace = true
@@ -31,5 +31,10 @@ result_large_err = "allow"
arrow-schema = { workspace = true }
arrow-array = { workspace = true }
arrow-buffer = { workspace = true }
+datafusion-common = { workspace = true }
+datafusion-expr = { workspace = true }
sedona-common = { path = "../sedona-common" }
+sedona-expr = { path = "../sedona-expr" }
+sedona-raster = { path = "../sedona-raster" }
sedona-schema = { path = "../sedona-schema" }
+sedona-testing = { path = "../sedona-testing" }
diff --git a/rust/sedona-raster-functions/src/executor.rs
b/rust/sedona-raster-functions/src/executor.rs
new file mode 100644
index 0000000..75123e5
--- /dev/null
+++ b/rust/sedona-raster-functions/src/executor.rs
@@ -0,0 +1,187 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow_array::{Array, ArrayRef, StructArray};
+use datafusion_common::error::Result;
+use datafusion_common::{DataFusionError, ScalarValue};
+use datafusion_expr::ColumnarValue;
+use sedona_common::sedona_internal_err;
+use sedona_raster::array::{RasterRefImpl, RasterStructArray};
+use sedona_schema::datatypes::SedonaType;
+use sedona_schema::datatypes::RASTER;
+
+/// Helper for writing raster kernel implementations
+///
+/// The [RasterExecutor] provides a simplified interface for executing
functions
+/// on raster arrays, handling the common pattern of downcasting to
StructArray,
+/// creating raster iterators, and handling null values.
+pub struct RasterExecutor<'a, 'b> {
+ pub arg_types: &'a [SedonaType],
+ pub args: &'b [ColumnarValue],
+ num_iterations: usize,
+}
+
+impl<'a, 'b> RasterExecutor<'a, 'b> {
+ /// Create a new [RasterExecutor]
+ pub fn new(arg_types: &'a [SedonaType], args: &'b [ColumnarValue]) -> Self
{
+ Self {
+ arg_types,
+ args,
+ num_iterations: Self::calc_num_iterations(args),
+ }
+ }
+
+ /// Return the number of iterations that will be performed
+ pub fn num_iterations(&self) -> usize {
+ self.num_iterations
+ }
+
+ /// Execute a function by iterating over rasters in the first argument
+ ///
+ /// This handles the common pattern of:
+ /// 1. Downcasting array to StructArray
+ /// 2. Creating raster iterator
+ /// 3. Iterating with null checks
+ /// 4. Calling the provided function with each raster
+ pub fn execute_raster_void<F>(&self, mut func: F) -> Result<()>
+ where
+ F: FnMut(usize, Option<RasterRefImpl<'_>>) -> Result<()>,
+ {
+ if self.arg_types[0] != RASTER {
+ return sedona_internal_err!("First argument must be a raster
type");
+ }
+ let raster_array = match &self.args[0] {
+ ColumnarValue::Array(array) => array,
+ ColumnarValue::Scalar(_) => {
+ return Err(DataFusionError::NotImplemented(
+ "Scalar raster input not yet supported".to_string(),
+ ));
+ }
+ };
+
+ // Downcast to StructArray (rasters are stored as structs)
+ let raster_struct = raster_array
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .ok_or_else(|| {
+ DataFusionError::Internal("Expected StructArray for raster
data".to_string())
+ })?;
+
+ // Create raster iterator
+ let raster_array = RasterStructArray::new(raster_struct);
+
+ // Iterate through each raster in the array
+ for i in 0..self.num_iterations {
+ if raster_array.is_null(i) {
+ func(i, None)?;
+ continue;
+ }
+ let raster = raster_array.get(i)?;
+
+ func(i, Some(raster))?;
+ }
+
+ Ok(())
+ }
+
+ /// Finish an [ArrayRef] output as the appropriate [ColumnarValue]
+ ///
+ /// Converts the output into a [ColumnarValue::Scalar] if all arguments
were scalars,
+ /// or a [ColumnarValue::Array] otherwise.
+ pub fn finish(&self, out: ArrayRef) -> Result<ColumnarValue> {
+ for arg in self.args {
+ match arg {
+ // If any argument was an array, we return an array
+ ColumnarValue::Array(_) => {
+ return Ok(ColumnarValue::Array(out));
+ }
+ ColumnarValue::Scalar(_) => {}
+ }
+ }
+
+ // For all scalar arguments, we return a scalar
+ Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(&out, 0)?))
+ }
+
+ /// Calculates the number of iterations that should happen based on the
+ /// argument ColumnarValue types
+ fn calc_num_iterations(args: &[ColumnarValue]) -> usize {
+ for arg in args {
+ match arg {
+ // If any argument is an array, we have to iterate array.len()
times
+ ColumnarValue::Array(array) => {
+ return array.len();
+ }
+ ColumnarValue::Scalar(_) => {}
+ }
+ }
+
+ // All scalars: we iterate once
+ 1
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::builder::UInt64Builder;
+ use arrow_array::UInt64Array;
+ use sedona_raster::traits::RasterRef;
+ use sedona_schema::datatypes::RASTER;
+ use sedona_testing::rasters::generate_test_rasters;
+ use std::sync::Arc;
+
+ #[test]
+ fn test_raster_executor_execute_raster_void() {
+ // 3 rasters, second one is null
+ let rasters = generate_test_rasters(3, Some(1)).unwrap();
+ let args = [ColumnarValue::Array(Arc::new(rasters))];
+ let arg_types = vec![RASTER];
+
+ let executor = RasterExecutor::new(&arg_types, &args);
+ assert_eq!(executor.num_iterations(), 3);
+
+ let mut builder =
UInt64Builder::with_capacity(executor.num_iterations());
+ executor
+ .execute_raster_void(|_i, raster_opt| {
+ match raster_opt {
+ None => builder.append_null(),
+ Some(raster) => {
+ let width = raster.metadata().width();
+ builder.append_value(width);
+ }
+ }
+ Ok(())
+ })
+ .unwrap();
+
+ let result = executor.finish(Arc::new(builder.finish())).unwrap();
+
+ let width_array = match &result {
+ ColumnarValue::Array(array) => array
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .expect("Expected UInt64Array"),
+ ColumnarValue::Scalar(_) => panic!("Expected array, got scalar"),
+ };
+
+ assert_eq!(width_array.len(), 3);
+ assert_eq!(width_array.value(0), 1);
+ assert!(width_array.is_null(1));
+ assert_eq!(width_array.value(2), 3);
+ }
+}
diff --git a/rust/sedona-testing/src/lib.rs
b/rust/sedona-raster-functions/src/lib.rs
similarity index 85%
copy from rust/sedona-testing/src/lib.rs
copy to rust/sedona-raster-functions/src/lib.rs
index 6652955..86aea00 100644
--- a/rust/sedona-testing/src/lib.rs
+++ b/rust/sedona-raster-functions/src/lib.rs
@@ -14,11 +14,7 @@
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
-pub mod benchmark_util;
-pub mod compare;
-pub mod create;
-pub mod data;
-pub mod datagen;
-pub mod fixtures;
-pub mod read;
-pub mod testers;
+
+mod executor;
+pub mod register;
+pub mod rs_size;
diff --git a/rust/sedona-testing/src/lib.rs
b/rust/sedona-raster-functions/src/register.rs
similarity index 51%
copy from rust/sedona-testing/src/lib.rs
copy to rust/sedona-raster-functions/src/register.rs
index 6652955..7499892 100644
--- a/rust/sedona-testing/src/lib.rs
+++ b/rust/sedona-raster-functions/src/register.rs
@@ -14,11 +14,31 @@
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
-pub mod benchmark_util;
-pub mod compare;
-pub mod create;
-pub mod data;
-pub mod datagen;
-pub mod fixtures;
-pub mod read;
-pub mod testers;
+use sedona_expr::function_set::FunctionSet;
+
+/// Export the set of functions defined in this crate
+pub fn default_function_set() -> FunctionSet {
+ let mut function_set = FunctionSet::new();
+
+ macro_rules! register_scalar_udfs {
+ ($function_set:expr, $($udf:expr),* $(,)?) => {
+ $(
+ $function_set.insert_scalar_udf($udf());
+ )*
+ };
+ }
+
+ macro_rules! register_aggregate_udfs {
+ ($function_set:expr, $($udf:expr),* $(,)?) => {
+ $(
+ $function_set.insert_aggregate_udf($udf());
+ )*
+ };
+ }
+
+ register_scalar_udfs!(function_set, crate::rs_size::rs_width_udf,);
+
+ register_aggregate_udfs!(function_set,);
+
+ function_set
+}
diff --git a/rust/sedona-raster-functions/src/rs_size.rs
b/rust/sedona-raster-functions/src/rs_size.rs
new file mode 100644
index 0000000..6b6d200
--- /dev/null
+++ b/rust/sedona-raster-functions/src/rs_size.rs
@@ -0,0 +1,128 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+use std::{sync::Arc, vec};
+
+use crate::executor::RasterExecutor;
+use arrow_array::builder::UInt64Builder;
+use arrow_schema::DataType;
+use datafusion_common::error::Result;
+use datafusion_expr::{
+ scalar_doc_sections::DOC_SECTION_OTHER, ColumnarValue, Documentation,
Volatility,
+};
+use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF};
+use sedona_raster::traits::RasterRef;
+use sedona_schema::{datatypes::SedonaType, matchers::ArgMatcher};
+
+/// RS_Width() scalar UDF implementation
+///
+/// Extract the width of the raster
+pub fn rs_width_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "rs_width",
+ vec![Arc::new(RsWidth {})],
+ Volatility::Immutable,
+ Some(rs_width_doc()),
+ )
+}
+
+fn rs_width_doc() -> Documentation {
+ Documentation::builder(
+ DOC_SECTION_OTHER,
+ "Return the width component of a raster".to_string(),
+ "RS_Width(raster: Raster)".to_string(),
+ )
+ .with_argument("raster", "Raster: Input raster")
+ .with_sql_example("SELECT RS_Width(raster)".to_string())
+ .build()
+}
+
+#[derive(Debug)]
+struct RsWidth {}
+
+impl SedonaScalarKernel for RsWidth {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ let matcher = ArgMatcher::new(
+ vec![ArgMatcher::is_raster()],
+ SedonaType::Arrow(DataType::UInt64),
+ );
+
+ matcher.match_args(args)
+ }
+
+ fn invoke_batch(
+ &self,
+ arg_types: &[SedonaType],
+ args: &[ColumnarValue],
+ ) -> Result<ColumnarValue> {
+ let executor = RasterExecutor::new(arg_types, args);
+ let mut builder =
UInt64Builder::with_capacity(executor.num_iterations());
+
+ executor.execute_raster_void(|_i, raster_opt| {
+ match raster_opt {
+ None => builder.append_null(),
+ Some(raster) => {
+ let width = raster.metadata().width();
+ builder.append_value(width);
+ }
+ }
+ Ok(())
+ })?;
+
+ executor.finish(Arc::new(builder.finish()))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::{Array, UInt64Array};
+ use datafusion_expr::ScalarUDF;
+ use sedona_schema::datatypes::RASTER;
+ use sedona_testing::rasters::generate_test_rasters;
+
+ #[test]
+ fn udf_size() {
+ let udf: ScalarUDF = rs_width_udf().into();
+ assert_eq!(udf.name(), "rs_width");
+ assert!(udf.documentation().is_some());
+ }
+
+ #[test]
+ fn udf_invoke() {
+ // 3 rasters, second one is null
+ let rasters = generate_test_rasters(3, Some(1)).unwrap();
+
+ // Create the UDF and invoke it
+ let kernel = RsWidth {};
+ let args = [ColumnarValue::Array(Arc::new(rasters))];
+ let arg_types = vec![RASTER];
+
+ let result = kernel.invoke_batch(&arg_types, &args).unwrap();
+
+ // Check the result
+ if let ColumnarValue::Array(result_array) = result {
+ let width_array =
result_array.as_any().downcast_ref::<UInt64Array>().unwrap();
+
+ assert_eq!(width_array.len(), 3);
+ assert_eq!(width_array.value(0), 1); // First raster width
+ assert!(width_array.is_null(1)); // Second raster is null
+ assert_eq!(width_array.value(2), 3); // Third raster width
+ } else {
+ panic!("Expected array result");
+ }
+ }
+}
diff --git a/rust/sedona-raster/Cargo.toml b/rust/sedona-raster/Cargo.toml
index 91f35dc..1a8562f 100644
--- a/rust/sedona-raster/Cargo.toml
+++ b/rust/sedona-raster/Cargo.toml
@@ -33,3 +33,6 @@ arrow-array = { workspace = true }
arrow-buffer = { workspace = true }
sedona-common = { path = "../sedona-common" }
sedona-schema = { path = "../sedona-schema" }
+
+[dev-dependencies]
+sedona-testing = { path = "../sedona-testing" }
diff --git a/rust/sedona-raster/src/array.rs b/rust/sedona-raster/src/array.rs
index 6e572e4..f7236a3 100644
--- a/rust/sedona-raster/src/array.rs
+++ b/rust/sedona-raster/src/array.rs
@@ -467,12 +467,19 @@ impl<'a> RasterStructArray<'a> {
}
/// Get a specific raster by index without consuming the iterator
- pub fn get(&self, index: usize) -> Option<RasterRefImpl<'a>> {
+ pub fn get(&self, index: usize) -> Result<RasterRefImpl<'a>, ArrowError> {
if index >= self.raster_array.len() {
- return None;
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "Invalid raster index: {}",
+ index
+ )));
}
- Some(RasterRefImpl::new(self.raster_array, index))
+ Ok(RasterRefImpl::new(self.raster_array, index))
+ }
+
+ pub fn is_null(&self, index: usize) -> bool {
+ self.raster_array.is_null(index)
}
}
@@ -482,6 +489,7 @@ mod tests {
use crate::builder::RasterBuilder;
use crate::traits::{BandMetadata, RasterMetadata};
use sedona_schema::raster::{BandDataType, StorageType};
+ use sedona_testing::rasters::generate_test_rasters;
#[test]
fn test_array_basic_functionality() {
@@ -621,4 +629,13 @@ mod tests {
assert_eq!(band_values, vec![0, 1, 2]);
}
+
+ #[test]
+ fn test_raster_is_null() {
+ let raster_array = generate_test_rasters(2, Some(1)).unwrap();
+ let rasters = RasterStructArray::new(&raster_array);
+ assert_eq!(rasters.len(), 2);
+ assert!(!rasters.is_null(0));
+ assert!(rasters.is_null(1));
+ }
}
diff --git a/rust/sedona-raster/src/builder.rs
b/rust/sedona-raster/src/builder.rs
index 357486a..3ad1e37 100644
--- a/rust/sedona-raster/src/builder.rs
+++ b/rust/sedona-raster/src/builder.rs
@@ -256,7 +256,7 @@ impl RasterBuilder {
self.band_offsets.push(current_offset);
// Mark raster as null
- self.raster_validity.append_value(false);
+ self.raster_validity.append_null();
Ok(())
}
diff --git a/rust/sedona-schema/src/matchers.rs
b/rust/sedona-schema/src/matchers.rs
index 4935f43..ca7536c 100644
--- a/rust/sedona-schema/src/matchers.rs
+++ b/rust/sedona-schema/src/matchers.rs
@@ -21,7 +21,7 @@ use arrow_schema::DataType;
use datafusion_common::{plan_err, Result};
use sedona_common::sedona_internal_err;
-use crate::datatypes::{Edges, SedonaType, WKB_GEOGRAPHY, WKB_GEOMETRY};
+use crate::datatypes::{Edges, SedonaType, RASTER, WKB_GEOGRAPHY, WKB_GEOMETRY};
/// Helper to match arguments and compute return types
#[derive(Debug)]
@@ -173,6 +173,11 @@ impl ArgMatcher {
Arc::new(IsGeography {})
}
+ /// Matches any raster argument
+ pub fn is_raster() -> Arc<dyn TypeMatcher + Send + Sync> {
+ Self::is_exact(RASTER)
+ }
+
/// Matches a null argument
pub fn is_null() -> Arc<dyn TypeMatcher + Send + Sync> {
Arc::new(IsNull {})
@@ -506,6 +511,10 @@ mod tests {
ArgMatcher::is_boolean().type_if_null(),
Some(SedonaType::Arrow(DataType::Boolean))
);
+
+ assert!(ArgMatcher::is_raster().match_type(&RASTER));
+
assert!(!ArgMatcher::is_raster().match_type(&SedonaType::Arrow(DataType::Int32)));
+ assert!(!ArgMatcher::is_raster().match_type(&WKB_GEOMETRY));
}
#[test]
diff --git a/rust/sedona-testing/Cargo.toml b/rust/sedona-testing/Cargo.toml
index d0c058f..f467aad 100644
--- a/rust/sedona-testing/Cargo.toml
+++ b/rust/sedona-testing/Cargo.toml
@@ -50,6 +50,7 @@ rand = { workspace = true }
sedona-common = { path = "../sedona-common" }
sedona-geometry = { path = "../sedona-geometry" }
sedona-expr = { path = "../sedona-expr" }
+sedona-raster = { path = "../sedona-raster" }
sedona-schema = { path = "../sedona-schema" }
wkb = { workspace = true }
wkt = { workspace = true }
diff --git a/rust/sedona-testing/src/lib.rs b/rust/sedona-testing/src/lib.rs
index 6652955..6fcd0d6 100644
--- a/rust/sedona-testing/src/lib.rs
+++ b/rust/sedona-testing/src/lib.rs
@@ -20,5 +20,6 @@ pub mod create;
pub mod data;
pub mod datagen;
pub mod fixtures;
+pub mod rasters;
pub mod read;
pub mod testers;
diff --git a/rust/sedona-testing/src/rasters.rs
b/rust/sedona-testing/src/rasters.rs
new file mode 100644
index 0000000..826024f
--- /dev/null
+++ b/rust/sedona-testing/src/rasters.rs
@@ -0,0 +1,118 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+use arrow_array::StructArray;
+use arrow_schema::ArrowError;
+use sedona_raster::builder::RasterBuilder;
+use sedona_raster::traits::{BandMetadata, RasterMetadata};
+use sedona_schema::raster::{BandDataType, StorageType};
+
+/// Generate a StructArray of rasters with sequentially increasing dimensions
and pixel values
+/// These tiny rasters are to provide fast, easy and predictable test data for
unit tests.
+pub fn generate_test_rasters(
+ count: usize,
+ null_raster_index: Option<usize>,
+) -> Result<StructArray, ArrowError> {
+ let mut builder = RasterBuilder::new(count);
+ for i in 0..count {
+ // If a null raster index is specified and that matches the current
index,
+ // append a null raster
+ if matches!(null_raster_index, Some(index) if index == i) {
+ builder.append_null()?;
+ continue;
+ }
+
+ let raster_metadata = RasterMetadata {
+ width: i as u64 + 1,
+ height: i as u64 + 2,
+ upperleft_x: i as f64 + 1.0,
+ upperleft_y: i as f64 + 2.0,
+ scale_x: i as f64 * 0.1,
+ scale_y: i as f64 * 0.2,
+ skew_x: i as f64 * 0.3,
+ skew_y: i as f64 * 0.4,
+ };
+ builder.start_raster(&raster_metadata, None)?;
+ builder.start_band(BandMetadata {
+ datatype: BandDataType::UInt16,
+ nodata_value: Some(vec![0u8; 2]),
+ storage_type: StorageType::InDb,
+ outdb_url: None,
+ outdb_band_id: None,
+ })?;
+
+ let pixel_count = (i + 1) * (i + 2); // width * height
+ let mut band_data = Vec::with_capacity(pixel_count * 2); // 2 bytes
per u16
+ for pixel_value in 0..pixel_count as u16 {
+ band_data.extend_from_slice(&pixel_value.to_le_bytes());
+ }
+
+ builder.band_data_writer().append_value(&band_data);
+ builder.finish_band()?;
+ builder.finish_raster()?;
+ }
+
+ builder.finish()
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use sedona_raster::array::RasterStructArray;
+ use sedona_raster::traits::RasterRef;
+
+ #[test]
+ fn test_generate_test_rasters() {
+ let count = 5;
+ let struct_array = generate_test_rasters(count, None).unwrap();
+ let raster_array = RasterStructArray::new(&struct_array);
+ assert_eq!(raster_array.len(), count);
+
+ for i in 0..count {
+ let raster = raster_array.get(i).unwrap();
+ let metadata = raster.metadata();
+ assert_eq!(metadata.width(), i as u64 + 1);
+ assert_eq!(metadata.height(), i as u64 + 2);
+ assert_eq!(metadata.upper_left_x(), i as f64 + 1.0);
+ assert_eq!(metadata.upper_left_y(), i as f64 + 2.0);
+ assert_eq!(metadata.scale_x(), (i as f64) * 0.1);
+ assert_eq!(metadata.scale_y(), (i as f64) * 0.2);
+ assert_eq!(metadata.skew_x(), (i as f64) * 0.3);
+ assert_eq!(metadata.skew_y(), (i as f64) * 0.4);
+
+ let bands = raster.bands();
+ let band = bands.band(1).unwrap();
+ let band_metadata = band.metadata();
+ assert_eq!(band_metadata.data_type(), BandDataType::UInt16);
+ assert_eq!(band_metadata.nodata_value(), Some(&[0u8, 0u8][..]));
+ assert_eq!(band_metadata.storage_type(), StorageType::InDb);
+ assert_eq!(band_metadata.outdb_url(), None);
+ assert_eq!(band_metadata.outdb_band_id(), None);
+
+ let band_data = band.data();
+ let expected_pixel_count = (i + 1) * (i + 2); // width * height
+
+ // Convert raw bytes back to u16 values for comparison
+ let mut actual_pixel_values = Vec::new();
+ for chunk in band_data.chunks_exact(2) {
+ let value = u16::from_le_bytes([chunk[0], chunk[1]]);
+ actual_pixel_values.push(value);
+ }
+ let expected_pixel_values: Vec<u16> = (0..expected_pixel_count as
u16).collect();
+ assert_eq!(actual_pixel_values, expected_pixel_values);
+ }
+ }
+}
