inherit zipstores from common parent and reduce buffers in memory in loadBuffer

brokkoli71 · brokkoli71 · commit 7d4f4873e31f · 2025-12-18T18:00:40.000+01:00
diff --git a/src/main/java/dev/zarr/zarrjava/store/BufferedZipStore.java b/src/main/java/dev/zarr/zarrjava/store/BufferedZipStore.java
@@ -8,58 +8,53 @@
 import java.nio.ByteBuffer;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Comparator;
 import java.util.stream.Stream;
 
-import com.fasterxml.jackson.databind.util.ByteBufferBackedInputStream;
 import org.apache.commons.compress.archivers.zip.*;
 
 import java.util.zip.CRC32;
 import java.util.zip.ZipEntry; // for STORED constant
 
-import static dev.zarr.zarrjava.utils.ZipUtils.getZipCommentFromBuffer;
-
 
 /** A Store implementation that buffers reads and writes and flushes them to an underlying Store as a zip file.
  */
-public class BufferedZipStore implements Store, Store.ListableStore {
+public class BufferedZipStore extends ZipStore {
 
-    private final StoreHandle underlyingStore;
     private final Store.ListableStore bufferStore;
     private String archiveComment;
-    private boolean flushOnWrite;
+    private final boolean flushOnWrite;
+
+    private final Comparator<String[]> zipEntryComparator = (a, b) -> {
+        boolean aIsZarr = a.length > 0 && a[a.length - 1].equals("zarr.json");
+        boolean bIsZarr = b.length > 0 && b[b.length - 1].equals("zarr.json");
+        // first all zarr.json files
+        if (aIsZarr && !bIsZarr) {
+            return -1;
+        } else if (!aIsZarr && bIsZarr) {
+            return 1;
+        } else if (aIsZarr && bIsZarr) {
+            // sort zarr.json in BFS order within same depth by lexicographical order
+            if (a.length != b.length) {
+                return Integer.compare(a.length, b.length);
+            } else {
+                return String.join("/", a).compareTo(String.join("/", b));
+            }
+        } else {
+            // then all other files in lexicographical order
+            return String.join("/", a).compareTo(String.join("/", b));
+        }
+    };
 
-    private void writeBuffer() throws IOException{
+    private void writeBuffer() throws IOException {
         // create zip file bytes from buffer store and write to underlying store
         ByteArrayOutputStream baos = new ByteArrayOutputStream();
         try (ZipArchiveOutputStream zos = new ZipArchiveOutputStream(baos)) {
             zos.setUseZip64(Zip64Mode.AsNeeded);
             if (archiveComment != null) {
                 zos.setComment(archiveComment);
             }
-            Stream<String[]> entries = bufferStore.list().sorted(
-                    (a, b) -> {
-                        boolean aIsZarr = a.length > 0 && a[a.length - 1].equals("zarr.json");
-                        boolean bIsZarr = b.length > 0 && b[b.length - 1].equals("zarr.json");
-                        // first all zarr.json files
-                        if (aIsZarr && !bIsZarr) {
-                            return -1;
-                        } else if (!aIsZarr && bIsZarr) {
-                            return 1;
-                        } else if (aIsZarr && bIsZarr) {
-                            // sort zarr.json in BFS order within same depth by lexicographical order
-                            if (a.length != b.length) {
-                                return Integer.compare(a.length, b.length);
-                            } else {
-                                return String.join("/", a).compareTo(String.join("/", b));
-                            }
-                        } else {
-                            // then all other files in lexicographical order
-                            return String.join("/", a).compareTo(String.join("/", b));
-                        }
-                    }
-            );
-
-            entries.forEach(keys -> {
+            bufferStore.list().sorted(zipEntryComparator).forEach(keys -> {
                 try {
                     if (keys == null || keys.length == 0) {
                         // skip root entry
@@ -116,22 +111,32 @@ private void writeBuffer() throws IOException{
         underlyingStore.set(ByteBuffer.wrap(zipBytes));
     }
 
+    public void setArchiveComment(@Nullable String archiveComment) throws IOException {
+        this.archiveComment = archiveComment;
+        if (flushOnWrite) {
+            writeBuffer();
+        }
+    }
+
+    public void deleteArchiveComment() throws IOException {
+        this.setArchiveComment(null);
+    }
 
-    private void loadBuffer() throws IOException{
-        // read zip file bytes from underlying store and populate buffer store
-        ByteBuffer buffer = underlyingStore.read();
-        if (buffer == null) {
-            return;
+    /**
+     * Loads the buffer from the underlying store zip file.
+     */
+    private void loadBuffer() throws IOException {
+        String loadedArchiveComment = super.getArchiveComment();
+        if (loadedArchiveComment != null && this.archiveComment == null) {
+            // don't overwrite existing archiveComment
+            this.archiveComment = loadedArchiveComment;
         }
-        byte[] bufArray;
-        if (buffer.hasArray()) {
-            bufArray = buffer.array();
-        } else {
-            bufArray = new byte[buffer.remaining()];
-            buffer.duplicate().get(bufArray);
+
+        InputStream inputStream = underlyingStore.getInputStream();
+        if (inputStream == null) {
+            return;
         }
-        this.archiveComment = getZipCommentFromBuffer(bufArray);
-        try (ZipArchiveInputStream zis = new ZipArchiveInputStream(new ByteBufferBackedInputStream(buffer))) {
+        try (ZipArchiveInputStream zis = new ZipArchiveInputStream(inputStream)) {
             ZipArchiveEntry entry;
             while ((entry = zis.getNextEntry()) != null) {
                 if (entry.isDirectory()) {
@@ -150,7 +155,7 @@ private void loadBuffer() throws IOException{
     }
 
     public BufferedZipStore(@Nonnull StoreHandle underlyingStore, @Nonnull Store.ListableStore bufferStore, @Nullable String archiveComment, boolean flushOnWrite) {
-        this.underlyingStore = underlyingStore;
+        super(underlyingStore);
         this.bufferStore = bufferStore;
         this.archiveComment = archiveComment;
         this.flushOnWrite = flushOnWrite;
@@ -229,6 +234,7 @@ public void flush() throws IOException {
         writeBuffer();
     }
 
+    @Override
     public String getArchiveComment() {
         return archiveComment;
     }
diff --git a/src/main/java/dev/zarr/zarrjava/store/FilesystemStore.java b/src/main/java/dev/zarr/zarrjava/store/FilesystemStore.java
@@ -153,6 +153,9 @@ public String toString() {
     public InputStream getInputStream(String[] keys, long start, long end) {
         Path keyPath = resolveKeys(keys);
         try {
+            if (!Files.exists(keyPath)) {
+                return null;
+            }
             InputStream inputStream = Files.newInputStream(keyPath);
             if (start > 0) {
                 long skipped = inputStream.skip(start);
@@ -170,6 +173,7 @@ public InputStream getInputStream(String[] keys, long start, long end) {
             throw new RuntimeException(e);
         }
     }
+
     public long getSize(String[] keys) {
         try {
             return Files.size(resolveKeys(keys));
diff --git a/src/main/java/dev/zarr/zarrjava/store/ReadOnlyZipStore.java b/src/main/java/dev/zarr/zarrjava/store/ReadOnlyZipStore.java
@@ -15,16 +15,12 @@
 import java.nio.file.Paths;
 import java.util.stream.Stream;
 
-import static dev.zarr.zarrjava.utils.ZipUtils.getZipCommentFromBuffer;
-
 
 /** A Store implementation that provides read-only access to a zip archive stored in an underlying Store.
  * Compared to BufferedZipStore, this implementation reads directly from the zip archive without parsing
  * its contents into a buffer store first making it more efficient for read-only access to large zip archives.
  */
-public class ReadOnlyZipStore implements Store, Store.ListableStore {
-
-    private final StoreHandle underlyingStore;
+public class ReadOnlyZipStore extends ZipStore {
 
     String resolveKeys(String[] keys) {
         return String.join("/", keys);
@@ -51,42 +47,6 @@ public ByteBuffer get(String[] keys, long start) {
         return get(keys, start, -1);
     }
 
-    public String getArchiveComment() throws IOException {
-        // Attempt to read from the end of the file to find the EOCD record.
-        // We try a small chunk first (1KB) which covers most short comments (or no comment),
-        // then the maximum possible EOCD size (approx 65KB).
-        int[] readSizes = {1024, 65535 + 22};
-
-        for (int size : readSizes) {
-            ByteBuffer buffer;
-            long fileSize = underlyingStore.getSize();
-
-            if (fileSize < size){
-                buffer = underlyingStore.read();
-            }
-            else {
-                buffer = underlyingStore.read(fileSize - size);
-            }
-
-            if (buffer == null) {
-                return null;
-            }
-
-            byte[] bufArray;
-            if (buffer.hasArray()) {
-                bufArray = buffer.array();
-            } else {
-                bufArray = new byte[buffer.remaining()];
-                buffer.duplicate().get(bufArray);
-            }
-
-            String comment = getZipCommentFromBuffer(bufArray);
-            if (comment != null) {
-                return comment;
-            }
-        }
-        return null;
-    }
     @Nullable
     @Override
     public ByteBuffer get(String[] keys, long start, long end) {
@@ -153,7 +113,7 @@ public String toString() {
     }
 
     public ReadOnlyZipStore(@Nonnull StoreHandle underlyingStore) {
-        this.underlyingStore = underlyingStore;
+        super(underlyingStore);
     }
 
     public ReadOnlyZipStore(@Nonnull Path underlyingStore) {
diff --git a/src/main/java/dev/zarr/zarrjava/store/ZipStore.java b/src/main/java/dev/zarr/zarrjava/store/ZipStore.java
@@ -0,0 +1,81 @@
+package dev.zarr.zarrjava.store;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+public abstract class ZipStore implements Store, Store.ListableStore {
+    protected final StoreHandle underlyingStore;
+
+    public ZipStore(@Nonnull StoreHandle underlyingStore) {
+        this.underlyingStore = underlyingStore;
+    }
+
+    public String getArchiveComment() throws IOException {
+        // Attempt to read from the end of the file to find the EOCD record.
+        // We try a small chunk first (1KB) which covers most short comments (or no comment),
+        // then the maximum possible EOCD size (approx 65KB).
+        if (!underlyingStore.exists()) {
+            return null;
+        }
+        int[] readSizes = {1024, 65535 + 22};
+
+        for (int size : readSizes) {
+            ByteBuffer buffer;
+            long fileSize = underlyingStore.getSize();
+
+            if (fileSize < size){
+                buffer = underlyingStore.read();
+            }
+            else {
+                buffer = underlyingStore.read(fileSize - size);
+            }
+
+            if (buffer == null) {
+                return null;
+            }
+
+            byte[] bufArray;
+            if (buffer.hasArray()) {
+                bufArray = buffer.array();
+            } else {
+                bufArray = new byte[buffer.remaining()];
+                buffer.duplicate().get(bufArray);
+            }
+
+            String comment = getZipCommentFromBuffer(bufArray);
+            if (comment != null) {
+                return comment;
+            }
+        }
+        return null;
+    }
+
+    // adopted from https://stackoverflow.com/a/9918966
+    @Nullable
+    public static String getZipCommentFromBuffer(byte[] bufArray) throws IOException {
+        // End of Central Directory (EOCD) record magic number
+        byte[]  EOCD = {0x50, 0x4b, 0x05, 0x06};
+        int buffLen = bufArray.length;
+        // Check the buffer from the end
+        search:
+        for (int i = buffLen - EOCD.length - 22; i >= 0; i--) {
+            for (int k = 0; k < EOCD.length; k++) {
+                if (bufArray[i + k] != EOCD[k]) {
+                    continue search;
+                }
+            }
+            // End of Central Directory found!
+            int commentLen = bufArray[i + 20] + bufArray[i + 21] * 256;
+            int realLen = buffLen - i - 22;
+            if (commentLen != realLen) {
+                throw new IOException("ZIP comment size mismatch: "
+                        + "directory says len is " + commentLen
+                        + ", but file ends after " + realLen + " bytes!");
+            }
+            return new String(bufArray, i + 22, commentLen);
+        }
+        return null;
+    }
+}
diff --git a/src/main/java/dev/zarr/zarrjava/utils/ZipUtils.java b/src/main/java/dev/zarr/zarrjava/utils/ZipUtils.java
