GEOMESA-3565 FSDS - Add step to datetime partition scheme (#3525)

Author: elahrvivaz

docs/user/filesystem/partition_schemes.rst

@@ -26,6 +26,7 @@ The following options are supported:
 
 * ``attribute`` - The name of a ``Date``\ -type attribute from the SimpleFeatureType to use. If not specified, the default
   date attribute is used.
+* ``step`` - The number of time units (hours, days, etc) to include in each partition.
 
 Spatial Schemes
 ---------------

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-common/src/main/scala/org/locationtech/geomesa/fs/storage/common/partitions/AttributeScheme.scala

@@ -141,8 +141,7 @@ object AttributeScheme {
       if (opts.name != Name) { None } else {
         val attribute = opts.getSingle("attribute").orNull
         require(attribute != null, s"Attribute scheme requires an attribute to be specified with 'attribute=<attribute>'")
-        val index = sft.indexOf(attribute)
-        require(index != -1, s"Attribute '$attribute' does not exist in schema '${sft.getTypeName}'")
+        val index = attributeIndex(sft, attribute)
         val binding = sft.getDescriptor(index).getType.getBinding
         require(AttributeIndexKey.encodable(binding), s"Invalid type binding '${binding.getName}' of attribute '$attribute'")
 

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-common/src/main/scala/org/locationtech/geomesa/fs/storage/common/partitions/DateTimeScheme.scala

@@ -22,17 +22,16 @@ import java.time.temporal.ChronoUnit
 import java.time.{Instant, ZoneOffset, ZonedDateTime}
 import java.util.{Date, Locale}
 
-case class DateTimeScheme(
-    dtg: String,
-    dtgIndex: Int,
-    unit: ChronoUnit,
-  ) extends PartitionScheme {
+case class DateTimeScheme(dtg: String, dtgIndex: Int, unit: ChronoUnit, step: Int = 1) extends PartitionScheme {
 
   import FilterHelper.ff
 
   private val encoder = LexiTypeEncoders.integerEncoder()
 
-  override val name: String = s"${unit.name().toLowerCase(Locale.US)}:attribute=$dtg"
+  override val name: String = {
+    val stepOpt = if (step == 1) { "" } else { s":step=$step"}
+    s"${unit.name().toLowerCase(Locale.US)}:attribute=$dtg$stepOpt"
+  }
 
   override def getPartition(feature: SimpleFeature): PartitionKey = {
     val instant = feature.getAttribute(dtgIndex).asInstanceOf[Date].toInstant
@@ -65,16 +64,13 @@ case class DateTimeScheme(
   }
 
   override def getCoveringFilter(partition: PartitionKey): Filter = {
-    val offset = encoder.decode(partition.value)
+    val offset = encoder.decode(partition.value) * step
     val start = DateTimeScheme.Epoch.plus(offset.longValue(), unit)
-    val end = ff.literal(DateParsing.format(start.plus(1, unit)))
+    val end = ff.literal(DateParsing.format(start.plus(step, unit)))
     ff.and(ff.greaterOrEqual(ff.property(dtg), ff.literal(DateParsing.format(start))), ff.less(ff.property(dtg), end))
   }
 
-  private def toPartition(dt: ZonedDateTime): Int = {
-    require(!dt.isBefore(DateTimeScheme.Epoch), s"Date exceeds minimum indexable value (${DateTimeScheme.Epoch}): $dt")
-    unit.between(DateTimeScheme.Epoch, dt).toInt
-  }
+  private def toPartition(dt: ZonedDateTime): Int = unit.between(DateTimeScheme.Epoch, dt).toInt / step
 
   private def getBounds(filter: Filter): Option[Seq[Bounds[ZonedDateTime]]] = {
     val bounds = FilterHelper.extractIntervals(filter, dtg)
@@ -107,7 +103,6 @@ object DateTimeScheme {
 
   private val UnboundedUpper = "zzz"
 
-  // TODO allow for 2 hours, etc
   class DateTimePartitionSchemeFactory extends PartitionSchemeFactory {
     override def load(sft: SimpleFeatureType, scheme: String): Option[PartitionScheme] = {
       val opts = SchemeOpts(scheme)
@@ -122,10 +117,9 @@ object DateTimeScheme {
       unit.map { u =>
         val dtg = opts.getSingle("attribute").orElse(sft.getDtgField).orNull
         require(dtg != null, s"Date scheme requires an attribute to be specified with 'attribute=<attribute>'")
-        val index = sft.indexOf(dtg)
-        require(index != -1, s"Attribute '$dtg' does not exist in schema '${sft.getTypeName}'")
-        require(classOf[Date].isAssignableFrom(sft.getDescriptor(index).getType.getBinding), s"Attribute '$dtg' is not a date")
-        DateTimeScheme(dtg, index, u)
+        val index = attributeIndex(sft, dtg, Some(classOf[Date]))
+        val step = opts.getSingle("step").map(_.toInt).getOrElse(1)
+        DateTimeScheme(dtg, index, u, step)
       }
     }
   }

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-common/src/main/scala/org/locationtech/geomesa/fs/storage/common/partitions/HierarchicalDateTimeScheme.scala

@@ -219,9 +219,7 @@ object HierarchicalDateTimeScheme extends PartitionSchemeFactory {
     lazy val dtg = opts.opts.get(Config.DtgAttribute).orElse(sft.getDtgField).getOrElse {
       throw new IllegalArgumentException(s"DateTime scheme requires valid attribute '${Config.DtgAttribute}'")
     }
-    lazy val dtgIndex = Some(sft.indexOf(dtg)).filter(_ != -1).getOrElse {
-      throw new IllegalArgumentException(s"Attribute '$dtg' does not exist in feature type ${sft.getTypeName}")
-    }
+    lazy val dtgIndex = attributeIndex(sft, dtg, Some(classOf[Date]))
 
     if (opts.name == Name) {
       val unit = opts.opts.get(Config.StepUnitOpt).map(c => ChronoUnit.valueOf(c.toUpperCase)).getOrElse {

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-common/src/main/scala/org/locationtech/geomesa/fs/storage/common/partitions/SpatialScheme.scala

@@ -16,8 +16,10 @@ import org.locationtech.geomesa.fs.storage.api.StorageMetadata.PartitionKey
 import org.locationtech.geomesa.fs.storage.api.{PartitionScheme, PartitionSchemeFactory}
 import org.locationtech.geomesa.utils.geotools.GeometryUtils
 import org.locationtech.geomesa.zorder.sfcurve.IndexRange
+import org.locationtech.jts.geom.Geometry
 
 import java.util.regex.Pattern
+import scala.reflect.ClassTag
 
 abstract class SpatialScheme(id: String, bits: Int, geom: String) extends PartitionScheme {
 
@@ -69,7 +71,7 @@ object SpatialScheme {
 
   import org.locationtech.geomesa.utils.geotools.RichSimpleFeatureType.RichSimpleFeatureType
 
-  abstract class SpatialPartitionSchemeFactory(name: String) extends PartitionSchemeFactory {
+  abstract class SpatialPartitionSchemeFactory[T <: Geometry : ClassTag](name: String) extends PartitionSchemeFactory {
 
     private val namePattern: Pattern = Pattern.compile(s"$name-([0-9]+)bits?:?")
 
@@ -80,8 +82,7 @@ object SpatialScheme {
       def build(resolution: Short): Option[PartitionScheme] = {
         val geom = opts.getSingle("attribute").orElse(Option(sft.getGeomField)).orNull
         require(geom != null, s"Spatial schemes requires an attribute to be specified with 'attribute=<attribute>'")
-        val index = sft.indexOf(geom)
-        require(index != -1, s"Attribute '$geom' does not exist in schema '${sft.getTypeName}'")
+        val index = attributeIndex(sft, geom, Some(implicitly[ClassTag[T]].runtimeClass))
         Some(buildPartitionScheme(resolution, geom, index))
       }
 

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-common/src/main/scala/org/locationtech/geomesa/fs/storage/common/partitions/XZ2Scheme.scala

@@ -41,7 +41,7 @@ object XZ2Scheme {
 
   val Name = "xz2"
 
-  class XZ2PartitionSchemeFactory extends SpatialPartitionSchemeFactory(Name) {
+  class XZ2PartitionSchemeFactory extends SpatialPartitionSchemeFactory[Geometry](Name) {
     override def buildPartitionScheme(bits: Int, geom: String, geomIndex: Int): SpatialScheme =
       XZ2Scheme(bits, geom, geomIndex)
   }

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-common/src/main/scala/org/locationtech/geomesa/fs/storage/common/partitions/Z2Scheme.scala

@@ -56,7 +56,7 @@ object Z2Scheme {
 
   val Name = "z2"
 
-  class Z2PartitionSchemeFactory extends SpatialPartitionSchemeFactory(Name) {
+  class Z2PartitionSchemeFactory extends SpatialPartitionSchemeFactory[Point](Name) {
     override def buildPartitionScheme(bits: Int, geom: String, geomIndex: Int): SpatialScheme =
       Z2Scheme(bits, geom, geomIndex)
   }

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-common/src/main/scala/org/locationtech/geomesa/fs/storage/common/partitions/package.scala

@@ -8,6 +8,8 @@
 
 package org.locationtech.geomesa.fs.storage.common
 
+import org.geotools.api.feature.simple.SimpleFeatureType
+
 import java.nio.charset.StandardCharsets
 import java.util.Locale
 
@@ -17,6 +19,15 @@ package object partitions {
   // note: we use 1 instead of 0 b/c 0 is not a valid char in postgres so breaks jdbc metadata filtering
   val ZeroChar = new String(Array[Byte](1), StandardCharsets.UTF_8)
 
+  private[partitions] def attributeIndex(sft: SimpleFeatureType, name: String, binding: Option[Class[_]] = None): Int = {
+    val index = sft.indexOf(name)
+    require(index != -1, s"Attribute '$name' does not exist in schema '${sft.getTypeName}'")
+    binding.foreach { b =>
+      require(b.isAssignableFrom(sft.getDescriptor(index).getType.getBinding), s"Attribute '$name' is not a ${b.getSimpleName}")
+    }
+    index
+  }
+
   case class SchemeOpts(name: String, opts: Map[String, String], multiOpts: Map[String, Seq[String]]) {
     def getSingle(k: String): Option[String] = {
       if (multiOpts.contains(k)) {

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-common/src/test/scala/org/locationtech/geomesa/fs/storage/common/partitions/DateTimeSchemeTest.scala

@@ -36,6 +36,22 @@ class DateTimeSchemeTest extends Specification {
 
   "DateTimeScheme" should {
 
+    "partition based on hours" >> {
+      val ps = DateTimeScheme("dtg", 0, ChronoUnit.HOURS)
+      val partition = ps.getPartition(sf)
+      partition.value mustEqual "80064c8a"
+      val hours = AttributeIndexKey.decode("integer", partition.value).asInstanceOf[Int]
+      ZonedDateTime.ofInstant(Instant.EPOCH, ZoneOffset.UTC).plusHours(hours) mustEqual truncate(date, ChronoUnit.HOURS)
+    }
+
+    "partition based on hours with step" >> {
+      val ps = DateTimeScheme("dtg", 0, ChronoUnit.HOURS, step = 2)
+      val partition = ps.getPartition(sf)
+      partition.value mustEqual "80032645"
+      val hours = AttributeIndexKey.decode("integer", partition.value).asInstanceOf[Int] * 2
+      ZonedDateTime.ofInstant(Instant.EPOCH, ZoneOffset.UTC).plusHours(hours) mustEqual truncate(date, ChronoUnit.HOURS, 2)
+    }
+
     "partition based on days" >> {
       val ps = DateTimeScheme("dtg", 0, ChronoUnit.DAYS)
       val partition = ps.getPartition(sf)
@@ -44,12 +60,12 @@ class DateTimeSchemeTest extends Specification {
       ZonedDateTime.ofInstant(Instant.EPOCH, ZoneOffset.UTC).plusDays(days) mustEqual truncate(date, ChronoUnit.DAYS)
     }
 
-    "partition based on hours" >> {
-      val ps = DateTimeScheme("dtg", 0, ChronoUnit.HOURS)
+    "partition based on days with step" >> {
+      val ps = DateTimeScheme("dtg", 0, ChronoUnit.DAYS, step = 2)
       val partition = ps.getPartition(sf)
-      partition.value mustEqual "80064c8a"
-      val hours = AttributeIndexKey.decode("integer", partition.value).asInstanceOf[Int]
-      ZonedDateTime.ofInstant(Instant.EPOCH, ZoneOffset.UTC).plusHours(hours) mustEqual truncate(date, ChronoUnit.HOURS)
+      partition.value mustEqual "80002198"
+      val days = AttributeIndexKey.decode("integer", partition.value).asInstanceOf[Int] * 2
+      ZonedDateTime.ofInstant(Instant.EPOCH, ZoneOffset.UTC).plusDays(days) mustEqual truncate(date, ChronoUnit.DAYS, 2)
     }
 
     "partition based on week" >> {
@@ -61,13 +77,29 @@ class DateTimeSchemeTest extends Specification {
       ZonedDateTime.ofInstant(Instant.EPOCH, ZoneOffset.UTC).plusWeeks(weeks) mustEqual truncate(date, ChronoUnit.WEEKS)
     }
 
+    "partition based on week with step" >> {
+      val ps = PartitionSchemeFactory.load(sft, "weekly:step=2")
+      ps must beAnInstanceOf[DateTimeScheme]
+      val partition = ps.getPartition(sf)
+      partition.value mustEqual "800004cc"
+      val weeks = AttributeIndexKey.decode("integer", partition.value).asInstanceOf[Int] * 2
+      ZonedDateTime.ofInstant(Instant.EPOCH, ZoneOffset.UTC).plusWeeks(weeks) mustEqual truncate(date, ChronoUnit.WEEKS, 2)
+    }
+
     "enumerate partition ranges" >> {
       val ps = DateTimeScheme("dtg", 0, ChronoUnit.HOURS)
       val partitions = ps.getPartitionsForFilter(ECQL.toFilter("dtg >= '2017-02-03T10:15:00Z' AND dtg < '2017-02-03T11:18:00Z'"))
       partitions must beSome
       partitions.get.map(_.value) mustEqual Seq("80064c8a", "80064c8b")
     }
 
+    "enumerate partition ranges with step" >> {
+      val ps = DateTimeScheme("dtg", 0, ChronoUnit.HOURS, step = 2)
+      val partitions = ps.getPartitionsForFilter(ECQL.toFilter("dtg >= '2017-02-03T10:15:00Z' AND dtg < '2017-02-03T13:18:00Z'"))
+      partitions must beSome
+      partitions.get.map(_.value) mustEqual Seq("80032645", "80032646")
+    }
+
     "calculate covering filters for partitions" >> {
       foreach(Seq(ChronoUnit.HOURS, ChronoUnit.DAYS, ChronoUnit.WEEKS, ChronoUnit.MONTHS, ChronoUnit.YEARS)) { unit =>
         val ps = DateTimeScheme("dtg", 0, unit)
@@ -82,6 +114,20 @@ class DateTimeSchemeTest extends Specification {
       }
     }
 
+    "calculate covering filters for partitions with step" >> {
+      foreach(Seq(ChronoUnit.HOURS, ChronoUnit.DAYS, ChronoUnit.WEEKS, ChronoUnit.MONTHS, ChronoUnit.YEARS)) { unit =>
+        val ps = DateTimeScheme("dtg", 0, unit, step = 2)
+        val partition = ps.getPartition(sf)
+        val covering = ps.getCoveringFilter(partition)
+        val expected = {
+          val start = truncate(date, unit, 2)
+          val end = start.plus(2, unit)
+          ECQL.toFilter(s"dtg >= '${DateParsing.format(start)}' AND dtg < '${DateParsing.format(end)}'")
+        }
+        decomposeAnd(covering) must containTheSameElementsAs(decomposeAnd(expected))
+      }
+    }
+
     "calculate intersecting partitions for filters" >> {
       foreach(Seq(ChronoUnit.HOURS, ChronoUnit.DAYS, ChronoUnit.WEEKS, ChronoUnit.MONTHS, ChronoUnit.YEARS)) { unit =>
         val ps = DateTimeScheme("dtg", 0, unit)
@@ -96,8 +142,21 @@ class DateTimeSchemeTest extends Specification {
       }
     }
 
+    "calculate intersecting partitions for filters with step" >> {
+      foreach(Seq(ChronoUnit.HOURS, ChronoUnit.DAYS, ChronoUnit.WEEKS, ChronoUnit.MONTHS, ChronoUnit.YEARS)) { unit =>
+        val ps = DateTimeScheme("dtg", 0, unit, step = 2)
+        val partition = ps.getPartition(sf)
+        val expectedEndPartition = java.lang.Long.toHexString(java.lang.Long.parseLong(partition.value, 16) + 1)
+        val start = truncate(date, unit, 2)
+        val end = start.plus(2, unit)
+        val filter = ECQL.toFilter(s"dtg >= '${DateParsing.format(start)}' AND dtg < '${DateParsing.format(end)}'")
+        val partitions = ps.getRangesForFilter(filter).orNull
+        partitions must not(beNull)
+        partitions mustEqual Seq(PartitionRange(ps.name, partition.value, expectedEndPartition))
+      }
+    }
+
     "handle edge boundaries" >> {
-      // note: these will change when we fix simplified filters
       val dtScheme = DateTimeScheme("dtg", 0, ChronoUnit.DAYS)
       val startpoint = dtScheme.getPartition(ScalaSimpleFeature.create(sft, "1", "2017-01-02T00:00:00.000Z", "POINT (10 10)"))
       val endpoint = dtScheme.getPartition(ScalaSimpleFeature.create(sft, "1", "2017-01-04T00:00:00.000Z", "POINT (10 10)"))
@@ -108,7 +167,26 @@ class DateTimeSchemeTest extends Specification {
       twoDays.get.head.contains(startpoint.value) must beTrue
       twoDays.get.head.contains(endpoint.value) must beFalse
 
-      val inclusive = ECQL.toFilter("dtg >= '2017-01-02T00:00:00.000Z' and dtg <= '2017-01-04T00:00:00.001Z'")
+      val inclusive = ECQL.toFilter("dtg >= '2017-01-02T00:00:00.000Z' and dtg <= '2017-01-04T00:00:00.000Z'")
+      val threeDays = dtScheme.getRangesForFilter(inclusive)
+      threeDays must beSome
+      threeDays.get must haveSize(1)
+      threeDays.get.head.contains(startpoint.value) must beTrue
+      threeDays.get.head.contains(endpoint.value) must beTrue
+    }
+
+    "handle edge boundaries with step" >> {
+      val dtScheme = DateTimeScheme("dtg", 0, ChronoUnit.DAYS, step = 2)
+      val startpoint = dtScheme.getPartition(ScalaSimpleFeature.create(sft, "1", "2017-01-02T00:00:00.000Z", "POINT (10 10)"))
+      val endpoint = dtScheme.getPartition(ScalaSimpleFeature.create(sft, "1", "2017-01-04T00:00:00.000Z", "POINT (10 10)"))
+      val exclusive = ECQL.toFilter("dtg > '2017-01-02T00:00:00.000Z' and dtg < '2017-01-04T00:00:00.000Z'")
+      val twoDays = dtScheme.getRangesForFilter(exclusive)
+      twoDays must beSome
+      twoDays.get must haveSize(1)
+      twoDays.get.head.contains(startpoint.value) must beTrue
+      twoDays.get.head.contains(endpoint.value) must beFalse
+
+      val inclusive = ECQL.toFilter("dtg >= '2017-01-02T00:00:00.000Z' and dtg <= '2017-01-04T00:00:00.000Z'")
       val threeDays = dtScheme.getRangesForFilter(inclusive)
       threeDays must beSome
       threeDays.get must haveSize(1)
@@ -117,10 +195,13 @@ class DateTimeSchemeTest extends Specification {
     }
   }
 
-  private def truncate(date: ZonedDateTime, unit: ChronoUnit): ZonedDateTime = {
-    unit match {
+  private def truncate(date: ZonedDateTime, unit: ChronoUnit, step: Int = 1): ZonedDateTime = {
+    val base = unit match {
       case ChronoUnit.HOURS | ChronoUnit.DAYS => date.truncatedTo(unit)
       case _ => epoch.plus(unit.between(epoch, date), unit)
     }
+    if (step == 1) { base } else {
+      base.minus(unit.between(epoch, base).toInt % step, unit)
+    }
   }
 }