Add rf_local_is_in function

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

Author: vpipkt

core/src/main/scala/org/locationtech/rasterframes/RasterFunctions.scala

@@ -389,6 +389,9 @@ trait RasterFunctions {
   /** Cellwise inequality comparison between a tile and a scalar. */
   def rf_local_unequal[T: Numeric](tileCol: Column, value: T): Column = Unequal(tileCol, value)
 
+  /** Test if each cell value is in provided array */
+  def rf_local_is_in(tileCol: Column, arrayCol: Column) = IsIn(tileCol, arrayCol)
+
   /** Return a tile with ones where the input is NoData, otherwise zero */
   def rf_local_no_data(tileCol: Column): Column = Undefined(tileCol)
 

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

@@ -0,0 +1,88 @@
+/*
+ * This software is licensed under the Apache 2 license, quoted below.
+ *
+ * Copyright 2019 Astraea, 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.expressions.localops
+
+import geotrellis.raster.Tile
+import geotrellis.raster.mapalgebra.local.IfCell
+import org.apache.spark.sql.Column
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{TypeCheckFailure, TypeCheckSuccess}
+import org.apache.spark.sql.types.{ArrayType, DataType}
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.catalyst.expressions.{BinaryExpression, Expression, ExpressionDescription}
+import org.apache.spark.sql.catalyst.util.ArrayData
+import org.apache.spark.sql.rf.TileUDT
+import org.locationtech.rasterframes.encoders.CatalystSerializer._
+import org.locationtech.rasterframes.expressions.DynamicExtractors._
+import org.locationtech.rasterframes.expressions._
+
+@ExpressionDescription(
+  usage = "_FUNC_(tile, rhs) - In each cell of `tile`, return true if the value is in rhs.",
+  arguments = """
+  Arguments:
+    * tile - tile column to apply abs
+    * rhs - array to test against
+  """,
+  examples = """
+  Examples:
+    > SELECT  _FUNC_(tile, array);
+       ..."""
+)
+case class IsIn(left: Expression, right: Expression) extends BinaryExpression with CodegenFallback {
+  override val nodeName: String = "rf_local_is_in"
+
+  override def dataType: DataType = left.dataType
+
+  @transient private lazy val elementType: DataType = left.dataType.asInstanceOf[ArrayType].elementType
+
+  override def checkInputDataTypes(): TypeCheckResult =
+    if(!tileExtractor.isDefinedAt(left.dataType)) {
+      TypeCheckFailure(s"Input type '${left.dataType}' does not conform to a raster type.")
+    } else right.dataType match {
+      case _: ArrayType ⇒ TypeCheckSuccess
+      case _ ⇒ TypeCheckFailure(s"Input type '${right.dataType}' does not conform to ArrayType.")
+    }
+
+  override protected def nullSafeEval(input1: Any, input2: Any): Any = {
+    implicit val tileSer = TileUDT.tileSerializer
+    val (childTile, childCtx) = tileExtractor(left.dataType)(row(input1))
+
+    val arr = input2.asInstanceOf[ArrayData].toArray[AnyRef](elementType)
+
+    childCtx match {
+      case Some(ctx) => ctx.toProjectRasterTile(op(childTile, arr)).toInternalRow
+      case None => op(childTile, arr).toInternalRow
+    }
+
+  }
+
+  protected def op(left: Tile, right: IndexedSeq[AnyRef]): Tile = {
+    def fn(i: Int): Boolean = right.contains(i)
+    IfCell(left, fn(_), 1, 0)
+  }
+
+}
+
+object IsIn {
+  def apply(left: Column, right: Column): Column =
+    new Column(IsIn(left.expr, right.expr))
+}

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

@@ -87,6 +87,7 @@ package object expressions {
     registry.registerExpression[GreaterEqual]("rf_local_greater_equal")
     registry.registerExpression[Equal]("rf_local_equal")
     registry.registerExpression[Unequal]("rf_local_unequal")
+    registry.registerExpression[IsIn]("rf_local_is_in")
     registry.registerExpression[Undefined]("rf_local_no_data")
     registry.registerExpression[Defined]("rf_local_data")
     registry.registerExpression[Sum]("rf_tile_sum")

core/src/test/scala/org/locationtech/rasterframes/RasterFunctionsSpec.scala

@@ -972,4 +972,25 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
     val dResult = df.select($"ld").as[Tile].first()
     dResult should be (randNDPRT.localDefined())
   }
+
+  it("should check values isin"){
+    checkDocs("rf_local_is_in")
+
+    // tile is 3 by 3 with values, 1 to 9
+    val df = Seq((byteArrayTile, lit(1), lit(5), lit(10))).toDF("t", "one", "five", "ten")
+      .withColumn("in_expect_2", rf_local_is_in($"t", array($"one", $"five")))
+      .withColumn("in_expect_1", rf_local_is_in($"t", array($"ten", $"five")))
+      .withColumn("in_expect_0", rf_local_is_in($"t", array($"ten")))
+
+    val e2Result = df.select(rf_tile_sum($"in_expect_2")).as[Double].first()
+    e2Result should be (2.0)
+
+    val e1Result = df.select(rf_tile_sum($"in_expect_1")).as[Double].first()
+    e1Result should be (1.0)
+
+    val e0Result = df.select($"in_expect_1").as[Tile].first()
+    e0Result.toArray() should contain only (0)
+
+//    lazy val invalid = df.select(rf_local_is_in($"t", lit("foobar"))).as[Tile].first()
+  }
 }