Read Binary Files Lazily With Xarray and Dask

Contents

Read Binary Files Lazily With Xarray and Dask#

This notebook shows the main xarray-binfile workflow:

describe a binary-file convention with FileSpecsGetter
open individual files or a whole collection with engine="binfile"
keep the resulting arrays lazy and Dask-backed
inspect the task graph before executing a calculation
call .compute() only when you want concrete results in memory

import pathlib
import re
import tempfile
from shutil import which

import numpy as np
import xarray as xr
from IPython.display import Markdown, display

import xarray_binfile  # noqa: F401  (registers the .binary_engine accessors)
from xarray_binfile.tutorial import DatasetGenerator, FileSpecsGetter

temporary_directory = tempfile.TemporaryDirectory()
data_directory = pathlib.Path(temporary_directory.name)
data_directory

PosixPath('/tmp/tmpmok4uc8w')

base_coords = {
    "x": np.linspace(0.0, 3.0, num=16, dtype=np.float32),
    "y": np.linspace(-1.0, 1.0, num=12, dtype=np.float32),
    "z": np.linspace(0.0, 2.0, num=8, dtype=np.float32),
}

file_specs_getter = FileSpecsGetter(
    base_coords=base_coords,
    dtype=np.float32,
    filename_template="{name}-{digits:04}.bin",
    filename_regex=re.compile(r"(?P<name>\w+)-(?P<digits>\d{4})\.bin"),
)

dataset_generator = DatasetGenerator(file_specs_getter.reader)
filenames = (
    file_specs_getter.filename_template.format(name=name, digits=step)
    for name in ("ux", "uy", "uz")
    for step in range(2)
)
source_dataset = dataset_generator(map(pathlib.Path, filenames))
source_dataset.binary_engine.to_file(file_specs_getter.writer, data_directory)
sorted(path.name for path in data_directory.glob("*.bin"))[:6]

['ux-0000.bin',
 'ux-0001.bin',
 'uy-0000.bin',
 'uy-0001.bin',
 'uz-0000.bin',
 'uz-0001.bin']

Open many files lazily#

xr.open_mfdataset combines the files by coordinates and keeps the arrays lazy because we pass Dask chunks. The example data here is intentionally small enough for a fast docs build, but the same chunked pattern is the one you use when the full dataset no longer fits comfortably in memory.

lazy_dataset = xr.open_mfdataset(
    sorted(data_directory.glob("*.bin")),
    engine="binfile",
    read_specs_getter=file_specs_getter.reader,
    chunks={"time": 1},
    parallel=True,
)
lazy_dataset

<xarray.Dataset> Size: 37kB
Dimensions:  (x: 16, y: 12, z: 8, time: 2)
Coordinates:
  * x        (x) float32 64B 0.0 0.2 0.4 0.6 0.8 1.0 ... 2.0 2.2 2.4 2.6 2.8 3.0
  * y        (y) float32 48B -1.0 -0.8182 -0.6364 -0.4545 ... 0.6364 0.8182 1.0
  * z        (z) float32 32B 0.0 0.2857 0.5714 0.8571 1.143 1.429 1.714 2.0
  * time     (time) int64 16B 0 1
Data variables:
    ux       (x, y, z, time) float32 12kB dask.array<chunksize=(16, 12, 8, 1), meta=np.ndarray>
    uy       (x, y, z, time) float32 12kB dask.array<chunksize=(16, 12, 8, 1), meta=np.ndarray>
    uz       (x, y, z, time) float32 12kB dask.array<chunksize=(16, 12, 8, 1), meta=np.ndarray>

ux = lazy_dataset["ux"]
print(type(ux.data))
print(ux.chunks)
print(
    f"Estimated in-memory size if fully loaded: {lazy_dataset.nbytes / 1024**2:.2f} MiB"
)
ux.data

<class 'dask.array.core.Array'>
((16,), (12,), (8,), (1, 1))
Estimated in-memory size if fully loaded: 0.04 MiB

	Array	Chunk
Bytes	12.00 kiB	6.00 kiB
Shape	(16, 12, 8, 2)	(16, 12, 8, 1)
Dask graph	2 chunks in 5 graph layers
Data type	float32 numpy.ndarray

The dask.array representation is the important part here. At this stage the notebook has built a task graph, not loaded every chunk into memory. That is what lets the same workflow scale to much larger datasets.

lazy_speed = np.sqrt(
    lazy_dataset["ux"] ** 2 + lazy_dataset["uy"] ** 2 + lazy_dataset["uz"] ** 2
).rename("speed")
lazy_mean_speed = lazy_speed.mean(dim="time")
lazy_mean_speed

<xarray.DataArray 'speed' (x: 16, y: 12, z: 8)> Size: 6kB
dask.array<mean_agg-aggregate, shape=(16, 12, 8), dtype=float32, chunksize=(16, 12, 8), chunktype=numpy.ndarray>
Coordinates:
  * x        (x) float32 64B 0.0 0.2 0.4 0.6 0.8 1.0 ... 2.0 2.2 2.4 2.6 2.8 3.0
  * y        (y) float32 48B -1.0 -0.8182 -0.6364 -0.4545 ... 0.6364 0.8182 1.0
  * z        (z) float32 32B 0.0 0.2857 0.5714 0.8571 1.143 1.429 1.714 2.0

Visualize the Dask task graph#

This graph represents the operations that will be executed when the result is finally computed. It is still a plan at this point, not a loaded array.

The dataset here is deliberately small so the graph stays readable. Real workloads produce much larger graphs with the same overall shape: read each chunk, combine the chunks, then reduce.

if which("dot"):
    display(lazy_mean_speed.data.visualize(rankdir="LR", format="svg", filename=None))
else:
    display(
        Markdown(
            "Graphviz is not available in this environment, so the task graph image cannot be rendered here. "
            "The published documentation installs Graphviz and renders this cell as an SVG task graph."
        )
    )

../_images/866bc8c380e0c645f367042293791e88593662d9fdfed0174e3c0adfd715b8bb.svg