diff --git a/src/main/java/dev/zarr/zarrjava/store/ByteRangeCoalescer.java b/src/main/java/dev/zarr/zarrjava/store/ByteRangeCoalescer.java new file mode 100644 index 00000000..7100ae69 --- /dev/null +++ b/src/main/java/dev/zarr/zarrjava/store/ByteRangeCoalescer.java @@ -0,0 +1,73 @@ +package dev.zarr.zarrjava.store; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Comparator; +import java.util.List; + +/** + * Plans which byte ranges of one value can be fetched together, mirroring zarr-python's + * {@code zarr.core._coalesce.coalesce_ranges}. + */ +final class ByteRangeCoalescer { + + private ByteRangeCoalescer() { + } + + /** + * Checks that {@code starts} and {@code ends} describe well-formed ranges {@code [starts[i], ends[i])}. + */ + static void validate(long[] starts, long[] ends) { + if (starts.length != ends.length) { + throw new IllegalArgumentException("'starts' and 'ends' need to have the same length."); + } + for (int i = 0; i < starts.length; i++) { + if (starts[i] < 0 || ends[i] < starts[i]) { + throw new IllegalArgumentException( + String.format("Invalid byte range [%d, %d).", starts[i], ends[i])); + } + } + } + + /** + * Groups the ranges {@code [starts[i], ends[i])} so that each group can be served by a single + * fetch. Ranges are sorted by start; a range joins the current group if the gap to the group's + * running end is at most {@code maxGapBytes} and the merged span stays within + * {@code maxCoalescedBytes}. + * + * @return the groups as lists of input indices, each sorted by start + */ + static List coalesce(long[] starts, long[] ends, long maxGapBytes, long maxCoalescedBytes) { + Integer[] order = new Integer[starts.length]; + for (int i = 0; i < order.length; i++) { + order[i] = i; + } + Arrays.sort(order, Comparator.comparingLong(i -> starts[i])); + + List groups = new ArrayList<>(); + List group = new ArrayList<>(); + long groupStart = 0; + long groupEnd = 0; + for (int i : order) { + if (!group.isEmpty() && starts[i] - groupEnd <= maxGapBytes) { + long prospectiveEnd = Math.max(groupEnd, ends[i]); + if (prospectiveEnd - groupStart <= maxCoalescedBytes) { + group.add(i); + groupEnd = prospectiveEnd; + continue; + } + } + if (!group.isEmpty()) { + groups.add(group.stream().mapToInt(Integer::intValue).toArray()); + } + group = new ArrayList<>(); + group.add(i); + groupStart = starts[i]; + groupEnd = ends[i]; + } + if (!group.isEmpty()) { + groups.add(group.stream().mapToInt(Integer::intValue).toArray()); + } + return groups; + } +} diff --git a/src/main/java/dev/zarr/zarrjava/store/FilesystemStore.java b/src/main/java/dev/zarr/zarrjava/store/FilesystemStore.java index 3644121c..5982d982 100644 --- a/src/main/java/dev/zarr/zarrjava/store/FilesystemStore.java +++ b/src/main/java/dev/zarr/zarrjava/store/FilesystemStore.java @@ -102,6 +102,34 @@ public ByteBuffer get(String[] keys, long start, long end) { } + /** + * Opens the file once and reads each range from it directly. Coalescing buys nothing locally, + * so {@code maxGapBytes} and {@code maxCoalescedBytes} are ignored. + */ + @Nullable + @Override + public ByteBuffer[] getRanges(String[] keys, long[] starts, long[] ends, + long maxGapBytes, long maxCoalescedBytes) { + ByteRangeCoalescer.validate(starts, ends); + try (SeekableByteChannel byteChannel = Files.newByteChannel(resolveKeys(keys))) { + ByteBuffer[] result = new ByteBuffer[starts.length]; + for (int i = 0; i < starts.length; i++) { + ByteBuffer bytes = Utils.allocateNative((int) (ends[i] - starts[i])); + byteChannel.position(starts[i]); + while (bytes.hasRemaining() && byteChannel.read(bytes) >= 0) { + // read() may return fewer bytes than requested + } + bytes.rewind(); + result[i] = bytes; + } + return result; + } catch (NoSuchFileException e) { + return null; + } catch (IOException e) { + throw StoreException.readFailed(this.toString(), keys, e); + } + } + @Override public void set(String[] keys, ByteBuffer bytes) { Path keyPath = resolveKeys(keys); diff --git a/src/main/java/dev/zarr/zarrjava/store/Store.java b/src/main/java/dev/zarr/zarrjava/store/Store.java index 7d478e57..c13b28aa 100644 --- a/src/main/java/dev/zarr/zarrjava/store/Store.java +++ b/src/main/java/dev/zarr/zarrjava/store/Store.java @@ -4,6 +4,7 @@ import javax.annotation.Nullable; import java.io.InputStream; import java.nio.ByteBuffer; +import java.util.Arrays; import java.util.stream.Stream; public interface Store { @@ -19,6 +20,65 @@ public interface Store { @Nullable ByteBuffer get(String[] keys, long start, long end); + /** + * Default for the {@code maxGapBytes} argument of {@link #getRanges(String[], long[], long[], long, long)} + * (1 MiB, as in zarr-python). + */ + long DEFAULT_MAX_GAP_BYTES = 1 << 20; + + /** + * Default for the {@code maxCoalescedBytes} argument of + * {@link #getRanges(String[], long[], long[], long, long)} (16 MiB, as in zarr-python). + */ + long DEFAULT_MAX_COALESCED_BYTES = 16 << 20; + + /** + * Reads many byte ranges of the value at the given keys, using the default coalescing limits. + * + * @see #getRanges(String[], long[], long[], long, long) + */ + @Nullable + default ByteBuffer[] getRanges(String[] keys, long[] starts, long[] ends) { + return getRanges(keys, starts, ends, DEFAULT_MAX_GAP_BYTES, DEFAULT_MAX_COALESCED_BYTES); + } + + /** + * Reads many byte ranges {@code [starts[i], ends[i])} of the value at the given keys. + *

+ * The default implementation coalesces nearby ranges into fewer {@link #get(String[], long, long)} + * calls and slices the results back apart. Stores with a cheaper way to read many ranges may + * override it. + * + * @param keys the keys identifying the value + * @param starts the start offsets (inclusive) of the ranges + * @param ends the end offsets (exclusive) of the ranges + * @param maxGapBytes two ranges separated by at most this many bytes may be fetched together + * @param maxCoalescedBytes upper bound on the size of a single coalesced fetch + * @return one buffer per range, in input order, or null if the value does not exist + */ + @Nullable + default ByteBuffer[] getRanges(String[] keys, long[] starts, long[] ends, + long maxGapBytes, long maxCoalescedBytes) { + ByteRangeCoalescer.validate(starts, ends); + ByteBuffer[] result = new ByteBuffer[starts.length]; + for (int[] group : ByteRangeCoalescer.coalesce(starts, ends, maxGapBytes, maxCoalescedBytes)) { + long groupStart = starts[group[0]]; + long groupEnd = Arrays.stream(group).mapToLong(i -> ends[i]).max().getAsLong(); + ByteBuffer groupBytes = get(keys, groupStart, groupEnd); + if (groupBytes == null) { + return null; + } + for (int i : group) { + // clamp to what was returned, in case the value ends before groupEnd + ByteBuffer slice = groupBytes.duplicate(); + slice.position((int) Math.min(groupBytes.position() + (starts[i] - groupStart), groupBytes.limit())); + slice.limit((int) Math.min(slice.position() + (ends[i] - starts[i]), groupBytes.limit())); + result[i] = slice.slice().order(groupBytes.order()); + } + } + return result; + } + void set(String[] keys, ByteBuffer bytes); void delete(String[] keys); diff --git a/src/test/java/dev/zarr/zarrjava/store/ByteRangeCoalescerTest.java b/src/test/java/dev/zarr/zarrjava/store/ByteRangeCoalescerTest.java new file mode 100644 index 00000000..930d163e --- /dev/null +++ b/src/test/java/dev/zarr/zarrjava/store/ByteRangeCoalescerTest.java @@ -0,0 +1,84 @@ +package dev.zarr.zarrjava.store; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; + +public class ByteRangeCoalescerTest { + + @Test + public void testMergesNearbyRangesInStartOrder() { + List groups = ByteRangeCoalescer.coalesce( + new long[]{9000, 100, 210, 305}, new long[]{9100, 200, 300, 400}, 10, 1 << 20); + Assertions.assertEquals(2, groups.size()); + Assertions.assertArrayEquals(new int[]{1, 2, 3}, groups.get(0)); + Assertions.assertArrayEquals(new int[]{0}, groups.get(1)); + } + + @Test + public void testRespectsGapLimit() { + // gap of exactly maxGapBytes merges, one byte more does not + Assertions.assertEquals(1, ByteRangeCoalescer.coalesce( + new long[]{0, 20}, new long[]{10, 30}, 10, 1 << 20).size()); + Assertions.assertEquals(2, ByteRangeCoalescer.coalesce( + new long[]{0, 21}, new long[]{10, 31}, 10, 1 << 20).size()); + } + + @Test + public void testRespectsSizeLimit() { + List groups = ByteRangeCoalescer.coalesce( + new long[]{0, 10, 20}, new long[]{10, 20, 30}, 0, 20); + Assertions.assertEquals(2, groups.size()); + Assertions.assertArrayEquals(new int[]{0, 1}, groups.get(0)); + Assertions.assertArrayEquals(new int[]{2}, groups.get(1)); + } + + @Test + public void testMergesOverlappingRanges() { + List groups = ByteRangeCoalescer.coalesce( + new long[]{0, 5, 2}, new long[]{50, 10, 60}, 0, 1 << 20); + Assertions.assertEquals(1, groups.size()); + Assertions.assertArrayEquals(new int[]{0, 2, 1}, groups.get(0)); + } + + @Test + public void testRejectsInvalidRanges() { + Assertions.assertThrows(IllegalArgumentException.class, + () -> ByteRangeCoalescer.validate(new long[]{0}, new long[]{10, 20})); + Assertions.assertThrows(IllegalArgumentException.class, + () -> ByteRangeCoalescer.validate(new long[]{-1}, new long[]{10})); + Assertions.assertThrows(IllegalArgumentException.class, + () -> ByteRangeCoalescer.validate(new long[]{10}, new long[]{5})); + } + + @Test + public void testDefaultGetRangesIssuesOneGetPerGroup() { + List calls = new ArrayList<>(); + MemoryStore store = new MemoryStore() { + @Override + public ByteBuffer get(String[] keys, long start, long end) { + calls.add(new long[]{start, end}); + return super.get(keys, start, end); + } + }; + byte[] data = new byte[10_000]; + for (int i = 0; i < data.length; i++) { + data[i] = (byte) i; + } + store.set(new String[]{"shard"}, ByteBuffer.wrap(data)); + + ByteBuffer[] ranges = store.getRanges(new String[]{"shard"}, + new long[]{9000, 100, 210, 305}, new long[]{9100, 200, 300, 400}, 10, 1 << 20); + + Assertions.assertEquals(2, calls.size()); + Assertions.assertArrayEquals(new long[]{100, 400}, calls.get(0)); + Assertions.assertArrayEquals(new long[]{9000, 9100}, calls.get(1)); + Assertions.assertEquals(ByteBuffer.wrap(data, 9000, 100), ranges[0]); + Assertions.assertEquals(ByteBuffer.wrap(data, 100, 100), ranges[1]); + Assertions.assertEquals(ByteBuffer.wrap(data, 210, 90), ranges[2]); + Assertions.assertEquals(ByteBuffer.wrap(data, 305, 95), ranges[3]); + } +} diff --git a/src/test/java/dev/zarr/zarrjava/store/StoreTest.java b/src/test/java/dev/zarr/zarrjava/store/StoreTest.java index 4f1c80a1..94899f27 100644 --- a/src/test/java/dev/zarr/zarrjava/store/StoreTest.java +++ b/src/test/java/dev/zarr/zarrjava/store/StoreTest.java @@ -128,6 +128,39 @@ public void testGetWithStartEnd() { Assertions.assertArrayEquals(expectedBytes, actualBytes); } + @Test + public void testGetRanges() { + StoreHandle storeHandle = storeHandleWithData(); + long size = storeHandle.getSize(); + if (size < 20) { + Assertions.fail("Store size is too small to test getRanges"); + } + // unsorted, adjacent, overlapping and far-apart ranges + long[] starts = {size - 5, 4, 8, 6, 0, size / 2}; + long[] ends = {size, 8, 12, 10, 2, size / 2 + 3}; + ByteBuffer[] ranges = storeHandle.store.getRanges(storeHandle.keys, starts, ends); + Assertions.assertNotNull(ranges); + Assertions.assertEquals(starts.length, ranges.length); + for (int i = 0; i < starts.length; i++) { + Assertions.assertEquals(storeHandle.read(starts[i], ends[i]), ranges[i], + "Range " + i + " differs from get(start, end)"); + } + + // a gap of 0 bytes still allows merging adjacent ranges, a size limit of 1 byte forbids any merge + for (long[] limits : new long[][]{{0, Store.DEFAULT_MAX_COALESCED_BYTES}, {Store.DEFAULT_MAX_GAP_BYTES, 1}}) { + ranges = storeHandle.store.getRanges(storeHandle.keys, starts, ends, limits[0], limits[1]); + for (int i = 0; i < starts.length; i++) { + Assertions.assertEquals(storeHandle.read(starts[i], ends[i]), ranges[i]); + } + } + } + + @Test + public void testGetRangesMissingKey() { + StoreHandle storeHandle = storeHandleWithoutData(); + Assertions.assertNull(storeHandle.store.getRanges(storeHandle.keys, new long[]{0}, new long[]{10})); + } + @Test public abstract void testList() throws ZarrException, IOException;