【JavaEE】大量データを扱うJPA、どうにかしたくてカーソルを試してみた話

初めての記事投稿です。
至らない点が多々あると思いますがご容赦ください。

結論

• 実行速度を改善したいなら、IN句の処理を Java側で行いクエリ結果はカーソルを使って取得する
• 実行速度がどんなに遅くてもメモリ消費の抑制を最優先にするなら、IN句付きクエリを発行する

環境

• Java 21
• JPA の実装：Eclipcelink 4.0.4
• APサーバー：WIldFly 32.0.1 Final

サンプルプロジェクト

今回使用したプロジェクトをGitHubに公開しています。
詳しい実行方法などはこちらのReadmeを参照してください。

背景

ある時、JPAを使ったDBアクセスのレスポンス改善に取り組む機会がありました。

主な対応内容は、「IN句に当たる処理をJava側で行う」というものです。

当時担当していたシステムは一度に数千件のIDに対応するのデータを取り扱うことがあり、IN句のサイズが非常に大きくなってしまうことがありました。

調査したところ、(少なくとも)Eclipcelinkを利用してDBからデータを取得する時は、IN句を使わずにデータを取得した後にJava側でフィルター処理を行う方が速いということが分かりました。（理由は深く追えていませんが。。。）

// 遅い
String in = "WHERE t.id IN :idList";
entityManager.createQuery(selectQuery + in).setParameter("idList", idList).getResultList();

//こっちのほうが速い
Set<String> idSet = new HashSet<>(idList);
entityManager.createQuery(selectQuery, TestTable.class)
    .getResultList()
    .stream()
    .filter(e -> idSet.contains(e.getId()))
    .toList();

問題点

IN句を使わずにデータを取得するということは、生成されるEntityのインスタンスの数が増えるということ。
メモリ消費が気になります。

特に問題になるのが、一次キャッシュ(後述)です。
こいつの生存期間はトランザクション中であるため、トランザクションが終了するまでは該当のEntityを保持し続けます。
最初はトランザクション範囲を狭めることで対応していましたが、取得件数が極端に多い場合、そのトランザクション範囲内だけでメモリを使い果たしてしまうことも考えられました。

何か対処方法は無いかと調査をしていたところ、カーソル(Cursor) を取得する方法にたどり着きました。

カーソルとは

カーソルとは、結果セットの行を一行ずつ巡回するためのデータです。
「ポインタ」と表現してもよいかもしれません。

一括で結果を取得するのではなく、1 件ずつ順番に処理することで、メモリのオーバーフローを防ぐことができるみたいです。
また、このような仕組みを指して「カーソル」と呼ばれたりもするらしいです。

JPAやJDBCでも、カーソル型の取得方式が提供されています。

一次キャッシュとは

こちらも軽く説明します。

一次キャッシュとは、Entity Managerに紐づいたキャッシュです。

Entity Managerを通じてDBから取得したEntity(DBのテーブルに対応するオブジェクト)は、その後Entity Managerによって管理されるようになります。
このEntityの集合は、「Persistence Context(永続性コンテキスト)」と呼ばれます。

Persistence Contextは、Entityのライフサイクルや変更状態を管理する他、Entityクラスのインスタンスの一時的な保持(キャッシュ)も行います。

このキャッシュが、「一次キャッシュ」と呼ばれるものです。

「L1キャッシュ」や「first-level cache」とも書かれます。
Eclipselinkでは「独立キャッシュ」とも表現されます。

一次キャッシュは、Entity Managerのスコープ内で再利用されます。
そのため、Entity Managerが生きている間は一次キャッシュも生き続けます。

戦略

JPAのEntityManagerインターフェースにはdetachというメソッドがあり、引数のEntityを一次キャッシュから明示的に削除することができます。
これを使い、カーソルで 1 件ずつ取得 → EntityManagerからdetachしてキャッシュから除外という実装にすることで、大量データ取得時の一次キャッシュの肥大化を防げると考えました。

while (cursor.hasNext()) {
    //1件ずつ取得
    T entity = (T) cursor.next();
    result.add(entity);

    // detach
    entityManager.detach(entity);
}

サンプルコード

以下を行う処理を実装してみました。
・クエリ結果をカーソルとして取得し、
・結果の取得後にEntityManager#detatchを実行する

CacheSafeEntityLoader.java

public class CacheSafeEntityLoader {

    private static final int CURSOR_CACHE_CLEAR_THRESHOLD = 500;

    public static <T> List<T> getUsingCursor(EntityManager entityManager, TypedQuery<T> query) {
        query.setHint(QueryHints.CURSOR, true);

        CursoredStream cursor = (CursoredStream) query.getSingleResult();

        List<T> result = new ArrayList<>();
        Boolean isEntity = null;
        int processedRow = 0;
        try {
            while (cursor.hasNext()) {
                @SuppressWarnings("unchecked")
                T entity = (T) cursor.next();

                result.add(entity);

                // TupleやIntegerなどのEntityクラスではないオブジェクトをdetachに渡すとIllegalArgumentExceptionとなるため事前にチェック
                if (Objects.nonNull(isEntity) ? isEntity : (isEntity = hasEntityAnnotation(entity.getClass()))) {
                    entityManager.detach(entity);
                }

                if(++processedRow % CURSOR_CACHE_CLEAR_THRESHOLD == 0) {
                    cursor.clear();
                    processedRow = 0;
                }
            }
        }
        finally {
            cursor.close();
        }

        return result;
    }

    private static boolean hasEntityAnnotation(Class<?> clazz) {
        Class<?> _clazz = clazz;
        while (_clazz != null && _clazz != Object.class) {
            // Entityアノテーションの有無で判断する
            if (_clazz.getAnnotation(Entity.class) != null) {
                return true;
            }
            _clazz = _clazz.getSuperclass();
        }
        return false;
    }

}

この処理は引数のentityManagerとqueryに対して破壊的です。
特にqueryについては、呼び出し後にそのまま別の場所で使用する時は注意してください。

テストデータ

以下の構造のテーブルにランダムなデータを100万件格納しています。

CREATE TABLE TEST_TABLE (
    ID CHAR(36) NOT NULL
    , NAME VARCHAR(200) NOT NULL
    , PRIMARY KEY (ID)
);

実行テスト

まずは、以下をそれぞれ実装したRepositoryクラスを作成しました。

• IN 句を使わずに取得し、IN 句に当たる処理を Java 側で行う処理（find）
• IN 句を含むクエリを発行する処理（findUsingIn）
• Cursor を使って結果を取得する処理（findUsingCursor）

加えて、テスト用のIDを用意するために、全件取得を行うgetメソッドも実装しています。

Repository.java

@Stateless
// 各メソッドでトランザクションを分割するために、TransactionAttributeTypeをREQUIRES_NEWに設定
@TransactionAttribute(TransactionAttributeType.REQUIRES_NEW)
public class Repository {

    @PersistenceContext(unitName = "TEST")
    private EntityManager entityManager;

    private static final String selectQuery = """
            SELECT t
            FROM TestTable t
            """;

    // TestTable内の全てのデータを取得する。
    public List<TestDto> get() {
        TypedQuery<TestTable> typedQuery = this.entityManager.createQuery(selectQuery, TestTable.class);

        List<TestTable> resultList = typedQuery.getResultList();

        return convert(resultList);
    }

    // TestTable内の全てのデータを取得する。
    // 全件取得後、Stream#filterを利用してidListでの絞り込みを行う。
    public List<TestDto> find(List<String> idList) {
        TypedQuery<TestTable> typedQuery = this.entityManager.createQuery(selectQuery, TestTable.class);

        Set<String> idSet = new HashSet<>(idList);

        List<TestTable> reusltList = typedQuery.getResultList()
                .stream()
                .filter(e -> idSet.contains(e.getId()))
                .toList();

        return convert(reusltList);
    }

    // TestTable内の、idListのIDを持つデータを取得する。
    // クエリにIN句を追記して取得する。
    public List<TestDto> findUsingIn(List<String> idList) {
        String in = "WHERE t.id IN :idList";

        String query = selectQuery + in;
        TypedQuery<TestTable> typedQuery = this.entityManager.createQuery(query, TestTable.class)
                .setParameter("idList", idList);

        // SQL Serverはパラメータの数が2100を超えるとエラーになるため、IN句内のパラメータ数を2000ずつに分けてクエリを実行する
        List<TestTable> reusltList = split(idList, 2000, splitIdList -> {
            typedQuery.setParameter("idList", splitIdList);
            return typedQuery.getResultList();
        });

        return convert(reusltList);
    }

    // TestTable内の全てのデータを、CacheSafeEntityLoader#getUsingCursorを利用して取得する。
    // 全件取得後、Stream#filterを利用してidListでの絞り込みを行う。
    public List<TestDto> findUsingCursor(List<String> idList) {
        TypedQuery<TestTable> typedQuery = this.entityManager.createQuery(selectQuery, TestTable.class);

        Set<String> idSet = new HashSet<>(idList);

        List<TestTable> reusltList = CacheSafeEntityLoader.getUsingCursor(this.entityManager, typedQuery)
                .stream()
                .filter(entity -> idSet.contains(entity.getId()))
                .toList();

        return convert(reusltList);
    }

    // sourceListを用いるfunctionを、sourceListの要素をunitSize個ずつに分割して実行する。
    private static <T, E> List<E> split(List<T> sourceList, int unitSize, Function<List<T>, List<E>> function) {
        int index = 0;
        int sourceSize = sourceList.size();

        List<E> result = new ArrayList<>();

        while (index < sourceSize) {
            int nextIndex = Math.min(index + unitSize, sourceSize);
            List<T> subList = sourceList.subList(index, nextIndex);
            result.addAll(function.apply(subList));
            index = nextIndex;
        }

        return result;
    }

    // TestTableをTestDtoに変換する。
    private static List<TestDto> convert(List<TestTable> entityList) {
        return entityList.stream().map(entity ->
                new TestDto(
                        entity.getId(),
                        entity.getName()
                )
        ).toList();
    }
}

そして、Repositoryクラスの各メソッドが正しく動いているかを、以下のコードを使って確認します。

Application.java

@Stateless
public class Application {

    private static final int ID_LIST_SIZE = 10000;

    @Inject
    private Repository repository;

    public void execute() {
        List<String> idList = repository.get().stream()
                .map(TestDto::getId)
                .limit(ID_LIST_SIZE)
                .toList();

        System.out.println("size of idList = " + idList.size());

        List<TestDto> findResult = repository.find(idList);
        System.out.println("size of findResult = " + findResult.size());

        List<TestDto> usingInResult = repository.findUsingIn(idList);
        System.out.println("size of usingInResult = " + usingInResult.size());

        List<TestDto> usingCursorResult = repository.findUsingCursor(idList);
        System.out.println("size of usingCursorResult = " + usingCursorResult.size());

        System.out.println("findResult == usingInResult : " + isSameList(findResult, usingInResult));
        System.out.println("findResult == usingCursorResult : " + isSameList(findResult, usingCursorResult));
    }

    private static <T> boolean isSameList(List<T> list1, List<T> list2) {
        return list1.size() == list2.size() && new HashSet<>(list1).containsAll(list2);
    }

}

それぞれが結果を正しく取得できていれば、以下のような出力を確認できます。

size of idList = 10000
size of findResult = 10000
size of usingInResult = 10000
size of usingCursorResult = 10000
findResult == usingInResult : true
findResult == usingCursorResult : true

比較

Repositoryクラスのfind、findUsingIn、findUsingCursorそれぞれの実行時間とメモリ使用量を比較してみます。
比較には「VisualVM」というツールを使います。

実行時間

各メソッドを 5 回ずつ実行し、最長と最短を除いた 3 回の平均実行時間を記録しました。
※単位は[秒]です。

• ID_LIST_SIZE = 10000

findUsingIn	find	findUsingCursor
810.369	3.307	2.193

findUsingInが突出して長い時間をかけていることが分かります。
findとfindWithCursorでは、Cursor を使った処理の方が約 1 秒短いですね。

ここから、ID_LIST_SIZEの値を増やして計測をしてみます。
※実行時間が長すぎるfindUsingInは除外しました。

• ID_LIST_SIZE = 100000

find	findUsingCursor
3.054	2.404

• ID_LIST_SIZE = 500000

find	findUsingCursor
3.212	2.701

• ID_LIST_SIZE = 1000000

find	findUsingCursor
3.234	2.952

いずれもfindWithCursorの実行時間が短いですが、取得件数が増えるにつれて差が縮まっていますね。

どこかで逆転するのかな？と思い、DBに格納するデータ数とID_LIST_SIZEを1000万にして計測してみたところ、以下のようになりました。

find	find(実質)	findUsingCursor
620.292	59.799	53.899

この時のfindですが、実行時間を見てみると以下のような分岐がありました。
※めちゃめちゃ端折って書いています

jpacursortest.infra.Repository$Proxy$_$$_Weld$EnterpriseProxy$.find () 620,292 ms
├─ org.jboss.as.ejb3.tx.CMTTxInterceptor.endTransaction ()             560,493 ms
└─ jpacursortest.infra.Repository.find ()                              59,799  ms

endTransactionはアプリケーションサーバー側にある EJB の実装と思われます。
この処理の実行時間を引いた実行時間を「find(実質)」と表記しています。
これでも、findUsingCursorの方が若干速いですね。

endTransactionの影響でfindがめちゃめちゃ遅くなっているように見えてしまっていますが、findUsingCursorの方は特に何もありませんでした。
いずれ詳しく調査したいですね・・・

なんにせ確認できた範囲に関しては、両者の実行時間はほぼ同じと見て良いでしょう。

メモリ使用量

Application#execute内でのRepositoryの呼び出し(get以外)を以下のように1つを残してコメントアウトし、アプリケーションサーバー起動後初回実行時のメモリ使用量を比較します。

Application#execute

public void execute() {
    List<String> idList = repository.get().stream()
            .map(TestDto::getId)
            .limit(ID_LIST_SIZE)
            .toList();

    System.out.println("size of idList = " + idList.size());

    List<TestDto> findResult = repository.find(idList);
    System.out.println("size of findResult = " + findResult.size());

//   List<TestDto> usingInResult = repository.findUsingIn(idList);
//   System.out.println("size of usingInResult = " + usingInResult.size());
//
//   List<TestDto> usingCursorResult = repository.findUsingCursor(idList);
//   System.out.println("size of usingCursorResult = " + usingCursorResult.size());
//
//   System.out.println("findResult == usingInResult : " + isSameList(findResult, usingInResult));
//   System.out.println("findResult == usingCursorResult : " + isSameList(findResult, usingCursorResult));
}

Repository側のメソッドではブレークポイントをreturn文に貼り、停止したタイミングでヒープダンプの取得を行います。
ヒープダンプの取得に時間がかかることがあったため、必要に応じてRepository#convertを以下のように変更しています。

Repository#convert

private List<TestDto> convert(List<TestTable> entityList) {
    try {
        Thread.sleep(20000); // スリープ (この間にヒープダンプの出力が終わればOK)
    } catch (InterruptedException e) {
        throw new RuntimeException(e);
    }
    return entityList.stream().map(entity ->
            new TestDto(
                    entity.getId(),
                    entity.getName()
            )
    ).toList();
}

結果

VisualVMでヒープダンプを取得すると、以下のような情報を確認することができます。

ここから得られる全体のヒープ使用量は、それぞれ以下の通りでした。

findUsingIn	find	findUsingCursor
84MB	712MB	512MB

findUsingIn < findWithCursor < find という結果になりました。

findUsingInは余計なことをしていないからか、他の2つと比べてメモリ消費が非常に少ないです。
findUsingCursorと findでは、前者の方が約30%少ない結果となりました。

次に、1次キャッシュの削減がちゃんとできているのかを確認します。

以下は、javaオブジェクトのインスタンスの一覧とそれぞれの個数・合計サイズを示したものです。
VisualVMでのヒープダンプの表示を「Summary」から「Objects」にすると見れるものです。

上から7行目にjpacursortest.infra.TestTableとあり、「Count」列の数字が1000001と表示されています。
SQLの実行結果の行数+1になっていますね。（この+1の正体は分かりませんが、今はスルーします。）

また、上から9行目のorg.eclipse.persistence.internal.identitymaps.CacheKeyが、一次キャッシュと呼ばれているものです。
こちらは「Count」列の数字が1000000になっています。SQL の実行結果の行数ピッタリですね。

「+」をクリックして展開していくと、「object」というフィールドにTestTableのインスタンスが格納されていることが分かります。

では、findWithCursorの「Objects」を確認してみます。

TestTableは確認できますが、CacheKeyが見当たりません。

Ctrl + fで検索してみたところ、「Count」が 0 の状態で発見されました。

狙い通り、一次キャッシュをため込まないように処理を実行できているようです。

おまけ

getWithCursorメソッドに以下のようなオーバーロードを追加してみました。

CacheSafeEntityLoader#getUsingCursor

public static <T> List<T> getUsingCursor(EntityManager entityManager, TypedQuery<T> query, Predicate<T> predicate) {
    query.setHint(QueryHints.CURSOR, true);

    CursoredStream cursor = (CursoredStream) query.getSingleResult();

    List<T> result = new ArrayList<>();
    Boolean isEntity = null;
    int processedRow = 0;
    try {
        while (cursor.hasNext()) {
            @SuppressWarnings("unchecked")
            T entity = (T) cursor.next();

            if(predicate.test(entity)) {
                result.add(entity);
            }

            if (Objects.nonNull(isEntity) ? isEntity : (isEntity = hasEntityAnnotation(entity.getClass()))) {
                entityManager.detach(entity);
            }

            if(++processedRow % CURSOR_CACHE_CLEAR_THRESHOLD == 0) {
                cursor.clear();
                processedRow = 0;
            }
        }
    }
    finally {
        cursor.close();
    }

    return result;
}

以下の点が変わっています。
・引数にPredicate<T>を追加
・result.add()前に条件分岐を追加

呼び出す側は、引数を以下のようにして実行します。

CacheSafeEntityLoader.getUsingCursor(this.entityManager, typedQuery, entity -> idSet.contains(entity.getId()))

これでIN句にあたる処理もCacheSafeEntityLoaderに任せることができ、余分な参照の削減もできると思います。

さいごに

カーソルを利用して結果を取得することで、IN句を使わなくても実行速度を守りつつメモリ使用量をある程度抑えることができました。
メモリ消費はIN句を使う方法に大きく劣りますが、実行速度を考えるとカーソルを利用する方法が最も良い選択肢のように思います。

Cursor の利用についてなにか知見や指摘があれば、ぜひコメントで教えていただけると幸いです。

参考

一次キャッシュ
https://docs.oracle.com/cd/F32751_01/toplink/14.1.1.0/concepts/understanding-caching.html
https://enterprisegeeks.hatenablog.com/entry/2014/11/24/110723
https://qiita.com/haroya01/items/e2fbb38a90a686855b06

カーソル
https://ja.wikipedia.org/wiki/%E3%82%AB%E3%83%BC%E3%82%BD%E3%83%AB_(%E3%83%87%E3%83%BC%E3%82%BF%E3%83%99%E3%83%BC%E3%82%B9)
https://w3schools.tech/ja/tutorial/sql/sql-cursors