From 7b98a743ad4b8b2b08f47169a8361b953533c17b Mon Sep 17 00:00:00 2001 From: Konstantin Date: Wed, 26 Aug 2026 18:39:39 +0200 Subject: [PATCH] consolidated metadata --- USERGUIDE.md | 32 ++ .../zarrjava/v3/ConsolidatedMetadata.java | 127 +++++ src/main/java/dev/zarr/zarrjava/v3/Group.java | 223 +++++++- .../dev/zarr/zarrjava/v3/GroupMetadata.java | 37 +- .../zarrjava/ConsolidatedMetadataTest.java | 498 ++++++++++++++++++ .../dev/zarr/zarrjava/ZarrPythonTests.java | 35 ++ .../python-scripts/zarr_python_consolidate.py | 34 ++ 7 files changed, 979 insertions(+), 7 deletions(-) create mode 100644 src/main/java/dev/zarr/zarrjava/v3/ConsolidatedMetadata.java create mode 100644 src/test/java/dev/zarr/zarrjava/ConsolidatedMetadataTest.java create mode 100644 src/test/python-scripts/zarr_python_consolidate.py diff --git a/USERGUIDE.md b/USERGUIDE.md index 65bf9944..0bc97da2 100644 --- a/USERGUIDE.md +++ b/USERGUIDE.md @@ -288,6 +288,38 @@ Array array = group.createArray( .build() ); ``` +### Consolidated Metadata (v3) +A group can keep a copy of the metadata of all of its descendants inside its own `zarr.json`, so that +the whole hierarchy can be opened with a single read instead of one read per node. This matters most +over HTTP and S3, where every node otherwise costs a request. + +```java +// Write the cache. This walks the hierarchy once and stores the metadata of every +// descendant in the metadata of this group. +Group root = Group.open(storeHandle).consolidateMetadata(); + +// Later reads are answered from the cache, without touching the store. +Group sub = (Group) root.get("sub"); +Array array = (Array) sub.get("nested"); + +// Remove the cache again +root.dropConsolidatedMetadata(); + +// Ignore a cache that is present, for example when the hierarchy may have changed +Group fresh = Group.open(storeHandle, false); +``` + +The cache is written in the same format as `zarr.consolidate_metadata()` in zarr-python, so both +libraries can read each other's output. + +**The cache is a snapshot.** Nothing invalidates it when a node is added, removed or changed +afterwards, so `consolidateMetadata()` has to be called again after modifying the hierarchy. Reading a +node that is missing from the cache logs a warning and falls back to reading the node itself, but a +node that was *modified* after consolidating is served from the cache and cannot be detected. Open the +group with `Group.open(storeHandle, false)` if in doubt. + +Consolidated metadata is a Zarr v3 feature here; the v2 `.zmetadata` file is not supported. + ### Hierarchical Example ```java Group root = Group.create( diff --git a/src/main/java/dev/zarr/zarrjava/v3/ConsolidatedMetadata.java b/src/main/java/dev/zarr/zarrjava/v3/ConsolidatedMetadata.java new file mode 100644 index 00000000..6de4a41f --- /dev/null +++ b/src/main/java/dev/zarr/zarrjava/v3/ConsolidatedMetadata.java @@ -0,0 +1,127 @@ +package dev.zarr.zarrjava.v3; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.JsonNode; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; + +/** + * An optional cache of the metadata of all descendants of a group, stored inside that group's own + * {@code zarr.json} under the {@code consolidated_metadata} key. It allows a reader to open a whole + * hierarchy with a single request instead of one request per node. + *

+ * The keys of {@link #metadata} are flat, {@code "/"}-joined paths relative to the group holding the + * cache, for example {@code "ocean"} and {@code "ocean/salinity"}. + *

+ * The cached node metadata is deliberately kept as raw {@link JsonNode} rather than as parsed + * {@link ArrayMetadata} / {@link GroupMetadata}. The cache is declared with + * {@code must_understand: false}, so a reader that cannot interpret an entry has to ignore it rather + * than fail. Parsing entries eagerly would mean that a single node written by another implementation + * with a field this library does not model would make the whole group unopenable. Keeping the raw + * JSON also lets {@link Group#consolidateMetadata()} copy each node's metadata verbatim, so the cache + * never silently loses information that is present in the node's own {@code zarr.json}. + *

+ * The cache is a snapshot taken at the time of consolidation. Nothing invalidates it when a + * descendant changes, so {@link Group#consolidateMetadata()} has to be re-run after modifying the + * hierarchy. + */ +@JsonIgnoreProperties(ignoreUnknown = true) +public final class ConsolidatedMetadata { + + /** + * The only cache kind defined so far: the metadata is stored inline in the group's metadata + * document. A cache of any other kind is ignored by this library. + */ + public static final String KIND_INLINE = "inline"; + + @Nonnull + @JsonProperty("kind") + public final String kind; + + @JsonProperty("must_understand") + public final boolean mustUnderstand; + + /** + * The cached metadata documents, keyed by their {@code "/"}-joined path relative to the group + * holding this cache. + */ + @Nonnull + @JsonProperty("metadata") + public final Map metadata; + + @JsonCreator(mode = JsonCreator.Mode.PROPERTIES) + public ConsolidatedMetadata( + @Nullable @JsonProperty("kind") String kind, + @Nullable @JsonProperty("must_understand") Boolean mustUnderstand, + @Nullable @JsonProperty("metadata") Map metadata + ) { + this.kind = kind == null ? KIND_INLINE : kind; + this.mustUnderstand = mustUnderstand != null && mustUnderstand; + this.metadata = metadata == null + ? Collections.emptyMap() + : Collections.unmodifiableMap(new LinkedHashMap<>(metadata)); + } + + public ConsolidatedMetadata(@Nonnull Map metadata) { + this(KIND_INLINE, false, metadata); + } + + /** + * An empty inline cache, used for a consolidated subgroup whose own entries have been hoisted + * into the cache of an ancestor. + */ + public static ConsolidatedMetadata empty() { + return new ConsolidatedMetadata(Collections.emptyMap()); + } + + /** + * Whether this cache is stored inline and can therefore be used by this library. + */ + @JsonIgnore + public boolean isInline() { + return KIND_INLINE.equals(kind); + } + + @JsonIgnore + public boolean isEmpty() { + return metadata.isEmpty(); + } + + /** + * Returns the cached metadata document for a node, or null if this cache does not hold it. + * + * @param key the path of the node relative to the group holding this cache + */ + @Nullable + public JsonNode get(String[] key) { + if (!isInline()) { + return null; + } + return metadata.get(String.join("/", key)); + } + + /** + * Returns the entries below {@code prefix} with the prefix stripped from their keys, so that the + * result can serve as the cache of the subgroup at {@code prefix}. + */ + public ConsolidatedMetadata sub(String[] prefix) { + if (!isInline()) { + return empty(); + } + String keyPrefix = String.join("/", prefix) + "/"; + Map sub = new LinkedHashMap<>(); + for (Map.Entry entry : metadata.entrySet()) { + if (entry.getKey().startsWith(keyPrefix)) { + sub.put(entry.getKey().substring(keyPrefix.length()), entry.getValue()); + } + } + return new ConsolidatedMetadata(sub); + } +} diff --git a/src/main/java/dev/zarr/zarrjava/v3/Group.java b/src/main/java/dev/zarr/zarrjava/v3/Group.java index 8b1a81bb..536ed90f 100644 --- a/src/main/java/dev/zarr/zarrjava/v3/Group.java +++ b/src/main/java/dev/zarr/zarrjava/v3/Group.java @@ -1,6 +1,9 @@ package dev.zarr.zarrjava.v3; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.ObjectWriter; +import com.fasterxml.jackson.databind.node.ObjectNode; import dev.zarr.zarrjava.ZarrException; import dev.zarr.zarrjava.core.Attributes; import dev.zarr.zarrjava.store.FilesystemStore; @@ -15,8 +18,17 @@ import java.nio.file.NoSuchFileException; import java.nio.file.Path; import java.nio.file.Paths; +import java.text.Normalizer; +import java.util.ArrayList; import java.util.Arrays; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; import java.util.function.Function; +import java.util.logging.Logger; +import java.util.stream.Collectors; import java.util.stream.Stream; import static dev.zarr.zarrjava.v3.Node.makeObjectMapper; @@ -25,11 +37,34 @@ public class Group extends dev.zarr.zarrjava.core.Group implements Node { + private static final Logger LOGGER = Logger.getLogger(Group.class.getName()); + + /** + * The order in which entries are written into the consolidated metadata: shallow paths first, + * then case-insensitively by name. This only affects the byte layout of the written metadata + * document, but it makes consolidating the same hierarchy twice produce an identical file. + */ + private static final Comparator CONSOLIDATED_KEY_ORDER = Comparator + .comparingInt((String key) -> (int) key.chars().filter(c -> c == '/').count()) + .thenComparing(key -> Normalizer.normalize(key, Normalizer.Form.NFKC).toLowerCase(Locale.ROOT)) + .thenComparing(Comparator.naturalOrder()); + public GroupMetadata metadata; + /** + * Whether {@link #get} may be answered from the consolidated metadata of this group. + */ + private final boolean useConsolidated; + protected Group(@Nonnull StoreHandle storeHandle, @Nonnull GroupMetadata groupMetadata) throws IOException { + this(storeHandle, groupMetadata, true); + } + + protected Group(@Nonnull StoreHandle storeHandle, @Nonnull GroupMetadata groupMetadata, + boolean useConsolidated) throws IOException { super(storeHandle); this.metadata = groupMetadata; + this.useConsolidated = useConsolidated; } /** @@ -39,9 +74,25 @@ protected Group(@Nonnull StoreHandle storeHandle, @Nonnull GroupMetadata groupMe * @throws IOException if the metadata cannot be read */ public static Group open(@Nonnull StoreHandle storeHandle) throws IOException { + return open(storeHandle, true); + } + + /** + * Opens an existing Zarr group at a specified storage location. + * + * @param storeHandle the storage location of the Zarr group + * @param useConsolidated whether the consolidated metadata of the group, if it has any, may be + * used to look up its descendants. Pass false to always read every node + * from the store, for example when the hierarchy may have been modified + * since it was consolidated. + * @throws IOException if the metadata cannot be read + */ + public static Group open(@Nonnull StoreHandle storeHandle, boolean useConsolidated) throws IOException { StoreHandle metadataHandle = storeHandle.resolve(ZARR_JSON); ByteBuffer metadataBytes = metadataHandle.readNonNull(); - return new Group(storeHandle, makeObjectMapper().readValue(Utils.toArray(metadataBytes), GroupMetadata.class)); + GroupMetadata groupMetadata = + makeObjectMapper().readValue(Utils.toArray(metadataBytes), GroupMetadata.class); + return new Group(storeHandle, groupMetadata, useConsolidated); } @@ -184,9 +235,43 @@ public static Group create(String path) throws IOException, ZarrException { */ @Nullable public Node get(String[] key) throws ZarrException, IOException { + ConsolidatedMetadata consolidated = useConsolidated ? metadata.consolidatedMetadata : null; + if (consolidated == null || !consolidated.isInline()) { + return openFromStore(key); + } + JsonNode cached = consolidated.get(key); + if (cached != null) { + Node node = nodeFromConsolidatedMetadata(key, cached, consolidated); + if (node != null) { + return node; + } + // The cached document could not be interpreted, fall back to the node itself. + return openFromStore(key); + } + Node node = openFromStore(key); + if (node != null) { + LOGGER.warning("The node '" + String.join("/", key) + "' below " + storeHandle + + " is missing from the consolidated metadata of the group. The consolidated" + + " metadata is a snapshot and does not track later changes to the hierarchy;" + + " call consolidateMetadata() again to refresh it."); + } + return node; + } + + /** + * Opens the node at {@code key} by reading its metadata from the store, ignoring any consolidated + * metadata. + */ + @Nullable + private Node openFromStore(String[] key) throws ZarrException, IOException { StoreHandle keyHandle = storeHandle.resolve(key); try { - return Node.open(keyHandle); + Node node = Node.open(keyHandle); + if (!useConsolidated && node instanceof Group) { + Group group = (Group) node; + return new Group(group.storeHandle, group.metadata, false); + } + return node; } catch (NoSuchFileException e) { return null; } @@ -211,6 +296,135 @@ public Stream list() { } + /** + * Builds a node from a cached metadata document, or returns null if the document cannot be + * interpreted. The consolidated metadata is declared with {@code must_understand: false}, so an + * entry this library does not understand is skipped in favour of reading the node itself rather + * than failing. + */ + @Nullable + private Node nodeFromConsolidatedMetadata(String[] key, JsonNode cached, + ConsolidatedMetadata consolidated) { + StoreHandle keyHandle = storeHandle.resolve(key); + JsonNode nodeTypeNode = cached.get("node_type"); + String nodeType = nodeTypeNode == null ? null : nodeTypeNode.asText(); + try { + ObjectMapper objectMapper = makeObjectMapper(); + if (ArrayMetadata.NODE_TYPE.equals(nodeType)) { + return new Array(keyHandle, objectMapper.treeToValue(cached, ArrayMetadata.class)); + } + if (GroupMetadata.NODE_TYPE.equals(nodeType)) { + GroupMetadata groupMetadata = objectMapper.treeToValue(cached, GroupMetadata.class); + // The entries of a consolidated subgroup are hoisted into the cache of this group, so + // hand the subgroup its own slice of them instead of the emptied cache it carries. + return new Group(keyHandle, + groupMetadata.withConsolidatedMetadata(consolidated.sub(key)), true); + } + LOGGER.warning("Ignoring the consolidated metadata of '" + String.join("/", key) + + "' below " + storeHandle + ", it has an unsupported node type '" + nodeType + "'."); + return null; + } catch (Exception e) { + LOGGER.warning("Ignoring the consolidated metadata of '" + String.join("/", key) + + "' below " + storeHandle + ", it could not be parsed: " + e.getMessage()); + return null; + } + } + + /** + * Writes the metadata of all descendants of this group into the metadata of this group, so that + * the whole hierarchy can afterwards be opened with a single read. + *

+ * The metadata of each descendant is copied verbatim, keyed by its {@code "/"}-joined path + * relative to this group. The copy of a subgroup is given an empty cache of its own, marking it as + * covered by the cache written here; a subgroup that was consolidated itself therefore does not + * have its entries stored twice. + *

+ * The result is a snapshot. Nothing invalidates it when a node is added, removed or modified + * afterwards, so this has to be called again after changing the hierarchy. Reading a node that is + * missing from the cache logs a warning and falls back to reading the node itself, but a node that + * was modified after consolidating is served from the cache and cannot be detected. + * + * @return this group, with the consolidated metadata written + * @throws IOException if the metadata cannot be read or written + * @throws UnsupportedOperationException if the underlying store does not support listing + */ + public Group consolidateMetadata() throws IOException { + Map entries = new LinkedHashMap<>(); + collectDescendantMetadata(new String[0], entries); + + List keys = new ArrayList<>(entries.keySet()); + keys.sort(CONSOLIDATED_KEY_ORDER); + Map sorted = new LinkedHashMap<>(); + for (String key : keys) { + sorted.put(key, entries.get(key)); + } + return writeMetadata( + metadata.withConsolidatedMetadata(new ConsolidatedMetadata(sorted))); + } + + /** + * Removes the consolidated metadata of this group, so that its descendants are read from the store + * again. + * + * @return this group, with the consolidated metadata removed + * @throws IOException if the metadata cannot be written + */ + public Group dropConsolidatedMetadata() throws IOException { + if (metadata.consolidatedMetadata == null) { + return this; + } + return writeMetadata(metadata.withConsolidatedMetadata(null)); + } + + /** + * Collects the metadata documents of all nodes below {@code prefix} into {@code out}, keyed by + * their path relative to this group. + */ + private void collectDescendantMetadata(String[] prefix, Map out) + throws IOException { + List children; + try (Stream stream = storeHandle.resolve(prefix).listChildren()) { + children = stream.filter(name -> !ZARR_JSON.equals(name)).collect(Collectors.toList()); + } + for (String child : children) { + String[] key = Utils.concatArrays(prefix, new String[]{child}); + ByteBuffer metadataBytes = storeHandle.resolve(key).resolve(ZARR_JSON).read(); + if (metadataBytes == null) { + // Not a node itself, but it may still contain nodes further down. + collectDescendantMetadata(key, out); + continue; + } + JsonNode nodeMetadata = makeObjectMapper().readTree(Utils.toArray(metadataBytes)); + JsonNode nodeTypeNode = nodeMetadata.get("node_type"); + boolean isGroup = nodeTypeNode != null && GroupMetadata.NODE_TYPE.equals(nodeTypeNode.asText()); + if (isGroup) { + markSubgroupAsConsolidated(nodeMetadata); + } + out.put(String.join("/", key), nodeMetadata); + if (isGroup) { + collectDescendantMetadata(key, out); + } + } + } + + /** + * Gives the cached metadata of a subgroup an empty consolidated metadata cache of its own. The + * empty cache marks the subgroup as covered by the cache being written here, which is where its + * entries live. A subgroup that carried a cache of its own loses it in this copy, so that the same + * entries are not held twice and cannot drift apart. This mirrors what zarr-python writes. + */ + private static void markSubgroupAsConsolidated(JsonNode nodeMetadata) { + if (!(nodeMetadata instanceof ObjectNode)) { + return; + } + ObjectNode metadataObject = (ObjectNode) nodeMetadata; + ObjectNode nested = metadataObject.objectNode(); + nested.put("kind", ConsolidatedMetadata.KIND_INLINE); + nested.put("must_understand", false); + nested.set("metadata", metadataObject.objectNode()); + metadataObject.set("consolidated_metadata", nested); + } + /** * Creates a new subgroup with the provided metadata at the specified key. * @@ -302,7 +516,10 @@ public Group updateAttributes(Function attributeMapper) * @throws IOException if the metadata cannot be serialized */ public Group setAttributes(Attributes newAttributes) throws ZarrException, IOException { - GroupMetadata newGroupMetadata = new GroupMetadata(newAttributes); + // The consolidated metadata describes the descendants of this group, which are unaffected by + // a change to the attributes of the group itself. + GroupMetadata newGroupMetadata = + new GroupMetadata(newAttributes, metadata.consolidatedMetadata); return writeMetadata(newGroupMetadata); } diff --git a/src/main/java/dev/zarr/zarrjava/v3/GroupMetadata.java b/src/main/java/dev/zarr/zarrjava/v3/GroupMetadata.java index 56414a2e..3ec5f151 100644 --- a/src/main/java/dev/zarr/zarrjava/v3/GroupMetadata.java +++ b/src/main/java/dev/zarr/zarrjava/v3/GroupMetadata.java @@ -16,21 +16,35 @@ public final class GroupMetadata extends dev.zarr.zarrjava.core.GroupMetadata { public final int zarrFormat = ZARR_FORMAT; @JsonProperty("node_type") public final String nodeType = "group"; + + /** + * An optional cache of the metadata of all descendants of this group, or null if this group has + * not been consolidated. See {@link ConsolidatedMetadata} and {@link Group#consolidateMetadata()}. + */ + @Nullable @JsonProperty("consolidated_metadata") - public final Object consolidatedMetadata = null; + public final ConsolidatedMetadata consolidatedMetadata; @Nullable public final Attributes attributes; public GroupMetadata(@Nullable Attributes attributes) throws ZarrException { - this(ZARR_FORMAT, NODE_TYPE, attributes); + this(ZARR_FORMAT, NODE_TYPE, attributes, null); + } + + public GroupMetadata( + @Nullable Attributes attributes, + @Nullable ConsolidatedMetadata consolidatedMetadata + ) throws ZarrException { + this(ZARR_FORMAT, NODE_TYPE, attributes, consolidatedMetadata); } @JsonCreator(mode = JsonCreator.Mode.PROPERTIES) public GroupMetadata( @JsonProperty(value = "zarr_format", required = true) int zarrFormat, @JsonProperty(value = "node_type", required = true) String nodeType, - @Nullable @JsonProperty(value = "attributes") Attributes attributes + @Nullable @JsonProperty(value = "attributes") Attributes attributes, + @Nullable @JsonProperty(value = "consolidated_metadata") ConsolidatedMetadata consolidatedMetadata ) throws ZarrException { if (zarrFormat != this.zarrFormat) { throw new ZarrException( @@ -41,11 +55,12 @@ public GroupMetadata( "Expected node type '" + this.nodeType + "', got '" + nodeType + "'."); } this.attributes = attributes; + this.consolidatedMetadata = consolidatedMetadata; } public static GroupMetadata defaultValue() { try { - return new GroupMetadata(ZARR_FORMAT, NODE_TYPE, new Attributes()); + return new GroupMetadata(ZARR_FORMAT, NODE_TYPE, new Attributes(), null); } catch (ZarrException e) { // This should never happen with default values throw new IllegalStateException( @@ -53,6 +68,20 @@ public static GroupMetadata defaultValue() { } } + /** + * Returns a copy of this metadata with a different consolidated metadata cache, or without one if + * {@code newConsolidatedMetadata} is null. + */ + public GroupMetadata withConsolidatedMetadata(@Nullable ConsolidatedMetadata newConsolidatedMetadata) { + try { + return new GroupMetadata(zarrFormat, nodeType, attributes, newConsolidatedMetadata); + } catch (ZarrException e) { + // This should never happen, the format and node type are copied from a valid instance + throw new IllegalStateException( + "Failed to copy GroupMetadata - this indicates a programming error", e); + } + } + @Override public @Nonnull Attributes attributes() throws ZarrException { if (attributes == null) { diff --git a/src/test/java/dev/zarr/zarrjava/ConsolidatedMetadataTest.java b/src/test/java/dev/zarr/zarrjava/ConsolidatedMetadataTest.java new file mode 100644 index 00000000..9094bda0 --- /dev/null +++ b/src/test/java/dev/zarr/zarrjava/ConsolidatedMetadataTest.java @@ -0,0 +1,498 @@ +package dev.zarr.zarrjava; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import dev.zarr.zarrjava.core.Attributes; +import dev.zarr.zarrjava.store.MemoryStore; +import dev.zarr.zarrjava.store.Store; +import dev.zarr.zarrjava.store.StoreHandle; +import dev.zarr.zarrjava.utils.Utils; +import dev.zarr.zarrjava.v3.ConsolidatedMetadata; +import dev.zarr.zarrjava.v3.DataType; +import dev.zarr.zarrjava.v3.Group; +import dev.zarr.zarrjava.v3.Node; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.io.IOException; +import java.io.InputStream; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Stream; + +import static dev.zarr.zarrjava.core.Node.ZARR_JSON; + +/** + * Tests for the {@code consolidated_metadata} cache of a v3 group: writing it with + * {@link Group#consolidateMetadata()}, answering {@link Group#get} from it, and tolerating caches + * this library cannot fully interpret. + */ +public class ConsolidatedMetadataTest { + + private static final Set EXPECTED_ENTRIES = new HashSet<>(Arrays.asList( + "arr", + "sub", + "sub/nested", + "sub/deep", + "sub/deep/deepArray" + )); + + /** + * A {@link MemoryStore} that counts how often it is asked to list or read, so that tests can + * assert on the number of store operations a group traversal costs. + */ + static final class CountingStore implements Store, Store.ListableStore { + + private final MemoryStore delegate = new MemoryStore(); + final AtomicInteger listCalls = new AtomicInteger(); + final AtomicInteger listChildrenCalls = new AtomicInteger(); + final AtomicInteger readCalls = new AtomicInteger(); + + void resetCounters() { + listCalls.set(0); + listChildrenCalls.set(0); + readCalls.set(0); + } + + @Override + public Stream list(String[] prefix) { + listCalls.incrementAndGet(); + return delegate.list(prefix); + } + + @Override + public Stream listChildren(String[] prefix) { + listChildrenCalls.incrementAndGet(); + return delegate.listChildren(prefix); + } + + @Override + public boolean exists(String[] keys) { + readCalls.incrementAndGet(); + return delegate.exists(keys); + } + + @Nullable + @Override + public ByteBuffer get(String[] keys) { + readCalls.incrementAndGet(); + return delegate.get(keys); + } + + @Nullable + @Override + public ByteBuffer get(String[] keys, long start) { + readCalls.incrementAndGet(); + return delegate.get(keys, start); + } + + @Nullable + @Override + public ByteBuffer get(String[] keys, long start, long end) { + readCalls.incrementAndGet(); + return delegate.get(keys, start, end); + } + + @Override + public void set(String[] keys, ByteBuffer bytes) { + delegate.set(keys, bytes); + } + + @Override + public void delete(String[] keys) { + delegate.delete(keys); + } + + @Nonnull + @Override + public StoreHandle resolve(String... keys) { + return new StoreHandle(this, keys); + } + + @Override + public InputStream getInputStream(String[] keys, long start, long end) { + readCalls.incrementAndGet(); + return delegate.getInputStream(keys, start, end); + } + + @Override + public long getSize(String[] keys) { + return delegate.getSize(keys); + } + + @Override + public String toString() { + return ""; + } + } + + /** + * Writes a v3 hierarchy: + * 

+     * /            group
+     * /arr         array (chunked)
+     * /sub         group
+     * /sub/nested  array
+     * /sub/deep    group
+     * /sub/deep/deepArray array
+     *
+ */ + static Group writeTreeV3(StoreHandle storeHandle) throws IOException, ZarrException { + Group root = Group.create(storeHandle); + byte[] data = new byte[64 * 64]; + for (int i = 0; i < data.length; i++) { + data[i] = (byte) i; + } + dev.zarr.zarrjava.v3.Array array = root.createArray("arr", b -> b + .withShape(64, 64) + .withDataType(DataType.UINT8) + .withChunkShape(8, 8)); + array.write(ucar.ma2.Array.factory(ucar.ma2.DataType.BYTE, new int[]{64, 64}, data)); + + Group sub = root.createGroup("sub"); + sub.createArray("nested", b -> b + .withShape(8, 8) + .withDataType(DataType.UINT8) + .withChunkShape(8, 8)); + Group deep = sub.createGroup("deep"); + deep.createArray("deepArray", b -> b + .withShape(8, 8) + .withDataType(DataType.UINT8) + .withChunkShape(8, 8)); + return root; + } + + private static ObjectNode readJson(StoreHandle handle) throws IOException { + ByteBuffer bytes = handle.resolve(ZARR_JSON).readNonNull(); + return (ObjectNode) new ObjectMapper().readTree(Utils.toArray(bytes)); + } + + private static void writeJson(StoreHandle handle, JsonNode json) throws IOException { + handle.resolve(ZARR_JSON).set(ByteBuffer.wrap(new ObjectMapper().writeValueAsBytes(json))); + } + + @Test + public void testConsolidateWritesAllDescendants() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + writeTreeV3(store.resolve()).consolidateMetadata(); + + ObjectNode written = readJson(store.resolve()); + JsonNode consolidated = written.get("consolidated_metadata"); + Assertions.assertEquals("inline", consolidated.get("kind").asText()); + Assertions.assertFalse(consolidated.get("must_understand").asBoolean()); + + Set keys = new HashSet<>(); + consolidated.get("metadata").fieldNames().forEachRemaining(keys::add); + Assertions.assertEquals(EXPECTED_ENTRIES, keys); + + Assertions.assertEquals("array", + consolidated.get("metadata").get("sub/deep/deepArray").get("node_type").asText()); + } + + @Test + public void testConsolidatedEntriesMatchTheNodesThemselves() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + writeTreeV3(store.resolve()).consolidateMetadata(); + + JsonNode entries = readJson(store.resolve()).get("consolidated_metadata").get("metadata"); + for (String key : EXPECTED_ENTRIES) { + ObjectNode fromNode = readJson(store.resolve(key.split("/"))); + fromNode.remove("consolidated_metadata"); + ObjectNode fromCache = (ObjectNode) entries.get(key).deepCopy(); + fromCache.remove("consolidated_metadata"); + Assertions.assertEquals(fromNode, fromCache, "the cached metadata of '" + key + + "' must be a verbatim copy of the metadata of the node"); + } + } + + @Test + public void testGetIsAnsweredWithoutReadingTheStore() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + writeTreeV3(store.resolve()).consolidateMetadata(); + + Group root = Group.open(store.resolve()); + store.resetCounters(); + + Assertions.assertNotNull(root.get("arr")); + Assertions.assertNotNull(root.get(new String[]{"sub", "deep", "deepArray"})); + Assertions.assertEquals(0, store.readCalls.get(), + "a node held by the consolidated metadata must not be read from the store"); + } + + @Test + public void testSubgroupsAreAlsoConsolidated() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + writeTreeV3(store.resolve()).consolidateMetadata(); + + Group root = Group.open(store.resolve()); + store.resetCounters(); + + Group sub = (Group) root.get("sub"); + Assertions.assertNotNull(sub); + Group deep = (Group) sub.get("deep"); + Assertions.assertNotNull(deep); + Assertions.assertNotNull(deep.get("deepArray")); + Assertions.assertEquals(0, store.readCalls.get(), + "walking into a subgroup must keep using the consolidated metadata of the root"); + } + + @Test + public void testListUsesTheConsolidatedMetadata() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + writeTreeV3(store.resolve()).consolidateMetadata(); + + Group root = Group.open(store.resolve()); + store.resetCounters(); + + Assertions.assertEquals(EXPECTED_ENTRIES.size(), root.listAsArray().length); + // Listing still has to discover the keys, but none of the metadata is read again. + Assertions.assertEquals(0, store.readCalls.get()); + } + + @Test + public void testNodeAddedAfterConsolidatingIsStillFound() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + Group root = writeTreeV3(store.resolve()).consolidateMetadata(); + + root.createArray("late", b -> b + .withShape(4, 4) + .withDataType(DataType.UINT8) + .withChunkShape(4, 4)); + + Group reopened = Group.open(store.resolve()); + Assertions.assertNotNull(reopened.get("late"), + "a node missing from the stale cache must be read from the store instead"); + Assertions.assertNull(reopened.get("doesNotExist")); + } + + @Test + public void testSubgroupEntriesAreMarkedAsConsolidated() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + writeTreeV3(store.resolve()).consolidateMetadata(); + + JsonNode entries = readJson(store.resolve()).get("consolidated_metadata").get("metadata"); + for (String key : Arrays.asList("sub", "sub/deep")) { + JsonNode nested = entries.get(key).get("consolidated_metadata"); + Assertions.assertNotNull(nested, "the cached metadata of the subgroup '" + key + + "' must carry an empty cache, marking it as covered by the cache above it"); + Assertions.assertEquals("inline", nested.get("kind").asText()); + Assertions.assertFalse(nested.get("must_understand").asBoolean()); + Assertions.assertEquals(0, nested.get("metadata").size()); + } + // The subgroups themselves are untouched by consolidating the group above them. + Assertions.assertFalse(readJson(store.resolve("sub")).has("consolidated_metadata")); + } + + @Test + public void testNestedConsolidatedMetadataIsEmptied() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + Group root = writeTreeV3(store.resolve()); + ((Group) root.get("sub")).consolidateMetadata(); + root.consolidateMetadata(); + + JsonNode entries = readJson(store.resolve()).get("consolidated_metadata").get("metadata"); + JsonNode nested = entries.get("sub").get("consolidated_metadata"); + Assertions.assertNotNull(nested, "the key must be kept, so that it stays visible that the" + + " subgroup is consolidated"); + Assertions.assertEquals(0, nested.get("metadata").size(), + "the entries of a consolidated subgroup must not be duplicated inside the cache of" + + " the group above it"); + + // The subgroup keeps its own cache in its own metadata document. + Assertions.assertEquals(3, + readJson(store.resolve("sub")).get("consolidated_metadata").get("metadata").size()); + } + + @Test + public void testUnknownKindIsIgnored() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + writeTreeV3(store.resolve()).consolidateMetadata(); + + ObjectNode written = readJson(store.resolve()); + ((ObjectNode) written.get("consolidated_metadata")).put("kind", "something_else"); + writeJson(store.resolve(), written); + + Group root = Group.open(store.resolve()); + Assertions.assertNotNull(root.metadata.consolidatedMetadata); + Assertions.assertFalse(root.metadata.consolidatedMetadata.isInline()); + + store.resetCounters(); + Assertions.assertNotNull(root.get("arr")); + Assertions.assertTrue(store.readCalls.get() > 0, + "a cache of an unknown kind must be ignored, not used"); + } + + @Test + public void testUnknownFieldInACachedEntryDoesNotBreakTheGroup() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + writeTreeV3(store.resolve()).consolidateMetadata(); + + ObjectNode written = readJson(store.resolve()); + ObjectNode entry = (ObjectNode) written.get("consolidated_metadata").get("metadata").get("arr"); + entry.putArray("some_future_field").add("value"); + writeJson(store.resolve(), written); + + // Opening the group must not fail because of a cache entry it cannot interpret. + Group root = Group.open(store.resolve()); + store.resetCounters(); + Assertions.assertNotNull(root.get("arr"), + "an entry that cannot be parsed must fall back to reading the node itself"); + Assertions.assertTrue(store.readCalls.get() > 0); + + // The other entries are unaffected. + store.resetCounters(); + Assertions.assertNotNull(root.get(new String[]{"sub", "nested"})); + Assertions.assertEquals(0, store.readCalls.get()); + } + + @Test + public void testUnknownFieldSurvivesConsolidation() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + Group root = writeTreeV3(store.resolve()); + + ObjectNode arrayMetadata = readJson(store.resolve("arr")); + arrayMetadata.putArray("some_future_field").add("value"); + writeJson(store.resolve("arr"), arrayMetadata); + + root.consolidateMetadata(); + + JsonNode cached = readJson(store.resolve()) + .get("consolidated_metadata").get("metadata").get("arr"); + Assertions.assertEquals(arrayMetadata, cached, + "consolidating must copy the metadata of a node verbatim, including fields this" + + " library does not model"); + } + + @Test + public void testUseConsolidatedFalseIgnoresTheCache() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + writeTreeV3(store.resolve()).consolidateMetadata(); + + Group root = Group.open(store.resolve(), false); + store.resetCounters(); + + Assertions.assertNotNull(root.get("arr")); + Assertions.assertTrue(store.readCalls.get() > 0); + + // The opt-out is inherited by subgroups. + store.resetCounters(); + Group sub = (Group) root.get("sub"); + Assertions.assertNotNull(sub.get("nested")); + Assertions.assertTrue(store.readCalls.get() > 0); + } + + @Test + public void testExplicitNullIsParsedAndNotWrittenBack() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + Group root = writeTreeV3(store.resolve()); + + ObjectNode written = readJson(store.resolve()); + written.putNull("consolidated_metadata"); + writeJson(store.resolve(), written); + + Group reopened = Group.open(store.resolve()); + Assertions.assertNull(reopened.metadata.consolidatedMetadata); + + reopened.setAttributes(new Attributes().set("a", 1)); + Assertions.assertFalse(readJson(store.resolve()).has("consolidated_metadata"), + "an absent cache must be omitted, never written as null"); + } + + @Test + public void testDropConsolidatedMetadata() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + Group root = writeTreeV3(store.resolve()).consolidateMetadata(); + + root.dropConsolidatedMetadata(); + Assertions.assertNull(root.metadata.consolidatedMetadata); + Assertions.assertFalse(readJson(store.resolve()).has("consolidated_metadata")); + + store.resetCounters(); + Assertions.assertNotNull(root.get("arr")); + Assertions.assertTrue(store.readCalls.get() > 0); + } + + @Test + public void testAttributesUpdateKeepsTheCache() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + Group root = writeTreeV3(store.resolve()).consolidateMetadata(); + + root.setAttributes(new Attributes().set("answer", 42)); + + ObjectNode written = readJson(store.resolve()); + Assertions.assertEquals(42, written.get("attributes").get("answer").asInt()); + Assertions.assertEquals(EXPECTED_ENTRIES.size(), + written.get("consolidated_metadata").get("metadata").size(), + "changing the attributes of the group does not change its descendants"); + } + + @Test + public void testConsolidatingIsReproducible() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + Group root = writeTreeV3(store.resolve()); + + root.consolidateMetadata(); + byte[] first = Utils.toArray(store.resolve().resolve(ZARR_JSON).readNonNull()); + Group.open(store.resolve()).consolidateMetadata(); + byte[] second = Utils.toArray(store.resolve().resolve(ZARR_JSON).readNonNull()); + + Assertions.assertArrayEquals(first, second); + } + + @Test + public void testEntriesAreOrderedByDepthThenName() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + writeTreeV3(store.resolve()).consolidateMetadata(); + + List keys = new ArrayList<>(); + readJson(store.resolve()).get("consolidated_metadata").get("metadata") + .fieldNames().forEachRemaining(keys::add); + Assertions.assertEquals( + Arrays.asList("arr", "sub", "sub/deep", "sub/nested", "sub/deep/deepArray"), keys); + } + + @Test + public void testConsolidatedMetadataOfAnEmptyGroup() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + Group root = Group.create(store.resolve()).consolidateMetadata(); + + Assertions.assertNotNull(root.metadata.consolidatedMetadata); + Assertions.assertTrue(root.metadata.consolidatedMetadata.isEmpty()); + Assertions.assertEquals(0, + readJson(store.resolve()).get("consolidated_metadata").get("metadata").size(), + "a consolidated group without descendants keeps the key with an empty cache, so that" + + " it stays distinguishable from a group that was never consolidated"); + } + + @Test + public void testConsolidatedMetadataIsExposedOnTheMetadata() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + Group root = writeTreeV3(store.resolve()).consolidateMetadata(); + + ConsolidatedMetadata consolidated = root.metadata.consolidatedMetadata; + Assertions.assertNotNull(consolidated); + Assertions.assertEquals(EXPECTED_ENTRIES, consolidated.metadata.keySet()); + Assertions.assertNotNull(consolidated.get(new String[]{"sub", "deep"})); + Assertions.assertEquals(new HashSet<>(Arrays.asList("nested", "deep", "deep/deepArray")), + consolidated.sub(new String[]{"sub"}).metadata.keySet()); + } + + @Test + public void testNodeOpenIgnoresTheCacheOfTheGroupItself() throws IOException, ZarrException { + CountingStore store = new CountingStore(); + writeTreeV3(store.resolve()).consolidateMetadata(); + + // Opening the group through the generic entry point must give the same, usable group. + Group root = (Group) Node.open(store.resolve()); + Assertions.assertNotNull(root.metadata.consolidatedMetadata); + Assertions.assertNotNull(root.get(new String[]{"sub", "nested"})); + } +} diff --git a/src/test/java/dev/zarr/zarrjava/ZarrPythonTests.java b/src/test/java/dev/zarr/zarrjava/ZarrPythonTests.java index 5c22b5fc..b0111b8e 100644 --- a/src/test/java/dev/zarr/zarrjava/ZarrPythonTests.java +++ b/src/test/java/dev/zarr/zarrjava/ZarrPythonTests.java @@ -8,6 +8,7 @@ import dev.zarr.zarrjava.v2.Group; import dev.zarr.zarrjava.v3.Array; import dev.zarr.zarrjava.v3.ArrayMetadataBuilder; +import dev.zarr.zarrjava.v3.ConsolidatedMetadata; import dev.zarr.zarrjava.v3.DataType; import dev.zarr.zarrjava.v3.codec.CodecBuilder; import org.junit.jupiter.api.Assertions; @@ -24,6 +25,7 @@ import java.nio.ByteBuffer; import java.nio.file.Path; import java.nio.file.Paths; +import java.util.ArrayList; import java.util.Arrays; import java.util.stream.Stream; @@ -355,4 +357,37 @@ public void testGroupReadWriteV3() throws Exception { Assertions.assertArrayEquals(new int[]{16, 16, 16}, result.getShape()); assertIsTestdata(result, dataType); } + + /** + * Checks that the consolidated metadata written by zarr-java is understood by zarr-python and the + * other way round. + */ + @Test + public void testConsolidatedMetadataReadWriteV3() throws Exception { + StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT).resolve("testConsolidatedMetadataV3", "write"); + StoreHandle storeHandle2 = new FilesystemStore(TESTOUTPUT).resolve("testConsolidatedMetadataV3", "read"); + + ConsolidatedMetadataTest.writeTreeV3(storeHandle).consolidateMetadata(); + + run_python_script("zarr_python_consolidate.py", storeHandle.toPath().toString(), + storeHandle2.toPath().toString()); + + dev.zarr.zarrjava.v3.Group group = dev.zarr.zarrjava.v3.Group.open(storeHandle2); + ConsolidatedMetadata consolidated = group.metadata.consolidatedMetadata; + Assertions.assertNotNull(consolidated, "zarr-java did not pick up the consolidated metadata" + + " written by zarr-python"); + Assertions.assertEquals( + Arrays.asList("arr", "sub", "sub/deep", "sub/nested", "sub/deep/deepArray"), + new ArrayList<>(consolidated.metadata.keySet())); + + dev.zarr.zarrjava.v3.Array array = + (dev.zarr.zarrjava.v3.Array) group.get(new String[]{"sub", "deep", "deepArray"}); + Assertions.assertNotNull(array); + Assertions.assertArrayEquals(new long[]{8, 8}, array.metadata().shape); + + dev.zarr.zarrjava.v3.Array topArray = (dev.zarr.zarrjava.v3.Array) group.get("arr"); + Assertions.assertNotNull(topArray); + Assertions.assertArrayEquals(new int[]{64, 64}, topArray.read().getShape()); + Assertions.assertEquals(5, group.listAsArray().length); + } } diff --git a/src/test/python-scripts/zarr_python_consolidate.py b/src/test/python-scripts/zarr_python_consolidate.py new file mode 100644 index 00000000..28009f21 --- /dev/null +++ b/src/test/python-scripts/zarr_python_consolidate.py @@ -0,0 +1,34 @@ +import sys +from pathlib import Path + +import numpy as np +import zarr +from zarr.storage import LocalStore + +store_path_read = Path(sys.argv[1]) +store_path_write = Path(sys.argv[2]) + +expected_members = ["arr", "sub", "sub/deep", "sub/deep/deepArray", "sub/nested"] + +# Read a hierarchy that zarr-java consolidated. +g = zarr.open_group(store=LocalStore(store_path_read), zarr_format=3, use_consolidated=True) +consolidated = g.metadata.consolidated_metadata +assert consolidated is not None, "zarr-python did not pick up the consolidated metadata" +# zarr-python re-nests the flat keys when it reads them, so the top level only holds the direct +# children and the rest lives in the caches it builds for the subgroups. +assert list(consolidated.metadata.keys()) == ["arr", "sub"], list(consolidated.metadata.keys()) +members = sorted(k for k, _ in g.members(max_depth=None)) +assert members == expected_members, f"got {members}, expected {expected_members}" +assert g["arr"].shape == (64, 64), g["arr"].shape +assert g["sub"]["deep"]["deepArray"].shape == (8, 8) + +# Write a hierarchy of the same shape and consolidate it, for zarr-java to read. +g2 = zarr.create_group(store=LocalStore(store_path_write), zarr_format=3) +g2.attrs["attr"] = "value" +arr = g2.create_array(name="arr", shape=(64, 64), chunks=(8, 8), dtype="uint8", fill_value=0) +arr[:] = np.arange(64 * 64, dtype="uint8").reshape(64, 64) +sub = g2.create_group("sub") +sub.create_array(name="nested", shape=(8, 8), chunks=(8, 8), dtype="uint8", fill_value=0) +deep = sub.create_group("deep") +deep.create_array(name="deepArray", shape=(8, 8), chunks=(8, 8), dtype="uint8", fill_value=0) +zarr.consolidate_metadata(g2.store)