Skip to content

Commit

Permalink
feat(ZarrArray): add store get options passthrough (#146)
Browse files Browse the repository at this point in the history
  • Loading branch information
@az0uz
az0uz committed Oct 1, 2023
1 parent 15e3a3f commit fba232d
Show file tree
Hide file tree
Showing 6 changed files with 50 additions and 49 deletions.
59 changes: 30 additions & 29 deletions src/core/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,13 @@ import { getCodec } from "../compression/registry";
import type { Codec } from 'numcodecs';
import PQueue from 'p-queue';

export interface GetOptions {
export interface GetOptions<StoreGetOptions> {
concurrencyLimit?: number;
progressCallback?: (progressUpdate: {
progress: number;
queueSize: number;
}) => void;
storeOptions?: StoreGetOptions;
}

export interface SetOptions {
Expand All @@ -34,20 +35,20 @@ export interface SetOptions {
}) => void;
}

export interface GetRawChunkOptions<O> {
storeOptions: O;
export interface GetRawChunkOptions<StoreGetOptions> {
storeOptions: StoreGetOptions;
}

export class ZarrArray {
export class ZarrArray<StoreGetOptions = any> {

public store: Store;
public store: Store<StoreGetOptions>;
private compressor: Promise<Codec> | null;

private _chunkStore: Store | null;
private _chunkStore: Store<StoreGetOptions> | null;
/**
* A `Store` providing the underlying storage for array chunks.
*/
public get chunkStore(): Store {
public get chunkStore(): Store<StoreGetOptions> {
if (this._chunkStore) {
return this._chunkStore;
}
Expand Down Expand Up @@ -193,7 +194,7 @@ export class ZarrArray {
* @param cacheAttrs If true (default), user attributes will be cached for attribute read operations.
* If false, user attributes are reloaded from the store prior to all attribute read operations.
*/
public static async create(store: Store, path: null | string = null, readOnly = false, chunkStore: Store | null = null, cacheMetadata = true, cacheAttrs = true) {
public static async create<StoreGetOptions>(store: Store<StoreGetOptions>, path: null | string = null, readOnly = false, chunkStore: Store<StoreGetOptions> | null = null, cacheMetadata = true, cacheAttrs = true) {
const metadata = await this.loadMetadataForConstructor(store, path);
return new ZarrArray(store, path, metadata as ZarrArrayMetadata, readOnly, chunkStore, cacheMetadata, cacheAttrs);
}
Expand Down Expand Up @@ -224,7 +225,7 @@ export class ZarrArray {
* @param cacheAttrs If true (default), user attributes will be cached for attribute read operations.
* If false, user attributes are reloaded from the store prior to all attribute read operations.
*/
private constructor(store: Store, path: null | string = null, metadata: ZarrArrayMetadata, readOnly = false, chunkStore: Store | null = null, cacheMetadata = true, cacheAttrs = true) {
private constructor(store: Store<StoreGetOptions>, path: null | string = null, metadata: ZarrArrayMetadata, readOnly = false, chunkStore: Store<StoreGetOptions> | null = null, cacheMetadata = true, cacheAttrs = true) {
// N.B., expect at this point store is fully initialized with all
// configuration metadata fully specified and normalized

Expand Down Expand Up @@ -263,26 +264,26 @@ export class ZarrArray {
}
}

public get(selection?: undefined | Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], opts?: GetOptions): Promise<NestedArray<TypedArray> | number>;
public get(selection?: ArraySelection, opts?: GetOptions): Promise<NestedArray<TypedArray> | number>;
public get(selection: ArraySelection = null, opts: GetOptions = {}): Promise<NestedArray<TypedArray> | number> {
public get(selection?: undefined | Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], opts?: GetOptions<StoreGetOptions>): Promise<NestedArray<TypedArray> | number>;
public get(selection?: ArraySelection, opts?: GetOptions<StoreGetOptions>): Promise<NestedArray<TypedArray> | number>;
public get(selection: ArraySelection = null, opts: GetOptions<StoreGetOptions> = {}): Promise<NestedArray<TypedArray> | number> {
return this.getBasicSelection(selection, false, opts);
}

public getRaw(selection?: undefined | Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], opts?: GetOptions): Promise<RawArray | number>;
public getRaw(selection?: ArraySelection, opts?: GetOptions): Promise<RawArray | number>;
public getRaw(selection: ArraySelection = null, opts: GetOptions = {}): Promise<RawArray | number> {
public getRaw(selection?: undefined | Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], opts?: GetOptions<StoreGetOptions>): Promise<RawArray | number>;
public getRaw(selection?: ArraySelection, opts?: GetOptions<StoreGetOptions>): Promise<RawArray | number>;
public getRaw(selection: ArraySelection = null, opts: GetOptions<StoreGetOptions> = {}): Promise<RawArray | number> {
return this.getBasicSelection(selection, true, opts);
}

// asRaw = false
public async getBasicSelection(selection: Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], asRaw?: false, opts?: GetOptions): Promise<NestedArray<TypedArray> | number>;
public async getBasicSelection(selection: ArraySelection, asRaw?: false, opts?: GetOptions): Promise<NestedArray<TypedArray> | number>;
public async getBasicSelection(selection: Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], asRaw?: false, opts?: GetOptions<StoreGetOptions>): Promise<NestedArray<TypedArray> | number>;
public async getBasicSelection(selection: ArraySelection, asRaw?: false, opts?: GetOptions<StoreGetOptions>): Promise<NestedArray<TypedArray> | number>;
// asRaw = true
public async getBasicSelection(selection: Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], asRaw?: true, opts?: GetOptions): Promise<RawArray | number>;
public async getBasicSelection(selection: ArraySelection, asRaw?: true, opts?: GetOptions): Promise<RawArray | number>;
public async getBasicSelection(selection: Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], asRaw?: true, opts?: GetOptions<StoreGetOptions>): Promise<RawArray | number>;
public async getBasicSelection(selection: ArraySelection, asRaw?: true, opts?: GetOptions<StoreGetOptions>): Promise<RawArray | number>;

public async getBasicSelection(selection: ArraySelection, asRaw = false, { concurrencyLimit = 10, progressCallback }: GetOptions = {}): Promise<NestedArray<TypedArray> | RawArray | number> {
public async getBasicSelection(selection: ArraySelection, asRaw = false, { concurrencyLimit = 10, progressCallback, storeOptions }: GetOptions<StoreGetOptions> = {}): Promise<NestedArray<TypedArray> | RawArray | number> {
// Refresh metadata
if (!this.cacheMetadata) {
await this.reloadMetadata();
Expand All @@ -292,16 +293,16 @@ export class ZarrArray {
if (this.shape.length === 0) {
throw new Error("Shape [] indexing is not supported yet");
} else {
return this.getBasicSelectionND(selection, asRaw, concurrencyLimit, progressCallback);
return this.getBasicSelectionND(selection, asRaw, concurrencyLimit, progressCallback, storeOptions);
}
}

private getBasicSelectionND(selection: ArraySelection, asRaw: boolean, concurrencyLimit: number, progressCallback?: (progressUpdate: { progress: number; queueSize: number }) => void): Promise<number | NestedArray<TypedArray> | RawArray> {
private getBasicSelectionND(selection: ArraySelection, asRaw: boolean, concurrencyLimit: number, progressCallback?: (progressUpdate: { progress: number; queueSize: number }) => void, storeOptions?: StoreGetOptions): Promise<number | NestedArray<TypedArray> | RawArray> {
const indexer = new BasicIndexer(selection, this);
return this.getSelection(indexer, asRaw, concurrencyLimit, progressCallback);
return this.getSelection(indexer, asRaw, concurrencyLimit, progressCallback, storeOptions);
}

private async getSelection(indexer: BasicIndexer, asRaw: boolean, concurrencyLimit: number, progressCallback?: (progressUpdate: { progress: number; queueSize: number }) => void): Promise<number | NestedArray<TypedArray> | RawArray> {
private async getSelection(indexer: BasicIndexer, asRaw: boolean, concurrencyLimit: number, progressCallback?: (progressUpdate: { progress: number; queueSize: number }) => void, storeOptions?: StoreGetOptions): Promise<number | NestedArray<TypedArray> | RawArray> {
// We iterate over all chunks which overlap the selection and thus contain data
// that needs to be extracted. Each chunk is processed in turn, extracting the
// necessary data and storing into the correct location in the output array.
Expand Down Expand Up @@ -346,15 +347,15 @@ export class ZarrArray {
progressCallback({ progress: 0, queueSize: queueSize });
for (const proj of indexer.iter()) {
(async () => {
await queue.add(() => this.chunkGetItem(proj.chunkCoords, proj.chunkSelection, out, proj.outSelection, indexer.dropAxes));
await queue.add(() => this.chunkGetItem(proj.chunkCoords, proj.chunkSelection, out, proj.outSelection, indexer.dropAxes, storeOptions));
progress += 1;
progressCallback({ progress: progress, queueSize: queueSize });
})();
}

} else {
for (const proj of indexer.iter()) {
queue.add(() => this.chunkGetItem(proj.chunkCoords, proj.chunkSelection, out, proj.outSelection, indexer.dropAxes));
queue.add(() => this.chunkGetItem(proj.chunkCoords, proj.chunkSelection, out, proj.outSelection, indexer.dropAxes, storeOptions));
}
}

Expand All @@ -377,14 +378,14 @@ export class ZarrArray {
* @param outSelection Location of region within output array to store results in.
* @param dropAxes Axes to squeeze out of the chunk.
*/
private async chunkGetItem<T extends TypedArray>(chunkCoords: number[], chunkSelection: DimensionSelection[], out: NestedArray<T> | RawArray, outSelection: DimensionSelection[], dropAxes: null | number[]) {
private async chunkGetItem<T extends TypedArray>(chunkCoords: number[], chunkSelection: DimensionSelection[], out: NestedArray<T> | RawArray, outSelection: DimensionSelection[], dropAxes: null | number[], storeOptions?: StoreGetOptions) {
if (chunkCoords.length !== this._chunkDataShape.length) {
throw new ValueError(`Inconsistent shapes: chunkCoordsLength: ${chunkCoords.length}, cDataShapeLength: ${this.chunkDataShape.length}`);
}

const cKey = this.chunkKey(chunkCoords);
try {
const cdata = await this.chunkStore.getItem(cKey);
const cdata = await this.chunkStore.getItem(cKey, storeOptions);
const decodedChunk = await this.decodeChunk(cdata);

if (out instanceof NestedArray) {
Expand Down Expand Up @@ -431,7 +432,7 @@ export class ZarrArray {
}
}

public async getRawChunk<O>(chunkCoords: number[], opts?: GetRawChunkOptions<O>): Promise<RawArray> {
public async getRawChunk(chunkCoords: number[], opts?: GetRawChunkOptions<StoreGetOptions>): Promise<RawArray> {
if (chunkCoords.length !== this.shape.length) {
throw new Error(`Chunk coordinates ${chunkCoords.join(".")} do not correspond to shape ${this.shape}.`);
}
Expand Down
12 changes: 6 additions & 6 deletions src/creation.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ export type CreateArrayOptions = {
compressor?: CompressorConfig | null;
fillValue?: FillType;
order?: Order;
store?: Store;
store?: Store | string;
overwrite?: boolean;
path?: string | null;
chunkStore?: Store;
Expand Down Expand Up @@ -54,10 +54,10 @@ export type CreateArrayOptions = {
* @param dimensionSeparator if specified, defines an alternate string separator placed between the dimension chunks.
*/
export async function create(
{ shape, chunks = true, dtype = "<i4", compressor = null, fillValue = null, order = "C", store, overwrite = false, path, chunkStore, filters, cacheMetadata = true, cacheAttrs = true, readOnly = false, dimensionSeparator }: CreateArrayOptions,
{ shape, chunks = true, dtype = "<i4", compressor = null, fillValue = null, order = "C", store: storeArgument, overwrite = false, path, chunkStore, filters, cacheMetadata = true, cacheAttrs = true, readOnly = false, dimensionSeparator }: CreateArrayOptions,
): Promise<ZarrArray> {

store = normalizeStoreArgument(store);
const store = normalizeStoreArgument(storeArgument);

await initArray(store, shape, chunks, dtype, path, compressor, fillValue, order, overwrite, chunkStore, filters, dimensionSeparator);
const z = await ZarrArray.create(store, path, readOnly, chunkStore, cacheMetadata, cacheAttrs);
Expand Down Expand Up @@ -127,9 +127,9 @@ export async function array(data: Buffer | ArrayBuffer | NestedArray<TypedArray>
type OpenArrayOptions = Partial<CreateArrayOptions & { mode: PersistenceMode }>;

export async function openArray(
{ shape, mode = "a", chunks = true, dtype = "<i4", compressor = null, fillValue = null, order = "C", store, overwrite = false, path = null, chunkStore, filters, cacheMetadata = true, cacheAttrs = true, dimensionSeparator }: OpenArrayOptions = {},
{ shape, mode = "a", chunks = true, dtype = "<i4", compressor = null, fillValue = null, order = "C", store: storeArgument, overwrite = false, path = null, chunkStore, filters, cacheMetadata = true, cacheAttrs = true, dimensionSeparator }: OpenArrayOptions = {},
) {
store = normalizeStoreArgument(store);
const store = normalizeStoreArgument(storeArgument);
if (chunkStore === undefined) {
chunkStore = normalizeStoreArgument(store);
}
Expand Down
22 changes: 11 additions & 11 deletions src/hierarchy.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,11 @@ import { TypedArray } from './nestedArray/types';
import { ZarrArray } from './core';


export class Group implements AsyncMutableMapping<Group | ZarrArray> {
export class Group<StoreGetOptions = any> implements AsyncMutableMapping<Group<StoreGetOptions> | ZarrArray<StoreGetOptions>> {
/**
* A `Store` providing the underlying storage for the group.
*/
public store: Store;
public store: Store<StoreGetOptions>;

/**
* Storage path.
Expand Down Expand Up @@ -52,11 +52,11 @@ export class Group implements AsyncMutableMapping<Group | ZarrArray> {
public attrs: Attributes<UserAttributes>;


private _chunkStore: Store | null;
private _chunkStore: Store<StoreGetOptions> | null;
/**
* A `Store` providing the underlying storage for array chunks.
*/
public get chunkStore(): Store {
public get chunkStore(): Store<StoreGetOptions> {
if (this._chunkStore) {
return this._chunkStore;
}
Expand All @@ -67,12 +67,12 @@ export class Group implements AsyncMutableMapping<Group | ZarrArray> {
public readOnly: boolean;
private meta: ZarrGroupMetadata;

public static async create(store: Store, path: string | null = null, readOnly = false, chunkStore: Store | null = null, cacheAttrs = true) {
public static async create<StoreGetOptions>(store: Store<StoreGetOptions>, path: string | null = null, readOnly = false, chunkStore: Store<StoreGetOptions> | null = null, cacheAttrs = true) {
const metadata = await this.loadMetadataForConstructor(store, path);
return new Group(store, path, metadata as ZarrGroupMetadata, readOnly, chunkStore, cacheAttrs);
}

private static async loadMetadataForConstructor(store: Store, path: null | string) {
private static async loadMetadataForConstructor<StoreGetOptions>(store: Store<StoreGetOptions>, path: null | string) {
path = normalizeStoragePath(path);
const keyPrefix = pathToPrefix(path);
try {
Expand All @@ -86,7 +86,7 @@ export class Group implements AsyncMutableMapping<Group | ZarrArray> {
}
}

private constructor(store: Store, path: string | null = null, metadata: ZarrGroupMetadata, readOnly = false, chunkStore: Store | null = null, cacheAttrs = true) {
private constructor(store: Store<StoreGetOptions>, path: string | null = null, metadata: ZarrGroupMetadata, readOnly = false, chunkStore: Store<StoreGetOptions> | null = null, cacheAttrs = true) {
this.store = store;
this._chunkStore = chunkStore;
this.path = normalizeStoragePath(path);
Expand Down Expand Up @@ -202,7 +202,7 @@ export class Group implements AsyncMutableMapping<Group | ZarrArray> {
}
opts = this.getOptsForArrayCreation(name, opts);

let z: Promise<ZarrArray>;
let z: Promise<ZarrArray<StoreGetOptions>>;
if (data === undefined) {
if (shape === undefined) {
throw new ValueError("Shape must be set if no data is passed to CreateDataset");
Expand Down Expand Up @@ -242,7 +242,7 @@ export class Group implements AsyncMutableMapping<Group | ZarrArray> {
return await containsArray(this.store, path) || containsGroup(this.store, path);
}

proxy(): AsyncMutableMappingProxy<Group> {
proxy(): AsyncMutableMappingProxy<Group<StoreGetOptions>> {
return createProxy(this);
}
}
Expand All @@ -256,7 +256,7 @@ export class Group implements AsyncMutableMapping<Group | ZarrArray> {
* @param cacheAttrs If `true` (default), user attributes will be cached for attribute read operations.
* If `false`, user attributes are reloaded from the store prior to all attribute read operations.
*/
export async function group(store?: Store | string, path: string | null = null, chunkStore?: Store, overwrite = false, cacheAttrs = true) {
export async function group<StoreGetOptions>(store?: Store<StoreGetOptions> | string, path: string | null = null, chunkStore?: Store<StoreGetOptions>, overwrite = false, cacheAttrs = true) {
store = normalizeStoreArgument(store);
path = normalizeStoragePath(path);

Expand All @@ -277,7 +277,7 @@ export async function group(store?: Store | string, path: string | null = null,
* If False, user attributes are reloaded from the store prior to all attribute read operations.
*
*/
export async function openGroup(store?: Store | string, path: string | null = null, mode: PersistenceMode = "a", chunkStore?: Store, cacheAttrs = true) {
export async function openGroup<StoreGetOptions>(store?: Store<StoreGetOptions> | string, path: string | null = null, mode: PersistenceMode = "a", chunkStore?: Store<StoreGetOptions>, cacheAttrs = true) {
store = normalizeStoreArgument(store);
if (chunkStore !== undefined) {
chunkStore = normalizeStoreArgument(store);
Expand Down
2 changes: 1 addition & 1 deletion src/storage/httpStore.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ interface HTTPStoreOptions {
supportedMethods?: HTTPMethod[];
}

export class HTTPStore<UrlRoot extends string | URL=string> implements AsyncStore<ArrayBuffer> {
export class HTTPStore<UrlRoot extends string | URL=string> implements AsyncStore<ArrayBuffer, RequestInit> {
listDir?: undefined;
rmDir?: undefined;
getSize?: undefined;
Expand Down
2 changes: 1 addition & 1 deletion src/storage/memoryStore.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ import { SyncStore, ValidStoreType } from "./types";
import { createProxy, MutableMappingProxy } from "../mutableMapping";
import { KeyError } from "../errors";

export class MemoryStore<T extends ValidStoreType> implements SyncStore<T> {
export class MemoryStore<T extends ValidStoreType> implements SyncStore<T, undefined> {
listDir?: undefined;
rmDir?: undefined;
getSize?: undefined;
Expand Down
2 changes: 1 addition & 1 deletion src/storage/types.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ import { MutableMapping, AsyncMutableMapping } from "../mutableMapping";
export type ValidStoreType = Buffer | string | ArrayBuffer;


export type Store = SyncStore<ValidStoreType> | AsyncStore<ValidStoreType>;
export type Store<StoreGetOptions = any> = SyncStore<ValidStoreType, StoreGetOptions> | AsyncStore<ValidStoreType, StoreGetOptions>;

/**
* This module contains storage classes for use with Zarr arrays and groups.
Expand Down

0 comments on commit fba232d

Please sign in to comment.