Skip to content

Commit 4bf1a7c

Browse files
tpietzschtpietzsch
committed
Add N5Writer.writeRegion methods
writeRegion() takes a DataBlockSupplier and a region in pixel coordinates that should be written. It decised which shards/blocks need to be touched and asks the DataBlockSupplier for the respective blocks. Currently only minimal parallelization (over shards) is implemented. This could be improved in the future, if called with a ForkJoinPool executor. Then we could further parallelize the production and encoding of DataBlocks.
1 parent 59eed63 commit 4bf1a7c

6 files changed

Lines changed: 585 additions & 10 deletions

File tree

‎src/main/java/org/janelia/saalfeldlab/n5/GsonKeyValueN5Writer.java‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@
3535
import java.util.Map;
3636

3737
import com.google.gson.JsonSyntaxException;
38+
import java.util.concurrent.ExecutionException;
39+
import java.util.concurrent.ExecutorService;
3840
import org.janelia.saalfeldlab.n5.N5Exception.N5IOException;
3941
import org.janelia.saalfeldlab.n5.shard.PositionValueAccess;
4042

@@ -210,6 +212,41 @@ default boolean removeAttributes(final String pathName, final List<String> attri
210212
return removed;
211213
}
212214

215+
@Override
216+
default <T> void writeRegion(
217+
final String datasetPath,
218+
final DatasetAttributes datasetAttributes,
219+
final long[] min,
220+
final long[] size,
221+
final DataBlockSupplier<T> dataBlocks,
222+
final boolean writeFully) throws N5Exception {
223+
try {
224+
final PositionValueAccess posKva = PositionValueAccess.fromKva(getKeyValueAccess(), getURI(), N5URI.normalizeGroupPath(datasetPath), datasetAttributes);
225+
datasetAttributes.<T>getDatasetAccess().writeRegion(posKva, min, size, dataBlocks, datasetAttributes.getDimensions(), writeFully);
226+
} catch (final UncheckedIOException e) {
227+
throw new N5IOException(
228+
"Failed to write blocks into dataset " + datasetPath, e);
229+
}
230+
}
231+
232+
@Override
233+
default <T> void writeRegion(
234+
final String datasetPath,
235+
final DatasetAttributes datasetAttributes,
236+
final long[] min,
237+
final long[] size,
238+
final DataBlockSupplier<T> dataBlocks,
239+
final boolean writeFully,
240+
final ExecutorService exec) throws N5Exception, InterruptedException, ExecutionException {
241+
try {
242+
final PositionValueAccess posKva = PositionValueAccess.fromKva(getKeyValueAccess(), getURI(), N5URI.normalizeGroupPath(datasetPath), datasetAttributes);
243+
datasetAttributes.<T>getDatasetAccess().writeRegion(posKva, min, size, dataBlocks, datasetAttributes.getDimensions(), writeFully, exec);
244+
} catch (final UncheckedIOException e) {
245+
throw new N5IOException(
246+
"Failed to write blocks into dataset " + datasetPath, e);
247+
}
248+
}
249+
213250
@Override
214251
default <T> void writeBlocks(
215252
final String datasetPath,
@@ -239,7 +276,6 @@ default <T> void writeBlock(
239276
"Failed to write block " + Arrays.toString(dataBlock.getGridPosition()) + " into dataset " + path,
240277
e);
241278
}
242-
243279
}
244280

245281
@Override

‎src/main/java/org/janelia/saalfeldlab/n5/N5Writer.java‎

Lines changed: 52 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,6 @@
2828
*/
2929
package org.janelia.saalfeldlab.n5;
3030

31-
import org.janelia.saalfeldlab.n5.codec.BlockCodecInfo;
32-
import org.janelia.saalfeldlab.n5.codec.DataCodecInfo;
33-
3431
import java.io.ByteArrayOutputStream;
3532
import java.io.IOException;
3633
import java.io.ObjectOutputStream;
@@ -39,6 +36,8 @@
3936
import java.util.Collections;
4037
import java.util.List;
4138
import java.util.Map;
39+
import java.util.concurrent.ExecutionException;
40+
import java.util.concurrent.ExecutorService;
4241

4342
/**
4443
* A simple structured container API for hierarchies of chunked
@@ -265,6 +264,56 @@ default <T> void writeBlocks(
265264
writeBlock(datasetPath, datasetAttributes, block);
266265
}
267266

267+
@FunctionalInterface
268+
interface DataBlockSupplier<T> {
269+
270+
/**
271+
*
272+
* @param gridPos
273+
* @param existingDataBlock
274+
* existing data to be merged into the new data block (maybe {@code null})
275+
*
276+
* @return data block at the given gridPos
277+
*/
278+
DataBlock<T> get(long[] gridPos, final DataBlock<T> existingDataBlock);
279+
}
280+
281+
/**
282+
* @param datasetPath the dataset path
283+
* @param datasetAttributes the dataset attributes
284+
* @param min min pixel coordinate of region to write
285+
* @param size size in pixels of region to write
286+
* @param dataBlocks is asked to create blocks within the given region
287+
* @param writeFully if false, merge existing data in shards/blocks that overlap the region boundary. if true, override everything.
288+
* @throws N5Exception the exception
289+
*/
290+
<T> void writeRegion(
291+
String datasetPath,
292+
DatasetAttributes datasetAttributes,
293+
long[] min,
294+
long[] size,
295+
DataBlockSupplier<T> dataBlocks,
296+
boolean writeFully) throws N5Exception;
297+
298+
/**
299+
* @param datasetPath the dataset path
300+
* @param datasetAttributes the dataset attributes
301+
* @param min min pixel coordinate of region to write
302+
* @param size size in pixels of region to write
303+
* @param dataBlocks is asked to create blocks within the given region
304+
* @param writeFully if false, merge existing data in shards/blocks that overlap the region boundary. if true, override everything.
305+
* @param exec used to parallelize over blocks and shards
306+
* @throws N5Exception the exception
307+
*/
308+
<T> void writeRegion(
309+
String datasetPath,
310+
DatasetAttributes datasetAttributes,
311+
long[] min,
312+
long[] size,
313+
DataBlockSupplier<T> dataBlocks,
314+
boolean writeFully,
315+
ExecutorService exec) throws N5Exception, InterruptedException, ExecutionException;
316+
268317
/**
269318
* Deletes the block at {@code gridPosition}.
270319
*

‎src/main/java/org/janelia/saalfeldlab/n5/shard/DatasetAccess.java‎

Lines changed: 64 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,18 @@
11
package org.janelia.saalfeldlab.n5.shard;
22

3-
import org.janelia.saalfeldlab.n5.DataBlock;
4-
import org.janelia.saalfeldlab.n5.N5Exception.N5IOException;
5-
import org.janelia.saalfeldlab.n5.shard.Nesting.NestedGrid;
6-
import org.janelia.saalfeldlab.n5.shard.Nesting.NestedPosition;
7-
83
import java.util.ArrayList;
94
import java.util.Collection;
105
import java.util.Collections;
116
import java.util.List;
127
import java.util.TreeMap;
8+
import java.util.concurrent.ExecutionException;
9+
import java.util.concurrent.ExecutorService;
10+
import org.janelia.saalfeldlab.n5.DataBlock;
11+
import org.janelia.saalfeldlab.n5.N5Exception;
12+
import org.janelia.saalfeldlab.n5.N5Exception.N5IOException;
13+
import org.janelia.saalfeldlab.n5.N5Writer.DataBlockSupplier;
14+
import org.janelia.saalfeldlab.n5.shard.Nesting.NestedGrid;
15+
import org.janelia.saalfeldlab.n5.shard.Nesting.NestedPosition;
1316

1417
/**
1518
* Wrap an instantiated DataBlock/shard codec hierarchy to implement (single and
@@ -26,10 +29,66 @@ public interface DatasetAccess<T> {
2629

2730
boolean deleteBlock(PositionValueAccess kva, long[] gridPosition) throws N5IOException;
2831

32+
2933
List<DataBlock<T>> readBlocks(PositionValueAccess kva, List<long[]> positions) throws N5IOException;
3034

3135
void writeBlocks(PositionValueAccess kva, List<DataBlock<T>> blocks) throws N5IOException;
3236

37+
38+
/**
39+
*
40+
* @param pva
41+
* @param min
42+
* min pixel coordinate of region to write
43+
* @param size
44+
* size in pixels of region to write
45+
* @param blocks
46+
* is asked to create blocks within the given region
47+
* @param datasetDimensions
48+
* @param writeFully
49+
* if false, merge existing data in shards/blocks that overlap the region boundary. if true, override everything.
50+
*
51+
* @throws N5IOException
52+
*/
53+
void writeRegion(
54+
PositionValueAccess pva,
55+
long[] min,
56+
long[] size,
57+
DataBlockSupplier<T> blocks,
58+
long[] datasetDimensions,
59+
boolean writeFully
60+
) throws N5IOException;
61+
62+
/**
63+
*
64+
* @param pva
65+
* @param min
66+
* min pixel coordinate of region to write
67+
* @param size
68+
* size in pixels of region to write
69+
* @param blocks
70+
* is asked to create blocks within the given region. must be thread-safe.
71+
* @param datasetDimensions
72+
* @param writeFully
73+
* if false, merge existing data in shards/blocks that overlap the region boundary. if true, override everything.
74+
* @param exec
75+
* used to parallelize over blocks and shards
76+
*
77+
* @throws N5Exception
78+
* @throws InterruptedException
79+
* @throws ExecutionException
80+
*/
81+
void writeRegion(
82+
PositionValueAccess pva,
83+
long[] min,
84+
long[] size,
85+
DataBlockSupplier<T> blocks,
86+
long[] datasetDimensions,
87+
boolean writeFully,
88+
ExecutorService exec
89+
) throws N5Exception, InterruptedException, ExecutionException;
90+
91+
3392
NestedGrid getGrid();
3493

3594
/**

‎src/main/java/org/janelia/saalfeldlab/n5/shard/DefaultDatasetAccess.java‎

Lines changed: 94 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,20 @@
55
import java.util.Collection;
66
import java.util.Collections;
77
import java.util.List;
8+
import java.util.concurrent.ExecutionException;
9+
import java.util.concurrent.ExecutorService;
810
import java.util.stream.Collectors;
911
import org.janelia.saalfeldlab.n5.DataBlock;
1012
import org.janelia.saalfeldlab.n5.N5Exception;
1113
import org.janelia.saalfeldlab.n5.N5Exception.N5IOException;
1214
import org.janelia.saalfeldlab.n5.N5Exception.N5NoSuchKeyException;
15+
import org.janelia.saalfeldlab.n5.N5Writer.DataBlockSupplier;
1316
import org.janelia.saalfeldlab.n5.codec.BlockCodec;
1417
import org.janelia.saalfeldlab.n5.readdata.ReadData;
1518
import org.janelia.saalfeldlab.n5.shard.Nesting.NestedGrid;
1619
import org.janelia.saalfeldlab.n5.shard.Nesting.NestedPosition;
1720

18-
import static org.janelia.saalfeldlab.n5.shard.DatasetAccess.*;
21+
import static org.janelia.saalfeldlab.n5.shard.DatasetAccess.groupInnerPositions;
1922

2023
public class DefaultDatasetAccess<T> implements DatasetAccess<T> {
2124

@@ -253,6 +256,96 @@ private ReadData writeShardRecursive(
253256
return codec.encode(new RawShardDataBlock(shardPosition, shard));
254257
}
255258

259+
public void writeRegion(
260+
final PositionValueAccess pva,
261+
final long[] min,
262+
final long[] size,
263+
final DataBlockSupplier<T> blocks,
264+
final long[] datasetDimensions,
265+
final boolean writeFully
266+
) throws N5IOException {
267+
268+
final Region region = new Region(min, size, grid, datasetDimensions);
269+
270+
for (long[] key : Region.gridPositions(region.minPos().key(), region.maxPos().key())) {
271+
final NestedPosition pos = grid.nestedPosition(key, grid.numLevels() - 1);
272+
final boolean nestedWriteFully = writeFully || region.fullyContains(pos);
273+
final ReadData existingData = nestedWriteFully ? null : getExistingReadData(pva, key);
274+
final ReadData modifiedData = writeRegionRecursive(existingData, region, blocks, pos);
275+
pva.put(key, modifiedData);
276+
}
277+
}
278+
279+
@Override
280+
public void writeRegion(
281+
final PositionValueAccess pva,
282+
final long[] min,
283+
final long[] size,
284+
final DataBlockSupplier<T> blocks,
285+
final long[] datasetDimensions,
286+
final boolean writeFully,
287+
final ExecutorService exec) throws N5Exception, InterruptedException, ExecutionException {
288+
289+
final Region region = new Region(min, size, grid, datasetDimensions);
290+
291+
for (long[] key : Region.gridPositions(region.minPos().key(), region.maxPos().key())) {
292+
exec.submit(() -> {
293+
final NestedPosition pos = grid.nestedPosition(key, grid.numLevels() - 1);
294+
final boolean nestedWriteFully = writeFully || region.fullyContains(pos);
295+
final ReadData existingData = nestedWriteFully ? null : getExistingReadData(pva, key);
296+
final ReadData modifiedData = writeRegionRecursive(existingData, region, blocks, pos);
297+
pva.put(key, modifiedData);
298+
});
299+
}
300+
}
301+
302+
private ReadData writeRegionRecursive(
303+
final ReadData existingReadData, // may be null
304+
final Region region,
305+
final DataBlockSupplier<T> blocks,
306+
final NestedPosition position
307+
) {
308+
final boolean writeFully = existingReadData == null;
309+
final int level = position.level();
310+
String indent = "";
311+
for (int i = 0; i < 2-level; ++i)
312+
indent += " ";
313+
System.out.println(indent + "position = " + position + (writeFully ? " (writeFully)" : ""));
314+
if ( level == 0 ) {
315+
316+
@SuppressWarnings("unchecked")
317+
final BlockCodec<T> codec = (BlockCodec<T>) codecs[0];
318+
final long[] gridPosition = position.absolute(0);
319+
320+
// If the DataBlock is not fully contained in the region, we will
321+
// get existingReadData != null. In that case, we decode the
322+
// existing DataBlock and pass it to the BlockSupplier for modification.
323+
final DataBlock<T> existingDataBlock = (existingReadData == null)
324+
? null
325+
: codec.decode(existingReadData, gridPosition);
326+
final DataBlock<T> dataBlock = blocks.get(gridPosition, existingDataBlock);
327+
328+
return codec.encode(dataBlock);
329+
} else {
330+
331+
@SuppressWarnings("unchecked")
332+
final BlockCodec<RawShard> codec = (BlockCodec<RawShard>) codecs[level];
333+
final long[] gridPos = position.absolute(level);
334+
final RawShard shard = existingReadData == null ?
335+
new RawShard(grid.relativeBlockSize(level)) :
336+
codec.decode(existingReadData, gridPos).getData();
337+
for (NestedPosition pos : region.containedNestedPositions(position)) {
338+
final boolean nestedWriteFully = writeFully || region.fullyContains(pos);
339+
final long[] elementPos = pos.relative();
340+
final ReadData existingElementData = nestedWriteFully ? null : shard.getElementData(elementPos);
341+
final ReadData modifiedElementData = writeRegionRecursive(existingElementData, region, blocks, pos);
342+
shard.setElementData(modifiedElementData, elementPos);
343+
}
344+
345+
return codec.encode(new RawShardDataBlock(gridPos, shard));
346+
}
347+
}
348+
256349
@Override
257350
public boolean deleteBlock(final PositionValueAccess kva, final long[] gridPosition) throws N5IOException {
258351
final NestedPosition position = grid.nestedPosition(gridPosition);

0 commit comments

Comments
 (0)