Use multi-get; native bucket scanner in segment stats computation

CHANGELOG.unreleased.md

@@ -15,6 +15,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 - The opacity of meshes can be adjusted using the 'Change Segment Color' context menu entry in the segments tab. [#8443](https://github.com/scalableminds/webknossos/pull/8443)
 - The maximum available storage of an organization is now enforced during upload. [#8385](https://github.com/scalableminds/webknossos/pull/8385)
 - Performance improvements for volume annotation save requests. [#8460](https://github.com/scalableminds/webknossos/pull/8460)
+- Performance improvements for segment statistics (volume + bounding box in context menu). [#8469](https://github.com/scalableminds/webknossos/pull/8469)
 
 ### Changed
 

util/src/main/scala/com/scalableminds/util/tools/Fox.scala

@@ -88,7 +88,7 @@ object Fox extends FoxImplicits {
   }
 
   // run serially, return individual results in list of box
-  def serialSequenceBox[A, B](l: List[A])(f: A => Fox[B])(implicit ec: ExecutionContext): Future[List[Box[B]]] = {
+  def serialSequenceBox[A, B](l: Seq[A])(f: A => Fox[B])(implicit ec: ExecutionContext): Future[List[Box[B]]] = {
     def runNext(remaining: List[A], results: List[Box[B]]): Future[List[Box[B]]] =
       remaining match {
         case head :: tail =>
@@ -99,7 +99,7 @@ object Fox extends FoxImplicits {
         case Nil =>
           Future.successful(results.reverse)
       }
-    runNext(l, Nil)
+    runNext(l.toList, Nil)
   }
 
   def sequence[T](l: List[Fox[T]])(implicit ec: ExecutionContext): Future[List[Box[T]]] =

webknossos-datastore/app/com/scalableminds/webknossos/datastore/DataStoreModule.scala

@@ -25,7 +25,7 @@ class DataStoreModule extends AbstractModule {
     bind(classOf[AgglomerateService]).asEagerSingleton()
     bind(classOf[AdHocMeshServiceHolder]).asEagerSingleton()
     bind(classOf[ApplicationHealthService]).asEagerSingleton()
-    bind(classOf[DatasetErrorLoggingService]).asEagerSingleton()
+    bind(classOf[DSDatasetErrorLoggingService]).asEagerSingleton()
     bind(classOf[MeshFileService]).asEagerSingleton()
   }
 }

webknossos-datastore/app/com/scalableminds/webknossos/datastore/controllers/DataSourceController.scala

@@ -43,7 +43,7 @@ class DataSourceController @Inject()(
     connectomeFileService: ConnectomeFileService,
     segmentIndexFileService: SegmentIndexFileService,
     storageUsageService: DSUsedStorageService,
-    datasetErrorLoggingService: DatasetErrorLoggingService,
+    datasetErrorLoggingService: DSDatasetErrorLoggingService,
     exploreRemoteLayerService: ExploreRemoteLayerService,
     uploadService: UploadService,
     composeService: ComposeService,

webknossos-datastore/app/com/scalableminds/webknossos/datastore/dataformats/BucketProvider.scala

@@ -3,9 +3,16 @@ package com.scalableminds.webknossos.datastore.dataformats
 import com.scalableminds.util.accesscontext.TokenContext
 import com.scalableminds.util.tools.Fox
 import com.scalableminds.webknossos.datastore.models.requests.DataReadInstruction
+import net.liftweb.common.Box
 
 import scala.concurrent.ExecutionContext
 
 trait BucketProvider {
   def load(readInstruction: DataReadInstruction)(implicit ec: ExecutionContext, tc: TokenContext): Fox[Array[Byte]]
+
+  def loadMultiple(readInstructions: Seq[DataReadInstruction])(implicit ec: ExecutionContext,
+                                                               tc: TokenContext): Fox[Seq[Box[Array[Byte]]]] =
+    Fox.serialSequenceBox(readInstructions) { readInstruction =>
+      load(readInstruction)
+    }
 }

webknossos-datastore/app/com/scalableminds/webknossos/datastore/helpers/NativeBucketScanner.scala

@@ -8,4 +8,24 @@ class NativeBucketScanner() {
                                 bytesPerElement: Int,
                                 isSigned: Boolean,
                                 skipZeroes: Boolean): Array[Long]
+
+  @native def countSegmentVoxels(bucketBytes: Array[Byte],
+                                 bytesPerElement: Int,
+                                 isSigned: Boolean,
+                                 segmentId: Long): Long
+
+  @native def extendSegmentBoundingBox(bucketBytes: Array[Byte],
+                                       bytesPerElement: Int,
+                                       isSigned: Boolean,
+                                       bucketLength: Int,
+                                       segmentId: Long,
+                                       bucketTopLeftX: Int,
+                                       bucketTopLeftY: Int,
+                                       bucketTopLeftZ: Int,
+                                       existingBBoxTopLeftX: Int,
+                                       existingBBoxTopLeftY: Int,
+                                       existingBBoxTopLeftZ: Int,
+                                       existingBBoxBottomRightX: Int,
+                                       existingBBoxBottomRightY: Int,
+                                       existingBBoxBottomRightZ: Int): Array[Int]
 }

webknossos-datastore/app/com/scalableminds/webknossos/datastore/helpers/SegmentStatistics.scala

@@ -3,8 +3,10 @@ package com.scalableminds.webknossos.datastore.helpers
 import com.scalableminds.util.geometry.{BoundingBox, Vec3Int}
 import com.scalableminds.util.tools.{Fox, FoxImplicits}
 import com.scalableminds.webknossos.datastore.geometry.Vec3IntProto
-import com.scalableminds.webknossos.datastore.models.datasource.DataLayer
-import com.scalableminds.webknossos.datastore.models.{AdditionalCoordinate, SegmentInteger}
+import com.scalableminds.webknossos.datastore.models.datasource.{DataLayer, ElementClass}
+import com.scalableminds.webknossos.datastore.models.AdditionalCoordinate
+import net.liftweb.common.{Box, Empty, Failure, Full}
+import net.liftweb.common.Box.tryo
 import play.api.libs.json.{Json, OFormat}
 
 import scala.concurrent.ExecutionContext
@@ -19,35 +21,44 @@ object SegmentStatisticsParameters {
 
 trait SegmentStatistics extends ProtoGeometryImplicits with FoxImplicits {
 
-  def calculateSegmentVolume(segmentId: Long,
-                             mag: Vec3Int,
-                             additionalCoordinates: Option[Seq[AdditionalCoordinate]],
-                             getBucketPositions: (Long, Vec3Int) => Fox[Set[Vec3IntProto]],
-                             getTypedDataForBucketPosition: (
-                                 Vec3Int,
-                                 Vec3Int,
-                                 Option[Seq[AdditionalCoordinate]]) => Fox[Array[SegmentInteger]])(
+  protected def bucketScanner: NativeBucketScanner
+
+  def calculateSegmentVolume(
+      segmentId: Long,
+      mag: Vec3Int,
+      additionalCoordinates: Option[Seq[AdditionalCoordinate]],
+      getBucketPositions: (Long, Vec3Int) => Fox[Set[Vec3IntProto]],
+      getDataForBucketPositions: (
+          Seq[Vec3Int],
+          Vec3Int,
+          Option[Seq[AdditionalCoordinate]]) => Fox[(Seq[Box[Array[Byte]]], ElementClass.Value)])(
       implicit ec: ExecutionContext): Fox[Long] =
     for {
       bucketPositionsProtos: Set[Vec3IntProto] <- getBucketPositions(segmentId, mag)
       bucketPositionsInMag = bucketPositionsProtos.map(vec3IntFromProto)
-      volumeBoxes <- Fox.serialSequence(bucketPositionsInMag.toList)(bucketPosition =>
-        for {
-          dataTyped: Array[SegmentInteger] <- getTypedDataForBucketPosition(bucketPosition, mag, additionalCoordinates)
-          count = dataTyped.count(segmentInteger => segmentInteger.toLong == segmentId)
-        } yield count.toLong)
-      counts <- Fox.combined(volumeBoxes.map(_.toFox))
+      (bucketBoxes, elementClass) <- getDataForBucketPositions(bucketPositionsInMag.toSeq, mag, additionalCoordinates)
+      counts <- Fox.serialCombined(bucketBoxes.toList) {
+        case Full(bucketBytes) =>
+          tryo(
+            bucketScanner.countSegmentVoxels(bucketBytes,
+                                             ElementClass.bytesPerElement(elementClass),
+                                             ElementClass.isSigned(elementClass),
+                                             segmentId)).toFox
+        case Empty      => Full(0L).toFox
+        case f: Failure => f.toFox
+      }
     } yield counts.sum
 
   // Returns the bounding box in voxels in the target mag
-  def calculateSegmentBoundingBox(segmentId: Long,
-                                  mag: Vec3Int,
-                                  additionalCoordinates: Option[Seq[AdditionalCoordinate]],
-                                  getBucketPositions: (Long, Vec3Int) => Fox[Set[Vec3IntProto]],
-                                  getTypedDataForBucketPosition: (
-                                      Vec3Int,
-                                      Vec3Int,
-                                      Option[Seq[AdditionalCoordinate]]) => Fox[Array[SegmentInteger]])(
+  def calculateSegmentBoundingBox(
+      segmentId: Long,
+      mag: Vec3Int,
+      additionalCoordinates: Option[Seq[AdditionalCoordinate]],
+      getBucketPositions: (Long, Vec3Int) => Fox[Set[Vec3IntProto]],
+      getDataForBucketPositions: (
+          Seq[Vec3Int],
+          Vec3Int,
+          Option[Seq[AdditionalCoordinate]]) => Fox[(Seq[Box[Array[Byte]]], ElementClass.Value)])(
       implicit ec: ExecutionContext): Fox[BoundingBox] =
     for {
       allBucketPositions: Set[Vec3IntProto] <- getBucketPositions(segmentId, mag)
@@ -58,14 +69,16 @@ trait SegmentStatistics extends ProtoGeometryImplicits with FoxImplicits {
                                                                     Int.MinValue,
                                                                     Int.MinValue,
                                                                     Int.MinValue) //topleft, bottomright
-      _ <- Fox.serialCombined(relevantBucketPositions.iterator)(
-        bucketPosition =>
-          extendBoundingBoxByData(mag,
-                                  segmentId,
-                                  boundingBoxMutable,
-                                  bucketPosition,
-                                  additionalCoordinates,
-                                  getTypedDataForBucketPosition))
+      (bucketBoxes, elementClass) <- getDataForBucketPositions(relevantBucketPositions.toSeq,
+                                                               mag,
+                                                               additionalCoordinates)
+      _ <- Fox.serialCombined(relevantBucketPositions.zip(bucketBoxes)) {
+        case (bucketPosition, Full(bucketData)) =>
+          Fox.successful(
+            extendBoundingBoxByData(segmentId, boundingBoxMutable, bucketPosition, bucketData, elementClass))
+        case (_, Empty)      => Fox.successful(())
+        case (_, f: Failure) => f.toFox
+      }
     } yield
       if (boundingBoxMutable.exists(item => item == Int.MaxValue || item == Int.MinValue)) {
         BoundingBox.empty
@@ -93,44 +106,30 @@ trait SegmentStatistics extends ProtoGeometryImplicits with FoxImplicits {
         .map(vec3IntFromProto)
     }
 
-  private def extendBoundingBoxByData(mag: Vec3Int,
-                                      segmentId: Long,
-                                      mutableBoundingBox: scala.collection.mutable.ListBuffer[Int],
+  private def extendBoundingBoxByData(segmentId: Long,
+                                      boundingBoxMutable: scala.collection.mutable.ListBuffer[Int],
                                       bucketPosition: Vec3Int,
-                                      additionalCoordinates: Option[Seq[AdditionalCoordinate]],
-                                      getTypedDataForBucketPosition: (
-                                          Vec3Int,
-                                          Vec3Int,
-                                          Option[Seq[AdditionalCoordinate]]) => Fox[Array[SegmentInteger]]): Fox[Unit] =
-    for {
-      dataTyped: Array[SegmentInteger] <- getTypedDataForBucketPosition(bucketPosition, mag, additionalCoordinates)
-      bucketTopLeftInTargetMagVoxels = bucketPosition * DataLayer.bucketLength
-      _ = scanDataAndExtendBoundingBox(dataTyped, bucketTopLeftInTargetMagVoxels, segmentId, mutableBoundingBox)
-    } yield ()
-
-  private def scanDataAndExtendBoundingBox(dataTyped: Array[SegmentInteger],
-                                           bucketTopLeftInTargetMagVoxels: Vec3Int,
-                                           segmentId: Long,
-                                           mutableBoundingBox: scala.collection.mutable.ListBuffer[Int]): Unit =
-    for {
-      x <- 0 until DataLayer.bucketLength
-      y <- 0 until DataLayer.bucketLength
-      z <- 0 until DataLayer.bucketLength
-      index = z * DataLayer.bucketLength * DataLayer.bucketLength + y * DataLayer.bucketLength + x
-    } yield {
-      if (dataTyped(index).toLong == segmentId) {
-        val voxelPosition = bucketTopLeftInTargetMagVoxels + Vec3Int(x, y, z)
-        extendBoundingBoxByPosition(mutableBoundingBox, voxelPosition)
-      }
-    }
-
-  private def extendBoundingBoxByPosition(mutableBoundingBox: scala.collection.mutable.ListBuffer[Int],
-                                          position: Vec3Int): Unit = {
-    mutableBoundingBox(0) = Math.min(mutableBoundingBox(0), position.x)
-    mutableBoundingBox(1) = Math.min(mutableBoundingBox(1), position.y)
-    mutableBoundingBox(2) = Math.min(mutableBoundingBox(2), position.z)
-    mutableBoundingBox(3) = Math.max(mutableBoundingBox(3), position.x)
-    mutableBoundingBox(4) = Math.max(mutableBoundingBox(4), position.y)
-    mutableBoundingBox(5) = Math.max(mutableBoundingBox(5), position.z)
+                                      bucketBytes: Array[Byte],
+                                      elementClass: ElementClass.Value): Unit = {
+    val bucketTopLeftInTargetMagVoxels = bucketPosition * DataLayer.bucketLength
+    val extendedBBArray: Array[Int] =
+      bucketScanner.extendSegmentBoundingBox(
+        bucketBytes,
+        ElementClass.bytesPerElement(elementClass),
+        ElementClass.isSigned(elementClass),
+        DataLayer.bucketLength,
+        segmentId,
+        bucketTopLeftInTargetMagVoxels.x,
+        bucketTopLeftInTargetMagVoxels.y,
+        bucketTopLeftInTargetMagVoxels.z,
+        boundingBoxMutable(0),
+        boundingBoxMutable(1),
+        boundingBoxMutable(2),
+        boundingBoxMutable(3),
+        boundingBoxMutable(4),
+        boundingBoxMutable(5)
+      )
+    (0 until 6).foreach(i => boundingBoxMutable(i) = extendedBBArray(i))
   }
+
 }

webknossos-datastore/app/com/scalableminds/webknossos/datastore/models/requests/DataServiceRequests.scala

@@ -19,7 +19,9 @@ case class DataServiceDataRequest(
     dataLayer: DataLayer,
     cuboid: Cuboid,
     settings: DataServiceRequestSettings
-)
+) {
+  def isSingleBucket: Boolean = cuboid.isSingleBucket(DataLayer.bucketLength)
+}
 
 case class DataReadInstruction(
     baseDir: Path,