You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@sedona.apache.org by ja...@apache.org on 2026/10/21 19:28:27 UTC
(sedona-db) branch main updated: feat(rust/sedona-raster-functions): add RS_Stack_Aggr (#1376)
This is an automated email from the ASF dual-hosted git repository.
james-willis 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 9879abca9 feat(rust/sedona-raster-functions): add RS_Stack_Aggr (#1376)
9879abca9 is described below
commit 9879abca99584477c2421a13caef24e426b9f84f
Author: James Willis <ja...@users.noreply.github.com>
AuthorDate: Wed Oct 21 12:28:21 2026 -0700
feat(rust/sedona-raster-functions): add RS_Stack_Aggr (#1376)
---
docs/reference/sql/rs_stack_aggr.qmd | 63 +++
integration/spark-parity/test_rs_stack_aggr.py | 141 ++++++
.../sedonadb/tests/functions/test_rs_stack_aggr.py | 142 ++++++
rust/sedona-raster-functions/src/lib.rs | 1 +
rust/sedona-raster-functions/src/register.rs | 2 +-
rust/sedona-raster-functions/src/rs_stack.rs | 49 +-
rust/sedona-raster-functions/src/rs_stack_aggr.rs | 519 +++++++++++++++++++++
rust/sedona-testing/src/testers.rs | 13 +-
8 files changed, 912 insertions(+), 18 deletions(-)
diff --git a/docs/reference/sql/rs_stack_aggr.qmd b/docs/reference/sql/rs_stack_aggr.qmd
new file mode 100644
index 000000000..aa63d9442
--- /dev/null
+++ b/docs/reference/sql/rs_stack_aggr.qmd
@@ -0,0 +1,63 @@
+---
+# 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_Stack_Aggr
+description: >
+ Aggregates a column of rasters into one raster holding all of their bands,
+ ordered by an index column.
+kernels:
+ - returns: raster
+ args:
+ - raster
+ - name: index
+ type: integer
+ description: >
+ The position of this row's raster in the result. Indexes must be
+ distinct and evenly spaced, such as 1, 2, 3 or 0, 10, 20.
+---
+
+::: callout-warning
+**Experimental.** This function is experimental; its behavior may change without notice.
+:::
+
+## Description
+
+An aggregate function: the aggregate form of [RS_Stack](rs_stack.qmd). Sedona
+Spark 1.9 calls it `RS_Union_Aggr`.
+
+The result holds every band of the raster with the lowest index, then every band
+of the next, and so on, under the georeference of the lowest-index raster. The
+index decides the order because an aggregate sees its rows in no particular
+order. Indexes must be distinct and evenly spaced: a gap usually means a raster
+is missing, which would silently shift every later band, so it is an error.
+
+As with RS_Stack, every raster must be on one grid (the same width, height,
+geotransform and CRS), each band keeps its own pixel type, nodata value and
+name, and a band's spatial dimensions take the lowest-index raster's names.
+Rows with a NULL raster or a NULL index are skipped; a group with no other rows
+gives NULL.
+
+No pixel is read: bands are carried over as they are, so stacking rasters that
+are still out-of-database does not load them.
+
+## Examples
+
+```sql
+SELECT RS_NumBands(RS_Stack_Aggr(rast, i))
+FROM (VALUES (RS_Example(), 1), (RS_Example(), 2)) AS t(rast, i);
+```
diff --git a/integration/spark-parity/test_rs_stack_aggr.py b/integration/spark-parity/test_rs_stack_aggr.py
new file mode 100644
index 000000000..0f6792e3a
--- /dev/null
+++ b/integration/spark-parity/test_rs_stack_aggr.py
@@ -0,0 +1,141 @@
+# 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.
+"""SedonaDB vs Sedona Spark parity for RS_Stack_Aggr.
+
+The Sedona Spark 1.9.1 release the suite pins calls this aggregate
+RS_Union_Aggr, and the rename to RS_Stack_Aggr (apache/sedona#3427) is not
+released yet. Every case that needs a result from Sedona Spark is an xfail until
+the pin moves to a release that has RS_Stack_Aggr; drop those markers then. The
+refusal cases carry no marker: Sedona Spark already refuses the name it does not
+know, and it keeps refusing those inputs once it has the aggregate.
+
+Each raster is registered as its own view and the rows are gathered with
+UNION ALL, each with its index, listed out of index order so the index and not
+the row order decides the band order. Each case compares the whole output
+raster, decoded on both engines, and anchors it to the inputs' bands stacked in
+index order under the lowest-index raster's grid. As in test_rs_stack.py, the
+fixtures carry no nodata value, and Sedona Spark ignoring every raster's
+georeference but the first, where SedonaDB requires one grid, is pinned down as
+a test of what each engine does, so it fails loudly if either engine changes.
+"""
+
+import numpy as np
+import pytest
+
+from sedonadb.raster_testing import DecodedRaster, assert_decoded_equal, write_geotiff
+from sedonadb.testing import SedonaDB, compare
+from sedonadb.testing_spark import SedonaSpark
+
+no_rs_stack_aggr = pytest.mark.xfail(
+ reason="Sedona Spark 1.9.1 has no RS_Stack_Aggr; it calls the aggregate "
+ "RS_Union_Aggr until the rename in apache/sedona#3427 is released"
+)
+
+BBOX = (100, 482, 114, 500)
+
+
+def _views(tmp_path, rasters):
+ """Register each `(name, DecodedRaster)` as a view on both engines."""
+ sedona, spark = SedonaDB(), SedonaSpark()
+ for name, raster in rasters:
+ path = tmp_path / f"{name}.tif"
+ write_geotiff(
+ path, raster.pixels, gdal_transform=raster.gdal_transform, nodata=None
+ )
+ for eng in (sedona, spark):
+ eng.create_raster_view(name, path)
+ return sedona, spark
+
+
+def _raster(dtype="uint8", bands=1, plant=0, bbox=BBOX, width=7, height=6):
+ """A random raster, made distinct from its siblings by a planted value."""
+ raster = DecodedRaster.random(
+ dtype,
+ bands=bands,
+ width=width,
+ height=height,
+ bbox=bbox,
+ plants={(2, 3): plant},
+ )
+ raster.nodata = [None] * bands
+ return raster
+
+
+def _stacked(*rasters):
+ return DecodedRaster(
+ np.concatenate([r.pixels for r in rasters]),
+ gdal_transform=rasters[0].gdal_transform,
+ nodata=[None] * sum(len(r.pixels) for r in rasters),
+ )
+
+
+def _sql(rows, select="RS_Stack_Aggr(rast, i)"):
+ """`select` over `(view, index)` rows gathered with UNION ALL"""
+ union = " UNION ALL ".join(f"SELECT rast, {i} AS i FROM {view}" for view, i in rows)
+ return f"SELECT {select} FROM ({union}) AS t"
+
+
+@no_rs_stack_aggr
+def test_rs_stack_aggr(tmp_path):
+ a, b, c = _raster(plant=1), _raster(bands=2, plant=2), _raster(plant=3)
+ sedona, spark = _views(tmp_path, [("sa_a", a), ("sa_b", b), ("sa_c", c)])
+ sql = _sql([("sa_c", 3), ("sa_a", 1), ("sa_b", 2)])
+ compare(sql, sedona, spark, expected=_stacked(a, b, c))
+
+
+@no_rs_stack_aggr
+def test_rs_stack_aggr_index_step(tmp_path):
+ """Evenly spaced indexes need not step by one."""
+ a, b = _raster(plant=1), _raster(plant=2)
+ sedona, spark = _views(tmp_path, [("sa_a", a), ("sa_b", b)])
+ sql = _sql([("sa_b", 20), ("sa_a", 10)])
+ compare(sql, sedona, spark, expected=_stacked(a, b))
+
+
+@pytest.mark.parametrize("indexes", [(1, 1), (1, 2, 4)], ids=["repeated", "uneven"])
+def test_rs_stack_aggr_bad_indexes(indexes, tmp_path):
+ """Both engines refuse a repeated index and unevenly spaced indexes. Error
+ types differ, so parity here is parity on refusal."""
+ rasters = [(f"sa_{n}", _raster(plant=n)) for n in range(len(indexes))]
+ sedona, spark = _views(tmp_path, rasters)
+ sql = _sql([(view, i) for (view, _), i in zip(rasters, indexes)])
+ for eng in (sedona, spark):
+ with pytest.raises(Exception):
+ eng.decode_raster_result(sql)
+
+
+@no_rs_stack_aggr
+def test_rs_stack_aggr_another_grid(tmp_path):
+ """The second raster has the same shape but is georeferenced elsewhere.
+ The engines diverge: SedonaDB rejects it, Sedona Spark keeps the
+ lowest-index raster's grid and ignores the others' georeference."""
+ a, b = _raster(plant=1), _raster(plant=2, bbox=(0, 0, 7, 6))
+ sedona, spark = _views(tmp_path, [("sa_a", a), ("sa_b", b)])
+ sql = _sql([("sa_a", 1), ("sa_b", 2)])
+ with pytest.raises(Exception):
+ sedona.decode_raster_result(sql)
+ assert_decoded_equal(spark.decode_raster_result(sql), _stacked(a, b), context=sql)
+
+
+def test_rs_stack_aggr_shape_mismatch(tmp_path):
+ """Both engines refuse rasters whose width or height differ."""
+ a, b = _raster(plant=1), _raster(plant=2, width=5, height=4, bbox=(0, 0, 5, 4))
+ sedona, spark = _views(tmp_path, [("sa_a", a), ("sa_b", b)])
+ sql = _sql([("sa_a", 1), ("sa_b", 2)])
+ for eng in (sedona, spark):
+ with pytest.raises(Exception):
+ eng.decode_raster_result(sql)
diff --git a/python/sedonadb/tests/functions/test_rs_stack_aggr.py b/python/sedonadb/tests/functions/test_rs_stack_aggr.py
new file mode 100644
index 000000000..7084ddd16
--- /dev/null
+++ b/python/sedonadb/tests/functions/test_rs_stack_aggr.py
@@ -0,0 +1,142 @@
+# 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_Stack_Aggr against a numpy reference.
+
+The inputs are GeoTIFFs read through RS_FromPath, which stay out-of-database:
+RS_Stack_Aggr carries their bands over by reference, so each test loads the
+result with RS_EnsureLoaded before decoding it and comparing with the input
+pixels, stacked in index order.
+"""
+
+import numpy as np
+import pyarrow as pa
+import pytest
+import sedonadb
+
+from sedonadb.raster_testing import (
+ DecodedRaster,
+ assert_decoded_equal,
+ decode_raster,
+ random_raster_data,
+ write_geotiff,
+)
+
+pytest.importorskip("rasterio")
+
+BBOX = (100.0, 482.0, 114.0, 500.0)
+
+
+def _tiff(tmp_path, name, *, bbox=BBOX, crs=None, seed=42):
+ data = random_raster_data("uint8", bands=1, height=6, width=7, seed=seed)
+ path = tmp_path / f"{name}.tif"
+ write_geotiff(path, data, bbox=bbox, crs=crs)
+ return path, data
+
+
+def _stack_aggr(con, rows, group_by=False):
+ """RS_Stack_Aggr over (group, path, index) rows; returns {group: raster}"""
+ groups, paths, indexes = zip(*rows) if rows else ([], [], [])
+ con.create_data_frame(
+ pa.table(
+ {
+ "g": pa.array(groups, pa.string()),
+ "p": pa.array([p and str(p) for p in paths], pa.string()),
+ "i": pa.array(indexes, pa.int64()),
+ }
+ )
+ ).to_view("stack_aggr_src", overwrite=True)
+ stack = "RS_EnsureLoaded(RS_Stack_Aggr(RS_FromPath(p), i))"
+ if group_by:
+ sql = f"SELECT g, {stack} AS r FROM stack_aggr_src GROUP BY g ORDER BY g"
+ else:
+ sql = f"SELECT 'all' AS g, {stack} AS r FROM stack_aggr_src"
+ table = con.sql(sql).to_arrow_table()
+ return dict(zip(table.column("g").to_pylist(), table.column("r")))
+
+
+def test_rs_stack_aggr_stacks_each_group_in_index_order(tmp_path):
+ """Rows arrive out of index order. Four partitions split the aggregate into
+ partial and final stages, so each group's state is serialized and merged
+ across a repartition before the rasters are stacked."""
+ con = sedonadb.connect()
+ con.sql("SET datafusion.execution.target_partitions = 4").execute()
+ r1, d1 = _tiff(tmp_path, "r1", seed=1)
+ r2, d2 = _tiff(tmp_path, "r2", seed=2)
+ r3, d3 = _tiff(tmp_path, "r3", seed=3)
+ got = _stack_aggr(
+ con,
+ [
+ ("a", r3, 3),
+ ("a", r1, 1),
+ ("b", r2, 20),
+ ("a", r2, 2),
+ ("b", r1, 10),
+ ],
+ group_by=True,
+ )
+ assert_decoded_equal(
+ decode_raster(got["a"]),
+ DecodedRaster(np.concatenate([d1, d2, d3]), bbox=BBOX, nodata=[None] * 3),
+ )
+ assert_decoded_equal(
+ decode_raster(got["b"]),
+ DecodedRaster(np.concatenate([d1, d2]), bbox=BBOX, nodata=[None] * 2),
+ )
+
+
+def test_rs_stack_aggr_skips_null_rows(con, tmp_path):
+ """A row with a NULL raster or a NULL index is left out, and a group with
+ no other rows gives NULL."""
+ r1, d1 = _tiff(tmp_path, "r1", seed=1)
+ r2, _ = _tiff(tmp_path, "r2", seed=2)
+ got = _stack_aggr(
+ con,
+ [
+ ("a", r1, 1),
+ ("a", None, 2),
+ ("a", r2, None),
+ ("b", r2, None),
+ ],
+ group_by=True,
+ )
+ assert_decoded_equal(
+ decode_raster(got["a"]), DecodedRaster(d1, bbox=BBOX, nodata=[None])
+ )
+ assert not got["b"].is_valid
+ assert not _stack_aggr(con, [])["all"].is_valid
+
+
+def test_rs_stack_aggr_rejects_bad_indexes(con, tmp_path):
+ r1, _ = _tiff(tmp_path, "r1", seed=1)
+ r2, _ = _tiff(tmp_path, "r2", seed=2)
+ r3, _ = _tiff(tmp_path, "r3", seed=3)
+ with pytest.raises(Exception, match="index 1 is given to more than one raster"):
+ _stack_aggr(con, [("a", r1, 1), ("a", r2, 1)])
+ with pytest.raises(Exception, match="indexes must be evenly spaced"):
+ _stack_aggr(con, [("a", r1, 1), ("a", r2, 2), ("a", r3, 4)])
+
+
+def test_rs_stack_aggr_rejects_another_grid(con, tmp_path):
+ """A raster of the same shape but georeferenced elsewhere, or in another
+ CRS, is not on the lowest-index raster's grid."""
+ r1, _ = _tiff(tmp_path, "r1", seed=1)
+ moved, _ = _tiff(tmp_path, "moved", bbox=(0.0, 0.0, 7.0, 6.0), seed=2)
+ with pytest.raises(Exception, match="the raster with index 2 has geotransform"):
+ _stack_aggr(con, [("a", r1, 1), ("a", moved, 2)])
+ mercator, _ = _tiff(tmp_path, "mercator", crs="EPSG:3857", seed=3)
+ with pytest.raises(Exception, match="the raster with index 2 has a different CRS"):
+ _stack_aggr(con, [("a", r1, 1), ("a", mercator, 2)])
diff --git a/rust/sedona-raster-functions/src/lib.rs b/rust/sedona-raster-functions/src/lib.rs
index 49371ade1..254dc4938 100644
--- a/rust/sedona-raster-functions/src/lib.rs
+++ b/rust/sedona-raster-functions/src/lib.rs
@@ -49,6 +49,7 @@ pub mod rs_slice;
pub mod rs_spatial_predicates;
pub mod rs_srid;
pub mod rs_stack;
+pub mod rs_stack_aggr;
pub mod rs_summarystats;
pub mod rs_value;
pub mod rs_values;
diff --git a/rust/sedona-raster-functions/src/register.rs b/rust/sedona-raster-functions/src/register.rs
index c4d950de7..fd4b379b1 100644
--- a/rust/sedona-raster-functions/src/register.rs
+++ b/rust/sedona-raster-functions/src/register.rs
@@ -95,7 +95,7 @@ pub fn default_function_set() -> FunctionSet {
crate::rs_worldcoordinate::rs_rastertoworldcoordy_udf,
);
- register_aggregate_udfs!(function_set,);
+ register_aggregate_udfs!(function_set, crate::rs_stack_aggr::rs_stack_aggr_udf,);
function_set
}
diff --git a/rust/sedona-raster-functions/src/rs_stack.rs b/rust/sedona-raster-functions/src/rs_stack.rs
index cc266738a..f1ba64240 100644
--- a/rust/sedona-raster-functions/src/rs_stack.rs
+++ b/rust/sedona-raster-functions/src/rs_stack.rs
@@ -102,7 +102,10 @@ impl SedonaScalarKernel for RsStack {
.map(|r| r.get(i))
.collect::<std::result::Result<Vec<_>, _>>()?;
let row: Vec<&dyn RasterRef> = row.iter().map(|r| r as &dyn RasterRef).collect();
- stack(&mut builder, &row)?;
+ stack("RS_Stack", &mut builder, &row, &|k| match k {
+ 0 => "the first raster".to_string(),
+ k => format!("raster {}", k + 1),
+ })?;
}
RasterExecutor::finish_over(args, Arc::new(builder.finish()?))
@@ -110,11 +113,17 @@ impl SedonaScalarKernel for RsStack {
}
/// Append one raster holding every band of `rasters`, in order, under the first
-/// raster's header.
-fn stack(builder: &mut RasterBuilder, rasters: &[&dyn RasterRef]) -> Result<()> {
+/// raster's header. `func` names the calling function in errors, and
+/// `label(k)` names the `k`th raster (e.g. "raster 2").
+pub(crate) fn stack(
+ func: &str,
+ builder: &mut RasterBuilder,
+ rasters: &[&dyn RasterRef],
+ label: &dyn Fn(usize) -> String,
+) -> Result<()> {
let first = rasters[0];
for (k, raster) in rasters.iter().enumerate().skip(1) {
- check_same_grid(first, *raster, k + 1)?;
+ check_same_grid(func, first, *raster, &label(0), &label(k))?;
}
let (x_dim, y_dim) = (first.x_dim(), first.y_dim());
@@ -134,10 +143,11 @@ fn stack(builder: &mut RasterBuilder, rasters: &[&dyn RasterRef]) -> Result<()>
d if d == from_x => Ok(x_dim),
d if d == from_y => Ok(y_dim),
d if d == x_dim || d == y_dim => exec_err!(
- "RS_Stack: band {} of raster {} has a non-spatial dimension named \
- '{d}', the name of one of the first raster's spatial dimensions",
+ "{func}: band {} of {} has a non-spatial dimension named '{d}', the \
+ name of one of {}'s spatial dimensions",
band_idx + 1,
- k + 1
+ label(k),
+ label(0)
),
d => Ok(d),
})
@@ -157,21 +167,28 @@ fn stack(builder: &mut RasterBuilder, rasters: &[&dyn RasterRef]) -> Result<()>
Ok(())
}
-/// Error unless `raster` (the `k`th argument) is on the first raster's grid:
-/// the same width, height, geotransform and CRS.
-fn check_same_grid(first: &dyn RasterRef, raster: &dyn RasterRef, k: usize) -> Result<()> {
+/// Error unless `raster` is on the `first` raster's grid: the same width,
+/// height, geotransform and CRS. `first_label` and `label` name the two in
+/// errors.
+fn check_same_grid(
+ func: &str,
+ first: &dyn RasterRef,
+ raster: &dyn RasterRef,
+ first_label: &str,
+ label: &str,
+) -> Result<()> {
let (width, height) = (first.width()?, first.height()?);
let (w, h) = (raster.width()?, raster.height()?);
if (w, h) != (width, height) {
return exec_err!(
- "RS_Stack: raster {k} is {w} x {h}, but the first raster is {width} x {height}; \
- every raster must be on the same grid"
+ "{func}: {label} is {w} x {h}, but {first_label} is {width} x {height}; every \
+ raster must be on the same grid"
);
}
if raster.transform() != first.transform() {
return exec_err!(
- "RS_Stack: raster {k} has geotransform {:?}, but the first raster has {:?}; \
- every raster must be on the same grid",
+ "{func}: {label} has geotransform {:?}, but {first_label} has {:?}; every raster \
+ must be on the same grid",
raster.transform(),
first.transform()
);
@@ -186,8 +203,8 @@ fn check_same_grid(first: &dyn RasterRef, raster: &dyn RasterRef, k: usize) -> R
};
if !same_crs {
return exec_err!(
- "RS_Stack: raster {k} has a different CRS than the first raster; every raster must \
- be on the same grid"
+ "{func}: {label} has a different CRS than {first_label}; every raster must be on \
+ the same grid"
);
}
Ok(())
diff --git a/rust/sedona-raster-functions/src/rs_stack_aggr.rs b/rust/sedona-raster-functions/src/rs_stack_aggr.rs
new file mode 100644
index 000000000..040cdcf82
--- /dev/null
+++ b/rust/sedona-raster-functions/src/rs_stack_aggr.rs
@@ -0,0 +1,519 @@
+// 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_Stack_Aggr` — the bands of a column of rasters, ordered by an index
+//! column, as one raster.
+//!
+//! ```text
+//! RS_Stack_Aggr(raster, index) -> Raster
+//! ```
+//!
+//! Sedona Spark 1.9 calls this aggregate `RS_Union_Aggr`. It is the aggregate
+//! form of [`RS_Stack`](crate::rs_stack): every band of the raster with the
+//! lowest index, then every band of the next, and so on, under the header of
+//! the lowest-index raster. The rasters must all be on one grid and their bands
+//! are carried over as they are, exactly as `RS_Stack` does.
+//!
+//! The index orders the rasters, since an aggregate sees its rows in no
+//! particular order. As in Sedona Spark, the indexes must be distinct and
+//! evenly spaced (1, 2, 3 or 0, 10, 20): a gap usually means a missing raster,
+//! which would silently shift every later band. Rows with a NULL raster or a
+//! NULL index are skipped, and a group with no other rows gives NULL.
+
+use std::sync::Arc;
+
+use arrow_array::{
+ Array, ArrayRef, BooleanArray, Int64Array, ListArray, StructArray, new_empty_array,
+};
+use arrow_buffer::OffsetBuffer;
+use arrow_schema::{DataType, Field, FieldRef};
+use datafusion_common::arrow::compute::{cast, concat, filter};
+use datafusion_common::cast::as_list_array;
+use datafusion_common::{Result, ScalarValue, exec_err, internal_datafusion_err};
+use datafusion_expr::{Accumulator, Volatility, utils::AggregateOrderSensitivity};
+use sedona_common::sedona_internal_err;
+use sedona_expr::aggregate_udf::{SedonaAccumulator, SedonaAggregateUDF};
+use sedona_raster::array::RasterStructArray;
+use sedona_raster::builder::RasterBuilder;
+use sedona_raster::traits::RasterRef;
+use sedona_schema::datatypes::SedonaType;
+use sedona_schema::matchers::ArgMatcher;
+
+use crate::rs_stack::stack;
+
+const FUNC: &str = "RS_Stack_Aggr";
+
+/// `RS_Stack_Aggr()` aggregate UDF — the bands of a column of rasters as one
+/// raster, ordered by an index column.
+pub fn rs_stack_aggr_udf() -> SedonaAggregateUDF {
+ SedonaAggregateUDF::new(
+ "rs_stack_aggr",
+ vec![Arc::new(RsStackAggr)],
+ Volatility::Immutable,
+ )
+ // The index, not the row order, decides the band order.
+ .with_order_sensitivity(AggregateOrderSensitivity::Insensitive)
+}
+
+#[derive(Debug)]
+struct RsStackAggr;
+
+impl SedonaAccumulator for RsStackAggr {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ ArgMatcher::new(
+ vec![ArgMatcher::is_raster(), ArgMatcher::is_integer()],
+ SedonaType::Raster,
+ )
+ .match_args(args)
+ }
+
+ fn accumulator(
+ &self,
+ _args: &[SedonaType],
+ _output_type: &SedonaType,
+ ) -> Result<Box<dyn Accumulator>> {
+ Ok(Box::new(StackAccumulator::default()))
+ }
+
+ fn state_fields(&self, _args: &[SedonaType]) -> Result<Vec<FieldRef>> {
+ Ok(vec![
+ Arc::new(Field::new("rasters", DataType::List(raster_item()?), true)),
+ Arc::new(Field::new("indexes", DataType::List(index_item()), true)),
+ ])
+ }
+}
+
+/// The element field of the state's raster list
+fn raster_item() -> Result<FieldRef> {
+ Ok(Arc::new(SedonaType::Raster.to_storage_field("item", true)?))
+}
+
+/// The element field of the state's index list
+fn index_item() -> FieldRef {
+ Arc::new(Field::new("item", DataType::Int64, true))
+}
+
+/// Collects the non-NULL (raster, index) rows of a group, as the arrays they
+/// arrived in, and stacks them by index on evaluation.
+#[derive(Debug, Default)]
+struct StackAccumulator {
+ rasters: Vec<ArrayRef>,
+ indexes: Vec<ArrayRef>,
+}
+
+impl StackAccumulator {
+ fn push(&mut self, rasters: ArrayRef, indexes: ArrayRef) {
+ if !rasters.is_empty() {
+ self.rasters.push(rasters);
+ self.indexes.push(indexes);
+ }
+ }
+
+ /// Every collected raster as one array, with its indexes
+ fn collected(&self) -> Result<(ArrayRef, ArrayRef)> {
+ let rasters = if self.rasters.is_empty() {
+ new_empty_array(SedonaType::Raster.storage_type())
+ } else {
+ let parts: Vec<&dyn Array> = self.rasters.iter().map(|a| a.as_ref()).collect();
+ concat(&parts)?
+ };
+ let indexes = if self.indexes.is_empty() {
+ new_empty_array(&DataType::Int64)
+ } else {
+ let parts: Vec<&dyn Array> = self.indexes.iter().map(|a| a.as_ref()).collect();
+ concat(&parts)?
+ };
+ Ok((rasters, indexes))
+ }
+}
+
+impl Accumulator for StackAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
+ let [rasters, indexes] = values else {
+ return sedona_internal_err!("{FUNC} expects two arguments, got {}", values.len());
+ };
+ let indexes = cast(indexes, &DataType::Int64)?;
+ // Skip a row whose raster or index is NULL, as a two-argument
+ // aggregate such as CORR skips a row with either argument NULL.
+ let keep = BooleanArray::from_iter(
+ (0..rasters.len()).map(|i| Some(rasters.is_valid(i) && indexes.is_valid(i))),
+ );
+ self.push(filter(rasters, &keep)?, filter(&indexes, &keep)?);
+ Ok(())
+ }
+
+ fn evaluate(&mut self) -> Result<ScalarValue> {
+ let (rasters, indexes) = self.collected()?;
+ if rasters.is_empty() {
+ return ScalarValue::try_from(SedonaType::Raster.storage_type());
+ }
+ let indexes = indexes
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .ok_or_else(|| internal_datafusion_err!("Expected Int64Array for {FUNC} indexes"))?;
+ let rasters = rasters
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .ok_or_else(|| internal_datafusion_err!("Expected StructArray for raster"))?;
+ let rasters = RasterStructArray::try_new(rasters)?;
+
+ let mut order: Vec<usize> = (0..indexes.len()).collect();
+ order.sort_by_key(|&i| indexes.value(i));
+ let sorted: Vec<i64> = order.iter().map(|&i| indexes.value(i)).collect();
+ check_indexes(&sorted)?;
+
+ let refs = order
+ .iter()
+ .map(|&i| rasters.get(i))
+ .collect::<std::result::Result<Vec<_>, _>>()?;
+ let refs: Vec<&dyn RasterRef> = refs.iter().map(|r| r as &dyn RasterRef).collect();
+ let mut builder = RasterBuilder::new(1);
+ stack(FUNC, &mut builder, &refs, &|k| {
+ format!("the raster with index {}", sorted[k])
+ })?;
+ ScalarValue::try_from_array(&builder.finish()?, 0)
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self)
+ + self
+ .rasters
+ .iter()
+ .chain(&self.indexes)
+ .map(|a| a.get_array_memory_size())
+ .sum::<usize>()
+ }
+
+ fn state(&mut self) -> Result<Vec<ScalarValue>> {
+ let (rasters, indexes) = self.collected()?;
+ let one_list = |item: FieldRef, values: ArrayRef| {
+ let offsets = OffsetBuffer::from_lengths([values.len()]);
+ ScalarValue::List(Arc::new(ListArray::new(item, offsets, values, None)))
+ };
+ Ok(vec![
+ one_list(raster_item()?, rasters),
+ one_list(index_item(), indexes),
+ ])
+ }
+
+ fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
+ let [rasters, indexes] = states else {
+ return sedona_internal_err!("{FUNC} expects two state fields, got {}", states.len());
+ };
+ let rasters = as_list_array(rasters)?;
+ let indexes = as_list_array(indexes)?;
+ for i in 0..rasters.len() {
+ if rasters.is_valid(i) {
+ self.push(rasters.value(i), indexes.value(i));
+ }
+ }
+ Ok(())
+ }
+}
+
+/// Error unless the sorted `indexes` are distinct and evenly spaced
+fn check_indexes(indexes: &[i64]) -> Result<()> {
+ let steps: Vec<(i64, i64)> = indexes.windows(2).map(|w| (w[0], w[1])).collect();
+ if let Some((index, _)) = steps.iter().find(|(a, b)| a == b) {
+ return exec_err!(
+ "{FUNC}: index {index} is given to more than one raster; each raster needs its own \
+ index"
+ );
+ }
+ let step = |(a, b): (i64, i64)| i128::from(b) - i128::from(a);
+ if let Some(&first) = steps.first()
+ && let Some(&(a, b)) = steps.iter().find(|&&pair| step(pair) != step(first))
+ {
+ return exec_err!(
+ "{FUNC}: indexes must be evenly spaced, but {a} to {b} is a step of {} where {} to {} \
+ is a step of {}; is a raster missing?",
+ step((a, b)),
+ first.0,
+ first.1,
+ step(first)
+ );
+ }
+ Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use datafusion_expr::AggregateUDF;
+ use sedona_schema::datatypes::RASTER;
+ use sedona_testing::raster_spec::{RasterSpec, assert_raster_scalar_equals, raster_array};
+ use sedona_testing::testers::AggregateUdfTester;
+
+ #[test]
+ fn udf_metadata() {
+ let udf: AggregateUDF = rs_stack_aggr_udf().into();
+ assert_eq!(udf.name(), "rs_stack_aggr");
+ let tester = AggregateUdfTester::new(udf, vec![RASTER, SedonaType::Arrow(DataType::Int64)]);
+ assert_eq!(tester.return_type().unwrap(), SedonaType::Raster);
+ }
+
+ #[test]
+ fn stacks_bands_in_index_order_across_batches() {
+ // Rows arrive out of order and split across batches (and so across
+ // partial states); the index alone decides the band order.
+ let tester = AggregateUdfTester::new(
+ rs_stack_aggr_udf().into(),
+ vec![RASTER, SedonaType::Arrow(DataType::Int64)],
+ );
+ let first = RasterSpec::d2(3, 2).band_values(&[0u8, 1, 2, 3, 4, 5]);
+ let second = RasterSpec::d2(3, 2).band_values(&[10u8, 11, 12, 13, 14, 15]);
+ let third = RasterSpec::d2(3, 2).band_values(&[20u8, 21, 22, 23, 24, 25]);
+ let result = tester
+ .aggregate_columns(&[
+ vec![
+ Arc::new(raster_array(vec![Some(third), Some(first)])) as ArrayRef,
+ Arc::new(Int64Array::from(vec![3, 1])) as ArrayRef,
+ ],
+ vec![
+ Arc::new(raster_array(vec![Some(second)])) as ArrayRef,
+ Arc::new(Int64Array::from(vec![2])) as ArrayRef,
+ ],
+ ])
+ .unwrap();
+ let expected = RasterSpec::d2(3, 2)
+ .band_values(&[0u8, 1, 2, 3, 4, 5])
+ .band_values(&[10u8, 11, 12, 13, 14, 15])
+ .band_values(&[20u8, 21, 22, 23, 24, 25]);
+ assert_raster_scalar_equals(&result, &expected);
+ }
+
+ #[test]
+ fn keeps_every_band_of_each_raster() {
+ let tester = AggregateUdfTester::new(
+ rs_stack_aggr_udf().into(),
+ vec![RASTER, SedonaType::Arrow(DataType::Int64)],
+ );
+ let two_bands = RasterSpec::d2(3, 2)
+ .band_values(&[0u8, 1, 2, 3, 4, 5])
+ .band_values(&[6u8, 7, 8, 9, 10, 11]);
+ let one_band = RasterSpec::d2(3, 2).band_values(&[20u8, 21, 22, 23, 24, 25]);
+ let result = tester
+ .aggregate_columns(&[vec![
+ Arc::new(raster_array(vec![Some(one_band), Some(two_bands)])) as ArrayRef,
+ Arc::new(Int64Array::from(vec![1, 0])) as ArrayRef,
+ ]])
+ .unwrap();
+ let expected = RasterSpec::d2(3, 2)
+ .band_values(&[0u8, 1, 2, 3, 4, 5])
+ .band_values(&[6u8, 7, 8, 9, 10, 11])
+ .band_values(&[20u8, 21, 22, 23, 24, 25]);
+ assert_raster_scalar_equals(&result, &expected);
+ }
+
+ #[test]
+ fn evenly_spaced_indexes_may_step_by_more_than_one() {
+ let tester = AggregateUdfTester::new(
+ rs_stack_aggr_udf().into(),
+ vec![RASTER, SedonaType::Arrow(DataType::Int64)],
+ );
+ let rasters = vec![
+ Some(RasterSpec::d2(3, 2).band_values(&[10u8, 11, 12, 13, 14, 15])),
+ Some(RasterSpec::d2(3, 2).band_values(&[0u8, 1, 2, 3, 4, 5])),
+ Some(RasterSpec::d2(3, 2).band_values(&[20u8, 21, 22, 23, 24, 25])),
+ ];
+ let result = tester
+ .aggregate_columns(&[vec![
+ Arc::new(raster_array(rasters)) as ArrayRef,
+ Arc::new(Int64Array::from(vec![10, 0, 20])) as ArrayRef,
+ ]])
+ .unwrap();
+ let expected = RasterSpec::d2(3, 2)
+ .band_values(&[0u8, 1, 2, 3, 4, 5])
+ .band_values(&[10u8, 11, 12, 13, 14, 15])
+ .band_values(&[20u8, 21, 22, 23, 24, 25]);
+ assert_raster_scalar_equals(&result, &expected);
+ }
+
+ #[test]
+ fn a_null_raster_or_index_skips_the_row() {
+ let tester = AggregateUdfTester::new(
+ rs_stack_aggr_udf().into(),
+ vec![RASTER, SedonaType::Arrow(DataType::Int64)],
+ );
+ let rasters = vec![
+ Some(RasterSpec::d2(3, 2).band_values(&[0u8, 1, 2, 3, 4, 5])),
+ None,
+ Some(RasterSpec::d2(3, 2).band_values(&[10u8, 11, 12, 13, 14, 15])),
+ Some(RasterSpec::d2(3, 2).band_values(&[20u8, 21, 22, 23, 24, 25])),
+ ];
+ let result = tester
+ .aggregate_columns(&[vec![
+ Arc::new(raster_array(rasters)) as ArrayRef,
+ Arc::new(Int64Array::from(vec![Some(1), Some(2), None, Some(2)])) as ArrayRef,
+ ]])
+ .unwrap();
+ let expected = RasterSpec::d2(3, 2)
+ .band_values(&[0u8, 1, 2, 3, 4, 5])
+ .band_values(&[20u8, 21, 22, 23, 24, 25]);
+ assert_raster_scalar_equals(&result, &expected);
+ }
+
+ #[test]
+ fn no_rows_gives_null() {
+ let tester = AggregateUdfTester::new(
+ rs_stack_aggr_udf().into(),
+ vec![RASTER, SedonaType::Arrow(DataType::Int64)],
+ );
+ let rasters = vec![
+ None,
+ Some(RasterSpec::d2(3, 2).band_values(&[0u8, 1, 2, 3, 4, 5])),
+ ];
+ let result = tester
+ .aggregate_columns(&[vec![
+ Arc::new(raster_array(rasters)) as ArrayRef,
+ Arc::new(Int64Array::from(vec![Some(1), None])) as ArrayRef,
+ ]])
+ .unwrap();
+ assert!(result.is_null());
+
+ let result = tester
+ .aggregate_columns(&[vec![
+ Arc::new(raster_array(vec![])) as ArrayRef,
+ Arc::new(Int64Array::from(Vec::<i64>::new())) as ArrayRef,
+ ]])
+ .unwrap();
+ assert!(result.is_null());
+ }
+
+ #[test]
+ fn a_repeated_index_errors() {
+ // The two rows meet only when their batches' states are merged.
+ let tester = AggregateUdfTester::new(
+ rs_stack_aggr_udf().into(),
+ vec![RASTER, SedonaType::Arrow(DataType::Int64)],
+ );
+ let err = tester
+ .aggregate_columns(&[
+ vec![
+ Arc::new(raster_array(vec![Some(
+ RasterSpec::d2(3, 2).band_values(&[0u8, 1, 2, 3, 4, 5]),
+ )])) as ArrayRef,
+ Arc::new(Int64Array::from(vec![1])) as ArrayRef,
+ ],
+ vec![
+ Arc::new(raster_array(vec![Some(
+ RasterSpec::d2(3, 2).band_values(&[10u8, 11, 12, 13, 14, 15]),
+ )])) as ArrayRef,
+ Arc::new(Int64Array::from(vec![1])) as ArrayRef,
+ ],
+ ])
+ .unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("index 1 is given to more than one raster"),
+ "{err}"
+ );
+ }
+
+ #[test]
+ fn unevenly_spaced_indexes_error() {
+ let tester = AggregateUdfTester::new(
+ rs_stack_aggr_udf().into(),
+ vec![RASTER, SedonaType::Arrow(DataType::Int64)],
+ );
+ let rasters = vec![
+ Some(RasterSpec::d2(3, 2).band_values(&[0u8, 1, 2, 3, 4, 5])),
+ Some(RasterSpec::d2(3, 2).band_values(&[10u8, 11, 12, 13, 14, 15])),
+ Some(RasterSpec::d2(3, 2).band_values(&[20u8, 21, 22, 23, 24, 25])),
+ ];
+ let err = tester
+ .aggregate_columns(&[vec![
+ Arc::new(raster_array(rasters)) as ArrayRef,
+ Arc::new(Int64Array::from(vec![1, 2, 4])) as ArrayRef,
+ ]])
+ .unwrap_err();
+ assert!(
+ err.to_string().contains(
+ "indexes must be evenly spaced, but 2 to 4 is a step of 2 where 1 to 2 is a \
+ step of 1; is a raster missing?"
+ ),
+ "{err}"
+ );
+ }
+
+ #[test]
+ fn rasters_off_the_lowest_index_grid_error() {
+ let tester = AggregateUdfTester::new(
+ rs_stack_aggr_udf().into(),
+ vec![RASTER, SedonaType::Arrow(DataType::Int64)],
+ );
+ let moved = vec![
+ Some(RasterSpec::d2(3, 2).band_values(&[0u8, 1, 2, 3, 4, 5])),
+ Some(
+ RasterSpec::d2(3, 2)
+ .transform([5.0, 1.0, 0.0, 0.0, 0.0, -1.0])
+ .band_values(&[10u8, 11, 12, 13, 14, 15]),
+ ),
+ ];
+ let err = tester
+ .aggregate_columns(&[vec![
+ Arc::new(raster_array(moved)) as ArrayRef,
+ Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef,
+ ]])
+ .unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("RS_Stack_Aggr: the raster with index 2 has geotransform"),
+ "{err}"
+ );
+ assert!(
+ err.to_string().contains("but the raster with index 1 has"),
+ "{err}"
+ );
+
+ let reshaped = vec![
+ Some(RasterSpec::d2(3, 2).band_values(&[0u8, 1, 2, 3, 4, 5])),
+ Some(RasterSpec::d2(2, 3).band_values(&[0u8; 6])),
+ ];
+ let err = tester
+ .aggregate_columns(&[vec![
+ Arc::new(raster_array(reshaped)) as ArrayRef,
+ Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef,
+ ]])
+ .unwrap_err();
+ assert!(
+ err.to_string()
+ .contains("the raster with index 2 is 2 x 3, but the raster with index 1 is 3 x 2"),
+ "{err}"
+ );
+ }
+
+ #[test]
+ fn indexes_far_apart_do_not_overflow() {
+ let tester = AggregateUdfTester::new(
+ rs_stack_aggr_udf().into(),
+ vec![RASTER, SedonaType::Arrow(DataType::Int64)],
+ );
+ let rasters = vec![
+ Some(RasterSpec::d2(3, 2).band_values(&[0u8, 1, 2, 3, 4, 5])),
+ Some(RasterSpec::d2(3, 2).band_values(&[10u8, 11, 12, 13, 14, 15])),
+ Some(RasterSpec::d2(3, 2).band_values(&[20u8, 21, 22, 23, 24, 25])),
+ ];
+ let err = tester
+ .aggregate_columns(&[vec![
+ Arc::new(raster_array(rasters)) as ArrayRef,
+ Arc::new(Int64Array::from(vec![i64::MIN, 0, i64::MAX])) as ArrayRef,
+ ]])
+ .unwrap_err();
+ assert!(err.to_string().contains("evenly spaced"), "{err}");
+ }
+}
diff --git a/rust/sedona-testing/src/testers.rs b/rust/sedona-testing/src/testers.rs
index f25b1764e..45b862ad6 100644
--- a/rust/sedona-testing/src/testers.rs
+++ b/rust/sedona-testing/src/testers.rs
@@ -110,12 +110,23 @@ impl AggregateUdfTester {
/// in batches of one. This has the effect of testing all the pieces of
/// an aggregator in a somewhat configurable/predictable way.
pub fn aggregate(&self, batches: &Vec<ArrayRef>) -> Result<ScalarValue> {
+ let batches = batches
+ .iter()
+ .map(|batch| vec![batch.clone()])
+ .collect::<Vec<_>>();
+ self.aggregate_columns(&batches)
+ }
+
+ /// Perform a simple aggregation of a multi-argument aggregate
+ ///
+ /// As [Self::aggregate], but each batch holds one array per argument.
+ pub fn aggregate_columns(&self, batches: &[Vec<ArrayRef>]) -> Result<ScalarValue> {
let state_schema = Arc::new(Schema::new(self.state_fields()?));
let mut state_accumulator = self.new_accumulator()?;
for batch in batches {
let mut batch_accumulator = self.new_accumulator()?;
- batch_accumulator.update_batch(std::slice::from_ref(batch))?;
+ batch_accumulator.update_batch(batch)?;
let state_batch_of_one = RecordBatch::try_new(
state_schema.clone(),
batch_accumulator