Implement viewing sharded neuroglancer precomputed datasets (#6920)

CHANGELOG.unreleased.md

@@ -14,6 +14,7 @@ For upgrade instructions, please check the [migration guide](MIGRATIONS.released
 - Added support for datasets where layers are transformed individually (with an affine matrix). Transformations can be specified via datasource-properties.json or via JS API (will be ephemeral, then). [#6748](https://github.com/scalableminds/webknossos/pull/6748)
 - Added list of all respective team members to the administration page for teams. [#6915](https://github.com/scalableminds/webknossos/pull/6915)
 - Added email notifications for WK worker jobs. [#6918](https://github.com/scalableminds/webknossos/pull/6918)
+- Added support for viewing sharded neuroglancer precomputed datasets. [#6920](https://github.com/scalableminds/webknossos/pull/6920)
 
 ### Changed
 - Interpolation during rendering is now more performance intensive, since the rendering approach was changed. Therefore, interpolation is disabled by default. On the flip side, the rendered quality is often higher than it used to be. [#6748](https://github.com/scalableminds/webknossos/pull/6748)

MIGRATIONS.released.md

@@ -1,5 +1,6 @@
 # Migration Guide (Released)
-All migrations of WEBKNOSOSS are documented in this file.
+
+All migrations of WEBKNOSSOS are documented in this file.
 See `MIGRATIONS.unreleased.md` for the changes which are not yet part of an official release.
 
 This project adheres to [Calendar Versioning](http://calver.org/) `0Y.0M.MICRO`.

app/models/binary/explore/PrecomputedExplorer.scala

@@ -30,7 +30,6 @@ class PrecomputedExplorer extends RemoteLayerExplorer {
     for {
       name <- guessNameFromPath(remotePath)
       firstScale <- precomputedHeader.scales.headOption.toFox
-      _ <- bool2Fox(firstScale.sharding.isEmpty) ?~> "Failed to read dataset: sharding not supported"
       boundingBox <- BoundingBox.fromSizeArray(firstScale.size).toFox
       elementClass: ElementClass.Value <- elementClassFromPrecomputedDataType(precomputedHeader.data_type) ?~> "Unknown data type"
       smallestResolution = firstScale.resolution

test/backend/CompressedMortonCodeTestSuite.scala

@@ -0,0 +1,37 @@
+package backend
+
+import com.scalableminds.webknossos.datastore.datareaders.precomputed.CompressedMortonCode
+import org.scalatestplus.play.PlaySpec
+
+class CompressedMortonCodeTestSuite extends PlaySpec {
+
+  "Compressed Morton Code" when {
+    "Grid size = 10,10,10" should {
+      val grid_size = Array(10, 10, 10)
+      "encode 0,0,0" in {
+        assert(CompressedMortonCode.encode(Array(0, 0, 0), grid_size) == 0)
+      }
+      "encode 1,2,3" in {
+        assert(CompressedMortonCode.encode(Array(1, 2, 3), grid_size) == 53)
+      }
+      "encode 9,9,9" in {
+        assert(CompressedMortonCode.encode(Array(9, 9, 9), grid_size) == 3591)
+      }
+      "encode 10,10,10" in {
+        assert(CompressedMortonCode.encode(Array(10, 10, 10), grid_size) == 3640)
+      }
+    }
+    "Grid size = 2048,2048,1024" should {
+      val grid_size = Array(2048, 2048, 1024)
+      "encode 0,0,0" in {
+        assert(CompressedMortonCode.encode(Array(0, 0, 0), grid_size) == 0)
+      }
+      "encode 1,2,3" in {
+        assert(CompressedMortonCode.encode(Array(1, 2, 3), grid_size) == 53)
+      }
+      "encode 1024, 512, 684" in {
+        assert(CompressedMortonCode.encode(Array(1024, 512, 684), grid_size) == 1887570176)
+      }
+    }
+  }
+}

util/src/main/scala/com/scalableminds/util/geometry/Vec3Int.scala

@@ -63,6 +63,9 @@ case class Vec3Int(x: Int, y: Int, z: Int) {
     } yield Vec3Int(x, y, z)
 
   def product: Int = x * y * z
+
+  def alignWithGridFloor(gridCellSize: Vec3Int): Vec3Int =
+    this / gridCellSize * gridCellSize
 }
 
 object Vec3Int {

webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/ChunkReader.scala

@@ -6,6 +6,7 @@ import ucar.ma2.{Array => MultiArray, DataType => MADataType}
 
 import java.io.{ByteArrayInputStream, ByteArrayOutputStream, IOException}
 import javax.imageio.stream.MemoryCacheImageInputStream
+import scala.collection.immutable.NumericRange
 import scala.concurrent.Future
 import scala.util.Using
 
@@ -28,16 +29,17 @@ class ChunkReader(val header: DatasetHeader, val vaultPath: VaultPath, val chunk
   lazy val chunkSize: Int = header.chunkSize.toList.product
 
   @throws[IOException]
-  def read(path: String, chunkShape: Array[Int]): Future[MultiArray] = {
-    val chunkBytesAndShape = readChunkBytesAndShape(path)
+  def read(path: String, chunkShape: Array[Int], range: Option[NumericRange[Long]]): Future[MultiArray] = {
+    val chunkBytesAndShape = readChunkBytesAndShape(path, range)
     chunkTyper.wrapAndType(chunkBytesAndShape.map(_._1), chunkBytesAndShape.flatMap(_._2).getOrElse(chunkShape))
   }
 
   // Returns bytes (optional, None may later be replaced with fill value)
   // and chunk shape (optional, only for data formats where each chunk reports its own shape, e.g. N5)
-  protected def readChunkBytesAndShape(path: String): Option[(Array[Byte], Option[Array[Int]])] =
+  protected def readChunkBytesAndShape(path: String,
+                                       range: Option[NumericRange[Long]]): Option[(Array[Byte], Option[Array[Int]])] =
     Using.Manager { use =>
-      (vaultPath / path).readBytes().map { bytes =>
+      (vaultPath / path).readBytes(range).map { bytes =>
         val is = use(new ByteArrayInputStream(bytes))
         val os = use(new ByteArrayOutputStream())
         header.compressorImpl.uncompress(is, os)

webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/ChunkUtils.scala

@@ -23,7 +23,7 @@ object ChunkUtils extends LazyLogging {
     }
     if (numChunks < 0) {
       logger.warn(
-        s"Failed to compute zarr chunk indices. array shape ${arrayShape.toList}, chunkShape: ${arrayChunkSize.toList}, requested ${selectedShape.toList} at ${selectedOffset.toList}")
+        s"Failed to compute chunk indices. array shape ${arrayShape.toList}, chunkShape: ${arrayChunkSize.toList}, requested ${selectedShape.toList} at ${selectedOffset.toList}")
     }
     val chunkIndices = new Array[Array[Int]](numChunks)
     val currentIdx = start.clone

webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/DatasetArray.scala

@@ -13,6 +13,7 @@ import ucar.ma2.{InvalidRangeException, Array => MultiArray}
 import java.io.IOException
 import java.nio.ByteOrder
 import java.util
+import scala.collection.immutable.NumericRange
 import scala.concurrent.{ExecutionContext, Future}
 
 class DatasetArray(relativePath: DatasetPath,
@@ -73,7 +74,7 @@ class DatasetArray(relativePath: DatasetPath,
       val targetBuffer = MultiArrayUtils.createDataBuffer(header.resolvedDataType, shape)
       val targetInCOrder: MultiArray =
         MultiArrayUtils.orderFlippedView(MultiArrayUtils.createArrayWithGivenStorage(targetBuffer, shape.reverse))
-      val wasCopiedFox = Fox.serialCombined(chunkIndices) { chunkIndex: Array[Int] =>
+      val copiedFuture = Future.sequence(chunkIndices.map { chunkIndex: Array[Int] =>
         for {
           sourceChunk: MultiArray <- getSourceChunkDataWithCache(axisOrder.permuteIndices(chunkIndex))
           offsetInChunk = computeOffsetInChunk(chunkIndex, offset)
@@ -82,21 +83,33 @@ class DatasetArray(relativePath: DatasetPath,
                                                                              flip = header.order != ArrayOrder.C)
           _ = MultiArrayUtils.copyRange(offsetInChunk, sourceChunkInCOrder, targetInCOrder)
         } yield ()
-      }
+      })
       for {
-        _ <- wasCopiedFox
+        _ <- copiedFuture
       } yield targetBuffer
     }
   }
 
-  private def getSourceChunkDataWithCache(chunkIndex: Array[Int]): Future[MultiArray] = {
-    val chunkFilename = getChunkFilename(chunkIndex)
-    val chunkFilePath = relativePath.resolve(chunkFilename)
-    val storeKey = chunkFilePath.storeKey
-    val chunkShape = header.chunkSizeAtIndex(chunkIndex)
+  protected def getShardedChunkPathAndRange(chunkIndex: Array[Int])(
+      implicit ec: ExecutionContext): Future[(VaultPath, NumericRange[Long])] = ???
 
-    chunkContentsCache.getOrLoad(storeKey, key => chunkReader.read(key, chunkShape))
-  }
+  private def getSourceChunkDataWithCache(chunkIndex: Array[Int])(implicit ec: ExecutionContext): Future[MultiArray] =
+    chunkContentsCache.getOrLoad(chunkIndex.mkString(","), _ => readSourceChunkData(chunkIndex))
+
+  private def readSourceChunkData(chunkIndex: Array[Int])(implicit ec: ExecutionContext): Future[MultiArray] =
+    if (header.isSharded) {
+      for {
+        (shardPath, chunkRange) <- getShardedChunkPathAndRange(chunkIndex)
+        chunkShape = header.chunkSizeAtIndex(chunkIndex)
+        multiArray <- chunkReader.read(shardPath.toString, chunkShape, Some(chunkRange))
+      } yield multiArray
+    } else {
+      val chunkFilename = getChunkFilename(chunkIndex)
+      val chunkFilePath = relativePath.resolve(chunkFilename)
+      val storeKey = chunkFilePath.storeKey
+      val chunkShape = header.chunkSizeAtIndex(chunkIndex)
+      chunkReader.read(storeKey, chunkShape, None)
+    }
 
   protected def getChunkFilename(chunkIndex: Array[Int]): String =
     chunkIndex.mkString(header.dimension_separator.toString)

webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/DatasetHeader.scala

@@ -45,4 +45,6 @@ trait DatasetHeader {
   lazy val rank: Int = datasetShape.length
 
   def chunkSizeAtIndex(chunkIndex: Array[Int]): Array[Int] = chunkSize
+
+  def isSharded = false
 }

webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/n5/N5ChunkReader.scala

@@ -5,6 +5,7 @@ import com.scalableminds.webknossos.datastore.datavault.VaultPath
 import com.typesafe.scalalogging.LazyLogging
 
 import java.io.{ByteArrayInputStream, ByteArrayOutputStream}
+import scala.collection.immutable.NumericRange
 import scala.util.Using
 
 object N5ChunkReader {
@@ -23,7 +24,9 @@ class N5ChunkReader(header: DatasetHeader, vaultPath: VaultPath, typedChunkReade
 
   val dataExtractor: N5DataExtractor = new N5DataExtractor
 
-  override protected def readChunkBytesAndShape(path: String): Option[(Array[Byte], Option[Array[Int]])] =
+  override protected def readChunkBytesAndShape(
+      path: String,
+      range: Option[NumericRange[Long]]): Option[(Array[Byte], Option[Array[Int]])] =
     Using.Manager { use =>
       def processBytes(bytes: Array[Byte], expectedElementCount: Int): Array[Byte] = {
         val is = use(new ByteArrayInputStream(bytes))
@@ -37,7 +40,7 @@ class N5ChunkReader(header: DatasetHeader, vaultPath: VaultPath, typedChunkReade
       }
 
       for {
-        bytes <- (vaultPath / path).readBytes()
+        bytes <- (vaultPath / path).readBytes(range)
         (blockHeader, data) = dataExtractor.readBytesAndHeader(bytes)
         paddedChunkBytes = processBytes(data, blockHeader.blockSize.product)
       } yield (paddedChunkBytes, Some(blockHeader.blockSize))

webknossos-datastore/app/com/scalableminds/webknossos/datastore/datareaders/precomputed/CompressedMortonCode.scala

@@ -0,0 +1,39 @@
+package com.scalableminds.webknossos.datastore.datareaders.precomputed
+
+import scala.math.log10
+
+object CompressedMortonCode {
+
+  def log2(x: Double): Double = log10(x) / log10(2.0)
+
+  def encode(position: Array[Int], gridSize: Array[Int]): Long = {
+    /*
+        Computes the compressed morton code as per
+        https://github.com/google/neuroglancer/blob/master/src/neuroglancer/datasource/precomputed/volume.md#compressed-morton-code
+        https://github.com/google/neuroglancer/blob/162b698f703c86e0b3e92b8d8e0cacb0d3b098df/src/neuroglancer/util/zorder.ts#L72
+     */
+    val bits = gridSize.map(log2(_).ceil.toInt)
+    val maxBits = bits.max
+    var outputBit = 0L
+    val one = 1L
+
+    var output = 0L
+    for (bit <- 0 to maxBits) {
+      if (bit < bits(0)) {
+        output |= (((position(0) >> bit) & one) << outputBit)
+        outputBit += 1
+      }
+      if (bit < bits(1)) {
+        output |= (((position(1) >> bit) & one) << outputBit)
+        outputBit += 1
+      }
+      if (bit < bits(2)) {
+        output |= (((position(2) >> bit) & one) << outputBit)
+        outputBit += 1
+      }
+    }
+
+    output
+  }
+
+}