From fba232d6bf408a53d05fbaa805b0223798b2ba2b Mon Sep 17 00:00:00 2001 From: Antoine de Maleprade Date: Sun, 1 Oct 2023 15:28:09 +0200 Subject: [PATCH] feat(ZarrArray): add store get options passthrough (#146) --- src/core/index.ts | 59 +++++++++++++++++++------------------- src/creation.ts | 12 ++++---- src/hierarchy.ts | 22 +++++++------- src/storage/httpStore.ts | 2 +- src/storage/memoryStore.ts | 2 +- src/storage/types.ts | 2 +- 6 files changed, 50 insertions(+), 49 deletions(-) diff --git a/src/core/index.ts b/src/core/index.ts index 5aaafbd..a421d14 100644 --- a/src/core/index.ts +++ b/src/core/index.ts @@ -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 { concurrencyLimit?: number; progressCallback?: (progressUpdate: { progress: number; queueSize: number; }) => void; + storeOptions?: StoreGetOptions; } export interface SetOptions { @@ -34,20 +35,20 @@ export interface SetOptions { }) => void; } -export interface GetRawChunkOptions { - storeOptions: O; +export interface GetRawChunkOptions { + storeOptions: StoreGetOptions; } -export class ZarrArray { +export class ZarrArray { - public store: Store; + public store: Store; private compressor: Promise | null; - private _chunkStore: Store | null; + private _chunkStore: Store | null; /** * A `Store` providing the underlying storage for array chunks. */ - public get chunkStore(): Store { + public get chunkStore(): Store { if (this._chunkStore) { return this._chunkStore; } @@ -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(store: Store, path: null | string = null, readOnly = false, chunkStore: Store | 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); } @@ -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, path: null | string = null, metadata: ZarrArrayMetadata, readOnly = false, chunkStore: Store | null = null, cacheMetadata = true, cacheAttrs = true) { // N.B., expect at this point store is fully initialized with all // configuration metadata fully specified and normalized @@ -263,26 +264,26 @@ export class ZarrArray { } } - public get(selection?: undefined | Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], opts?: GetOptions): Promise | number>; - public get(selection?: ArraySelection, opts?: GetOptions): Promise | number>; - public get(selection: ArraySelection = null, opts: GetOptions = {}): Promise | number> { + public get(selection?: undefined | Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], opts?: GetOptions): Promise | number>; + public get(selection?: ArraySelection, opts?: GetOptions): Promise | number>; + public get(selection: ArraySelection = null, opts: GetOptions = {}): Promise | number> { return this.getBasicSelection(selection, false, opts); } - public getRaw(selection?: undefined | Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], opts?: GetOptions): Promise; - public getRaw(selection?: ArraySelection, opts?: GetOptions): Promise; - public getRaw(selection: ArraySelection = null, opts: GetOptions = {}): Promise { + public getRaw(selection?: undefined | Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], opts?: GetOptions): Promise; + public getRaw(selection?: ArraySelection, opts?: GetOptions): Promise; + public getRaw(selection: ArraySelection = null, opts: GetOptions = {}): Promise { return this.getBasicSelection(selection, true, opts); } // asRaw = false - public async getBasicSelection(selection: Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], asRaw?: false, opts?: GetOptions): Promise | number>; - public async getBasicSelection(selection: ArraySelection, asRaw?: false, opts?: GetOptions): Promise | number>; + public async getBasicSelection(selection: Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], asRaw?: false, opts?: GetOptions): Promise | number>; + public async getBasicSelection(selection: ArraySelection, asRaw?: false, opts?: GetOptions): Promise | number>; // asRaw = true - public async getBasicSelection(selection: Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], asRaw?: true, opts?: GetOptions): Promise; - public async getBasicSelection(selection: ArraySelection, asRaw?: true, opts?: GetOptions): Promise; + public async getBasicSelection(selection: Slice | ":" | "..." | null | (Slice | null | ":" | "...")[], asRaw?: true, opts?: GetOptions): Promise; + public async getBasicSelection(selection: ArraySelection, asRaw?: true, opts?: GetOptions): Promise; - public async getBasicSelection(selection: ArraySelection, asRaw = false, { concurrencyLimit = 10, progressCallback }: GetOptions = {}): Promise | RawArray | number> { + public async getBasicSelection(selection: ArraySelection, asRaw = false, { concurrencyLimit = 10, progressCallback, storeOptions }: GetOptions = {}): Promise | RawArray | number> { // Refresh metadata if (!this.cacheMetadata) { await this.reloadMetadata(); @@ -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 | RawArray> { + private getBasicSelectionND(selection: ArraySelection, asRaw: boolean, concurrencyLimit: number, progressCallback?: (progressUpdate: { progress: number; queueSize: number }) => void, storeOptions?: StoreGetOptions): Promise | 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 | RawArray> { + private async getSelection(indexer: BasicIndexer, asRaw: boolean, concurrencyLimit: number, progressCallback?: (progressUpdate: { progress: number; queueSize: number }) => void, storeOptions?: StoreGetOptions): Promise | 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. @@ -346,7 +347,7 @@ 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 }); })(); @@ -354,7 +355,7 @@ export class ZarrArray { } 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)); } } @@ -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(chunkCoords: number[], chunkSelection: DimensionSelection[], out: NestedArray | RawArray, outSelection: DimensionSelection[], dropAxes: null | number[]) { + private async chunkGetItem(chunkCoords: number[], chunkSelection: DimensionSelection[], out: NestedArray | 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) { @@ -431,7 +432,7 @@ export class ZarrArray { } } - public async getRawChunk(chunkCoords: number[], opts?: GetRawChunkOptions): Promise { + public async getRawChunk(chunkCoords: number[], opts?: GetRawChunkOptions): Promise { if (chunkCoords.length !== this.shape.length) { throw new Error(`Chunk coordinates ${chunkCoords.join(".")} do not correspond to shape ${this.shape}.`); } diff --git a/src/creation.ts b/src/creation.ts index 35660f4..0e08d4a 100644 --- a/src/creation.ts +++ b/src/creation.ts @@ -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; @@ -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 = " { - - 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); @@ -127,9 +127,9 @@ export async function array(data: Buffer | ArrayBuffer | NestedArray type OpenArrayOptions = Partial; export async function openArray( - { shape, mode = "a", chunks = true, dtype = " { +export class Group implements AsyncMutableMapping | ZarrArray> { /** * A `Store` providing the underlying storage for the group. */ - public store: Store; + public store: Store; /** * Storage path. @@ -52,11 +52,11 @@ export class Group implements AsyncMutableMapping { public attrs: Attributes; - private _chunkStore: Store | null; + private _chunkStore: Store | null; /** * A `Store` providing the underlying storage for array chunks. */ - public get chunkStore(): Store { + public get chunkStore(): Store { if (this._chunkStore) { return this._chunkStore; } @@ -67,12 +67,12 @@ export class Group implements AsyncMutableMapping { 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(store: Store, path: string | null = null, readOnly = false, chunkStore: Store | 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(store: Store, path: null | string) { path = normalizeStoragePath(path); const keyPrefix = pathToPrefix(path); try { @@ -86,7 +86,7 @@ export class Group implements AsyncMutableMapping { } } - private constructor(store: Store, path: string | null = null, metadata: ZarrGroupMetadata, readOnly = false, chunkStore: Store | null = null, cacheAttrs = true) { + private constructor(store: Store, path: string | null = null, metadata: ZarrGroupMetadata, readOnly = false, chunkStore: Store | null = null, cacheAttrs = true) { this.store = store; this._chunkStore = chunkStore; this.path = normalizeStoragePath(path); @@ -202,7 +202,7 @@ export class Group implements AsyncMutableMapping { } opts = this.getOptsForArrayCreation(name, opts); - let z: Promise; + let z: Promise>; if (data === undefined) { if (shape === undefined) { throw new ValueError("Shape must be set if no data is passed to CreateDataset"); @@ -242,7 +242,7 @@ export class Group implements AsyncMutableMapping { return await containsArray(this.store, path) || containsGroup(this.store, path); } - proxy(): AsyncMutableMappingProxy { + proxy(): AsyncMutableMappingProxy> { return createProxy(this); } } @@ -256,7 +256,7 @@ export class Group implements AsyncMutableMapping { * @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(store?: Store | string, path: string | null = null, chunkStore?: Store, overwrite = false, cacheAttrs = true) { store = normalizeStoreArgument(store); path = normalizeStoragePath(path); @@ -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(store?: Store | string, path: string | null = null, mode: PersistenceMode = "a", chunkStore?: Store, cacheAttrs = true) { store = normalizeStoreArgument(store); if (chunkStore !== undefined) { chunkStore = normalizeStoreArgument(store); diff --git a/src/storage/httpStore.ts b/src/storage/httpStore.ts index fa09a3a..629eb78 100644 --- a/src/storage/httpStore.ts +++ b/src/storage/httpStore.ts @@ -15,7 +15,7 @@ interface HTTPStoreOptions { supportedMethods?: HTTPMethod[]; } -export class HTTPStore implements AsyncStore { +export class HTTPStore implements AsyncStore { listDir?: undefined; rmDir?: undefined; getSize?: undefined; diff --git a/src/storage/memoryStore.ts b/src/storage/memoryStore.ts index ef36684..7fa98da 100644 --- a/src/storage/memoryStore.ts +++ b/src/storage/memoryStore.ts @@ -2,7 +2,7 @@ import { SyncStore, ValidStoreType } from "./types"; import { createProxy, MutableMappingProxy } from "../mutableMapping"; import { KeyError } from "../errors"; -export class MemoryStore implements SyncStore { +export class MemoryStore implements SyncStore { listDir?: undefined; rmDir?: undefined; getSize?: undefined; diff --git a/src/storage/types.ts b/src/storage/types.ts index 304b120..740c356 100644 --- a/src/storage/types.ts +++ b/src/storage/types.ts @@ -3,7 +3,7 @@ import { MutableMapping, AsyncMutableMapping } from "../mutableMapping"; export type ValidStoreType = Buffer | string | ArrayBuffer; -export type Store = SyncStore | AsyncStore; +export type Store = SyncStore | AsyncStore; /** * This module contains storage classes for use with Zarr arrays and groups.