python/tests/raster/test_serde.py - sedona - Git at Google

 # 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.

 import numpy as np
 import pyspark
 import pytest
 import rasterio
 from pyspark.sql.functions import expr
 from tests import world_map_raster_input_location
 from tests.test_base import TestBase

 from sedona.spark.sql.types import RasterType


 class TestRasterSerde(TestBase):
     def test_empty_raster(self):
         df = TestRasterSerde.spark.sql(
             "SELECT RS_MakeEmptyRaster(2, 100, 200, 1000, 2000, 1) as raster"
         )
         raster = df.first()[0]
         assert (
             raster.width == 100 and raster.height == 200 and len(raster.bands_meta) == 2
         )
         assert raster.affine_trans.ip_x == 1000 and raster.affine_trans.ip_y == 2000
         assert raster.affine_trans.scale_x == 1 and raster.affine_trans.scale_y == -1

     def test_banded_sample_model(self):
         df = TestRasterSerde.spark.sql(
             "SELECT RS_MakeRasterForTesting(3, 'I', 'BandedSampleModel', 10, 8, 100, 100, 10, -10, 0, 0, 3857) as raster"
         )
         raster = df.first()[0]
         assert raster.width == 10 and raster.height == 8 and len(raster.bands_meta) == 3
         self.validate_test_raster(raster)

     def test_pixel_interleaved_sample_model(self):
         df = TestRasterSerde.spark.sql(
             "SELECT RS_MakeRasterForTesting(3, 'I', 'PixelInterleavedSampleModel', 10, 10, 100, 100, 10, -10, 0, 0, 3857) as raster"
         )
         raster = df.first()[0]
         assert (
             raster.width == 10 and raster.height == 10 and len(raster.bands_meta) == 3
         )
         self.validate_test_raster(raster)
         df = TestRasterSerde.spark.sql(
             "SELECT RS_MakeRasterForTesting(4, 'I', 'PixelInterleavedSampleModelComplex', 8, 10, 100, 100, 10, -10, 0, 0, 3857) as raster"
         )
         raster = df.first()[0]
         assert raster.width == 8 and raster.height == 10 and len(raster.bands_meta) == 4
         self.validate_test_raster(raster)

     def test_component_sample_model(self):
         for pixel_type in ["B", "S", "US", "I", "F", "D"]:
             df = TestRasterSerde.spark.sql(
                 "SELECT RS_MakeRasterForTesting(4, '{}', 'ComponentSampleModel', 10, 10, 100, 100, 10, -10, 0, 0, 3857) as raster".format(
                     pixel_type
                 )
             )
             raster = df.first()[0]
             assert (
                 raster.width == 10
                 and raster.height == 10
                 and len(raster.bands_meta) == 4
             )
             self.validate_test_raster(raster)

     def test_multi_pixel_packed_sample_model(self):
         df = TestRasterSerde.spark.sql(
             "SELECT RS_MakeRasterForTesting(1, 'B', 'MultiPixelPackedSampleModel', 10, 10, 100, 100, 10, -10, 0, 0, 3857) as raster"
         )
         raster = df.first()[0]
         assert (
             raster.width == 10 and raster.height == 10 and len(raster.bands_meta) == 1
         )
         self.validate_test_raster(raster, packed=True)

     def test_single_pixel_packed_sample_model(self):
         df = TestRasterSerde.spark.sql(
             "SELECT RS_MakeRasterForTesting(4, 'I', 'SinglePixelPackedSampleModel', 10, 10, 100, 100, 10, -10, 0, 0, 3857) as raster"
         )
         raster = df.first()[0]
         assert (
             raster.width == 10 and raster.height == 10 and len(raster.bands_meta) == 4
         )
         self.validate_test_raster(raster, packed=True)

     def test_raster_read_from_geotiff(self):
         raster_path = world_map_raster_input_location
         r_orig = rasterio.open(raster_path)
         band = r_orig.read(1)
         band_masked = np.where(band == 0, np.nan, band)
         df = (
             TestRasterSerde.spark.read.format("binaryFile")
             .load(raster_path)
             .selectExpr("RS_FromGeoTiff(content) as raster")
         )
         raster = df.first()[0]
         assert raster.width == r_orig.width
         assert raster.height == r_orig.height
         assert raster.bands_meta[0].nodata == 0

         # test as_rasterio
         assert (band == raster.as_numpy()[0, :, :]).all()
         ds = raster.as_rasterio()
         assert ds.crs is not None
         band_actual = ds.read(1)
         assert (band == band_actual).all()

         # test as_numpy
         arr = raster.as_numpy()
         assert (arr[0, :, :] == band).all()

         # test as_numpy_masked
         arr = raster.as_numpy_masked()[0, :, :]
         assert np.array_equal(arr, band_masked) or np.array_equal(
             np.isnan(arr), np.isnan(band_masked)
         )

         raster.close()
         r_orig.close()

     def test_to_pandas(self):
         spark = TestRasterSerde.spark
         df = spark.sql(
             "SELECT RS_MakeRasterForTesting(3, 'I', 'BandedSampleModel', 10, 8, 100, 100, 10, -10, 0, 0, 3857) as raster"
         )
         pandas_df = df.toPandas()
         raster = pandas_df.iloc[0]["raster"]
         self.validate_test_raster(raster)

     def validate_test_raster(self, raster, packed=False):
         arr = raster.as_numpy()
         ds = raster.as_rasterio()
         bands, height, width = arr.shape
         assert bands > 0 and width > 0 and height > 0
         assert ds.crs is not None
         for b in range(bands):
             band = ds.read(b + 1)
             for y in range(height):
                 for x in range(width):
                     expected = b + y * width + x
                     if packed:
                         expected = expected % 16
                     assert arr[b, y, x] == expected
                     assert band[y, x] == expected

     @pytest.mark.skipif(
         pyspark.__version__ < "3.4", reason="requires Spark 3.4 or higher"
     )
     def test_serialize_round_trip_via_udf(self):
         """Deserialize → serialize round-trip through a raster-returning UDF.
         Verifies that the Python-serialized bytes can be deserialized back
         to produce correct pixel values.
         """
         from sedona.spark.raster import raster_serde

         spark = self.spark

         df = spark.sql(
             "SELECT RS_MakeRasterForTesting(3, 'I', 'BandedSampleModel', "
             "10, 8, 100, 100, 10, -10, 0, 0, 3857) as raster"
         )
         # Collect the deserialized raster from Spark
         original = df.first()["raster"]

         # Serialize on Python side and deserialize back
         serialized = raster_serde.serialize(original)
         round_tripped = raster_serde.deserialize(serialized)

         # Verify pixel values for all 3 bands
         arr = round_tripped.as_numpy()
         assert arr.shape == (3, 8, 10)
         for b in range(3):
             for y in range(8):
                 for x in range(10):
                     expected = b + y * 10 + x
                     assert arr[b, y, x] == expected

         round_tripped.close()
         original.close()

     @pytest.mark.skipif(
         pyspark.__version__ < "3.4", reason="requires Spark 3.4 or higher"
     )
     def test_serialize_preserves_metadata(self):
         """Verify CRS, affine, and dimensions survive serialize round-trip."""
         from sedona.spark.raster import raster_serde

         spark = self.spark

         df = spark.sql(
             "SELECT RS_MakeRasterForTesting(1, 'I', 'BandedSampleModel', "
             "4, 3, 100, 100, 10, -10, 0, 0, 3857) as raster"
         )
         original = df.first()["raster"]

         # Serialize and deserialize in Python
         serialized = raster_serde.serialize(original)
         round_tripped = raster_serde.deserialize(serialized)

         # Verify spatial metadata survived the round-trip
         assert round_tripped.width == 4
         assert round_tripped.height == 3
         assert round_tripped.affine_trans.scale_x == 10.0
         assert round_tripped.affine_trans.scale_y == -10.0
         assert round_tripped.affine_trans.ip_x == 100.0
         assert round_tripped.affine_trans.ip_y == 100.0
         assert "3857" in round_tripped.crs_wkt or "Mercator" in round_tripped.crs_wkt

         round_tripped.close()
         original.close()
	# 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.

	import numpy as np
	import pyspark
	import pytest
	import rasterio
	from pyspark.sql.functions import expr
	from tests import world_map_raster_input_location
	from tests.test_base import TestBase

	from sedona.spark.sql.types import RasterType


	class TestRasterSerde(TestBase):
	def test_empty_raster(self):
	df = TestRasterSerde.spark.sql(
	"SELECT RS_MakeEmptyRaster(2, 100, 200, 1000, 2000, 1) as raster"
	)
	raster = df.first()[0]
	assert (
	raster.width == 100 and raster.height == 200 and len(raster.bands_meta) == 2
	)
	assert raster.affine_trans.ip_x == 1000 and raster.affine_trans.ip_y == 2000
	assert raster.affine_trans.scale_x == 1 and raster.affine_trans.scale_y == -1

	def test_banded_sample_model(self):
	df = TestRasterSerde.spark.sql(
	"SELECT RS_MakeRasterForTesting(3, 'I', 'BandedSampleModel', 10, 8, 100, 100, 10, -10, 0, 0, 3857) as raster"
	)
	raster = df.first()[0]
	assert raster.width == 10 and raster.height == 8 and len(raster.bands_meta) == 3
	self.validate_test_raster(raster)

	def test_pixel_interleaved_sample_model(self):
	df = TestRasterSerde.spark.sql(
	"SELECT RS_MakeRasterForTesting(3, 'I', 'PixelInterleavedSampleModel', 10, 10, 100, 100, 10, -10, 0, 0, 3857) as raster"
	)
	raster = df.first()[0]
	assert (
	raster.width == 10 and raster.height == 10 and len(raster.bands_meta) == 3
	)
	self.validate_test_raster(raster)
	df = TestRasterSerde.spark.sql(
	"SELECT RS_MakeRasterForTesting(4, 'I', 'PixelInterleavedSampleModelComplex', 8, 10, 100, 100, 10, -10, 0, 0, 3857) as raster"
	)
	raster = df.first()[0]
	assert raster.width == 8 and raster.height == 10 and len(raster.bands_meta) == 4
	self.validate_test_raster(raster)

	def test_component_sample_model(self):
	for pixel_type in ["B", "S", "US", "I", "F", "D"]:
	df = TestRasterSerde.spark.sql(
	"SELECT RS_MakeRasterForTesting(4, '{}', 'ComponentSampleModel', 10, 10, 100, 100, 10, -10, 0, 0, 3857) as raster".format(
	pixel_type
	)
	)
	raster = df.first()[0]
	assert (
	raster.width == 10
	and raster.height == 10
	and len(raster.bands_meta) == 4
	)
	self.validate_test_raster(raster)

	def test_multi_pixel_packed_sample_model(self):
	df = TestRasterSerde.spark.sql(
	"SELECT RS_MakeRasterForTesting(1, 'B', 'MultiPixelPackedSampleModel', 10, 10, 100, 100, 10, -10, 0, 0, 3857) as raster"
	)
	raster = df.first()[0]
	assert (
	raster.width == 10 and raster.height == 10 and len(raster.bands_meta) == 1
	)
	self.validate_test_raster(raster, packed=True)

	def test_single_pixel_packed_sample_model(self):
	df = TestRasterSerde.spark.sql(
	"SELECT RS_MakeRasterForTesting(4, 'I', 'SinglePixelPackedSampleModel', 10, 10, 100, 100, 10, -10, 0, 0, 3857) as raster"
	)
	raster = df.first()[0]
	assert (
	raster.width == 10 and raster.height == 10 and len(raster.bands_meta) == 4
	)
	self.validate_test_raster(raster, packed=True)

	def test_raster_read_from_geotiff(self):
	raster_path = world_map_raster_input_location
	r_orig = rasterio.open(raster_path)
	band = r_orig.read(1)
	band_masked = np.where(band == 0, np.nan, band)
	df = (
	TestRasterSerde.spark.read.format("binaryFile")
	.load(raster_path)
	.selectExpr("RS_FromGeoTiff(content) as raster")
	)
	raster = df.first()[0]
	assert raster.width == r_orig.width
	assert raster.height == r_orig.height
	assert raster.bands_meta[0].nodata == 0

	# test as_rasterio
	assert (band == raster.as_numpy()[0, :, :]).all()
	ds = raster.as_rasterio()
	assert ds.crs is not None
	band_actual = ds.read(1)
	assert (band == band_actual).all()

	# test as_numpy
	arr = raster.as_numpy()
	assert (arr[0, :, :] == band).all()

	# test as_numpy_masked
	arr = raster.as_numpy_masked()[0, :, :]
	assert np.array_equal(arr, band_masked) or np.array_equal(
	np.isnan(arr), np.isnan(band_masked)
	)

	raster.close()
	r_orig.close()

	def test_to_pandas(self):
	spark = TestRasterSerde.spark
	df = spark.sql(
	"SELECT RS_MakeRasterForTesting(3, 'I', 'BandedSampleModel', 10, 8, 100, 100, 10, -10, 0, 0, 3857) as raster"
	)
	pandas_df = df.toPandas()
	raster = pandas_df.iloc[0]["raster"]
	self.validate_test_raster(raster)

	def validate_test_raster(self, raster, packed=False):
	arr = raster.as_numpy()
	ds = raster.as_rasterio()
	bands, height, width = arr.shape
	assert bands > 0 and width > 0 and height > 0
	assert ds.crs is not None
	for b in range(bands):
	band = ds.read(b + 1)
	for y in range(height):
	for x in range(width):
	expected = b + y * width + x
	if packed:
	expected = expected % 16
	assert arr[b, y, x] == expected
	assert band[y, x] == expected

	@pytest.mark.skipif(
	pyspark.__version__ < "3.4", reason="requires Spark 3.4 or higher"
	)
	def test_serialize_round_trip_via_udf(self):
	"""Deserialize → serialize round-trip through a raster-returning UDF.
	Verifies that the Python-serialized bytes can be deserialized back
	to produce correct pixel values.
	"""
	from sedona.spark.raster import raster_serde

	spark = self.spark

	df = spark.sql(
	"SELECT RS_MakeRasterForTesting(3, 'I', 'BandedSampleModel', "
	"10, 8, 100, 100, 10, -10, 0, 0, 3857) as raster"
	)
	# Collect the deserialized raster from Spark
	original = df.first()["raster"]

	# Serialize on Python side and deserialize back
	serialized = raster_serde.serialize(original)
	round_tripped = raster_serde.deserialize(serialized)

	# Verify pixel values for all 3 bands
	arr = round_tripped.as_numpy()
	assert arr.shape == (3, 8, 10)
	for b in range(3):
	for y in range(8):
	for x in range(10):
	expected = b + y * 10 + x
	assert arr[b, y, x] == expected

	round_tripped.close()
	original.close()

	@pytest.mark.skipif(
	pyspark.__version__ < "3.4", reason="requires Spark 3.4 or higher"
	)
	def test_serialize_preserves_metadata(self):
	"""Verify CRS, affine, and dimensions survive serialize round-trip."""
	from sedona.spark.raster import raster_serde

	spark = self.spark

	df = spark.sql(
	"SELECT RS_MakeRasterForTesting(1, 'I', 'BandedSampleModel', "
	"4, 3, 100, 100, 10, -10, 0, 0, 3857) as raster"
	)
	original = df.first()["raster"]

	# Serialize and deserialize in Python
	serialized = raster_serde.serialize(original)
	round_tripped = raster_serde.deserialize(serialized)

	# Verify spatial metadata survived the round-trip
	assert round_tripped.width == 4
	assert round_tripped.height == 3
	assert round_tripped.affine_trans.scale_x == 10.0
	assert round_tripped.affine_trans.scale_y == -10.0
	assert round_tripped.affine_trans.ip_x == 100.0
	assert round_tripped.affine_trans.ip_y == 100.0
	assert "3857" in round_tripped.crs_wkt or "Mercator" in round_tripped.crs_wkt

	round_tripped.close()
	original.close()