diff --git a/docs/reference/sql/rs_zonalstats.qmd b/docs/reference/sql/rs_zonalstats.qmd new file mode 100644 index 0000000000..91fa536442 --- /dev/null +++ b/docs/reference/sql/rs_zonalstats.qmd @@ -0,0 +1,136 @@ +--- +# 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_ZonalStats +description: > + Computes a single summary statistic of the raster pixels covered by a region of interest + geometry. +kernels: + - returns: double + args: + - raster + - {name: roi, type: geometry} + - {name: stat_type, type: string} + - returns: double + args: + - raster + - {name: roi, type: geometry} + - {name: band, type: integer} + - {name: stat_type, type: string} + - returns: double + args: + - raster + - {name: roi, type: geometry} + - {name: band, type: integer} + - {name: stat_type, type: string} + - {name: all_touched, type: boolean} + - returns: double + args: + - raster + - {name: roi, type: geometry} + - {name: band, type: integer} + - {name: stat_type, type: string} + - {name: all_touched, type: boolean} + - {name: exclude_no_data, type: boolean} + - returns: double + args: + - raster + - name: roi + type: geometry + description: > + Region-of-interest geometry. Reprojected into the + raster's CRS when both carry one; it is an error for exactly one side to + have a CRS. + - name: band + type: integer + description: > + 1-based band index to compute over. Required for a multiband raster; a + single-band raster may omit it via the band-less overload. + - name: stat_type + type: string + description: > + Statistic to return (case-insensitive): count, sum, mean, median, mode, + stddev, variance, min, or max. `avg`/`average` alias mean and `sd` + aliases stddev. + - name: all_touched + type: boolean + description: > + If true, include every pixel the roi touches; otherwise only pixels + whose center falls inside it. Defaults to false. + - name: exclude_no_data + type: boolean + description: > + If true (the default), skip pixels equal to the band's nodata value. + - name: lenient + type: boolean + description: > + If true (the default), return NULL when the roi does not intersect the + raster; if false, raise an error. +--- + +::: callout-warning +**Experimental.** This function is experimental; its behavior may change without notice. +::: + +## Description + +`RS_ZonalStats` returns one summary statistic of the pixels of a single band +that a region of interest (ROI) geometry covers. A pixel +is included when its center falls inside the roi (or, with `all_touched`, when +the roi touches it at all). By default the band's nodata pixels are excluded. + +The statistic is one of `count`, `sum`, `mean`, `median`, `mode`, `stddev`, +`variance`, `min`, or `max`. `count` is returned as a whole number; all other +statistics are floating point. Variance and standard deviation are the sample +(n-1) values, and `mode` breaks ties toward the larger value. + +When the ROI overlaps the raster but selects no pixel, `count` is 0 and every +other statistic is NULL. When the ROI does not intersect the raster at all, the +result is NULL under the default `lenient` behavior, or an error when +`lenient` is set to false. + +The `band`, `all_touched`, `exclude_no_data`, and `lenient` arguments are added +one at a time by the wider overloads; `all_touched` defaults to false, +`exclude_no_data` to true, and `lenient` to true. The band-less overload does +not default to band 1 on a multiband raster: naming the band is required there. This function operates on 2-D `(y, x)` bands; computing a +statistic per non-spatial plane of an N-D band is not supported. + +Use [`RS_ZonalStatsAll`](rs_zonalstatsall.qmd) to compute every statistic at +once. + +## Examples + +```sql +SELECT RS_ZonalStats( + RS_Example(), + ST_GeomFromText('POLYGON ((60 90, 160 90, 160 190, 60 190, 60 90))', 'OGC:CRS84'), + 1, + 'mean' +); +``` + +```sql +SELECT RS_ZonalStats( + RS_Example(), + ST_GeomFromText('POLYGON ((60 90, 160 90, 160 190, 60 190, 60 90))', 'OGC:CRS84'), + 1, + 'count', + false, + false +); +``` diff --git a/docs/reference/sql/rs_zonalstatsall.qmd b/docs/reference/sql/rs_zonalstatsall.qmd new file mode 100644 index 0000000000..2db54b6e9f --- /dev/null +++ b/docs/reference/sql/rs_zonalstatsall.qmd @@ -0,0 +1,116 @@ +--- +# 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_ZonalStatsAll +description: > + Computes every summary statistic of the raster pixels covered by a region of interest + geometry and returns them as a struct. +kernels: + - returns: struct + args: + - raster + - {name: roi, type: geometry} + - returns: struct + args: + - raster + - {name: roi, type: geometry} + - {name: band, type: integer} + - returns: struct + args: + - raster + - {name: roi, type: geometry} + - {name: band, type: integer} + - {name: all_touched, type: boolean} + - returns: struct + args: + - raster + - {name: roi, type: geometry} + - {name: band, type: integer} + - {name: all_touched, type: boolean} + - {name: exclude_no_data, type: boolean} + - returns: struct + args: + - raster + - name: roi + type: geometry + description: > + Region-of-interest geometry. Reprojected into the + raster's CRS when both carry one; it is an error for exactly one side to + have a CRS. + - name: band + type: integer + description: > + 1-based band index to compute over. Required for a multiband raster; a + single-band raster may omit it via the band-less overload. + - name: all_touched + type: boolean + description: > + If true, include every pixel the roi touches; otherwise only pixels + whose center falls inside it. Defaults to false. + - name: exclude_no_data + type: boolean + description: > + If true (the default), skip pixels equal to the band's nodata value. + - name: lenient + type: boolean + description: > + If true (the default), return NULL when the roi does not intersect the + raster; if false, raise an error. +--- + +::: callout-warning +**Experimental.** This function is experimental; its behavior may change without notice. +::: + +## Description + +`RS_ZonalStatsAll` returns every summary statistic of the pixels of a single +band that a roi geometry covers, as a struct with fields `count`, `sum`, +`mean`, `median`, `mode`, `stddev`, `variance`, `min`, and `max`. A pixel is +included when its +center falls inside the roi (or, with `all_touched`, when the roi touches it +at all), and the band's nodata pixels are excluded by default. + +`count` is a whole number (a 64-bit integer); every other field is floating +point. Variance and standard deviation are the sample (n-1) values, and `mode` +breaks ties toward the larger value. + +When the roi overlaps the raster but selects no pixel, `count` is 0 and every +other field is NULL. When the roi does not intersect the raster at all, the +whole struct is NULL under the default `lenient` behavior, or the call raises an +error when `lenient` is set to false. + +The overloads are the same ladder as `RS_ZonalStats` without `stat_type`. The +`band`, `all_touched`, `exclude_no_data`, and `lenient` arguments are added one +at a time by the wider overloads; `all_touched` defaults to false, +`exclude_no_data` to true, and `lenient` to true. The band-less overload does +not default to band 1 on a multiband raster: naming the band is required there. +This function operates on 2-D `(y, x)` bands; computing statistics per +non-spatial plane of an N-D band is not supported. + +Use [`RS_ZonalStats`](rs_zonalstats.qmd) to compute a single statistic. + +## Examples + +```sql +SELECT RS_ZonalStatsAll( + RS_Example(), + ST_GeomFromText('POLYGON ((60 90, 160 90, 160 190, 60 190, 60 90))', 'OGC:CRS84'), + 1 +); +``` diff --git a/python/sedonadb/tests/functions/test_rs_zonalstats.py b/python/sedonadb/tests/functions/test_rs_zonalstats.py new file mode 100644 index 0000000000..d2f965a847 --- /dev/null +++ b/python/sedonadb/tests/functions/test_rs_zonalstats.py @@ -0,0 +1,373 @@ +# 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_ZonalStats / RS_ZonalStatsAll cross-checked against a numpy reference. + +Both functions mirror Apache Sedona Spark's positional overloads, so the tests +call them positionally: `(raster, roi, stat_type)` / +`(raster, roi, band, stat_type[, all_touched[, exclude_no_data[, lenient]]])` +for RS_ZonalStats and the same ladder without `stat_type` for RS_ZonalStatsAll. + +The fixture raster is CRS-less (so nothing reprojects and pixel selection is +bit-comparable). The reference rasterizes the roi with `rasterio.features` +(the same GDAL rasterizer the kernel uses) and reduces the selected pixels with +numpy; exact-selection statistics (count, sum, min, max, median, mode) are +compared exactly and the float-accumulation ones (mean, variance, stddev) with +a tolerance. + +rasterio is required to write the fixture GeoTIFF, so the whole module skips +when it is unavailable rather than importing it at module scope. +""" + +import math + +import numpy as np +import pyarrow as pa +import pytest +import shapely +from shapely.geometry import box + +pytest.importorskip("rasterio") + +from sedonadb.raster_testing import ( + random_raster_data, + write_geotiff, +) + +# GDAL-order geotransform: origin (100, 500), 2-wide by 3-tall north-up pixels; +# a 6x7 raster then spans x in [100, 114], y in [482, 500]. +GDAL_TRANSFORM = (100.0, 2.0, 0.0, 500.0, 0.0, -3.0) +BANDS, HEIGHT, WIDTH = 1, 6, 7 +NODATA = -9999 + +# A rectangle well inside the raster that selects a block of pixels. +GEOM_RECT = ( + "POLYGON ((102.6 495.8, 109.3 495.8, 109.3 485.9, 102.6 485.9, 102.6 495.8))" +) +# Entirely outside the raster extent. +GEOM_DISJOINT = "POLYGON ((900 900, 910 900, 910 890, 900 890, 900 900))" +# Bounding box overlaps the raster, but the geometry itself is disjoint: the +# triangle sits in the far corner of its bounding box, clear of the raster. A +# bounding-box gate would burn no pixels and report count 0; a true-geometry gate +# (matching Sedona Spark's rsIntersects) treats it as a no-intersection case. +GEOM_DISJOINT_BBOX = "POLYGON ((124 490, 124 510, 108 510, 124 490))" +# A thin strip crossing the x = 104 pixel boundary but covering no pixel center +# (centers sit at odd x): selects nothing unless all_touched. +GEOM_SLIVER = "POLYGON ((103.6 499, 104.4 499, 104.4 483, 103.6 483, 103.6 499))" + +STATS = ["count", "sum", "mean", "median", "mode", "stddev", "variance", "min", "max"] +EXACT_STATS = {"count", "sum", "min", "max", "median", "mode"} + + +def fixture_raster(tmp_path): + """A single-band int32 raster with planted nodata and a repeated value. + + Returns `(path, band)` where `band` is the `(HEIGHT, WIDTH)` numpy array. + Two interior pixels hold the nodata value and three hold a repeated value + (66) so the mode is unambiguous and nodata exclusion is observable. + """ + data = random_raster_data( + "int32", + bands=BANDS, + height=HEIGHT, + width=WIDTH, + seed=7, + plants={(1, 1): NODATA, (2, 2): NODATA, (1, 2): 66, (2, 3): 66, (3, 1): 66}, + ) + path = tmp_path / "zonal.tif" + write_geotiff(path, data, gdal_transform=GDAL_TRANSFORM, nodata=NODATA) + return path, data[0] + + +def float_fixture_raster(tmp_path, *, planted_value, name): + """A single-band float64 raster with `planted_value` (NaN or inf) at an + interior pixel that GEOM_RECT selects, so the roi statistics must reckon with + it. The planted value is not the nodata sentinel, so nodata exclusion (on by + default) does not drop it. + + Returns `(path, band)` where `band` is the `(HEIGHT, WIDTH)` numpy array. + """ + data = random_raster_data( + "float64", bands=BANDS, height=HEIGHT, width=WIDTH, seed=7 + ) + # Pixel centre (105, 492.5) sits inside GEOM_RECT. + data[0][2, 2] = planted_value + path = tmp_path / name + write_geotiff(path, data, gdal_transform=GDAL_TRANSFORM, nodata=NODATA) + return path, data[0] + + +def numpy_reference(band, wkt, *, all_touched, exclude_no_data): + """Reference statistics over the pixels the roi selects, via rasterio+numpy. + + Returns a dict of every statistic, or the sentinel string ``"empty"`` when + the selection is empty (the caller maps that to count 0 / NULLs). + """ + import rasterio.features + from rasterio.transform import Affine + + geom = shapely.from_wkt(wkt) + mask = rasterio.features.rasterize( + [(geom, 1)], + out_shape=band.shape, + transform=Affine.from_gdal(*GDAL_TRANSFORM), + all_touched=all_touched, + fill=0, + dtype="uint8", + ) + sel = band[mask == 1].astype(np.float64) + if exclude_no_data: + sel = sel[sel != NODATA] + if sel.size == 0: + return "empty" + + values, counts = np.unique(sel, return_counts=True) + mode = float(values[counts == counts.max()].max()) # ties -> largest + n = sel.size + return { + "count": float(n), + "sum": float(sel.sum()), + "mean": float(sel.mean()), + "median": float(np.median(sel)), + "mode": mode, + "stddev": float(sel.std(ddof=1)) if n > 1 else 0.0, + "variance": float(sel.var(ddof=1)) if n > 1 else 0.0, + "min": float(sel.min()), + "max": float(sel.max()), + } + + +def _assert_stat_equal(stat, got, expected): + """Compare one statistic against the numpy reference with NaN/inf-aware + equality: NaN never equals itself, so match it explicitly; inf compares + exactly; otherwise use exact equality for the integer-exact statistics and a + tolerance for the float-accumulating ones.""" + if math.isnan(expected): + assert got is not None and math.isnan(got), f"{stat}: expected NaN, got {got!r}" + elif math.isinf(expected): + assert got == expected, f"{stat}: expected {expected}, got {got!r}" + elif stat in EXACT_STATS: + assert got == expected, f"{stat}: {got!r} != {expected!r}" + else: + assert got == pytest.approx(expected), f"{stat}: {got!r} !~ {expected!r}" + + +def _one_row(con, path, wkt): + """A one-row frame with the raster and roi as columns, so the kernel runs + its real per-row array path rather than constant-folding a scalar.""" + df = con.create_data_frame( + pa.table( + { + "path": pa.array([str(path)], pa.utf8()), + "wkt": pa.array([wkt], pa.utf8()), + } + ) + ) + return df, df.path.funcs.rs_frompath(), con.funcs.st_geomfromtext(df.wkt) + + +def zonal_stat(con, path, wkt, trailing): + """RS_ZonalStats over a one-row table. + + `trailing` is the positional argument list after `(raster, roi)` — e.g. + `["mean"]` (band-less overload) or `[1, "mean", all_touched, exclude_no_data, + lenient]`. Raster and roi travel as columns; the trailing scalars are + literals, matching how the SQL reads. + """ + df, raster, geom = _one_row(con, path, wkt) + table = df.select(r=raster.funcs.rs_zonalstats(geom, *trailing)).to_arrow_table() + return table["r"][0].as_py() + + +def zonal_stats_all(con, path, wkt, trailing): + """RS_ZonalStatsAll over a one-row table; returns the struct as a dict. + + `trailing` is the positional argument list after `(raster, roi)` — e.g. + `[]` (band-less overload) or `[1, all_touched, exclude_no_data, lenient]`. + """ + df, raster, geom = _one_row(con, path, wkt) + table = df.select(r=raster.funcs.rs_zonalstatsall(geom, *trailing)).to_arrow_table() + return table["r"][0].as_py() + + +@pytest.mark.parametrize("stat", STATS) +@pytest.mark.parametrize("all_touched", [False, True]) +def test_single_stat_matches_numpy(con, tmp_path, stat, all_touched): + path, band = fixture_raster(tmp_path) + expected = numpy_reference( + band, GEOM_RECT, all_touched=all_touched, exclude_no_data=True + ) + assert expected != "empty", "GEOM_RECT should select pixels" + + # (raster, roi, band, stat_type, all_touched) — the 5-arg overload. + got = zonal_stat(con, path, GEOM_RECT, [1, stat, all_touched]) + if stat in EXACT_STATS: + assert got == expected[stat] + else: + assert got == pytest.approx(expected[stat]) + + +def test_all_struct_matches_numpy(con, tmp_path): + path, band = fixture_raster(tmp_path) + expected = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_no_data=True) + # (raster, roi, band) — all_touched defaults to false. + got = zonal_stats_all(con, path, GEOM_RECT, [1]) + + # count is an integer field (Int64); every other field is floating point. + assert isinstance(got["count"], int) + assert isinstance(got["sum"], float) + assert got["count"] == expected["count"] + for stat in EXACT_STATS - {"count"}: + assert got[stat] == expected[stat] + for stat in ("mean", "variance", "stddev"): + assert got[stat] == pytest.approx(expected[stat]) + + +def test_nan_pixel_poisons_every_statistic_like_numpy(con, tmp_path): + """A NaN pixel that is not the nodata sentinel poisons every statistic (numpy + semantics): count stays a real tally, everything else is NaN. Pinned against + rasterio+numpy over the same masked selection, so the NaN handling is + validated against a trusted reference rather than asserted on faith.""" + path, band = float_fixture_raster( + tmp_path, planted_value=float("nan"), name="zonal_nan.tif" + ) + with np.errstate(all="ignore"): + expected = numpy_reference( + band, GEOM_RECT, all_touched=False, exclude_no_data=True + ) + assert expected != "empty", "GEOM_RECT should select the planted pixel" + got = zonal_stats_all(con, path, GEOM_RECT, [1]) + for stat in STATS: + _assert_stat_equal(stat, got[stat], expected[stat]) + + +def test_infinity_pixel_matches_numpy(con, tmp_path): + """A +inf pixel flows through (not the nodata sentinel, not NaN): sum, mean, + max and mode go to +inf, min and median stay finite, and variance/stddev + become NaN (inf - inf). Pinned against rasterio+numpy over the same + selection.""" + path, band = float_fixture_raster( + tmp_path, planted_value=float("inf"), name="zonal_inf.tif" + ) + with np.errstate(all="ignore"): + expected = numpy_reference( + band, GEOM_RECT, all_touched=False, exclude_no_data=True + ) + assert expected != "empty", "GEOM_RECT should select the planted pixel" + got = zonal_stats_all(con, path, GEOM_RECT, [1]) + for stat in STATS: + _assert_stat_equal(stat, got[stat], expected[stat]) + + +def test_exclude_no_data_default_and_disabled(con, tmp_path): + path, band = fixture_raster(tmp_path) + # Default excludes nodata; disabling it keeps those pixels, raising count. + excluded = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_no_data=True) + included = numpy_reference( + band, GEOM_RECT, all_touched=False, exclude_no_data=False + ) + assert included["count"] > excluded["count"] + + # Default (4-arg (raster, roi, band, stat_type)) excludes nodata. + assert zonal_stat(con, path, GEOM_RECT, [1, "count"]) == excluded["count"] + # exclude_no_data => false keeps it: the 6-arg overload trails (all_touched, + # exclude_no_data). + assert ( + zonal_stat(con, path, GEOM_RECT, [1, "count", False, False]) + == included["count"] + ) + + +def test_sliver_selects_nothing_unless_all_touched(con, tmp_path): + path, _ = fixture_raster(tmp_path) + # The roi overlaps the raster but covers no pixel center: count 0, rest NULL. + assert zonal_stat(con, path, GEOM_SLIVER, [1, "count"]) == 0.0 + assert zonal_stat(con, path, GEOM_SLIVER, [1, "sum"]) is None + # all_touched (5-arg overload) picks up the pixels it crosses. + touched = zonal_stat(con, path, GEOM_SLIVER, [1, "count", True]) + assert touched > 0.0 + + +def test_no_intersection_is_null_when_lenient_and_errors_when_strict(con, tmp_path): + path, _ = fixture_raster(tmp_path) + # Lenient (default): NULL, including count. + assert zonal_stat(con, path, GEOM_DISJOINT, [1, "count"]) is None + assert zonal_stats_all(con, path, GEOM_DISJOINT, [1]) is None + # Strict (lenient => false): the 7-arg overload trails (all_touched, + # exclude_no_data, lenient). + with pytest.raises(Exception, match="does not intersect"): + zonal_stat(con, path, GEOM_DISJOINT, [1, "count", False, True, False]) + + +def test_bbox_overlapping_but_geometry_disjoint_is_no_intersection(con, tmp_path): + path, _ = fixture_raster(tmp_path) + + # Premise: the roi's bounding box overlaps the raster extent, but the + # geometry is disjoint from it (unlike GEOM_DISJOINT, whose bbox misses too). + ox, px, _, oy, _, py = GDAL_TRANSFORM + raster_extent = box(ox, oy + py * HEIGHT, ox + px * WIDTH, oy) + geom = shapely.from_wkt(GEOM_DISJOINT_BBOX) + assert box(*geom.bounds).intersects(raster_extent), "bbox must overlap the raster" + assert geom.disjoint(raster_extent), "geometry must be disjoint from the raster" + + # Lenient (default): NULL, not count 0 — a true no-intersection case. + assert zonal_stat(con, path, GEOM_DISJOINT_BBOX, [1, "count"]) is None + assert zonal_stats_all(con, path, GEOM_DISJOINT_BBOX, [1]) is None + # Strict: errors, exactly like the fully-disjoint roi. + with pytest.raises(Exception, match="does not intersect"): + zonal_stat(con, path, GEOM_DISJOINT_BBOX, [1, "count", False, True, False]) + + +def test_unknown_statistic_errors(con, tmp_path): + path, _ = fixture_raster(tmp_path) + with pytest.raises(Exception, match="unknown statistic"): + zonal_stat(con, path, GEOM_RECT, [1, "nonsense"]) + + +def test_implicit_band_on_multiband_raster_errors(con, tmp_path): + # A 2-band raster: the band-less overloads must error rather than default to + # band 1 (this deliberately diverges from Sedona Spark). + data = random_raster_data("int32", bands=2, height=HEIGHT, width=WIDTH, seed=3) + path = tmp_path / "multiband.tif" + write_geotiff(path, data, gdal_transform=GDAL_TRANSFORM) + + # RS_ZonalStats 3-arg (raster, roi, stat_type): band implicit. + with pytest.raises(Exception, match="2 bands"): + zonal_stat(con, path, GEOM_RECT, ["count"]) + # RS_ZonalStatsAll 2-arg (raster, roi): band implicit. + with pytest.raises(Exception, match="2 bands"): + zonal_stats_all(con, path, GEOM_RECT, []) + # Naming the band resolves the ambiguity. + assert zonal_stat(con, path, GEOM_RECT, [1, "count"]) > 0.0 + + +def test_sql_text_smoke(con, tmp_path): + """One raw-SQL invocation per function keeps the parser path covered.""" + path, band = fixture_raster(tmp_path) + expected = numpy_reference(band, GEOM_RECT, all_touched=False, exclude_no_data=True) + + single = con.sql( + "SELECT RS_ZonalStats(RS_FromPath($1), ST_GeomFromText($2), 1, 'sum') AS r", + params=(str(path), GEOM_RECT), + ).to_arrow_table() + assert single["r"][0].as_py() == expected["sum"] + + everything = con.sql( + "SELECT RS_ZonalStatsAll(RS_FromPath($1), ST_GeomFromText($2), 1) AS r", + params=(str(path), GEOM_RECT), + ).to_arrow_table() + assert everything["r"][0].as_py()["count"] == expected["count"] diff --git a/rust/sedona-raster-functions/src/rs_spatial_predicates.rs b/rust/sedona-raster-functions/src/rs_spatial_predicates.rs index 5e8b0ab21d..e866beef25 100644 --- a/rust/sedona-raster-functions/src/rs_spatial_predicates.rs +++ b/rust/sedona-raster-functions/src/rs_spatial_predicates.rs @@ -394,6 +394,20 @@ fn evaluate_predicate(wkb_a: &[u8], wkb_b: &[u8]) -> Re /// = 9 + 4 + 80 = 93 const CONVEXHULL_WKB_SIZE: usize = 93; +/// Test whether a geometry intersects a raster's footprint (its convex-hull +/// polygon), using the same true geometry intersection as RS_Intersects rather +/// than a bounding-box overlap. This is the gate Sedona Spark's zonal statistics +/// use (`RasterPredicates.rsIntersects`): a zone whose bounding box overlaps the +/// raster but whose geometry is disjoint is treated as not intersecting. +/// +/// `geom_wkb` must already be in the raster's CRS; this performs no CRS +/// transformation (the raster footprint is built in the raster's own CRS). +pub fn raster_intersects_geom_wkb(raster: &dyn RasterRef, geom_wkb: &[u8]) -> Result { + let mut raster_wkb = Vec::with_capacity(CONVEXHULL_WKB_SIZE); + write_convexhull_wkb(raster, &mut raster_wkb)?; + evaluate_predicate::(&raster_wkb, geom_wkb) +} + #[cfg(test)] mod tests { use super::*; diff --git a/rust/sedona-raster-gdal/Cargo.toml b/rust/sedona-raster-gdal/Cargo.toml index 591eaef94f..b89c967792 100644 --- a/rust/sedona-raster-gdal/Cargo.toml +++ b/rust/sedona-raster-gdal/Cargo.toml @@ -104,3 +104,8 @@ path = "benches/rs_resample.rs" harness = false name = "rs_tile" path = "benches/rs_tile.rs" + +[[bench]] +harness = false +name = "rs_zonalstats" +path = "benches/rs_zonalstats.rs" diff --git a/rust/sedona-raster-gdal/benches/rs_zonalstats.rs b/rust/sedona-raster-gdal/benches/rs_zonalstats.rs new file mode 100644 index 0000000000..7f5122682a --- /dev/null +++ b/rust/sedona-raster-gdal/benches/rs_zonalstats.rs @@ -0,0 +1,184 @@ +// 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 the RS_ZonalStats / RS_ZonalStatsAll UDFs. +//! +//! Both functions rasterize the zone geometry into a mask, walk the masked +//! window collecting the selected pixel values, and reduce them to statistics. +//! +//! Each case builds a raster whose world extent is exactly the zone-polygon +//! generator's `[-10, 10]²` bounds at the requested resolution, so every +//! generated polygon lands on the raster and the full mask/collect/reduce path +//! runs. `all_touched = true` (the trailing boolean argument) guarantees a +//! polygon smaller than a cell still burns at least one pixel rather than +//! hitting the empty-zone early return. +//! +//! Axes: +//! - **Raster resolution** (`64²`, `256²`, `1024²`) with a small polygon: +//! rasterization + window scan dominate. +//! - **Zone polygon complexity** (vertex count) at a fixed resolution, driving +//! the GDAL rasterization cost. +//! - **Large zone**: a polygon covering most of the raster, so collecting the +//! masked values and the reduction (sort for median, frequency map for mode) +//! dominate — the `RS_ZonalStatsAll` case is the heaviest since it computes +//! every statistic. +//! +//! Numerical correctness against a reference (rasterio / numpy) is pinned by +//! the Python parity tests, not here; this bench only measures throughput. + +use std::sync::Arc; + +use arrow_array::{ArrayRef, BinaryArray, BooleanArray, Int64Array, StringArray}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion_expr::ScalarUDF; +use sedona_schema::datatypes::{SedonaType, RASTER, WKB_GEOMETRY}; +use sedona_testing::{ + benchmark_util::BenchmarkArgSpec, create::make_wkb, raster_spec::RasterSpec, + testers::ScalarUdfTester, +}; + +fn criterion_benchmark(c: &mut Criterion) { + let f = sedona_raster_gdal::register::default_function_set(); + let stats_udf: ScalarUDF = f + .scalar_udf("rs_zonalstats") + .expect("rs_zonalstats is registered") + .clone() + .into(); + let stats_all_udf: ScalarUDF = f + .scalar_udf("rs_zonalstatsall") + .expect("rs_zonalstatsall is registered") + .clone() + .into(); + + // RS_ZonalStats(raster, zone, band, stat, all_touched) and + // RS_ZonalStatsAll(raster, zone, band, all_touched). + let stats_tester = ScalarUdfTester::new( + stats_udf, + vec![ + RASTER, + WKB_GEOMETRY, + SedonaType::Arrow(arrow_schema::DataType::Int64), + SedonaType::Arrow(arrow_schema::DataType::Utf8), + SedonaType::Arrow(arrow_schema::DataType::Boolean), + ], + ); + let stats_all_tester = ScalarUdfTester::new( + stats_all_udf, + vec![ + RASTER, + WKB_GEOMETRY, + SedonaType::Arrow(arrow_schema::DataType::Int64), + SedonaType::Arrow(arrow_schema::DataType::Boolean), + ], + ); + + let band: ArrayRef = Arc::new(Int64Array::from(vec![1])); + let mean_stat: ArrayRef = Arc::new(StringArray::from(vec!["mean"])); + let all_touched: ArrayRef = Arc::new(BooleanArray::from(vec![true])); + + // A north-up raster covering exactly the polygon generator's [-10, 10]² + // bounds at the requested resolution, so every generated polygon overlaps. + let build_raster = |w: i64, h: i64| -> ArrayRef { + let transform = [-10.0, 20.0 / w as f64, 0.0, 10.0, 0.0, -20.0 / h as f64]; + let values: Vec = (0..(w * h)).map(|v| v as f64).collect(); + Arc::new( + RasterSpec::d2(w, h) + .band_values(&values) + .crs(None) + .transform(transform) + .build(), + ) + }; + + let gen_polygon = |vertices: usize| -> ArrayRef { + BenchmarkArgSpec::Polygon(vertices) + .build_arrays(0, 1, 1) + .expect("build zone polygon") + .remove(0) + }; + + let run_single = |c: &mut Criterion, label: &str, raster: ArrayRef, geom: ArrayRef| { + c.bench_function(label, |b| { + b.iter(|| { + stats_tester + .invoke_arrays(vec![ + raster.clone(), + geom.clone(), + band.clone(), + mean_stat.clone(), + all_touched.clone(), + ]) + .unwrap() + }) + }); + }; + + let run_all = |c: &mut Criterion, label: &str, raster: ArrayRef, geom: ArrayRef| { + c.bench_function(label, |b| { + b.iter(|| { + stats_all_tester + .invoke_arrays(vec![ + raster.clone(), + geom.clone(), + band.clone(), + all_touched.clone(), + ]) + .unwrap() + }) + }); + }; + + // Resolution sweep (simple 8-vertex polygon), single-stat mean. + for (w, h) in [(64i64, 64i64), (256, 256), (1024, 1024)] { + let label = + format!("raster-gdal rs_zonalstats ZonalStats(Raster({w}x{h}), Polygon(8), mean)"); + run_single(c, &label, build_raster(w, h), gen_polygon(8)); + } + + // Zone-complexity axis at a fixed 64×64 resolution. + run_single( + c, + "raster-gdal rs_zonalstats ZonalStats(Raster(64x64), Polygon(50), mean)", + build_raster(64, 64), + gen_polygon(50), + ); + + // Large zone: the polygon covers nearly the whole raster, so collecting the + // masked values and the reduction dominate. RS_ZonalStatsAll is the heaviest + // (median sort + mode frequency map over every selected pixel). + let big_geom = || -> ArrayRef { + Arc::new(BinaryArray::from_iter_values([make_wkb( + "POLYGON ((-9.5 -9.5, 9.5 -9.5, 9.5 9.5, -9.5 9.5, -9.5 -9.5))", + ) + .as_slice()])) + }; + run_single( + c, + "raster-gdal rs_zonalstats ZonalStats(Raster(1024x1024), Polygon(large), mean)", + build_raster(1024, 1024), + big_geom(), + ); + run_all( + c, + "raster-gdal rs_zonalstats ZonalStatsAll(Raster(1024x1024), Polygon(large))", + build_raster(1024, 1024), + big_geom(), + ); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/rust/sedona-raster-gdal/src/gdal_common.rs b/rust/sedona-raster-gdal/src/gdal_common.rs index ebb9d6277c..f082fa1f86 100644 --- a/rust/sedona-raster-gdal/src/gdal_common.rs +++ b/rust/sedona-raster-gdal/src/gdal_common.rs @@ -65,6 +65,16 @@ impl ToGdalGeoTransform for T { } } +/// A raster's stored six-coefficient GDAL geo-transform as a fixed array, +/// erroring when the transform is not exactly six elements. +/// +/// GDAL geo-transforms are +/// `[origin_x, pixel_width, rotation_x, origin_y, rotation_y, pixel_height]`. +pub fn raster_geo_transform(raster: &R) -> Result { + <[f64; 6]>::try_from(raster.transform()) + .map_err(|_| exec_datafusion_err!("expected a 6-element geotransform")) +} + /// Reconstruct raster metadata from a GDAL six-element geo-transform and raster dimensions. pub(crate) trait RasterMetadataFromGdalGeoTransform { fn to_raster_metadata(&self, width: usize, height: usize) -> RasterMetadata; diff --git a/rust/sedona-raster-gdal/src/lib.rs b/rust/sedona-raster-gdal/src/lib.rs index 28bb22a922..08673ef718 100644 --- a/rust/sedona-raster-gdal/src/lib.rs +++ b/rust/sedona-raster-gdal/src/lib.rs @@ -32,6 +32,7 @@ mod gdal_common; #[allow(dead_code)] mod gdal_dataset_provider; +mod mask; mod raster_loader; mod rs_as_geotiff; mod rs_as_raster; @@ -43,6 +44,7 @@ mod rs_polygonize; mod rs_reproject_match; mod rs_resample; mod rs_tile; +mod rs_zonal_stats; mod source_uri; mod utils; diff --git a/rust/sedona-raster-gdal/src/mask.rs b/rust/sedona-raster-gdal/src/mask.rs new file mode 100644 index 0000000000..ede4f04aa2 --- /dev/null +++ b/rust/sedona-raster-gdal/src/mask.rs @@ -0,0 +1,162 @@ +// 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. + +//! Geometry masking machinery shared by the raster functions that select the +//! pixels a geometry covers (RS_Clip, RS_ZonalStats). +//! +//! A mask is built in two steps: [`envelope_window`] clamps the geometry's +//! envelope to a rectangular pixel window on the raster grid, and +//! [`rasterize_geometry_mask`] burns the geometry into a window-sized 0/1 `u8` +//! mask. Callers then interpret the mask however they need — RS_Clip writes +//! nodata outside it, RS_ZonalStats reads the selected pixel values — so this +//! module owns only the window addressing and rasterization, not the +//! per-pixel consumption. + +use datafusion_common::{exec_datafusion_err, Result}; +use sedona_gdal::gdal::Gdal; +use sedona_gdal::geo_transform::{GeoTransform, GeoTransformEx}; +use sedona_gdal::mem::MemDatasetBuilder; +use sedona_gdal::raster::types::GdalDataType; +use sedona_gdal::vector::geometry::Geometry; + +/// A rectangular pixel window (offset + size) into a raster grid. +#[derive(Debug, Clone, Copy)] +pub struct PixelWindow { + pub col_off: usize, + pub row_off: usize, + pub width: usize, + pub height: usize, +} + +/// The geometry's envelope intersected with the raster extent, snapped outward +/// to whole pixels. `None` when the clamped window has no area (the envelope +/// falls entirely outside the raster, or only touches its boundary). +/// +/// This is the window PostGIS ST_Clip, `gdalwarp -crop_to_cutline`, and Sedona +/// Spark's raster functions use. All four envelope corners are mapped through +/// the inverse geotransform (so a skewed/rotated raster still gets a correct +/// superset window) and the resulting pixel-space bbox is floored/ceiled to +/// whole pixels. A degenerate envelope (point/line) landing exactly on a grid +/// line is widened to one pixel so the rasterizer — not the snapping — decides +/// whether it burns. +pub fn envelope_window( + geometry: &Geometry, + transform: &GeoTransform, + width: usize, + height: usize, +) -> Result> { + let env = geometry.envelope(); + let inverse = transform + .invert() + .map_err(|e| exec_datafusion_err!("raster mask: geotransform is not invertible: {e}"))?; + + let corners = [ + (env.MinX, env.MinY), + (env.MinX, env.MaxY), + (env.MaxX, env.MinY), + (env.MaxX, env.MaxY), + ]; + let mut min_col = f64::INFINITY; + let mut max_col = f64::NEG_INFINITY; + let mut min_row = f64::INFINITY; + let mut max_row = f64::NEG_INFINITY; + for (x, y) in corners { + let (col, row) = inverse.apply(x, y); + min_col = min_col.min(col); + max_col = max_col.max(col); + min_row = min_row.min(row); + max_row = max_row.max(row); + } + + let col0 = min_col.floor(); + let row0 = min_row.floor(); + let col1 = max_col.ceil().max(col0 + 1.0); + let row1 = max_row.ceil().max(row0 + 1.0); + + // Intersect with the raster extent. `>=` also rejects the NaN envelope of + // an empty geometry. + let col0 = col0.max(0.0); + let row0 = row0.max(0.0); + let col1 = col1.min(width as f64); + let row1 = row1.min(height as f64); + if !(col0 < col1 && row0 < row1) { + return Ok(None); + } + + Ok(Some(PixelWindow { + col_off: col0 as usize, + row_off: row0 as usize, + width: (col1 - col0) as usize, + height: (row1 - row0) as usize, + })) +} + +/// Rasterize `geometry` into a window-sized `u8` mask: 1 where the geometry +/// covers a pixel, 0 elsewhere. +/// +/// The mask is written into the caller-owned `out` buffer (cleared first), whose +/// allocation is reused across calls so a per-row rasterization does not allocate +/// a fresh `Vec` each time. It ends up `window.width * window.height` bytes long. +/// +/// The mask is a MEM UInt8 dataset covering only `window`, with the raster +/// geotransform shifted to the window's upper-left corner so pixel indices in +/// the mask line up with the same-offset pixels of the source raster. GDAL's +/// MEM driver zero-fills the band on creation, so only the burn (value 1, +/// inside the geometry) has to be written. `geometry` is consumed, since the +/// burn is its only remaining use. +pub fn rasterize_geometry_mask( + gdal: &Gdal, + geometry: Geometry, + transform: &GeoTransform, + window: &PixelWindow, + all_touched: bool, + out: &mut Vec, +) -> Result<()> { + let mask_dataset = + MemDatasetBuilder::create(gdal, window.width, window.height, 1, GdalDataType::UInt8) + .map_err(|e| exec_datafusion_err!("raster mask: failed to create mask dataset: {e}"))?; + let (window_ulx, window_uly) = transform.apply(window.col_off as f64, window.row_off as f64); + let mask_transform = [ + window_ulx, + transform[1], + transform[2], + window_uly, + transform[4], + transform[5], + ]; + mask_dataset + .set_geo_transform(&mask_transform) + .map_err(|e| exec_datafusion_err!("raster mask: failed to set mask geotransform: {e}"))?; + + gdal.rasterize_affine(&mask_dataset, &[1], &[geometry], &[1.0], all_touched) + .map_err(|e| exec_datafusion_err!("raster mask: failed to rasterize geometry: {e}"))?; + + let mask_band = mask_dataset + .rasterband(1) + .map_err(|e| exec_datafusion_err!("raster mask: failed to read mask band: {e}"))?; + let mask_buffer = mask_band + .read_as::( + (0, 0), + (window.width, window.height), + (window.width, window.height), + None, + ) + .map_err(|e| exec_datafusion_err!("raster mask: failed to read mask: {e}"))?; + out.clear(); + out.extend_from_slice(mask_buffer.data()); + Ok(()) +} diff --git a/rust/sedona-raster-gdal/src/register.rs b/rust/sedona-raster-gdal/src/register.rs index c2fb8d5ada..e06f459ff9 100644 --- a/rust/sedona-raster-gdal/src/register.rs +++ b/rust/sedona-raster-gdal/src/register.rs @@ -30,5 +30,7 @@ pub fn default_function_set() -> FunctionSet { function_set.insert_scalar_udf(crate::rs_reproject_match::rs_reproject_match_udf()); function_set.insert_scalar_udf(crate::rs_resample::rs_resample_udf()); function_set.insert_scalar_udf(crate::rs_tile::rs_tile_udf()); + function_set.insert_scalar_udf(crate::rs_zonal_stats::rs_zonal_stats_udf()); + function_set.insert_scalar_udf(crate::rs_zonal_stats::rs_zonal_stats_all_udf()); function_set } diff --git a/rust/sedona-raster-gdal/src/rs_clip.rs b/rust/sedona-raster-gdal/src/rs_clip.rs index 8134c01097..f6467210c6 100644 --- a/rust/sedona-raster-gdal/src/rs_clip.rs +++ b/rust/sedona-raster-gdal/src/rs_clip.rs @@ -34,10 +34,6 @@ use datafusion_common::{exec_datafusion_err, ScalarValue}; use datafusion_expr::{ColumnarValue, Volatility}; use sedona_common::sedona_internal_err; use sedona_gdal::gdal::Gdal; -use sedona_gdal::geo_transform::{GeoTransform, GeoTransformEx}; -use sedona_gdal::mem::MemDatasetBuilder; -use sedona_gdal::raster::types::GdalDataType; -use sedona_gdal::vector::geometry::Geometry; use arrow_schema::DataType; use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; @@ -53,8 +49,9 @@ use sedona_schema::datatypes::{SedonaType, RASTER}; use sedona_schema::matchers::ArgMatcher; use sedona_schema::raster::BandDataType; -use crate::gdal_common::with_gdal; +use crate::gdal_common::{raster_geo_transform, with_gdal}; use crate::gdal_dataset_provider::configure_thread_local_options; +use crate::mask::{envelope_window, rasterize_geometry_mask, PixelWindow}; use sedona_raster::traits::nodata_f64_to_bytes; /// RS_Clip() scalar UDF implementation @@ -364,16 +361,7 @@ struct ClippedRasterData { /// the geometry's envelope intersected with the raster extent, snapped /// outward to the pixel grid. `None` means the full original raster /// extent was kept (crop=false). - crop_window: Option, -} - -/// A rectangular crop window in pixel coordinates. -#[derive(Debug, Clone, Copy)] -struct CropWindow { - col_off: usize, - row_off: usize, - width: usize, - height: usize, + crop_window: Option, } /// Clip a raster to a geometry. @@ -399,14 +387,8 @@ fn clip_raster( .geometry_from_wkb(geom_wkb) .map_err(|e| exec_datafusion_err!("Failed to parse geometry from WKB: {}", e))?; - let geotransform = [ - metadata.upper_left_x(), - metadata.scale_x(), - metadata.skew_x(), - metadata.upper_left_y(), - metadata.skew_y(), - metadata.scale_y(), - ]; + // GDAL geotransform: [upper_left_x, scale_x, skew_x, upper_left_y, skew_y, scale_y]. + let geotransform = raster_geo_transform(raster)?; // The clip window is the geometry's envelope intersected with the raster // extent, snapped outward to the pixel grid — the window PostGIS ST_Clip, @@ -417,49 +399,17 @@ fn clip_raster( return Ok(None); }; - // Create a mask raster covering only the clip window, with the geotransform - // shifted to the window's upper-left corner. - let mask_dataset = - MemDatasetBuilder::create(gdal, window.width, window.height, 1, GdalDataType::UInt8) - .map_err(|e| exec_datafusion_err!("Failed to create mask dataset: {}", e))?; - let (window_ulx, window_uly) = geotransform.apply(window.col_off as f64, window.row_off as f64); - let mask_geotransform = [ - window_ulx, - geotransform[1], - geotransform[2], - window_uly, - geotransform[4], - geotransform[5], - ]; - mask_dataset - .set_geo_transform(&mask_geotransform) - .map_err(|e| exec_datafusion_err!("Failed to set geotransform: {}", e))?; - - // GDAL's MEM driver zero-fills owned band buffers at creation, so the mask - // already reads 0 (outside) everywhere; rasterize_affine burns 1 inside the - // geometry. No explicit zero-init write needed. - gdal.rasterize_affine( - &mask_dataset, - &[1], // band 1 - &[geometry], - &[1.0], // burn value = 1 (inside) + // Rasterize the geometry into a window-sized 0/1 mask: 1 inside, 0 outside + // (moves `geometry`, whose only remaining use is the burn). + let mut mask = Vec::new(); + rasterize_geometry_mask( + gdal, + geometry, + &geotransform, + &window, all_touched, - ) - .map_err(|e| exec_datafusion_err!("Failed to rasterize geometry: {}", e))?; - - // Read the (window-sized) mask - let mask_band = mask_dataset - .rasterband(1) - .map_err(|e| exec_datafusion_err!("Failed to get mask band: {}", e))?; - let mask_buffer = mask_band - .read_as::( - (0, 0), - (window.width, window.height), - (window.width, window.height), - None, - ) - .map_err(|e| exec_datafusion_err!("Failed to read mask: {}", e))?; - let mask = mask_buffer.data(); + &mut mask, + )?; // The envelope may overlap the raster while the geometry itself selects no // pixel (e.g. it falls between pixel centers); that is still the @@ -586,7 +536,7 @@ fn clip_raster( if let Some(cw) = crop_window { apply_mask_and_crop( plane_bytes, - mask, + &mask, width, &data_type, &nodata_bytes, @@ -596,7 +546,7 @@ fn clip_raster( } else { apply_mask_to_band( plane_bytes, - mask, + &mask, width, &data_type, &nodata_bytes, @@ -632,67 +582,6 @@ fn clip_raster( })) } -/// Compute the clip window: the geometry's envelope intersected with the -/// raster extent, snapped outward to the pixel grid. Returns `None` when the -/// envelope is disjoint from the raster extent (no clipping possible). -/// -/// The envelope corners are mapped through the inverse geotransform (all four, -/// so a skewed/rotated raster still gets a correct superset window) and the -/// resulting pixel-space bbox is floored/ceiled to whole pixels. A degenerate -/// envelope (point/line) landing exactly on a grid line is widened to one -/// pixel so the rasterizer — not the snapping — decides whether it burns. -fn envelope_window( - geometry: &Geometry, - geotransform: &GeoTransform, - width: usize, - height: usize, -) -> Result> { - let env = geometry.envelope(); - let inverse = geotransform - .invert() - .map_err(|e| exec_datafusion_err!("RS_Clip: geotransform is not invertible: {}", e))?; - - let corners = [ - (env.MinX, env.MinY), - (env.MinX, env.MaxY), - (env.MaxX, env.MinY), - (env.MaxX, env.MaxY), - ]; - let mut min_col = f64::INFINITY; - let mut max_col = f64::NEG_INFINITY; - let mut min_row = f64::INFINITY; - let mut max_row = f64::NEG_INFINITY; - for (x, y) in corners { - let (col, row) = inverse.apply(x, y); - min_col = min_col.min(col); - max_col = max_col.max(col); - min_row = min_row.min(row); - max_row = max_row.max(row); - } - - let col0 = min_col.floor(); - let row0 = min_row.floor(); - let col1 = max_col.ceil().max(col0 + 1.0); - let row1 = max_row.ceil().max(row0 + 1.0); - - // Intersect with the raster extent. `>=` also rejects the NaN envelope of - // an empty geometry. - let col0 = col0.max(0.0); - let row0 = row0.max(0.0); - let col1 = col1.min(width as f64); - let row1 = row1.min(height as f64); - if !(col0 < col1 && row0 < row1) { - return Ok(None); - } - - Ok(Some(CropWindow { - col_off: col0 as usize, - row_off: row0 as usize, - width: (col1 - col0) as usize, - height: (row1 - row0) as usize, - })) -} - /// Apply mask to band data (no cropping — preserves original dimensions). /// The mask covers only `window`; every pixel outside it is outside the /// geometry's envelope and therefore nodata. The plane's bytes are appended @@ -704,7 +593,7 @@ fn apply_mask_to_band( width: usize, data_type: &BandDataType, nodata_bytes: &[u8], - window: &CropWindow, + window: &PixelWindow, out: &mut Vec, ) -> Result<()> { let byte_size = data_type.byte_size(); @@ -754,7 +643,7 @@ fn apply_mask_and_crop( full_width: usize, data_type: &BandDataType, nodata_bytes: &[u8], - cw: &CropWindow, + cw: &PixelWindow, out: &mut Vec, ) -> Result<()> { let byte_size = data_type.byte_size(); diff --git a/rust/sedona-raster-gdal/src/rs_zonal_stats.rs b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs new file mode 100644 index 0000000000..a82bb80285 --- /dev/null +++ b/rust/sedona-raster-gdal/src/rs_zonal_stats.rs @@ -0,0 +1,1756 @@ +// 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_ZonalStats / RS_ZonalStatsAll UDFs — summary statistics of the raster +//! pixels covered by a roi geometry. +//! +//! Both mirror Apache Sedona Spark's positional overloads verbatim so that +//! Spark SQL tends to run unchanged. `RS_ZonalStats` returns one statistic as a +//! `Float64`: +//! +//! - `RS_ZonalStats(raster, roi, stat)` +//! - `RS_ZonalStats(raster, roi, band, stat)` +//! - `RS_ZonalStats(raster, roi, band, stat, all_touched)` +//! - `RS_ZonalStats(raster, roi, band, stat, all_touched, exclude_no_data)` +//! - `RS_ZonalStats(raster, roi, band, stat, all_touched, exclude_no_data, lenient)` +//! +//! `RS_ZonalStatsAll` returns every statistic as a struct, with the same ladder +//! minus `stat`: +//! +//! - `RS_ZonalStatsAll(raster, roi)` +//! - `RS_ZonalStatsAll(raster, roi, band)` +//! - `RS_ZonalStatsAll(raster, roi, band, all_touched)` +//! - `RS_ZonalStatsAll(raster, roi, band, all_touched, exclude_no_data)` +//! - `RS_ZonalStatsAll(raster, roi, band, all_touched, exclude_no_data, lenient)` +//! +//! A pixel is included when its centre falls inside the roi (or that the roi +//! merely touches, with `all_touched`), optionally excluding the band's nodata +//! value. `all_touched` defaults to false, `exclude_no_data` to true, and +//! `lenient` to true. Unlike Sedona Spark, the band-less overloads do not +//! default to band 1 on a multiband raster: naming the band is required there +//! (a single-band raster resolves unambiguously). +//! +//! These functions operate on 2-D `(y, x)` bands. A band that is not a 2-D +//! spatial grid is rejected; computing a statistic per non-spatial plane of an +//! N-D band is not supported. + +use std::sync::Arc; + +use arrow_array::builder::{Float64Builder, Int64Builder}; +use arrow_array::{ArrayRef, BooleanArray, Int64Array, StringArray, StructArray}; +use arrow_buffer::{BooleanBufferBuilder, NullBuffer}; +use arrow_schema::{DataType, Field, Fields}; +use datafusion_common::cast::{as_boolean_array, as_int64_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_err; +use sedona_expr::scalar_udf::{SedonaScalarKernel, SedonaScalarUDF}; +use sedona_gdal::gdal::Gdal; +use sedona_raster::array::RasterRefImpl; +use sedona_raster::traits::RasterRef; +use sedona_raster_functions::crs_utils::{align_wkb_to_crs, resolve_crs, with_crs_engine}; +use sedona_raster_functions::rs_ensure_loaded::NEEDS_PIXELS_METADATA_KEY; +use sedona_raster_functions::rs_spatial_predicates::raster_intersects_geom_wkb; +use sedona_raster_functions::RasterExecutor; +use sedona_schema::datatypes::SedonaType; +use sedona_schema::matchers::ArgMatcher; +use sedona_schema::raster::BandDataType; + +use crate::gdal_common::{raster_geo_transform, with_gdal}; +use crate::gdal_dataset_provider::configure_thread_local_options; +use crate::mask::{envelope_window, rasterize_geometry_mask, PixelWindow}; + +/// The statistics RS_ZonalStatsAll returns, in the order Sedona Spark reports +/// them. RS_ZonalStats selects one of these by name. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StatType { + Count, + Sum, + Mean, + Median, + Mode, + StdDev, + Variance, + Min, + Max, +} + +impl StatType { + /// Parse a statistic name (case-insensitive). Aliases match Sedona Spark + /// (`avg`/`average` for mean, `sd` for stddev). + fn from_str(s: &str) -> Option { + match s.to_lowercase().as_str() { + "count" => Some(StatType::Count), + "sum" => Some(StatType::Sum), + "mean" | "avg" | "average" => Some(StatType::Mean), + "median" => Some(StatType::Median), + "mode" => Some(StatType::Mode), + "stddev" | "sd" => Some(StatType::StdDev), + "variance" => Some(StatType::Variance), + "min" => Some(StatType::Min), + "max" => Some(StatType::Max), + _ => None, + } + } +} + +/// Defaults for the trailing flags, applied by the narrower overloads that omit +/// them (matching Sedona Spark). +const DEFAULT_ALL_TOUCHED: bool = false; +const DEFAULT_EXCLUDE_NODATA: bool = true; +const DEFAULT_LENIENT: bool = true; + +/// The resolved parameters for one row's zonal-stats computation, assembled from +/// the positional arguments the matched overload carried. +#[derive(Debug, Clone)] +struct ZonalStatsParams { + /// 1-based band to compute over. `None` means "resolve the implicit band": + /// band 1 for a single-band raster, an error for a multiband raster (naming + /// the band is required rather than silently getting band 1). Only the + /// band-less overloads leave this `None`. + band: Option, + /// Include every pixel the roi touches, not only those whose centre it + /// covers. + all_touched: bool, + /// Skip pixels equal to the band's nodata value. + exclude_no_data: bool, + /// Return NULL when the roi does not intersect the raster, rather than + /// erroring. Only the no-intersection case is softened; malformed geometry + /// or an unreadable band always errors. + lenient: bool, +} + +/// Every statistic for a roi. `count` is always present (0 when the roi +/// selects no pixels); the remaining fields are `None` in exactly that +/// no-pixel case and `Some` otherwise, mirroring Sedona Spark (which returns +/// `count = 0` and NULL for the rest). +#[derive(Debug, Clone, PartialEq)] +struct ZonalStatistics { + count: i64, + sum: Option, + mean: Option, + median: Option, + mode: Option, + stddev: Option, + variance: Option, + min: Option, + max: Option, +} + +/// Whether a roi geometry intersects the raster. `NoIntersection` mirrors Sedona +/// Spark's `rsIntersects` gate (the caller turns it into NULL when `lenient`, an +/// error otherwise); `Collected` means the selected pixel values are in the +/// caller's scratch buffer, possibly empty for a roi that intersects the +/// footprint but selects no pixel centre (a `count = 0` result). +enum RoiCoverage { + NoIntersection, + Collected, +} + +// ============================================================================= +// RS_ZonalStats +// ============================================================================= + +/// `RS_ZonalStats` — one statistic as a `Float64`. `stat` is a statistic name +/// (`count`, `sum`, `mean`, `median`, `mode`, `stddev`, `variance`, `min`, +/// `max`). See the module docs for the full positional overload ladder. +pub fn rs_zonal_stats_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_zonalstats", + vec![ + Arc::new(RsZonalStats { arg_count: 3 }), // (raster, roi, stat) + Arc::new(RsZonalStats { arg_count: 4 }), // (raster, roi, band, stat) + Arc::new(RsZonalStats { arg_count: 5 }), // + all_touched + Arc::new(RsZonalStats { arg_count: 6 }), // + exclude_no_data + Arc::new(RsZonalStats { arg_count: 7 }), // + lenient + ], + Volatility::Immutable, + ) + // Reads band pixels, so the planner materializes OutDb rasters via + // RS_EnsureLoaded first. + .with_metadata(NEEDS_PIXELS_METADATA_KEY, "true") +} + +#[derive(Debug)] +struct RsZonalStats { + /// Number of arguments in the matched signature (3..=7). + arg_count: usize, +} + +impl SedonaScalarKernel for RsZonalStats { + fn return_type(&self, args: &[SedonaType]) -> Result> { + // Argument order mirrors Sedona Spark: (raster, roi, [band,] stat, + // [all_touched, [exclude_no_data, [lenient]]]). The 3-arg overload omits + // band (its stat is at index 2); the 4+-arg overloads carry band at + // index 2 and stat at index 3. + let matchers = match self.arg_count { + 3 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_string(), + ], + 4 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_string(), + ], + 5 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_string(), + ArgMatcher::is_boolean(), + ], + 6 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_string(), + ArgMatcher::is_boolean(), + ArgMatcher::is_boolean(), + ], + 7 => vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ArgMatcher::is_integer(), + ArgMatcher::is_string(), + ArgMatcher::is_boolean(), + ArgMatcher::is_boolean(), + ArgMatcher::is_boolean(), + ], + _ => { + return sedona_internal_err!( + "RS_ZonalStats: unexpected arg_count {}", + self.arg_count + ); + } + }; + let matcher = ArgMatcher::new(matchers, SedonaType::Arrow(DataType::Float64)); + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result { + 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 { + let num_iterations = RasterExecutor::num_iterations_over(args); + + // band (index 2) only exists in the 4+-arg overloads; the 3-arg overload + // leaves it implicit. stat is at index 2 (3-arg) or 3 (4+-arg). + let has_band = self.arg_count >= 4; + let stat_idx = if has_band { 3 } else { 2 }; + let stat_array = expand_string_arg(&args[stat_idx], num_iterations)?; + let mut stat_iter = stat_array.iter(); + + let band_array = has_band + .then(|| expand_int64_arg(&args[2], num_iterations)) + .transpose()?; + let mut band_iter = band_array.as_ref().map(|a| a.iter()); + + // all_touched (index 4), exclude_no_data (index 5), lenient (index 6): + // read from the column when the overload carries it, else the default. + let all_touched_array = expand_flag( + args, + 4, + self.arg_count >= 5, + DEFAULT_ALL_TOUCHED, + num_iterations, + )?; + let exclude_no_data_array = expand_flag( + args, + 5, + self.arg_count >= 6, + DEFAULT_EXCLUDE_NODATA, + num_iterations, + )?; + let lenient_array = expand_flag( + args, + 6, + self.arg_count >= 7, + DEFAULT_LENIENT, + num_iterations, + )?; + let mut all_touched_iter = all_touched_array.iter(); + let mut exclude_no_data_iter = exclude_no_data_array.iter(); + let mut lenient_iter = lenient_array.iter(); + + let mut builder = Float64Builder::with_capacity(num_iterations); + let mut scratch: Vec = Vec::new(); + let mut mask_scratch: Vec = Vec::new(); + + // The executor only sees (raster, roi); the option columns are advanced + // in lockstep below. + let exec_arg_types = [arg_types[0].clone(), arg_types[1].clone()]; + let exec_args = [args[0].clone(), args[1].clone()]; + let executor = + RasterExecutor::new_with_num_iterations(&exec_arg_types, &exec_args, num_iterations); + + with_gdal(|gdal| { + configure_thread_local_options(gdal, config_options)?; + with_crs_engine(config_options, |engine| { + executor.execute_raster_wkb_crs_void(|raster_opt, wkb_opt, geom_crs| { + let stat_str = stat_iter.next().flatten(); + let Some(params) = next_params( + &mut band_iter, + &mut all_touched_iter, + &mut exclude_no_data_iter, + &mut lenient_iter, + ) else { + builder.append_null(); + return Ok(()); + }; + + // A NULL stat, raster, or roi propagates to a NULL row. + let (Some(stat_str), Some(raster), Some(wkb)) = (stat_str, raster_opt, wkb_opt) + else { + builder.append_null(); + return Ok(()); + }; + let stat_type = StatType::from_str(stat_str).ok_or_else(|| { + exec_datafusion_err!("RS_ZonalStats: unknown statistic {stat_str:?}") + })?; + + // Reproject the roi into the raster's CRS, borrowing it + // unchanged when the CRSes already match; a CRS on exactly + // one side is an error, since it would mislocate the roi. + let raster_crs = resolve_crs(raster.crs())?; + let geom_wkb = align_wkb_to_crs( + wkb, + geom_crs, + raster_crs.as_deref(), + "geometry", + "raster", + engine, + )?; + match collect_zonal_values( + gdal, + raster, + &geom_wkb, + ¶ms, + &mut scratch, + &mut mask_scratch, + )? { + // Compute only the requested statistic, not all of them. + RoiCoverage::Collected => { + match compute_single_statistic(&mut scratch, stat_type) { + Some(value) => builder.append_value(value), + None => builder.append_null(), + } + } + // The roi does not intersect the raster: NULL when + // lenient (the default), an error otherwise. + RoiCoverage::NoIntersection if params.lenient => builder.append_null(), + RoiCoverage::NoIntersection => return no_intersection_err(), + } + Ok(()) + }) + })?; + + let out: ArrayRef = Arc::new(builder.finish()); + RasterExecutor::finish_over(args, out) + }) + } +} + +// ============================================================================= +// RS_ZonalStatsAll +// ============================================================================= + +/// `RS_ZonalStatsAll` — every statistic as a struct with fields `count, sum, +/// mean, median, mode, stddev, variance, min, max`. See the module docs for the +/// full positional overload ladder. +pub fn rs_zonal_stats_all_udf() -> SedonaScalarUDF { + SedonaScalarUDF::new( + "rs_zonalstatsall", + vec![ + Arc::new(RsZonalStatsAll { arg_count: 2 }), // (raster, roi) + Arc::new(RsZonalStatsAll { arg_count: 3 }), // (raster, roi, band) + Arc::new(RsZonalStatsAll { arg_count: 4 }), // + all_touched + Arc::new(RsZonalStatsAll { arg_count: 5 }), // + exclude_no_data + Arc::new(RsZonalStatsAll { arg_count: 6 }), // + lenient + ], + Volatility::Immutable, + ) + .with_metadata(NEEDS_PIXELS_METADATA_KEY, "true") +} + +#[derive(Debug)] +struct RsZonalStatsAll { + /// Number of arguments in the matched signature (2..=6). + arg_count: usize, +} + +impl SedonaScalarKernel for RsZonalStatsAll { + fn return_type(&self, args: &[SedonaType]) -> Result> { + // Argument order mirrors Sedona Spark: (raster, roi, [band, + // [all_touched, [exclude_no_data, [lenient]]]]). The 2-arg overload omits + // band; the 3+-arg overloads carry it at index 2. + let mut matchers = vec![ + ArgMatcher::is_raster(), + ArgMatcher::is_geometry_or_geography(), + ]; + if self.arg_count >= 3 { + matchers.push(ArgMatcher::is_integer()); // band + } + for _ in 4..=self.arg_count { + matchers.push(ArgMatcher::is_boolean()); // all_touched, exclude_no_data, lenient + } + if self.arg_count < 2 || self.arg_count > 6 { + return sedona_internal_err!( + "RS_ZonalStatsAll: unexpected arg_count {}", + self.arg_count + ); + } + let matcher = ArgMatcher::new(matchers, SedonaType::Arrow(zonal_stats_struct_type())); + matcher.match_args(args) + } + + fn invoke_batch( + &self, + arg_types: &[SedonaType], + args: &[ColumnarValue], + ) -> Result { + 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 { + let num_iterations = RasterExecutor::num_iterations_over(args); + + // band (index 2) only exists in the 3+-arg overloads; the 2-arg overload + // leaves it implicit. all_touched (index 3), exclude_no_data (index 4), + // and lenient (index 5) follow. + let band_array = (self.arg_count >= 3) + .then(|| expand_int64_arg(&args[2], num_iterations)) + .transpose()?; + let mut band_iter = band_array.as_ref().map(|a| a.iter()); + + let all_touched_array = expand_flag( + args, + 3, + self.arg_count >= 4, + DEFAULT_ALL_TOUCHED, + num_iterations, + )?; + let exclude_no_data_array = expand_flag( + args, + 4, + self.arg_count >= 5, + DEFAULT_EXCLUDE_NODATA, + num_iterations, + )?; + let lenient_array = expand_flag( + args, + 5, + self.arg_count >= 6, + DEFAULT_LENIENT, + num_iterations, + )?; + let mut all_touched_iter = all_touched_array.iter(); + let mut exclude_no_data_iter = exclude_no_data_array.iter(); + let mut lenient_iter = lenient_array.iter(); + + let mut builders = ZonalStatsBuilders::with_capacity(num_iterations); + let mut scratch: Vec = Vec::new(); + let mut mask_scratch: Vec = Vec::new(); + + let exec_arg_types = [arg_types[0].clone(), arg_types[1].clone()]; + let exec_args = [args[0].clone(), args[1].clone()]; + let executor = + RasterExecutor::new_with_num_iterations(&exec_arg_types, &exec_args, num_iterations); + + with_gdal(|gdal| { + configure_thread_local_options(gdal, config_options)?; + with_crs_engine(config_options, |engine| { + executor.execute_raster_wkb_crs_void(|raster_opt, wkb_opt, geom_crs| { + let Some(params) = next_params( + &mut band_iter, + &mut all_touched_iter, + &mut exclude_no_data_iter, + &mut lenient_iter, + ) else { + builders.push_null(); + return Ok(()); + }; + + let (Some(raster), Some(wkb)) = (raster_opt, wkb_opt) else { + builders.push_null(); + return Ok(()); + }; + + let raster_crs = resolve_crs(raster.crs())?; + let geom_wkb = align_wkb_to_crs( + wkb, + geom_crs, + raster_crs.as_deref(), + "geometry", + "raster", + engine, + )?; + match collect_zonal_values( + gdal, + raster, + &geom_wkb, + ¶ms, + &mut scratch, + &mut mask_scratch, + )? { + RoiCoverage::Collected => { + builders.push_stats(&compute_statistics(&mut scratch)) + } + RoiCoverage::NoIntersection if params.lenient => builders.push_null(), + RoiCoverage::NoIntersection => return no_intersection_err(), + } + Ok(()) + }) + })?; + + let out: ArrayRef = Arc::new(builders.finish()); + RasterExecutor::finish_over(args, out) + }) + } +} + +/// Struct data type RS_ZonalStatsAll returns. +fn zonal_stats_struct_type() -> DataType { + DataType::Struct(zonal_stats_struct_fields()) +} + +/// Fields of the RS_ZonalStatsAll struct, in Sedona Spark order. `count` is an +/// `Int64` (a whole pixel count); every other statistic is a `Float64`. +fn zonal_stats_struct_fields() -> Fields { + Fields::from(vec![ + Field::new("count", DataType::Int64, true), + Field::new("sum", DataType::Float64, true), + Field::new("mean", DataType::Float64, true), + Field::new("median", DataType::Float64, true), + Field::new("mode", DataType::Float64, true), + Field::new("stddev", DataType::Float64, true), + Field::new("variance", DataType::Float64, true), + Field::new("min", DataType::Float64, true), + Field::new("max", DataType::Float64, true), + ]) +} + +/// Column builders for the RS_ZonalStatsAll struct output: one typed builder per +/// field plus an outer struct-level validity buffer. +/// +/// Building the columns directly and assembling the [`StructArray`] once at the +/// end avoids downcasting a `StructBuilder`'s boxed field builders on every row. +struct ZonalStatsBuilders { + count: Int64Builder, + /// `sum, mean, median, mode, stddev, variance, min, max` — Sedona Spark field + /// order, matching [`zonal_stats_struct_fields`] after `count`. + floats: [Float64Builder; 8], + /// Struct-level null bitmap: `false` for a fully-NULL row (a NULL input or a + /// non-intersecting roi under `lenient`). + validity: BooleanBufferBuilder, +} + +impl ZonalStatsBuilders { + fn with_capacity(capacity: usize) -> Self { + Self { + count: Int64Builder::with_capacity(capacity), + floats: std::array::from_fn(|_| Float64Builder::with_capacity(capacity)), + validity: BooleanBufferBuilder::new(capacity), + } + } + + /// Append one computed-stats row (the struct itself is valid). The float + /// fields carry through the `Option`, so an empty roi records `count = 0` + /// with the rest NULL. + fn push_stats(&mut self, stats: &ZonalStatistics) { + self.count.append_value(stats.count); + let values = [ + stats.sum, + stats.mean, + stats.median, + stats.mode, + stats.stddev, + stats.variance, + stats.min, + stats.max, + ]; + for (builder, value) in self.floats.iter_mut().zip(values) { + builder.append_option(value); + } + self.validity.append(true); + } + + /// Append a fully-NULL struct row (a NULL input, or a non-intersecting roi + /// when `lenient`). Every field is null and the struct itself is null. + fn push_null(&mut self) { + self.count.append_null(); + for builder in &mut self.floats { + builder.append_null(); + } + self.validity.append(false); + } + + /// Assemble the accumulated columns into the struct array. + fn finish(mut self) -> StructArray { + let mut arrays: Vec = Vec::with_capacity(9); + arrays.push(Arc::new(self.count.finish())); + for builder in &mut self.floats { + arrays.push(Arc::new(builder.finish())); + } + let nulls = NullBuffer::new(self.validity.finish()); + StructArray::new(zonal_stats_struct_fields(), arrays, Some(nulls)) + } +} + +// ============================================================================= +// Core computation +// ============================================================================= + +/// Collect the pixel values a roi geometry selects on one band into `scratch`. +/// +/// Returns [`RoiCoverage::NoIntersection`] when the roi geometry does not +/// intersect the raster's footprint — a true geometry intersection (matching +/// Sedona Spark's `rsIntersects` gate), not a bounding-box overlap: a roi whose +/// envelope overlaps the raster but whose geometry is disjoint is a +/// no-intersection case. The caller turns that into NULL when `lenient`, an +/// error otherwise. A roi that intersects the footprint but whose selected +/// pixels are all outside the geometry or all nodata returns +/// [`RoiCoverage::Collected`] with `scratch` left empty (a `count = 0` result). +/// +/// The caller computes the statistic(s) it needs from `scratch` — every one for +/// RS_ZonalStatsAll, only the requested one for RS_ZonalStats — so this shared +/// collection never computes statistics the caller would discard. +/// +/// `scratch` is a reused buffer for the selected pixel values and `mask_scratch` +/// for the rasterized roi mask, both reused so the per-row computation does not +/// allocate a fresh `Vec` each call. +fn collect_zonal_values( + gdal: &Gdal, + raster: &RasterRefImpl<'_>, + geom_wkb: &[u8], + params: &ZonalStatsParams, + scratch: &mut Vec, + mask_scratch: &mut Vec, +) -> Result { + let num_bands = raster.num_bands(); + let band_num = resolve_band(params.band, num_bands)?; + + let band = raster + .bands() + .band(band_num) + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to read band {band_num}: {e}"))?; + if !band.is_spatial_2d() { + return exec_err!( + "RS_ZonalStats supports 2-D rasters only; band {band_num} is not a 2-D (y, x) grid" + ); + } + let data_type = band.data_type(); + let byte_size = data_type.byte_size(); + + let metadata = raster.metadata(); + let transform = raster_geo_transform(raster)?; + let width = usize::try_from(metadata.width()) + .map_err(|_| exec_datafusion_err!("RS_ZonalStats: negative raster width"))?; + let height = usize::try_from(metadata.height()) + .map_err(|_| exec_datafusion_err!("RS_ZonalStats: negative raster height"))?; + + // No-intersection gate: a true geometry intersection between the roi and + // the raster footprint (matching Sedona Spark's rsIntersects gate), not a + // bounding-box overlap. A roi whose envelope overlaps the raster but whose + // geometry is disjoint is a no-intersection case, not a count-0 case. The + // roi is already in the raster's CRS here, so no transform is needed. + if !raster_intersects_geom_wkb(raster, geom_wkb)? { + return Ok(RoiCoverage::NoIntersection); + } + + // Parse the roi and clamp its envelope to the raster grid for the pixel + // window to rasterize. The gate above already established overlap; a + // degenerate window (the roi only touches the raster boundary) selects no + // pixels, so it is count 0 rather than no-intersection. + let geometry = gdal + .geometry_from_wkb(geom_wkb) + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to parse geometry: {e}"))?; + let Some(window) = envelope_window(&geometry, &transform, width, height)? else { + scratch.clear(); + return Ok(RoiCoverage::Collected); + }; + + // Rasterize the roi into a window-sized 0/1 mask (moves `geometry`, whose + // only remaining use is the burn). The mask reuses `mask_scratch` across + // rows rather than allocating a fresh buffer each call. + rasterize_geometry_mask( + gdal, + geometry, + &transform, + &window, + params.all_touched, + mask_scratch, + )?; + + // Read the band once (zero-copy borrow) and collect the selected values. + let nd_buffer = band + .nd_buffer() + .map_err(|e| exec_datafusion_err!("RS_ZonalStats: failed to read band {band_num}: {e}"))?; + let band_bytes = nd_buffer.as_contiguous().map_err(|e| { + exec_datafusion_err!("RS_ZonalStats: band {band_num} is not contiguous: {e}") + })?; + let expected = width + .checked_mul(height) + .and_then(|n| n.checked_mul(byte_size)) + .ok_or_else(|| exec_datafusion_err!("RS_ZonalStats: raster dimensions overflow"))?; + if band_bytes.len() != expected { + return sedona_internal_err!( + "RS_ZonalStats: band {band_num} byte length {} does not match {width}x{height} of {data_type:?}", + band_bytes.len() + ); + } + + // Nodata is compared in the band's own byte representation, never through + // f64 — an Int64/UInt64 nodata beyond 2^53 must not alias a nearby pixel. + let nodata = if params.exclude_no_data { + band.nodata() + } else { + None + }; + if let Some(nd) = nodata { + if nd.len() != byte_size { + return sedona_internal_err!( + "RS_ZonalStats: band {band_num} nodata is {} bytes, expected {byte_size} for {data_type:?}", + nd.len() + ); + } + } + + scratch.clear(); + collect_masked_values( + band_bytes, + data_type, + width, + &window, + mask_scratch, + nodata, + scratch, + ); + + Ok(RoiCoverage::Collected) +} + +/// Resolve the 1-based band to use. `Some(b)` must be a valid 1-based index; +/// `None` defaults to band 1 for a single-band raster and errors for a +/// multiband raster (matching the codebase's `default_band` convention, which +/// refuses to silently pick band 1 when the choice is ambiguous). +fn resolve_band(band: Option, num_bands: usize) -> Result { + match band { + Some(b) => { + if b < 1 { + return exec_err!("RS_ZonalStats: band must be >= 1, got {b}"); + } + let b = b as usize; + if b > num_bands { + return exec_err!("RS_ZonalStats: band {b} is out of range (1-{num_bands})"); + } + Ok(b) + } + None => { + if num_bands == 1 { + Ok(1) + } else { + exec_err!( + "RS_ZonalStats: raster has {num_bands} bands; pass the band argument to \ + choose one (only a single-band raster may omit it)" + ) + } + } + } +} + +/// Append every selected pixel value (masked in, and — when `nodata` is set — +/// not byte-equal to the nodata sentinel) to `out` as `f64`. +/// +/// The data type is dispatched once, outside the loop, so the per-pixel body is +/// a fixed-width little-endian read plus the mask/nodata comparisons rather +/// than a per-pixel type match. +fn collect_masked_values( + band_bytes: &[u8], + data_type: BandDataType, + width: usize, + window: &PixelWindow, + mask: &[u8], + nodata: Option<&[u8]>, + out: &mut Vec, +) { + macro_rules! collect { + ($t:ty, $n:literal) => {{ + for row in 0..window.height { + let src_row = window.row_off + row; + let mask_row = row * window.width; + for col in 0..window.width { + if mask[mask_row + col] == 0 { + continue; + } + let idx = (src_row * width + window.col_off + col) * $n; + let px = &band_bytes[idx..idx + $n]; + if let Some(nd) = nodata { + if px == nd { + continue; + } + } + let mut arr = [0u8; $n]; + arr.copy_from_slice(px); + out.push(<$t>::from_le_bytes(arr) as f64); + } + } + }}; + } + + match data_type { + BandDataType::UInt8 => collect!(u8, 1), + BandDataType::Int8 => collect!(i8, 1), + BandDataType::UInt16 => collect!(u16, 2), + BandDataType::Int16 => collect!(i16, 2), + BandDataType::UInt32 => collect!(u32, 4), + BandDataType::Int32 => collect!(i32, 4), + BandDataType::UInt64 => collect!(u64, 8), + BandDataType::Int64 => collect!(i64, 8), + BandDataType::Float32 => collect!(f32, 4), + BandDataType::Float64 => collect!(f64, 8), + } +} + +/// Compute every statistic from the selected pixel values (for +/// RS_ZonalStatsAll, which returns all of them). +/// +/// An empty slice yields `count = 0` and NULL for the rest (Sedona Spark's +/// empty-roi shortcut). Variance is the sample (n-1) variance, matching Spark; +/// for a single pixel it is 0. Median is the linear-interpolated 50th +/// percentile, which reduces to the middle element (odd n) or the mean of the +/// two central elements (even n). Mode is the most frequent value, breaking ties +/// toward the larger value. +/// +/// `values` is sorted in place (for min, max, median, and mode); the caller owns +/// it as reusable scratch. +fn compute_statistics(values: &mut [f64]) -> ZonalStatistics { + let count = values.len() as i64; + if values.is_empty() { + return ZonalStatistics { + count: 0, + sum: None, + mean: None, + median: None, + mode: None, + stddev: None, + variance: None, + min: None, + max: None, + }; + } + + // A NaN pixel (e.g. a float band whose NaN nodata was not excluded) poisons + // every statistic under numpy semantics. Return NaN for all of them rather + // than letting f64::min / f64::max silently skip NaN while sum and mean + // propagate it — an internally inconsistent, reference-diverging result. + if values.iter().any(|v| v.is_nan()) { + return ZonalStatistics { + count, + sum: Some(f64::NAN), + mean: Some(f64::NAN), + median: Some(f64::NAN), + mode: Some(f64::NAN), + stddev: Some(f64::NAN), + variance: Some(f64::NAN), + min: Some(f64::NAN), + max: Some(f64::NAN), + }; + } + + let sum: f64 = values.iter().sum(); + let mean = sum / count as f64; + let variance = sample_variance(values, mean); + let stddev = variance.sqrt(); + + // min, max, median, and mode all read the values in sorted order: sort once + // in place, then take the extremes from the ends instead of folding again. + sort_values(values); + let min = values[0]; + let max = values[values.len() - 1]; + let median = median_of_sorted(values); + let mode = mode_of_sorted(values); + + ZonalStatistics { + count, + sum: Some(sum), + mean: Some(mean), + median: Some(median), + mode: Some(mode), + stddev: Some(stddev), + variance: Some(variance), + min: Some(min), + max: Some(max), + } +} + +/// Compute only `stat` from the selected pixel values (for RS_ZonalStats, which +/// returns a single statistic). `values` is sorted in place only when the +/// requested statistic — median or mode — needs ordered data, so the simple +/// statistics do not pay for a sort. +/// +/// Mirrors [`compute_statistics`]' empty/NaN semantics: `count` is always +/// defined (0 for an empty roi); every other statistic is NULL (`None`) for an +/// empty roi, and NaN when a NaN pixel is present. +fn compute_single_statistic(values: &mut [f64], stat: StatType) -> Option { + if stat == StatType::Count { + return Some(values.len() as f64); + } + if values.is_empty() { + return None; + } + if values.iter().any(|v| v.is_nan()) { + return Some(f64::NAN); + } + Some(match stat { + // Count is handled before the value checks above. + StatType::Count => unreachable!("count returns early"), + StatType::Sum => values.iter().sum(), + StatType::Mean => mean_of(values), + StatType::Min => values.iter().copied().fold(f64::INFINITY, f64::min), + StatType::Max => values.iter().copied().fold(f64::NEG_INFINITY, f64::max), + StatType::Variance => sample_variance(values, mean_of(values)), + StatType::StdDev => sample_variance(values, mean_of(values)).sqrt(), + StatType::Median => { + sort_values(values); + median_of_sorted(values) + } + StatType::Mode => { + sort_values(values); + mode_of_sorted(values) + } + }) +} + +/// The arithmetic mean of `values`, which must be non-empty. +fn mean_of(values: &[f64]) -> f64 { + values.iter().sum::() / values.len() as f64 +} + +/// The sample (n-1) variance of `values` about their precomputed `mean`, +/// matching Sedona Spark; 0 for a single value. Taking the mean as an argument +/// lets the all-statistics path reuse the mean it already computed rather than +/// summing again; the squared-deviation form (rather than the naive +/// sum-of-squares) avoids catastrophic cancellation. `values` must be non-empty +/// and NaN-free. +fn sample_variance(values: &[f64], mean: f64) -> f64 { + let n = values.len(); + if n <= 1 { + return 0.0; + } + let sum_sq: f64 = values.iter().map(|&v| (v - mean).powi(2)).sum(); + sum_sq / (n as f64 - 1.0) +} + +/// Sort `values` ascending in place. Callers exclude NaN beforehand, so the +/// `partial_cmp` fallback is never exercised. +fn sort_values(values: &mut [f64]) { + values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); +} + +/// The linear-interpolated 50th percentile of `sorted` (ascending): the middle +/// element for an odd count, the mean of the two central elements for an even +/// count. `sorted` must be non-empty. +fn median_of_sorted(sorted: &[f64]) -> f64 { + let mid = sorted.len() / 2; + if sorted.len().is_multiple_of(2) { + (sorted[mid - 1] + sorted[mid]) / 2.0 + } else { + sorted[mid] + } +} + +/// The most frequent value in `sorted` (ascending), breaking ties toward the +/// larger value (matching Sedona Spark's `StatUtils.mode`). Equal values are +/// adjacent once sorted, so a single pass over the runs finds the mode without a +/// map. `sorted` must be non-empty. +fn mode_of_sorted(sorted: &[f64]) -> f64 { + let mut best_val = sorted[0]; + let mut best_len = 1usize; + let mut run_len = 1usize; + for i in 1..sorted.len() { + run_len = if sorted[i] == sorted[i - 1] { + run_len + 1 + } else { + 1 + }; + // `>=` keeps the later — and, since ascending, larger — value on a tie. + if run_len >= best_len { + best_len = run_len; + best_val = sorted[i]; + } + } + best_val +} + +// ============================================================================= +// Argument helpers +// ============================================================================= + +/// The error returned for a non-intersecting roi when `lenient` is off. +fn no_intersection_err() -> Result { + exec_err!( + "RS_ZonalStats: the roi geometry does not intersect the raster; \ + pass lenient => true to return NULL instead" + ) +} + +/// Advance the option iterators one row (keeping every column in lockstep) and +/// assemble the resolved params, or return `None` when an explicit-but-NULL +/// value makes this a NULL output row. +/// +/// A `band_iter` of `None` is the band-less overload, whose implicit band stays +/// `None` (resolved to band 1 for a single-band raster, an error otherwise). A +/// `Some` band iterator carrying a NULL, or any NULL flag, is a NULL row. +fn next_params( + band_iter: &mut Option, + all_touched_iter: &mut F, + exclude_no_data_iter: &mut F, + lenient_iter: &mut F, +) -> Option +where + B: Iterator>, + F: Iterator>, +{ + // Advance every iterator first so a NULL-driven early return does not desync + // the columns on the next row. + let band_cell = band_iter.as_mut().map(|iter| iter.next().flatten()); + let all_touched = all_touched_iter.next().flatten(); + let exclude_no_data = exclude_no_data_iter.next().flatten(); + let lenient = lenient_iter.next().flatten(); + + let band = match band_cell { + None => None, // band-less overload: implicit band + Some(Some(b)) => Some(b), // explicit band + Some(None) => return None, // explicit NULL band -> NULL row + }; + Some(ZonalStatsParams { + band, + all_touched: all_touched?, + exclude_no_data: exclude_no_data?, + lenient: lenient?, + }) +} + +/// The boolean flag column at `args[index]` when the overload carries it +/// (`present`), otherwise a constant array of `default`. +fn expand_flag( + args: &[ColumnarValue], + index: usize, + present: bool, + default: bool, + num_iterations: usize, +) -> Result { + if present { + let array = args[index] + .clone() + .cast_to(&DataType::Boolean, None)? + .into_array(num_iterations)?; + Ok(as_boolean_array(&array)?.clone()) + } else { + let array = ScalarValue::Boolean(Some(default)).to_array_of_size(num_iterations)?; + Ok(as_boolean_array(&array)?.clone()) + } +} + +/// Cast a column to `Int64` and materialize it so its values can be iterated in +/// lockstep with the raster/roi rows. +fn expand_int64_arg(arg: &ColumnarValue, num_iterations: usize) -> Result { + let array = arg + .clone() + .cast_to(&DataType::Int64, None)? + .into_array(num_iterations)?; + Ok(as_int64_array(&array)?.clone()) +} + +/// Cast a column to `Utf8` and materialize it to a `StringArray` so its values +/// can be iterated in lockstep with the raster/roi rows. +fn expand_string_arg(arg: &ColumnarValue, num_iterations: usize) -> Result { + let array = arg + .clone() + .cast_to(&DataType::Utf8, None)? + .into_array(num_iterations)?; + Ok(as_string_array(&array)?.clone()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn stat_type_from_str_matches_spark_aliases() { + assert_eq!(StatType::from_str("count"), Some(StatType::Count)); + assert_eq!(StatType::from_str("COUNT"), Some(StatType::Count)); + assert_eq!(StatType::from_str("mean"), Some(StatType::Mean)); + assert_eq!(StatType::from_str("avg"), Some(StatType::Mean)); + assert_eq!(StatType::from_str("average"), Some(StatType::Mean)); + assert_eq!(StatType::from_str("stddev"), Some(StatType::StdDev)); + assert_eq!(StatType::from_str("sd"), Some(StatType::StdDev)); + assert_eq!(StatType::from_str("variance"), Some(StatType::Variance)); + assert_eq!(StatType::from_str("min"), Some(StatType::Min)); + assert_eq!(StatType::from_str("max"), Some(StatType::Max)); + assert_eq!(StatType::from_str("nonsense"), None); + } + + #[test] + fn resolve_band_defaults_and_bounds() { + // Single-band raster may omit the band. + assert_eq!(resolve_band(None, 1).unwrap(), 1); + // Multiband raster must name the band. + let err = resolve_band(None, 3).unwrap_err().to_string(); + assert!(err.contains("has 3 bands"), "{err}"); + // Explicit band is range-checked (1-based). + assert_eq!(resolve_band(Some(2), 3).unwrap(), 2); + assert!(resolve_band(Some(0), 3) + .unwrap_err() + .to_string() + .contains(">= 1")); + assert!(resolve_band(Some(4), 3) + .unwrap_err() + .to_string() + .contains("out of range")); + } + + #[test] + fn statistics_of_one_to_five() { + // count, sum, mean, min, max, median are exact; variance/stddev are the + // sample (n-1) values: ((1-3)^2+..+(5-3)^2)/4 = 10/4 = 2.5. + let mut values = vec![1.0, 2.0, 3.0, 4.0, 5.0]; + let s = compute_statistics(&mut values); + assert_eq!(s.count, 5); + assert_eq!(s.sum, Some(15.0)); + assert_eq!(s.mean, Some(3.0)); + assert_eq!(s.min, Some(1.0)); + assert_eq!(s.max, Some(5.0)); + assert_eq!(s.median, Some(3.0)); + assert_eq!(s.variance, Some(2.5)); + assert_eq!(s.stddev, Some(2.5_f64.sqrt())); + } + + #[test] + fn single_statistic_matches_full_computation() { + // RS_ZonalStats' single-stat path must return exactly what the full + // RS_ZonalStatsAll path computes, only without the others. A tie (two + // 4s) and an even count exercise mode and median. + let sample = [4.0, 1.0, 4.0, 2.0, 5.0, 3.0]; + let mut all = sample; + let full = compute_statistics(&mut all); + let cases = [ + (StatType::Count, Some(full.count as f64)), + (StatType::Sum, full.sum), + (StatType::Mean, full.mean), + (StatType::Median, full.median), + (StatType::Mode, full.mode), + (StatType::StdDev, full.stddev), + (StatType::Variance, full.variance), + (StatType::Min, full.min), + (StatType::Max, full.max), + ]; + for (stat, expected) in cases { + let mut work = sample; + assert_eq!( + compute_single_statistic(&mut work, stat), + expected, + "single statistic {stat:?} diverged from the full computation" + ); + } + } + + #[test] + fn statistics_empty_is_zero_count_and_nulls() { + let mut values: Vec = vec![]; + let s = compute_statistics(&mut values); + assert_eq!(s.count, 0); + assert_eq!(s.sum, None); + assert_eq!(s.mean, None); + assert_eq!(s.median, None); + assert_eq!(s.min, None); + assert_eq!(s.max, None); + assert_eq!(s.variance, None); + // The single-stat path returns 0 for count, NULL for the rest. + assert_eq!( + compute_single_statistic(&mut values, StatType::Count), + Some(0.0) + ); + assert_eq!(compute_single_statistic(&mut values, StatType::Sum), None); + assert_eq!(compute_single_statistic(&mut values, StatType::Mean), None); + } + + #[test] + fn statistics_single_value_has_zero_variance() { + let mut values = vec![42.0]; + let s = compute_statistics(&mut values); + assert_eq!(s.count, 1); + assert_eq!(s.mean, Some(42.0)); + assert_eq!(s.median, Some(42.0)); + assert_eq!(s.variance, Some(0.0)); + assert_eq!(s.stddev, Some(0.0)); + assert_eq!(s.mode, Some(42.0)); + } + + #[test] + fn median_even_count_averages_the_middle_pair() { + let mut values = vec![4.0, 1.0, 3.0, 2.0]; + let s = compute_statistics(&mut values); + assert_eq!(s.median, Some(2.5)); + } + + #[test] + fn mode_breaks_ties_toward_the_larger_value() { + // 1 and 3 each appear twice; the tie resolves to the larger, 3. + let mut values = vec![1.0, 1.0, 3.0, 3.0, 2.0]; + assert_eq!(compute_statistics(&mut values).mode, Some(3.0)); + // A clear winner is returned as-is. + let mut values = vec![7.0, 7.0, 7.0, 1.0, 2.0]; + assert_eq!(compute_statistics(&mut values).mode, Some(7.0)); + } + + #[test] + fn statistics_with_nan_are_all_nan_except_count() { + // A NaN pixel poisons every statistic (numpy semantics); count still + // reflects the number of selected values. Guards against min/max + // silently skipping NaN while sum/mean propagate it. + let mut values = vec![1.0, f64::NAN, 3.0]; + let s = compute_statistics(&mut values); + assert_eq!(s.count, 3); + assert!(s.sum.unwrap().is_nan()); + assert!(s.mean.unwrap().is_nan()); + assert!(s.median.unwrap().is_nan()); + assert!(s.mode.unwrap().is_nan()); + assert!(s.min.unwrap().is_nan()); + assert!(s.max.unwrap().is_nan()); + assert!(s.variance.unwrap().is_nan()); + assert!(s.stddev.unwrap().is_nan()); + } +} + +/// UDF-level tests: exercise the kernels end to end and pin the numbers against +/// values computed by hand (which agree with numpy — see the Python parity +/// tests for the rasterio/numpy cross-check). +#[cfg(test)] +mod udf_tests { + use super::*; + + use std::sync::Arc; + + use arrow_array::cast::AsArray; + use arrow_array::types::{Float64Type, Int64Type}; + use arrow_array::{Array, StructArray}; + use datafusion_expr::ScalarUDF; + use sedona_proj::transform::{with_global_proj_engine, LazyProjEngine}; + use sedona_raster_functions::crs_utils::crs_transform_wkb; + use sedona_schema::crs::deserialize_crs; + use sedona_schema::datatypes::{Edges, RASTER}; + use sedona_testing::create::make_wkb; + use sedona_testing::raster_spec::RasterSpec; + use sedona_testing::testers::ScalarUdfTester; + + // Struct field positions (Sedona Spark order). + const COUNT: usize = 0; + const SUM: usize = 1; + const MEAN: usize = 2; + const MEDIAN: usize = 3; + const MODE: usize = 4; + const STDDEV: usize = 5; + const VARIANCE: usize = 6; + const MIN: usize = 7; + const MAX: usize = 8; + + /// A 4×2 UInt8 raster with pixel values 1..=8 (row-major), world extent + /// x ∈ [0, 4], y ∈ [0, 2] with 1×1 north-up pixels. Pixel centres: + /// row y=1.5 → 1,2,3,4 at x=0.5,1.5,2.5,3.5; row y=0.5 → 5,6,7,8. + fn small_raster() -> RasterSpec { + RasterSpec::d2(4, 2) + .band_values(&[1u8, 2, 3, 4, 5, 6, 7, 8]) + .crs(None) + .transform([0.0, 1.0, 0.0, 2.0, 0.0, -1.0]) + } + + /// The left half of `small_raster` (x ∈ [0, 2]) selects the four pixels + /// {1, 2, 5, 6}. + const LEFT_HALF: &str = "POLYGON((0 0, 2 0, 2 2, 0 2, 0 0))"; + + // ScalarValue constructors for the positional trailing arguments, so a call + // reads close to its Sedona Spark SQL form: `[band(1), stat("sum"), + // flag(true), flag(false)]` is `(raster, roi, 1, 'sum', true, false)`. + fn band(b: i64) -> ScalarValue { + ScalarValue::Int64(Some(b)) + } + fn stat(s: &str) -> ScalarValue { + ScalarValue::Utf8(Some(s.to_string())) + } + fn flag(b: bool) -> ScalarValue { + ScalarValue::Boolean(Some(b)) + } + + /// Invoke a zonal-stats UDF on a scalar raster + roi with the given + /// positional trailing arguments. Routing through the UDF (rather than a + /// hand-picked kernel) exercises overload selection by argument count and + /// type; the raw `ScalarValue` is returned so both value and error paths are + /// testable. + fn invoke_udf( + udf: SedonaScalarUDF, + spec: &RasterSpec, + geom: ScalarValue, + trailing: Vec, + ) -> Result { + let mut arg_types = vec![RASTER, SedonaType::Wkb(Edges::Planar, None)]; + let mut args = vec![ + ColumnarValue::Scalar(spec.scalar()), + ColumnarValue::Scalar(geom), + ]; + for value in trailing { + arg_types.push(SedonaType::Arrow(value.data_type())); + args.push(ColumnarValue::Scalar(value)); + } + match ScalarUdfTester::new(udf.into(), arg_types).invoke(args)? { + ColumnarValue::Scalar(s) => Ok(s), + other => panic!("expected a scalar result, got {other:?}"), + } + } + + /// RS_ZonalStats over a scalar raster + roi with the given trailing args. + fn call_stats(spec: &RasterSpec, wkt: &str, trailing: Vec) -> Result { + let geom = ScalarValue::Binary(Some(make_wkb(wkt))); + invoke_udf(rs_zonal_stats_udf(), spec, geom, trailing) + } + + /// RS_ZonalStatsAll over a scalar raster + roi with the given trailing args. + fn call_all(spec: &RasterSpec, wkt: &str, trailing: Vec) -> Result { + let geom = ScalarValue::Binary(Some(make_wkb(wkt))); + invoke_udf(rs_zonal_stats_all_udf(), spec, geom, trailing) + } + + fn cv_f64(s: ScalarValue) -> Option { + match s { + ScalarValue::Float64(v) => v, + other => panic!("expected a Float64 scalar, got {other:?}"), + } + } + + fn cv_struct(s: ScalarValue) -> Arc { + match s { + ScalarValue::Struct(s) => s, + other => panic!("expected a struct scalar, got {other:?}"), + } + } + + fn f64_field(s: &StructArray, col: usize) -> Option { + let c = s.column(col); + (!c.is_null(0)).then(|| c.as_primitive::().value(0)) + } + + fn i64_field(s: &StructArray, col: usize) -> Option { + let c = s.column(col); + (!c.is_null(0)).then(|| c.as_primitive::().value(0)) + } + + #[test] + fn single_stats_match_hand_computed_values() { + let spec = small_raster(); + // Selected pixels {1, 2, 5, 6}: exact for the integer-selection stats. + // The 3-arg overload leaves the band implicit (unambiguous single band). + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("count")]).unwrap()), + Some(4.0) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("sum")]).unwrap()), + Some(14.0) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("mean")]).unwrap()), + Some(3.5) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("min")]).unwrap()), + Some(1.0) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("max")]).unwrap()), + Some(6.0) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("median")]).unwrap()), + Some(3.5) + ); + // All four values are unique, so every one is a mode; the tie resolves + // to the largest (6), matching Sedona Spark. + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("mode")]).unwrap()), + Some(6.0) + ); + + // Sample (n-1) variance / stddev: float accumulation, so approximate. + let var = cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("variance")]).unwrap()).unwrap(); + assert!((var - 17.0 / 3.0).abs() < 1e-9, "variance was {var}"); + let sd = cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("stddev")]).unwrap()).unwrap(); + assert!( + (sd - (17.0_f64 / 3.0).sqrt()).abs() < 1e-9, + "stddev was {sd}" + ); + } + + #[test] + fn all_returns_full_struct() { + // The 2-arg overload leaves the band implicit. + let s = cv_struct(call_all(&small_raster(), LEFT_HALF, vec![]).unwrap()); + assert!(!s.is_null(0), "the struct itself is valid"); + assert_eq!(i64_field(&s, COUNT), Some(4)); + assert_eq!(f64_field(&s, SUM), Some(14.0)); + assert_eq!(f64_field(&s, MEAN), Some(3.5)); + assert_eq!(f64_field(&s, MEDIAN), Some(3.5)); + assert_eq!(f64_field(&s, MODE), Some(6.0)); + assert_eq!(f64_field(&s, MIN), Some(1.0)); + assert_eq!(f64_field(&s, MAX), Some(6.0)); + assert!((f64_field(&s, VARIANCE).unwrap() - 17.0 / 3.0).abs() < 1e-9); + assert!((f64_field(&s, STDDEV).unwrap() - (17.0_f64 / 3.0).sqrt()).abs() < 1e-9); + } + + #[test] + fn overloads_dispatch_by_arg_count() { + // The 3-arg (raster, roi, stat) and 4-arg (raster, roi, band, stat) + // overloads resolve by argument count and by the type at position 2 (a + // stat string vs. a band integer). On a single-band raster both compute + // the same mean over {1, 2, 5, 6}. + let spec = small_raster(); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, vec![stat("mean")]).unwrap()), + Some(3.5) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF, vec![band(1), stat("mean")]).unwrap()), + Some(3.5) + ); + // RS_ZonalStatsAll: the 2-arg and 3-arg (band) overloads agree too. + assert_eq!( + i64_field( + &cv_struct(call_all(&spec, LEFT_HALF, vec![]).unwrap()), + COUNT + ), + Some(4) + ); + assert_eq!( + i64_field( + &cv_struct(call_all(&spec, LEFT_HALF, vec![band(1)]).unwrap()), + COUNT + ), + Some(4) + ); + } + + #[test] + fn roi_that_selects_no_pixel_centre_is_count_zero_not_null() { + // A tiny roi inside the top-left pixel (centre 0.5, 1.5) but not + // covering that centre: with all_touched off, no pixel is selected. The + // roi still overlaps the raster extent, so count is 0 (not NULL). + let tiny = "POLYGON((0.1 1.6, 0.4 1.6, 0.4 1.9, 0.1 1.9, 0.1 1.6))"; + assert_eq!( + cv_f64(call_stats(&small_raster(), tiny, vec![stat("count")]).unwrap()), + Some(0.0) + ); + assert_eq!( + cv_f64(call_stats(&small_raster(), tiny, vec![stat("sum")]).unwrap()), + None + ); + assert_eq!( + cv_f64(call_stats(&small_raster(), tiny, vec![stat("mean")]).unwrap()), + None + ); + + let s = cv_struct(call_all(&small_raster(), tiny, vec![]).unwrap()); + assert!( + !s.is_null(0), + "an intersecting-but-empty roi is a valid row" + ); + assert_eq!(i64_field(&s, COUNT), Some(0)); + assert_eq!(f64_field(&s, SUM), None); + assert_eq!(f64_field(&s, MEAN), None); + } + + #[test] + fn all_touched_selects_the_touched_pixel() { + // The same tiny roi, with all_touched, burns the pixel it lies inside + // (value 1) even though it misses the centre. all_touched first appears + // in the 5-arg overload (raster, roi, band, stat, all_touched), so the + // band must be named to reach it. + let tiny = "POLYGON((0.1 1.6, 0.4 1.6, 0.4 1.9, 0.1 1.9, 0.1 1.6))"; + assert_eq!( + cv_f64( + call_stats( + &small_raster(), + tiny, + vec![band(1), stat("count"), flag(true)] + ) + .unwrap() + ), + Some(1.0) + ); + assert_eq!( + cv_f64( + call_stats( + &small_raster(), + tiny, + vec![band(1), stat("sum"), flag(true)] + ) + .unwrap() + ), + Some(1.0) + ); + } + + #[test] + fn no_intersection_is_null_when_lenient_and_errors_when_strict() { + let far = "POLYGON((100 100, 101 100, 101 101, 100 101, 100 100))"; + // Lenient (default): the whole value is NULL, including count. + assert_eq!( + cv_f64(call_stats(&small_raster(), far, vec![stat("count")]).unwrap()), + None + ); + assert!(cv_struct(call_all(&small_raster(), far, vec![]).unwrap()).is_null(0)); + + // Strict (lenient => false): both functions error. RS_ZonalStats reaches + // lenient only in its 7-arg overload, whose trailing flags are + // (all_touched, exclude_no_data, lenient). + let err = call_stats( + &small_raster(), + far, + vec![band(1), stat("count"), flag(false), flag(true), flag(false)], + ) + .unwrap_err() + .to_string(); + assert!( + err.contains("does not intersect"), + "unexpected error: {err}" + ); + // RS_ZonalStatsAll's 6-arg overload trails (all_touched, exclude_no_data, + // lenient) after the band. + let err = call_all( + &small_raster(), + far, + vec![band(1), flag(false), flag(true), flag(false)], + ) + .unwrap_err() + .to_string(); + assert!( + err.contains("does not intersect"), + "unexpected error: {err}" + ); + } + + #[test] + fn bbox_overlapping_but_geometry_disjoint_roi_is_no_intersection() { + // small_raster covers x ∈ [0, 4], y ∈ [0, 2]. This triangle lives on the + // far side of the line x + y = 7, so its geometry is disjoint from the + // raster (every raster point has x + y ≤ 6), yet its bounding box + // [0, 7] × [0, 7] contains the whole raster. A bounding-box gate would + // burn zero pixels and report count 0; the true-geometry gate (matching + // Sedona Spark's rsIntersects) treats it as a no-intersection case. + let disjoint = "POLYGON((7 0, 0 7, 7 7, 7 0))"; + + // Lenient (default): NULL, not count 0. + assert_eq!( + cv_f64(call_stats(&small_raster(), disjoint, vec![stat("count")]).unwrap()), + None + ); + assert!(cv_struct(call_all(&small_raster(), disjoint, vec![]).unwrap()).is_null(0)); + + // Strict (lenient => false): both functions error. + let err = call_stats( + &small_raster(), + disjoint, + vec![band(1), stat("count"), flag(false), flag(true), flag(false)], + ) + .unwrap_err() + .to_string(); + assert!( + err.contains("does not intersect"), + "unexpected error: {err}" + ); + let err = call_all( + &small_raster(), + disjoint, + vec![band(1), flag(false), flag(true), flag(false)], + ) + .unwrap_err() + .to_string(); + assert!( + err.contains("does not intersect"), + "unexpected error: {err}" + ); + } + + #[test] + fn nodata_pixels_are_excluded_by_default_and_kept_when_asked() { + // A 2×2 UInt8 raster [10, 255, 20, 30] with nodata 255, world extent + // x ∈ [0, 2], y ∈ [0, 2]; the roi covers all four pixels. + let spec = RasterSpec::d2(2, 2) + .band_values(&[10u8, 255, 20, 30]) + .nodata(255u8) + .crs(None) + .transform([0.0, 1.0, 0.0, 2.0, 0.0, -1.0]); + // Default excludes the nodata pixel: {10, 20, 30}. + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF_FULL, vec![stat("count")]).unwrap()), + Some(3.0) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF_FULL, vec![stat("sum")]).unwrap()), + Some(60.0) + ); + // exclude_no_data => false keeps it: {10, 255, 20, 30}. It first appears + // in the 6-arg overload, whose trailing flags are (all_touched, + // exclude_no_data). + assert_eq!( + cv_f64( + call_stats( + &spec, + LEFT_HALF_FULL, + vec![band(1), stat("count"), flag(false), flag(false)] + ) + .unwrap() + ), + Some(4.0) + ); + assert_eq!( + cv_f64( + call_stats( + &spec, + LEFT_HALF_FULL, + vec![band(1), stat("sum"), flag(false), flag(false)] + ) + .unwrap() + ), + Some(315.0) + ); + } + + /// A roi covering the whole 2×2 nodata raster above. + const LEFT_HALF_FULL: &str = "POLYGON((0 0, 2 0, 2 2, 0 2, 0 0))"; + + #[test] + fn multiband_raster_requires_the_band_argument() { + let spec = RasterSpec::d2(2, 2) + .band_values(&[1u8, 2, 3, 4]) + .band_values(&[10u8, 20, 30, 40]) + .crs(None) + .transform([0.0, 1.0, 0.0, 2.0, 0.0, -1.0]); + // Omitting the band on a multiband raster errors rather than defaulting + // to band 1 (this deliberately diverges from Sedona Spark). The 3-arg + // RS_ZonalStats overload and the 2-arg RS_ZonalStatsAll overload both + // leave the band implicit. + let err = call_stats(&spec, LEFT_HALF_FULL, vec![stat("sum")]) + .unwrap_err() + .to_string(); + assert!(err.contains("2 bands"), "unexpected error: {err}"); + let err = call_all(&spec, LEFT_HALF_FULL, vec![]) + .unwrap_err() + .to_string(); + assert!(err.contains("2 bands"), "unexpected error: {err}"); + // Naming the band selects it (band 1 sums to 10, band 2 to 100). + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF_FULL, vec![band(1), stat("sum")]).unwrap()), + Some(10.0) + ); + assert_eq!( + cv_f64(call_stats(&spec, LEFT_HALF_FULL, vec![band(2), stat("sum")]).unwrap()), + Some(100.0) + ); + // An out-of-range band errors. + let err = call_stats(&spec, LEFT_HALF_FULL, vec![band(3), stat("sum")]) + .unwrap_err() + .to_string(); + assert!(err.contains("out of range"), "unexpected error: {err}"); + } + + #[test] + fn unknown_statistic_errors() { + let err = call_stats(&small_raster(), LEFT_HALF, vec![stat("bogus")]) + .unwrap_err() + .to_string(); + assert!(err.contains("unknown statistic"), "unexpected error: {err}"); + } + + #[test] + fn null_raster_or_roi_yields_null() { + // A NULL roi geometry propagates to a NULL result (3-arg overload). + let null_roi = invoke_udf( + rs_zonal_stats_udf(), + &small_raster(), + ScalarValue::Binary(None), + vec![stat("count")], + ) + .unwrap(); + assert_eq!(cv_f64(null_roi), None); + + // A NULL statistic name also yields NULL. + let null_stat = + call_stats(&small_raster(), LEFT_HALF, vec![ScalarValue::Utf8(None)]).unwrap(); + assert_eq!(cv_f64(null_stat), None); + } + + #[test] + fn reprojects_the_roi_into_the_raster_crs() { + // The raster is EPSG:4326; the roi is supplied in EPSG:3857 (the + // reprojected LEFT_HALF polygon). Reprojecting it back to the raster CRS + // must recover the same four-pixel selection. + let spec = small_raster().crs(Some("EPSG:4326")); + let crs_4326 = deserialize_crs("EPSG:4326").unwrap().unwrap(); + let crs_3857 = deserialize_crs("EPSG:3857").unwrap().unwrap(); + let wkb_4326 = make_wkb(LEFT_HALF); + let wkb_3857 = with_global_proj_engine(|engine| { + crs_transform_wkb(&wkb_4326, crs_4326.as_ref(), crs_3857.as_ref(), engine) + }) + .unwrap(); + + let udf: ScalarUDF = rs_zonal_stats_udf().into(); + let arg_types = vec![ + RASTER, + SedonaType::Wkb(Edges::Planar, Some(crs_3857)), + SedonaType::Arrow(DataType::Utf8), + ]; + let tester = ScalarUdfTester::new(udf, arg_types).with_crs_engine(Arc::new(LazyProjEngine)); + let result = tester + .invoke_scalar_scalar_scalar(&spec, ScalarValue::Binary(Some(wkb_3857)), "count") + .unwrap(); + assert_eq!(result, ScalarValue::Float64(Some(4.0))); + } +}