fix(rust/sedona-query-planner): Ensure user provided RS_EnsureLoaded call preserves metadata (#969)

Author: paleolimbot

python/sedonadb/python/sedonadb/raster.py

@@ -15,6 +15,9 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import struct
+import math
+
 from typing import List, Optional, TYPE_CHECKING, Tuple, Any, Iterable
 import geoarrow.types as gat
 import pyarrow as pa
@@ -229,10 +232,36 @@ def source_data(self) -> memoryview:
         view_scalar = self._array.field("data")[0]
         return memoryview(view_scalar.as_buffer())
 
+    @property
+    def source_data_size(self) -> int:
+        """The number of bytes consumed by soure_data if it were loaded"""
+        buffer_type_id = self._py_field("data_type")
+        buffer_type_char = BAND_DATA_TYPE_STRUCT_CHARS[buffer_type_id]
+        element_size = struct.calcsize(buffer_type_char)
+        return math.prod(self.source_shape) * element_size
+
+    @property
+    def data_size(self) -> int:
+        """The number of bytes consumed by data if it were loaded"""
+        buffer_type_id = self._py_field("data_type")
+        buffer_type_char = BAND_DATA_TYPE_STRUCT_CHARS[buffer_type_id]
+        element_size = struct.calcsize(buffer_type_char)
+        return math.prod(self.shape) * element_size
+
     @property
     def data(self) -> memoryview:
         """The band data as a typed, shaped memoryview."""
-        if self.outdb_uri is not None:
+        buffer_type_id = self._py_field("data_type")
+        buffer_type_char = BAND_DATA_TYPE_STRUCT_CHARS[buffer_type_id]
+
+        # This is not quite right, but shapes that contain zeroes are not well
+        # supported by the memoryview yet. Callers should check data_size for
+        # empty handling with non-numpy views.
+        if self.data_size == 0:
+            return memoryview(b"")
+
+        source_data = self.source_data
+        if self.outdb_uri is not None and len(source_data) == 0:
             raise ValueError("Can't extract buffer from a reference to external data.")
 
         # When views are supported, we would need to calculate the striding
@@ -241,14 +270,15 @@ def data(self) -> memoryview:
         if views:
             raise NotImplementedError("Lazy views are not yet supported")
 
-        buffer_type_id = self._py_field("data_type")
-        buffer_type_char = BAND_DATA_TYPE_STRUCT_CHARS[buffer_type_id]
         return self.source_data.cast(buffer_type_char, self.shape)
 
     def to_numpy(self) -> "np.ndarray":
         """Convert this band's data to a numpy array."""
         import numpy as np
 
+        if self.data_size == 0:
+            return np.empty(self.shape, dtype=self.data_type)
+
         return np.array(self.data)
 
     def __repr__(self) -> str:

python/sedonadb/tests/functions/conftest.py

@@ -15,43 +15,44 @@
 # specific language governing permissions and limitations
 # under the License.
 
-"""Table-driven tests for RS_ accessor functions over an in-DB example raster.
-
-The `rasters` view (see `raster_con` in conftest.py) holds a single
-`RS_Example()` raster: 64x32, three UInt8 bands, nodata 127, with a fixed
-geotransform (origin (43.08, 79.07), scale 2, skew 1). These exercise the RS_
-accessor kernels against the raster Arrow type with no zarr dependency; the
-zarr reader path (OutDb chunk anchors, fill_value->nodata, RS_EnsureLoaded) is
-tested in the sedonadb-zarr package.
-
-There is no PostGIS twin for these (unlike the geometry function tests), so
-plain `con.sql(...)` assertions are the right altitude.
-"""
-
 import pytest
 
-
-def query_value(con, expr):
-    """Evaluate `expr` over the single example raster row and return the value."""
-    table = con.sql(f"SELECT {expr} AS v FROM rasters").to_arrow_table()
-    return table["v"][0].as_py()
+from sedonadb.testing import SedonaDB
 
 
 @pytest.mark.parametrize(
     ("expr", "expected"),
     [
-        ("RS_NumBands(raster)", 3),
-        ("RS_Width(raster)", 64),
-        ("RS_Height(raster)", 32),
-        ("RS_BandPixelType(raster, 1)", "UNSIGNED_8BITS"),
-        ("RS_BandNoDataValue(raster, 1)", 127.0),
-        ("RS_ScaleX(raster)", 2.0),
-        ("RS_ScaleY(raster)", 2.0),
-        ("RS_SkewX(raster)", 1.0),
-        ("RS_SkewY(raster)", 1.0),
-        ("RS_UpperLeftX(raster)", 43.08),
-        ("RS_UpperLeftY(raster)", 79.07),
+        ("RS_NumBands(RS_Example())", 3),
+        ("RS_Width(RS_Example())", 64),
+        ("RS_Height(RS_Example())", 32),
+        ("RS_BandPixelType(RS_Example(), 1)", "UNSIGNED_8BITS"),
+        ("RS_BandNoDataValue(RS_Example(), 1)", 127.0),
+        ("RS_ScaleX(RS_Example())", 2.0),
+        ("RS_ScaleY(RS_Example())", 2.0),
+        ("RS_SkewX(RS_Example())", 1.0),
+        ("RS_SkewY(RS_Example())", 1.0),
+        ("RS_UpperLeftX(RS_Example())", 43.08),
+        ("RS_UpperLeftY(RS_Example())", 79.07),
     ],
 )
-def test_rs_function(raster_con, expr, expected):
-    assert query_value(raster_con, expr) == expected
+def test_rs_function(expr, expected):
+    eng = SedonaDB()
+    eng.assert_query_result(f"SELECT {expr}", expected)
+
+
+def test_rs_ensureloaded(con, sedona_testing):
+    path = sedona_testing / "data/raster/sentinel2.tif"
+    t = con.sql("SELECT RS_FromPath($1) AS raster", params=(str(path),))
+    tab = t.select(raster=t.raster.funcs.rs_ensureloaded()).to_arrow_table()
+    r = tab["raster"][0].as_py()
+    assert r.height == 512
+    assert r.width == 512
+
+    assert len(r.bands) == 1
+    b = r.bands[0]
+    assert b.shape == (512, 512)
+    arr = b.to_numpy()
+    assert arr.shape == (512, 512)
+    assert arr.dtype == "uint16"
+    assert arr[0, 0] == 2324

python/sedonadb/tests/functions/test_raster_functions.py

@@ -49,6 +49,8 @@ def test_raster_accessors(con):
     assert b.source_shape == (32, 64)
     assert b.outdb_uri is None
     assert b.data_type == "uint8"
+    assert b.source_data_size == 32 * 64 * 1  # uint8 = 1 byte
+    assert b.data_size == 32 * 64 * 1
     assert repr(b) == "<Band uint8 32x64>"
 
     arr = b.to_numpy()
@@ -87,6 +89,8 @@ def test_raster_lazy():
     assert b.source_shape == (512, 1024)
     assert b.data_type == "float32"
     assert b.outdb_uri == "s3://bucket/path/to/data.zarr"
+    assert b.source_data_size == 512 * 1024 * 4  # float32 = 4 bytes
+    assert b.data_size == 512 * 1024 * 4
 
     # Lazy raster should have empty data buffer
     assert len(b.source_data) == 0
@@ -116,3 +120,22 @@ def test_raster_lazy_invalid_shape():
 
     with pytest.raises(ValueError, match="exactly two dimensions"):
         Raster.lazy(uri="s3://bucket/data.zarr", shape=(10, 20, 30), dtype="UInt8")
+
+
+def test_raster_lazy_zero_size():
+    """Test that a raster with zero-size shape returns an empty memoryview."""
+    r = Raster.lazy(
+        uri="s3://bucket/empty.zarr",
+        shape=(0, 64),
+        dtype="float32",
+    )
+
+    b = r.bands[0]
+    assert b.source_shape == (0, 64)
+    assert b.data_size == 0
+    assert b.source_data_size == 0
+    assert b.data == memoryview(b"")
+
+    arr = b.to_numpy()
+    assert arr.shape == (0, 64)
+    assert arr.dtype == "float32"