Updated handling of Python TileUDT to handle RasterRefs better.

Author: metasim

core/src/main/scala/org/locationtech/rasterframes/ref/DelegatingRasterSource.scala

@@ -43,7 +43,7 @@ abstract class DelegatingRasterSource(source: URI, delegateBuilder: () => GTRast
       f(_delRef)
     }
     catch {
-      // On this exeception we attempt to recreate the delegate and read again.
+      // On this Exeception we attempt to recreate the delegate and read again.
       case _: java.nio.BufferUnderflowException =>
         _delRef = null
         val newDel = delegateBuilder()
@@ -62,7 +62,7 @@ abstract class DelegatingRasterSource(source: URI, delegateBuilder: () => GTRast
   override def hashCode(): Int = source.hashCode()
 
   // This helps reduce header reads between serializations
-  def info: SimpleRasterInfo = SimpleRasterInfo.cache.get(source.toASCIIString, _ =>
+  def info: SimpleRasterInfo = SimpleRasterInfo(source.toASCIIString, _ =>
     retryableRead(rs => SimpleRasterInfo(rs))
   )
 

core/src/main/scala/org/locationtech/rasterframes/ref/GDALRasterSource.scala

@@ -48,7 +48,7 @@ case class GDALRasterSource(source: URI) extends RasterSource with URIRasterSour
     VLMRasterSource(tweaked)
   }
 
-  protected def tiffInfo = SimpleRasterInfo.cache.get(source.toASCIIString, _ => SimpleRasterInfo(gdal))
+  protected def tiffInfo = SimpleRasterInfo(source.toASCIIString, _ => SimpleRasterInfo(gdal))
 
   override def crs: CRS = tiffInfo.crs
 

core/src/main/scala/org/locationtech/rasterframes/ref/SimpleRasterInfo.scala

@@ -44,6 +44,11 @@ case class SimpleRasterInfo(
 )
 
 object SimpleRasterInfo {
+  // Not a fan of this.... need a better abstraction that doesn't put the
+  // type-specific logic in here.
+  def apply(key: String, builder: String => SimpleRasterInfo): SimpleRasterInfo =
+    cache.get(key, builder)
+
   def apply(info: GeoTiffReader.GeoTiffInfo): SimpleRasterInfo =
     SimpleRasterInfo(
       info.segmentLayout.totalCols,
@@ -76,10 +81,9 @@ object SimpleRasterInfo {
     )
   }
 
-  private[rasterframes]
-  lazy val cache = Scaffeine()
+  private lazy val cache = Scaffeine()
     .recordStats()
     .build[String, SimpleRasterInfo]
 
   def cacheStats = cache.stats()
-}
+}

pyrasterframes/src/main/python/pyrasterframes/rf_context.py

@@ -49,15 +49,15 @@ def build_info(self):
         return self._jrfctx.buildInfo()
 
     # NB: Tightly coupled to `org.locationtech.rasterframes.py.PyRFContext._resolveRasterRef`
-    def _resolve_raster_ref(self, ref):
+    def _resolve_raster_ref(self, ref_struct):
         f = self.lookup("_resolveRasterRef")
         return f(
-            ref.source.raster_source_kryo,
-            ref.bandIndex,
-            ref.subextent.xmin,
-            ref.subextent.ymin,
-            ref.subextent.xmax,
-            ref.subextent.ymax,
+            ref_struct.source.raster_source_kryo,
+            ref_struct.bandIndex,
+            ref_struct.subextent.xmin,
+            ref_struct.subextent.ymin,
+            ref_struct.subextent.xmax,
+            ref_struct.subextent.ymax,
         )
 
     @staticmethod

pyrasterframes/src/main/python/pyrasterframes/rf_types.py

@@ -421,16 +421,22 @@ def deserialize(self, datum):
         :param datum:
         :return: A Tile object from row data.
         """
-        cell_type = CellType(datum.cell_context.cellType.cellTypeName)
-        cols = datum.cell_context.dimensions.cols
-        rows = datum.cell_context.dimensions.rows
+
         cell_data_bytes = datum.cell_data.cells
         if cell_data_bytes is None:
             if datum.cell_data.ref is None:
                 raise Exception("Invalid Tile structure. Missing cells and reference")
             else:
                 payload = datum.cell_data.ref
-                cell_data_bytes = RFContext.active()._resolve_raster_ref(payload)
+                ref = RFContext.active()._resolve_raster_ref(payload)
+                cell_type = CellType(ref.cellType().name())
+                cols = ref.cols()
+                rows = ref.rows()
+                cell_data_bytes = ref.tile().toBytes()
+        else:
+            cell_type = CellType(datum.cell_context.cellType.cellTypeName)
+            cols = datum.cell_context.dimensions.cols
+            rows = datum.cell_context.dimensions.rows
 
         if cell_data_bytes is None:
             raise Exception("Unable to fetch cell data from: " + repr(datum))

pyrasterframes/src/main/scala/org/locationtech/rasterframes/py/PyRFContext.scala

@@ -228,8 +228,7 @@ class PyRFContext(implicit sparkSession: SparkSession) extends RasterFunctions
   def _resolveRasterRef(srcBin: Array[Byte], bandIndex: jInt, xmin: jDouble, ymin: jDouble, xmax: jDouble, ymax: jDouble): AnyRef = {
     val src = KryoSupport.deserialize[RasterSource](ByteBuffer.wrap(srcBin))
     val extent = Extent(xmin, ymin, xmax, ymax)
-    val ref = RasterRef(src, bandIndex, Some(extent), None)
-    ref.tile.toArrayTile().toBytes()
+    RasterRef(src, bandIndex, Some(extent), None)
   }
 
   def _dfToMarkdown(df: DataFrame, numRows: Int, truncate: Boolean): String = {