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 faa3ed72 feat(rust/sedona-raster-functions): Add RS_Height (#302)
faa3ed72 is described below
commit faa3ed722d7505533cdaca323c3c20d2ff1d6606
Author: jp <[email protected]>
AuthorDate: Sat Nov 15 13:34:28 2025 -0800
feat(rust/sedona-raster-functions): Add RS_Height (#302)
Co-authored-by: Dewey Dunnington <[email protected]>
---
Cargo.lock | 2 +
rust/sedona-raster-functions/Cargo.toml | 1 +
.../benches/native-raster-functions.rs | 1 +
rust/sedona-raster-functions/src/register.rs | 3 +-
rust/sedona-raster-functions/src/rs_size.rs | 85 ++++++++++++++++++----
rust/sedona/Cargo.toml | 1 +
rust/sedona/src/context.rs | 3 +
7 files changed, 80 insertions(+), 16 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index b6ad1c25..45e771cd 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4839,6 +4839,7 @@ dependencies = [
"sedona-geoparquet",
"sedona-geos",
"sedona-proj",
+ "sedona-raster-functions",
"sedona-s2geography",
"sedona-schema",
"sedona-spatial-join",
@@ -5164,6 +5165,7 @@ dependencies = [
"criterion",
"datafusion-common",
"datafusion-expr",
+ "rstest",
"sedona-common",
"sedona-expr",
"sedona-raster",
diff --git a/rust/sedona-raster-functions/Cargo.toml
b/rust/sedona-raster-functions/Cargo.toml
index d0a2ea17..b674b136 100644
--- a/rust/sedona-raster-functions/Cargo.toml
+++ b/rust/sedona-raster-functions/Cargo.toml
@@ -41,6 +41,7 @@ sedona-schema = { path = "../sedona-schema" }
[dev-dependencies]
criterion = { workspace = true }
sedona-testing = { path = "../sedona-testing", features = ["criterion"] }
+rstest = { workspace = true }
[[bench]]
harness = false
diff --git a/rust/sedona-raster-functions/benches/native-raster-functions.rs
b/rust/sedona-raster-functions/benches/native-raster-functions.rs
index d729ccc5..3d1901f8 100644
--- a/rust/sedona-raster-functions/benches/native-raster-functions.rs
+++ b/rust/sedona-raster-functions/benches/native-raster-functions.rs
@@ -20,6 +20,7 @@ use sedona_testing::benchmark_util::{benchmark,
BenchmarkArgSpec::*};
fn criterion_benchmark(c: &mut Criterion) {
let f = sedona_raster_functions::register::default_function_set();
+ benchmark::scalar(c, &f, "native", "rs_height", Raster(64, 64));
benchmark::scalar(c, &f, "native", "rs_width", Raster(64, 64));
}
diff --git a/rust/sedona-raster-functions/src/register.rs
b/rust/sedona-raster-functions/src/register.rs
index 1945372d..fca2f862 100644
--- a/rust/sedona-raster-functions/src/register.rs
+++ b/rust/sedona-raster-functions/src/register.rs
@@ -38,8 +38,9 @@ pub fn default_function_set() -> FunctionSet {
register_scalar_udfs!(
function_set,
- crate::rs_size::rs_width_udf,
crate::rs_example::rs_example_udf,
+ crate::rs_size::rs_height_udf,
+ crate::rs_size::rs_width_udf,
);
register_aggregate_udfs!(function_set,);
diff --git a/rust/sedona-raster-functions/src/rs_size.rs
b/rust/sedona-raster-functions/src/rs_size.rs
index 6b6d200f..ff16edc9 100644
--- a/rust/sedona-raster-functions/src/rs_size.rs
+++ b/rust/sedona-raster-functions/src/rs_size.rs
@@ -33,12 +33,28 @@ use sedona_schema::{datatypes::SedonaType,
matchers::ArgMatcher};
pub fn rs_width_udf() -> SedonaScalarUDF {
SedonaScalarUDF::new(
"rs_width",
- vec![Arc::new(RsWidth {})],
+ vec![Arc::new(RsSize {
+ size_type: SizeType::Width,
+ })],
Volatility::Immutable,
Some(rs_width_doc()),
)
}
+/// RS_Height() scalar UDF documentation
+///
+/// Extract the height of the raster
+pub fn rs_height_udf() -> SedonaScalarUDF {
+ SedonaScalarUDF::new(
+ "rs_height",
+ vec![Arc::new(RsSize {
+ size_type: SizeType::Height,
+ })],
+ Volatility::Immutable,
+ Some(rs_height_doc()),
+ )
+}
+
fn rs_width_doc() -> Documentation {
Documentation::builder(
DOC_SECTION_OTHER,
@@ -50,10 +66,29 @@ fn rs_width_doc() -> Documentation {
.build()
}
+fn rs_height_doc() -> Documentation {
+ Documentation::builder(
+ DOC_SECTION_OTHER,
+ "Return the height component of a raster".to_string(),
+ "RS_Height(raster: Raster)".to_string(),
+ )
+ .with_argument("raster", "Raster: Input raster")
+ .with_sql_example("SELECT RS_Height(raster)".to_string())
+ .build()
+}
+
+#[derive(Debug, Clone)]
+enum SizeType {
+ Width,
+ Height,
+}
+
#[derive(Debug)]
-struct RsWidth {}
+struct RsSize {
+ size_type: SizeType,
+}
-impl SedonaScalarKernel for RsWidth {
+impl SedonaScalarKernel for RsSize {
fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
let matcher = ArgMatcher::new(
vec![ArgMatcher::is_raster()],
@@ -74,10 +109,16 @@ impl SedonaScalarKernel for RsWidth {
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);
- }
+ Some(raster) => match self.size_type {
+ SizeType::Width => {
+ let width = raster.metadata().width();
+ builder.append_value(width);
+ }
+ SizeType::Height => {
+ let height = raster.metadata().height();
+ builder.append_value(height);
+ }
+ },
}
Ok(())
})?;
@@ -91,6 +132,7 @@ mod tests {
use super::*;
use arrow_array::{Array, UInt64Array};
use datafusion_expr::ScalarUDF;
+ use rstest::rstest;
use sedona_schema::datatypes::RASTER;
use sedona_testing::rasters::generate_test_rasters;
@@ -99,15 +141,21 @@ mod tests {
let udf: ScalarUDF = rs_width_udf().into();
assert_eq!(udf.name(), "rs_width");
assert!(udf.documentation().is_some());
+
+ let udf: ScalarUDF = rs_height_udf().into();
+ assert_eq!(udf.name(), "rs_height");
+ assert!(udf.documentation().is_some());
}
- #[test]
- fn udf_invoke() {
+ #[rstest]
+ fn udf_invoke(#[values(SizeType::Width, SizeType::Height)] st: SizeType) {
+ let kernel = RsSize {
+ size_type: st.clone(),
+ };
// 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];
@@ -115,12 +163,19 @@ mod tests {
// Check the result
if let ColumnarValue::Array(result_array) = result {
- let width_array =
result_array.as_any().downcast_ref::<UInt64Array>().unwrap();
+ let size_array =
result_array.as_any().downcast_ref::<UInt64Array>().unwrap();
+
+ assert_eq!(size_array.len(), 3);
- 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
+ match st.clone() {
+ SizeType::Width => assert_eq!(size_array.value(0), 1), //
First raster width
+ SizeType::Height => assert_eq!(size_array.value(0), 2), //
First raster height
+ }
+ assert!(size_array.is_null(1)); // Second raster is null
+ match st.clone() {
+ SizeType::Width => assert_eq!(size_array.value(2), 3), //
Third raster width
+ SizeType::Height => assert_eq!(size_array.value(2), 4), //
Third raster height
+ }
} else {
panic!("Expected array result");
}
diff --git a/rust/sedona/Cargo.toml b/rust/sedona/Cargo.toml
index c7cb95a2..66f4324e 100644
--- a/rust/sedona/Cargo.toml
+++ b/rust/sedona/Cargo.toml
@@ -70,6 +70,7 @@ sedona-geometry = { path = "../sedona-geometry" }
sedona-geoparquet = { path = "../sedona-geoparquet" }
sedona-geos = { path = "../../c/sedona-geos", optional = true }
sedona-proj = { path = "../../c/sedona-proj", default-features = false }
+sedona-raster-functions = { path = "../../rust/sedona-raster-functions" }
sedona-schema = { path = "../sedona-schema" }
sedona-spatial-join = { path = "../sedona-spatial-join", optional = true }
sedona-s2geography = { path = "../../c/sedona-s2geography", optional = true }
diff --git a/rust/sedona/src/context.rs b/rust/sedona/src/context.rs
index bd503e5d..3903cfa9 100644
--- a/rust/sedona/src/context.rs
+++ b/rust/sedona/src/context.rs
@@ -153,6 +153,9 @@ impl SedonaContext {
// is called).
out.register_scalar_kernels(sedona_proj::register::scalar_kernels().into_iter())?;
+ // Always register raster functions
+
out.register_function_set(sedona_raster_functions::register::default_function_set());
+
Ok(out)
}
