netcdf: track spatial keys (#421)

jdries · web-flow · commit 95aa2b1282e4 · 2025-04-17T18:09:49.000+02:00
* merge_cubes: support cube partitioned by tile, as provided by netcdf collection Open-EO/openeo-geopyspark-driver#1146 #43 * consistently implement spatial key provider Open-EO/openeo-geopyspark-driver#1146 #43 * consistently implement spatial key provider Open-EO/openeo-geopyspark-driver#1146 #43 * more testing of spatial key provider Open-EO/openeo-geopyspark-driver#1146 #43 * more testing of partitioner Open-EO/openeo-geopyspark-driver#1146 #43 * more testing of partitioner Open-EO/openeo-geopyspark-driver#1146 #43 * more testing of partitioner Open-EO/openeo-geopyspark-driver#1146 #43 * fix netcdf sample geometry lookup Open-EO/openeo-geopyspark-driver#1146
diff --git a/geotrellis-common/src/main/scala/org/openeo/geotrelliscommon/package.scala b/geotrellis-common/src/main/scala/org/openeo/geotrelliscommon/package.scala
@@ -46,19 +46,23 @@ package object geotrelliscommon {
 
   }
 
-  object ByTileSpacetimePartitioner extends PartitionerIndex[SpaceTimeKey] {
+  class ByTileSpacetimePartitioner(val theKeys: Option[Array[SpatialKey]] = Option.empty) extends PartitionerIndex[SpaceTimeKey] with SpatialKeysProvider {
     private def toZ(key: SpaceTimeKey): Z2 = Z2(key.col, key.row)
 
     def toIndex(key: SpaceTimeKey): BigInt = toZ(key).z
 
     def indexRanges(keyRange: (SpaceTimeKey, SpaceTimeKey)): Seq[(BigInt, BigInt)] =
       Z2.zranges(ZRange(toZ(keyRange._1), toZ(keyRange._2))).map(r => (BigInt(r.lower), BigInt(r.upper)))
 
+    override def spatialKeys: Option[Array[SpatialKey]] = {
+      theKeys
+    }
   }
 
   object SparseSpaceOnlyPartitioner {
     // Shift by 8 removes the last 8 bytes: 256 tiles max in one partition.
     def toIndex(key: SpaceTimeKey, indexReduction:Int = 8): BigInt = Z2(key.col,key.row).z >> indexReduction
+    def toIndex(key: SpatialKey, indexReduction:Int): BigInt = Z2(key.col,key.row).z >> indexReduction
   }
 
   object SparseSpaceTimePartitioner {
@@ -68,7 +72,12 @@ package object geotrelliscommon {
     def toIndex(key: SpaceTimeKey, indexReduction:Int = 8): BigInt = keyIndex.toIndex(key) >> indexReduction
   }
 
-  class SparseSpaceTimePartitioner (val indices: Array[BigInt], val indexReduction:Int = 8, val theKeys: Option[Array[SpaceTimeKey]] = Option.empty) extends PartitionerIndex[SpaceTimeKey] {
+
+  trait SpatialKeysProvider {
+    def spatialKeys: Option[Array[SpatialKey]]
+  }
+
+  class SparseSpaceTimePartitioner (val indices: Array[BigInt], val indexReduction:Int = 8, val theKeys: Option[Array[SpaceTimeKey]] = Option.empty) extends PartitionerIndex[SpaceTimeKey] with SpatialKeysProvider {
 
     def toIndex(key: SpaceTimeKey): BigInt = SparseSpaceTimePartitioner.toIndex(key, indexReduction)
 
@@ -99,9 +108,14 @@ package object geotrelliscommon {
 
 
     override def toString = s"SparseSpaceTimePartitioner ${indices.length} ${theKeys.isDefined}"
+
+    override def spatialKeys: Option[Array[SpatialKey]] = {
+      theKeys.map(_.map(_.spatialKey).distinct)
+    }
+
   }
 
-  class SparseSpaceOnlyPartitioner (val indices: Array[BigInt], val indexReduction:Int = 8, val theKeys: Option[Array[SpaceTimeKey]] = Option.empty ) extends PartitionerIndex[SpaceTimeKey] {
+  class SparseSpaceOnlyPartitioner (val indices: Array[BigInt], val indexReduction:Int = 8, val theKeys: Option[Array[SpaceTimeKey]] = Option.empty ) extends PartitionerIndex[SpaceTimeKey] with SpatialKeysProvider {
 
     def toIndex(key: SpaceTimeKey): BigInt = SparseSpaceOnlyPartitioner.toIndex(key, indexReduction)
 
@@ -123,9 +137,11 @@ package object geotrelliscommon {
       val state = Seq( indexReduction)
       state.map(_.hashCode()).foldLeft(0)((a, b) => 31 * a + b)
     }
+
+    override def spatialKeys: Option[Array[SpatialKey]] = theKeys.map(_.map(_.spatialKey).distinct)
   }
 
-  class SparseSpatialPartitioner (val indices: Array[BigInt], val indexReduction:Int = 8, val theKeys: Option[Array[SpatialKey]] = Option.empty ) extends PartitionerIndex[SpatialKey] {
+  class SparseSpatialPartitioner (val indices: Array[BigInt], val indexReduction:Int = 8, val theKeys: Option[Array[SpatialKey]] = Option.empty ) extends PartitionerIndex[SpatialKey] with SpatialKeysProvider {
 
     def toIndex(key: SpatialKey): BigInt = Z2(key.col,key.row).z >> indexReduction
 
@@ -147,6 +163,8 @@ package object geotrelliscommon {
       val state = Seq( indexReduction)
       state.map(_.hashCode()).foldLeft(0)((a, b) => 31 * a + b)
     }
+
+    override def spatialKeys: Option[Array[SpatialKey]] = theKeys
   }
 
   class ConfigurableSpatialPartitioner(val indexReduction:Int = 4) extends PartitionerIndex[SpatialKey] {
diff --git a/openeo-geotrellis/src/main/scala/org/openeo/geotrellis/OpenEOProcesses.scala b/openeo-geotrellis/src/main/scala/org/openeo/geotrellis/OpenEOProcesses.scala
@@ -27,7 +27,7 @@ import org.apache.spark.{Partitioner, SparkContext}
 import org.openeo.geotrellis.OpenEOProcessScriptBuilder.{MaxIgnoreNoData, MinIgnoreNoData, OpenEOProcess, safeConvert}
 import org.openeo.geotrellis.focal._
 import org.openeo.geotrellis.netcdf.NetCDFRDDWriter.ContextSeq
-import org.openeo.geotrelliscommon.{ByTileSpacetimePartitioner, ByTileSpatialPartitioner, ConfigurableSpaceTimePartitioner, ConfigurableSpatialPartitionerReduceZ, DatacubeSupport, FFTConvolve, OpenEORasterCube, OpenEORasterCubeMetadata, SCLConvolutionFilter, SpaceTimeByMonthPartitioner, SparseSpaceOnlyPartitioner, SparseSpaceTimePartitioner, SparseSpatialPartitioner}
+import org.openeo.geotrelliscommon.{ByTileSpacetimePartitioner, ByTileSpatialPartitioner, ConfigurableSpaceTimePartitioner, ConfigurableSpatialPartitionerReduceZ, DatacubeSupport, FFTConvolve, OpenEORasterCube, OpenEORasterCubeMetadata, SCLConvolutionFilter, SpaceTimeByMonthPartitioner, SparseSpaceOnlyPartitioner, SparseSpaceTimePartitioner, SparseSpatialPartitioner, SpatialKeysProvider}
 import org.slf4j.LoggerFactory
 
 import java.io.File
@@ -173,8 +173,6 @@ class OpenEOProcesses extends Serializable {
   }
 
     private def transformTimeDimension[KT](datacube: MultibandTileLayerRDD[SpaceTimeKey], scriptBuilder: OpenEOProcessScriptBuilder, context: util.Map[String, Any], reduce:Boolean=false): RDD[(KT, MultibandTile)] = {
-
-
       val expectedCellType = datacube.metadata.cellType
       val applyToTimeseries: Iterable[(SpaceTimeKey, MultibandTile)] => Map[KT, MultibandTile] =
         if(reduce){
@@ -191,8 +189,6 @@ class OpenEOProcesses extends Serializable {
     }
 
   private def transformTimeDimension[KT](datacube: MultibandTileLayerRDD[SpaceTimeKey],applyToTimeseries: Iterable[(SpaceTimeKey, MultibandTile)] => Map[KT, MultibandTile],  reduce:Boolean ): RDD[(KT, MultibandTile)] = {
-
-
     val index: Option[PartitionerIndex[SpaceTimeKey]] =
       if (datacube.partitioner.isDefined && datacube.partitioner.get.isInstanceOf[SpacePartitioner[SpaceTimeKey]]) {
         Some(datacube.partitioner.get.asInstanceOf[SpacePartitioner[SpaceTimeKey]].index)
@@ -201,15 +197,15 @@ class OpenEOProcesses extends Serializable {
       }
     logger.info(s"Applying callback on time dimension of cube with partitioner: ${datacube.partitioner.getOrElse("no partitioner")} - index: ${index.getOrElse("no index")} and metadata ${datacube.metadata}")
     val rdd: RDD[(SpaceTimeKey, MultibandTile)] =
-      if (index.isDefined && (index.get.isInstanceOf[SparseSpaceOnlyPartitioner] || index.get == ByTileSpacetimePartitioner )) {
+      if (index.isDefined && (index.get.isInstanceOf[SparseSpaceOnlyPartitioner] || index.get.isInstanceOf[ByTileSpacetimePartitioner] )) {
         datacube
       } else {
-        val keys: Option[Array[SpaceTimeKey]] = findPartitionerKeys(datacube)
+        val keys: Option[Array[SpatialKey]] = findPartitionerSpatialKeys(datacube)
         val spatiallyGroupingIndex =
           if(keys.isDefined){
-            new SparseSpaceOnlyPartitioner(keys.get.map(SparseSpaceOnlyPartitioner.toIndex(_, indexReduction = 0)).distinct.sorted, 0, keys)
+            new SparseSpaceOnlyPartitioner(keys.get.map(SparseSpaceOnlyPartitioner.toIndex(_, indexReduction = 0)).distinct.sorted, 0, findPartitionerKeys(datacube))
           }else{
-            ByTileSpacetimePartitioner
+            new ByTileSpacetimePartitioner()
           }
         logger.info(f"Regrouping data cube along the time dimension, with index $spatiallyGroupingIndex. Cube metadata: ${datacube.metadata}")
         val partitioner: Partitioner = new SpacePartitioner(datacube.metadata.bounds)(implicitly, implicitly, spatiallyGroupingIndex)
@@ -289,10 +285,10 @@ class OpenEOProcesses extends Serializable {
 
     val newBounds = retiled.metadata.bounds.asInstanceOf[KeyBounds[SpaceTimeKey]].toSpatial
 
-    val keys = findPartitionerKeys(datacube)
+    val keys = findPartitionerSpatialKeys(datacube)
     val spatiallyGroupingIndex =
       if(keys.isDefined){
-        val spatialKeys: Array[SpatialKey] = keys.get.map(_.spatialKey).distinct
+        val spatialKeys: Array[SpatialKey] = keys.get
         new SparseSpatialPartitioner(spatialKeys.map(ByTileSpatialPartitioner.toIndex).distinct.sorted, 0, Some(spatialKeys))
       }else{
         ByTileSpatialPartitioner
@@ -306,11 +302,11 @@ class OpenEOProcesses extends Serializable {
   private def groupOnTimeDimension(datacube: MultibandTileLayerRDD[SpaceTimeKey]) = {
     val targetBounds = datacube.metadata.bounds.asInstanceOf[KeyBounds[SpaceTimeKey]].toSpatial
 
-    val keys: Option[Array[SpaceTimeKey]] = findPartitionerKeys(datacube)
+    val keys: Option[Array[SpatialKey]] = findPartitionerSpatialKeys(datacube)
 
     val index =
       if (keys.isDefined) {
-        new SparseSpatialPartitioner(keys.get.map(SparseSpaceOnlyPartitioner.toIndex(_, indexReduction = 0)).distinct.sorted, 0, keys.map(_.map(_.spatialKey)))
+        new SparseSpatialPartitioner(keys.get.map(SparseSpaceOnlyPartitioner.toIndex(_, indexReduction = 0)).distinct.sorted, 0, keys)
       } else {
         ByTileSpatialPartitioner
       }
@@ -322,6 +318,20 @@ class OpenEOProcesses extends Serializable {
     groupedOnTime
   }
 
+  def findPartitionerSpatialKeys(datacube: MultibandTileLayerRDD[SpaceTimeKey]): Option[Array[SpatialKey]] = {
+    val keys: Option[Array[SpatialKey]] = if (datacube.partitioner.isDefined && (datacube.partitioner.get.isInstanceOf[SpacePartitioner[SpaceTimeKey]] || datacube.partitioner.get.isInstanceOf[SpacePartitioner[SpatialKey]])) {
+      val index = datacube.partitioner.get.asInstanceOf[SpacePartitioner[SpaceTimeKey]].index
+      index match {
+        case value: SpatialKeysProvider =>
+          value.spatialKeys
+        case _ =>
+          Option.empty
+      }
+    } else {
+      Option.empty
+    }
+    keys
+  }
 
   def findPartitionerKeys(datacube: MultibandTileLayerRDD[SpaceTimeKey]): Option[Array[SpaceTimeKey]] = {
     val keys: Option[Array[SpaceTimeKey]] = if (datacube.partitioner.isDefined && datacube.partitioner.get.isInstanceOf[SpacePartitioner[SpaceTimeKey]]) {
@@ -461,9 +471,9 @@ class OpenEOProcesses extends Serializable {
 
     val filteredCube = filterNegativeSpatialKeys(datacube)
 
-    val keys = findPartitionerKeys(filteredCube).map(_.filter( k => k.row >= 0 && k.col >= 0 ))
+    val keys = findPartitionerSpatialKeys(filteredCube).map(_.filter( k => k.row >= 0 && k.col >= 0 ))
     val allPossibleKeys: immutable.Seq[SpatialKey] = if(keys.isDefined) {
-      keys.get.map(_.spatialKey).distinct.toList
+      keys.get.distinct.toList
     } else{
       filteredCube.metadata.tileBounds
         .coordsIter
@@ -480,7 +490,7 @@ class OpenEOProcesses extends Serializable {
       }else{
         if (datacube.partitioner.isDefined && datacube.partitioner.get.isInstanceOf[SpacePartitioner[SpaceTimeKey]]) {
           val index = datacube.partitioner.get.asInstanceOf[SpacePartitioner[SpaceTimeKey]].index
-          if (index.isInstanceOf[SparseSpaceOnlyPartitioner] || index == ByTileSpacetimePartitioner) {
+          if (index.isInstanceOf[SparseSpaceOnlyPartitioner] || index.isInstanceOf[ByTileSpacetimePartitioner]) {
             index//a space only partitioner does not care about time, so can be reused as-is
           } else {
             SpaceTimeByMonthPartitioner
@@ -537,7 +547,7 @@ class OpenEOProcesses extends Serializable {
 
     val tilesByInterval: RDD[(SpaceTimeKey, MultibandTile)] =
       if(reduce) {
-        if(datacube.partitioner.isDefined && datacube.partitioner.get.isInstanceOf[SpacePartitioner[SpaceTimeKey]] && (datacube.partitioner.get.asInstanceOf[SpacePartitioner[SpaceTimeKey]].index.isInstanceOf[SparseSpaceOnlyPartitioner] || datacube.partitioner.get.asInstanceOf[SpacePartitioner[SpaceTimeKey]].index == ByTileSpacetimePartitioner )) {
+        if(datacube.partitioner.isDefined && datacube.partitioner.get.isInstanceOf[SpacePartitioner[SpaceTimeKey]] && (datacube.partitioner.get.asInstanceOf[SpacePartitioner[SpaceTimeKey]].index.isInstanceOf[SparseSpaceOnlyPartitioner] || datacube.partitioner.get.asInstanceOf[SpacePartitioner[SpaceTimeKey]].index.isInstanceOf[ByTileSpacetimePartitioner] )) {
           filteredCube.mapPartitions(elements =>{
             val byNewKey= elements.flatMap(mapToNewKey).toStream.groupBy(_._1)
             byNewKey.mapValues(v=>aggregateTiles(v.map(_._2))).iterator
@@ -717,7 +727,7 @@ class OpenEOProcesses extends Serializable {
         implicit val newIndex: PartitionerIndex[K] = new SparseSpaceOnlyPartitioner(newIndices,leftPart.index.asInstanceOf[SparseSpaceOnlyPartitioner].indexReduction).asInstanceOf[PartitionerIndex[K]]
         SpacePartitioner[K](kb)(implicitly,implicitly,newIndex)
       }
-      else if(leftPart.index == rightPart.index && (leftPart.index == ByTileSpatialPartitioner || leftPart.index == ByTileSpacetimePartitioner)) {
+      else if(leftPart.index == rightPart.index && (leftPart.index == ByTileSpatialPartitioner || leftPart.index.isInstanceOf[ByTileSpacetimePartitioner])) {
         leftPart
       }
       else if(leftPart.index == rightPart.index && leftPart.index.isInstanceOf[ConfigurableSpaceTimePartitioner] ) {
diff --git a/openeo-geotrellis/src/main/scala/org/openeo/geotrellis/layers/NetCDFCollection.scala b/openeo-geotrellis/src/main/scala/org/openeo/geotrellis/layers/NetCDFCollection.scala
@@ -107,7 +107,8 @@ object NetCDFCollection {
     val spatialBounds = KeyBounds(layout.mapTransform(extent))
     val temporalBounds = KeyBounds(SpaceTimeKey(spatialBounds.minKey,TemporalKey(LocalDate.of(1990,1,1).atStartOfDay(ZoneId.of("UTC")))),SpaceTimeKey(spatialBounds.maxKey,TemporalKey(LocalDate.now().atStartOfDay(ZoneId.of("UTC")))))
 
-    val partitioner: Partitioner = new SpacePartitioner(temporalBounds)(implicitly, implicitly, ByTileSpacetimePartitioner)
+    val keys: Array[SpatialKey] =  items.map(i => i.geometry.getOrElse(i.bbox.toPolygon())).clipToGrid(layout).map(_._1).distinct().collect()
+    val partitioner: Partitioner = new SpacePartitioner(temporalBounds)(implicitly, implicitly, new ByTileSpacetimePartitioner(Some(keys)))
 
     val metadata = TileLayerMetadata[SpaceTimeKey](cellType, layout, extent, crs(0), temporalBounds)
     val retiled: RDD[(SpaceTimeKey, MultibandTile)] = features.tileToLayout(metadata).partitionBy(partitioner)
diff --git a/openeo-geotrellis/src/test/scala/org/openeo/geotrellis/OpenEOProcessesSpec.scala b/openeo-geotrellis/src/test/scala/org/openeo/geotrellis/OpenEOProcessesSpec.scala
