atlas-python 0.9.0__cp310-abi3-win_amd64.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- atlas/__init__.py +4 -0
- atlas/__init__.pyi +332 -0
- atlas/_atlas.pyd +0 -0
- atlas/py.typed +0 -0
- atlas/xarray.py +578 -0
- atlas_python-0.9.0.dist-info/METADATA +17 -0
- atlas_python-0.9.0.dist-info/RECORD +9 -0
- atlas_python-0.9.0.dist-info/WHEEL +4 -0
- atlas_python-0.9.0.dist-info/sboms/atlas-python.cyclonedx.json +8212 -0
atlas/xarray.py
ADDED
|
@@ -0,0 +1,578 @@
|
|
|
1
|
+
"""xarray integration for atlas.
|
|
2
|
+
|
|
3
|
+
`xarray` and `dask` are required dependencies.
|
|
4
|
+
|
|
5
|
+
Reads automatically return dask-backed variables when an array was stored with a
|
|
6
|
+
non-trivial chunk shape (`chunk_shape != shape`); full-shape arrays come back
|
|
7
|
+
eager as numpy. The dask chunks mirror the on-disk chunk grid one-to-one.
|
|
8
|
+
|
|
9
|
+
Bulk ingestion — Atlas itself batches; explicit flush on the atlas persists:
|
|
10
|
+
with atlas:
|
|
11
|
+
for nc_path in nc_paths:
|
|
12
|
+
ds = xr.open_dataset(nc_path)
|
|
13
|
+
atlas.add_xr_dataset(ds, name=Path(nc_path).stem)
|
|
14
|
+
# atlas.close() (== flush) runs on __exit__.
|
|
15
|
+
"""
|
|
16
|
+
from __future__ import annotations
|
|
17
|
+
|
|
18
|
+
import itertools
|
|
19
|
+
import json
|
|
20
|
+
import time
|
|
21
|
+
from concurrent.futures import Future, ThreadPoolExecutor
|
|
22
|
+
from typing import TYPE_CHECKING, Any, Iterator, Optional, Sequence
|
|
23
|
+
|
|
24
|
+
import numpy as np
|
|
25
|
+
|
|
26
|
+
from ._atlas import log_chunk_event as _log_chunk_event
|
|
27
|
+
|
|
28
|
+
if TYPE_CHECKING:
|
|
29
|
+
import xarray as xr
|
|
30
|
+
|
|
31
|
+
from ._atlas import Atlas, DatasetView
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
_COORDS_ATTR = "_pyatlas_coords"
|
|
35
|
+
_JSON_PREFIX = "json:"
|
|
36
|
+
|
|
37
|
+
_NUMPY_TO_ATLAS = {
|
|
38
|
+
np.dtype("int8"): "int8",
|
|
39
|
+
np.dtype("int16"): "int16",
|
|
40
|
+
np.dtype("int32"): "int32",
|
|
41
|
+
np.dtype("int64"): "int64",
|
|
42
|
+
np.dtype("uint8"): "uint8",
|
|
43
|
+
np.dtype("uint16"): "uint16",
|
|
44
|
+
np.dtype("uint32"): "uint32",
|
|
45
|
+
np.dtype("uint64"): "uint64",
|
|
46
|
+
np.dtype("float32"): "float32",
|
|
47
|
+
np.dtype("float64"): "float64",
|
|
48
|
+
np.dtype("datetime64[ns]"): "timestamp_nanoseconds",
|
|
49
|
+
}
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def _np_to_atlas_dtype(np_dtype: np.dtype) -> str:
|
|
53
|
+
if np_dtype in _NUMPY_TO_ATLAS:
|
|
54
|
+
return _NUMPY_TO_ATLAS[np_dtype]
|
|
55
|
+
# Object (Python str/bytes) and fixed-size byte/unicode strings all
|
|
56
|
+
# become variable-length atlas strings.
|
|
57
|
+
if np_dtype.kind in ("O", "S", "U"):
|
|
58
|
+
return "string"
|
|
59
|
+
supported = ", ".join(sorted(set(_NUMPY_TO_ATLAS.values())))
|
|
60
|
+
raise NotImplementedError(
|
|
61
|
+
f"numpy dtype {np_dtype!r} is not supported by atlas "
|
|
62
|
+
f"(supported: {supported}, plus object/bytes/unicode → string)"
|
|
63
|
+
)
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def _sanitize_str(s: str) -> str:
|
|
67
|
+
"""Strip lone Unicode surrogates from a Python str.
|
|
68
|
+
|
|
69
|
+
NetCDF backends often surface byte attrs as Python strs that were decoded
|
|
70
|
+
with ``errors='surrogateescape'``; the resulting strs hold pseudo-codepoints
|
|
71
|
+
in U+DC80..U+DCFF which Rust's UTF-8 strs can't represent. We try to recover
|
|
72
|
+
the original bytes via surrogateescape and re-decode as UTF-8 (the common
|
|
73
|
+
case: bytes that *were* valid UTF-8 all along but were treated as Latin-1
|
|
74
|
+
upstream); if that fails we fall back to lossy replacement.
|
|
75
|
+
"""
|
|
76
|
+
try:
|
|
77
|
+
s.encode("utf-8")
|
|
78
|
+
return s
|
|
79
|
+
except UnicodeEncodeError:
|
|
80
|
+
pass
|
|
81
|
+
raw = s.encode("utf-8", errors="surrogateescape")
|
|
82
|
+
try:
|
|
83
|
+
return raw.decode("utf-8")
|
|
84
|
+
except UnicodeDecodeError:
|
|
85
|
+
return raw.decode("utf-8", errors="replace")
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
def _encode_attr_value(value: Any) -> Any:
|
|
89
|
+
"""Coerce an xarray attr value to something atlas can store.
|
|
90
|
+
|
|
91
|
+
Primitives (bool/int/float/str) are returned as-is. Everything else is
|
|
92
|
+
JSON-encoded and prefixed with ``json:`` so it can be decoded losslessly
|
|
93
|
+
on read. Numpy scalars are unwrapped first. Numpy arrays are converted to
|
|
94
|
+
lists. Values that can't be JSON-serialised raise ``TypeError``.
|
|
95
|
+
"""
|
|
96
|
+
# Unwrap numpy scalars
|
|
97
|
+
if isinstance(value, np.generic):
|
|
98
|
+
value = value.item()
|
|
99
|
+
|
|
100
|
+
if isinstance(value, str):
|
|
101
|
+
return _sanitize_str(value)
|
|
102
|
+
if isinstance(value, (bool, int, float)):
|
|
103
|
+
return value
|
|
104
|
+
|
|
105
|
+
# Convert numpy arrays to nested lists
|
|
106
|
+
if isinstance(value, np.ndarray):
|
|
107
|
+
value = value.tolist()
|
|
108
|
+
|
|
109
|
+
try:
|
|
110
|
+
return _JSON_PREFIX + json.dumps(value)
|
|
111
|
+
except (TypeError, ValueError) as exc:
|
|
112
|
+
raise TypeError(
|
|
113
|
+
f"attribute value of type {type(value).__name__} is not JSON-serialisable "
|
|
114
|
+
f"and cannot be stored: {value!r}"
|
|
115
|
+
) from exc
|
|
116
|
+
|
|
117
|
+
|
|
118
|
+
def _decode_attr_value(value: Any) -> Any:
|
|
119
|
+
"""Inverse of :func:`_encode_attr_value`."""
|
|
120
|
+
if isinstance(value, str) and value.startswith(_JSON_PREFIX):
|
|
121
|
+
return json.loads(value[len(_JSON_PREFIX):])
|
|
122
|
+
return value
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
def _normalize_fill_value(value: Any, np_dtype: np.dtype) -> Any:
|
|
126
|
+
"""Coerce an xarray `_FillValue` attribute into a Python scalar matching the array dtype.
|
|
127
|
+
|
|
128
|
+
The binding's `define_array(fill_value=...)` expects a plain Python scalar
|
|
129
|
+
that is type-consistent with the array dtype. xarray often stores `_FillValue`
|
|
130
|
+
as a 0-D numpy array or numpy scalar; this unwraps that and (for datetime64
|
|
131
|
+
arrays) reinterprets the value as nanoseconds since the epoch. For
|
|
132
|
+
string-kind arrays (object/bytes/unicode), `bytes` are decoded to `str`
|
|
133
|
+
since NetCDF stores fixed-width string fill values as bytes.
|
|
134
|
+
"""
|
|
135
|
+
if value is None:
|
|
136
|
+
return None
|
|
137
|
+
if isinstance(value, np.ndarray) and value.ndim == 0:
|
|
138
|
+
value = value.item() if np_dtype.kind != "M" else value.view(np.int64).item()
|
|
139
|
+
elif isinstance(value, np.generic):
|
|
140
|
+
if isinstance(value, np.datetime64):
|
|
141
|
+
return value.astype("datetime64[ns]").view(np.int64).item()
|
|
142
|
+
value = value.item()
|
|
143
|
+
if np_dtype.kind in ("O", "S", "U") and isinstance(value, bytes):
|
|
144
|
+
return value.decode("utf-8", errors="replace")
|
|
145
|
+
return value
|
|
146
|
+
|
|
147
|
+
|
|
148
|
+
def _is_dask_array(arr: Any) -> bool:
|
|
149
|
+
"""Return True if `arr` is a `dask.array.Array`. False if dask isn't installed."""
|
|
150
|
+
try:
|
|
151
|
+
import dask.array as da
|
|
152
|
+
except ImportError:
|
|
153
|
+
return False
|
|
154
|
+
return isinstance(arr, da.Array)
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
def _dask_chunk_shape(arr: Any) -> list[int]:
|
|
158
|
+
"""First chunk size along each dim of a dask array — used as the atlas chunk_shape."""
|
|
159
|
+
return [c[0] for c in arr.chunks]
|
|
160
|
+
|
|
161
|
+
|
|
162
|
+
def _contiguous(a: np.ndarray) -> np.ndarray:
|
|
163
|
+
# np.ascontiguousarray promotes 0-D arrays to 1-D (ndmin=1 default),
|
|
164
|
+
# which breaks scalar-array writes. Force-copy non-contiguous arrays
|
|
165
|
+
# via np.asarray + .copy(order='C') instead, which preserves rank.
|
|
166
|
+
if a.flags["C_CONTIGUOUS"]:
|
|
167
|
+
return a
|
|
168
|
+
return a.copy(order="C")
|
|
169
|
+
|
|
170
|
+
|
|
171
|
+
# Tuned for the "many small NetCDF chunks" case. batch_size controls how many
|
|
172
|
+
# dask blocks ride one scheduler invocation (lower dask plumbing overhead);
|
|
173
|
+
# prefetch_depth bounds peak memory at batch_size * prefetch_depth chunks.
|
|
174
|
+
_DEFAULT_BATCH_SIZE = 8
|
|
175
|
+
_DEFAULT_PREFETCH_DEPTH = 2
|
|
176
|
+
|
|
177
|
+
|
|
178
|
+
def _iter_blocks(
|
|
179
|
+
arr: Any,
|
|
180
|
+
var_name: str = "",
|
|
181
|
+
batch_size: int = _DEFAULT_BATCH_SIZE,
|
|
182
|
+
prefetch_depth: int = _DEFAULT_PREFETCH_DEPTH,
|
|
183
|
+
) -> Iterator[tuple[list[int], np.ndarray]]:
|
|
184
|
+
"""Yield ``(start_index, block_np)`` for each chunk in `arr`.
|
|
185
|
+
|
|
186
|
+
Numpy-backed inputs are yielded as one full-shape block. Dask-backed inputs
|
|
187
|
+
are prefetched in batches: a single background thread runs
|
|
188
|
+
``dask.compute(*K)`` on the next ``batch_size`` blocks while the main thread
|
|
189
|
+
consumes already-materialised blocks, capping in-flight work at
|
|
190
|
+
``prefetch_depth`` batches. This collapses N per-chunk dask scheduler
|
|
191
|
+
invocations into N/batch_size, and overlaps NetCDF I/O with atlas writes.
|
|
192
|
+
|
|
193
|
+
Emits one ``event=dask_compute`` debug event per fulfilled batch via the
|
|
194
|
+
Rust tracing subscriber (see ``log_chunk_event``).
|
|
195
|
+
"""
|
|
196
|
+
if not _is_dask_array(arr):
|
|
197
|
+
a = np.asarray(arr)
|
|
198
|
+
yield [0] * a.ndim, _contiguous(a)
|
|
199
|
+
return
|
|
200
|
+
|
|
201
|
+
import dask
|
|
202
|
+
|
|
203
|
+
chunks = arr.chunks
|
|
204
|
+
offsets = [[0, *itertools.accumulate(c)][:-1] for c in chunks]
|
|
205
|
+
pairs: list[tuple[list[int], Any]] = []
|
|
206
|
+
for block_idx in itertools.product(*[range(len(c)) for c in chunks]):
|
|
207
|
+
start = [offsets[d][i] for d, i in enumerate(block_idx)]
|
|
208
|
+
pairs.append((start, arr.blocks[block_idx]))
|
|
209
|
+
|
|
210
|
+
def _fetch(batch: list[tuple[list[int], Any]]) -> list[tuple[list[int], np.ndarray]]:
|
|
211
|
+
starts = [s for s, _ in batch]
|
|
212
|
+
delayed = [b for _, b in batch]
|
|
213
|
+
t0 = time.perf_counter_ns()
|
|
214
|
+
materialised = dask.compute(*delayed)
|
|
215
|
+
elapsed_us = (time.perf_counter_ns() - t0) // 1000
|
|
216
|
+
_log_chunk_event("dask_compute", var_name, elapsed_us, chunks=len(batch))
|
|
217
|
+
return [(s, _contiguous(b)) for s, b in zip(starts, materialised)]
|
|
218
|
+
|
|
219
|
+
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
220
|
+
in_flight: list[Future] = []
|
|
221
|
+
cursor = 0
|
|
222
|
+
while cursor < len(pairs) and len(in_flight) < prefetch_depth:
|
|
223
|
+
in_flight.append(executor.submit(_fetch, pairs[cursor : cursor + batch_size]))
|
|
224
|
+
cursor += batch_size
|
|
225
|
+
while in_flight:
|
|
226
|
+
future = in_flight.pop(0)
|
|
227
|
+
for start, block in future.result():
|
|
228
|
+
yield start, block
|
|
229
|
+
if cursor < len(pairs):
|
|
230
|
+
in_flight.append(
|
|
231
|
+
executor.submit(_fetch, pairs[cursor : cursor + batch_size])
|
|
232
|
+
)
|
|
233
|
+
cursor += batch_size
|
|
234
|
+
|
|
235
|
+
|
|
236
|
+
def _write_xarray_to_view(
|
|
237
|
+
view: "DatasetView",
|
|
238
|
+
ds: "xr.Dataset",
|
|
239
|
+
chunks: Optional[dict[str, Sequence[int]]] = None,
|
|
240
|
+
) -> None:
|
|
241
|
+
"""Populate an empty `DatasetView` with the contents of an xarray Dataset.
|
|
242
|
+
|
|
243
|
+
Writes every coordinate and data variable as an atlas array, the
|
|
244
|
+
coordinate names as ``_pyatlas_coords`` (JSON list), all dataset attrs,
|
|
245
|
+
and all per-variable attrs flattened as ``{var}.{attr}``.
|
|
246
|
+
"""
|
|
247
|
+
coord_names = [str(n) for n in ds.coords.keys()]
|
|
248
|
+
|
|
249
|
+
# Write coords first, then data_vars. Order doesn't matter to atlas but
|
|
250
|
+
# makes the on-disk file layout predictable.
|
|
251
|
+
for var_name in coord_names + [str(n) for n in ds.data_vars.keys()]:
|
|
252
|
+
var = ds[var_name]
|
|
253
|
+
atlas_dtype = _np_to_atlas_dtype(np.dtype(var.dtype))
|
|
254
|
+
dims = [str(d) for d in var.dims]
|
|
255
|
+
shape = [int(s) for s in var.shape]
|
|
256
|
+
|
|
257
|
+
# Pick the atlas chunk_shape:
|
|
258
|
+
# 1. explicit user override via the `chunks=` kwarg, else
|
|
259
|
+
# 2. the dask chunk shape if the variable is dask-backed, else
|
|
260
|
+
# 3. None (atlas defaults to a single full-shape chunk).
|
|
261
|
+
if chunks is not None and var_name in chunks:
|
|
262
|
+
chunk_shape: Optional[list[int]] = [int(s) for s in chunks[var_name]]
|
|
263
|
+
elif _is_dask_array(var.data):
|
|
264
|
+
chunk_shape = _dask_chunk_shape(var.data)
|
|
265
|
+
else:
|
|
266
|
+
chunk_shape = None
|
|
267
|
+
|
|
268
|
+
# Extract the CF/netCDF `_FillValue` attribute (if any) so it's passed
|
|
269
|
+
# to define_array as a typed fill value, not stored as a flattened atlas
|
|
270
|
+
# attribute. Copy attrs first to avoid mutating the user's Dataset.
|
|
271
|
+
var_attrs = dict(var.attrs)
|
|
272
|
+
fill_value = _normalize_fill_value(var_attrs.pop("_FillValue", None), var.dtype)
|
|
273
|
+
|
|
274
|
+
view.define_array(
|
|
275
|
+
var_name,
|
|
276
|
+
dtype=atlas_dtype,
|
|
277
|
+
dims=dims,
|
|
278
|
+
shape=shape,
|
|
279
|
+
chunk_shape=chunk_shape,
|
|
280
|
+
fill_value=fill_value,
|
|
281
|
+
)
|
|
282
|
+
|
|
283
|
+
# Stream blocks: prefetched batches for dask-backed data, a single
|
|
284
|
+
# full-shape block for numpy-backed data.
|
|
285
|
+
for start, block in _iter_blocks(var.data, var_name=var_name):
|
|
286
|
+
# TimestampNs columns: the bindings accept np.int64 only; cast the
|
|
287
|
+
# numpy datetime64 view to int64 without copying.
|
|
288
|
+
if block.dtype.kind == "M":
|
|
289
|
+
block = block.view(np.int64)
|
|
290
|
+
t0 = time.perf_counter_ns()
|
|
291
|
+
view.write_array(var_name, start=start, data=block)
|
|
292
|
+
_log_chunk_event(
|
|
293
|
+
"write",
|
|
294
|
+
var_name,
|
|
295
|
+
(time.perf_counter_ns() - t0) // 1000,
|
|
296
|
+
bytes=int(block.nbytes),
|
|
297
|
+
)
|
|
298
|
+
|
|
299
|
+
# Per-variable attrs → flattened as `{var}.{attr}` (sans `_FillValue`).
|
|
300
|
+
for attr_key, attr_val in var_attrs.items():
|
|
301
|
+
encoded = _encode_attr_value(attr_val)
|
|
302
|
+
view.set_attribute(_sanitize_str(f"{var_name}.{attr_key}"), encoded)
|
|
303
|
+
|
|
304
|
+
# Dataset-level attrs
|
|
305
|
+
for attr_key, attr_val in ds.attrs.items():
|
|
306
|
+
encoded = _encode_attr_value(attr_val)
|
|
307
|
+
view.set_attribute(_sanitize_str(str(attr_key)), encoded)
|
|
308
|
+
|
|
309
|
+
# Marker so we can faithfully restore coord/var distinction on read.
|
|
310
|
+
view.set_attribute(_COORDS_ATTR, json.dumps(coord_names))
|
|
311
|
+
|
|
312
|
+
|
|
313
|
+
_ATLAS_TO_NUMPY = {atlas: np_dt for np_dt, atlas in _NUMPY_TO_ATLAS.items()}
|
|
314
|
+
|
|
315
|
+
|
|
316
|
+
def _atlas_to_numpy_dtype(atlas_dtype: str) -> np.dtype:
|
|
317
|
+
"""Numpy dtype that `view.read_array` returns for a given atlas dtype string."""
|
|
318
|
+
if atlas_dtype in _ATLAS_TO_NUMPY:
|
|
319
|
+
return _ATLAS_TO_NUMPY[atlas_dtype]
|
|
320
|
+
if atlas_dtype == "string":
|
|
321
|
+
return np.dtype("object")
|
|
322
|
+
raise NotImplementedError(
|
|
323
|
+
f"atlas dtype {atlas_dtype!r} is not supported on the dask read path"
|
|
324
|
+
)
|
|
325
|
+
|
|
326
|
+
|
|
327
|
+
def _dask_chunks_for(shape: Sequence[int], chunk_shape: Sequence[int]) -> tuple:
|
|
328
|
+
"""Per-dim chunk-length tuples in the form dask expects."""
|
|
329
|
+
chunks: list[tuple[int, ...]] = []
|
|
330
|
+
for dim_size, dim_chunk in zip(shape, chunk_shape):
|
|
331
|
+
if dim_chunk <= 0 or dim_size == 0:
|
|
332
|
+
chunks.append((dim_size,))
|
|
333
|
+
continue
|
|
334
|
+
full = dim_size // dim_chunk
|
|
335
|
+
rem = dim_size - full * dim_chunk
|
|
336
|
+
c = (dim_chunk,) * full + ((rem,) if rem else ())
|
|
337
|
+
chunks.append(c if c else (0,))
|
|
338
|
+
return tuple(chunks)
|
|
339
|
+
|
|
340
|
+
|
|
341
|
+
def _view_to_dask_array(view: "DatasetView", name: str) -> Any:
|
|
342
|
+
"""Build a `dask.array.Array` that lazily reads `name` chunk-by-chunk.
|
|
343
|
+
|
|
344
|
+
Each on-disk chunk becomes one dask task; values are fetched via
|
|
345
|
+
`view.read_array(name, start, block_shape)` on demand. Used by
|
|
346
|
+
`_view_to_xarray` when an array's `chunk_shape != shape`.
|
|
347
|
+
"""
|
|
348
|
+
import dask
|
|
349
|
+
import dask.array as da
|
|
350
|
+
|
|
351
|
+
meta = view.array_meta(name)
|
|
352
|
+
assert meta is not None, f"array {name!r} not found in view {view.name!r}"
|
|
353
|
+
shape: list[int] = list(meta["shape"])
|
|
354
|
+
chunk_shape: list[int] = list(meta["chunk_shape"])
|
|
355
|
+
np_dtype = _atlas_to_numpy_dtype(meta["dtype"])
|
|
356
|
+
|
|
357
|
+
per_dim_chunks = _dask_chunks_for(shape, chunk_shape)
|
|
358
|
+
offsets = [
|
|
359
|
+
[0, *itertools.accumulate(dim_chunks)][:-1] for dim_chunks in per_dim_chunks
|
|
360
|
+
]
|
|
361
|
+
|
|
362
|
+
def _read_block(start: list[int], block_shape: list[int]) -> np.ndarray:
|
|
363
|
+
return view.read_array(name, start=start, shape=block_shape)
|
|
364
|
+
|
|
365
|
+
def _nested_blocks(axis: int, prefix_start: list[int], prefix_shape: list[int]):
|
|
366
|
+
if axis == len(shape):
|
|
367
|
+
block_shape = list(prefix_shape)
|
|
368
|
+
block_start = list(prefix_start)
|
|
369
|
+
delayed = dask.delayed(_read_block)(block_start, block_shape)
|
|
370
|
+
return da.from_delayed(delayed, shape=tuple(block_shape), dtype=np_dtype)
|
|
371
|
+
return [
|
|
372
|
+
_nested_blocks(
|
|
373
|
+
axis + 1,
|
|
374
|
+
prefix_start + [offsets[axis][i]],
|
|
375
|
+
prefix_shape + [per_dim_chunks[axis][i]],
|
|
376
|
+
)
|
|
377
|
+
for i in range(len(per_dim_chunks[axis]))
|
|
378
|
+
]
|
|
379
|
+
|
|
380
|
+
nested = _nested_blocks(0, [], [])
|
|
381
|
+
return da.block(nested)
|
|
382
|
+
|
|
383
|
+
|
|
384
|
+
def _view_to_xarray(view: "DatasetView", force_lazy: bool = False) -> "xr.Dataset":
|
|
385
|
+
"""Convert an atlas `DatasetView` into an xarray Dataset.
|
|
386
|
+
|
|
387
|
+
Variables stored with `chunk_shape != shape` come back dask-backed (one task
|
|
388
|
+
per on-disk chunk); full-shape arrays come back eager as numpy unless
|
|
389
|
+
``force_lazy=True``, in which case they're wrapped in a single-chunk
|
|
390
|
+
`dask.array` so the returned Dataset is uniformly lazy (used by
|
|
391
|
+
`_atlas_to_xarray_many` so concat returns a lazy graph regardless of the
|
|
392
|
+
source chunk_shape).
|
|
393
|
+
"""
|
|
394
|
+
import xarray as xr
|
|
395
|
+
|
|
396
|
+
array_names = list(view.list_arrays())
|
|
397
|
+
array_name_set = set(array_names)
|
|
398
|
+
|
|
399
|
+
# Coord/var assignment
|
|
400
|
+
coords_marker = view.get_attribute(_COORDS_ATTR)
|
|
401
|
+
if isinstance(coords_marker, str):
|
|
402
|
+
try:
|
|
403
|
+
coord_names = set(json.loads(coords_marker))
|
|
404
|
+
except (TypeError, ValueError):
|
|
405
|
+
coord_names = set()
|
|
406
|
+
else:
|
|
407
|
+
# Fallback heuristic: 1-D array whose single dim matches its name.
|
|
408
|
+
coord_names = set()
|
|
409
|
+
for name in array_names:
|
|
410
|
+
meta = view.array_meta(name)
|
|
411
|
+
if meta is None:
|
|
412
|
+
continue
|
|
413
|
+
if (
|
|
414
|
+
len(meta["dimension_names"]) == 1
|
|
415
|
+
and meta["dimension_names"][0] == name
|
|
416
|
+
):
|
|
417
|
+
coord_names.add(name)
|
|
418
|
+
|
|
419
|
+
# Pull array data. Chunked arrays (chunk_shape != shape) come back as a
|
|
420
|
+
# lazy dask.array; full-shape and 0-D arrays come back eager.
|
|
421
|
+
data_vars: dict[str, tuple] = {}
|
|
422
|
+
coords: dict[str, tuple] = {}
|
|
423
|
+
for name in array_names:
|
|
424
|
+
meta = view.array_meta(name)
|
|
425
|
+
if meta is None:
|
|
426
|
+
continue
|
|
427
|
+
shape = list(meta["shape"])
|
|
428
|
+
chunk_shape = list(meta["chunk_shape"])
|
|
429
|
+
if not shape or chunk_shape == shape:
|
|
430
|
+
arr = view.read_array(name)
|
|
431
|
+
if arr is None:
|
|
432
|
+
continue
|
|
433
|
+
if force_lazy and shape:
|
|
434
|
+
import dask.array as da
|
|
435
|
+
arr = da.from_array(arr, chunks=tuple(shape))
|
|
436
|
+
else:
|
|
437
|
+
arr = _view_to_dask_array(view, name)
|
|
438
|
+
dims = list(meta["dimension_names"])
|
|
439
|
+
entry = (dims, arr, {}) # placeholder for per-var attrs; filled below
|
|
440
|
+
if name in coord_names:
|
|
441
|
+
coords[name] = entry
|
|
442
|
+
else:
|
|
443
|
+
data_vars[name] = entry
|
|
444
|
+
|
|
445
|
+
# Split dataset attrs vs flattened per-var attrs
|
|
446
|
+
raw_attrs = dict(view.attributes())
|
|
447
|
+
raw_attrs.pop(_COORDS_ATTR, None)
|
|
448
|
+
|
|
449
|
+
dataset_attrs: dict[str, Any] = {}
|
|
450
|
+
per_var_attrs: dict[str, dict[str, Any]] = {n: {} for n in array_names}
|
|
451
|
+
for key, value in raw_attrs.items():
|
|
452
|
+
if "." in key:
|
|
453
|
+
var, rest = key.split(".", 1)
|
|
454
|
+
if var in array_name_set:
|
|
455
|
+
per_var_attrs[var][rest] = _decode_attr_value(value)
|
|
456
|
+
continue
|
|
457
|
+
dataset_attrs[key] = _decode_attr_value(value)
|
|
458
|
+
|
|
459
|
+
# Restore _FillValue for any array that was defined with one.
|
|
460
|
+
for name in array_names:
|
|
461
|
+
fv = view.array_fill_value(name)
|
|
462
|
+
if fv is None:
|
|
463
|
+
continue
|
|
464
|
+
per_var_attrs.setdefault(name, {})
|
|
465
|
+
per_var_attrs[name]["_FillValue"] = fv
|
|
466
|
+
|
|
467
|
+
# Inject per-var attrs into the (dims, data, attrs) triples
|
|
468
|
+
def _with_attrs(name: str, triple: tuple) -> tuple:
|
|
469
|
+
dims, arr, _ = triple
|
|
470
|
+
return (dims, arr, per_var_attrs.get(name, {}))
|
|
471
|
+
|
|
472
|
+
data_vars = {n: _with_attrs(n, t) for n, t in data_vars.items()}
|
|
473
|
+
coords = {n: _with_attrs(n, t) for n, t in coords.items()}
|
|
474
|
+
|
|
475
|
+
return xr.Dataset(data_vars=data_vars, coords=coords, attrs=dataset_attrs)
|
|
476
|
+
|
|
477
|
+
|
|
478
|
+
def _atlas_to_xarray_many(
|
|
479
|
+
atlas: "Atlas",
|
|
480
|
+
names: list[str],
|
|
481
|
+
concat_dim: str = "dataset",
|
|
482
|
+
parallel: bool = True, # noqa: ARG001 — kept for API compat; ignored
|
|
483
|
+
) -> "xr.Dataset":
|
|
484
|
+
"""Open many atlas datasets and stack them into one xr.Dataset along
|
|
485
|
+
`concat_dim`. atlas-native equivalent of `xr.open_mfdataset(...)`.
|
|
486
|
+
|
|
487
|
+
Implementation: opens the first dataset to discover the schema (vars,
|
|
488
|
+
dims, dtypes, coords, per-var attrs), then for each data variable calls
|
|
489
|
+
`Atlas.read_array_across` to bulk-read across all `names` in one Rust
|
|
490
|
+
call. The N reads share one `RwLock::read` guard on the shared physical
|
|
491
|
+
file and dispatch concurrently on the tokio runtime — avoids the N
|
|
492
|
+
Python ↔ Rust round-trips the prior dask-delayed implementation paid.
|
|
493
|
+
|
|
494
|
+
Returns eager numpy-backed arrays of shape `(len(names), *original_shape)`.
|
|
495
|
+
Wrap with `.chunk(...)` downstream if you need dask laziness.
|
|
496
|
+
|
|
497
|
+
The `parallel` parameter is accepted for API compatibility but no longer
|
|
498
|
+
selects an implementation — the bulk path is always taken.
|
|
499
|
+
"""
|
|
500
|
+
import numpy as np
|
|
501
|
+
import xarray as xr
|
|
502
|
+
|
|
503
|
+
if not names:
|
|
504
|
+
raise ValueError("to_xarray_many: `names` is empty")
|
|
505
|
+
|
|
506
|
+
# Schema discovery from the first dataset (cheap: in-memory meta lookup).
|
|
507
|
+
first_view = atlas.open_dataset(names[0])
|
|
508
|
+
template = _view_to_xarray(first_view, force_lazy=False)
|
|
509
|
+
|
|
510
|
+
# Bulk-read each data variable across all datasets in one Rust call,
|
|
511
|
+
# returning a pre-stacked (N, *shape) numpy array. Skips the Python-side
|
|
512
|
+
# `np.stack` copy that the list-returning `read_array_across` would
|
|
513
|
+
# require — significant on big workloads (a 1000-dataset gridded run
|
|
514
|
+
# saves several seconds of memory bandwidth).
|
|
515
|
+
#
|
|
516
|
+
# Variables are processed serially because each Rust call internally
|
|
517
|
+
# parallelises N reads up to num_cpus — running multiple vars in
|
|
518
|
+
# parallel would oversubscribe CPU.
|
|
519
|
+
data_vars: dict[str, xr.DataArray] = {}
|
|
520
|
+
for var in template.data_vars:
|
|
521
|
+
stacked = atlas.read_array_across_stacked(var, names)
|
|
522
|
+
original_dims = list(template[var].dims)
|
|
523
|
+
original_attrs = dict(template[var].attrs)
|
|
524
|
+
data_vars[var] = xr.DataArray(
|
|
525
|
+
stacked,
|
|
526
|
+
dims=[concat_dim, *original_dims],
|
|
527
|
+
attrs=original_attrs,
|
|
528
|
+
)
|
|
529
|
+
|
|
530
|
+
# Coords + dataset-level attrs come from the first dataset, matching
|
|
531
|
+
# xarray.open_mfdataset(coords="minimal", compat="override") semantics.
|
|
532
|
+
coords = {name: template.coords[name] for name in template.coords}
|
|
533
|
+
coords[concat_dim] = xr.DataArray(np.asarray(names), dims=[concat_dim])
|
|
534
|
+
|
|
535
|
+
return xr.Dataset(data_vars=data_vars, coords=coords, attrs=dict(template.attrs))
|
|
536
|
+
|
|
537
|
+
|
|
538
|
+
def _write_xarray_new_dataset(
|
|
539
|
+
atlas: "Atlas",
|
|
540
|
+
ds: "xr.Dataset",
|
|
541
|
+
name: str,
|
|
542
|
+
chunks: Optional[dict[str, Sequence[int]]] = None,
|
|
543
|
+
) -> None:
|
|
544
|
+
"""Rust-delegated helper: create a fresh atlas dataset and populate it.
|
|
545
|
+
|
|
546
|
+
Both `atlas.add_xr_dataset` and the `ds.atlas.write` accessor route through
|
|
547
|
+
this function.
|
|
548
|
+
"""
|
|
549
|
+
view = atlas.create_dataset(name)
|
|
550
|
+
_write_xarray_to_view(view, ds, chunks=chunks)
|
|
551
|
+
|
|
552
|
+
|
|
553
|
+
# --- xarray accessor ----------------------------------------------------------
|
|
554
|
+
# Registered as `ds.atlas` once `atlas` is imported. The whole module is
|
|
555
|
+
# side-effect-imported from `atlas/__init__.py` so importing `atlas` is
|
|
556
|
+
# enough to activate this.
|
|
557
|
+
|
|
558
|
+
import xarray as _xr # noqa: E402
|
|
559
|
+
|
|
560
|
+
|
|
561
|
+
@_xr.register_dataset_accessor("atlas")
|
|
562
|
+
class _AtlasAccessor:
|
|
563
|
+
"""Methods exposed on an `xr.Dataset` as `ds.atlas.*`."""
|
|
564
|
+
|
|
565
|
+
def __init__(self, ds: "_xr.Dataset") -> None:
|
|
566
|
+
self._ds = ds
|
|
567
|
+
|
|
568
|
+
def write(
|
|
569
|
+
self,
|
|
570
|
+
atlas: "Atlas",
|
|
571
|
+
name: str,
|
|
572
|
+
chunks: Optional[dict[str, Sequence[int]]] = None,
|
|
573
|
+
) -> None:
|
|
574
|
+
"""Append this Dataset to the open atlas store under `name`.
|
|
575
|
+
|
|
576
|
+
Equivalent to `atlas.add_xr_dataset(self_ds, name, chunks)`.
|
|
577
|
+
"""
|
|
578
|
+
atlas.add_xr_dataset(self._ds, name, chunks)
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: atlas-python
|
|
3
|
+
Version: 0.9.0
|
|
4
|
+
Requires-Dist: numpy>=1.23
|
|
5
|
+
Requires-Dist: xarray>=2023.1
|
|
6
|
+
Requires-Dist: dask>=2023.1
|
|
7
|
+
Requires-Dist: zarr>=3 ; extra == 'bench'
|
|
8
|
+
Requires-Dist: numcodecs ; extra == 'bench'
|
|
9
|
+
Requires-Dist: netcdf4 ; extra == 'bench'
|
|
10
|
+
Requires-Dist: obstore>=0.9 ; extra == 'cloud'
|
|
11
|
+
Requires-Dist: pytest ; extra == 'test'
|
|
12
|
+
Requires-Dist: netcdf4 ; extra == 'test'
|
|
13
|
+
Provides-Extra: bench
|
|
14
|
+
Provides-Extra: cloud
|
|
15
|
+
Provides-Extra: test
|
|
16
|
+
Summary: Python bindings for the ATLAS array store
|
|
17
|
+
Requires-Python: >=3.10
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
atlas/__init__.py,sha256=Q2om8rCNbD_dYWw0C4Av_Jfvi7dazDqFxAbJ5dQRqtU,222
|
|
2
|
+
atlas/__init__.pyi,sha256=m5Hcnj7tlLii4B_UIToVQIndd74ke1qvMYG96-_t6Ss,12708
|
|
3
|
+
atlas/_atlas.pyd,sha256=xTxpL7r2rYy2a5s3gsT_PjukYK_GELjHThwdJXY576s,19614208
|
|
4
|
+
atlas/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
5
|
+
atlas/xarray.py,sha256=A80ty9shMNn5ynER48NA9StltSxAxPZ1L0-6zU1syiU,22826
|
|
6
|
+
atlas_python-0.9.0.dist-info/METADATA,sha256=aYbfSEtT97GKtS-zj9zyJJU5-WO3b2jucixZv8U5VBM,537
|
|
7
|
+
atlas_python-0.9.0.dist-info/WHEEL,sha256=OUT0XP5TL9Hq-6CIgsb5m6BAU8pfcNqYjx0xnFDWhNs,96
|
|
8
|
+
atlas_python-0.9.0.dist-info/sboms/atlas-python.cyclonedx.json,sha256=zJj1b8tP1-cpy0i50fjXrXZ7WP9QWEpZTMfqddsTVno,265661
|
|
9
|
+
atlas_python-0.9.0.dist-info/RECORD,,
|