Data loading: save mem + cpu by re-using allocated fill-value chunks (#…

…8271) * Data loading: save mem + cpu by re-using allocated fill-value chunks * changelog
scalableminds · Dec 11, 2024 · c18c5e3 · c18c5e3
1 parent 002c915
commit c18c5e3
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 4 deletions.
diff --git a/CHANGELOG.unreleased.md b/CHANGELOG.unreleased.md
@@ -12,6 +12,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 
 ### Added
 - Added the total volume of a dataset to a tooltip in the dataset info tab. [#8229](https://github.com/scalableminds/webknossos/pull/8229)
+- Optimized performance of data loading with “fill value“ chunks. [#8271](https://github.com/scalableminds/webknossos/pull/8271)
 
 ### Changed
 - Renamed "resolution" to "magnification" in more places within the codebase, including local variables. [#8168](https://github.com/scalableminds/webknossos/pull/8168)

diff --git a/...nossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/ChunkReader.scala b/...nossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/ChunkReader.scala
@@ -26,11 +26,11 @@ class ChunkReader(header: DatasetHeader) {
         case Full(chunkBytes) if useSkipTypingShortcut =>
           shortcutChunkTyper.wrapAndType(chunkBytes, chunkShape).toFox ?~> "chunk.shortcutWrapAndType.failed"
         case Empty if useSkipTypingShortcut =>
-          shortcutChunkTyper.createFromFillValue(chunkShape).toFox ?~> "chunk.shortcutCreateFromFillValue.failed"
+          shortcutChunkTyper.createFromFillValueCached(chunkShape).toFox ?~> "chunk.shortcutCreateFromFillValue.failed"
         case Full(chunkBytes) =>
           chunkTyper.wrapAndType(chunkBytes, chunkShape).toFox ?~> "chunk.wrapAndType.failed"
         case Empty =>
-          chunkTyper.createFromFillValue(chunkShape).toFox ?~> "chunk.createFromFillValue.failed"
+          chunkTyper.createFromFillValueCached(chunkShape).toFox ?~> "chunk.createFromFillValue.failed"
         case f: Failure =>
           f.toFox ?~> s"Reading chunk at $path failed"
       }

diff --git a/webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/ChunkTyper.scala b/webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/ChunkTyper.scala
@@ -1,5 +1,8 @@
 package com.scalableminds.webknossos.datastore.datareaders
 
+import com.scalableminds.util.cache.AlfuCache
+import com.scalableminds.util.tools.Fox
+import com.scalableminds.util.tools.Fox.box2Fox
 import net.liftweb.common.Box
 import net.liftweb.common.Box.tryo
 
@@ -8,6 +11,8 @@ import javax.imageio.stream.MemoryCacheImageInputStream
 import scala.util.Using
 import ucar.ma2.{Array => MultiArray, DataType => MADataType}
 
+import scala.concurrent.ExecutionContext
+
 object ChunkTyper {
   def createFromHeader(header: DatasetHeader): ChunkTyper = header.resolvedDataType match {
     case ArrayDataType.i1 | ArrayDataType.u1 => new ByteChunkTyper(header)
@@ -25,7 +30,14 @@ abstract class ChunkTyper {
   def ma2DataType: MADataType
   def wrapAndType(bytes: Array[Byte], chunkShape: Array[Int]): Box[MultiArray]
 
-  def createFromFillValue(chunkShape: Array[Int]): Box[MultiArray] =
+  // If large areas of the array use the fill value, the same chunk shape will be requested often.
+  // This cache implements a flyweight pattern by returning the same instance of the fill-valued array.
+  private val fillValueChunkCache: AlfuCache[String, MultiArray] = AlfuCache(maxCapacity = 1)
+
+  def createFromFillValueCached(chunkShape: Array[Int])(implicit ec: ExecutionContext): Fox[MultiArray] =
+    fillValueChunkCache.getOrLoad(chunkShape.mkString(","), _ => createFromFillValue(chunkShape).toFox)
+
+  protected def createFromFillValue(chunkShape: Array[Int]): Box[MultiArray] =
     MultiArrayUtils.createFilledArray(ma2DataType, chunkShapeOrdered(chunkShape), header.fillValueNumber)
 
   // Chunk shape in header is in C-Order (XYZ), but data may be in F-Order (ZYX), so the chunk shape
@@ -127,7 +139,7 @@ class ShortcutChunkTyper(val header: DatasetHeader) extends ChunkTyper {
     MultiArray.factory(ma2DataType, flatShape, bytes)
   }
 
-  override def createFromFillValue(chunkShape: Array[Int]): Box[MultiArray] = {
+  override protected def createFromFillValue(chunkShape: Array[Int]): Box[MultiArray] = {
     val flatShape = Array(chunkShape.product * header.bytesPerElement)
     MultiArrayUtils.createFilledArray(ma2DataType, flatShape, header.fillValueNumber)
   }