Introduction

Managing petabyte-scale multi-dimensional arrays is a recurring bottleneck in modern machine learning and scientific computing. When datasets exceed available RAM, traditional loading methods fail, and I/O latency becomes the primary constraint on GPU utilization. TensorStore, an open-source C++ and Python library from Google, solves this by providing a high-performance engine for reading and writing large n-dimensional arrays across diverse storage backends. By decoupling the logical view of the data from its physical storage, TensorStore enables researchers to manipulate massive datasets without needing to load the entire array into memory.

What Is TensorStore?

TensorStore is a C++ and Python software library designed for the storage and manipulation of large multi-dimensional arrays. It provides a uniform API for reading and writing multiple array formats, including Zarr and N5, while natively supporting multiple storage systems such as local and network filesystems, Google Cloud Storage (GCS), Amazon S3-compatible object stores, and HTTP servers. Licensed under the Apache License 2.0, it is maintained by Google and designed to handle data at a scale where traditional NumPy arrays cannot operate.

The core abstraction of the library is the TensorStore object, which acts as an asynchronous view of a multi-dimensional array. Every store is backed by a driver that that connects the high-level interface to the underlying storage mechanism. This architecture allows users to open or create stores using a JSON specification (spec), which is analogous to a connection string or file path, providing a highly flexible way to define how data is accessed and stored.

Why TensorStore Matters

In fields like neuroscience, weather modeling, and large-scale AI training, datasets often reach the petabyte scale. For example, Google researchers used TensorStore to handle the h01 dataset, which contains approximately 1.4 petabytes of imaging data of human brain tissue. Traditional data loaders often struggle with these scales because they lack the necessary concurrency and asynchronous I/O capabilities to saturate network bandwidth.

TensorStore matters because it eliminates the “memory wall” by using lazy loading and virtual views. Instead of loading a full dataset, it only fetches the specific slices or chunks required for a computation. This allows a developer to work with a 56-trillion-voxel image of a fly’s brain on a standard workstation without crashing the system. Furthermore, its C++ core ensures that encoding, decoding, and I/O operations are performed with minimal overhead, making it an ideal choice for feeding data into high-performance accelerators like GPUs.

As AI models grow in complexity, the need for efficient checkpointing and parameter management becomes critical. TensorStore has been utilized in the creation of large-scale models like PaLM, addressing the challenge of managing model parameters during distributed training. This traction among top-tier AI research teams signals its importance as a foundational piece of infrastructure for the next generation of scalable ML.

Key Features

• Uniform API for Multiple Formats: TensorStore provides a single interface to read and write data in Zarr, N5, and Neuroglancer precomputed formats, reducing the need for format-specific libraries.
• Native Multi-Backend Support: It supports a wide array of storage drivers, including local filesystems, Google Cloud Storage, Amazon S3, and in-memory storage, allowing seamless transitions between local development and cloud deployment.
• Asynchronous I/O API: The library offers an asynchronous API that enables high-throughput access even to high-latency remote storage, preventing the CPU from idling while waiting for data.
• Composable Indexing and Virtual Views: Users can create virtual views of data through indexing, downsampling, and data type conversion. These operations are lazy and fully composable, meaning they don’t trigger data movement until the data is actually requested.
• ACID Transactions: TensorStore supports read/writeback caching and transactions with strong atomicity, isolation, consistency, and durability (ACID) guarantees, ensuring data integrity during parallel writes.
• Optimistic Concurrency: It enables safe and efficient access from multiple processes and machines via optimistic concurrency, which is critical for distributed training and large-scale data processing.
• High-Performance C++ Core: The core implementation in C++ automatically leverages multiple CPU cores for encoding and decoding, saturating network bandwidth to maximize I/O throughput.
• Dimension Labels: Support for dimension labels allows users to index data using meaningful names rather than just integer coordinates, improving code readability and reducing errors in high-dimensional data.

How TensorStore Compares

When comparing TensorStore to Zarr-Python, the primary differentiator is performance and concurrency. While Zarr is an excellent format and widely used in the scientific community, the TensorStore library is often faster in read/write operations because it is implemented in C++ and handles I/O asynchronously. Benchmarks indicate that TensorStore is consistently faster when reading and writing Zarr format data, as it can better utilize multiple CPU cores and saturate network bandwidth without needing external libraries like Dask for parallelism.

Compared to HDF5, TensorStore is significantly more cloud-native. HDF5 was designed for local filesystems and often struggles with object storage (S3/GCS) due to its file-locking mechanisms and metadata overhead. TensorStore, however, is built from the ground up for the cloud, treating object stores as first-class citizens. This makes it the superior choice for teams moving their data pipelines from local servers to cloud-based GPU clusters.

The tradeoff is that Zarr-Python has a larger, more established ecosystem of higher-level libraries (like Xarray) that integrate with it. While TensorStore provides a Python API, it is a lower-level storage engine. Therefore, the best approach for many users is to use TensorStore as the backend for Zarr data, leveraging the Zarr format for compatibility and TensorStore’s engine for raw performance.

Getting Started: Installation

TensorStore provides both a Python API and a C++ API. For most users, the Python API is the fastest way to get started.

Python API via PyPI

TensorStore requires Python 3.11 or later. It is recommended to use a virtual environment to avoid dependency conflicts.

python3 -m pip install tensorstore

Python API from Source

If you need to modify the C++ core or use a feature not yet in the PyPI release, you can install from the git repository. This requires the Bazel build system to be installed on your machine.

git clone https://github.com/google/tensorstore
cd tensorstore
python3 -m pip install .

C++ API Installation

The C++ API is intended for high-performance applications that require the absolute minimum overhead. Installation involves using Bazel to build the library from source.

For detailed build instructions, refer to the official Building and Installing guide.

How to Use TensorStore

The basic workflow in TensorStore involves defining a specification (spec) for your array, opening the store, and then performing indexing operations to retrieve data.

To open an existing Zarr array on a local filesystem, you would use ts.open. Because TensorStore is asynchronous, the .result() method is used to block and wait for the future to resolve into a actual store object.

Once the store is opened, you can use NumPy-like indexing to slice the data. Because the indexing is lazy, no data is actually read from disk until you call a method that triggers a view resolution, such as .read(). This allows you to chain multiple indexing operations together without any immediate I/O overhead.

Code Examples

Below are examples of how to use the TensorStore Python API to manage large arrays.

Opening and Reading a Store

import tensorstore as ts

# Open a Zarr store from a local path
# The .result() call blocks until the store is opened successfully
store = ts.open({
    'driver': 'zarr3',
    'kvstore': 'file://data/my_array.zarr',
},).result()

# Read a 100x100 slice of the data
# This is a lazy operation; it returns a view
slice_view = store[0:100, 0:100]

# Actually fetch the data into a NumPy array
data = slice_view.read().result()
print(data)

This example demonstrates the basic read flow: open the store using a spec, create a lazy slice view, and then call .read().result() to perform the asynchronous I/O operation.

Creating a New Store

import tensorstore as ts

# Define the specification for a new array
spec = {
    'driver': 'zarr3',
    'kvstore': 'file://my_new_array.zarr',
    'metadata': {
        'dimensions': {
            'x': 1000, 
            'y': 1000, 
            'z': 1000
        },
        'chunks': {'x': 64, 'y': 64, 'z': 64},
        'codecs': [{'id': 'blosc',
                  'clevel': 5,
                  'shuffle': 1}]
    }
}

# Create the store based on the spec
store = ts.create(spec).result()
print(f"Store created at {spec['kvstore']}")

In this example, we define a JSON spec that describes the array’s dimensions, chunking and compression codecs. This allows us to create a persistent storage structure on disk that can be opened by other tools in the Zarr ecosystem.

Real-World Use Cases

TensorStore is designed for scenarios where data size exceeds the capacity of a single machine’s memory, and where high-throughput I/O is required.

• Neuroscience and Brain Mapping: Researchers use TensorStore to manage petabyte-scale 3D electron microscopy data. By using virtual views and downsampling, they can visualize and analyze a 1.4 petabyte dataset of human brain tissue without needing to load the entire volume into memory.
• Large-Scale AI Model Checkpointing: In the training of models like PaLM, TensorStore is used to manage model parameters (checkpoints) during distributed training. This ensures that the hundreds of billions of parameters are saved and saved efficiently across a distributed filesystem, avoiding bottlenecks in the training loop.
• Distributed Data Processing: When using parallel computing frameworks like Apache Beam or Dask, TensorStore can be used as the high-performance I/O layer that feeds data into these pipelines, ensuring that GPUs are not left idling while waiting for data from cloud storage.
• Climate and Weather Modeling: For atmospheric measurements over a spatial grid, TensorStore allows scientists to slice and dice multi-dimensional arrays of weather data across time and space, performing parallel analysis on numerous machines working in parallel.

Contributing to TensorStore

TensorStore is an open-source project maintained by Google. Contributions are welcome through the standard GitHub flow. To contribute, you should first sign the Google Contributor License Agreement (CLA) and sign the Google CLA. After that, you can report bugs via GitHub Issues and submit pull requests for new features or features improvements.

The project follows a professional code of conduct to ensure a positive community environment. New contributors are encouraged to look for issues labeled “good first issue” to get started with the project’s C++ or Python bindings.

Community and Support

The primary hub for TensorStore is its official GitHub repository, where developers can track issues, discuss features and report bugs. Documentation is hosted at google.github.io/tensorstore, which includes detailed tutorials and a full Python API reference.

While there is no dedicated Discord or Slack channel, the community is largely centered around the scientific computing and ML engineering community, particularly those using the Zarr and N5 formats. Because it is a Google-maintained project, it is maintained with a high level of activity and is actively updated to support new storage drivers and array formats.

Conclusion

TensorStore is the right choice for developers and researchers who are working with multi-dimensional arrays that are too large to fit in memory. It is particularly powerful when moving from local development to cloud-based GPU clusters, as its asynchronous I/O and native cloud storage support make it an efficient bridge between object storage and high-performance compute.

If you are already using Zarr or N5, TensorStore is a highly recommended upgrade for your I/O layer. However, if your datasets are small enough to fit in a single NumPy array, the added complexity of asynchronous futures and JSON specs may be not necessary. Use TensorStore when the scale of the ldata is the primary bottleneck in your ML pipeline.

Star the repo, try the quickstart, and join the community of researchers pushing the boundaries of scalable array storage.

Resources

• TensorStore GitHub Repository
• Official TensorStore Documentation
• Google AI Blog: TensorStore Announcement