Skip to content

Commit 417c046

Browse files
lazarusAlazarusA
authored
icechunk fixes (#634)
1 parent 3d6831d commit 417c046

4 files changed

Lines changed: 667 additions & 478 deletions

File tree

‎src/components/zarr/ZarrLoaderLRU.ts‎

Lines changed: 119 additions & 107 deletions
Original file line numberDiff line numberDiff line change
@@ -1,140 +1,152 @@
11
import * as zarr from "zarrita";
22
import { IcechunkStore } from "icechunk-js";
3-
import { ZarrMetadata, ZarrTitleDescription, FetchStoreOptions } from "./Interfaces";
3+
import {
4+
ZarrMetadata,
5+
ZarrTitleDescription,
6+
FetchStoreOptions,
7+
} from "./Interfaces";
48
import { IcechunkStoreOptions } from "./Interfaces";
59
import { useGlobalStore } from "@/GlobalStates/GlobalStore";
6-
import { useCacheStore } from '@/GlobalStates/CacheStore';
10+
import { useCacheStore } from "@/GlobalStates/CacheStore";
711
import { useZarrStore } from "@/GlobalStates/ZarrStore";
8-
import {
9-
getFetchStore,
10-
getFetchStoreMetadata,
11-
getFetchStoreTitleDescription,
12-
getFetchStoreAttributes,
13-
getFetchStoreDims} from "./zarrita-store";
14-
import {
15-
getIcechunkStore,
16-
getIcechunkMetadata,
17-
getIcechunkTitleDescription,
18-
getIcechunkAttributes,
19-
getIcechunkDims} from "./icechunk-store";
12+
import {
13+
getFetchStore,
14+
getFetchStoreMetadata,
15+
getFetchStoreTitleDescription,
16+
getFetchStoreAttributes,
17+
getFetchStoreDims,
18+
} from "./zarrita-store";
19+
import {
20+
getIcechunkStore,
21+
getIcechunkMetadata,
22+
getIcechunkTitleDescription,
23+
getIcechunkAttributes,
24+
getIcechunkDims,
25+
} from "./icechunk-store";
2026
import { GetNCMetadata } from "./NCGetters";
2127

22-
type GroupType = zarr.Group<zarr.FetchStore | zarr.Listable<zarr.FetchStore> | IcechunkStore>;
28+
type GroupType = zarr.Group<
29+
zarr.FetchStore | zarr.Listable<zarr.FetchStore> | IcechunkStore
30+
>;
2331

2432
export async function GetStore(
25-
storePath: string,
26-
fetchOptions?: FetchStoreOptions,
27-
icechunkOptions?: IcechunkStoreOptions
33+
storePath: string,
34+
fetchOptions?: FetchStoreOptions,
35+
icechunkOptions?: IcechunkStoreOptions,
2836
): Promise<GroupType | undefined> {
29-
if (icechunkOptions) {
30-
return getIcechunkStore(storePath, icechunkOptions);
31-
}
32-
return getFetchStore(storePath, fetchOptions);
37+
if (icechunkOptions) {
38+
return getIcechunkStore(storePath, icechunkOptions);
39+
}
40+
return getFetchStore(storePath, fetchOptions);
3341
}
3442

3543
export async function GetZarrMetadata(
36-
groupStore: Promise<GroupType>
44+
groupStore: Promise<GroupType>,
3745
): Promise<ZarrMetadata[]> {
38-
const group = await groupStore;
39-
if (group.store instanceof IcechunkStore) {
40-
console.log('Fetching metadata from icechunk store');
41-
return getIcechunkMetadata(group as zarr.Group<IcechunkStore>);
42-
}
43-
console.log('Fetching metadata from fetch store');
44-
return getFetchStoreMetadata(
45-
group as zarr.Group<zarr.FetchStore | zarr.Listable<zarr.FetchStore>>
46-
);
46+
const group = await groupStore;
47+
if (group.store instanceof IcechunkStore) {
48+
return getIcechunkMetadata(group as zarr.Group<IcechunkStore>);
49+
}
50+
return getFetchStoreMetadata(
51+
group as zarr.Group<zarr.FetchStore | zarr.Listable<zarr.FetchStore>>,
52+
);
4753
}
4854

4955
export async function GetTitleDescription(
50-
groupStore: Promise<GroupType>
56+
groupStore: Promise<GroupType>,
5157
): Promise<ZarrTitleDescription> {
52-
const group = await groupStore;
53-
if (group.store instanceof IcechunkStore) {
54-
console.log('Fetching title and description from icechunk store');
55-
return getIcechunkTitleDescription(group as zarr.Group<IcechunkStore>);
56-
}
57-
console.log('Fetching title and description from fetch store');
58-
return getFetchStoreTitleDescription(
59-
group as zarr.Group<zarr.FetchStore | zarr.Listable<zarr.FetchStore>>
60-
);
58+
const group = await groupStore;
59+
if (group.store instanceof IcechunkStore) {
60+
return getIcechunkTitleDescription(group as zarr.Group<IcechunkStore>);
61+
}
62+
return getFetchStoreTitleDescription(
63+
group as zarr.Group<zarr.FetchStore | zarr.Listable<zarr.FetchStore>>,
64+
);
6165
}
6266

6367
export async function GetZarrAttributes(thisVariable?: string) {
64-
const { initStore, variable } = useGlobalStore.getState();
65-
const { cache } = useCacheStore.getState();
66-
const { currentStore } = useZarrStore.getState();
67-
const resolvedVariable = thisVariable ?? variable;
68-
const cacheName = `${initStore}_${resolvedVariable}_meta`;
68+
const { initStore, variable } = useGlobalStore.getState();
69+
const { cache } = useCacheStore.getState();
70+
const { currentStore } = useZarrStore.getState();
71+
const resolvedVariable = thisVariable ?? variable;
72+
const cacheName = `${initStore}_${resolvedVariable}_meta`;
6973

70-
if (cache.has(cacheName)) return cache.get(cacheName);
74+
if (cache.has(cacheName)) return cache.get(cacheName);
7175

72-
const group = await currentStore;
73-
if (!group) throw new Error(`Failed to open store: ${initStore}`);
76+
const group = await currentStore;
77+
if (!group) throw new Error(`Failed to open store: ${initStore}`);
7478

75-
if (group.store instanceof IcechunkStore) {
76-
return getIcechunkAttributes(
77-
group as zarr.Group<IcechunkStore>,
78-
resolvedVariable,
79-
cacheName
80-
);
81-
}
82-
return getFetchStoreAttributes(group, resolvedVariable, cacheName);
79+
if (group.store instanceof IcechunkStore) {
80+
return getIcechunkAttributes(
81+
group as zarr.Group<IcechunkStore>,
82+
resolvedVariable,
83+
cacheName,
84+
);
85+
}
86+
return getFetchStoreAttributes(group, resolvedVariable, cacheName);
8387
}
8488

8589
export async function GetZarrDims(variable: string) {
86-
const { cache } = useCacheStore.getState();
87-
const { initStore } = useGlobalStore.getState();
88-
const cacheName = `${initStore}_${variable}_meta`;
89-
console.log(`Getting dimensions for variable "${variable}" with cache key "${cacheName}"`);
90-
// Cache hit — same for both store types
91-
if (cache.has(cacheName)) {
92-
const meta = cache.get(cacheName);
93-
const dimNames = (meta._ARRAY_DIMENSIONS ?? meta.dimensionNames) as string[] | undefined;
94-
console.log('Cache hit for dimensions:', { meta, dimNames });
95-
const dimArrays: unknown[] = [];
96-
const dimUnits: unknown[] = [];
90+
const { cache } = useCacheStore.getState();
91+
const { initStore } = useGlobalStore.getState();
92+
const cacheName = `${initStore}_${variable}_meta`;
93+
// Cache hit — same for both store types
94+
if (cache.has(cacheName)) {
95+
const meta = cache.get(cacheName);
96+
const dimNames = (meta._ARRAY_DIMENSIONS ?? meta.dimensionNames) as
97+
| string[]
98+
| undefined;
99+
const dimArrays: unknown[] = [];
100+
const dimUnits: unknown[] = [];
97101

98-
if (dimNames) {
99-
for (const dim of dimNames) {
100-
dimArrays.push(cache.get(`${initStore}_${dim}`) ?? [0]);
101-
dimUnits.push(cache.get(`${initStore}_${dim}_meta`)?.units ?? null);
102-
}
103-
} else {
104-
for (const dimLength of meta.shape) {
105-
dimArrays.push(Array(dimLength).fill(0));
106-
dimUnits.push('Default');
107-
}
108-
}
109-
return { dimNames: dimNames ?? Array(meta.shape.length).fill('Default'), dimArrays, dimUnits };
110-
}
102+
if (dimNames) {
103+
for (const dim of dimNames) {
104+
dimArrays.push(cache.get(`${initStore}_${dim}`) ?? [0]);
105+
dimUnits.push(
106+
cache.get(`${initStore}_${dim}_meta`)?.units ?? null,
107+
);
108+
}
109+
} else {
110+
for (const dimLength of meta.shape) {
111+
dimArrays.push(Array(dimLength).fill(0));
112+
dimUnits.push("Default");
113+
}
114+
}
115+
return {
116+
dimNames: dimNames ?? Array(meta.shape.length).fill("Default"),
117+
dimArrays,
118+
dimUnits,
119+
};
120+
}
111121

112-
const group = await useZarrStore.getState().currentStore;
113-
if (!group) throw new Error(`Failed to open store: ${initStore}`);
122+
const group = await useZarrStore.getState().currentStore;
123+
if (!group) throw new Error(`Failed to open store: ${initStore}`);
114124

115-
if (group.store instanceof IcechunkStore) {
116-
console.log('Fetching dimensions from icechunk store for variable:', variable);
117-
return getIcechunkDims(group as zarr.Group<IcechunkStore>, variable, initStore);
118-
}
119-
return getFetchStoreDims(group, variable, initStore);
125+
if (group.store instanceof IcechunkStore) {
126+
return getIcechunkDims(
127+
group as zarr.Group<IcechunkStore>,
128+
variable,
129+
initStore,
130+
);
131+
}
132+
return getFetchStoreDims(group, variable, initStore);
120133
}
121134

122-
export async function GetAttributes(thisVariable? : string){
123-
const {initStore, variable } = useGlobalStore.getState();
124-
const {cache} = useCacheStore.getState();
125-
const {useNC} = useZarrStore.getState();
126-
const cacheName = `${initStore}_${thisVariable?? variable}_meta`
127-
if (cache.has(cacheName)){
128-
const meta = cache.get(cacheName)
129-
return useNC ? meta.attributes : meta;
130-
}
131-
else {
132-
if (useNC){
133-
const meta = await GetNCMetadata(thisVariable)
134-
return meta.attributes
135-
} else {
136-
const meta = await GetZarrAttributes(thisVariable)
137-
return meta
138-
}
139-
}
135+
export async function GetAttributes(thisVariable?: string) {
136+
const { initStore, variable } = useGlobalStore.getState();
137+
const { cache } = useCacheStore.getState();
138+
const { useNC } = useZarrStore.getState();
139+
const cacheName = `${initStore}_${thisVariable ?? variable}_meta`;
140+
if (cache.has(cacheName)) {
141+
const meta = cache.get(cacheName);
142+
return useNC ? meta.attributes : meta;
143+
} else {
144+
if (useNC) {
145+
const meta = await GetNCMetadata(thisVariable);
146+
return meta.attributes;
147+
} else {
148+
const meta = await GetZarrAttributes(thisVariable);
149+
return meta;
150+
}
151+
}
140152
}

0 commit comments

Comments
 (0)