【Python】巨大なNumPy配列を必要な範囲だけ読む

PythonのTopに戻る

結論:npyファイルをmmap_mode="r"で開く

巨大なNumPy配列の必要な範囲だけを扱いたいときは、配列をnpyファイルとして保存し、np.load(path, mmap_mode="r", allow_pickle=False)で開く方法がある。メモリマップされた配列へ基本スライスでアクセスすれば、配列全体を通常のndarrayとして先に読み込まずに範囲を参照できる。すでにメモリ上にある任意の配列を、後から部分読込へ変えるという意味ではない。

対象はディスク上のnpyである。CSVのテキスト解析や、npzのアーカイブから配列を取り出す処理とは別の仕組みなので、同じようにmmapが働くと考えないこと。大きな配列を何度も部分参照する予定なら、前処理の段階でnpyへ保存しておくと便利である。shapeやdtypeを持つ形式なので、生のバイナリを手動で解釈する負担も減る。

そのまま動かせる例

例では仕組みを小さなデータで試すため、10000行×6列のfloat64配列を一時フォルダーへ作る。約48万バイトの要素データであり、巨大ファイルを実際に用意する負荷試験ではない。open_memmapで書込用のnpyを作り、1000行ずつ値を入れる。作成時にも全行分の中間配列を一度に用意しない形にしている。

読込側は読取専用のrモードで開き、100〜103行・2〜4列を参照する。その後、基本スライスと整数配列による選択の違い、全行を小分けに合計する処理を確認する。ファイルは一時的なもので、例の終了時に削除される。NumPy以外の追加ライブラリや外部データは不要である。

from pathlib import Path
from tempfile import TemporaryDirectory
import gc
import numpy as np
from numpy.lib.format import open_memmap

with TemporaryDirectory() as folder:
    path = Path(folder) / "large.npy"
    writer = open_memmap(path, mode="w+", dtype="float64", shape=(10000, 6))
    for start in range(0, 10000, 1000):
        rows = np.arange(start, start + 1000)[:, None]
        writer[start:start + 1000] = rows + np.arange(6)
    writer.flush()
    del writer
    gc.collect()
    mapped = np.load(path, mmap_mode="r", allow_pickle=False)
    block = mapped[100:104, 2:5]
    np.testing.assert_array_equal(block, [[102, 103, 104], [103, 104, 105],
                                         [104, 105, 106], [105, 106, 107]])
    assert isinstance(mapped, np.memmap)
    assert np.shares_memory(mapped, block)
    picked = mapped[[100, 102], :]
    assert not np.shares_memory(mapped, picked)
    total = 0.0
    for start in range(0, len(mapped), 1000):
        total += float(mapped[start:start + 1000].sum())
    assert total == 300120000.0
    print("shape:", mapped.shape, "dtype:", str(mapped.dtype))
    print("block:", block.tolist())
    print("basic slice shares memory:", bool(np.shares_memory(mapped, block)))
    print("advanced indexing shares memory:", bool(np.shares_memory(mapped, picked)))
    print("full scan sum:", total)
    del block, picked, mapped
    gc.collect()

実行結果

shape: (10000, 6) dtype: float64
block: [[102.0, 103.0, 104.0], [103.0, 104.0, 105.0], [104.0, 105.0, 106.0], [105.0, 106.0, 107.0]]
basic slice shares memory: True
advanced indexing shares memory: False
full scan sum: 300120000.0

ビューとコピーの違いを確かめる

出力では基本スライスが元のマップとメモリを共有し、整数配列による高度なインデックス指定は共有しない。後者は選んだデータのコピーを作るため、対象が大きいと追加のメモリを使う。型の名前だけから判断するより、どの操作がビューでどの操作がコピーかを理解する方が重要である。小さなコピーを明示的に作って下流へ渡す方が扱いやすい場合もある。

mapped + 1のような全体演算、全要素へのastype、並べ替えや大きなブール選択などを続ければ、大きな中間配列を作ったり全ファイルへアクセスしたりする。入口をmmapにしただけで、後の計算が全て省メモリになるわけではない。処理を必要な範囲へ絞り、各段階のshapeとコピー発生箇所を確認しよう。

小分け集計も全データを読む

例の合計は1000行ずつ計算するので大きな結果配列を保持しないが、最終的には全行へアクセスする。部分読込の例と全件集計の例は目的が違う。OSはページ単位でファイルを読み込み、先読みやキャッシュも使うため、指定した要素のバイト数だけが厳密にディスクから読まれるとは限らない。メモリ使用量や速度はアクセス順序とOSの状態にも左右される。

C順の2次元配列なら、行方向に連続した範囲を読む方が、離れた要素を飛び飛びに読むより入出力の局所性を保ちやすい。列だけを繰り返し読む用途なら、保存時の配置やデータ分割も設計の対象となる。実際の配列サイズ、ストレージ、同時処理数で測定し、例の小さなデータの速度を巨大データへそのまま当てはめないようにしたい。

モードとファイルの寿命に注意

rは読取専用なので、解析中に元のファイルを書き換える事故を避けやすい。r+はファイルへ変更を書き込む用途で、気軽に置き換えるべき設定ではない。書込側のflushは変更内容の反映に使うが、オブジェクトを閉じる操作そのものではない。別プロセスが読んでいる途中でファイルを上書き・切り詰めることも避ける。

NumPyのmemmapには、共有ビューを含めて安全に扱える一般的な公開close APIが用意されていない。例では参照を削除して回収を促してから一時フォルダーを終了しており、Linux環境で実行確認した。ファイルの削除可否はOSにも依存するため、Windows等でも同じ寿命管理を実際に確認してほしい。厳密に資源寿命を分離したい処理では、読込を独立したプロセスにまとめる設計も選択肢となる。

確認環境と参考資料

例はLinux・CPython 3.12.14・NumPy 2.3.5で実行した。掲載した出力はこの環境での結果である。公式資料のstable版や最新版は更新されるため、手元のバージョンと対応する仕様も確認してほしい。

関連項目:測定配列を形と型を保って保存する / tracemallocでメモリが増えた場所を調べる

PythonのTopに戻る