Small area stats (#597)

* updated the logic of some of the statistical raster functions to work with small clipping outputs
databrickslabs · Nov 13, 2024 · 851e69d · 851e69d
1 parent a8ef2d3
commit 851e69d
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 7 deletions.
diff --git a/src/main/scala/com/databricks/labs/mosaic/core/raster/gdal/MosaicRasterBandGDAL.scala b/src/main/scala/com/databricks/labs/mosaic/core/raster/gdal/MosaicRasterBandGDAL.scala
@@ -148,9 +148,19 @@ case class MosaicRasterBandGDAL(band: Band, id: Int) {
       */
     def computeMinMax: Seq[Double] = {
         val minMaxVals = Array.fill[Double](2)(0)
-        Try(band.ComputeRasterMinMax(minMaxVals, 0))
-            .map(_ => minMaxVals.toSeq)
-            .getOrElse(Seq(Double.NaN, Double.NaN))
+        // will GDAL refuse to compute these stats?
+        if (band.GetXSize() == 1 || band.GetYSize() == 1) {
+            val validPixels = values.filter(_ != noDataValue)
+            if (validPixels.isEmpty) {
+                return Seq(Double.NaN, Double.NaN)
+            } else {
+                Seq(validPixels.min, validPixels.max)
+            }
+        } else {
+            Try(band.ComputeRasterMinMax(minMaxVals, 0))
+                .map(_ => minMaxVals.toSeq)
+                .getOrElse(Seq(Double.NaN, Double.NaN))
+        }
     }
 
     /**

diff --git a/src/main/scala/com/databricks/labs/mosaic/datasource/multiread/RasterAsGridReader.scala b/src/main/scala/com/databricks/labs/mosaic/datasource/multiread/RasterAsGridReader.scala
@@ -64,7 +64,7 @@ class RasterAsGridReader(sparkSession: SparkSession) extends MosaicDataFrameRead
         } else {
             lit(config("convertToFormat"))
         }
-      
+
         val rasterToGridCombiner = getRasterToGridFunc(config("combiner"))
 
         val loadedDf = retiledDf

diff --git a/src/test/scala/com/databricks/labs/mosaic/expressions/raster/RST_MedianBehaviors.scala b/src/test/scala/com/databricks/labs/mosaic/expressions/raster/RST_MedianBehaviors.scala
@@ -9,7 +9,8 @@ import org.scalatest.matchers.should.Matchers._
 
 trait RST_MedianBehaviors extends QueryTest {
 
-    def behavior(indexSystem: IndexSystem, geometryAPI: GeometryAPI): Unit = {
+    def largeAreaBehavior(indexSystem: IndexSystem, geometryAPI: GeometryAPI): Unit = {
+
         val mc = MosaicContext.build(indexSystem, geometryAPI)
         mc.register()
         val sc = spark
@@ -49,4 +50,38 @@ trait RST_MedianBehaviors extends QueryTest {
 
     }
 
+    def smallAreaBehavior(indexSystem: IndexSystem, geometryAPI: GeometryAPI): Unit = {
+
+        val mc = MosaicContext.build(indexSystem, geometryAPI)
+        mc.register()
+        val sc = spark
+        import mc.functions._
+        import sc.implicits._
+
+        val testRegion = "Polygon ((-8151524 1216659, -8151061 1216659, -8151061 1217123, -8151524 1217123, -8151524 1216659))"
+
+        val rastersInMemory = spark.read
+            .format("gdal")
+            .option("raster_storage", "in-memory")
+            .load("src/test/resources/modis")
+
+        rastersInMemory.createOrReplaceTempView("source")
+
+        val df = rastersInMemory
+            .withColumn("tile", rst_clip($"tile", st_buffer(lit(testRegion), lit(-20))))
+            .withColumn("result", rst_median($"tile"))
+            .select("result")
+            .select(explode($"result").as("result"))
+
+        noException should be thrownBy spark.sql("""
+                                                   |select rst_median(tile) from source
+                                                   |""".stripMargin)
+
+        val result = df.as[Double].collect().max
+
+        result should be > 0.0
+
+    }
+
+
 }
diff --git a/src/test/scala/com/databricks/labs/mosaic/expressions/raster/RST_MedianTest.scala b/src/test/scala/com/databricks/labs/mosaic/expressions/raster/RST_MedianTest.scala
@@ -22,11 +22,18 @@ class RST_MedianTest extends QueryTest with SharedSparkSessionGDAL with RST_Medi
 
     // These tests are not index system nor geometry API specific.
     // Only testing one pairing is sufficient.
-    test("Testing rst_median behavior with H3IndexSystem and JTS") {
+    test("Testing rst_median behavior with H3IndexSystem and JTS (tessellation case)") {
         noCodegen {
             assume(System.getProperty("os.name") == "Linux")
-            behavior(H3IndexSystem, JTS)
+            largeAreaBehavior(H3IndexSystem, JTS)
         }
     }
 
+        test("Testing rst_median behavior with H3IndexSystem and JTS (small area case)") {
+            noCodegen {
+                assume(System.getProperty("os.name") == "Linux")
+                smallAreaBehavior(H3IndexSystem, JTS)
+            }
+    }
+
 }