tensorblob

A lightweight, dynamic-sized, memory-mapped tensor storage with file-like APIs, while also supporting integer indexing and slicing, built with MemoryMappedTensor from tensordict.

Features

🔗 Memory-mapped storage: Efficient storage of large collections of same-shaped tensors
💾 File-like APIs: Read, write, and seek like a file, while also supporting integer indexing and slicing
⚡ Dynamic-sized: No need to specify the total number of tensors upfront
🔄 Extend and truncate: Extend the blob with another blob or truncate the blob to a specific position
🚀 LRU cache: Automatic management of memory-mapped blocks for scalability with large blobs
🧩 Multi-field databases: TensorDB manages several row-aligned blobs for heterogeneous data, e.g., multivariate time series or event streams

Installation

From PyPI:

pip install tensorblob

If you are interested in the experimental (i.e., unstable and undertested) version, you can install it from GitHub:

pip install git+https://github.com/Guest400123064/tensorblob.git

Core Use Cases

Quick Start

The example below shows how to create a new storage for a collection of randomly generated fake embeddings, and how to access them by index. Since the storage is memory-mapped, no need to read all tensors into memory; just access them by index.

import torch
from tensorblob import TensorBlob

# Create a new storage for a collection of randomly generated fake embeddings;
# need to specify the data type and shape of each tensor for creation
with TensorBlob.open("embeddings.blob", "w", dtype="float32", shape=768) as blob:
    blob.write(torch.randn(100_000, 768))
    print(f"Wrote {len(blob)} embeddings")

# No need to specify the configurations again after creation
with TensorBlob.open("embeddings.blob", "r") as blob:
    e1 = blob[42]
    e2 = blob[-1:16384:-12345]
    print(f"Similarity: {torch.cosine_similarity(e1, e2)}")

Processing Large Datasets

Store and preprocess datasets larger than RAM using memory mapping can be useful to accelerate the training process by reducing the time spent on data loading and transformation.

with TensorBlob.open("data/images.blob", "w", dtype="float32", shape=(3, 224, 224)) as blob:
    for image_batch in data_loader:
        blob.write(preprocess(image_batch))

with TensorBlob.open("data/images.blob", "r") as blob:
    for image in blob:
        result = model(image)

Incremental Data Collection

Append new data to existing blobs can be useful with streaming data collection.

with TensorBlob.open("positions.blob", "w", dtype="float32", shape=3) as blob:
    blob.write(initial_position)

# Later: append more data by opening the blob in append mode
with TensorBlob.open("positions.blob", "a") as blob:
    for pos in trajectory_queue.get():
        blob.write(pos)
    print(f"Total trajectory recorded: {len(blob)}")

Random Access and Updates with File-Like APIs

Read and modify specific tensors starting from a specific position.

import io

with TensorBlob.open("data/features.blob", "r+") as blob:
    blob.seek(1000)
    print(f"Current position: {blob.tell()}")

    batch = blob.read(size=100)
    print(f"Read {batch.shape} tensors")

    # Update specific positions, whence is also supported
    blob.seek(-500, whence=io.SEEK_END)
    blob.write(updated_features)

    # Append new data
    blob.seek(len(blob))
    blob.write(additional_features)

Extend and Truncate

Extend the blob with another blob or truncate the blob to a specific position. Extension could be useful if we want to merge two blobs into one, e.g., results from two different processes. Note that extension operation does not delete the original data.

with TensorBlob.open("data/features.blob", "a") as blob:
    blob.extend(other_blob)

# Extension without maintaining the order is faster
with TensorBlob.open("data/features.blob", "r+") as blob:
    blob.extend(other_blob, maintain_order=False)

with TensorBlob.open("data/features.blob", "r+") as blob:
    blob.truncate(1000)
    print(f"Truncated to {len(blob)} tensors")

Heterogeneous Data with TensorDB

For multi-modal or multi-field data (e.g., multivariate time series, event streams), TensorDB manages several TensorBlobs under the hood — one per field — with row orders always aligned. Each field has its own dtype and shape, and each field's storage gets its own independent LRU cache and block files.

from tensorblob import TensorDB

# Create a database with a fixed schema mapping field names to (dtype, shape)
with TensorDB.open("events.db", "w",
                   schema={"price": ("float32", 1),
                           "embed": ("float16", 768)}) as db:
    # Rows are dense: every write must supply every field with the same row count
    db.write({"price": torch.randn(100_000, 1),
              "embed": torch.randn(100_000, 768).half()})
    print(f"Wrote {len(db)} rows")

# No need to specify the schema again after creation
with TensorDB.open("events.db", "r") as db:
    row = db[42]          # {"price": tensor of shape (1,), "embed": (768,)}
    batch = db[10:100]    # {"price": (90, 1), "embed": (90, 768)}
    print(f"Fields: {list(batch)}, price range: {batch['price'].min()}..{batch['price'].max()}")

TensorDB supports the same file-like APIs as TensorBlob, applied row-wise across all fields:

with TensorDB.open("events.db", "r+") as db:
    db.seek(1000)
    batch = db.read(size=100)                 # dict of (100, ...) tensors

    db.seek(-500, whence=io.SEEK_END)
    db.write({"price": new_prices, "embed": new_embeds})  # overwrite in place

    db.truncate(10_000)                       # truncate all fields at once
    db.extend(other_db, maintain_order=False) # merge another db with the same schema

# Cleanup removes the whole database directory
TensorDB.unlink("events.db")

Consistency guarantee: a write commits the row count only after all fields are written. If a crash interrupts a write mid-way, the next open reports the last committed (fully written) row count, and writable opens automatically truncate the stray partial rows, so row alignment is always preserved.

Performance and Scalability

Memory Management

TensorBlob uses an LRU (Least Recently Used) cache to manage memory-mapped blocks efficiently. This allows you to work with blobs containing millions of tensors without loading everything into memory.

Default behavior:

Automatically caches up to ~4,000 blocks (1/16 of system's VMA limit)
Blocks loaded on-demand when accessed
Least recently used blocks automatically evicted when cache is full

For large-scale workloads:

# Increase cache for better random access performance
with TensorBlob.open("large.blob", "r", max_cached_blocks=10_000) as blob:
    for idx in random_indices:
        tensor = blob[idx]  # Cached blocks reused efficiently

# Decrease cache for memory-constrained environments
with TensorBlob.open("data.blob", "r", max_cached_blocks=100) as blob:
    for tensor in blob:  # Sequential access works fine with small cache
        process(tensor)

Performance tips:

Sequential access patterns work well with any cache size
Random access benefits from larger cache sizes — but do not undersize the cache for random workloads: when the random working set exceeds max_cached_blocks, every access evicts and remaps a block, degrading lookup latency several-fold (~30 µs → ~140 µs in our benchmarks). If in doubt, increase the cache or the block size
Each cached block consumes ~200 bytes of kernel memory (VMA overhead)
System limit: typically ~65,000 memory-mapped regions per process
To avoid frequent cache evictions, one can also increase the block size to reduce the total number of blocks
For random batches, use vectorized batch indexing blob[idxs] (list, tuple, or 1-D torch.Tensor of row indices) instead of gathering row by row — it is ~7x faster; contiguous slices are faster still, so pre-sorting indices helps when order is flexible

Benchmarks

Headline numbers from the synthetic benchmark suite (500k × 768-dim float32 rows, 12-core x86_64, 16 GiB RAM; see benchmarks/ for full analysis and reproducible scripts):

Measurement	Result
Sequential write throughput	~165 MB/s (~54k rows/s)
Sequential read throughput	~2.2 GB/s (~730k rows/s), vs ~7.2 GB/s in-memory upper bound
Random single-row lookup	~31 µs median (in-memory: ~5 µs)
Preprocessing offload (5 epochs)	~3.5x faster than re-preprocessing; breaks even after ~1.2 epochs
Memory footprint	bounded by `max_cached_blocks`; +16 VMAs / +1.3 MiB RSS at cache size 16
TensorDB column projection	reading one cheap field only is ~6x cheaper than full-row reads

Contributing

Contributions welcome! Please submit a Pull Request.

License

Apache License 2.0 - see LICENSE file for details.

View Source

 1"""
 2.. include:: ../../README.md
 3"""
 4
 5from ._blob import TensorBlob
 6from ._db import TensorDB
 7
 8__version__ = "0.2.1"
 9
10__all__ = [
11    "TensorBlob",
12    "TensorDB",
13]

class TensorBlob(configmixin._core.ConfigMixin): View Source

 43class TensorBlob(ConfigMixin):
 44    _m_rd = False
 45    _m_wr = False
 46    _m_ap = False
 47
 48    status_name = ".stat"
 49    config_name = ".conf"
 50    ignore_for_config: ClassVar[list[str]] = ["filename", "mode", "max_cached_blocks"]
 51
 52    @classmethod
 53    def open(
 54        cls,
 55        filename,
 56        mode="r",
 57        *,
 58        dtype=None,
 59        shape=None,
 60        block_size=8192,
 61        max_cached_blocks=None,
 62    ):
 63        r"""Open a TensorBlob with file-like interface for tensor storage.
 64
 65        TensorBlob provides persistent, memory-mapped storage for large collections
 66        of same-shaped tensors. It uses a block-based architecture where tensors are
 67        organized into fixed-size blocks for efficient I/O and memory management.
 68
 69        The blob is stored as a directory containing:
 70        - ``.conf``: Configuration file (dtype, shape, block_size)
 71        - ``.stat``: State file (length, block list)
 72        - Block files: UUID-named memory-mapped tensor files
 73
 74        Parameters
 75        ----------
 76        filename : str or Path
 77            Directory path for blob storage. Supports tilde expansion (~) and
 78            relative paths.
 79        mode : str, default="r"
 80            File access mode ('r', 'w', 'a', 'r+', 'w+', 'a+'). See below for details.
 81        dtype : str or torch.dtype, optional
 82            Data type for tensors. Required for new blobs (modes 'w', 'w+').
 83        shape : tuple of int or int, optional
 84            Shape of individual tensors. Required for new blobs (modes 'w', 'w+').
 85        block_size : int, default=8192
 86            Number of tensors per memory-mapped block file.
 87        max_cached_blocks : int, optional
 88            Maximum number of memory-mapped blocks to keep cached. When exceeded,
 89            least recently used blocks are unmapped. If None (default), uses 1/16
 90            of system's max_map_count limit (typically ~4000). This limits kernel
 91            VMA overhead for blobs with many blocks.
 92
 93        Returns
 94        -------
 95        TensorBlob
 96            Opened blob object. Use with context manager for automatic cleanup.
 97
 98        Raises
 99        ------
100        FileNotFoundError
101            If mode is 'r', 'r+', 'a', or 'a+' and blob doesn't exist.
102        ValueError
103            If creating new blob without dtype or shape, or if mode is invalid.
104        TypeError
105            If dtype is neither string nor torch.dtype.
106
107        Examples
108        --------
109        Creating a new blob and writing data:
110
111        >>> import torch
112        >>> from tensorblob import TensorBlob
113        >>>
114        >>> with TensorBlob.open("data/embeddings", "w",
115        ...                       dtype="float32", shape=(768,)) as blob:
116        ...     embeddings = torch.randn(1000, 768)
117        ...     blob.write(embeddings)
118        ...     print(f"Wrote {len(blob)} tensors")
119        Wrote 1000 tensors
120
121        Reading from existing blob:
122
123        >>> with TensorBlob.open("data/embeddings", "r") as blob:
124        ...     all_data = blob.read()
125        ...     print(all_data.shape)
126        torch.Size([1000, 768])
127
128        Appending to existing blob:
129
130        >>> with TensorBlob.open("data/embeddings", "a") as blob:
131        ...     new_data = torch.randn(100, 768)
132        ...     blob.write(new_data)
133        ...     print(f"Total: {len(blob)}")
134        Total: 1100
135
136        Read and update with r+ mode:
137
138        >>> with TensorBlob.open("data/embeddings", "r+") as blob:
139        ...     first_10 = blob.read(size=10)
140        ...     blob.seek(5)
141        ...     blob.write(torch.ones(3, 768))  # Overwrite at position 5
142
143        Custom block size for large tensors:
144
145        >>> with TensorBlob.open("data/images", "w",
146        ...                       dtype=torch.float32,
147        ...                       shape=(3, 1024, 1024),
148        ...                       block_size=256) as blob:
149        ...     images = torch.randn(1000, 3, 1024, 1024)
150        ...     blob.write(images)
151
152        Custom cache size for large-scale random access:
153
154        >>> # Increase cache for better random access performance
155        >>> with TensorBlob.open("data/embeddings", "r",
156        ...                       max_cached_blocks=10000) as blob:
157        ...     for idx in random_indices:
158        ...         embedding = blob[idx]  # Frequently accessed blocks stay cached
159
160        >>> # Decrease cache for memory-constrained environments
161        >>> with TensorBlob.open("data/features", "r",
162        ...                       max_cached_blocks=100) as blob:
163        ...     for feature in blob:  # Sequential access works fine
164        ...         process(feature)
165
166        File Access Modes
167        -----------------
168        Similar to Python's built-in open(), supports the following modes:
169
170        Basic modes:
171        - 'r'  : Read-only. Blob must exist. Position starts at beginning.
172        - 'w'  : Write-only. Creates new or truncates existing. Position at start. **If the blob already exists,
173                   truncation will ignore any other parameters supplied and rely on existing configuration.**
174        - 'a'  : Append-only. Blob must exist. Position starts at end.
175                All writes go to end regardless of seek position.
176
177        Update modes (with '+'):
178        - 'r+' : Read and write. Blob must exist. Position at start.
179                   Can overwrite existing data or extend at end.
180        - 'w+' : Read and write. Creates new or truncates existing. Position at start.
181        - 'a+' : Read and append. Blob must exist. Position at end.
182                   Reads allowed anywhere, writes always append to end.
183
184        Data Type and Shape
185        -------------------
186        All tensors in a blob must have the same dtype and shape. These are
187        specified when creating a new blob (modes 'w', 'w+') and stored in
188        the configuration file. When opening existing blobs, dtype and shape
189        are loaded automatically.
190
191        Supported dtypes: "float32", "float64", "int32", "int64", "bool", etc.
192        Can also use torch.dtype objects like torch.float32.
193
194        Shape can be:
195        - Single integer: shape=10 creates 1D tensors of shape (10,)
196        - Tuple: shape=(3, 224, 224) creates 3D tensors
197        """
198        modes = set(mode)
199        if modes - set("raw+") or len(mode) > len(modes):
200            raise ValueError(f"Invalid mode: {mode}")
201        if sum(c in "raw" for c in mode) != 1 or mode.count("+") > 1:
202            raise ValueError(
203                f"Must have exactly one of read/write/append mode and at most one plus: {mode}"
204            )
205
206        filename = Path(filename).expanduser().resolve()
207        if not filename.exists():
208            if "r" in modes or "a" in modes:
209                raise FileNotFoundError(f"Blob not found: {filename!r}")
210            if dtype is None or shape is None:
211                raise ValueError(
212                    f"Arguments ``dtype`` and ``shape`` are required for new blob; got: {dtype!r} and {shape!r}"
213                )
214            if isinstance(dtype, torch.dtype):
215                dtype = str(dtype).split(".").pop()
216            elif not isinstance(dtype, str):
217                raise TypeError(
218                    f"dtype must be str or torch.dtype, got {type(dtype).__name__!r}"
219                )
220            shape = (shape,) if isinstance(shape, int) else tuple(shape)
221            return cls(
222                os.fspath(filename), dtype, shape, block_size, mode, max_cached_blocks
223            )
224
225        return cls.from_config(
226            save_directory=filename,
227            runtime_kwargs={
228                "mode": mode,
229                "filename": os.fspath(filename),
230                "max_cached_blocks": max_cached_blocks,
231            },
232        )
233
234    @classmethod
235    def unlink(cls, filename):
236        filename = Path(filename).expanduser().resolve()
237        if filename.exists():
238            try:
239                with cls.open(filename, "w") as _:
240                    pass
241                os.unlink(filename / cls.config_name)
242                os.unlink(filename / cls.status_name)
243                os.rmdir(os.fspath(filename))
244            except (OSError, ValueError) as exc:
245                warnings.warn(f"Failed to unlink blob at {filename!r}: {exc}")
246                return False
247        return True
248
249    @classmethod
250    def apply_param_hooks(cls, jdict):
251        jdict["shape"] = tuple(jdict["shape"])
252        return jdict
253
254    @classmethod
255    def _getsyscachesize(cls) -> int:
256        # Get default cache size for memory-mapped blocks. Returns 1/16 of system's
257        # max_map_count to be conservative, typically ~4000, leaving room for other
258        # VMAs in the process.
259        maxsize = 65536
260        try:
261            with open("/proc/sys/vm/max_map_count", "r") as f:
262                maxsize = int(f.read().strip())
263        except (FileNotFoundError, ValueError, PermissionError):
264            pass
265        return max(maxsize // 16, 128)
266
267    @register_to_config
268    def __init__(
269        self,
270        filename: str,
271        dtype: str,
272        shape: tuple[int, ...],
273        block_size: int,
274        mode: str,
275        max_cached_blocks: int | None = None,
276    ) -> None:
277        self.filename = filename
278        self.dtype = dtype
279        self.shape = shape
280        self.block_size = block_size
281        self.mode = mode
282        self.max_cached_blocks = max_cached_blocks or self._getsyscachesize()
283
284        self._pos = 0
285        self._closed = False
286
287        if "+" in mode:
288            self._m_rd = True
289            self._m_wr = True
290        match mode.replace("+", ""):
291            case "r":
292                self._m_rd = True
293            case "w":
294                self._m_wr = True
295                self._trunc()
296            case "a":
297                self._m_wr = True
298                self._m_ap = True
299                self._create()
300
301        self._loadstatus()
302
303    @property
304    def configpath(self) -> str:
305        return os.path.join(self.filename, self.config_name)
306
307    @property
308    def statuspath(self) -> str:
309        return os.path.join(self.filename, self.status_name)
310
311    @property
312    def closed(self) -> bool:
313        return self._closed
314
315    def __enter__(self) -> Self:
316        return self
317
318    def __exit__(self, *_) -> None:
319        self.close()
320
321    def __len__(self) -> int:
322        return self._status.len
323
324    def __getitem__(
325        self, idx: int | slice | list | tuple | torch.Tensor
326    ) -> torch.Tensor:
327        if isinstance(idx, int):
328            if idx >= len(self) or idx < -len(self):
329                raise IndexError(f"Index out of bounds: {idx!r} (length: {len(self)})")
330            i, o = divmod(idx + len(self) if idx < 0 else idx, self.block_size)
331            return self._getblock(i)[o].clone()
332        if isinstance(idx, slice):
333            # Although the current implementation may not be efficient, it is very easy to
334            # understand and debug. More efficient implementation requires much more complex
335            # edge case handling and is error prone. Also, I think the primary cost here is
336            # still the I/O operations, not the Python code.
337            ret = [
338                self._getblock(bd)[[i % self.block_size for i in _is]]
339                for bd, _is in groupby(
340                    range(*idx.indices(len(self))), key=lambda i: i // self.block_size
341                )
342            ]
343            if not ret:
344                return torch.empty(0, *self.shape, dtype=getattr(torch, self.dtype))
345            return torch.cat(ret, dim=0)
346        if isinstance(idx, (list, tuple, torch.Tensor)) or hasattr(idx, "__array__"):
347            return self._getbatch(idx)
348        raise TypeError(
349            "Index must be int, slice, or a sequence of int, "
350            f"got {type(idx).__name__!r}!"
351        )
352
353    def _getbatch(self, idx) -> torch.Tensor:
354        # Vectorized fancy indexing: gather rows in block-sorted order (one
355        # vectorized lookup per distinct block), then scatter back to the
356        # original order, mirroring torch's fancy indexing semantics.
357        raw = idx
358        idx = torch.as_tensor(idx)
359        if idx.numel() and (
360            idx.dtype == torch.bool or idx.is_floating_point() or idx.is_complex()
361        ):
362            if isinstance(raw, torch.Tensor):
363                got = f"dtype {raw.dtype}"
364            else:
365                got = f"elements of type {type(raw[0]).__name__!r}"
366            raise TypeError(f"Batch index must have an integer dtype, got {got}!")
367        idx = idx.long()
368        if idx.ndim != 1:
369            raise ValueError(f"Batch index must be 1-dimensional, got {idx.ndim}!")
370        if not idx.numel():
371            return torch.empty(0, *self.shape, dtype=getattr(torch, self.dtype))
372
373        n = len(self)
374        idx = torch.where(idx < 0, idx + n, idx)
375        if bool(((idx >= n) | (idx < 0)).any()):
376            raise IndexError(f"Index out of bounds (length: {n})")
377
378        order = torch.argsort(
379            torch.div(idx, self.block_size, rounding_mode="floor"), stable=True
380        )
381        sidx = idx[order]
382        blk = torch.div(sidx, self.block_size, rounding_mode="floor")
383        off = sidx - blk * self.block_size
384
385        # Boundaries between consecutive runs of identical block ids
386        brks = (
387            [0] + (blk[1:] != blk[:-1]).nonzero().flatten().add(1).tolist() + [len(blk)]
388        )
389        gathered = torch.cat(
390            [self._getblock(int(blk[lo]))[off[lo:hi]] for lo, hi in pairwise(brks)]
391        )
392        return gathered[torch.argsort(order)]
393
394    def __iter__(self) -> Iterator[torch.Tensor]:
395        for i in range(self._pos, len(self)):
396            self._pos += 1
397            yield self[i]
398
399    def _trunc(self) -> None:
400        if os.path.exists(self.filename):
401            try:
402                st = TensorBlobStatus.load(self.statuspath)
403            except FileNotFoundError as exc:
404                raise FileNotFoundError(
405                    f"Status file missing for blob at {self.statuspath!r}; file corrupted!"
406                ) from exc
407            for bd in st.bds:
408                os.remove(os.path.join(self.filename, bd))
409        self.save_config(save_directory=self.filename, overwrite=True)
410        TensorBlobStatus().dump(self.statuspath)
411
412    def _create(self) -> None:
413        if not os.path.exists(self.filename):
414            self.save_config(save_directory=self.filename)
415            TensorBlobStatus().dump(self.statuspath)
416
417    def _getblock(self, bd: str | int = -1) -> MemoryMappedTensor:
418        if not self._status.bds:
419            self._addblock()
420        if isinstance(bd, int):
421            bd = self._status.bds[bd]
422        if bd in self._memmap:
423            return self._memmap[bd]
424
425        # If cache no hit, a block is lazy-loaded into the cache. We need to
426        # avoid the __getitem__ call during return here to not increase the
427        # cache hit count a second time.
428        block = self._memmap[bd] = MemoryMappedTensor.from_filename(
429            os.path.join(self.filename, bd),
430            dtype=getattr(torch, self.dtype),
431            shape=(self.block_size, *self.shape),
432        )
433        return block
434
435    def _isfull(self) -> bool:
436        return (not len(self) % self.block_size) and bool(len(self))
437
438    def _addblock(self) -> MemoryMappedTensor:
439        if self._status.bds and not self._isfull():
440            raise RuntimeError(
441                "Attempt to create a new block when working block "
442                f"is not full: length <{len(self) % self.block_size}> "
443                f"< capacity <{self.block_size}>."
444            )
445        name = str(uuid.uuid4())
446        mmap = MemoryMappedTensor.empty(
447            self.block_size,
448            *self.shape,
449            dtype=getattr(torch, self.dtype),
450            filename=os.path.join(self.filename, name),
451        )
452        self._status.bds.append(name)
453        self._memmap[name] = mmap
454        return mmap
455
456    def _loadstatus(self) -> None:
457        try:
458            self._status = TensorBlobStatus.load(self.statuspath)
459            self._memmap = LRUCache(maxsize=self.max_cached_blocks)
460            if self._m_ap:
461                self._pos = len(self)
462        except FileNotFoundError as exc:
463            raise FileNotFoundError(
464                f"status file missing for blob at {self.statuspath!r}; file corrupted!"
465            ) from exc
466
467    def _checkclosed(self) -> None:
468        if self._closed:
469            raise OSError("I/O operation on closed blob.")
470
471    def _checkwritable(self) -> None:
472        if not self._m_wr:
473            raise OSError(f"Blob is not open for writing (mode='{self.mode}')")
474        self._checkclosed()
475
476    def _checkreadable(self) -> None:
477        if not self._m_rd:
478            raise OSError(f"Blob is not open for reading (mode='{self.mode}')")
479        self._checkclosed()
480
481    def tell(self) -> int:
482        self._checkclosed()
483        return self._pos
484
485    def seek(self, pos: int = 0, whence: int = io.SEEK_SET) -> int:
486        self._checkclosed()
487        match whence:
488            case io.SEEK_SET:
489                _pos = pos
490            case io.SEEK_CUR:
491                _pos = self._pos + pos
492            case io.SEEK_END:
493                _pos = len(self) + pos
494            case _:
495                raise ValueError(f"Invalid whence: {whence!r}")
496        self._pos = max(min(_pos, len(self)), 0)
497        return self.tell()
498
499    def close(self) -> None:
500        if not self._closed and self._m_wr:
501            self.flush()
502        self._closed = True
503
504    def flush(self) -> None:
505        self._checkwritable()
506        self._status.dump(self.statuspath)
507
508    def read(self, size: int | None = None) -> torch.Tensor:
509        self._checkreadable()
510        end = min(self._pos + (size if size is not None else len(self)), len(self))
511        ret = self[self._pos : end]
512        self.seek(end)
513        return ret
514
515    def write(self, ts: torch.Tensor) -> int:
516        self._checkwritable()
517        if self._m_ap:
518            self.seek(whence=io.SEEK_END)
519        ts = ts.view(-1, *self.shape)
520        nt = ts.size(0)
521
522        cnt = 0
523        while cnt < nt:
524            if self._isfull() and self._pos >= len(self):
525                self._addblock()
526            i, o = divmod(self._pos, self.block_size)
527            incr = min(self.block_size - o, nt - cnt)
528            self._getblock(i)[o : o + incr] = ts[cnt : cnt + incr]
529
530            # Update status length for new tensors exceeding the original range only, because
531            # the cursor may not always be at the EOF and the number of tensors written could
532            # be smaller than change in length
533            self._pos += incr
534            self._status.len += max(0, self._pos - len(self))
535
536            cnt += incr
537
538        assert cnt == nt, f"Write incomplete: wrote {cnt} of {nt} tensors!"
539        return cnt
540
541    def truncate(self, pos: int | None = None) -> int:
542        self._checkwritable()
543        if pos is not None and pos < 0:
544            raise ValueError(f"Truncate position must be non-negative, got {pos}!")
545        self.seek(pos if pos is not None else self.tell())
546        brk = ceil(self.tell() / self.block_size)
547        for bd in self._status.bds[brk:]:
548            if bd in self._memmap:
549                del self._memmap[bd]
550            os.remove(os.path.join(self.filename, bd))
551        self._status.bds = self._status.bds[:brk]
552        self._status.len = self.tell()
553        self.flush()
554        return self.tell()
555
556    def extend(self, other: TensorBlob, maintain_order: bool = False) -> None:
557        if self.dtype != other.dtype or self.shape != other.shape:
558            raise ValueError("Blob data types and shapes must match to extend blobs!")
559
560        self._checkwritable()
561        self.seek(whence=io.SEEK_END)
562
563        # TODO: Honestly this is a bit inefficient but I think this is rarely used.
564        if maintain_order:
565            for i in range(len(other)):
566                self.write(other[i])
567            return
568
569        # If order is not important, we can simply copy over the complete blocks from
570        # the other blob and merge incomplete blocks.
571        if self.block_size != other.block_size:
572            raise ValueError(
573                "Block sizes must match to extend blobs in non-order-preserving mode!"
574            )
575
576        comb = []
577        sbrk = len(self) // self.block_size * self.block_size
578        if sbrk < len(self):
579            comb.append(self[sbrk:])
580        obrk = len(other) // other.block_size * other.block_size
581        if obrk < len(other):
582            comb.append(other[obrk:])
583
584        # TODO: We are directly accessing internal data structures of the other blob here.
585        self.truncate(sbrk)
586        for obd in other._status.bds[: len(other) // other.block_size]:
587            sbd = str(uuid.uuid4())
588            shutil.copy(
589                os.path.join(other.filename, obd), os.path.join(self.filename, sbd)
590            )
591            self._status.bds.append(sbd)
592            self._status.len += self.block_size
593            self._memmap[sbd] = MemoryMappedTensor.from_filename(
594                os.path.join(self.filename, sbd),
595                dtype=getattr(torch, self.dtype),
596                shape=(self.block_size, *self.shape),
597            )
598
599        self.seek(whence=io.SEEK_END)
600        if comb:
601            self.write(torch.cat(comb, dim=0))
602        self.flush()

Mixin class for automated configuration registration and IO.

Attributes

config_name (str, default=None): Class attribute that specifies the filename under which the config should be stored when calling save_config. Should be overridden by the subclass.
ignore_for_config (list[str], default=[]): Class attribute that specifies a list of attributes that should not be saved in the config. Should be overridden by the subclass.

Examples

In this example, we have a model with 3 arguments:

hidden_size: The hidden size of the model.
_num_layers: The number of layers in the model.
dropout: The dropout rate of the model.

Among the three arguments, the number of layers is implicitly ignored by the decorator because of the leading underscore; the dropout argument is explicitly based on the specification in ignore_for_config class variable. The hidden_size argument is registered to the config.

>>> class MyModel(ConfigMixin):
...     config_name = "my_model_config.json"
...     ignore_for_config = ["dropout"]
...
...     @register_to_config
...     def __init__(self, hidden_size: int = 768, _num_layers: int = 12, dropout: float = 0.1):
...         self.hidden_size = hidden_size
...         self.num_layers = _num_layers
...         self.dropout = dropout  # This will be ignored because of the specification in `ignore_for_config`
...
>>> model = MyModel(hidden_size=1024, _num_layers=20, dropout=0.2)
>>> model.config
mappingproxy({'__notes__': {'class_name': '__main__.MyModel', 'using_default_values': [], 'args': (), 'kwargs': {}}, 'hidden_size': 1024})
>>> model.num_layers
20
>>> model.dropout
0.2

@register_to_config

TensorBlob( filename: str, dtype: str, shape: tuple[int, ...], block_size: int, mode: str, max_cached_blocks: int | None = None) View Source

267    @register_to_config
268    def __init__(
269        self,
270        filename: str,
271        dtype: str,
272        shape: tuple[int, ...],
273        block_size: int,
274        mode: str,
275        max_cached_blocks: int | None = None,
276    ) -> None:
277        self.filename = filename
278        self.dtype = dtype
279        self.shape = shape
280        self.block_size = block_size
281        self.mode = mode
282        self.max_cached_blocks = max_cached_blocks or self._getsyscachesize()
283
284        self._pos = 0
285        self._closed = False
286
287        if "+" in mode:
288            self._m_rd = True
289            self._m_wr = True
290        match mode.replace("+", ""):
291            case "r":
292                self._m_rd = True
293            case "w":
294                self._m_wr = True
295                self._trunc()
296            case "a":
297                self._m_wr = True
298                self._m_ap = True
299                self._create()
300
301        self._loadstatus()

status_name = '.stat'

config_name = '.conf'

ignore_for_config: ClassVar[list[str]] = ['filename', 'mode', 'max_cached_blocks']

@classmethod

def open( cls, filename, mode='r', *, dtype=None, shape=None, block_size=8192, max_cached_blocks=None): View Source

 52    @classmethod
 53    def open(
 54        cls,
 55        filename,
 56        mode="r",
 57        *,
 58        dtype=None,
 59        shape=None,
 60        block_size=8192,
 61        max_cached_blocks=None,
 62    ):
 63        r"""Open a TensorBlob with file-like interface for tensor storage.
 64
 65        TensorBlob provides persistent, memory-mapped storage for large collections
 66        of same-shaped tensors. It uses a block-based architecture where tensors are
 67        organized into fixed-size blocks for efficient I/O and memory management.
 68
 69        The blob is stored as a directory containing:
 70        - ``.conf``: Configuration file (dtype, shape, block_size)
 71        - ``.stat``: State file (length, block list)
 72        - Block files: UUID-named memory-mapped tensor files
 73
 74        Parameters
 75        ----------
 76        filename : str or Path
 77            Directory path for blob storage. Supports tilde expansion (~) and
 78            relative paths.
 79        mode : str, default="r"
 80            File access mode ('r', 'w', 'a', 'r+', 'w+', 'a+'). See below for details.
 81        dtype : str or torch.dtype, optional
 82            Data type for tensors. Required for new blobs (modes 'w', 'w+').
 83        shape : tuple of int or int, optional
 84            Shape of individual tensors. Required for new blobs (modes 'w', 'w+').
 85        block_size : int, default=8192
 86            Number of tensors per memory-mapped block file.
 87        max_cached_blocks : int, optional
 88            Maximum number of memory-mapped blocks to keep cached. When exceeded,
 89            least recently used blocks are unmapped. If None (default), uses 1/16
 90            of system's max_map_count limit (typically ~4000). This limits kernel
 91            VMA overhead for blobs with many blocks.
 92
 93        Returns
 94        -------
 95        TensorBlob
 96            Opened blob object. Use with context manager for automatic cleanup.
 97
 98        Raises
 99        ------
100        FileNotFoundError
101            If mode is 'r', 'r+', 'a', or 'a+' and blob doesn't exist.
102        ValueError
103            If creating new blob without dtype or shape, or if mode is invalid.
104        TypeError
105            If dtype is neither string nor torch.dtype.
106
107        Examples
108        --------
109        Creating a new blob and writing data:
110
111        >>> import torch
112        >>> from tensorblob import TensorBlob
113        >>>
114        >>> with TensorBlob.open("data/embeddings", "w",
115        ...                       dtype="float32", shape=(768,)) as blob:
116        ...     embeddings = torch.randn(1000, 768)
117        ...     blob.write(embeddings)
118        ...     print(f"Wrote {len(blob)} tensors")
119        Wrote 1000 tensors
120
121        Reading from existing blob:
122
123        >>> with TensorBlob.open("data/embeddings", "r") as blob:
124        ...     all_data = blob.read()
125        ...     print(all_data.shape)
126        torch.Size([1000, 768])
127
128        Appending to existing blob:
129
130        >>> with TensorBlob.open("data/embeddings", "a") as blob:
131        ...     new_data = torch.randn(100, 768)
132        ...     blob.write(new_data)
133        ...     print(f"Total: {len(blob)}")
134        Total: 1100
135
136        Read and update with r+ mode:
137
138        >>> with TensorBlob.open("data/embeddings", "r+") as blob:
139        ...     first_10 = blob.read(size=10)
140        ...     blob.seek(5)
141        ...     blob.write(torch.ones(3, 768))  # Overwrite at position 5
142
143        Custom block size for large tensors:
144
145        >>> with TensorBlob.open("data/images", "w",
146        ...                       dtype=torch.float32,
147        ...                       shape=(3, 1024, 1024),
148        ...                       block_size=256) as blob:
149        ...     images = torch.randn(1000, 3, 1024, 1024)
150        ...     blob.write(images)
151
152        Custom cache size for large-scale random access:
153
154        >>> # Increase cache for better random access performance
155        >>> with TensorBlob.open("data/embeddings", "r",
156        ...                       max_cached_blocks=10000) as blob:
157        ...     for idx in random_indices:
158        ...         embedding = blob[idx]  # Frequently accessed blocks stay cached
159
160        >>> # Decrease cache for memory-constrained environments
161        >>> with TensorBlob.open("data/features", "r",
162        ...                       max_cached_blocks=100) as blob:
163        ...     for feature in blob:  # Sequential access works fine
164        ...         process(feature)
165
166        File Access Modes
167        -----------------
168        Similar to Python's built-in open(), supports the following modes:
169
170        Basic modes:
171        - 'r'  : Read-only. Blob must exist. Position starts at beginning.
172        - 'w'  : Write-only. Creates new or truncates existing. Position at start. **If the blob already exists,
173                   truncation will ignore any other parameters supplied and rely on existing configuration.**
174        - 'a'  : Append-only. Blob must exist. Position starts at end.
175                All writes go to end regardless of seek position.
176
177        Update modes (with '+'):
178        - 'r+' : Read and write. Blob must exist. Position at start.
179                   Can overwrite existing data or extend at end.
180        - 'w+' : Read and write. Creates new or truncates existing. Position at start.
181        - 'a+' : Read and append. Blob must exist. Position at end.
182                   Reads allowed anywhere, writes always append to end.
183
184        Data Type and Shape
185        -------------------
186        All tensors in a blob must have the same dtype and shape. These are
187        specified when creating a new blob (modes 'w', 'w+') and stored in
188        the configuration file. When opening existing blobs, dtype and shape
189        are loaded automatically.
190
191        Supported dtypes: "float32", "float64", "int32", "int64", "bool", etc.
192        Can also use torch.dtype objects like torch.float32.
193
194        Shape can be:
195        - Single integer: shape=10 creates 1D tensors of shape (10,)
196        - Tuple: shape=(3, 224, 224) creates 3D tensors
197        """
198        modes = set(mode)
199        if modes - set("raw+") or len(mode) > len(modes):
200            raise ValueError(f"Invalid mode: {mode}")
201        if sum(c in "raw" for c in mode) != 1 or mode.count("+") > 1:
202            raise ValueError(
203                f"Must have exactly one of read/write/append mode and at most one plus: {mode}"
204            )
205
206        filename = Path(filename).expanduser().resolve()
207        if not filename.exists():
208            if "r" in modes or "a" in modes:
209                raise FileNotFoundError(f"Blob not found: {filename!r}")
210            if dtype is None or shape is None:
211                raise ValueError(
212                    f"Arguments ``dtype`` and ``shape`` are required for new blob; got: {dtype!r} and {shape!r}"
213                )
214            if isinstance(dtype, torch.dtype):
215                dtype = str(dtype).split(".").pop()
216            elif not isinstance(dtype, str):
217                raise TypeError(
218                    f"dtype must be str or torch.dtype, got {type(dtype).__name__!r}"
219                )
220            shape = (shape,) if isinstance(shape, int) else tuple(shape)
221            return cls(
222                os.fspath(filename), dtype, shape, block_size, mode, max_cached_blocks
223            )
224
225        return cls.from_config(
226            save_directory=filename,
227            runtime_kwargs={
228                "mode": mode,
229                "filename": os.fspath(filename),
230                "max_cached_blocks": max_cached_blocks,
231            },
232        )

Open a TensorBlob with file-like interface for tensor storage.

TensorBlob provides persistent, memory-mapped storage for large collections of same-shaped tensors. It uses a block-based architecture where tensors are organized into fixed-size blocks for efficient I/O and memory management.

The blob is stored as a directory containing:

.conf: Configuration file (dtype, shape, block_size)
.stat: State file (length, block list)
Block files: UUID-named memory-mapped tensor files

Parameters

filename (str or Path): Directory path for blob storage. Supports tilde expansion (~) and relative paths.
mode (str, default="r"): File access mode ('r', 'w', 'a', 'r+', 'w+', 'a+'). See below for details.
dtype (str or torch.dtype, optional): Data type for tensors. Required for new blobs (modes 'w', 'w+').
shape (tuple of int or int, optional): Shape of individual tensors. Required for new blobs (modes 'w', 'w+').
block_size (int, default=8192): Number of tensors per memory-mapped block file.
max_cached_blocks (int, optional): Maximum number of memory-mapped blocks to keep cached. When exceeded, least recently used blocks are unmapped. If None (default), uses 1/16 of system's max_map_count limit (typically ~4000). This limits kernel VMA overhead for blobs with many blocks.

Returns

TensorBlob: Opened blob object. Use with context manager for automatic cleanup.

Raises

FileNotFoundError: If mode is 'r', 'r+', 'a', or 'a+' and blob doesn't exist.
ValueError: If creating new blob without dtype or shape, or if mode is invalid.
TypeError: If dtype is neither string nor torch.dtype.

Examples

Creating a new blob and writing data:

>>> import torch
>>> from tensorblob import TensorBlob
>>>
>>> with TensorBlob.open("data/embeddings", "w",
...                       dtype="float32", shape=(768,)) as blob:
...     embeddings = torch.randn(1000, 768)
...     blob.write(embeddings)
...     print(f"Wrote {len(blob)} tensors")
Wrote 1000 tensors

Reading from existing blob:

>>> with TensorBlob.open("data/embeddings", "r") as blob:
...     all_data = blob.read()
...     print(all_data.shape)
torch.Size([1000, 768])

Appending to existing blob:

>>> with TensorBlob.open("data/embeddings", "a") as blob:
...     new_data = torch.randn(100, 768)
...     blob.write(new_data)
...     print(f"Total: {len(blob)}")
Total: 1100

Read and update with r+ mode:

>>> with TensorBlob.open("data/embeddings", "r+") as blob:
...     first_10 = blob.read(size=10)
...     blob.seek(5)
...     blob.write(torch.ones(3, 768))  # Overwrite at position 5

Custom block size for large tensors:

>>> with TensorBlob.open("data/images", "w",
...                       dtype=torch.float32,
...                       shape=(3, 1024, 1024),
...                       block_size=256) as blob:
...     images = torch.randn(1000, 3, 1024, 1024)
...     blob.write(images)

Custom cache size for large-scale random access:

>>> # Increase cache for better random access performance
>>> with TensorBlob.open("data/embeddings", "r",
...                       max_cached_blocks=10000) as blob:
...     for idx in random_indices:
...         embedding = blob[idx]  # Frequently accessed blocks stay cached

>>> # Decrease cache for memory-constrained environments
>>> with TensorBlob.open("data/features", "r",
...                       max_cached_blocks=100) as blob:
...     for feature in blob:  # Sequential access works fine
...         process(feature)

File Access Modes

Similar to Python's built-in open(), supports the following modes:

Basic modes:

'r' : Read-only. Blob must exist. Position starts at beginning.
'w' : Write-only. Creates new or truncates existing. Position at start. If the blob already exists, truncation will ignore any other parameters supplied and rely on existing configuration.
'a' : Append-only. Blob must exist. Position starts at end. All writes go to end regardless of seek position.

Update modes (with '+'):

'r+' : Read and write. Blob must exist. Position at start. Can overwrite existing data or extend at end.
'w+' : Read and write. Creates new or truncates existing. Position at start.
'a+' : Read and append. Blob must exist. Position at end. Reads allowed anywhere, writes always append to end.

Data Type and Shape

All tensors in a blob must have the same dtype and shape. These are specified when creating a new blob (modes 'w', 'w+') and stored in the configuration file. When opening existing blobs, dtype and shape are loaded automatically.

Supported dtypes: "float32", "float64", "int32", "int64", "bool", etc. Can also use torch.dtype objects like torch.float32.

Shape can be:

Single integer: shape=10 creates 1D tensors of shape (10,)
Tuple: shape=(3, 224, 224) creates 3D tensors

@classmethod

def unlink(cls, filename): View Source

234    @classmethod
235    def unlink(cls, filename):
236        filename = Path(filename).expanduser().resolve()
237        if filename.exists():
238            try:
239                with cls.open(filename, "w") as _:
240                    pass
241                os.unlink(filename / cls.config_name)
242                os.unlink(filename / cls.status_name)
243                os.rmdir(os.fspath(filename))
244            except (OSError, ValueError) as exc:
245                warnings.warn(f"Failed to unlink blob at {filename!r}: {exc}")
246                return False
247        return True

@classmethod

def apply_param_hooks(cls, jdict): View Source

249    @classmethod
250    def apply_param_hooks(cls, jdict):
251        jdict["shape"] = tuple(jdict["shape"])
252        return jdict

Apply post-processing hooks to the JSON dictionary.

orjson.loads only decode configs to primitive types, which may not be directly consumable by the class initializer. For instance, a dataclass object will be loaded as a dictionary. Therefore, this method is intended to be overridden by the subclass to perform additional post-processing on the loaded config dictionary.

Note that, it is highly discouraged to abuse this method to deserialize complex objects and one should consider using runtime_kwargs argument of from_config instead, to explicitly pass the complex objects to the class initializer.

By default, this method returns the input dictionary unchanged.

Parameters

jdict (dict[str, Any]): The config dictionary after deserialization.

Returns

dict[str, Any]: The config dictionary after post-processing.

filename

dtype

shape

block_size

mode

max_cached_blocks

configpath: str View Source

303    @property
304    def configpath(self) -> str:
305        return os.path.join(self.filename, self.config_name)

statuspath: str View Source

307    @property
308    def statuspath(self) -> str:
309        return os.path.join(self.filename, self.status_name)

closed: bool View Source

311    @property
312    def closed(self) -> bool:
313        return self._closed

def tell(self) -> int: View Source

481    def tell(self) -> int:
482        self._checkclosed()
483        return self._pos

def seek(self, pos: int = 0, whence: int = 0) -> int: View Source

485    def seek(self, pos: int = 0, whence: int = io.SEEK_SET) -> int:
486        self._checkclosed()
487        match whence:
488            case io.SEEK_SET:
489                _pos = pos
490            case io.SEEK_CUR:
491                _pos = self._pos + pos
492            case io.SEEK_END:
493                _pos = len(self) + pos
494            case _:
495                raise ValueError(f"Invalid whence: {whence!r}")
496        self._pos = max(min(_pos, len(self)), 0)
497        return self.tell()

def close(self) -> None: View Source

499    def close(self) -> None:
500        if not self._closed and self._m_wr:
501            self.flush()
502        self._closed = True

def flush(self) -> None: View Source

504    def flush(self) -> None:
505        self._checkwritable()
506        self._status.dump(self.statuspath)

def read(self, size: int | None = None) -> torch.Tensor: View Source

508    def read(self, size: int | None = None) -> torch.Tensor:
509        self._checkreadable()
510        end = min(self._pos + (size if size is not None else len(self)), len(self))
511        ret = self[self._pos : end]
512        self.seek(end)
513        return ret

def write(self, ts: torch.Tensor) -> int: View Source

515    def write(self, ts: torch.Tensor) -> int:
516        self._checkwritable()
517        if self._m_ap:
518            self.seek(whence=io.SEEK_END)
519        ts = ts.view(-1, *self.shape)
520        nt = ts.size(0)
521
522        cnt = 0
523        while cnt < nt:
524            if self._isfull() and self._pos >= len(self):
525                self._addblock()
526            i, o = divmod(self._pos, self.block_size)
527            incr = min(self.block_size - o, nt - cnt)
528            self._getblock(i)[o : o + incr] = ts[cnt : cnt + incr]
529
530            # Update status length for new tensors exceeding the original range only, because
531            # the cursor may not always be at the EOF and the number of tensors written could
532            # be smaller than change in length
533            self._pos += incr
534            self._status.len += max(0, self._pos - len(self))
535
536            cnt += incr
537
538        assert cnt == nt, f"Write incomplete: wrote {cnt} of {nt} tensors!"
539        return cnt

def truncate(self, pos: int | None = None) -> int: View Source

541    def truncate(self, pos: int | None = None) -> int:
542        self._checkwritable()
543        if pos is not None and pos < 0:
544            raise ValueError(f"Truncate position must be non-negative, got {pos}!")
545        self.seek(pos if pos is not None else self.tell())
546        brk = ceil(self.tell() / self.block_size)
547        for bd in self._status.bds[brk:]:
548            if bd in self._memmap:
549                del self._memmap[bd]
550            os.remove(os.path.join(self.filename, bd))
551        self._status.bds = self._status.bds[:brk]
552        self._status.len = self.tell()
553        self.flush()
554        return self.tell()

def extend( self, other: TensorBlob, maintain_order: bool = False) -> None: View Source

556    def extend(self, other: TensorBlob, maintain_order: bool = False) -> None:
557        if self.dtype != other.dtype or self.shape != other.shape:
558            raise ValueError("Blob data types and shapes must match to extend blobs!")
559
560        self._checkwritable()
561        self.seek(whence=io.SEEK_END)
562
563        # TODO: Honestly this is a bit inefficient but I think this is rarely used.
564        if maintain_order:
565            for i in range(len(other)):
566                self.write(other[i])
567            return
568
569        # If order is not important, we can simply copy over the complete blocks from
570        # the other blob and merge incomplete blocks.
571        if self.block_size != other.block_size:
572            raise ValueError(
573                "Block sizes must match to extend blobs in non-order-preserving mode!"
574            )
575
576        comb = []
577        sbrk = len(self) // self.block_size * self.block_size
578        if sbrk < len(self):
579            comb.append(self[sbrk:])
580        obrk = len(other) // other.block_size * other.block_size
581        if obrk < len(other):
582            comb.append(other[obrk:])
583
584        # TODO: We are directly accessing internal data structures of the other blob here.
585        self.truncate(sbrk)
586        for obd in other._status.bds[: len(other) // other.block_size]:
587            sbd = str(uuid.uuid4())
588            shutil.copy(
589                os.path.join(other.filename, obd), os.path.join(self.filename, sbd)
590            )
591            self._status.bds.append(sbd)
592            self._status.len += self.block_size
593            self._memmap[sbd] = MemoryMappedTensor.from_filename(
594                os.path.join(self.filename, sbd),
595                dtype=getattr(torch, self.dtype),
596                shape=(self.block_size, *self.shape),
597            )
598
599        self.seek(whence=io.SEEK_END)
600        if comb:
601            self.write(torch.cat(comb, dim=0))
602        self.flush()

class TensorDB(configmixin._core.ConfigMixin): View Source

 38class TensorDB(ConfigMixin):
 39    _m_rd = False
 40    _m_wr = False
 41    _m_ap = False
 42
 43    status_name = ".stat"
 44    config_name = ".conf"
 45    ignore_for_config: ClassVar[list[str]] = ["filename", "mode", "max_cached_blocks"]
 46
 47    @classmethod
 48    def open(
 49        cls,
 50        filename,
 51        mode="r",
 52        *,
 53        schema=None,
 54        block_size=8192,
 55        max_cached_blocks=None,
 56    ):
 57        r"""Open a TensorDB with file-like interface for multi-field tensor storage.
 58
 59        TensorDB provides persistent, row-aligned storage for heterogeneous
 60        (multi-field) tensor collections. Each field is stored as a plain
 61        :class:`TensorBlob` in a subdirectory, and TensorDB keeps the row
 62        orders of all fields aligned. Rows are dense: every write must supply
 63        every field with the same row count.
 64
 65        The database is stored as a directory containing:
 66        - ``.conf``: Schema file (field names, dtypes, shapes)
 67        - ``.stat``: State file (committed row count)
 68        - Field-named subdirectories: One TensorBlob per field
 69
 70        Parameters
 71        ----------
 72        filename : str or Path
 73            Directory path for database storage. Supports tilde expansion (~)
 74            and relative paths.
 75        mode : str, default="r"
 76            File access mode ('r', 'w', 'a', 'r+', 'w+', 'a+'). Behaves like
 77            :meth:`TensorBlob.open`; the mode is applied to every field.
 78        schema : dict, optional
 79            Mapping of field names to ``(dtype, shape)`` pairs, e.g.,
 80            ``{"price": ("float32", 1), "embed": (torch.float16, (768,))}``.
 81            Required for new databases (modes 'w', 'w+'). Fixed at creation;
 82            loaded automatically when opening existing databases.
 83        block_size : int, default=8192
 84            Number of rows per memory-mapped block file, applied to all fields.
 85        max_cached_blocks : int, optional
 86            Maximum number of memory-mapped blocks to keep cached per field.
 87            If None (default), uses 1/16 of system's max_map_count limit.
 88
 89        Returns
 90        -------
 91        TensorDB
 92            Opened database object. Use with context manager for automatic
 93            cleanup.
 94
 95        Raises
 96        ------
 97        FileNotFoundError
 98            If mode is 'r', 'r+', 'a', or 'a+' and database doesn't exist.
 99        ValueError
100            If creating new database without schema, if the schema is
101            malformed, or if mode is invalid.
102        TypeError
103            If a dtype is neither string nor torch.dtype.
104
105        Examples
106        --------
107        Creating a new database and writing dense rows:
108
109        >>> import torch
110        >>> from tensorblob import TensorDB
111        >>>
112        >>> with TensorDB.open("events.db", "w",
113        ...                    schema={"price": ("float32", 1),
114        ...                            "embed": ("float32", 768)}) as db:
115        ...     db.write({"price": torch.randn(1000, 1),
116        ...               "embed": torch.randn(1000, 768)})
117        ...     print(f"Wrote {len(db)} rows")
118        Wrote 1000 rows
119
120        Reading rows back, aligned across fields:
121
122        >>> with TensorDB.open("events.db", "r") as db:
123        ...     row = db[42]        # {"price": (1,), "embed": (768,)}
124        ...     batch = db[10:100]  # {"price": (90, 1), "embed": (90, 768)}
125
126        Notes
127        -----
128        Writes commit the row count only after all fields are written. If a
129        crash leaves fields longer than the committed count, the next open
130        reports the committed (minimum) length and, for writable modes,
131        truncates the stray rows back to restore alignment.
132        """
133        modes = set(mode)
134        if modes - set("raw+") or len(mode) > len(modes):
135            raise ValueError(f"Invalid mode: {mode}")
136        if sum(c in "raw" for c in mode) != 1 or mode.count("+") > 1:
137            raise ValueError(
138                f"Must have exactly one of read/write/append mode and at most one plus: {mode}"
139            )
140
141        filename = Path(filename).expanduser().resolve()
142        if not filename.exists():
143            if "r" in modes or "a" in modes:
144                raise FileNotFoundError(f"Database not found: {filename!r}")
145            if schema is None:
146                raise ValueError("Argument ``schema`` is required for new database!")
147            schema = cls._normalize_schema(schema)
148            return cls(os.fspath(filename), schema, block_size, mode, max_cached_blocks)
149
150        return cls.from_config(
151            save_directory=filename,
152            runtime_kwargs={
153                "mode": mode,
154                "filename": os.fspath(filename),
155                "max_cached_blocks": max_cached_blocks,
156            },
157        )
158
159    @classmethod
160    def unlink(cls, filename):
161        filename = Path(filename).expanduser().resolve()
162        if filename.exists():
163            try:
164                shutil.rmtree(filename)
165            except OSError as exc:
166                warnings.warn(f"Failed to unlink database at {filename!r}: {exc}")
167                return False
168        return True
169
170    @classmethod
171    def _normalize_schema(cls, schema):
172        if not isinstance(schema, dict) or not schema:
173            raise ValueError(
174                "Schema must be a non-empty dict mapping field names to (dtype, shape)!"
175            )
176        norm = {}
177        for name, spec in schema.items():
178            if (
179                not isinstance(name, str)
180                or not name
181                or name.startswith(".")
182                or "/" in name
183                or os.sep in name
184                or (os.altsep and os.altsep in name)
185            ):
186                raise ValueError(f"Invalid field name: {name!r}")
187            dtype, shape = spec
188            if isinstance(dtype, torch.dtype):
189                dtype = str(dtype).split(".").pop()
190            elif not isinstance(dtype, str):
191                raise TypeError(
192                    f"dtype must be str or torch.dtype, got {type(dtype).__name__!r}"
193                )
194            shape = (shape,) if isinstance(shape, int) else tuple(shape)
195            norm[name] = (dtype, shape)
196        return norm
197
198    @classmethod
199    def apply_param_hooks(cls, jdict):
200        jdict["schema"] = {
201            name: (dtype, tuple(shape))
202            for name, (dtype, shape) in jdict["schema"].items()
203        }
204        return jdict
205
206    @register_to_config
207    def __init__(
208        self,
209        filename: str,
210        schema: dict[str, tuple[str, tuple[int, ...]]],
211        block_size: int,
212        mode: str,
213        max_cached_blocks: int | None = None,
214    ) -> None:
215        self.filename = filename
216        self.schema = schema
217        self.block_size = block_size
218        self.mode = mode
219        self.max_cached_blocks = max_cached_blocks
220
221        self._closed = False
222
223        if "+" in mode:
224            self._m_rd = True
225            self._m_wr = True
226        match mode.replace("+", ""):
227            case "r":
228                self._m_rd = True
229            case "w":
230                self._m_wr = True
231            case "a":
232                self._m_wr = True
233                self._m_ap = True
234
235        isnew = not os.path.exists(self.filename)
236        if isnew:
237            os.makedirs(self.filename)
238            self.save_config(save_directory=self.filename)
239        self._cols = {
240            name: TensorBlob.open(
241                os.path.join(self.filename, name),
242                mode,
243                dtype=dtype,
244                shape=shape,
245                block_size=block_size,
246                max_cached_blocks=max_cached_blocks,
247            )
248            for name, (dtype, shape) in self.schema.items()
249        }
250
251        # For new or truncated databases the committed count starts at zero; no
252        # repair is needed since the columns were just (re)initialized above.
253        if isnew or "w" in mode:
254            self._status = TensorDBStatus()
255            self._status.dump(self.statuspath)
256        else:
257            self._loadstatus()
258
259    @property
260    def configpath(self) -> str:
261        return os.path.join(self.filename, self.config_name)
262
263    @property
264    def statuspath(self) -> str:
265        return os.path.join(self.filename, self.status_name)
266
267    @property
268    def closed(self) -> bool:
269        return self._closed
270
271    def __enter__(self) -> Self:
272        return self
273
274    def __exit__(self, *_) -> None:
275        self.close()
276
277    def __len__(self) -> int:
278        return self._status.len
279
280    def __getitem__(
281        self, idx: int | slice | list | tuple | torch.Tensor
282    ) -> dict[str, torch.Tensor]:
283        if not isinstance(idx, (int, slice, list, tuple, torch.Tensor)) and not hasattr(
284            idx, "__array__"
285        ):
286            raise TypeError(
287                "Index must be int, slice, or a sequence of int, "
288                f"got {type(idx).__name__!r}!"
289            )
290        return {name: col[idx] for name, col in self._cols.items()}
291
292    def __iter__(self) -> Iterator[dict[str, torch.Tensor]]:
293        for i in range(self.tell(), len(self)):
294            self.seek(i + 1)
295            yield self[i]
296
297    def _loadstatus(self) -> None:
298        try:
299            self._status = TensorDBStatus.load(self.statuspath)
300        except FileNotFoundError as exc:
301            raise FileNotFoundError(
302                f"Status file missing for database at {self.statuspath!r}; file corrupted!"
303            ) from exc
304
305        # The committed row count is the source of truth. A crash mid-write can
306        # leave some columns longer than the committed count; report the minimum
307        # and, if writable, truncate stray rows back to restore alignment.
308        target = min([self._status.len] + [len(col) for col in self._cols.values()])
309        if target != self._status.len or any(
310            len(col) != target for col in self._cols.values()
311        ):
312            warnings.warn(
313                f"Inconsistent column lengths detected for database at {self.filename!r}; "
314                f"reporting {target} committed rows."
315            )
316            self._status.len = target
317            if self._m_wr:
318                for col in self._cols.values():
319                    if len(col) != target:
320                        col.truncate(target)
321                self._status.dump(self.statuspath)
322
323    def _checkclosed(self) -> None:
324        if self._closed:
325            raise OSError("I/O operation on closed database.")
326
327    def _checkwritable(self) -> None:
328        if not self._m_wr:
329            raise OSError(f"Database is not open for writing (mode='{self.mode}')")
330        self._checkclosed()
331
332    def _checkreadable(self) -> None:
333        if not self._m_rd:
334            raise OSError(f"Database is not open for reading (mode='{self.mode}')")
335        self._checkclosed()
336
337    def tell(self) -> int:
338        self._checkclosed()
339        return next(iter(self._cols.values())).tell()
340
341    def seek(self, pos: int = 0, whence: int = io.SEEK_SET) -> int:
342        self._checkclosed()
343        for col in self._cols.values():
344            col.seek(pos, whence)
345        return self.tell()
346
347    def close(self) -> None:
348        if self._closed:
349            return
350        for col in self._cols.values():
351            col.close()
352        if self._m_wr:
353            self._status.dump(self.statuspath)
354        self._closed = True
355
356    def flush(self) -> None:
357        self._checkwritable()
358        for col in self._cols.values():
359            col.flush()
360        self._status.dump(self.statuspath)
361
362    def read(self, size: int | None = None) -> dict[str, torch.Tensor]:
363        self._checkreadable()
364        # Clamp at the committed row count so uncommitted trailing rows left by
365        # an interrupted write are never visible.
366        remaining = len(self) - self.tell()
367        size = remaining if size is None else min(size, remaining)
368        if size <= 0:
369            return {
370                name: torch.empty(0, *shape, dtype=getattr(torch, dtype))
371                for name, (dtype, shape) in self.schema.items()
372            }
373        return {name: col.read(size) for name, col in self._cols.items()}
374
375    def write(self, rows: dict[str, torch.Tensor]) -> int:
376        self._checkwritable()
377        if not isinstance(rows, dict):
378            raise TypeError(
379                f"Rows must be a dict mapping field names to tensors, got {type(rows).__name__!r}!"
380            )
381        missing = sorted(self.schema.keys() - rows.keys())
382        extra = sorted(rows.keys() - self.schema.keys())
383        if missing or extra:
384            raise ValueError(
385                "Dense writes require exactly the schema fields; "
386                f"missing: {missing!r}, unexpected: {extra!r}"
387            )
388
389        nts = {
390            name: ts.view(-1, *self.schema[name][1]).size(0)
391            for name, ts in rows.items()
392        }
393        if len(set(nts.values())) != 1:
394            raise ValueError(f"All fields must have the same row count; got: {nts!r}")
395        nt = next(iter(nts.values()))
396
397        # Columns are written first and the committed row count is bumped only
398        # afterwards, so an interrupted write is rolled back on the next open.
399        for name, col in self._cols.items():
400            col.write(rows[name])
401        self._status.len = len(next(iter(self._cols.values())))
402        return nt
403
404    def truncate(self, pos: int | None = None) -> int:
405        self._checkwritable()
406        for col in self._cols.values():
407            col.truncate(pos)
408        self._status.len = self.tell()
409        self._status.dump(self.statuspath)
410        return self.tell()
411
412    def extend(self, other: TensorDB, maintain_order: bool = False) -> None:
413        if set(self.schema) != set(other.schema):
414            raise ValueError("Schema fields must match to extend databases!")
415        self._checkwritable()
416        for name, col in self._cols.items():
417            col.extend(other._cols[name], maintain_order=maintain_order)
418        self._status.len = len(next(iter(self._cols.values())))
419        self._status.dump(self.statuspath)

Mixin class for automated configuration registration and IO.

Attributes

config_name (str, default=None): Class attribute that specifies the filename under which the config should be stored when calling save_config. Should be overridden by the subclass.
ignore_for_config (list[str], default=[]): Class attribute that specifies a list of attributes that should not be saved in the config. Should be overridden by the subclass.

Examples

In this example, we have a model with 3 arguments:

hidden_size: The hidden size of the model.
_num_layers: The number of layers in the model.
dropout: The dropout rate of the model.

>>> class MyModel(ConfigMixin):
...     config_name = "my_model_config.json"
...     ignore_for_config = ["dropout"]
...
...     @register_to_config
...     def __init__(self, hidden_size: int = 768, _num_layers: int = 12, dropout: float = 0.1):
...         self.hidden_size = hidden_size
...         self.num_layers = _num_layers
...         self.dropout = dropout  # This will be ignored because of the specification in `ignore_for_config`
...
>>> model = MyModel(hidden_size=1024, _num_layers=20, dropout=0.2)
>>> model.config
mappingproxy({'__notes__': {'class_name': '__main__.MyModel', 'using_default_values': [], 'args': (), 'kwargs': {}}, 'hidden_size': 1024})
>>> model.num_layers
20
>>> model.dropout
0.2

@register_to_config

TensorDB( filename: str, schema: dict[str, tuple[str, tuple[int, ...]]], block_size: int, mode: str, max_cached_blocks: int | None = None) View Source

206    @register_to_config
207    def __init__(
208        self,
209        filename: str,
210        schema: dict[str, tuple[str, tuple[int, ...]]],
211        block_size: int,
212        mode: str,
213        max_cached_blocks: int | None = None,
214    ) -> None:
215        self.filename = filename
216        self.schema = schema
217        self.block_size = block_size
218        self.mode = mode
219        self.max_cached_blocks = max_cached_blocks
220
221        self._closed = False
222
223        if "+" in mode:
224            self._m_rd = True
225            self._m_wr = True
226        match mode.replace("+", ""):
227            case "r":
228                self._m_rd = True
229            case "w":
230                self._m_wr = True
231            case "a":
232                self._m_wr = True
233                self._m_ap = True
234
235        isnew = not os.path.exists(self.filename)
236        if isnew:
237            os.makedirs(self.filename)
238            self.save_config(save_directory=self.filename)
239        self._cols = {
240            name: TensorBlob.open(
241                os.path.join(self.filename, name),
242                mode,
243                dtype=dtype,
244                shape=shape,
245                block_size=block_size,
246                max_cached_blocks=max_cached_blocks,
247            )
248            for name, (dtype, shape) in self.schema.items()
249        }
250
251        # For new or truncated databases the committed count starts at zero; no
252        # repair is needed since the columns were just (re)initialized above.
253        if isnew or "w" in mode:
254            self._status = TensorDBStatus()
255            self._status.dump(self.statuspath)
256        else:
257            self._loadstatus()

status_name = '.stat'

config_name = '.conf'

ignore_for_config: ClassVar[list[str]] = ['filename', 'mode', 'max_cached_blocks']

@classmethod

def open( cls, filename, mode='r', *, schema=None, block_size=8192, max_cached_blocks=None): View Source

 47    @classmethod
 48    def open(
 49        cls,
 50        filename,
 51        mode="r",
 52        *,
 53        schema=None,
 54        block_size=8192,
 55        max_cached_blocks=None,
 56    ):
 57        r"""Open a TensorDB with file-like interface for multi-field tensor storage.
 58
 59        TensorDB provides persistent, row-aligned storage for heterogeneous
 60        (multi-field) tensor collections. Each field is stored as a plain
 61        :class:`TensorBlob` in a subdirectory, and TensorDB keeps the row
 62        orders of all fields aligned. Rows are dense: every write must supply
 63        every field with the same row count.
 64
 65        The database is stored as a directory containing:
 66        - ``.conf``: Schema file (field names, dtypes, shapes)
 67        - ``.stat``: State file (committed row count)
 68        - Field-named subdirectories: One TensorBlob per field
 69
 70        Parameters
 71        ----------
 72        filename : str or Path
 73            Directory path for database storage. Supports tilde expansion (~)
 74            and relative paths.
 75        mode : str, default="r"
 76            File access mode ('r', 'w', 'a', 'r+', 'w+', 'a+'). Behaves like
 77            :meth:`TensorBlob.open`; the mode is applied to every field.
 78        schema : dict, optional
 79            Mapping of field names to ``(dtype, shape)`` pairs, e.g.,
 80            ``{"price": ("float32", 1), "embed": (torch.float16, (768,))}``.
 81            Required for new databases (modes 'w', 'w+'). Fixed at creation;
 82            loaded automatically when opening existing databases.
 83        block_size : int, default=8192
 84            Number of rows per memory-mapped block file, applied to all fields.
 85        max_cached_blocks : int, optional
 86            Maximum number of memory-mapped blocks to keep cached per field.
 87            If None (default), uses 1/16 of system's max_map_count limit.
 88
 89        Returns
 90        -------
 91        TensorDB
 92            Opened database object. Use with context manager for automatic
 93            cleanup.
 94
 95        Raises
 96        ------
 97        FileNotFoundError
 98            If mode is 'r', 'r+', 'a', or 'a+' and database doesn't exist.
 99        ValueError
100            If creating new database without schema, if the schema is
101            malformed, or if mode is invalid.
102        TypeError
103            If a dtype is neither string nor torch.dtype.
104
105        Examples
106        --------
107        Creating a new database and writing dense rows:
108
109        >>> import torch
110        >>> from tensorblob import TensorDB
111        >>>
112        >>> with TensorDB.open("events.db", "w",
113        ...                    schema={"price": ("float32", 1),
114        ...                            "embed": ("float32", 768)}) as db:
115        ...     db.write({"price": torch.randn(1000, 1),
116        ...               "embed": torch.randn(1000, 768)})
117        ...     print(f"Wrote {len(db)} rows")
118        Wrote 1000 rows
119
120        Reading rows back, aligned across fields:
121
122        >>> with TensorDB.open("events.db", "r") as db:
123        ...     row = db[42]        # {"price": (1,), "embed": (768,)}
124        ...     batch = db[10:100]  # {"price": (90, 1), "embed": (90, 768)}
125
126        Notes
127        -----
128        Writes commit the row count only after all fields are written. If a
129        crash leaves fields longer than the committed count, the next open
130        reports the committed (minimum) length and, for writable modes,
131        truncates the stray rows back to restore alignment.
132        """
133        modes = set(mode)
134        if modes - set("raw+") or len(mode) > len(modes):
135            raise ValueError(f"Invalid mode: {mode}")
136        if sum(c in "raw" for c in mode) != 1 or mode.count("+") > 1:
137            raise ValueError(
138                f"Must have exactly one of read/write/append mode and at most one plus: {mode}"
139            )
140
141        filename = Path(filename).expanduser().resolve()
142        if not filename.exists():
143            if "r" in modes or "a" in modes:
144                raise FileNotFoundError(f"Database not found: {filename!r}")
145            if schema is None:
146                raise ValueError("Argument ``schema`` is required for new database!")
147            schema = cls._normalize_schema(schema)
148            return cls(os.fspath(filename), schema, block_size, mode, max_cached_blocks)
149
150        return cls.from_config(
151            save_directory=filename,
152            runtime_kwargs={
153                "mode": mode,
154                "filename": os.fspath(filename),
155                "max_cached_blocks": max_cached_blocks,
156            },
157        )

Open a TensorDB with file-like interface for multi-field tensor storage.

TensorDB provides persistent, row-aligned storage for heterogeneous (multi-field) tensor collections. Each field is stored as a plain TensorBlob in a subdirectory, and TensorDB keeps the row orders of all fields aligned. Rows are dense: every write must supply every field with the same row count.

The database is stored as a directory containing:

.conf: Schema file (field names, dtypes, shapes)
.stat: State file (committed row count)
Field-named subdirectories: One TensorBlob per field

Parameters

filename (str or Path): Directory path for database storage. Supports tilde expansion (~) and relative paths.
mode (str, default="r"): File access mode ('r', 'w', 'a', 'r+', 'w+', 'a+'). Behaves like TensorBlob.open(); the mode is applied to every field.
schema (dict, optional): Mapping of field names to (dtype, shape) pairs, e.g., {"price": ("float32", 1), "embed": (torch.float16, (768,))}. Required for new databases (modes 'w', 'w+'). Fixed at creation; loaded automatically when opening existing databases.
block_size (int, default=8192): Number of rows per memory-mapped block file, applied to all fields.
max_cached_blocks (int, optional): Maximum number of memory-mapped blocks to keep cached per field. If None (default), uses 1/16 of system's max_map_count limit.

Returns

TensorDB: Opened database object. Use with context manager for automatic cleanup.

Raises

FileNotFoundError: If mode is 'r', 'r+', 'a', or 'a+' and database doesn't exist.
ValueError: If creating new database without schema, if the schema is malformed, or if mode is invalid.
TypeError: If a dtype is neither string nor torch.dtype.

Examples

Creating a new database and writing dense rows:

>>> import torch
>>> from tensorblob import TensorDB
>>>
>>> with TensorDB.open("events.db", "w",
...                    schema={"price": ("float32", 1),
...                            "embed": ("float32", 768)}) as db:
...     db.write({"price": torch.randn(1000, 1),
...               "embed": torch.randn(1000, 768)})
...     print(f"Wrote {len(db)} rows")
Wrote 1000 rows

Reading rows back, aligned across fields:

>>> with TensorDB.open("events.db", "r") as db:
...     row = db[42]        # {"price": (1,), "embed": (768,)}
...     batch = db[10:100]  # {"price": (90, 1), "embed": (90, 768)}

Notes

Writes commit the row count only after all fields are written. If a crash leaves fields longer than the committed count, the next open reports the committed (minimum) length and, for writable modes, truncates the stray rows back to restore alignment.

@classmethod

def unlink(cls, filename): View Source

159    @classmethod
160    def unlink(cls, filename):
161        filename = Path(filename).expanduser().resolve()
162        if filename.exists():
163            try:
164                shutil.rmtree(filename)
165            except OSError as exc:
166                warnings.warn(f"Failed to unlink database at {filename!r}: {exc}")
167                return False
168        return True

@classmethod

def apply_param_hooks(cls, jdict): View Source

198    @classmethod
199    def apply_param_hooks(cls, jdict):
200        jdict["schema"] = {
201            name: (dtype, tuple(shape))
202            for name, (dtype, shape) in jdict["schema"].items()
203        }
204        return jdict

Apply post-processing hooks to the JSON dictionary.

By default, this method returns the input dictionary unchanged.

Parameters

jdict (dict[str, Any]): The config dictionary after deserialization.

Returns

dict[str, Any]: The config dictionary after post-processing.

filename

schema

block_size

mode

max_cached_blocks

configpath: str View Source

259    @property
260    def configpath(self) -> str:
261        return os.path.join(self.filename, self.config_name)

statuspath: str View Source

263    @property
264    def statuspath(self) -> str:
265        return os.path.join(self.filename, self.status_name)

closed: bool View Source

267    @property
268    def closed(self) -> bool:
269        return self._closed

def tell(self) -> int: View Source

337    def tell(self) -> int:
338        self._checkclosed()
339        return next(iter(self._cols.values())).tell()

def seek(self, pos: int = 0, whence: int = 0) -> int: View Source

341    def seek(self, pos: int = 0, whence: int = io.SEEK_SET) -> int:
342        self._checkclosed()
343        for col in self._cols.values():
344            col.seek(pos, whence)
345        return self.tell()

def close(self) -> None: View Source

347    def close(self) -> None:
348        if self._closed:
349            return
350        for col in self._cols.values():
351            col.close()
352        if self._m_wr:
353            self._status.dump(self.statuspath)
354        self._closed = True

def flush(self) -> None: View Source

356    def flush(self) -> None:
357        self._checkwritable()
358        for col in self._cols.values():
359            col.flush()
360        self._status.dump(self.statuspath)

def read(self, size: int | None = None) -> dict[str, torch.Tensor]: View Source

362    def read(self, size: int | None = None) -> dict[str, torch.Tensor]:
363        self._checkreadable()
364        # Clamp at the committed row count so uncommitted trailing rows left by
365        # an interrupted write are never visible.
366        remaining = len(self) - self.tell()
367        size = remaining if size is None else min(size, remaining)
368        if size <= 0:
369            return {
370                name: torch.empty(0, *shape, dtype=getattr(torch, dtype))
371                for name, (dtype, shape) in self.schema.items()
372            }
373        return {name: col.read(size) for name, col in self._cols.items()}

def write(self, rows: dict[str, torch.Tensor]) -> int: View Source

375    def write(self, rows: dict[str, torch.Tensor]) -> int:
376        self._checkwritable()
377        if not isinstance(rows, dict):
378            raise TypeError(
379                f"Rows must be a dict mapping field names to tensors, got {type(rows).__name__!r}!"
380            )
381        missing = sorted(self.schema.keys() - rows.keys())
382        extra = sorted(rows.keys() - self.schema.keys())
383        if missing or extra:
384            raise ValueError(
385                "Dense writes require exactly the schema fields; "
386                f"missing: {missing!r}, unexpected: {extra!r}"
387            )
388
389        nts = {
390            name: ts.view(-1, *self.schema[name][1]).size(0)
391            for name, ts in rows.items()
392        }
393        if len(set(nts.values())) != 1:
394            raise ValueError(f"All fields must have the same row count; got: {nts!r}")
395        nt = next(iter(nts.values()))
396
397        # Columns are written first and the committed row count is bumped only
398        # afterwards, so an interrupted write is rolled back on the next open.
399        for name, col in self._cols.items():
400            col.write(rows[name])
401        self._status.len = len(next(iter(self._cols.values())))
402        return nt

def truncate(self, pos: int | None = None) -> int: View Source

404    def truncate(self, pos: int | None = None) -> int:
405        self._checkwritable()
406        for col in self._cols.values():
407            col.truncate(pos)
408        self._status.len = self.tell()
409        self._status.dump(self.statuspath)
410        return self.tell()

def extend( self, other: TensorDB, maintain_order: bool = False) -> None: View Source

412    def extend(self, other: TensorDB, maintain_order: bool = False) -> None:
413        if set(self.schema) != set(other.schema):
414            raise ValueError("Schema fields must match to extend databases!")
415        self._checkwritable()
416        for name, col in self._cols.items():
417            col.extend(other._cols[name], maintain_order=maintain_order)
418        self._status.len = len(next(iter(self._cols.values())))
419        self._status.dump(self.statuspath)