Skip to content

Commit e07b7fe

Browse files
brokkoli71brokkoli71
committed
v2 and v3 Array implements Array Interface
1 parent 23164e0 commit e07b7fe

File tree

13 files changed

+540
-308
lines changed

13 files changed

+540
-308
lines changed

‎src/main/java/dev/zarr/zarrjava/interfaces/Array.java‎

Lines changed: 337 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,337 @@
1+
package dev.zarr.zarrjava.interfaces;
2+
3+
import dev.zarr.zarrjava.ZarrException;
4+
import dev.zarr.zarrjava.store.StoreHandle;
5+
import dev.zarr.zarrjava.utils.IndexingUtils;
6+
import dev.zarr.zarrjava.utils.MultiArrayUtils;
7+
import dev.zarr.zarrjava.utils.Utils;
8+
import dev.zarr.zarrjava.v3.codec.CodecPipeline;
9+
import ucar.ma2.InvalidRangeException;
10+
11+
import javax.annotation.Nonnull;
12+
import javax.annotation.Nullable;
13+
import java.nio.ByteBuffer;
14+
import java.util.Arrays;
15+
import java.util.stream.Stream;
16+
17+
public interface Array {
18+
19+
ArrayMetadata metadata();
20+
21+
/**
22+
* Writes a ucar.ma2.Array into the Zarr array at a specified offset. The shape of the Zarr array
23+
* needs be large enough for the write.
24+
*
25+
* @param offset
26+
* @param array
27+
* @param parallel
28+
*/
29+
default void write(long[] offset, ucar.ma2.Array array, boolean parallel) {
30+
ArrayMetadata metadata = metadata();
31+
if (offset.length != metadata.ndim()) {
32+
throw new IllegalArgumentException("'offset' needs to have rank '" + metadata.ndim() + "'.");
33+
}
34+
if (array.getRank() != metadata.ndim()) {
35+
throw new IllegalArgumentException("'array' needs to have rank '" + metadata.ndim() + "'.");
36+
}
37+
38+
int[] shape = array.getShape();
39+
40+
final int[] chunkShape = metadata.chunkShape();
41+
Stream<long[]> chunkStream = Arrays.stream(IndexingUtils.computeChunkCoords(metadata.shape(), chunkShape, offset, shape));
42+
if (parallel) {
43+
chunkStream = chunkStream.parallel();
44+
}
45+
chunkStream.forEach(
46+
chunkCoords -> {
47+
try {
48+
final IndexingUtils.ChunkProjection chunkProjection =
49+
IndexingUtils.computeProjection(chunkCoords, metadata.shape(), chunkShape, offset,
50+
shape
51+
);
52+
53+
ucar.ma2.Array chunkArray;
54+
if (IndexingUtils.isFullChunk(chunkProjection.chunkOffset, chunkProjection.shape,
55+
chunkShape
56+
)) {
57+
chunkArray = array.sectionNoReduce(chunkProjection.outOffset,
58+
chunkProjection.shape,
59+
null
60+
);
61+
} else {
62+
chunkArray = readChunk(chunkCoords);
63+
MultiArrayUtils.copyRegion(array, chunkProjection.outOffset, chunkArray,
64+
chunkProjection.chunkOffset, chunkProjection.shape
65+
);
66+
}
67+
writeChunk(chunkCoords, chunkArray);
68+
} catch (ZarrException | InvalidRangeException e) {
69+
throw new RuntimeException(e);
70+
}
71+
});
72+
73+
}
74+
75+
/**
76+
* Writes one chunk into the Zarr array as specified by the chunk coordinates. The shape of the
77+
* Zarr array needs be large enough for the write.
78+
*
79+
* @param chunkCoords
80+
* @param chunkArray
81+
* @throws ZarrException
82+
*/
83+
default void writeChunk(long[] chunkCoords, ucar.ma2.Array chunkArray) throws ZarrException {
84+
ArrayMetadata metadata = metadata();
85+
String[] chunkKeys = metadata.chunkKeyEncoding().encodeChunkKey(chunkCoords);
86+
StoreHandle chunkHandle = storeHandle().resolve(chunkKeys);
87+
88+
if (MultiArrayUtils.allValuesEqual(chunkArray, metadata.parsedFillValue())) {
89+
chunkHandle.delete();
90+
} else {
91+
ByteBuffer chunkBytes = codecPipeline().encode(chunkArray);
92+
chunkHandle.set(chunkBytes);
93+
}
94+
}
95+
96+
/**
97+
* Reads one chunk of the Zarr array as specified by the chunk coordinates into an
98+
* ucar.ma2.Array.
99+
*
100+
* @param chunkCoords The coordinates of the chunk as computed by the offset of the chunk divided
101+
* by the chunk shape.
102+
* @throws ZarrException
103+
*/
104+
@Nonnull
105+
default ucar.ma2.Array readChunk(long[] chunkCoords)
106+
throws ZarrException {
107+
ArrayMetadata metadata = metadata();
108+
if (!chunkIsInArray(chunkCoords)) {
109+
throw new ZarrException("Attempting to read data outside of the array's domain.");
110+
}
111+
112+
final String[] chunkKeys = metadata.chunkKeyEncoding().encodeChunkKey(chunkCoords);
113+
final StoreHandle chunkHandle = storeHandle().resolve(chunkKeys);
114+
115+
ByteBuffer chunkBytes = chunkHandle.read();
116+
if (chunkBytes == null) {
117+
return metadata.allocateFillValueChunk();
118+
}
119+
120+
return codecPipeline().decode(chunkBytes);
121+
}
122+
123+
124+
/**
125+
* Writes a ucar.ma2.Array into the Zarr array at the beginning of the Zarr array. The shape of
126+
* the Zarr array needs be large enough for the write.
127+
* Utilizes no parallelism.
128+
*
129+
* @param array
130+
*/
131+
default void write(ucar.ma2.Array array) {
132+
write(new long[metadata().ndim()], array);
133+
}
134+
135+
/**
136+
* Writes a ucar.ma2.Array into the Zarr array at a specified offset. The shape of the Zarr array
137+
* needs be large enough for the write.
138+
* Utilizes no parallelism.
139+
*
140+
* @param offset
141+
* @param array
142+
*/
143+
default void write(long[] offset, ucar.ma2.Array array) {
144+
write(offset, array, false);
145+
}
146+
147+
/**
148+
* Writes a ucar.ma2.Array into the Zarr array at the beginning of the Zarr array. The shape of
149+
* the Zarr array needs be large enough for the write.
150+
*
151+
* @param array
152+
* @param parallel
153+
*/
154+
default void write(ucar.ma2.Array array, boolean parallel) {
155+
write(new long[metadata().ndim()], array, parallel);
156+
}
157+
158+
/**
159+
* Reads the entire Zarr array into an ucar.ma2.Array.
160+
* Utilizes no parallelism.
161+
*
162+
* @throws ZarrException
163+
*/
164+
@Nonnull
165+
default ucar.ma2.Array read() throws ZarrException {
166+
return read(new long[metadata().ndim()], Utils.toIntArray(metadata().shape()));
167+
}
168+
169+
/**
170+
* Reads a part of the Zarr array based on a requested offset and shape into an ucar.ma2.Array.
171+
* Utilizes no parallelism.
172+
*
173+
* @param offset
174+
* @param shape
175+
* @throws ZarrException
176+
*/
177+
@Nonnull
178+
default ucar.ma2.Array read(final long[] offset, final int[] shape) throws ZarrException {
179+
return read(offset, shape, false);
180+
}
181+
182+
/**
183+
* Reads the entire Zarr array into an ucar.ma2.Array.
184+
*
185+
* @param parallel
186+
* @throws ZarrException
187+
*/
188+
@Nonnull
189+
default ucar.ma2.Array read(final boolean parallel) throws ZarrException {
190+
return read(new long[metadata().ndim()], Utils.toIntArray(metadata().shape()), parallel);
191+
}
192+
193+
default boolean chunkIsInArray(long[] chunkCoords) {
194+
final int[] chunkShape = metadata().chunkShape();
195+
for (int dimIdx = 0; dimIdx < metadata().ndim(); dimIdx++) {
196+
if (chunkCoords[dimIdx] < 0
197+
|| chunkCoords[dimIdx] * chunkShape[dimIdx] >= metadata().shape()[dimIdx]) {
198+
return false;
199+
}
200+
}
201+
return true;
202+
}
203+
204+
205+
StoreHandle storeHandle();
206+
207+
CodecPipeline codecPipeline();
208+
209+
/**
210+
* Reads a part of the Zarr array based on a requested offset and shape into an ucar.ma2.Array.
211+
*
212+
* @param offset
213+
* @param shape
214+
* @param parallel
215+
* @throws ZarrException
216+
*/
217+
@Nonnull
218+
default ucar.ma2.Array read(final long[] offset, final int[] shape, final boolean parallel) throws ZarrException {
219+
ArrayMetadata metadata = metadata();
220+
CodecPipeline codecPipeline = codecPipeline();
221+
if (offset.length != metadata.ndim()) {
222+
throw new IllegalArgumentException("'offset' needs to have rank '" + metadata.ndim() + "'.");
223+
}
224+
if (shape.length != metadata.ndim()) {
225+
throw new IllegalArgumentException("'shape' needs to have rank '" + metadata.ndim() + "'.");
226+
}
227+
for (int dimIdx = 0; dimIdx < metadata.ndim(); dimIdx++) {
228+
if (offset[dimIdx] < 0 || offset[dimIdx] + shape[dimIdx] > metadata.shape()[dimIdx]) {
229+
throw new ZarrException("Requested data is outside of the array's domain.");
230+
}
231+
}
232+
233+
final int[] chunkShape = metadata.chunkShape();
234+
if (IndexingUtils.isSingleFullChunk(offset, shape, chunkShape)) {
235+
return readChunk(IndexingUtils.computeSingleChunkCoords(offset, chunkShape));
236+
}
237+
238+
final ucar.ma2.Array outputArray = ucar.ma2.Array.factory(metadata.dataType().getMA2DataType(),
239+
shape);
240+
Stream<long[]> chunkStream = Arrays.stream(IndexingUtils.computeChunkCoords(metadata.shape(), chunkShape, offset, shape));
241+
if (parallel) {
242+
chunkStream = chunkStream.parallel();
243+
}
244+
chunkStream.forEach(
245+
chunkCoords -> {
246+
try {
247+
final IndexingUtils.ChunkProjection chunkProjection =
248+
IndexingUtils.computeProjection(chunkCoords, metadata.shape(), chunkShape, offset,
249+
shape
250+
);
251+
252+
if (chunkIsInArray(chunkCoords)) {
253+
MultiArrayUtils.copyRegion(metadata.allocateFillValueChunk(),
254+
chunkProjection.chunkOffset, outputArray, chunkProjection.outOffset,
255+
chunkProjection.shape
256+
);
257+
}
258+
259+
final String[] chunkKeys = metadata.chunkKeyEncoding().encodeChunkKey(chunkCoords);
260+
final StoreHandle chunkHandle = storeHandle().resolve(chunkKeys);
261+
if (!chunkHandle.exists()) {
262+
return;
263+
}
264+
if (codecPipeline.supportsPartialDecode()) {
265+
final ucar.ma2.Array chunkArray = codecPipeline.decodePartial(chunkHandle,
266+
Utils.toLongArray(chunkProjection.chunkOffset), chunkProjection.shape);
267+
MultiArrayUtils.copyRegion(chunkArray, new int[metadata.ndim()], outputArray,
268+
chunkProjection.outOffset, chunkProjection.shape
269+
);
270+
} else {
271+
MultiArrayUtils.copyRegion(readChunk(chunkCoords), chunkProjection.chunkOffset,
272+
outputArray, chunkProjection.outOffset, chunkProjection.shape
273+
);
274+
}
275+
276+
} catch (ZarrException e) {
277+
throw new RuntimeException(e);
278+
}
279+
});
280+
return outputArray;
281+
}
282+
283+
default ArrayAccessor access() {
284+
return new ArrayAccessor(this);
285+
}
286+
287+
final class ArrayAccessor {
288+
@Nullable
289+
long[] offset;
290+
@Nullable
291+
int[] shape;
292+
@Nonnull
293+
Array array;
294+
295+
public ArrayAccessor(@Nonnull Array array) {
296+
this.array = array;
297+
}
298+
299+
@Nonnull
300+
public ArrayAccessor withOffset(@Nonnull long... offset) {
301+
this.offset = offset;
302+
return this;
303+
}
304+
305+
306+
@Nonnull
307+
public ArrayAccessor withShape(@Nonnull int... shape) {
308+
this.shape = shape;
309+
return this;
310+
}
311+
312+
@Nonnull
313+
public ArrayAccessor withShape(@Nonnull long... shape) {
314+
this.shape = Utils.toIntArray(shape);
315+
return this;
316+
}
317+
318+
@Nonnull
319+
public ucar.ma2.Array read() throws ZarrException {
320+
if (offset == null) {
321+
throw new ZarrException("`offset` needs to be set.");
322+
}
323+
if (shape == null) {
324+
throw new ZarrException("`shape` needs to be set.");
325+
}
326+
return array.read(offset, shape);
327+
}
328+
329+
public void write(@Nonnull ucar.ma2.Array content) throws ZarrException {
330+
if (offset == null) {
331+
throw new ZarrException("`offset` needs to be set.");
332+
}
333+
array.write(offset, content);
334+
}
335+
336+
}
337+
}

‎src/main/java/dev/zarr/zarrjava/interfaces/ArrayMetadata.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
package dev.zarr.zarrjava.interfaces;
2+
3+
import dev.zarr.zarrjava.v3.chunkkeyencoding.ChunkKeyEncoding;
4+
import ucar.ma2.Array;
5+
6+
public interface ArrayMetadata {
7+
int ndim();
8+
9+
int[] chunkShape();
10+
11+
long[] shape();
12+
13+
DataType dataType();
14+
15+
Array allocateFillValueChunk();
16+
17+
ChunkKeyEncoding chunkKeyEncoding();
18+
19+
Object parsedFillValue();
20+
21+
}

‎src/main/java/dev/zarr/zarrjava/interfaces/DataType.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
package dev.zarr.zarrjava.interfaces;
2+
3+
public interface DataType {
4+
ucar.ma2.DataType getMA2DataType();
5+
}

0 commit comments

Comments
 (0)