Add Store.getRanges for reading many byte ranges of one value

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<int[]> 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<int[]> groups = new ArrayList<>();
+        List<Integer> 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;
+    }
+}

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);

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.
+     * <p>
+     * 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);

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<int[]> 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<int[]> 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<int[]> 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<long[]> 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]);
+    }
+}

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;