Memory-efficient writes¶
A chunk is written whole. Every vertex that lands in one chunk — positions and per-vertex attributes together — must be in memory at the moment that chunk is encoded, because its fragment index describes rows of one packed buffer. For a dataset that does not fit in RAM, the question is therefore not “how do I stream vertices?” but “how do I write one chunk at a time, and how big is a chunk?”
Two facts shape everything below:
Dataset.add_pointsis a one-shot write, not an append. A second call into the same store re-derives the grid from its own batch and overwrites what was there. Loopingadd_pointsover slabs does not accumulate — it leaves you with the last slab. See Quickstart.Chunk size is a decision you make before the store exists. Once written, changing it means rewriting the store.
So a partitioned write uses zarr_vectors.building — the supported surface
for code that constructs stores, whose per-chunk verbs are exactly “the
per-chunk verbs an ingest worker runs in parallel”. zarr_vectors.api
(zv.create / add_points) is the right surface when the batch fits in
memory; it does not have a partitioned writer.
Step 0: size the grid before writing anything¶
Grid.plan computes the grid a store with those bounds would get, without
creating one, and capacity() turns a vertex count into bytes per cell. This
is the only moment the answer can still change a decision — a pipeline that
discovers its chunks are 750 MB after the fact has already written a store it
must throw away.
import zarr_vectors as zv
grid = zv.Grid.plan(bounds=((0., 0., 0.), (10_000., 10_000., 10_000.)),
cell_size=(500., 500., 500.))
print(grid, grid.cells)
print(grid.capacity(n_vertices=500_000_000))
# The same 500M vertices in a coarse grid:
print(zv.Grid.plan(bounds=((0., 0., 0.), (10_000., 10_000., 10_000.)),
target_cells=2).capacity(n_vertices=500_000_000))
Grid(21x21x21 cells of (500.0, 500.0, 500.0)) 9261
21x21x21 = 9261 cells, ~0.6 MB/cell -- fits
3x3x3 = 27 cells, ~222.2 MB/cell -- does not fit: ~222 MB per cell exceeds the 67 MB target; use more cells
est_bytes_per_cell counts coordinates only (n_vertices / cells × ndim × 4). Add your per-vertex attributes to it by hand: five float32 attributes
alongside three float32 coordinates makes the real figure ⅔ larger again.
Peak memory during the write is roughly one partition — which is one or
more whole cells — not one cell, so size partitions with that headroom.
See Choosing a layout for what else pulls the number around.
Strategy 1: write in spatial partitions¶
The general shape of a bounded-memory write: a coordinator allocates the store’s arrays once, then each partition is loaded, written cell by cell, and freed. Peak memory is one partition, whatever the size of the dataset.
The coordinator allocates¶
import numpy as np
import zarr_vectors as zv
from zarr_vectors.building import (
LevelMetadata, create_store, create_resolution_level,
create_vertices_array, create_attribute_array, open_write_session,
open_store, get_resolution_level, write_chunk_vertices,
write_chunk_attributes, rebuild_presence, refresh_arrays_present,
update_level_metadata, write_multiscale_metadata,
)
STORE = "large_scan.zarrvectors"
BOUNDS = ([0., 0., 0.], [1000., 1000., 1000.])
CHUNK = (250., 250., 250.)
BIN = (62.5, 62.5, 62.5)
BINS_PER_CHUNK = tuple(int(round(c / b)) for c, b in zip(CHUNK, BIN))
N_BINS = int(np.prod(BINS_PER_CHUNK))
root = create_store(STORE, bounds=BOUNDS, chunk_shape=CHUNK,
base_bin_shape=BIN, geometry_types=["point_cloud"])
level = create_resolution_level(root, 0, LevelMetadata(
level=0, vertex_count=0, arrays_present=["vertices", "vertex_attributes"]))
with open_write_session(level, bounds=BOUNDS, chunk_shape=CHUNK):
create_vertices_array(level, dtype="float32")
create_attribute_array(level, "intensity", dtype="float32")
Warning
Allocate exactly once, in the coordinator. create_vertices_array and
create_attribute_array rewrite the array’s zarr.json, and inside a write
session that drops every cell already on disk. Calling them again from each
partition looks harmless and silently reduces a four-slab write to its last
slab. Workers write cells; they never create arrays.
bounds and chunk_shape are required by open_write_session because the
chunk grid’s extent — and its origin, for data with negative coordinates — is
the shape of each vlen array, and has to be sized up front. Every partition
must pass the same two values as the coordinator did, or their grids disagree.
Each partition writes its own cells¶
# NOTE: internal import -- see the admonition below.
from zarr_vectors.spatial.chunking import assign_bins, group_bins_by_chunk
rng = np.random.default_rng(0)
written = 0
for slab in range(4): # 4 z-slabs, one chunk deep each
z_lo, z_hi = slab * 250., (slab + 1) * 250.
positions, intensity = load_slab(z_lo, z_hi) # only this slab in memory
level = get_resolution_level(open_store(STORE, mode="r+"), 0)
with open_write_session(level, bounds=BOUNDS, chunk_shape=CHUNK):
per_chunk = group_bins_by_chunk(assign_bins(positions, BIN),
BINS_PER_CHUNK)
for cell, fragments in per_chunk.items():
verts = [np.zeros((0, 3), "float32") for _ in range(N_BINS)]
attrs = [np.zeros((0,), "float32") for _ in range(N_BINS)]
for fragment_index, rows in fragments.items():
verts[fragment_index] = positions[rows]
attrs[fragment_index] = intensity[rows]
write_chunk_vertices(level, cell, verts, dtype="float32",
record_presence=False)
write_chunk_attributes(level, "intensity", cell, attrs,
dtype="float32", record_presence=False)
written += sum(len(v) for v in verts)
del positions, intensity
record_presence=False is not an optimisation. nonempty_chunks is a single
attribute shared by every cell of an array, so a partition that stamps it
races every other partition. Partitions skip it; the coordinator rebuilds it
once, below.
One fragment per bin, not one per cell¶
The verts list above has one entry per bin, most of them empty. That is
deliberate. A fragment is the unit a coarsener picks representatives from, so
collapsing a cell into a single fragment — write_chunk_vertices(level, cell, [positions_in_cell]), which is the tempting one-liner — throws away the bin
structure that build_pyramid later needs. The same 16 000 vertices in the
same 4×4×4 grid:
fragments per cell |
level 1 |
level 2 |
|---|---|---|
64 (one per bin) |
512 |
64 |
1 (whole cell) |
8 |
1 |
Note
assign_bins and group_bins_by_chunk are imported from
zarr_vectors.spatial.chunking, which is internal — the import above
reaches past the contract, and may break between releases.
zarr_vectors.building exports assign_chunks (vertices → cells) but no
bin-level equivalent. That is a gap to report rather than a reason to treat
spatial as public. If you do not need per-bin fragments, use
building.assign_chunks and pass one fragment per cell, accepting the
coarsening cost above.
The coordinator finalises¶
root = open_store(STORE, mode="r+")
level = get_resolution_level(root, 0)
rebuild_presence(level) # the nonempty_chunks nobody stamped
print(refresh_arrays_present(level)) # what is actually on disk
update_level_metadata(level, vertex_count=written)
write_multiscale_metadata(root)
ds = zv.open(STORE)
print(ds.level(0).vertex_count, ds.level(0).scale, ds.level(0).resolution)
print(ds.read())
['vertex_attributes', 'vertex_fragments', 'vertices']
16000 (250.0, 250.0, 250.0) (62.5, 62.5, 62.5)
ReadResult(kind='point_cloud', vertices=16000, attributes=['intensity'])
refresh_arrays_present walks the level rather than trusting the list the
writers declared, which is why vertex_fragments appears even though the
LevelMetadata above never mentioned it. Both verbs are coordinator-only:
never run either while partitions are still writing.
Coarser levels come afterwards, from the finished store:
report = zv.open(STORE, mode="r+").build_pyramid(
factors=[(2., 1.), (2., 1.)], chunk_scale_factors=[2, 2])
print(report["levels_created"], [s["vertex_count"] for s in report["level_specs"]])
2 [512, 64]
Strategy 2: bound peak memory by partition size¶
There is no streaming-writer class. The bound on peak memory is the partition, and you set it by choosing how much of the volume each iteration loads:
peak ≈ cells_per_partition × vertices_per_cell × bytes_per_vertex
At 8 cells per partition, 50 000 vertices per cell and 32 bytes per vertex (three float32 coordinates plus five float32 attributes) that is about 13 MB of payload — call it double, for the encode buffer the write allocates alongside it. Partitions do not have to be slabs; any set of whole cells works, as long as no two partitions share one. Chunk-aligned partitions are the cheapest, because no cell is ever written twice.
Smaller partitions cost round-trips rather than correctness: each one reopens the store and opens its own write session. Against a high-latency object store that matters, and a partition covering a few hundred cells amortises it.
Strategy 3: generator-based ingest¶
For formats read by external libraries (LAS, TRX, SWC, …), the companion
package zarr-vectors-tools provides streaming converters that yield
vertex batches without loading the whole file. They drive the same
zarr_vectors.building per-chunk verbs shown in Strategy 1, so the memory
argument above applies to them unchanged.
Strategy 4: rechunk after the initial write¶
If your data is already in a format that can be read block by block, the
smallest change to an existing pipeline is to write it once with a coarse
chunk shape and re-cut the grid afterwards. rechunk with
RechunkSpec(by="spatial") changes the spatial chunk shape and keeps the
store’s topology, levels and object ids:
import numpy as np
import zarr_vectors as zv
from zarr_vectors.building import RechunkSpec, rechunk
# Initial write: one chunk for the whole volume.
temp = zv.create("scan_temp.zarrvectors", schema=zv.Schema(
bounds=((0., 0., 0.), (1000., 1000., 1000.)),
kind="point_cloud",
layout=zv.Layout(cells=1),
))
temp.add_points(positions, object_ids=object_ids)
print(zv.open("scan_temp.zarrvectors").level(0).grid)
summary = rechunk(
"scan_temp.zarrvectors",
RechunkSpec(by="spatial", spatial_chunk_shape=(250., 250., 250.)),
output="scan.zarrvectors",
)
print(summary)
print(zv.open("scan.zarrvectors").level(0).grid)
Grid(2x2x2 cells of (1000.0, 1000.0, 1000.0))
{'objects_rechunked': 200, 'bins_created': 1, 'total_vertices': 20000, 'rechunk_dims': ['spatial', 'x', 'y', 'z'], 'output_path': 'scan.zarrvectors'}
Grid(5x5x5 cells of (250.0, 250.0, 250.0))
This is the least memory-efficient route, and it is worth being explicit about why: the initial single-chunk write holds the entire dataset in memory at once, which is the problem the rest of this page exists to avoid. Reach for it when the dataset already fits and the chunk shape is merely wrong, not when it does not fit at all.
Two things to know about the result:
Rechunked stores gain a prefix dimension. Chunk keys become
(bin, z, y, x)—c/0/2/0/0rather thanc/2/0/0— becauserechunk’s output layout is shared with the non-spatial forms below.rechunk_dimsin the summary names the axes.output=Nonerechunks in place, by writing a temporary store and replacing the source. Leaveoutputset while you still want the original.
RechunkSpec also rechunks along a non-spatial dimension — by="group",
by="object_id", or by="attribute:<name>" — so that all objects sharing a
value land in one chunk. building.rechunk_by_attribute(store, "cell_type") is the
categorical shorthand: every distinct value gets its own bin, however many
there are.
Warning
Rechunking rewrites where every byte lives, so it is a coordinator operation. Never run it while anything else is writing to the store.
Monitoring memory usage¶
tracemalloc measures the write path directly:
import tracemalloc
import numpy as np
import zarr_vectors as zv
tracemalloc.start()
probe = zv.create("probe.zarrvectors", schema=zv.Schema(
bounds=((0., 0., 0.), (1000., 1000., 1000.)),
kind="point_cloud", layout=zv.Layout(cells=4)))
probe.add_points(np.random.default_rng(1)
.uniform(0., 1000., size=(100_000, 3)).astype("float32"))
current, peak = tracemalloc.get_traced_memory()
print(f"Peak memory: {peak / 1e6:.1f} MB")
tracemalloc.stop()
Peak memory: 8.0 MB
tracemalloc sees Python allocations only. NumPy’s own buffers are counted,
but a compressor’s native scratch space is not, so treat the figure as a
lower bound and leave headroom.
Next steps¶
HPC pipelines — the same partitioned write, run across SLURM array tasks or MPI ranks instead of a
forloop.Choosing a layout — how to pick the numbers
Grid.plantakes.Building stores — the full reference for the surface Strategy 1 uses.