The two obvious candidates are both wrong here. {@link Array#reshape} always allocates and + * copies. {@link Array#reshapeNoCopy} hands the raw backing store to the new shape and discards + * the input's strides and offset, so it silently reorders the elements of any view — such as + * the output of the {@code transpose} codec, or the strided section that {@code Array.write} + * passes in for every chunk of a multi-chunk write. + * + *
{@link Array#get1DJavaArray} instead returns the backing store itself when the input already + * walks it in lexicographical order (ma2 tracks this as {@code Index.fastIterator}), and a C-order + * copy when it does not. The result therefore shares its storage with the input whenever the + * elements are already laid out in {@code ravel} order, and is a correct copy otherwise. Either + * way {@code ravel(B) == ravel(A)} holds, and the reshaped array is itself in lexicographical + * order, so a following codec gets the cheap path too. + * + *
This is conservative compared to NumPy's {@code _attempt_nocopy_reshape}, which also keeps a + * view when splitting the axes of a strided array, and when the axes being merged happen to be + * internally contiguous ({@code stride[k] == shape[k+1] * stride[k+1]}). Matching that would need + * an {@link ucar.ma2.Index} with custom strides and a non-zero offset, which ma2's public API + * cannot build safely: {@code new Index(shape, stride)} forces {@code offset = 0} and leaves the + * internal {@code fastIterator} flag set, so the resulting array would later hand out its whole + * backing store as if it were the data. Decode is unaffected either way, because it always + * receives a freshly allocated array; on encode the extra copy is limited to multi-chunk writes, + * where the caller passes a strided section. + */ + private static Array reshapeView(Array chunkArray, int[] shape) { + ucar.ma2.DataType dataType = chunkArray.getDataType(); + return Array.factory(dataType, shape, chunkArray.get1DJavaArray(dataType)); } @Override diff --git a/src/test/java/dev/zarr/zarrjava/ReshapeCodecTest.java b/src/test/java/dev/zarr/zarrjava/ReshapeCodecTest.java index 8046c37..d0d45c7 100644 --- a/src/test/java/dev/zarr/zarrjava/ReshapeCodecTest.java +++ b/src/test/java/dev/zarr/zarrjava/ReshapeCodecTest.java @@ -12,6 +12,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; +import ucar.ma2.InvalidRangeException; import ucar.ma2.MAMath; import java.io.IOException; @@ -35,6 +36,20 @@ private static ReshapeCodec reshapeCodec(Object[] shape, int[] inputShape) throw return codec; } + /** + * The elements of {@code array} in lexicographical (C) order. Walks the array with an + * {@link ucar.ma2.IndexIterator} rather than {@code get1DJavaArray}, so that it stays an + * independent oracle for the codec, which uses the latter itself. + */ + private static int[] ravel(ucar.ma2.Array array) { + int[] elements = new int[(int) array.getSize()]; + ucar.ma2.IndexIterator iter = array.getIndexIterator(); + for (int i = 0; i < elements.length; i++) { + elements[i] = iter.getIntNext(); + } + return elements; + } + private static ucar.ma2.Array sequential(int[] shape) { int size = Arrays.stream(shape).reduce(1, (a, b) -> a * b); int[] data = new int[size]; @@ -81,9 +96,7 @@ public void testReshapeRoundTrip(int[] inputShape, Object[] shape, int[] expecte // Output shape matches the specification. Assertions.assertArrayEquals(expectedOutputShape, encoded.getShape()); // The lexicographical (C-order) ravel is preserved: ravel(B) == ravel(A). - Assertions.assertArrayEquals( - (int[]) input.get1DJavaArray(ucar.ma2.DataType.UINT), - (int[]) encoded.get1DJavaArray(ucar.ma2.DataType.UINT)); + Assertions.assertArrayEquals(ravel(input), ravel(encoded)); // decode is the inverse of encode. ucar.ma2.Array decoded = codec.decode(encoded); @@ -156,6 +169,52 @@ public void testReshapeInvalidConfig(int[] inputShape, Object[] shape) { assertThrows(ZarrException.class, () -> reshapeCodec(shape, inputShape)); } + @Test + public void testReshapeConstructsViewForLexicographicalInput() throws ZarrException { + // A freshly allocated array already walks its backing store in lexicographical order, so the + // reshape must be a virtual view: the result has to share storage with the input rather than + // copy it. This is the requirement the specification states as "implementations should, when + // possible, construct a virtual view rather than copy the array". + ReshapeCodec codec = reshapeCodec(new Object[]{new int[]{0, 1}, new int[]{2}}, new int[]{2, 3, 4}); + + ucar.ma2.Array input = sequential(new int[]{2, 3, 4}); + ucar.ma2.Array encoded = codec.encode(input); + + Assertions.assertArrayEquals(new int[]{6, 4}, encoded.getShape()); + Assertions.assertSame(input.getStorage(), encoded.getStorage()); + // The view is itself in lexicographical order, so decoding it stays a view as well. + Assertions.assertSame(input.getStorage(), codec.decode(encoded).getStorage()); + } + + @Test + public void testReshapeCopiesPermutedInputInRavelOrder() throws ZarrException { + // The transpose codec hands on a permuted view, whose elements are not laid out in ravel + // order. No shape/stride/offset descriptor over the original store can express that ravel, so + // a copy is mandatory here -- reshapeNoCopy would silently return the elements in store order. + ReshapeCodec codec = reshapeCodec(new Object[]{-1}, new int[]{4, 3, 2}); + + ucar.ma2.Array permuted = sequential(new int[]{2, 3, 4}).permute(new int[]{2, 1, 0}); + ucar.ma2.Array encoded = codec.encode(permuted); + + Assertions.assertArrayEquals(new int[]{24}, encoded.getShape()); + Assertions.assertNotSame(permuted.getStorage(), encoded.getStorage()); + Assertions.assertArrayEquals(ravel(permuted), ravel(encoded)); + } + + @Test + public void testReshapePreservesRavelOfStridedSection() throws ZarrException, InvalidRangeException { + // Array.write passes a section of the caller's array for every chunk of a multi-chunk write. + // Whether that can stay a view depends on the strides, so only the ravel is asserted here. + ReshapeCodec codec = reshapeCodec(new Object[]{new int[]{0, 1}, new int[]{2}}, new int[]{2, 3, 4}); + + ucar.ma2.Array section = sequential(new int[]{4, 3, 4}) + .sectionNoReduce(new int[]{1, 0, 0}, new int[]{2, 3, 4}, null); + ucar.ma2.Array encoded = codec.encode(section); + + Assertions.assertArrayEquals(new int[]{6, 4}, encoded.getShape()); + Assertions.assertArrayEquals(ravel(section), ravel(encoded)); + } + @Test public void testEncodeRejectsWrongInputShape() throws ZarrException { ReshapeCodec codec = reshapeCodec(new Object[]{new int[]{0, 1}}, new int[]{2, 3});