Skip to content

Commit a435117

Browse files
Copilotbrokkoli71
andcommitted
Add default chunk shape calculation for v3 arrays
Co-authored-by: brokkoli71 <44113112+brokkoli71@users.noreply.github.com>
1 parent 61659ff commit a435117

2 files changed

Lines changed: 84 additions & 2 deletions

File tree

src/main/java/dev/zarr/zarrjava/v3/ArrayMetadataBuilder.java

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -161,13 +161,48 @@ public ArrayMetadata build() throws ZarrException {
161161
if (dataType == null) {
162162
throw new ZarrException("Data type needs to be provided. Please call `.withDataType`.");
163163
}
164+
165+
// If chunk grid is not specified, calculate default chunks
164166
if (chunkGrid == null) {
165-
throw new ZarrException("Chunk grid needs to be provided. Please call `.withChunkShape`.");
167+
int[] defaultChunks = calculateDefaultChunks(shape);
168+
chunkGrid = new RegularChunkGrid(new RegularChunkGrid.Configuration(defaultChunks));
166169
}
170+
167171
return new ArrayMetadata(shape, dataType, chunkGrid, chunkKeyEncoding, fillValue, codecs,
168172
dimensionNames,
169173
attributes,
170174
storageTransformers
171175
);
172176
}
177+
178+
/**
179+
* Calculate default chunk shape when not specified.
180+
* This implements JZarr's ArrayParams.build() logic, targeting chunks of approximately 512 elements.
181+
*
182+
* The algorithm divides each dimension by 512 to determine the number of ~512-sized chunks,
183+
* then calculates chunk sizes that will cover the dimension. Note that the total coverage
184+
* may slightly exceed the dimension size (e.g., for shape=1024, chunks=342 results in
185+
* 3 chunks covering 1026 elements). This is intentional and matches JZarr behavior -
186+
* Zarr handles out-of-bounds gracefully, and the goal is approximate chunk sizes rather
187+
* than perfect tiling.
188+
*/
189+
private int[] calculateDefaultChunks(long[] shape) {
190+
int[] chunks = new int[shape.length];
191+
for (int i = 0; i < shape.length; i++) {
192+
long shapeDim = shape[i];
193+
int numChunks = (int) (shapeDim / 512);
194+
if (numChunks > 0) {
195+
int chunkDim = (int) (shapeDim / (numChunks + 1));
196+
if (shapeDim % chunkDim == 0) {
197+
chunks[i] = chunkDim;
198+
} else {
199+
chunks[i] = chunkDim + 1;
200+
}
201+
} else {
202+
// If dimension is smaller than 512, use the full dimension
203+
chunks[i] = (int) shapeDim;
204+
}
205+
}
206+
return chunks;
207+
}
173208
}

src/test/java/dev/zarr/zarrjava/ZarrV3Test.java

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -773,4 +773,51 @@ public void testUnalignedArrayAccess(int arrayShape, int chunkShape, int accessS
773773
Assertions.assertArrayEquals(expectedData, (int[]) result.get1DJavaArray(ucar.ma2.DataType.UINT));
774774
}
775775
}
776-
}
776+
777+
@Test
778+
public void testDefaultChunkShape() throws IOException, ZarrException {
779+
// Test with a small array (< 512 elements per dimension)
780+
Array smallArray = Array.create(
781+
new FilesystemStore(TESTOUTPUT).resolve("v3_default_chunks_small"),
782+
Array.metadataBuilder()
783+
.withShape(100, 50)
784+
.withDataType(DataType.UINT8)
785+
.build()
786+
);
787+
Assertions.assertEquals(2, smallArray.metadata().chunkShape().length);
788+
// Both dimensions < 512, so chunks should equal shape
789+
Assertions.assertEquals(100, smallArray.metadata().chunkShape()[0]);
790+
Assertions.assertEquals(50, smallArray.metadata().chunkShape()[1]);
791+
792+
// Test with a larger array (> 512 elements per dimension)
793+
Array largeArray = Array.create(
794+
new FilesystemStore(TESTOUTPUT).resolve("v3_default_chunks_large"),
795+
Array.metadataBuilder()
796+
.withShape(2000, 1500)
797+
.withDataType(DataType.UINT8)
798+
.build()
799+
);
800+
Assertions.assertEquals(2, largeArray.metadata().chunkShape().length);
801+
// Chunks should be calculated based on division by 512
802+
Assertions.assertTrue(largeArray.metadata().chunkShape()[0] > 0);
803+
Assertions.assertTrue(largeArray.metadata().chunkShape()[0] < 2000);
804+
Assertions.assertTrue(largeArray.metadata().chunkShape()[1] > 0);
805+
Assertions.assertTrue(largeArray.metadata().chunkShape()[1] < 1500);
806+
807+
// Test with mixed dimensions
808+
Array mixedArray = Array.create(
809+
new FilesystemStore(TESTOUTPUT).resolve("v3_default_chunks_mixed"),
810+
Array.metadataBuilder()
811+
.withShape(1024, 100, 2048)
812+
.withDataType(DataType.UINT8)
813+
.build()
814+
);
815+
Assertions.assertEquals(3, mixedArray.metadata().chunkShape().length);
816+
// Verify chunks are reasonable
817+
Assertions.assertTrue(mixedArray.metadata().chunkShape()[0] > 0);
818+
Assertions.assertTrue(mixedArray.metadata().chunkShape()[0] <= 1024);
819+
Assertions.assertEquals(100, mixedArray.metadata().chunkShape()[1]); // < 512, should equal shape
820+
Assertions.assertTrue(mixedArray.metadata().chunkShape()[2] > 0);
821+
Assertions.assertTrue(mixedArray.metadata().chunkShape()[2] <= 2048);
822+
}
823+
}

0 commit comments

Comments
 (0)