github.com/janelia-flyem/dvid@v1.0.0/datatype/labelarray/read.go

package labelarray

import (
	"encoding/binary"
	"fmt"
	"log"
	"sync"

	"github.com/janelia-flyem/dvid/datastore"
	"github.com/janelia-flyem/dvid/datatype/common/labels"
	"github.com/janelia-flyem/dvid/datatype/imageblk"
	"github.com/janelia-flyem/dvid/dvid"
	"github.com/janelia-flyem/dvid/server"
	"github.com/janelia-flyem/dvid/storage"
)

// ComputeTransform determines the block coordinate and beginning + ending voxel points
// for the data corresponding to the given Block.
func (v *Labels) ComputeTransform(tkey storage.TKey, blockSize dvid.Point) (blockBeg, dataBeg, dataEnd dvid.Point, err error) {
	var ptIndex *dvid.IndexZYX
	_, ptIndex, err = DecodeBlockTKey(tkey)
	if err != nil {
		return
	}

	// Get the bounding voxel coordinates for this block.
	minBlockVoxel := ptIndex.MinPoint(blockSize)
	maxBlockVoxel := ptIndex.MaxPoint(blockSize)

	// Compute the boundary voxel coordinates for the ExtData and adjust
	// to our block bounds.
	minDataVoxel := v.StartPoint()
	maxDataVoxel := v.EndPoint()
	begVolCoord, _ := minDataVoxel.Max(minBlockVoxel)
	endVolCoord, _ := maxDataVoxel.Min(maxBlockVoxel)

	// Adjust the DVID volume voxel coordinates for the data so that (0,0,0)
	// is where we expect this slice/subvolume's data to begin.
	dataBeg = begVolCoord.Sub(v.StartPoint())
	dataEnd = endVolCoord.Sub(v.StartPoint())

	// Compute block coord matching dataBeg
	blockBeg = begVolCoord.Sub(minBlockVoxel)

	return
}

// GetImage retrieves a 2d image from a version node given a geometry of labels.
func (d *Data) GetImage(v dvid.VersionID, vox *Labels, scale uint8, roiname dvid.InstanceName) (*dvid.Image, error) {
	r, err := imageblk.GetROI(v, roiname, vox)
	if err != nil {
		return nil, err
	}
	if err := d.GetLabels(v, scale, vox, r); err != nil {
		return nil, err
	}
	return vox.GetImage2d()
}

// GetVolume retrieves a n-d volume from a version node given a geometry of labels.
func (d *Data) GetVolume(v dvid.VersionID, vox *Labels, scale uint8, roiname dvid.InstanceName) ([]byte, error) {
	r, err := imageblk.GetROI(v, roiname, vox)
	if err != nil {
		return nil, err
	}
	if err := d.GetLabels(v, scale, vox, r); err != nil {
		return nil, err
	}
	return vox.Data(), nil
}

type getOperation struct {
	voxels      *Labels
	blocksInROI map[string]bool
	mapping     *labels.Mapping
}

// GetLabels copies labels from the storage engine to Labels, a requested subvolume or 2d image.
func (d *Data) GetLabels(v dvid.VersionID, scale uint8, vox *Labels, r *imageblk.ROI) error {
	store, err := datastore.GetOrderedKeyValueDB(d)
	if err != nil {
		return fmt.Errorf("Data type imageblk had error initializing store: %v\n", err)
	}
	if r != nil && scale != 0 {
		return fmt.Errorf("DVID does not currently support ROI masks on lower-scale GETs")
	}

	// Only do one request at a time, although each request can start many goroutines.
	if vox.NumVoxels() > 256*256*256 {
		server.LargeMutationMutex.Lock()
		defer server.LargeMutationMutex.Unlock()
	}

	ctx := datastore.NewVersionedCtx(d, v)

	iv := dvid.InstanceVersion{d.DataUUID(), v}
	mapping := labels.LabelMap(iv)

	wg := new(sync.WaitGroup)

	okv := store.(storage.BufferableOps)
	// extract buffer interface
	req, hasbuffer := okv.(storage.KeyValueRequester)
	if hasbuffer {
		okv = req.NewBuffer(ctx)
	}

	for it, err := vox.NewIndexIterator(d.BlockSize()); err == nil && it.Valid(); it.NextSpan() {
		indexBeg, indexEnd, err := it.IndexSpan()
		if err != nil {
			return err
		}
		begTKey := NewBlockTKey(scale, indexBeg)
		endTKey := NewBlockTKey(scale, indexEnd)

		// Get set of blocks in ROI if ROI provided
		var chunkOp *storage.ChunkOp
		if r != nil && r.Iter != nil {
			ptBeg := indexBeg.Duplicate().(dvid.ChunkIndexer)
			ptEnd := indexEnd.Duplicate().(dvid.ChunkIndexer)
			begX := ptBeg.Value(0)
			endX := ptEnd.Value(0)

			blocksInROI := make(map[string]bool, (endX - begX + 1))
			c := dvid.ChunkPoint3d{begX, ptBeg.Value(1), ptBeg.Value(2)}
			for x := begX; x <= endX; x++ {
				c[0] = x
				curIndex := dvid.IndexZYX(c)
				if r.Iter.InsideFast(curIndex) {
					indexString := string(curIndex.Bytes())
					blocksInROI[indexString] = true
				}
			}
			chunkOp = &storage.ChunkOp{&getOperation{vox, blocksInROI, mapping}, wg}
		} else {
			chunkOp = &storage.ChunkOp{&getOperation{vox, nil, mapping}, wg}
		}

		if !hasbuffer {
			// Send the entire range of key-value pairs to chunk processor
			err = okv.ProcessRange(ctx, begTKey, endTKey, chunkOp, storage.ChunkFunc(d.ReadChunk))
			if err != nil {
				return fmt.Errorf("Unable to GET data %s: %v", ctx, err)
			}
		} else {
			// Extract block list
			var tkeys []storage.TKey
			ptBeg := indexBeg.Duplicate().(dvid.ChunkIndexer)
			ptEnd := indexEnd.Duplicate().(dvid.ChunkIndexer)
			begX := ptBeg.Value(0)
			endX := ptEnd.Value(0)

			c := dvid.ChunkPoint3d{begX, ptBeg.Value(1), ptBeg.Value(2)}
			for x := begX; x <= endX; x++ {
				c[0] = x
				tk := NewBlockTKeyByCoord(scale, c.ToIZYXString())
				tkeys = append(tkeys, tk)
			}

			err = okv.(storage.RequestBuffer).ProcessList(ctx, tkeys, chunkOp, storage.ChunkFunc(d.ReadChunk))
			if err != nil {
				return fmt.Errorf("Unable to GET data %s: %v", ctx, err)
			}
		}
	}

	if hasbuffer {
		// submit the entire buffer to the DB
		err = okv.(storage.RequestBuffer).Flush()

		if err != nil {
			return fmt.Errorf("Unable to GET data %s: %v", ctx, err)
		}
	}

	if err != nil {
		return err
	}
	wg.Wait()
	return nil
}

// ReadChunk reads a chunk of data as part of a mapped operation.
// Only some multiple of the # of CPU cores can be used for chunk handling before
// it waits for chunk processing to abate via the buffered server.HandlerToken channel.
func (d *Data) ReadChunk(chunk *storage.Chunk) error {
	server.CheckChunkThrottling()
	go d.readChunk(chunk)
	return nil
}

func (d *Data) readChunk(chunk *storage.Chunk) {
	defer func() {
		// After processing a chunk, return the token.
		server.HandlerToken <- 1

		// Notify the requestor that this chunk is done.
		if chunk.Wg != nil {
			chunk.Wg.Done()
		}
	}()

	op, ok := chunk.Op.(*getOperation)
	if !ok {
		log.Fatalf("Illegal operation passed to readChunk() for data %s\n", d.DataName())
	}

	// Make sure our received chunk is valid.
	if chunk == nil {
		dvid.Errorf("Received nil chunk in readChunk.  Ignoring chunk.\n")
		return
	}
	if chunk.K == nil {
		dvid.Errorf("Received nil chunk key in readChunk.  Ignoring chunk.\n")
		return
	}

	// If there's an ROI, if outside ROI, use blank buffer.
	var zeroOut bool
	_, indexZYX, err := DecodeBlockTKey(chunk.K)
	if err != nil {
		dvid.Errorf("Error processing voxel block: %s\n", err)
		return
	}
	if op.blocksInROI != nil {
		indexString := string(indexZYX.Bytes())
		_, insideROI := op.blocksInROI[indexString]
		if !insideROI {
			zeroOut = true
		}
	}

	// Initialize the block buffer using the chunk of data.  For voxels, this chunk of
	// data needs to be uncompressed and deserialized.
	var block labels.Block
	if zeroOut || chunk.V == nil {
		blockSize, ok := d.BlockSize().(dvid.Point3d)
		if !ok {
			dvid.Errorf("Block size for data %q is not 3d: %s\n", d.DataName(), d.BlockSize())
			return
		}
		block = *labels.MakeSolidBlock(0, blockSize)
	} else {
		var data []byte
		data, _, err = dvid.DeserializeData(chunk.V, true)
		if err != nil {
			dvid.Errorf("Unable to deserialize block in %q: %v\n", d.DataName(), err)
			return
		}
		if err := block.UnmarshalBinary(data); err != nil {
			dvid.Errorf("Unable to unmarshal labels Block compression in %q: %v\n", d.DataName(), err)
			return
		}
	}

	// Perform the operation.
	if op.mapping != nil {
		err := op.voxels.readMappedBlock(chunk.K, block, d.BlockSize(), op.mapping)
		if err != nil {
			dvid.Errorf("readMappedBlock, key %v in %q: %v\n", chunk.K, d.DataName(), err)
		}
	} else {
		err := op.voxels.readBlock(chunk.K, block, d.BlockSize())
		if err != nil {
			dvid.Errorf("readBlock, key %v in %q: %v\n", chunk.K, d.DataName(), err)
		}
	}
}

func (v *Labels) readMappedBlock(tkey storage.TKey, block labels.Block, blockSize dvid.Point, m *labels.Mapping) error {
	if blockSize.NumDims() > 3 {
		return fmt.Errorf("DVID voxel blocks currently only supports up to 3d, not 4+ dimensions")
	}
	blockBeg, dataBeg, dataEnd, err := v.ComputeTransform(tkey, blockSize)
	if err != nil {
		return err
	}

	// TODO -- Refactor this function to make direct use of compressed label Block without
	// conversion into full label array.
	labels := v.Data()
	blabels, _ := block.MakeLabelVolume()

	// Compute the strides (in bytes)
	bX := int64(blockSize.Value(0)) * 8
	bY := int64(blockSize.Value(1)) * bX
	dX := int64(v.Stride())

	blockBegX := int64(blockBeg.Value(0))
	blockBegY := int64(blockBeg.Value(1))
	blockBegZ := int64(blockBeg.Value(2))

	// Do the transfers depending on shape of the external voxels.
	switch {
	case v.DataShape().Equals(dvid.XY):
		dataI := int64(dataBeg.Value(1))*dX + int64(dataBeg.Value(0))*8
		blockI := blockBegZ*bY + blockBegY*bX + blockBegX*8
		for y := int64(dataBeg.Value(1)); y <= int64(dataEnd.Value(1)); y++ {
			bI := blockI
			dI := dataI
			for x := dataBeg.Value(0); x <= dataEnd.Value(0); x++ {
				orig := binary.LittleEndian.Uint64(blabels[bI : bI+8])
				mapped, found := m.FinalLabel(orig)
				if found {
					binary.LittleEndian.PutUint64(labels[dI:dI+8], mapped)
				} else {
					copy(labels[dI:dI+8], blabels[bI:bI+8])
				}
				bI += 8
				dI += 8
			}
			blockI += bX
			dataI += dX
		}

	case v.DataShape().Equals(dvid.XZ):
		dataI := int64(dataBeg.Value(2))*dX + int64(dataBeg.Value(0))*8
		blockI := blockBegZ*bY + blockBegY*bX + blockBegX*8
		for y := int64(dataBeg.Value(2)); y <= int64(dataEnd.Value(2)); y++ {
			bI := blockI
			dI := dataI
			for x := dataBeg.Value(0); x <= dataEnd.Value(0); x++ {
				orig := binary.LittleEndian.Uint64(blabels[bI : bI+8])
				mapped, found := m.FinalLabel(orig)
				if found {
					binary.LittleEndian.PutUint64(labels[dI:dI+8], mapped)
				} else {
					copy(labels[dI:dI+8], blabels[bI:bI+8])
				}
				bI += 8
				dI += 8
			}
			blockI += bY
			dataI += dX
		}

	case v.DataShape().Equals(dvid.YZ):
		bz := blockBegZ
		for y := int64(dataBeg.Value(2)); y <= int64(dataEnd.Value(2)); y++ {
			dataI := y*dX + int64(dataBeg.Value(1))*8
			blockI := bz*bY + blockBegY*bX + blockBegX*8
			for x := dataBeg.Value(1); x <= dataEnd.Value(1); x++ {
				orig := binary.LittleEndian.Uint64(blabels[blockI : blockI+8])
				mapped, found := m.FinalLabel(orig)
				if found {
					binary.LittleEndian.PutUint64(labels[dataI:dataI+8], mapped)
				} else {
					copy(labels[dataI:dataI+8], blabels[blockI:blockI+8])
				}
				blockI += bX
				dataI += 8
			}
			bz++
		}

	case v.DataShape().ShapeDimensions() == 2:
		// TODO: General code for handling 2d ExtData in n-d space.
		return fmt.Errorf("DVID currently does not support 2d in n-d space.")

	case v.DataShape().Equals(dvid.Vol3d):
		blockOffset := blockBegX * 8
		dX := int64(v.Size().Value(0)) * 8
		dY := int64(v.Size().Value(1)) * dX
		dataOffset := int64(dataBeg.Value(0)) * 8
		blockZ := blockBegZ

		for dataZ := int64(dataBeg.Value(2)); dataZ <= int64(dataEnd.Value(2)); dataZ++ {
			blockY := blockBegY
			for dataY := int64(dataBeg.Value(1)); dataY <= int64(dataEnd.Value(1)); dataY++ {
				bI := blockZ*bY + blockY*bX + blockOffset
				dI := dataZ*dY + dataY*dX + dataOffset
				for x := dataBeg.Value(0); x <= dataEnd.Value(0); x++ {
					orig := binary.LittleEndian.Uint64(blabels[bI : bI+8])
					mapped, found := m.FinalLabel(orig)
					if found {
						binary.LittleEndian.PutUint64(labels[dI:dI+8], mapped)
					} else {
						copy(labels[dI:dI+8], blabels[bI:bI+8])
					}
					bI += 8
					dI += 8
				}
				blockY++
			}
			blockZ++
		}

	default:
		return fmt.Errorf("Cannot readBlock() unsupported voxels data shape %s", v.DataShape())
	}
	return nil
}

func (v *Labels) readBlock(tkey storage.TKey, block labels.Block, blockSize dvid.Point) error {
	if blockSize.NumDims() > 3 {
		return fmt.Errorf("DVID voxel blocks currently only supports up to 3d, not 4+ dimensions")
	}
	blockBeg, dataBeg, dataEnd, err := v.ComputeTransform(tkey, blockSize)
	if err != nil {
		return err
	}

	// TODO -- Refactor this function to make direct use of compressed label Block without
	// conversion into full label array.
	labels := v.Data()
	blabels, _ := block.MakeLabelVolume()

	// Compute the strides (in bytes)
	bX := int64(blockSize.Value(0)) * 8
	bY := int64(blockSize.Value(1)) * bX
	dX := int64(v.Stride())

	blockBegX := int64(blockBeg.Value(0))
	blockBegY := int64(blockBeg.Value(1))
	blockBegZ := int64(blockBeg.Value(2))

	// Do the transfers depending on shape of the external voxels.
	switch {
	case v.DataShape().Equals(dvid.XY):
		dataI := int64(dataBeg.Value(1))*dX + int64(dataBeg.Value(0))*8
		blockI := blockBegZ*bY + blockBegY*bX + blockBegX*8
		bytes := int64(dataEnd.Value(0)-dataBeg.Value(0)+1) * 8
		for y := dataBeg.Value(1); y <= dataEnd.Value(1); y++ {
			copy(labels[dataI:dataI+bytes], blabels[blockI:blockI+bytes])
			blockI += bX
			dataI += dX
		}

	case v.DataShape().Equals(dvid.XZ):
		dataI := int64(dataBeg.Value(2))*dX + int64(dataBeg.Value(0))*8
		blockI := blockBegZ*bY + blockBegY*bX + blockBegX*8
		bytes := int64(dataEnd.Value(0)-dataBeg.Value(0)+1) * 8
		for y := dataBeg.Value(2); y <= dataEnd.Value(2); y++ {
			copy(labels[dataI:dataI+bytes], blabels[blockI:blockI+bytes])
			blockI += bY
			dataI += dX
		}

	case v.DataShape().Equals(dvid.YZ):
		bz := blockBegZ
		for y := int64(dataBeg.Value(2)); y <= int64(dataEnd.Value(2)); y++ {
			dataI := y*dX + int64(dataBeg.Value(1))*8
			blockI := bz*bY + blockBegY*bX + blockBegX*8
			for x := dataBeg.Value(1); x <= dataEnd.Value(1); x++ {
				copy(labels[dataI:dataI+8], blabels[blockI:blockI+8])
				blockI += bX
				dataI += 8
			}
			bz++
		}

	case v.DataShape().ShapeDimensions() == 2:
		// TODO: General code for handling 2d ExtData in n-d space.
		return fmt.Errorf("DVID currently does not support 2d in n-d space.")

	case v.DataShape().Equals(dvid.Vol3d):
		blockOffset := blockBegX * 8
		dX := int64(v.Size().Value(0)) * 8
		dY := int64(v.Size().Value(1)) * dX
		dataOffset := int64(dataBeg.Value(0)) * 8
		bytes := int64(dataEnd.Value(0)-dataBeg.Value(0)+1) * 8
		blockZ := blockBegZ

		for dataZ := int64(dataBeg.Value(2)); dataZ <= int64(dataEnd.Value(2)); dataZ++ {
			blockY := blockBegY
			for dataY := int64(dataBeg.Value(1)); dataY <= int64(dataEnd.Value(1)); dataY++ {
				blockI := blockZ*bY + blockY*bX + blockOffset
				dataI := dataZ*dY + dataY*dX + dataOffset
				copy(labels[dataI:dataI+bytes], blabels[blockI:blockI+bytes])
				blockY++
			}
			blockZ++
		}

	default:
		return fmt.Errorf("Cannot readBlock() unsupported voxels data shape %s", v.DataShape())
	}
	return nil
}