Skip to content

Commit 8a828c7

Browse files
konstibobkonstibob
committed
add reshape codec
1 parent d5552e7 commit 8a828c7

4 files changed

Lines changed: 569 additions & 0 deletions

File tree

‎src/main/java/dev/zarr/zarrjava/v3/codec/CodecBuilder.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,15 @@ public CodecBuilder withTranspose(int[] order) {
6464
return this;
6565
}
6666

67+
/**
68+
* Adds a {@code reshape} codec. Each entry of {@code shape} must be a positive {@link Integer}, the
69+
* special value {@code -1} (at most once), or an {@code int[]} / array of input dimension indices.
70+
*/
71+
public CodecBuilder withReshape(Object[] shape) {
72+
codecs.add(new ReshapeCodec(new ReshapeCodec.Configuration(shape)));
73+
return this;
74+
}
75+
6776
public CodecBuilder withBytes(Endian endian) {
6877
if (dataType.getByteCount() <= 1)
6978
codecs.add(new BytesCodec());

‎src/main/java/dev/zarr/zarrjava/v3/codec/CodecRegistry.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ public class CodecRegistry {
1212

1313
static {
1414
addType("transpose", TransposeCodec.class);
15+
addType("reshape", ReshapeCodec.class);
1516
addType("bytes", BytesCodec.class);
1617
addType("blosc", BloscCodec.class);
1718
addType("gzip", GzipCodec.class);

‎src/main/java/dev/zarr/zarrjava/v3/codec/core/ReshapeCodec.java‎

Lines changed: 306 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,306 @@
1+
package dev.zarr.zarrjava.v3.codec.core;
2+
3+
import com.fasterxml.jackson.annotation.JsonCreator;
4+
import com.fasterxml.jackson.annotation.JsonIgnore;
5+
import com.fasterxml.jackson.annotation.JsonProperty;
6+
import dev.zarr.zarrjava.ZarrException;
7+
import dev.zarr.zarrjava.core.codec.ArrayArrayCodec;
8+
import dev.zarr.zarrjava.v3.ArrayMetadata;
9+
import dev.zarr.zarrjava.v3.codec.Codec;
10+
import ucar.ma2.Array;
11+
12+
import javax.annotation.Nonnull;
13+
import java.util.ArrayList;
14+
import java.util.Arrays;
15+
import java.util.List;
16+
17+
/**
18+
* {@code array -> array} codec that reshapes a chunk.
19+
*
20+
* <p>The codec preserves the (logical) lexicographical / C-order traversal of the elements, i.e.
21+
* {@code ravel(B) == ravel(A)}. It never transposes; combine it with the {@code transpose} codec if
22+
* a reordering is also required.
23+
*
24+
* <p>The target {@code shape} is described relative to the input (chunk) shape {@code A_shape}. Each
25+
* entry is one of:
26+
* <ul>
27+
* <li>a positive integer {@code size}, giving {@code B_shape[i] == size};</li>
28+
* <li>an array of input dimension indices {@code input_dims}, giving
29+
* {@code B_shape[i] == prod(A_shape[input_dims])};</li>
30+
* <li>the special value {@code -1} (at most once), inferred so that
31+
* {@code prod(B_shape) == prod(A_shape)}.</li>
32+
* </ul>
33+
*/
34+
public class ReshapeCodec extends ArrayArrayCodec implements Codec {
35+
36+
@JsonIgnore
37+
@Nonnull
38+
public final String name = "reshape";
39+
@Nonnull
40+
public final Configuration configuration;
41+
42+
@JsonCreator(mode = JsonCreator.Mode.PROPERTIES)
43+
public ReshapeCodec(
44+
@Nonnull @JsonProperty(value = "configuration", required = true) Configuration configuration
45+
) {
46+
this.configuration = configuration;
47+
}
48+
49+
@Override
50+
public Array encode(Array chunkArray) throws ZarrException {
51+
int[] inputShape = arrayMetadata.chunkShape;
52+
if (!Arrays.equals(chunkArray.getShape(), inputShape)) {
53+
throw new ZarrException(
54+
"reshape codec received an array of shape " + Arrays.toString(chunkArray.getShape())
55+
+ " but expected the chunk shape " + Arrays.toString(inputShape) + ".");
56+
}
57+
int[] outputShape = resolveOutputShape(inputShape);
58+
// Array.reshape copies the elements in lexicographical (C) order, hence ravel(B) == ravel(A)
59+
// even when the input array is a non-contiguous view.
60+
return chunkArray.reshape(outputShape);
61+
}
62+
63+
@Override
64+
public Array decode(Array chunkArray) throws ZarrException {
65+
int[] inputShape = arrayMetadata.chunkShape;
66+
int[] outputShape = resolveOutputShape(inputShape);
67+
if (!Arrays.equals(chunkArray.getShape(), outputShape)) {
68+
throw new ZarrException(
69+
"reshape codec received an array of shape " + Arrays.toString(chunkArray.getShape())
70+
+ " but expected the reshaped shape " + Arrays.toString(outputShape) + ".");
71+
}
72+
// Inverse operation: reshape back to the original chunk shape.
73+
return chunkArray.reshape(inputShape);
74+
}
75+
76+
@Override
77+
public long computeEncodedSize(long inputByteLength,
78+
ArrayMetadata.CoreArrayMetadata arrayMetadata) throws ZarrException {
79+
// Reshaping does not add or remove any elements.
80+
return inputByteLength;
81+
}
82+
83+
@Override
84+
public ArrayMetadata.CoreArrayMetadata resolveArrayMetadata() throws ZarrException {
85+
super.resolveArrayMetadata();
86+
87+
int[] inputChunkShape = arrayMetadata.chunkShape;
88+
int[] outputChunkShape = resolveOutputShape(inputChunkShape);
89+
90+
// Derive the output (grid) array shape so that number of chunks along each output dimension is
91+
// consistent with the input. For a merged output dimension this is the product of the input
92+
// chunk counts; a split input dimension keeps its chunk count on the outermost output dimension.
93+
long[] inputArrayShape = arrayMetadata.shape;
94+
long[] gridMultiplier = new long[outputChunkShape.length];
95+
Arrays.fill(gridMultiplier, 1L);
96+
97+
long[] outputStart = new long[outputChunkShape.length + 1];
98+
outputStart[0] = 1L;
99+
for (int i = 0; i < outputChunkShape.length; i++) {
100+
outputStart[i + 1] = outputStart[i] * outputChunkShape[i];
101+
}
102+
103+
long inputStart = 1L;
104+
for (int d = 0; d < inputChunkShape.length; d++) {
105+
long numChunks = inputArrayShape[d] / inputChunkShape[d];
106+
// Attach the chunk count of input dimension d to the output dimension whose flat range
107+
// contains the start of input dimension d.
108+
int target = outputChunkShape.length - 1;
109+
for (int i = 0; i < outputChunkShape.length; i++) {
110+
if (inputStart >= outputStart[i] && inputStart < outputStart[i + 1]) {
111+
target = i;
112+
break;
113+
}
114+
}
115+
gridMultiplier[target] *= numChunks;
116+
inputStart *= inputChunkShape[d];
117+
}
118+
119+
long[] outputArrayShape = new long[outputChunkShape.length];
120+
for (int i = 0; i < outputChunkShape.length; i++) {
121+
outputArrayShape[i] = gridMultiplier[i] * outputChunkShape[i];
122+
}
123+
124+
return new ArrayMetadata.CoreArrayMetadata(
125+
outputArrayShape,
126+
outputChunkShape,
127+
arrayMetadata.dataType,
128+
arrayMetadata.parsedFillValue
129+
);
130+
}
131+
132+
/**
133+
* Resolves the concrete output shape {@code B_shape} from the {@code shape} configuration and the
134+
* given input shape {@code A_shape}, validating all invariants mandated by the specification.
135+
*/
136+
int[] resolveOutputShape(int[] inputShape) throws ZarrException {
137+
int ndim = inputShape.length;
138+
long inputTotal = 1L;
139+
for (int size : inputShape) {
140+
inputTotal *= size;
141+
}
142+
143+
int n = configuration.shape.length;
144+
if (n == 0) {
145+
throw new ZarrException("reshape codec: 'shape' must not be empty.");
146+
}
147+
148+
long[] outputShape = new long[n];
149+
int[][] inputDimsPerOutput = new int[n][];
150+
int minusOnePos = -1;
151+
List<Integer> flatInputDims = new ArrayList<>();
152+
153+
for (int i = 0; i < n; i++) {
154+
Object element = configuration.shape[i];
155+
int[] inputDims = asInputDims(element);
156+
if (inputDims == null) {
157+
int size = asInteger(element);
158+
if (size == -1) {
159+
if (minusOnePos != -1) {
160+
throw new ZarrException("reshape codec: 'shape' may contain -1 at most once.");
161+
}
162+
minusOnePos = i;
163+
outputShape[i] = -1;
164+
} else if (size <= 0) {
165+
throw new ZarrException(
166+
"reshape codec: 'shape' entries must be a positive integer, -1, or an array of "
167+
+ "input dimensions, but got " + size + ".");
168+
} else {
169+
outputShape[i] = size;
170+
}
171+
} else {
172+
inputDimsPerOutput[i] = inputDims;
173+
long product = 1L;
174+
for (int d : inputDims) {
175+
if (d < 0 || d >= ndim) {
176+
throw new ZarrException(
177+
"reshape codec: input dimension " + d + " is out of range for an input array with "
178+
+ ndim + " dimensions.");
179+
}
180+
product *= inputShape[d];
181+
flatInputDims.add(d);
182+
}
183+
outputShape[i] = product;
184+
}
185+
}
186+
187+
// The flattened list of input dimensions must be strictly monotonically increasing. This rules
188+
// out configurations that would (incorrectly) suggest a transpose.
189+
for (int j = 1; j < flatInputDims.size(); j++) {
190+
if (flatInputDims.get(j) <= flatInputDims.get(j - 1)) {
191+
throw new ZarrException(
192+
"reshape codec: the flattened list of input dimensions must be strictly increasing, but got "
193+
+ flatInputDims + ".");
194+
}
195+
}
196+
197+
// Resolve the automatic dimension (-1).
198+
if (minusOnePos != -1) {
199+
long known = 1L;
200+
for (int i = 0; i < n; i++) {
201+
if (i != minusOnePos) {
202+
known *= outputShape[i];
203+
}
204+
}
205+
if (known == 0 || inputTotal % known != 0) {
206+
throw new ZarrException(
207+
"reshape codec: cannot infer the -1 dimension because prod(output shape) would not equal "
208+
+ "prod(input shape) (" + inputTotal + ").");
209+
}
210+
outputShape[minusOnePos] = inputTotal / known;
211+
}
212+
213+
// Invariant: prod(B_shape) == prod(A_shape).
214+
long outputTotal = 1L;
215+
for (long size : outputShape) {
216+
outputTotal *= size;
217+
}
218+
if (outputTotal != inputTotal) {
219+
throw new ZarrException(
220+
"reshape codec: prod(output shape)=" + outputTotal + " does not equal prod(input shape)="
221+
+ inputTotal + ".");
222+
}
223+
224+
// When an output dimension is specified by input_dims, verify that the coordinates along those
225+
// input dimensions actually correspond to the raveled index along that output dimension.
226+
for (int i = 0; i < n; i++) {
227+
int[] inputDims = inputDimsPerOutput[i];
228+
if (inputDims == null || inputDims.length == 0) {
229+
continue;
230+
}
231+
long outputPrefix = 1L;
232+
for (int j = 0; j < i; j++) {
233+
outputPrefix *= outputShape[j];
234+
}
235+
long outputSuffix = 1L;
236+
for (int j = i + 1; j < n; j++) {
237+
outputSuffix *= outputShape[j];
238+
}
239+
long inputPrefix = 1L;
240+
for (int d = 0; d < inputDims[0]; d++) {
241+
inputPrefix *= inputShape[d];
242+
}
243+
long inputSuffix = 1L;
244+
for (int d = inputDims[inputDims.length - 1] + 1; d < ndim; d++) {
245+
inputSuffix *= inputShape[d];
246+
}
247+
if (outputPrefix != inputPrefix || outputSuffix != inputSuffix) {
248+
throw new ZarrException(
249+
"reshape codec: output dimension " + i + " specified by input dimensions "
250+
+ Arrays.toString(inputDims) + " does not align with the raveled input array "
251+
+ "(prefix " + outputPrefix + " vs " + inputPrefix + ", suffix " + outputSuffix
252+
+ " vs " + inputSuffix + ").");
253+
}
254+
}
255+
256+
int[] result = new int[n];
257+
for (int i = 0; i < n; i++) {
258+
if (outputShape[i] > Integer.MAX_VALUE) {
259+
throw new ZarrException("reshape codec: output dimension " + i + " exceeds Integer.MAX_VALUE.");
260+
}
261+
result[i] = (int) outputShape[i];
262+
}
263+
return result;
264+
}
265+
266+
private static int[] asInputDims(Object element) throws ZarrException {
267+
if (element instanceof int[]) {
268+
return (int[]) element;
269+
}
270+
if (element instanceof Object[]) {
271+
Object[] array = (Object[]) element;
272+
int[] dims = new int[array.length];
273+
for (int i = 0; i < array.length; i++) {
274+
dims[i] = asInteger(array[i]);
275+
}
276+
return dims;
277+
}
278+
if (element instanceof List) {
279+
List<?> list = (List<?>) element;
280+
int[] dims = new int[list.size()];
281+
for (int i = 0; i < list.size(); i++) {
282+
dims[i] = asInteger(list.get(i));
283+
}
284+
return dims;
285+
}
286+
return null;
287+
}
288+
289+
private static int asInteger(Object element) throws ZarrException {
290+
if (element instanceof Number) {
291+
return ((Number) element).intValue();
292+
}
293+
throw new ZarrException(
294+
"reshape codec: 'shape' entries must be integers or arrays of integers, but got " + element + ".");
295+
}
296+
297+
public static final class Configuration {
298+
@Nonnull
299+
public final Object[] shape;
300+
301+
@JsonCreator(mode = JsonCreator.Mode.PROPERTIES)
302+
public Configuration(@Nonnull @JsonProperty(value = "shape", required = true) Object[] shape) {
303+
this.shape = shape;
304+
}
305+
}
306+
}

0 commit comments

Comments
 (0)