Skip to content

Commit 24f2603

Browse files
cmhulbertcmhulbert
committed
feat: support zarr3 sharded output for downsample output.
remaining n5-spark outputs remain TBD
1 parent 82573a8 commit 24f2603

5 files changed

Lines changed: 309 additions & 34 deletions

File tree

‎src/main/java/org/janelia/saalfeldlab/n5/spark/downsample/N5DownsamplerSpark.java‎

Lines changed: 72 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,8 @@
4747
import org.janelia.saalfeldlab.n5.DatasetAttributes;
4848
import org.janelia.saalfeldlab.n5.N5FSWriter;
4949
import org.janelia.saalfeldlab.n5.N5Writer;
50+
import org.janelia.saalfeldlab.n5.zarr.v3.ZarrV3DatasetAttributes;
51+
import org.janelia.saalfeldlab.n5.zarr.v3.ZarrV3KeyValueWriter;
5052
import org.janelia.saalfeldlab.n5.imglib2.N5Utils;
5153
import org.janelia.saalfeldlab.n5.spark.supplier.N5WriterSupplier;
5254
import org.janelia.saalfeldlab.n5.spark.util.CmdUtils;
@@ -176,6 +178,42 @@ public static <T extends NativeType<T> & RealType<T>> void downsample(
176178
final int[] blockSize,
177179
final boolean overwriteExisting) throws IOException {
178180

181+
downsample(
182+
sparkContext,
183+
n5Supplier,
184+
inputDatasetPath,
185+
outputDatasetPath,
186+
downsamplingFactors,
187+
blockSize,
188+
null,
189+
overwriteExisting);
190+
}
191+
192+
/**
193+
* Downsamples the given input dataset with respect to the given downsampling factors, storing the output blocks in
194+
* shards of {@code chunksPerShard} chunks per axis. A {@code null} {@code chunksPerShard} writes an unsharded dataset
195+
* (identical to the other overloads). Sharding requires a zarr3 ({@link ZarrV3KeyValueWriter}) writer.
196+
*
197+
* @param sparkContext
198+
* @param n5Supplier
199+
* @param inputDatasetPath
200+
* @param outputDatasetPath
201+
* @param downsamplingFactors
202+
* @param blockSize
203+
* @param chunksPerShard number of chunks per shard per axis, or {@code null} for an unsharded output
204+
* @param overwriteExisting
205+
* @throws IOException
206+
*/
207+
public static <T extends NativeType<T> & RealType<T>> void downsample(
208+
final JavaSparkContext sparkContext,
209+
final N5WriterSupplier n5Supplier,
210+
final String inputDatasetPath,
211+
final String outputDatasetPath,
212+
final int[] downsamplingFactors,
213+
final int[] blockSize,
214+
final int[] chunksPerShard,
215+
final boolean overwriteExisting) throws IOException {
216+
179217
final N5Writer n5 = n5Supplier.get();
180218
if (!n5.datasetExists(inputDatasetPath))
181219
throw new IllegalArgumentException("Input N5 dataset " + inputDatasetPath + " does not exist");
@@ -206,13 +244,26 @@ public static <T extends NativeType<T> & RealType<T>> void downsample(
206244
}
207245
}
208246

209-
n5.createDataset(
210-
outputDatasetPath,
211-
outputDimensions,
212-
outputBlockSize,
213-
inputAttributes.getDataType(),
214-
inputAttributes.getCompression()
215-
);
247+
if (chunksPerShard != null) {
248+
if (!(n5 instanceof ZarrV3KeyValueWriter))
249+
throw new IllegalArgumentException("Sharded downsampling requires a ZarrV3 writer");
250+
final int[] shardSize = new int[dim];
251+
for (int d = 0; d < dim; ++d)
252+
shardSize[d] = chunksPerShard[d] * outputBlockSize[d];
253+
254+
n5.createDataset(
255+
outputDatasetPath,
256+
new ZarrV3DatasetAttributes(outputDimensions, shardSize, outputBlockSize, inputAttributes.getDataType(), inputAttributes.getCompression())
257+
);
258+
} else {
259+
n5.createDataset(
260+
outputDatasetPath,
261+
outputDimensions,
262+
outputBlockSize,
263+
inputAttributes.getDataType(),
264+
inputAttributes.getCompression()
265+
);
266+
}
216267

217268
// set the downsampling factors attribute
218269
final int[] inputAbsoluteDownsamplingFactors = n5.getAttribute(inputDatasetPath, DOWNSAMPLING_FACTORS_ATTRIBUTE_KEY, int[].class);
@@ -221,14 +272,17 @@ public static <T extends NativeType<T> & RealType<T>> void downsample(
221272
outputAbsoluteDownsamplingFactors[d] = downsamplingFactors[d] * (inputAbsoluteDownsamplingFactors != null ? inputAbsoluteDownsamplingFactors[d] : 1);
222273
n5.setAttribute(outputDatasetPath, DOWNSAMPLING_FACTORS_ATTRIBUTE_KEY, outputAbsoluteDownsamplingFactors);
223274

224-
final CellGrid outputCellGrid = new CellGrid(outputDimensions, outputBlockSize);
275+
/* the parallel-write unit is DatasetAttributes#getBlockSize: the shard for a sharded dataset, the chunk otherwise */
276+
final int[] writeBlockSize = n5.getDatasetAttributes(outputDatasetPath).getBlockSize();
277+
278+
final CellGrid outputCellGrid = new CellGrid(outputDimensions, writeBlockSize);
225279
final long numDownsampledBlocks = Intervals.numElements(outputCellGrid.getGridDimensions());
226280
final List<Long> blockIndexes = LongStream.range(0, numDownsampledBlocks).boxed().collect(Collectors.toList());
227281

228282
sparkContext
229283
.parallelize(blockIndexes, Math.min(blockIndexes.size(), MAX_PARTITIONS))
230284
.foreach(blockIndex -> {
231-
final CellGrid cellGrid = new CellGrid(outputDimensions, outputBlockSize);
285+
final CellGrid cellGrid = new CellGrid(outputDimensions, writeBlockSize);
232286
final long[] blockGridPosition = new long[cellGrid.numDimensions()];
233287
cellGrid.getCellGridPositionFlat(blockIndex, blockGridPosition);
234288

@@ -279,12 +333,19 @@ public static <T extends NativeType<T> & RealType<T>> void downsample(
279333
final RandomAccessibleInterval<T> targetBlock = new ArrayImgFactory<>(defaultValue).create(targetInterval);
280334
Downsample.downsample(sourceBlock, targetBlock, downsamplingFactors);
281335

336+
final DatasetAttributes outputAttributes = n5Writer.getDatasetAttributes(outputDatasetPath);
337+
final int[] shardBlockSize = outputAttributes.getBlockSize();
338+
final int[] innerChunkSize = outputAttributes.getChunkSize();
339+
final long[] chunkGridOffset = new long[dim];
340+
for (int d = 0; d < dim; ++d)
341+
chunkGridOffset[d] = blockGridPosition[d] * (shardBlockSize[d] / innerChunkSize[d]);
342+
282343
if (overwriteExisting) {
283344
// Empty blocks will not be written out. Delete blocks to avoid remnant blocks if overwriting.
284-
N5Utils.deleteBlock(targetBlock, n5Writer, outputDatasetPath, blockGridPosition);
345+
N5Utils.deleteBlock(targetBlock, n5Writer, outputDatasetPath, chunkGridOffset);
285346
}
286347

287-
N5Utils.saveNonEmptyBlock(targetBlock, n5Writer, outputDatasetPath, blockGridPosition, defaultValue);
348+
N5Utils.saveNonEmptyBlock(targetBlock, n5Writer, outputDatasetPath, outputAttributes, chunkGridOffset, defaultValue);
288349
});
289350
}
290351

‎src/main/java/org/janelia/saalfeldlab/n5/spark/downsample/N5LabelDownsamplerSpark.java‎

Lines changed: 75 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -52,10 +52,13 @@
5252
import org.apache.spark.api.java.JavaSparkContext;
5353
import org.janelia.saalfeldlab.n5.DatasetAttributes;
5454
import org.janelia.saalfeldlab.n5.N5FSWriter;
55+
import org.janelia.saalfeldlab.n5.zarr.v3.ZarrV3DatasetAttributes;
56+
import org.janelia.saalfeldlab.n5.zarr.v3.ZarrV3KeyValueWriter;
5557
import org.janelia.saalfeldlab.n5.N5Writer;
5658
import org.janelia.saalfeldlab.n5.imglib2.N5Utils;
5759
import org.janelia.saalfeldlab.n5.spark.supplier.N5WriterSupplier;
5860
import org.janelia.saalfeldlab.n5.spark.util.CmdUtils;
61+
import org.janelia.scicomp.n5.zstandard.ZstandardCompression;
5962
import org.kohsuke.args4j.CmdLineException;
6063
import org.kohsuke.args4j.CmdLineParser;
6164
import org.kohsuke.args4j.Option;
@@ -159,6 +162,43 @@ public static <T extends NativeType<T> & IntegerType<T>> void downsampleLabel(
159162
final int[] blockSize,
160163
final boolean overwriteExisting) throws IOException {
161164

165+
downsampleLabel(
166+
sparkContext,
167+
n5Supplier,
168+
inputDatasetPath,
169+
outputDatasetPath,
170+
downsamplingFactors,
171+
blockSize,
172+
null,
173+
overwriteExisting
174+
);
175+
}
176+
177+
/**
178+
* Downsamples the given input dataset with respect to the given downsampling factors, storing the output blocks
179+
* in shards of {@code chunksPerShard} chunks per axis. A {@code null} {@code chunksPerShard} writes an unsharded
180+
* dataset (identical to the other overloads). Sharding requires a zarr3 ({@link ZarrV3KeyValueWriter}) writer.
181+
*
182+
* @param sparkContext
183+
* @param n5Supplier
184+
* @param inputDatasetPath
185+
* @param outputDatasetPath
186+
* @param downsamplingFactors
187+
* @param blockSize
188+
* @param chunksPerShard number of chunks per shard per axis, or {@code null} for an unsharded output
189+
* @param overwriteExisting
190+
* @throws IOException
191+
*/
192+
public static <T extends NativeType<T> & IntegerType<T>> void downsampleLabel(
193+
final JavaSparkContext sparkContext,
194+
final N5WriterSupplier n5Supplier,
195+
final String inputDatasetPath,
196+
final String outputDatasetPath,
197+
final int[] downsamplingFactors,
198+
final int[] blockSize,
199+
final int[] chunksPerShard,
200+
final boolean overwriteExisting) throws IOException {
201+
162202
final N5Writer n5 = n5Supplier.get();
163203
if (!n5.datasetExists(inputDatasetPath))
164204
throw new IllegalArgumentException("Input N5 dataset " + inputDatasetPath + " does not exist");
@@ -189,22 +229,38 @@ public static <T extends NativeType<T> & IntegerType<T>> void downsampleLabel(
189229
}
190230
}
191231

192-
n5.createDataset(
193-
outputDatasetPath,
194-
outputDimensions,
195-
outputBlockSize,
196-
inputAttributes.getDataType(),
197-
inputAttributes.getCompression()
198-
);
232+
if (chunksPerShard != null) {
233+
if (!(n5 instanceof ZarrV3KeyValueWriter))
234+
throw new IllegalArgumentException("Sharded downsampling requires a ZarrV3 writer");
235+
final int[] shardSize = new int[dim];
236+
for (int d = 0; d < dim; ++d)
237+
shardSize[d] = chunksPerShard[d] * outputBlockSize[d];
199238

200-
final CellGrid outputCellGrid = new CellGrid(outputDimensions, outputBlockSize);
239+
n5.createDataset(
240+
outputDatasetPath,
241+
new ZarrV3DatasetAttributes(outputDimensions, shardSize, outputBlockSize, inputAttributes.getDataType(), new ZstandardCompression())
242+
);
243+
} else {
244+
n5.createDataset(
245+
outputDatasetPath,
246+
outputDimensions,
247+
outputBlockSize,
248+
inputAttributes.getDataType(),
249+
new ZstandardCompression()
250+
);
251+
}
252+
253+
/* the parallel-write unit is DatasetAttributes#getBlockSize: the shard for a sharded dataset, the chunk otherwise */
254+
final int[] writeBlockSize = n5.getDatasetAttributes(outputDatasetPath).getBlockSize();
255+
256+
final CellGrid outputCellGrid = new CellGrid(outputDimensions, writeBlockSize);
201257
final long numDownsampledBlocks = Intervals.numElements(outputCellGrid.getGridDimensions());
202258
final List<Long> blockIndexes = LongStream.range(0, numDownsampledBlocks).boxed().collect(Collectors.toList());
203259

204260
sparkContext
205261
.parallelize(blockIndexes, Math.min(blockIndexes.size(), MAX_PARTITIONS))
206262
.foreach(blockIndex -> {
207-
final CellGrid cellGrid = new CellGrid(outputDimensions, outputBlockSize);
263+
final CellGrid cellGrid = new CellGrid(outputDimensions, writeBlockSize);
208264
final long[] blockGridPosition = new long[cellGrid.numDimensions()];
209265
cellGrid.getCellGridPositionFlat(blockIndex, blockGridPosition);
210266

@@ -253,17 +309,23 @@ public static <T extends NativeType<T> & IntegerType<T>> void downsampleLabel(
253309
final RandomAccessibleInterval<T> targetBlock = new ArrayImgFactory<>(defaultValue).create(targetInterval);
254310
downsampleLabel(sourceBlock, targetBlock, downsamplingFactors);
255311

312+
final DatasetAttributes outputAttributes = n5Writer.getDatasetAttributes(outputDatasetPath);
313+
final int[] shardBlockSize = outputAttributes.getBlockSize();
314+
final int[] innerChunkSize = outputAttributes.getChunkSize();
315+
final long[] chunkGridOffset = new long[dim];
316+
for (int d = 0; d < dim; ++d)
317+
chunkGridOffset[d] = blockGridPosition[d] * (shardBlockSize[d] / innerChunkSize[d]);
318+
256319
if (overwriteExisting) {
257320
// Empty blocks will not be written out. Delete blocks to avoid remnant blocks if overwriting.
258-
N5Utils.deleteBlock(targetBlock, n5Writer, outputDatasetPath, blockGridPosition);
321+
N5Utils.deleteBlock(targetBlock, n5Writer, outputDatasetPath, chunkGridOffset);
259322
}
260-
261-
N5Utils.saveNonEmptyBlock(targetBlock, n5Writer, outputDatasetPath, blockGridPosition, defaultValue);
323+
N5Utils.saveNonEmptyBlock(targetBlock, n5Writer, outputDatasetPath, outputAttributes, chunkGridOffset, defaultValue);
262324
});
263325
}
264326

265327
/**
266-
* Based on {@link bdv.export.Downsample}.
328+
* Based on bdv.export.Downsample.
267329
*/
268330
private static <T extends IntegerType<T>> void downsampleLabel(
269331
final RandomAccessible<T> input,

‎src/main/java/org/janelia/saalfeldlab/n5/spark/downsample/N5OffsetDownsamplerSpark.java‎

Lines changed: 76 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,8 @@
4343
import org.janelia.saalfeldlab.n5.DatasetAttributes;
4444
import org.janelia.saalfeldlab.n5.N5FSWriter;
4545
import org.janelia.saalfeldlab.n5.N5Writer;
46+
import org.janelia.saalfeldlab.n5.zarr.v3.ZarrV3DatasetAttributes;
47+
import org.janelia.saalfeldlab.n5.zarr.v3.ZarrV3KeyValueWriter;
4648
import org.janelia.saalfeldlab.n5.imglib2.N5Utils;
4749
import org.janelia.saalfeldlab.n5.spark.supplier.N5WriterSupplier;
4850
import org.janelia.saalfeldlab.n5.spark.util.CmdUtils;
@@ -133,6 +135,43 @@ public static < T extends NativeType< T > & RealType< T > > void downsampleWithO
133135
final int[] downsamplingFactors,
134136
final long[] offset,
135137
final int[] blockSize ) throws IOException
138+
{
139+
downsampleWithOffset(
140+
sparkContext,
141+
n5Supplier,
142+
inputDatasetPath,
143+
outputDatasetPath,
144+
downsamplingFactors,
145+
offset,
146+
blockSize,
147+
null
148+
);
149+
}
150+
151+
/**
152+
* Downsamples the given input dataset with respect to the given downsampling factors and offset, storing the output
153+
* blocks in shards of {@code chunksPerShard} chunks per axis. A {@code null} {@code chunksPerShard} writes an
154+
* unsharded dataset (identical to the other overloads). Sharding requires a zarr3 ({@link ZarrV3KeyValueWriter}) writer.
155+
*
156+
* @param sparkContext
157+
* @param n5Supplier
158+
* @param inputDatasetPath
159+
* @param outputDatasetPath
160+
* @param downsamplingFactors
161+
* @param offset
162+
* @param blockSize
163+
* @param chunksPerShard number of chunks per shard per axis, or {@code null} for an unsharded output
164+
* @throws IOException
165+
*/
166+
public static < T extends NativeType< T > & RealType< T > > void downsampleWithOffset(
167+
final JavaSparkContext sparkContext,
168+
final N5WriterSupplier n5Supplier,
169+
final String inputDatasetPath,
170+
final String outputDatasetPath,
171+
final int[] downsamplingFactors,
172+
final long[] offset,
173+
final int[] blockSize,
174+
final int[] chunksPerShard ) throws IOException
136175
{
137176
final N5Writer n5 = n5Supplier.get();
138177
if ( !n5.datasetExists( inputDatasetPath ) )
@@ -154,22 +193,42 @@ public static < T extends NativeType< T > & RealType< T > > void downsampleWithO
154193
throw new IllegalArgumentException( "Degenerate output dimensions: " + Arrays.toString( outputDimensions ) );
155194

156195
final int[] outputBlockSize = blockSize != null ? blockSize : inputAttributes.getBlockSize();
157-
n5.createDataset(
158-
outputDatasetPath,
159-
outputDimensions,
160-
outputBlockSize,
161-
inputAttributes.getDataType(),
162-
inputAttributes.getCompression()
163-
);
164196

165-
final CellGrid outputCellGrid = new CellGrid( outputDimensions, outputBlockSize );
197+
if ( chunksPerShard != null )
198+
{
199+
if ( !( n5 instanceof ZarrV3KeyValueWriter ) )
200+
throw new IllegalArgumentException( "Sharded downsampling requires a ZarrV3 writer" );
201+
final int[] shardSize = new int[ dim ];
202+
for ( int d = 0; d < dim; ++d )
203+
shardSize[ d ] = chunksPerShard[ d ] * outputBlockSize[ d ];
204+
205+
n5.createDataset(
206+
outputDatasetPath,
207+
new ZarrV3DatasetAttributes( outputDimensions, shardSize, outputBlockSize, inputAttributes.getDataType(), inputAttributes.getCompression() )
208+
);
209+
}
210+
else
211+
{
212+
n5.createDataset(
213+
outputDatasetPath,
214+
outputDimensions,
215+
outputBlockSize,
216+
inputAttributes.getDataType(),
217+
inputAttributes.getCompression()
218+
);
219+
}
220+
221+
/* the parallel-write unit is DatasetAttributes#getBlockSize: the shard for a sharded dataset, the chunk otherwise */
222+
final int[] writeBlockSize = n5.getDatasetAttributes( outputDatasetPath ).getBlockSize();
223+
224+
final CellGrid outputCellGrid = new CellGrid( outputDimensions, writeBlockSize );
166225
final long numDownsampledBlocks = Intervals.numElements( outputCellGrid.getGridDimensions() );
167226
final List< Long > blockIndexes = LongStream.range( 0, numDownsampledBlocks ).boxed().collect( Collectors.toList() );
168227

169228
sparkContext.parallelize( blockIndexes, Math.min( blockIndexes.size(), MAX_PARTITIONS ) ).foreach( blockIndex ->
170229
{
171230
// downsampled block index to grid position
172-
final CellGrid cellGrid = new CellGrid( outputDimensions, outputBlockSize );
231+
final CellGrid cellGrid = new CellGrid( outputDimensions, writeBlockSize );
173232
final long[] blockGridPosition = new long[ cellGrid.numDimensions() ];
174233
cellGrid.getCellGridPositionFlat( blockIndex, blockGridPosition );
175234

@@ -237,7 +296,14 @@ public static < T extends NativeType< T > & RealType< T > > void downsampleWithO
237296
else
238297
downsampleIntervalOutOfBoundsCheck( sourceBlock, targetBlock, downsamplingFactors, definedSourceBlockInterval );
239298

240-
N5Utils.saveNonEmptyBlock( targetBlock, n5Writer, outputDatasetPath, blockGridPosition, defaultValue );
299+
final DatasetAttributes outputAttributes = n5Writer.getDatasetAttributes( outputDatasetPath );
300+
final int[] shardBlockSize = outputAttributes.getBlockSize();
301+
final int[] innerChunkSize = outputAttributes.getChunkSize();
302+
final long[] chunkGridOffset = new long[ dim ];
303+
for ( int d = 0; d < dim; ++d )
304+
chunkGridOffset[ d ] = blockGridPosition[ d ] * ( shardBlockSize[ d ] / innerChunkSize[ d ] );
305+
306+
N5Utils.saveNonEmptyBlock( targetBlock, n5Writer, outputDatasetPath, outputAttributes, chunkGridOffset, defaultValue );
241307
} );
242308
}
243309

0 commit comments

Comments
 (0)