Azure Data Lake Storage 上のファイルを Python でメモリ節約して読み込む

Python で Azure DataLake 上の巨大なファイルを読み込んで処理する機会があり、メモリを圧迫しない方法をとることが求められました。 Azure SDK に少しクセがありどうすれば可能なのか悩みましたが、いろいろ試して方法が見つかったため纏めます。

## TL;DR
Python のストリーム（ファイルオブジェクト）を用いると、データを一度に読み込むのを避けて少しずつ処理でき、メモリ消費を抑えられます。

あとは DataLake に対するストリームをどう取得するかというだけですが、

* Azure SDK のクラス `StorageStreamDownloader` をそのままストリームとして扱うことはできません
* そのクラスの簡単なラッパーを自作するとストリームとして使えます
* Azure SDK でなく fsspec を用いるともっと簡単です

## サンプルアプリケーション
DataLake 上に保存してあるテキストファイルを読み込み、指定した語句を含む行を抽出して返す Azure Functions を作ってみました。

最初に何の工夫も無い方法を示します。 DataLake のファイルクライアントオブジェクトを取得したら、中身のデータをバイトオブジェクトとして取得し、文字列に変換し[^sdk-convert]、1行ずつに分解して処理します。

[^sdk-convert]: SDK 側で文字列に変換させるオプションもありますが、本筋ではないので割愛します。

```python:function_app.py
import os
from logging import getLogger

import azure.functions as func
from azure.storage.filedatalake import DataLakeServiceClient

logger = getLogger(__name__)
logger.setLevel("DEBUG")

app = func.FunctionApp(http_auth_level=func.AuthLevel.ANONYMOUS)


@app.route(route="http_trigger")
def http_trigger(req: func.HttpRequest) -> func.HttpResponse:

    # 検索キーワードをリクエストパラメータから取得する
    keyword = req.params.get("s", "")
    limit = int(req.params.get("limit", 10))
    logger.info(f"keyword: {keyword!r}, limit: {limit}")

    # ファイルクライアントを取得する
    conn_str = os.environ["AZURE_STORAGE_CONNECTION_STRING"]
    file_system = os.environ["DATALAKE_FILE_SYSTEM"]
    file_path = os.environ["DATALAKE_FILE_PATH"]
    service_client = DataLakeServiceClient.from_connection_string(conn_str)
    file_client = service_client.get_file_client(file_system, file_path)

    # バイナリデータをメモリ上にダウンロードする
    downloader = file_client.download_file()
    data = downloader.readall()

    # テキストとして処理する
    lines = str(data, encoding="utf-8").splitlines(keepends=True)
    filtered = []
    n = 0
    for line in lines:
        if keyword in line:
            logger.debug(line.rstrip("\n"))
            filtered.append(line)
            n += 1
            if n >= limit:
                break

    return func.HttpResponse("".join(filtered))
```

小惑星リストのテキストファイル（74万行、約 63MiB）を配置したときの実行例を以下に示します。

```text:実行例
$ curl http://localhost:7071/api/http_trigger?s=COIAS
(697402) Ao                  2017 BX232  2017 01 23* Maunakea                 COIAS
(718492)                     2017 FZ233  2017 03 22* Maunakea                 COIAS
(719612)                     2019 UW157  2019 10 27* Maunakea                 COIAS
```

（小惑星 Ao, Quro, Hoshizaki の命名、おめでとうございます！）

## 読み込み方法の変更
サンプルアプリケーションを段階的に改善していきます。

### ファイルオブジェクト（ストリーム）を用いる
[ファイルオブジェクト](https://docs.python.org/ja/3/glossary.html#term-file-object)はデータの読み書きをする際によく使います。一連のデータを順番に処理していけばいいときに便利で、大抵はメモリ上にデータの一部だけを保持してメモリを圧迫しないようになっています。
名前通りファイルが一番イメージしやすいですが、同じインタフェースになっていればファイル以外のデータも同じように扱えます。本記事ではローカルファイルは扱わないので、ストリームという呼称を用います。

Azure SDK の `StorageStreamDownloader` には `readinto` というメソッドがあり、データを他のストリームへ書き出すことができます。ローカルファイルへ書き出さずメモリ上にデータを保持するためには、 `BytesIO` のインスタンスを作成します。

バイナリから文字列への変換や、抽出した文字列の結合などもストリームへ置き換えれば、以下のようになります。ストリームのまま for 文で1行ずつ取り出せるというのがなかなか便利です。

```python:function_app.py（抜粋）
@app.route(route="http_trigger")
def http_trigger(req: func.HttpRequest) -> func.HttpResponse:
    ...
    file_client = service_client.get_file_client(file_system, file_path)

    # バイナリデータをメモリ上にダウンロードする
    buffer = BytesIO()
    downloader = file_client.download_file()
    downloader.readinto(buffer)
    buffer.seek(0)

    # テキストとして処理する
    stream = TextIOWrapper(buffer, encoding="utf-8")
    filtered = StringIO()
    n = 0
    for line in stream:
        if keyword in line:
            logger.debug(line.rstrip("\n"))
            filtered.write(line)
            n += 1
            if n >= limit:
                break

    return func.HttpResponse(filtered.getvalue())
```

しかしこの例では**全データをメモリに乗せてから後続の処理をしている**ため、根本的な問題は解決していません。 for 文のところで適宜データを追加ダウンロードしてもらうようにする必要があります。

### StorageStreamDownloader 自体をストリームとして扱う（失敗）
以前 AWS S3 のデータを Boto3 で読み込んだときは[^prev-article]、レスポンスから取得した [`StreamingBody`](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/response.html) は `read` するだけでなくそのままストリームとして扱えました。 Azure SDK の [`StorageStreamDownloader`](https://learn.microsoft.com/ja-jp/python/api/azure-storage-file-datalake/azure.storage.filedatalake.storagestreamdownloader) も同じように扱えれば楽です。

しかし、実際に `stream = TextIOWrapper(downloader)` と試してみたところ、エラーになってしまいました。

```text
AttributeError: 'StorageStreamDownloader' object has no attribute 'readable'
```

[^prev-article]: 記事「[pandas を用いて JSON Lines 形式で S3 に保存する](d60359df808c12746d0f)」

### ラッパーを自作する
`StorageStreamDownloader` クラスは、 `read` など各ストリームで実装すべきメソッドは実装しているものの、 `IOBase` を継承していません。そのためストリームのインタフェース仕様に沿っていなく、ファイルなどと同様には扱えません。

そこで、ラッパークラスを作成してストリームに適合させてみます。

```python:download_wrapper.py
from io import IOBase

from azure.storage.filedatalake import StorageStreamDownloader


# IOBase の実装クラスとしてラッパーを作成する
class DownloadWrapper(IOBase):

    def __init__(self, downloader: StorageStreamDownloader) -> None:
        self.downloader = downloader

    # このストリームが読み取り可能かどうか伝える
    def readable(self) -> bool:
        return True

    # 実際の処理は StorageStreamDownloader に任せる
    def read(self, size: int = -1) -> bytes:
        return self.downloader.read(size)

    # その他、必要なメソッドがあれば追加する
```

```python:function_app.py（抜粋）
@app.route(route="http_trigger")
def http_trigger(req: func.HttpRequest) -> func.HttpResponse:
    ...
    file_client = service_client.get_file_client(file_system, file_path)

    # StorageStreamDownloader をラップして IO に変える
    downloader = file_client.download_file()
    buffer = DownloadWrapper(downloader)

    # テキストとして処理する
    stream = TextIOWrapper(buffer, encoding="utf-8")
    ...
```

実際に動かしてみると、 Azure SDK による DataLake からのダウンロードの合間に for 文内の処理がログ出力されていて、**データが必要になってからダウンロードしている**ことが確認できます。
（あとはメモリが増大していないことの確認も必要ですが、 1GiB 近いファイルを全行処理しても Azure Functions のメモリが数百MBで済んでいることは確認しました。）

### fsspec を利用する
pandas のドキュメントを見ていると、パスにローカルファイルでなく `s3://...` のような URI も指定できることが書かれています。
[IO tools (text, CSV, HDF5, …) # Reading/writing remote files](https://pandas.pydata.org/pandas-docs/stable/user_guide/io.html#reading-writing-remote-files)

これは [fsspec](https://filesystem-spec.readthedocs.io/en/latest/) というライブラリを利用しているようです。　Microsoft の公式ドキュメントにも例が無いか探したところ、データ分析サービスのほうで見つかりました。（リンクは[参考](#参考)に記載しています）

fsspec 本体に加えて Azure 向けの実装が必要なので、 `pip install adlfs` でインストールして利用します。コードではファイルクライアントの取得は必要なくなり、組み込みの `open` に似た形でストリームを取得できます。

```python:function_app.py
import os
from io import StringIO, TextIOWrapper
from logging import getLogger

import azure.functions as func
import fsspec

logger = getLogger(__name__)
logger.setLevel("DEBUG")

app = func.FunctionApp(http_auth_level=func.AuthLevel.ANONYMOUS)


@app.route(route="http_trigger")
def http_trigger(req: func.HttpRequest) -> func.HttpResponse:

    # 検索キーワードをリクエストパラメータから取得する
    keyword = req.params.get("s", "")
    limit = int(req.params.get("limit", 10))
    logger.info(f"keyword: {keyword!r}, limit: {limit}")

    # リモートファイルの IO を取得する
    # （認証情報は環境変数 AZURE_STORAGE_CONNECTION_STRING に格納済み）
    file_system = os.environ["DATALAKE_FILE_SYSTEM"]
    file_path = os.environ["DATALAKE_FILE_PATH"].lstrip("/")

    with fsspec.open(f"az://{file_system}/{file_path}") as f:

        # テキストとして処理する
        stream = TextIOWrapper(f)
        filtered = StringIO()
        n = 0
        for line in stream:
            if keyword in line:
                logger.debug(line.rstrip("\n"))
                filtered.write(line)
                n += 1
                if n >= limit:
                    break

    return func.HttpResponse(filtered.getvalue())
```

## まとめ
Azure DataLake 上のファイルを Python でメモリ節約して読み込む方法を調査しました。

データを少しずつ読み込むにはストリームを使うのが便利です。しかし Azure SDK のクラスは残念ながら `IOBase` を継承していなく、そのままではストリームとして扱えませんでした。とはいえ簡単なラッパーを作成することでその問題を解消できました。

ストリームを取得する別の方法では fsspec がありました。 URI を指定してローカルファイルのように扱うので、 Azure SDK とは雰囲気が異なります。今回は深入りしていませんが、データ処理などのライブラリによっては URI を直接渡したりワイルドカード指定もできるようなので、ファイルの扱い全般が楽になりそうです。

## 参考

* Microsoft Learn
	* Azure SDK for Python 関連
		* [Python を使用して BLOB をダウンロードする](https://learn.microsoft.com/ja-jp/azure/storage/blobs/storage-blob-download-python)
		* [Python を使用して Azure Data Lake Storage でディレクトリとファイルを管理する](https://learn.microsoft.com/ja-jp/azure/storage/blobs/data-lake-storage-directory-file-acl-python)
		* [StorageStreamDownloader クラス](https://learn.microsoft.com/ja-jp/python/api/azure-storage-file-datalake/azure.storage.filedatalake.storagestreamdownloader) ※いくつかのモジュールに同名で存在します
	* fsspec 関連
		* [チュートリアル: FSSPEC を使用して Synapse Analytics のサーバーレス Apache Spark プール内の ADLS データの読み取りと書き込みを行う](https://learn.microsoft.com/ja-jp/azure/synapse-analytics/spark/tutorial-spark-pool-filesystem-spec)
		* [チュートリアル: Synapse Analytics のサーバーレス Apache Spark プールで、Pandas を使用した Azure Data Lake Storage Gen2 データの読み取り/書き込みを行う](https://learn.microsoft.com/ja-jp/azure/synapse-analytics/spark/tutorial-use-pandas-spark-pool)
	* Azure Functions
		* [クイックスタート: Visual Studio Code と Python を使用して Azure に関数を作成する](https://learn.microsoft.com/ja-jp/azure/azure-functions/create-first-function-vs-code-python)
* Python / ライブラリ
	* [io --- ストリームを扱うコアツール](https://docs.python.org/ja/3/library/io.html)
	* [`fsspec`: Filesystem interfaces for Python](https://filesystem-spec.readthedocs.io/en/latest/)
	* [Filesystem interface to Azure-Datalake Gen1 and Gen2 Storage](https://github.com/fsspec/adlfs)

小惑星リストのテキストファイルは、 IAU Minor Planet Center の `NumberedMPs.txt` （2024/11/14更新版）を使用しました。
https://www.minorplanetcenter.net/data