Skip to content

List group members level by level instead of enumerating all keys #92

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
konstibob wants to merge 1 commit into
base: main
Choose a base branch
from
Open

List group members level by level instead of enumerating all keys #92

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 98 additions & 1 deletion src/main/java/dev/zarr/zarrjava/core/Group.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,16 +3,32 @@
import dev.zarr.zarrjava.ZarrException;
import dev.zarr.zarrjava.store.FilesystemStore;
import dev.zarr.zarrjava.store.StoreHandle;
import dev.zarr.zarrjava.utils.Utils;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.AbstractMap;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public abstract class Group extends AbstractNode {

/**
* Keys that hold metadata of the group itself and never point at a child node.
*/
private static final Set<String> METADATA_KEYS = Collections.unmodifiableSet(
new HashSet<>(Arrays.asList(ZARR_JSON, ZARRAY, ZATTRS, ZGROUP)));

protected Group(@Nonnull StoreHandle storeHandle) {
super(storeHandle);
}
Expand Down Expand Up @@ -70,13 +86,94 @@ public Node get(String key) throws ZarrException, IOException {
return get(new String[]{key});
}

public abstract Stream<Node> list();
/**
* Lists the immediate children (arrays and subgroups) of this group.
* <p>
* This costs a single listing request on the underlying store, plus one metadata read per
* child. Keys that do not hold a Zarr node are skipped.
*
* @return a stream of the direct children of this group
* @throws UnsupportedOperationException if the underlying store does not support listing
*/
public Stream<Node> members() {
return childKeys(new String[0]).parallelStream()
.map(this::openChild)
.filter(Objects::nonNull)
.collect(Collectors.toList())
.stream();
}

public Node[] membersAsArray() {
try (Stream<Node> nodeStream = members()) {
return nodeStream.toArray(Node[]::new);
}
}

/**
* Lists all descendants (arrays and groups) of this group, at any depth.
* <p>
* The group hierarchy is walked one level at a time, so only group keys are listed and chunk
* keys are never enumerated. Descending into an array is not necessary and does not happen.
*
* @return a stream of all descendants of this group, excluding the group itself
* @throws UnsupportedOperationException if the underlying store does not support listing
*/
public Stream<Node> list() {
return listDescendants(new String[0]);
}

public Node[] listAsArray() {
try (Stream<Node> nodeStream = list()) {
return nodeStream.toArray(Node[]::new);
}
}

private Stream<Node> listDescendants(String[] prefix) {
List<Map.Entry<String[], Node>> children = childKeys(prefix).parallelStream()
.map(key -> new AbstractMap.SimpleEntry<String[], Node>(key, openChild(key)))
.collect(Collectors.toList());

return children.stream().flatMap(child -> {
Node node = child.getValue();
if (node == null) {
// Not a node itself, but it may still contain nodes further down.
return listDescendants(child.getKey());
}
if (node instanceof Group) {
return Stream.concat(Stream.of(node), listDescendants(child.getKey()));
}
return Stream.of(node);
});
}

/**
* Lists the keys directly below {@code prefix} that may hold a child node, relative to this
* group.
*/
private List<String[]> childKeys(String[] prefix) {
try (Stream<String> children = storeHandle.resolve(prefix).listChildren()) {
return children
.filter(name -> !METADATA_KEYS.contains(name))
.map(name -> Utils.concatArrays(prefix, new String[]{name}))
.collect(Collectors.toList());
}
}

/**
* Opens the node at {@code key}, or returns null if there is no node there.
*/
@Nullable
private Node openChild(String[] key) {
try {
return get(key);
} catch (IOException e) {
throw new RuntimeException(
"Failed to read node metadata for key '" + String.join("/", key) + "': " + e.getMessage(), e);
} catch (ZarrException e) {
throw new RuntimeException(
"Failed to parse node metadata for key '" + String.join("/", key) + "': " + e.getMessage(), e);
}
}

public abstract GroupMetadata metadata();
}
19 changes: 19 additions & 0 deletions src/main/java/dev/zarr/zarrjava/store/FilesystemStore.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,16 @@ public boolean exists(String[] keys) {
return Files.isRegularFile(resolveKeys(keys));
}

/**
* Whether there is no file at the given keys. Reading such a key does not always fail with a
* {@link NoSuchFileException}: if one of the path components is a file rather than a
* directory, the filesystem reports "Not a directory" instead. Either way the key holds no
* data, so {@link #get} has to return null for it.
*/
private boolean isMissing(String[] keys) {
return !Files.isRegularFile(resolveKeys(keys));
}

@Nullable
@Override
public ByteBuffer get(String[] keys) {
Expand All @@ -52,6 +62,9 @@ public ByteBuffer get(String[] keys) {
} catch (NoSuchFileException e) {
return null;
} catch (IOException e) {
if (isMissing(keys)) {
return null;
}
throw StoreException.readFailed(this.toString(), keys, e);
}
}
Expand All @@ -75,6 +88,9 @@ public ByteBuffer get(String[] keys, long start) {
} catch (NoSuchFileException e) {
return null;
} catch (IOException e) {
if (isMissing(keys)) {
return null;
}
throw StoreException.readFailed(this.toString(), keys, e);
}
}
Expand All @@ -97,6 +113,9 @@ public ByteBuffer get(String[] keys, long start, long end) {
} catch (NoSuchFileException e) {
return null;
} catch (IOException e) {
if (isMissing(keys)) {
return null;
}
throw StoreException.readFailed(this.toString(), keys, e);
}
}
Expand Down
22 changes: 14 additions & 8 deletions src/main/java/dev/zarr/zarrjava/store/S3Store.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Stream;

import javax.annotation.Nonnull;
Expand All @@ -12,7 +14,6 @@
import software.amazon.awssdk.core.ResponseInputStream;
import software.amazon.awssdk.core.sync.RequestBody;
import software.amazon.awssdk.services.s3.S3Client;
import software.amazon.awssdk.services.s3.model.CommonPrefix;
import software.amazon.awssdk.services.s3.model.DeleteObjectRequest;
import software.amazon.awssdk.services.s3.model.GetObjectRequest;
import software.amazon.awssdk.services.s3.model.GetObjectResponse;
Expand Down Expand Up @@ -184,15 +185,20 @@ public Stream<String> listChildren(String[] keys) {
.delimiter("/")
.build();

ListObjectsV2Response res = s3client.listObjectsV2(req);

// Combine CommonPrefixes (folders) and Contents (files)
Stream<String> folders = res.commonPrefixes().stream().map(CommonPrefix::prefix);
final String finalFullPrefix = fullPrefix;
Stream<String> files = res.contents().stream().map(S3Object::key)
.filter(key -> !key.equals(finalFullPrefix));
// Combine CommonPrefixes (folders) and Contents (files) across all pages. A single
// listObjectsV2 call returns at most 1000 entries, which would silently truncate the
// children of a large group.
List<String> children = new ArrayList<>();
for (ListObjectsV2Response res : s3client.listObjectsV2Paginator(req)) {
res.commonPrefixes().forEach(commonPrefix -> children.add(commonPrefix.prefix()));
res.contents().stream()
.map(S3Object::key)
.filter(key -> !key.equals(finalFullPrefix))
.forEach(children::add);
}

return Stream.concat(folders, files)
return children.stream()
.map(k -> keyToRelativeArray(k, finalFullPrefix)[0]);
}

Expand Down
23 changes: 0 additions & 23 deletions src/main/java/dev/zarr/zarrjava/v2/Group.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,7 @@
import java.nio.file.NoSuchFileException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.Objects;
import java.util.function.Function;
import java.util.stream.Stream;

import static dev.zarr.zarrjava.v2.Node.makeObjectMapper;
import static dev.zarr.zarrjava.v2.Node.makeObjectWriter;
Expand Down Expand Up @@ -182,26 +179,6 @@ public Node get(String[] key) throws ZarrException, IOException {
}
}

@Override
public Stream<dev.zarr.zarrjava.core.Node> list() {
return storeHandle.list().map(key -> {
if (key.length <= 1) return null; // exclude root from list
String fileName = key[key.length - 1];
StoreHandle parent = storeHandle.resolve(Arrays.copyOf(key, key.length - 1));
try {
if (fileName.equals(ZARRAY)) {
return Array.open(parent);
}
if (fileName.equals(ZGROUP)) {
return (dev.zarr.zarrjava.core.Node) Group.open(parent);
}
} catch (Exception e) {
throw new RuntimeException(e);
}
return null;
}).filter(Objects::nonNull);
}

/**
* Creates a new subgroup with default metadata at the specified key.
*
Expand Down
21 changes: 0 additions & 21 deletions src/main/java/dev/zarr/zarrjava/v3/Group.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,7 @@
import java.nio.file.NoSuchFileException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.function.Function;
import java.util.stream.Stream;

import static dev.zarr.zarrjava.v3.Node.makeObjectMapper;
import static dev.zarr.zarrjava.v3.Node.makeObjectWriter;
Expand Down Expand Up @@ -192,25 +190,6 @@ public Node get(String[] key) throws ZarrException, IOException {
}
}

@Override
public Stream<dev.zarr.zarrjava.core.Node> list() {
Stream<String[]> metadataKeys = storeHandle.list()
.filter(key -> key[key.length - 1].equals(ZARR_JSON))
.filter(key -> key.length > 1); // exclude root from list
return metadataKeys.map(key -> {
try {
return get(Arrays.copyOf(key, key.length - 1));
} catch (IOException e) {
throw new RuntimeException(
"Failed to read node metadata for key '" + String.join("/", key) + "': " + e.getMessage(), e);
} catch (ZarrException e) {
throw new RuntimeException(
"Failed to parse node metadata for key '" + String.join("/", key) + "': " + e.getMessage(), e);
}
});
}


/**
* Creates a new subgroup with the provided metadata at the specified key.
*
Expand Down
Loading
Loading