fixed failing zarr tests

databrickslabs · Nov 8, 2024 · 9db3d35 · 9db3d35
1 parent ff07ebe
commit 9db3d35
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/...in/scala/com/databricks/labs/mosaic/core/raster/operator/retile/BalancedSubdivision.scala b/...in/scala/com/databricks/labs/mosaic/core/raster/operator/retile/BalancedSubdivision.scala
@@ -82,6 +82,7 @@ object BalancedSubdivision {
         tile: MosaicRasterTile,
         sizeInMb: Int
     ): Seq[MosaicRasterTile] = {
+        if (sizeInMb <= 0) return Seq(tile)
         val numSplits = getNumSplits(tile.getRaster, sizeInMb)
         val (x, y) = tile.getRaster.getDimensions
         val (tileX, tileY) = getTileSize(x, y, numSplits)

diff --git a/src/main/scala/com/databricks/labs/mosaic/datasource/gdal/ReTileOnRead.scala b/src/main/scala/com/databricks/labs/mosaic/datasource/gdal/ReTileOnRead.scala
@@ -15,8 +15,6 @@ import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.types._
 
-import java.nio.file.{Files, Paths}
-
 /** An object defining the retiling read strategy for the GDAL file format. */
 object ReTileOnRead extends ReadStrategy {
 
@@ -141,7 +139,8 @@ object ReTileOnRead extends ReadStrategy {
             row
         })
 
-        Files.deleteIfExists(Paths.get(tmpPath))
+        PathUtils.cleanUpPath(tmpPath)
+//        Try(Files.deleteIfExists(Paths.get(tmpPath)))
 
         rows.iterator
     }

diff --git a/src/test/scala/com/databricks/labs/mosaic/datasource/multiread/RasterAsGridReaderTest.scala b/src/test/scala/com/databricks/labs/mosaic/datasource/multiread/RasterAsGridReaderTest.scala
@@ -50,10 +50,12 @@ class RasterAsGridReaderTest extends MosaicSpatialQueryTest with SharedSparkSess
 
         noException should be thrownBy MosaicContext.read
             .format("raster_to_grid")
+            .option("sizeInMB", "16")
+            .option("resolution", "0")
             .option("extensions", "grib")
             .option("combiner", "min")
             .option("retile", "true")
-            .option("tileSize", "10")
+            .option("tileSize", "100")
             .option("kRingInterpolate", "3")
             .load(filePath)
             .select("measure")
@@ -70,8 +72,11 @@ class RasterAsGridReaderTest extends MosaicSpatialQueryTest with SharedSparkSess
 
         noException should be thrownBy MosaicContext.read
             .format("raster_to_grid")
+            .option("sizeInMB", "16")
+            .option("resolution", "0")
+            .option("retile", "true")
+            .option("tileSize", "100")
             .option("combiner", "max")
-            .option("tileSize", "10")
             .option("kRingInterpolate", "3")
             .load(filePath)
             .select("measure")