Implemented CrsUDT in Python.

Author: metasim

core/src/main/scala/org/apache/spark/sql/rf/TileUDT.scala

@@ -45,7 +45,7 @@ class TileUDT extends UserDefinedType[Tile] {
   def userClass: Class[Tile] = classOf[Tile]
 
   def sqlType: StructType = StructType(Seq(
-    StructField("cell_type", StringType, false),
+    StructField("cellType", StringType, false),
     StructField("cols", IntegerType, false),
     StructField("rows", IntegerType, false),
     StructField("cells", BinaryType, true),

core/src/main/scala/org/locationtech/rasterframes/expressions/package.scala

@@ -82,7 +82,6 @@ package object expressions {
     registry.registerExpression[GetCRS]("rf_crs")
     registry.registerExpression[RealizeTile]("rf_tile")
     registry.registerExpression[CreateProjectedRaster]("rf_proj_raster")
-    registry.registerExpression[Subtract]("rf_local_subtract")
     registry.registerExpression[Multiply]("rf_local_multiply")
     registry.registerExpression[Divide]("rf_local_divide")
     registry.registerExpression[NormalizedDifference]("rf_normalized_difference")

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

@@ -1,3 +1,23 @@
+/*
+ * This software is licensed under the Apache 2 license, quoted below.
+ *
+ * Copyright 2021 Azavea, Inc.
+ *
+ * Licensed 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.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ */
 package org.locationtech.rasterframes.ref
 
 import geotrellis.raster.GridBounds

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

@@ -26,6 +26,8 @@ class here provides the PyRasterFrames entry point.
 """
 from itertools import product
 import functools, math
+
+import pyproj
 from pyspark import SparkContext
 from pyspark.sql import DataFrame, Column
 from pyspark.sql.types import (UserDefinedType, StructType, StructField, BinaryType, DoubleType, ShortType, IntegerType, StringType)
@@ -42,7 +44,8 @@ class here provides the PyRasterFrames entry point.
 
 from typing import List, Tuple
 
-__all__ = ['RasterFrameLayer', 'Tile', 'TileUDT', 'CellType', 'Extent', 'CRS', 'RasterSourceUDT', 'TileExploder', 'NoDataFilter']
+__all__ = ['RasterFrameLayer', 'Tile', 'TileUDT', 'CellType', 'Extent',
+           'CRS', 'CrsUDT', 'RasterSourceUDT', 'TileExploder', 'NoDataFilter']
 
 
 class cached_property(object):
@@ -227,7 +230,12 @@ def __str__(self):
 class CRS(object):
     # NB: The name `crsProj4` has to match what's used in StandardSerializers.crsSerializers
     def __init__(self, crsProj4):
-        self.crsProj4 = crsProj4
+        if isinstance(crsProj4, pyproj.CRS):
+            self.crsProj4 = crsProj4.to_proj4()
+        elif isinstance(crsProj4, str):
+            self.crsProj4 = crsProj4
+        else:
+            raise ValueError('Unexpected CRS definition type: {}'.format(type(crsProj4)))
 
     @cached_property
     def __jvm__(self):
@@ -242,9 +250,13 @@ def proj4_str(self):
         """Alias for `crsProj4`"""
         return self.crsProj4
 
+    def __eq__(self, other):
+        return isinstance(other, CRS) and self.crsProj4 == other.crsProj4
+
 
 class CellType(object):
     def __init__(self, cell_type_name):
+        assert(isinstance(cell_type_name, str))
         self.cell_type_name = cell_type_name
 
     @classmethod
@@ -443,29 +455,34 @@ def sqlType(cls):
         """
         Mirrors `schema` in scala companion object org.apache.spark.sql.rf.TileUDT
         """
+        extent = StructType([
+            StructField("xmin",DoubleType(), True),
+            StructField("ymin",DoubleType(), True),
+            StructField("xmax",DoubleType(), True),
+            StructField("ymax",DoubleType(), True)
+        ])
+        subgrid = StructType([
+            StructField("colMin", IntegerType(), True),
+            StructField("rowMin", IntegerType(), True),
+            StructField("colMax", IntegerType(), True),
+            StructField("rowMax", IntegerType() ,True)
+        ])
+
+        ref = StructType([
+            StructField("source", StructType([
+                StructField("raster_source_kryo", BinaryType(), False)
+            ]),True),
+            StructField("bandIndex", IntegerType(), True),
+            StructField("subextent", extent ,True),
+            StructField("subgrid", subgrid, True),
+        ])
+
         return StructType([
-            StructField("cell_context", StructType([
-                StructField("cellType", StructType([
-                    StructField("cellTypeName", StringType(), False)
-                ]), False),
-                StructField("dimensions", StructType([
-                    StructField("cols", ShortType(), False),
-                    StructField("rows", ShortType(), False)
-                ]), False),
-            ]), False),
-            StructField("cell_data", StructType([
-                StructField("cells", BinaryType(), True),
-                StructField("ref", StructType([
-                    StructField("source", RasterSourceUDT(), False),
-                    StructField("bandIndex", IntegerType(), False),
-                    StructField("subextent", StructType([
-                        StructField("xmin", DoubleType(), False),
-                        StructField("ymin", DoubleType(), False),
-                        StructField("xmax", DoubleType(), False),
-                        StructField("ymax", DoubleType(), False)
-                    ]), True)
-                ]), True)
-            ]), False)
+             StructField("cellType", StringType(), False),
+             StructField("cols", IntegerType(), False),
+             StructField("rows", IntegerType(), False),
+             StructField("cells", BinaryType(), True),
+             StructField("ref", ref, True)
         ])
 
     @classmethod
@@ -478,20 +495,14 @@ def scalaUDT(cls):
 
     def serialize(self, tile):
         cells = bytearray(tile.cells.flatten().tobytes())
-        row = [
-            # cell_context
-            [
-                [tile.cell_type.cell_type_name],
-                tile.dimensions()
-            ],
-            # cell_data
-            [
-                # cells
-                cells,
-                None
-            ]
+        dims = tile.dimensions()
+        return [
+            tile.cell_type.cell_type_name,
+            dims[0],
+            dims[1],
+            cells,
+            None
         ]
-        return row
 
     def deserialize(self, datum):
         """
@@ -500,21 +511,21 @@ def deserialize(self, datum):
         :return: A Tile object from row data.
         """
 
-        cell_data_bytes = datum.cell_data.cells
+        cell_data_bytes = datum.cells
         if cell_data_bytes is None:
-            if datum.cell_data.ref is None:
+            if datum.ref is None:
                 raise Exception("Invalid Tile structure. Missing cells and reference")
             else:
-                payload = datum.cell_data.ref
+                payload = datum.ref
                 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
+            cell_type = CellType(datum.cellType)
+            cols = datum.cols
+            rows = datum.rows
 
         if cell_data_bytes is None:
             raise Exception("Unable to fetch cell data from: " + repr(datum))
@@ -540,6 +551,34 @@ def deserialize(self, datum):
 Tile.__UDT__ = TileUDT()
 
 
+class CrsUDT(UserDefinedType):
+    @classmethod
+    def sqlType(cls):
+        """
+        Mirrors `schema` in scala companion object org.apache.spark.sql.rf.CrsUDT
+        """
+        return StringType()
+
+    @classmethod
+    def module(cls):
+        return 'pyrasterframes.rf_types'
+
+    @classmethod
+    def scalaUDT(cls):
+        return 'org.apache.spark.sql.rf.CrsUDT'
+
+    def serialize(self, crs):
+        return crs.proj4_str
+
+    def deserialize(self, datum):
+        return CRS(datum)
+
+    deserialize.__safe_for_unpickling__ = True
+
+
+CRS.__UDT__ = CrsUDT()
+
+
 class TileExploder(JavaTransformer, DefaultParamsReadable, DefaultParamsWritable):
     """
     Python wrapper for TileExploder.scala