rf_local_is_in python implementation

Signed-off-by: Jason T. Brown <[email protected]>

Author: vpipkt

core/src/main/scala/org/locationtech/rasterframes/expressions/localops/IsIn.scala

@@ -44,7 +44,7 @@ import org.locationtech.rasterframes.expressions._
   """,
   examples = """
   Examples:
-    > SELECT  _FUNC_(tile, array);
+    > SELECT  _FUNC_(tile, array(lit(33), lit(66), lit(99)));
        ..."""
 )
 case class IsIn(left: Expression, right: Expression) extends BinaryExpression with CodegenFallback {

docs/src/main/paradox/release-notes.md

@@ -6,6 +6,7 @@
 
 * _Breaking_ (potentially): removed `GeoTiffCollectionRelation` due to usage limitation and overlap with `RasterSourceDataSource` functionality.
 * Upgraded to Spark 2.4.4
+ * Add `rf_local_is_in` raster function
 
 ### 0.8.3
 

pyrasterframes/src/main/python/docs/reference.pymd

@@ -183,7 +183,7 @@ Parameters `tile_columns` and `tile_rows` are literals, not column expressions.
 
     Tile rf_array_to_tile(Array arrayCol, Int numCols, Int numRows)
 
-Python only. Create a `tile` from a Spark SQL [Array](http://spark.apache.org/docs/2.3.2/api/python/pyspark.sql.html#pyspark.sql.types.ArrayType), filling values in row-major order.
+Python only. Create a `tile` from a Spark SQL [Array][Array], filling values in row-major order.
 
 ### rf_assemble_tile
 
@@ -374,6 +374,13 @@ Returns a `tile` column containing the element-wise equality of `tile1` and `rhs
 
 Returns a `tile` column containing the element-wise inequality of `tile1` and `rhs`.
 
+### rf_local_is_in
+
+    Tile rf_local_is_in(Tile tile, Array array)
+    Tile rf_local_is_in(Tile tile, list l)
+
+Returns a `tile` column with cell values of 1 where the `tile` cell value is in the provided array or list. The `array` is a Spark SQL [Array][Array]. A python `list` of numeric values can also be passed.
+
 ### rf_round
 
     Tile rf_round(Tile tile)
@@ -621,13 +628,13 @@ Python only. As with @ref:[`rf_explode_tiles`](reference.md#rf-explode-tiles), b
 
     Array rf_tile_to_array_int(Tile tile)
 
-Convert Tile column to Spark SQL [Array](http://spark.apache.org/docs/2.3.2/api/python/pyspark.sql.html#pyspark.sql.types.ArrayType), in row-major order. Float cell types will be coerced to integral type by flooring.
+Convert Tile column to Spark SQL [Array][Array], in row-major order. Float cell types will be coerced to integral type by flooring.
 
 ### rf_tile_to_array_double
 
     Array rf_tile_to_arry_double(Tile tile)
 
-Convert tile column to Spark [Array](http://spark.apache.org/docs/2.3.2/api/python/pyspark.sql.html#pyspark.sql.types.ArrayType), in row-major order. Integral cell types will be coerced to floats.
+Convert tile column to Spark [Array][Array], in row-major order. Integral cell types will be coerced to floats.
 
 ### rf_render_ascii
 
@@ -657,3 +664,4 @@ Runs [`rf_rgb_composite`](reference.md#rf-rgb-composite) on the given tile colum
 
 [RasterFunctions]: org.locationtech.rasterframes.RasterFunctions
 [scaladoc]: latest/api/index.html
+[Array]: http://spark.apache.org/docs/latest/api/python/pyspark.sql.html#pyspark.sql.types.ArrayType

pyrasterframes/src/main/python/pyrasterframes/rasterfunctions.py

@@ -260,14 +260,24 @@ def rf_local_unequal_int(tile_col, scalar):
     """Return a Tile with values equal 1 if the cell is not equal to a scalar, otherwise 0"""
     return _apply_scalar_to_tile('rf_local_unequal_int', tile_col, scalar)
 
+
 def rf_local_no_data(tile_col):
     """Return a tile with ones where the input is NoData, otherwise zero."""
     return _apply_column_function('rf_local_no_data', tile_col)
 
+
 def rf_local_data(tile_col):
     """Return a tile with zeros where the input is NoData, otherwise one."""
     return _apply_column_function('rf_local_data', tile_col)
 
+def rf_local_is_in(tile_col, array):
+    """Return a tile with cell values of 1 where the `tile_col` cell is in the provided array."""
+    from pyspark.sql.functions import array as sql_array, lit
+    if isinstance(array, list):
+        array = sql_array([lit(v) for v in array])
+
+    return _apply_column_function('rf_local_is_in', tile_col, array)
+
 def _apply_column_function(name, *args):
     jfcn = RFContext.active().lookup(name)
     jcols = [_to_java_column(arg) for arg in args]

pyrasterframes/src/main/python/tests/PyRasterFramesTests.py

@@ -131,7 +131,7 @@ def test_tile_udt_serialization(self):
             cells[1][1] = nd
             a_tile = Tile(cells, ct.with_no_data_value(nd))
             round_trip = udt.fromInternal(udt.toInternal(a_tile))
-            self.assertEquals(a_tile, round_trip, "round-trip serialization for " + str(ct))
+            self.assertEqual(a_tile, round_trip, "round-trip serialization for " + str(ct))
 
             schema = StructType([StructField("tile", TileUDT(), False)])
             df = self.spark.createDataFrame([{"tile": a_tile}], schema)

pyrasterframes/src/main/python/tests/RasterFunctionsTests.py

@@ -347,8 +347,11 @@ def test_rf_local_data_and_no_data(self):
         import numpy as np
         from numpy.testing import assert_equal
 
-        t = Tile(np.array([[1, 3, 4], [5, 0, 3]]), CellType.uint8().with_no_data_value(5))
-        #note the convert is due to issue #188
+        nd = 5
+        t = Tile(
+            np.array([[1, 3, 4], [nd, 0, 3]]),
+            CellType.uint8().with_no_data_value(nd))
+        # note the convert is due to issue #188
         df = self.spark.createDataFrame([Row(t=t)])\
             .withColumn('lnd', rf_convert_cell_type(rf_local_no_data('t'), 'uint8')) \
             .withColumn('ld',  rf_convert_cell_type(rf_local_data('t'),    'uint8'))
@@ -359,3 +362,34 @@ def test_rf_local_data_and_no_data(self):
 
         result_d = result['ld']
         assert_equal(result_d.cells, np.invert(t.cells.mask))
+
+    def test_rf_local_is_in(self):
+        from pyspark.sql.functions import lit, array, col
+        from pyspark.sql import Row
+        import numpy as np
+        from numpy.testing import assert_equal
+
+        nd = 5
+        t = Tile(
+            np.array([[1, 3, 4], [nd, 0, 3]]),
+            CellType.uint8().with_no_data_value(nd))
+        # note the convert is due to issue #188
+        df = self.spark.createDataFrame([Row(t=t)]) \
+            .withColumn('a', array(lit(3), lit(4))) \
+            .withColumn('in2', rf_convert_cell_type(
+                rf_local_is_in(col('t'), array(lit(0), lit(4))),
+                'uint8')) \
+            .withColumn('in3', rf_convert_cell_type(rf_local_is_in('t', 'a'), 'uint8')) \
+            .withColumn('in4', rf_convert_cell_type(
+                rf_local_is_in('t', array(lit(0), lit(4), lit(3))),
+                'uint8')) \
+            .withColumn('in_list', rf_convert_cell_type(rf_local_is_in(col('t'), [4, 1]), 'uint8'))
+
+        result = df.first()
+        self.assertEqual(result['in2'].cells.sum(), 2)
+        assert_equal(result['in2'].cells, np.isin(t.cells, np.array([0, 4])))
+        self.assertEqual(result['in3'].cells.sum(), 3)
+        self.assertEqual(result['in4'].cells.sum(), 4)
+        self.assertEqual(result['in_list'].cells.sum(), 2,
+                         "Tile value {} should contain two 1s as: [[1, 0, 1],[0, 0, 0]]"
+                         .format(result['in_list'].cells))