GEOMESA-3575 FSDS - Move xz2 filter function into filter module (#3561)

* Allow for push-down predicate evaluation in Accumulo/HBase

Author: elahrvivaz

geomesa-filter/src/main/resources/META-INF/services/org.geotools.api.filter.expression.Function

@@ -5,3 +5,4 @@ org.locationtech.geomesa.filter.function.DateToLong
 org.locationtech.geomesa.filter.function.FastProperty
 org.locationtech.geomesa.filter.function.MurmurHashFunction
 org.locationtech.geomesa.filter.function.ProxyIdFunction
+org.locationtech.geomesa.filter.function.XZ2Function

geomesa-filter/src/main/scala/org/locationtech/geomesa/filter/function/XZ2Function.scala

@@ -0,0 +1,67 @@
+/***********************************************************************
+ * Copyright (c) 2013-2025 General Atomics Integrated Intelligence, Inc.
+ * All rights reserved. This program and the accompanying materials
+ * are made available under the terms of the Apache License, Version 2.0
+ * which accompanies this distribution and is available at
+ * https://www.apache.org/licenses/LICENSE-2.0
+ ***********************************************************************/
+
+package org.locationtech.geomesa.filter.function
+
+import org.geotools.api.feature.simple.SimpleFeature
+import org.geotools.api.filter.expression.{Expression, ExpressionVisitor}
+import org.geotools.filter.FunctionExpressionImpl
+import org.geotools.filter.capability.FunctionNameImpl
+import org.geotools.filter.capability.FunctionNameImpl.parameter
+import org.locationtech.geomesa.curve.XZ2SFC
+import org.locationtech.geomesa.filter.function.XZ2Function.GetDefaultGeometry
+import org.locationtech.jts.geom.Geometry
+
+/**
+ * Function to calculate an XZ2 hex-encoded value
+ */
+class XZ2Function extends FunctionExpressionImpl(XZ2Function.FunctionName) {
+
+  private var expression: Expression = _
+
+  override def setParameters(params: java.util.List[Expression]): Unit = {
+    super.setParameters(params)
+    if (params.isEmpty) {
+      expression = GetDefaultGeometry
+    } else {
+      expression = getExpression(0)
+    }
+  }
+
+  override def evaluate(o: AnyRef): AnyRef = {
+    if (o == null) {
+      return null
+    }
+    val value = expression.evaluate(o, classOf[Geometry])
+    if (value == null) {
+      return null
+    }
+    val env = value.getEnvelopeInternal
+    if (env.isNull) {
+      return null
+    }
+    XZ2SFC.hexEncode(XZ2SFC.index(env.getMinX, env.getMinY, env.getMaxX, env.getMaxY))
+  }
+}
+
+object XZ2Function {
+
+  val FunctionName = new FunctionNameImpl("xz2", classOf[String], parameter("geom", classOf[String], 0, 1))
+
+  private object GetDefaultGeometry extends Expression {
+
+    override def evaluate(obj: Any): Geometry = obj match {
+      case sf: SimpleFeature => sf.getDefaultGeometry.asInstanceOf[Geometry]
+      case _ => null
+    }
+
+    override def evaluate[T](obj: Any, context: Class[T]): T = evaluate(obj).asInstanceOf[T] // only called by our code, above
+
+    override def accept(visitor: ExpressionVisitor, extraData: Any): AnyRef = throw new UnsupportedOperationException()
+  }
+}

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-core/src/main/resources/META-INF/services/org.geotools.api.filter.expression.Function

@@ -221,7 +221,6 @@ object StorageMetadata {
 
     private val sfc = XZ2SFC
     private val factory = new GeometryFactory()
-    private val hexFormat = HexFormat.of()
 
     override val getAlias: String = "xz2"
 
@@ -235,11 +234,11 @@ object StorageMetadata {
       if (env.isNull) {
         throw new NullPointerException("Geometry has a null envelope")
       }
-      toHex(sfc.index(env.getMinX, env.getMinY, env.getMaxX, env.getMaxY))
+      sfc.hexEncode(sfc.index(env.getMinX, env.getMinY, env.getMaxX, env.getMaxY))
     }
 
     override def decode(value: String): Geometry = {
-      val (xmin, ymin, xmax, ymax) = sfc.invert(fromHex(value))
+      val (xmin, ymin, xmax, ymax) = sfc.invert(sfc.hexDecode(value))
       val ring = Array(
         new Coordinate(xmin, ymin),
         new Coordinate(xmin, ymax),
@@ -257,11 +256,6 @@ object StorageMetadata {
      * @param maxRanges a rough upper limit on the number of ranges to generate
      */
     def ranges(queries: Seq[(Double, Double, Double, Double)], maxRanges: Option[Int] = None): Seq[(String, String)] =
-      sfc.ranges(queries, maxRanges).map(r => toHex(r.lower) -> toHex(r.upper))
-
-    // our z values use 25 bits, so we only need 7 digits to hex encode the full value
-    // we bit-shift by 3 to move dead bits to the end for better prefix matching
-    private def toHex(z: Long): String = hexFormat.toHexDigits(z << 3, 7)
-    private def fromHex(hex: String): Long = HexFormat.fromHexDigitsToLong(hex) >>> 3
+      sfc.ranges(queries, maxRanges).map(r => sfc.hexEncode(r.lower) -> sfc.hexEncode(r.upper))
   }
 }

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-core/src/main/scala/org/locationtech/geomesa/fs/storage/core/StorageMetadata.scala

@@ -11,19 +11,14 @@ package schemes
 
 import org.geotools.api.feature.simple.SimpleFeature
 import org.geotools.api.filter.Filter
-import org.geotools.api.filter.expression.{Expression, ExpressionVisitor}
-import org.geotools.filter.FunctionExpressionImpl
-import org.geotools.filter.capability.FunctionNameImpl
-import org.geotools.filter.capability.FunctionNameImpl.parameter
 import org.locationtech.geomesa.curve.XZ2SFC
 import org.locationtech.geomesa.filter.FilterHelper
-import org.locationtech.geomesa.fs.storage.core.StorageMetadata.XZ2Encoder
+import org.locationtech.geomesa.filter.function.XZ2Function
 import org.locationtech.geomesa.fs.storage.core.schemes.SpatialScheme.SpatialPartitionSchemeFactory
 import org.locationtech.geomesa.fs.storage.core.schemes.XZ2Scheme.incrementHex
 import org.locationtech.geomesa.utils.geotools.GeometryUtils
 import org.locationtech.jts.geom.Geometry
 
-import java.util.HexFormat
 import scala.annotation.tailrec
 
 /**
@@ -42,18 +37,16 @@ import scala.annotation.tailrec
 case class XZ2Scheme(attribute: String, index: Int, bits: Int) extends PartitionScheme {
 
   import FilterHelper.ff
-  import XZ2Scheme.FullHexDigits
 
   require(bits % 4 == 0, s"Bit precision must be a multiple of 4, but received $bits")
 
   private val xz2 = XZ2SFC
-  private val hexFormat = HexFormat.of()
 
   // partition level derived from bits parameter
   // each level adds 2 bits (4 quadrants)
   private val partitionLevel = (bits / 2).toShort
   // number of hex digits used to represent our z value - bits = (xz2.g - partitionLevel) * 2, then divide by 4 to get hex
-  private val digits = FullHexDigits - ((xz2.g - partitionLevel) / 2)
+  private val digits = xz2.hexDigits - ((xz2.g - partitionLevel) / 2)
 
   lazy private val wholeWorldRanges = Some(generateRanges(Seq((-180, -90, 180, 90))))
 
@@ -68,9 +61,9 @@ case class XZ2Scheme(attribute: String, index: Int, bits: Int) extends Partition
   override def getCoveringFilter(partition: PartitionKey): Filter = {
     // TODO maybe we can improve this with *some* kind of bbox?
     val zPrefix = partition.value
-    val lower = zPrefix.padTo(FullHexDigits, '0')
-    val upper = zPrefix.padTo(FullHexDigits, 'f')
-    ff.between(ff.function(XZ2Scheme.FunctionName.getName), ff.literal(lower), ff.literal(upper))
+    val lower = zPrefix.padTo(xz2.hexDigits, '0')
+    val upper = zPrefix.padTo(xz2.hexDigits, 'f')
+    ff.between(ff.function(XZ2Function.FunctionName.getName), ff.literal(lower), ff.literal(upper))
   }
 
   override def getRangesForFilter(filter: Filter): Option[Seq[PartitionRange]] = {
@@ -104,16 +97,11 @@ case class XZ2Scheme(attribute: String, index: Int, bits: Int) extends Partition
   }
 
   // truncates a full-resolution index to a partition group ID
-  private def truncateToPartition(fullIndex: Long): String = hexFormat.toHexDigits(fullIndex << 3, FullHexDigits).take(digits)
+  private def truncateToPartition(fullIndex: Long): String = xz2.hexEncode(fullIndex).take(digits)
 }
 
 object XZ2Scheme extends SpatialPartitionSchemeFactory[Geometry]("xz2") {
 
-  // number of digits required to index a "full" (g == 12) xz value, which is 25 bits
-  private val FullHexDigits = 7
-
-  val FunctionName = new FunctionNameImpl("xz2", classOf[String], parameter("geom", classOf[String], 0, 1))
-
   override def buildPartitionScheme(bits: Int, geom: String, geomIndex: Int): PartitionScheme =
     XZ2Scheme(geom, geomIndex, bits)
 
@@ -130,44 +118,4 @@ object XZ2Scheme extends SpatialPartitionSchemeFactory[Geometry]("xz2") {
       incrementHex(hex.substring(0, pos) + '0' + hex.substring(pos + 1), pos - 1)
     }
   }
-
-  /**
-   * Function to calculate an XZ2 hex-encoded value
-   */
-  class XZ2Function extends FunctionExpressionImpl(FunctionName) {
-
-    private var expression: Expression = _
-
-    override def setParameters(params: java.util.List[Expression]): Unit = {
-      super.setParameters(params)
-      if (params.isEmpty) {
-        expression = GetDefaultGeometry
-      } else {
-        expression = getExpression(0)
-      }
-    }
-
-    override def evaluate(o: AnyRef): AnyRef = {
-      if (o == null) {
-        return null
-      }
-      val value = expression.evaluate(o, classOf[Geometry])
-      if (value == null) {
-        return null
-      }
-      XZ2Encoder.encode(value)
-    }
-  }
-
-  private object GetDefaultGeometry extends Expression {
-
-    override def evaluate(obj: Any): Geometry = obj match {
-      case sf: SimpleFeature => sf.getDefaultGeometry.asInstanceOf[Geometry]
-      case _ => null
-    }
-
-    override def evaluate[T](obj: Any, context: Class[T]): T = evaluate(obj).asInstanceOf[T] // only called by our code, above
-
-    override def accept(visitor: ExpressionVisitor, extraData: Any): AnyRef = throw new UnsupportedOperationException()
-  }
 }

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-core/src/main/scala/org/locationtech/geomesa/fs/storage/core/schemes/XZ2Scheme.scala

@@ -18,6 +18,7 @@ import java.util.Random
 class XZ2SchemeTest extends SpecificationWithJUnit {
 
   val sft = SimpleFeatureTypes.createType("test", "*geom:Point:srid=4326")
+  val sftPoly = SimpleFeatureTypes.createType("test", "*geom:Polygon:srid=4326")
 
   "XZ2Scheme" should {
 
@@ -39,13 +40,13 @@ class XZ2SchemeTest extends SpecificationWithJUnit {
       ps.getPartition(ScalaSimpleFeature.create(sft, "1", "POINT (-75 38)")).value mustEqual "66f2db"
     }
 
-    "calculate covering filters" in {
+    "calculate covering filters for points" in {
       val r = new Random(77)
       val features = Seq.tabulate(100) { i =>
         ScalaSimpleFeature.create(sft, s"$i", s"POINT (${r.nextDouble(-180, 180)} ${r.nextDouble(-90, 90)})")
       }
 
-      foreach(Seq(4, 8, 16)) { bits =>
+      foreach(Seq(4, 8, 12, 16)) { bits =>
         val ps = PartitionSchemeFactory.load(sft, s"xz2:bits=$bits")
         ps must beAnInstanceOf[XZ2Scheme]
         foreach(features.groupBy(ps.getPartition)) { case (partition, group) =>
@@ -54,5 +55,64 @@ class XZ2SchemeTest extends SpecificationWithJUnit {
         }
       }
     }
+
+    "calculate covering filters for polygons" in {
+      val r = new Random(42)
+
+      // generate polygons of various sizes
+      val features = Seq.tabulate(100) { i =>
+        val centerX = r.nextDouble(-180, 180)
+        val centerY = r.nextDouble(-90, 90)
+
+        // vary the size from small (0.1 degrees) to large (45 degrees)
+        val size = r.nextDouble(0.1, 45.0) / 2
+
+        val xmin = math.max(-180, centerX - size)
+        val ymin = math.max(-90, centerY - size)
+        val xmax = math.min(180, centerX + size)
+        val ymax = math.min(90, centerY + size)
+
+        ScalaSimpleFeature.create(sftPoly, s"$i", s"POLYGON (($xmin $ymin, $xmax $ymin, $xmax $ymax, $xmin $ymax, $xmin $ymin))")
+      }
+
+      foreach(Seq(4, 8, 12, 16)) { bits =>
+        val ps = PartitionSchemeFactory.load(sftPoly, s"xz2:bits=$bits")
+        ps must beAnInstanceOf[XZ2Scheme]
+        foreach(features.groupBy(ps.getPartition)) { case (partition, group) =>
+          val filter = ps.getCoveringFilter(partition)
+          features.filter(filter.evaluate) must containTheSameElementsAs(group)
+        }
+      }
+    }
+
+    "calculate covering filters for mixed geometry sizes" in {
+      val r = new Random(123)
+
+      // mix of very small and very large polygons
+      val features = Seq.tabulate(50) { i =>
+        val centerX = r.nextDouble(-180, 180)
+        val centerY = r.nextDouble(-90, 90)
+
+        // alternate between tiny and huge polygons
+        val size = (if (i % 2 == 0) { r.nextDouble(0.01, 0.1) } else { r.nextDouble(30.0, 90.0) }) / 2
+
+        val xmin = math.max(-180, centerX - size)
+        val ymin = math.max(-90, centerY - size)
+        val xmax = math.min(180, centerX + size)
+        val ymax = math.min(90, centerY + size)
+
+        ScalaSimpleFeature.create(sftPoly, s"$i", s"POLYGON (($xmin $ymin, $xmax $ymin, $xmax $ymax, $xmin $ymax, $xmin $ymin))")
+      }
+
+      foreach(Seq(4, 8, 12, 16)) { bits =>
+        val ps = PartitionSchemeFactory.load(sftPoly, s"xz2:bits=$bits")
+        ps must beAnInstanceOf[XZ2Scheme]
+        foreach(features.groupBy(ps.getPartition)) { case (partition, group) =>
+          val filter = ps.getCoveringFilter(partition)
+          val matched = features.filter(filter.evaluate)
+          matched must containTheSameElementsAs(group)
+        }
+      }
+    }
   }
 }

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-core/src/test/scala/org/locationtech/geomesa/fs/storage/core/schemes/XZ2SchemeTest.scala

@@ -11,6 +11,7 @@ package org.locationtech.geomesa.curve
 import org.locationtech.geomesa.curve.XZ2SFC.{QueryWindow, XElement}
 import org.locationtech.geomesa.zorder.sfcurve.IndexRange
 
+import java.util.HexFormat
 import scala.collection.mutable.ArrayBuffer
 
 /**
@@ -25,6 +26,8 @@ class XZ2SFC(val g: Short, val xBounds: (Double, Double), val yBounds: (Double,
 
   // TODO see what the max value of g can be where we can use Ints instead of Longs and possibly refactor to use Ints
 
+  private val hexFormat = HexFormat.of()
+
   private val xLo = xBounds._1
   private val xHi = xBounds._2
   private val yLo = yBounds._1
@@ -33,6 +36,31 @@ class XZ2SFC(val g: Short, val xBounds: (Double, Double), val yBounds: (Double,
   private val xSize = xHi - xLo
   private val ySize = yHi - yLo
 
+  // upper bound on the z values produces by this curve (for g values <= 30)
+  val maxZValue: Long = (math.pow(4, g + 1).toLong - 1) / 3
+
+  // number of bits we bit-shift hex values, in order to move the significant digits to the left and allow prefix matching
+  private val hexBitOffset = java.lang.Long.numberOfLeadingZeros(maxZValue) % 4
+
+  // number of digits used in hex encoding function
+  val hexDigits: Int = ((64 - java.lang.Long.numberOfLeadingZeros(maxZValue)) / 4) + (if (hexBitOffset == 0) { 0 } else { 1 })
+
+  /**
+   * Encodes a z value into hex, bit-shifting left as necessary so that prefix matching works
+   *
+   * @param z z value
+   * @return hex-encoded string
+   */
+  def hexEncode(z: Long): String = hexFormat.toHexDigits(z << hexBitOffset, hexDigits)
+
+  /**
+   * Decodes a previously encoded z value from hex
+   *
+   * @param hex hex string from `hexEncode`
+   * @return original z value
+   */
+  def hexDecode(hex: String): Long = HexFormat.fromHexDigitsToLong(hex) >>> hexBitOffset
+
   /**
     * Index a polygon by its bounding box
     *