Skip to content

Creates Zarr to b2nd conversion utility #738

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
VibhuAg wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Creates Zarr to b2nd conversion utility #738

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bfc76f1
Add shared helpers module for CLI array converters
VibhuAg Oct 13, 2026
c02e177
Add zarr-to-b2nd converter using shared conversion helpers
VibhuAg Oct 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ fsspec = ["fsspec[http]"]
[project.scripts]
parquet-to-blosc2 = "blosc2.cli.parquet_to_blosc2:main"
b2nd-to-zarr = "blosc2.cli.b2nd_to_zarr:main"
zarr-to-b2nd = "blosc2.cli.zarr_to_b2nd:main"
blosc2-to-zarr = "blosc2.cli.b2nd_to_zarr:main"
b2view = "blosc2.b2view.cli:main"

Expand Down
299 changes: 17 additions & 282 deletions src/blosc2/cli/b2nd_to_zarr.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,126 +11,25 @@

import argparse
import contextlib
import itertools
import math
import os
import shutil
import sys
import time
from pathlib import Path
from typing import Any

import numpy as np

import blosc2

try:
from rich.console import Console
from rich.progress import (
BarColumn,
Progress,
TaskProgressColumn,
TextColumn,
TimeElapsedColumn,
TimeRemainingColumn,
TransferSpeedColumn,
)
from rich.table import Table

HAVE_RICH = True
except ImportError:
HAVE_RICH = False

DEFAULT_BUFFER_SIZE_MB = 128


def require_zarr() -> Any:
"""Ensure the zarr package is installed."""
try:
import zarr

return zarr
except ImportError as exc:
raise ImportError("b2nd-to-zarr requires zarr; install it with: pip install 'blosc2[zarr]'") from exc


def format_bytes(size: float) -> str:
"""Format byte count into human-readable string."""
units = ["B", "KB", "MB", "GB", "TB"]
unit_idx = 0
while size >= 1024.0 and unit_idx < len(units) - 1:
size /= 1024.0
unit_idx += 1
return f"{size:,.2f} {units[unit_idx]}" if unit_idx > 0 else f"{int(size)} B"


def get_path_size(path: Path) -> int:
"""Get total byte size of a file or directory tree."""
if not path.exists():
return 0
if path.is_file():
return path.stat().st_size
total = 0
for root, _, files in os.walk(path):
for f in files:
fp = os.path.join(root, f)
with contextlib.suppress(OSError):
total += os.path.getsize(fp)
return total


def parse_shape(shape_str: str) -> tuple[int, ...]:
"""Parse comma- or whitespace-separated integer dimensions."""
cleaned = shape_str.replace("x", ",").replace(" ", ",").replace("(", "").replace(")", "")
dims = [int(p.strip()) for p in cleaned.split(",") if p.strip()]
if not dims or any(d <= 0 for d in dims):
raise ValueError(f"Invalid shape string '{shape_str}', dimensions must be positive integers.")
return tuple(dims)


def compute_copy_shape(
shape: tuple[int, ...],
target_chunks: tuple[int, ...],
itemsize: int,
max_buffer_bytes: int = DEFAULT_BUFFER_SIZE_MB * 1024 * 1024,
) -> tuple[int, ...]:
"""Compute batch copy shape aligned with chunk boundaries to optimize memory and throughput."""
copy_shape = list(target_chunks)
ndim = len(shape)

# First attempt: expand inner dimensions to full shape (from right to left)
# to maximize contiguous memory slicing in C-order arrays.
for d in reversed(range(ndim)):
test_shape = list(copy_shape)
test_shape[d] = shape[d]
test_bytes = math.prod(test_shape) * itemsize
if test_bytes <= max_buffer_bytes:
copy_shape[d] = shape[d]
else:
current_bytes = math.prod(copy_shape) * itemsize
if current_bytes < max_buffer_bytes:
bytes_per_chunk = current_bytes // math.ceil(copy_shape[d] / target_chunks[d])
max_chunks = max_buffer_bytes // bytes_per_chunk
num_chunks = math.ceil(shape[d] / target_chunks[d])
mult = min(max(1, max_chunks), num_chunks)
copy_shape[d] = min(shape[d], mult * target_chunks[d])
break

# Second attempt: if memory headroom remains, expand leading dimensions
for d in range(ndim):
current_bytes = math.prod(copy_shape) * itemsize
if current_bytes >= max_buffer_bytes:
break
num_chunks = math.ceil(shape[d] / target_chunks[d])
curr_chunks = math.ceil(copy_shape[d] / target_chunks[d])
if curr_chunks < num_chunks:
bytes_per_chunk = current_bytes // curr_chunks
max_chunks = max_buffer_bytes // bytes_per_chunk
mult = min(max(1, max_chunks), num_chunks)
if mult > curr_chunks:
copy_shape[d] = min(shape[d], mult * target_chunks[d])

return tuple(copy_shape)
from blosc2.cli.conversion_utils import (
DEFAULT_BUFFER_SIZE_MB,
compute_batches,
copy_batches,
format_bytes,
get_path_size,
parse_shape,
print_summary,
require_zarr,
verify_arrays,
)


def resolve_shuffle(arr: blosc2.NDArray, shuffle_mode: str) -> str:
Expand Down Expand Up @@ -251,113 +150,6 @@ def copy_metadata(arr: blosc2.NDArray, z_arr: Any) -> None:
z_arr.attrs[k] = v


def copy_batches(
arr: blosc2.NDArray,
z_arr: Any,
batch_slices: list[tuple[slice, ...]],
total_bytes: int,
quiet: bool = False,
verbose: bool = False,
) -> float:
"""Copy data slices from Blosc2 NDArray to Zarr array with progress tracking."""
total_batches = len(batch_slices)
if total_batches == 0:
return 0.0
t_start = time.perf_counter()

if HAVE_RICH and not quiet:
console = Console()
with Progress(
TextColumn("[bold blue]{task.description}"),
BarColumn(),
TaskProgressColumn(),
TransferSpeedColumn(),
TimeElapsedColumn(),
TimeRemainingColumn(),
console=console,
) as progress:
task = progress.add_task("Converting b2nd -> zarr...", total=total_bytes)
for s in batch_slices:
data = arr[s]
z_arr[s] = data
batch_size = data.nbytes
progress.update(task, advance=batch_size)
del data
else:
for idx, s in enumerate(batch_slices):
data = arr[s]
z_arr[s] = data
del data
if verbose and not quiet:
pct = (idx + 1) / total_batches * 100
print(f"Batch {idx + 1}/{total_batches} ({pct:.1f}%) written")

return time.perf_counter() - t_start


def _fast_sample_verify(
b2_arr: blosc2.NDArray, z_arr: Any, slices_list: list[tuple[slice, ...]], quiet: bool
) -> None:
corner_indices = [
tuple(0 for _ in b2_arr.shape),
tuple(-1 for _ in b2_arr.shape),
tuple(s // 2 for s in b2_arr.shape),
]
for idx in corner_indices:
v_b2 = b2_arr[idx]
v_z = z_arr[idx]
if not np.array_equal(v_b2, v_z, equal_nan=True):
raise ValueError(f"Verification failed at index {idx}: Blosc2={v_b2} != Zarr={v_z}")

if slices_list:
for s in (slices_list[0], slices_list[-1]):
b2_chunk = b2_arr[s]
z_chunk = z_arr[s]
if not np.array_equal(b2_chunk, z_chunk, equal_nan=True):
raise ValueError(f"Verification failed for slice {s}")
if not quiet:
print("✓ Fast sample verification passed!")


def _full_verify(
b2_arr: blosc2.NDArray, z_arr: Any, slices_list: list[tuple[slice, ...]], quiet: bool
) -> None:
if not quiet:
print("Performing full verification across all slices...")
for idx, s in enumerate(slices_list):
b2_data = b2_arr[s]
z_data = z_arr[s]
if not np.array_equal(b2_data, z_data, equal_nan=True):
raise ValueError(f"Full verification failed at slice {s} (batch {idx})")
if not quiet:
print("✓ Full element verification passed!")


def verify_arrays(
b2_arr: blosc2.NDArray,
z_arr: Any,
slices_list: list[tuple[slice, ...]],
full: bool = False,
quiet: bool = False,
) -> bool:
"""Verify that data in Zarr array matches source Blosc2 NDArray."""
if b2_arr.shape != z_arr.shape:
raise ValueError(f"Shape mismatch: Blosc2 {b2_arr.shape} vs Zarr {z_arr.shape}")
if b2_arr.dtype != z_arr.dtype:
raise ValueError(f"Dtype mismatch: Blosc2 {b2_arr.dtype} vs Zarr {z_arr.dtype}")

if math.prod(b2_arr.shape) == 0:
if not quiet:
print("✓ Empty array verification passed!")
return True

if not full:
_fast_sample_verify(b2_arr, z_arr, slices_list, quiet)
else:
_full_verify(b2_arr, z_arr, slices_list, quiet)
return True


def _determine_target_chunks_and_shards(
arr: blosc2.NDArray,
chunks: tuple[int, ...] | None,
Expand All @@ -381,26 +173,6 @@ def _determine_target_chunks_and_shards(
return target_chunks, target_shards


def _compute_batches(
shape: tuple[int, ...],
target_chunks: tuple[int, ...],
target_shards: tuple[int, ...] | None,
itemsize: int,
buffer_size_mb: int,
) -> tuple[tuple[int, ...], list[tuple[slice, ...]]]:
if math.prod(shape) == 0:
return shape, []
unit_shape = target_shards if target_shards is not None else target_chunks
copy_shape = compute_copy_shape(
shape, unit_shape, itemsize, max_buffer_bytes=buffer_size_mb * 1024 * 1024
)
slices_per_dim = [
[slice(i, min(i + step, s)) for i in range(0, s, step)]
for s, step in zip(shape, copy_shape, strict=True)
]
return copy_shape, list(itertools.product(*slices_per_dim))


def b2nd_to_zarr(
input_path: str | Path,
output_path: str | Path | None = None,
Expand All @@ -421,7 +193,7 @@ def b2nd_to_zarr(
verbose: bool = False,
) -> dict[str, Any]:
"""Convert a Blosc2 .b2nd file containing an NDArray to a Zarr array."""
zarr = require_zarr()
zarr = require_zarr("b2nd-to-zarr")

src_path = Path(input_path).expanduser().resolve()
if not src_path.exists():
Expand Down Expand Up @@ -465,9 +237,7 @@ def b2nd_to_zarr(
blocksize=blocksize,
zarr_format=zarr_format,
)
copy_shape, batch_slices = _compute_batches(
shape, target_chunks, target_shards, itemsize, buffer_size_mb
)
copy_shape, batch_slices = compute_batches(shape, target_chunks, target_shards, itemsize, buffer_size_mb)

if verbose and not quiet:
print(f"Source: {src_path} ({format_bytes(src_file_size)})")
Expand Down Expand Up @@ -504,7 +274,9 @@ def b2nd_to_zarr(
)

copy_metadata(arr, z_arr)
copy_duration = copy_batches(arr, z_arr, batch_slices, uncompressed_bytes, quiet=quiet, verbose=verbose)
copy_duration = copy_batches(
arr, z_arr, batch_slices, uncompressed_bytes, "b2nd -> zarr", quiet=quiet, verbose=verbose
)
if zip_store is not None:
zip_store.close()

Expand Down Expand Up @@ -706,44 +478,7 @@ def main(argv: list[str] | None = None) -> int:
return 1

if not args.quiet:
if HAVE_RICH:
console = Console()
table = Table(title="[bold green]Conversion Summary[/bold green]", box=None)
table.add_column("Property", style="bold cyan")
table.add_column("Value")
table.add_row("Source", f"{summary['src_path']} ({format_bytes(summary['src_size_bytes'])})")
table.add_row(
"Destination", f"{summary['dst_path']} ({format_bytes(summary['dst_size_bytes'])})"
)
table.add_row("Shape / Dtype", f"{summary['shape']} / {summary['dtype']}")
table.add_row("Chunks", str(summary["chunks"]))
if summary["shards"] is not None:
table.add_row("Shards", str(summary["shards"]))
table.add_row("Uncompressed Size", format_bytes(summary["uncompressed_bytes"]))
table.add_row("Source Compression", f"{summary['src_cratio']:.2f}x")
table.add_row("Zarr Compression", f"{summary['dst_cratio']:.2f}x")
table.add_row(
"Conversion Speed",
f"{summary['throughput_mb_s']:.1f} MB/s ({summary['copy_duration_sec']:.2f} s)",
)
table.add_row("Total Time", f"{summary['total_duration_sec']:.2f} s")
console.print()
console.print(table)
else:
print("\n--- Conversion Summary ---")
print(f"Source: {summary['src_path']} ({format_bytes(summary['src_size_bytes'])})")
print(f"Destination: {summary['dst_path']} ({format_bytes(summary['dst_size_bytes'])})")
print(f"Shape / Dtype: {summary['shape']} / {summary['dtype']}")
print(f"Chunks: {summary['chunks']}")
if summary["shards"] is not None:
print(f"Shards: {summary['shards']}")
print(f"Uncompressed: {format_bytes(summary['uncompressed_bytes'])}")
print(f"Source CRatio: {summary['src_cratio']:.2f}x")
print(f"Zarr CRatio: {summary['dst_cratio']:.2f}x")
print(
f"Conversion Speed: {summary['throughput_mb_s']:.1f} MB/s ({summary['copy_duration_sec']:.2f} s)"
)
print(f"Total Time: {summary['total_duration_sec']:.2f} s")
print_summary(summary, dst_label="Zarr")

return 0

Expand Down
Loading
Loading