WormholeというSorted MapをJavaで実装して性能改善した話

数ヶ月前に B-Trees Are Back: Engineering Fast and Pageable Node Layouts という論文を見かけまして、その中で

MassTree is a trie with a span of 8 bytes, where each
trie node is a B-Tree [40]. We do not compare against MassTree, as
it is significantly slower than Wormhole [55].

という記載があることに気が付きました。Masstreeといえばあの Silo: Speedy transactions in multicore in-memory databases
の屋台骨として採用された高速in-memoryなindex向けデータ構造で、以前実装してみようかと思いつつ面倒くさくて着手していなかったやつでした。それよりも速そうなWormholeというものに興味が湧いたので、Javaで実装してみました。

（最新バージョンはv0.2.0）

また、初回リリース時に https://medium.com/@komamitsu/wormhole4j-a-fast-ordered-map-for-java-724dc7f88e6e という記事を書いた際に「Wormholeとは何か？」について簡単に説明したので以下に引用しておきます（Claudeによる翻訳）

高速なインメモリデータシステムを構築する際、データ取得の速度は非常に重要です。ハッシュテーブルは単一項目の検索に最適なO(1)の時間計算量を提供します。しかし、ハッシュテーブルは「apple」から「banana」までのすべてのキーを取得するような範囲クエリを自然にサポートしていません。そのため、B+木やスキップリストのような順序付きインデックスが一般的に使用されますが、これらは通常O(log N)のルックアップコストを伴い、スケール時に高コストになる可能性があります。
トライ木は、ルックアップ時間がキー数(N)ではなくキー長(L)に依存する代替手段を提供します。これは特定のワークロードに役立ちますが、トライ木は多くのメモリを消費し、断片化されたレイアウトに悩まされることが多く、キーが中程度または著しく長い実世界のアプリケーションでは効率が低下します。
Wormholeは、2019年にXingbo Wu、Fan Ni、Song Jiangによって、研究論文「Wormhole: A Fast Ordered Index for In-memory Data Management」で紹介されました。ハッシュテーブルと従来の順序付きインデックスの間のギャップを埋めることを目的としています。ハッシュテーブルに近いルックアップ性能を提供しながら、範囲検索やプレフィックススキャンなどの順序付き操作もサポートします。

データ構造ですが、概念的には以下のようになっています。

実際には、以下のようにTrieの部分が物理的にはHashTableに格納されています。

で、ここからが本題なのですが、最近以下の対応を進める中で、抽象化を壊していったり泥臭いコードに置き換えていったり等、面白い学びがあったので本記事で紹介したいと思います。

• 数値型キーの対応（v0.1では文字列キーのみ対応していた）
  • 論理的にはTrie木として扱うため、先頭のトークン（e.g., 文字列キーの場合、文字）から辿れる必要があるります。そのため、数値そのままでは扱えません
• 全体的な性能改善

数値型キーの対応（手抜きで数値キー→16進数String、文字列キーは元のStringのまま）

Wormhole4jのv0.1では文字列キーのみ対応していたのは前述のとおりです。数値キーも対応させたいなぁと思っていましたがTrie的にキーを探索できるようにする必要があり、あまりやる気が出なかったので数値を16進数表現した文字列に変換する手抜き実装を試してみました（https://github.com/komamitsu/wormhole4j/tree/d7911a85d6c78cf42422de77a2925285eec3e5d6）。

public class WormholeForIntKey<V> extends Wormhole<Integer, V> {

  ...
  
  @Override
  String encodeKey(Integer key) {
    // Encode an int into a lexicographically ordered hexadecimal string.
    return String.format("%08x", 0x80000000 ^ key); // 42 -> "8000002a"
                                    // Integer.MIN_VALUE  -> "00000000"
                                    // Integer.MAX_VALUE  -> "ffffffff"
  }
}

気をつけないといけないのは辞書順で比較できるように最上位ビットをいじることくらいでしょうか。この手法は簡単に対応できる一方、単に数値→文字列変換のコストのみが増えるので、性能は文字列キーに比べて有意に下がったと記憶しています。java.util.TreeMapで数値キーを扱う場合に比べると1/3程度の性能しか出ていなかったと記憶しています。

数値型キーの対応（数値キー/文字列キー→byte[]に変換してラップ）

Wormhole4jでは LeafNode インスタンス内に128個の KeyValue インスタンスを保持・管理していました。v0.1ではKeyValueの内部で文字列キーをそのまま保持してTrieのキーとしても使っていましたが、数値キーも対応させようとすると、Trieとして扱えるよう元の数値とは別にエンコードされたキー型が欲しくなりました。

ここでWormhole特有の事情（Wormhole4jの事情も含む？）を以下に記しておきます：

• Wormholeはキーの探索でTrieを使っているのでString, byte[]と相性が良さそう（Stringはv0.1の時点で対応済み）
  • ちなみに、数値からbyte[]への変換はByteBuffer経由で可能
• エンコードされたキー型はHashMap的にO(1)でlookupできる必要がある
  • このHashMapは独自に作り込みたくないのでjava.util.HashMap<K, V>を使いたい
  • そうすると、byte[]はキーに使えない（equals(), hashCode()`がデフォルトのままの同一性ベースなので）

なので、シンプルにbyte[]をラップしてequals(), hashCode()を実装したEncodedKeyクラスを作ってエンコードされたキー型として扱うようにしてみました（https://github.com/komamitsu/wormhole4j/tree/512c9c09ac5f78a120e5a39971b9d06491c2bba8）。

class EncodedKey implements Comparable<EncodedKey> {
  private final byte[] content;

  @Override
  public boolean equals(Object o) {
    if (this == o) return true;
    if (o == null || getClass() != o.getClass()) return false;
    EncodedKey that = (EncodedKey) o;
    return Arrays.equals(content, that.content);
  }

  @Override
  public int hashCode() {
    // return Arrays.hashCode(content);
    return (short) (0x7FFF & LongHashFunction.xx3().hashBytes(content));
  }

  ...

で、数値からbyte[]を取得してEncodedKeyインスタンスを生成。

public class WormholeForIntKey<V> extends WormholeBase<Integer, V> {
  ...

  @Override
  EncodedKey encodeKey(Integer key) {
    ByteBuffer buf = ByteBuffer.allocate(4);
    buf.putInt(0x80000000 ^ key);
    return new EncodedKey(buf.array());
  }
}

性能はというと、流石に手抜きの数値→16進数文字列変換よりは性能は改善したものの、java.util.TreeMap等と比較すると有意に性能は低かったです。おそらく、数値→byte[]でのインスタンス生成、EncodedKeyインスタンス生成が影響していると推測していました（特にEncodedKeyインスタンスは一時的な用途で結構生成されるため塵も積もれば的に性能に影響しそう）。

また、文字列キーもStringからbyte[]をzero-copyで引っこ抜いてEncodedKeyとして統一的に扱えるようにしたのですが、String そのものをキーとして使っていたv0.1に比べるとやはり有意に性能が下がりました。こちらもEncodedKeyインスタンス生成が影響したのかと推測しています。

ちなみに

Stringからbyte[]をzero-copyで引っこ抜いて

という話については

で簡単にまとめてあります。

数値型キーの対応（数値キー→byte[]に変換してラップ、文字列キーは元のStringのまま）、すなわち抽象化の破壊

一つ前の手法で数値型キーの性能がv0.1に比べてある程度向上したのですが、文字列キーの性能が下がってしまうのは致命的な問題でした。そこで文字列キーの場合にはEncodedKeyという抽象化を挟まないようにしてみました。具体的にはエンコードされたキー型として泣く泣くObjectを使い、どちらの形式でエンコードされているのか
をenum EncodedKeyTypeで保持するようにしてみました（https://github.com/komamitsu/wormhole4j/tree/3f22fc52cd867f1873add5d6405abec7d6dba375）。

エンコードされたキーに対するアクセスはすべてこのユーティリティ関数群を用います。staticおじさんになってしまいました。

final class EncodedKeyUtils {
  private EncodedKeyUtils() {}

  static Object createEmpty(EncodedKeyType encodedKeyType) {
    switch (encodedKeyType) {
      case STRING:
        return "";
      case BYTE_ARRAY:
        return ByteArray.EMPTY_INSTANCE;
      default:
        throw new AssertionError();
    }
  }

  static int length(EncodedKeyType encodedKeyType, Object obj) {
    switch (encodedKeyType) {
      case STRING:
        assert obj instanceof String;
        return ((String) obj).length();
      case BYTE_ARRAY:
        assert obj instanceof ByteArray;
        return ((ByteArray) obj).length();
      default:
        throw new AssertionError();
    }
  }

  static Object slice(EncodedKeyType encodedKeyType, Object obj, int length) {
    switch (encodedKeyType) {
      case STRING:
        assert obj instanceof String;
        return ((String) obj).substring(0, length);
      case BYTE_ARRAY:
        assert obj instanceof ByteArray;
        return ((ByteArray) obj).slice(length);
      default:
        throw new AssertionError();
    }
  }

  static int get(EncodedKeyType encodedKeyType, Object obj, int pos) {
    switch (encodedKeyType) {
      case STRING:
        assert obj instanceof String;
        return ((String) obj).charAt(pos);
      case BYTE_ARRAY:
        assert obj instanceof ByteArray;
        return ((ByteArray) obj).get(pos);
      default:
        throw new AssertionError();
    }
  }

  ...

その結果、文字列キーの場合はEncodedKeyインスタンスの生成を回避できるようになり、性能はv0.1と同等まで回復しました。一安心です。

数値型キーの対応（数値キー→元の数値をラップ、文字列キーは元のStringのまま）

ここまでで数値キーの場合は内部的にはbyte[]にエンコードしていたのですが、数値をそのままラップしてTrieで探索可能なメソッドを用意してあげたらどうか？と思い試してみました（https://github.com/komamitsu/wormhole4j/tree/a36b28addb08b8410729d78e9c92c24c50778c63）

すなわちIntWrapper的なクラスを作ることになります。

package org.komamitsu.wormhole4j;

final class IntWrapper implements Comparable<IntWrapper> {
  private static final int SIZE = 4;
  private static final int[] BITMASKS =
      new int[] {
        0x00000000, // Length: 0
        0xFF000000, // Length: 1
        0xFFFF0000, // Length: 2
        0xFFFFFF00, // Length: 3
        0xFFFFFFFF, // Length: 4
      };
  static final IntWrapper EMPTY_INSTANCE = new IntWrapper(0, 0);
  private static final char[] HEX_CHARS = "0123456789ABCDEF".toCharArray();
  // The value is already masked based on the length.
  private final int value;
  private final int length;

  IntWrapper(int value) {
    this.length = SIZE;
    this.value = value;
  }

  private IntWrapper(int value, int length) {
    assert length <= SIZE;
    this.length = length;
    this.value = value & BITMASKS[length];
  }

  @Override
  public int compareTo(IntWrapper other) {
    int minLen = Math.min(length, other.length);
    int mask = BITMASKS[minLen];
    int x1 = value & mask;
    int x2 = other.value & mask;
    if (x1 != x2) {
      return Integer.compareUnsigned(x1, x2);
    }
    return length - other.length;
  }

  int get(int pos) {
    assert pos < length;
    return (value >>> ((SIZE - 1 - pos) << 3)) & 0xFF;
  }

  @Override
  public boolean equals(Object o) {
    if (this == o) {
      return true;
    }
    if (!(o instanceof IntWrapper)) {
      return false;
    }
    IntWrapper other = (IntWrapper) o;
    return length == other.length && value == other.value;
  }

  @Override
  public int hashCode() {
    return value * 31 + length;
  }

  ...

Trieのキーとして4バイト未満の長さのキーも必要になるので、論理的なキー長を別途管理することになります。内部表現としてbyte[]を使っていたバージョンに比べて、hashCode()やequals()等の効率が改善されてそうです（profile取っていないので勘）。

結果的に数値キーの性能が有意に向上し、java.util.TreeMapの性能に近づいてきました。

全体的な最適化（さらなる抽象化の破壊、そして泥臭コードへ）

前述の通り、Wormhole4jではLeafNodeインスタンス内に128個のKeyValueインスタンスを保持・管理していましたが、それに加え同数のTag,KeyReferenceインスタンスを保持・管理していました（https://github.com/komamitsu/wormhole4j/tree/0c8d65f549065ad6df3e97c3893e31625a0ae7f1）。

Tagsは整列済み「エンコードされたキーのハッシュ値」の配列でget(key)時のpoint lookupに利用されます。Key Referenceは一部整列済み「エンコードされたキー」の配列でscan(start_key, ...)時のrange queryに利用されます。

final class LeafNode<K, V> {
  private final EncodedKeyType encodedKeyType;
  final Object anchorKey;
  private final int maxSize;
  private final KeyValues<K, V> keyValues;
  // All references are always sorted by hash.
  private final Tags<K, V> tags;
  // Some references are sorted by key.
  private final KeyReferences<K, V> keyReferences;

  private static class KeyValues<K, V> {
    private int count;
    private final KeyValue<K, V>[] entries;
    ...
  }

  private static class Tags<K, V> {
    private int count;
    private final KeyValues<K, V> keyValues;
    private final int[] entries;
    ...
  }

  private static class KeyReferences<K, V> {
    private final EncodedKeyType encodedKeyType;
    private final KeyValues<K, V> keyValues;
    private int count;
    private final int[] entries;
    private int numOfSortedEntries;
    ...
  }

  ...

これまでの経験上、抽象化を壊してkey-value, tag, key-referenceのリストを直接したほうが汚くなるが速くなる予感がしたので、汚くしてみました（https://github.com/komamitsu/wormhole4j/tree/4cd3d33163c608d3e6e305c4a9b9392a2592d20e）。

final class LeafNode<K, V> {
  private static final int TUPLE_SIZE = 3;
  private static final int ENCODED_KEY_OFFSET = 0;
  private static final int KEY_OFFSET = 1;
  private static final int VALUE_OFFSET = 2;
  private final EncodedKeyType encodedKeyType;
  final Object anchorKey;
  private final int maxSize;

  ...

  // For KeyValues
  private int keyValuesCount;
  private final Object[] keyValues;

  ...

  // For Tags
  private int tagsCount;
  private final int[] tags;

  ...

  // For KeyRefs
  private int keyRefsCount;
  private final int[] keyRefs;
  private int numOfSortedKeyRefs;
}

TagsやKeyRefsは基本的には内部の配列をLeafNode直下に移動しただけですが、
KeyValuesについてはKeyValueインスタンス化も回避することにしたため、Object[] keyValuesの中にエンコードされたキー、元のキー、値が順番に格納されている悲しい状況になっています。世も末です。

keyValues array (Object[])
┌───────────────────────────────────────────────────────────────────────┐
│                                                                       │
├─────────────┬─────────────┬─────────────┬─────────────┬───────────────┤
│ EncodedKey₀ │    Key₀     │   Value₀    │ EncodedKey₁ │    Key₁ ...   │
├─────────────┼─────────────┼─────────────┼─────────────┼───────────────┤
│   [0]       │   [1]       │   [2]       │   [3]       │   [4]   ...   │
└─────────────┴─────────────┴─────────────┴─────────────┴───────────────┘
      ↑             ↑             ↑             ↑             ↑
      │             │             │             │             │
   ENCODED_      KEY_         VALUE_        ENCODED_       KEY_
   KEY_OFFSET    OFFSET       OFFSET        KEY_OFFSET     OFFSET
   (0)           (1)          (2)           (0)            (1)


Index calculation for n-th entry:
┌────────────────────────────────────────┐
│ base_index = n * TUPLE_SIZE (3)        │
│ encoded_key_index = base + 0           │
│ key_index         = base + 1           │
│ value_index       = base + 2           │
└────────────────────────────────────────┘

Example for 3 key-value pairs:
┌──────┬──────┬──────┬──────┬──────┬──────┬──────┬──────┬──────┐
│ EK₀  │  K₀  │  V₀  │ EK₁  │  K₁  │  V₁  │ EK₂  │  K₂  │  V₂  │
├──────┼──────┼──────┼──────┼──────┼──────┼──────┼──────┼──────┤
│  [0] │  [1] │  [2] │  [3] │  [4] │  [5] │  [6] │  [7] │  [8] │
└──────┴──────┴──────┴──────┴──────┴──────┴──────┴──────┴──────┘
   └──── KV entry 0 ────┘ └──── KV entry 1 ────┘ └──── KV entry 2 ────┘

心理的には凹みますがその甲斐あって全体的に性能が改善しました。java.util.TreeMap等と雑に比べると以下のような感じです。

• 数値キーの場合
  • Getはやや負け
  • Updateはやや勝ち
  • Scanは勝ち
• 文字列キーの場合
  • Getはやや勝ち
  • Updateはやや勝ち
  • Scanは大勝ち

性能改善的には一区切りついた感があるので、v0.2としてリリースしました。実用に耐えられるものになったのではないかと思います。

Java開発における性能改善tips

上記の機能追加・性能改善を経て、Javaの開発においていくつかの学びがあったので以下に記していきたいと思います。

（特定の状況下では）Javaインスタンス生成を避ける

通常の開発では性能的なボトルネックが発生したとしてもI/O起因だと思うので、基本的にきれいに抽象化すべきですが、ボトルネックが非I/Oの場合かつ計算量的な問題でもない場合は、Javaのインスタンスを生成するコストが影響してくる可能性があります。Javaのインスタンス生成コストはかなり低いものの、塵も積もれば有意な影響が現れてきます。その場合は（本当に必要であれば）、保守性等とのトレードオフですが抽象化を壊してJavaインスタンス生成を減らすことも選択肢の一つかもしれません。

可変引数は内部的に配列を生成している

自分で可変引数を実装する場合以外でも、java.util.Objects.hash(Object[]... values)等を気軽に使うと知らぬ間に配列オブジェクトを生成していることになります。前述の通り通常は問題ないので、個人的にも利用推奨派ですが、Javaインスタンス生成が性能に影響を与えるような特殊な条件下では利用の要否を検討したほうが良いかもです。

Object.hashCode()の上書きは慎重に

これは単なる僕のミスだったのですが、Integerキー対応後、これをベースに雑にLongキー対応したところ、性能がガタ落ちしてしまいました。原因はIntWrapper.hashCode()が

final class IntWrapper implements Comparable<IntWrapper> {
  private final int value;
  private final int length;
  ...
  @Override
  public int hashCode() {
    return value * 31 + length;
  }

となっていて、それをコピーしたLongWrapper.hashCode()が

final class LongWrapper implements Comparable<LongWrapper> {
  private final long value;
  private final int length;
　...
  @Override
  public int hashCode() {
    return (int)(value * 31 + length);
  }

と実装されていた（ように記憶している）のですが、この実装だと上位32ビットの情報を捨ててしまってハッシュ値の衝突が起こりやすくなっていました（これらのインスタンスはHashMapのキーになる）。正しくは

final class LongWrapper implements Comparable<LongWrapper> {
  private final long value;
  private final int length;
　...
  @Override
  public int hashCode() {
    return Long.hashCode(value) * 31 + length;
  }

等にすべきでした。ちなみに`java.lang.Long.hashCode()の実装は以下で、上位32ビットとのXORをとっています。

  public static int hashCode(long value) {
        return (int)(value ^ (value >>> 32));
    }

JMHを利用する場合、org.openjdk.jmh.infra.Blackhole.consume()の利用は最低限に控えること

Blackhole.consume()はdeadcode eliminationを防ぐ手段として推奨されることが多いですが、暗黙的にmemory barrierとして動作する等、強い厳格さにより予想外にボトルネックになる可能性があります。また、性能への影響の有無はconsume対象変数の由来等によるため、似たようなケースでも性能に影響が出たり出なかったりと問題の切り分けが難しい場合があります。頻繁に呼ばれる箇所ではローカル変数を更新するに止め、最後に当該ローカル変数をconsumeすることで、deadcode eliminationを最低限防ぎつつ、Blackhole.consume()の性能への影響の可能性を減らせるのではないかと思います。

最後に

ということで、まとまりのない文章になってしまいましたが、Wormhole4jというWormholeをJavaで実装したsorted mapの性能改善にまつわる話を書いてみました。