【Calcite源码学习】ImmutableBitSet介绍

Calcite中实现了一个ImmutableBitSet类，用于保存bit集合。在很多优化规则和物化视图相关的类中都使用了ImmutableBitSet来保存group by字段或者聚合函数参数字段对应的index，例如：

//MaterializedViewAggregateRule#compensateViewPartial()
ImmutableBitSet.Builder groupSet = ImmutableBitSet.builder();
groupSet.addAll(aggregateViewNode.getGroupSet());
groupSet.addAll(
    ImmutableBitSet.range(
        aggregateViewNode.getInput().getRowType().getFieldCount(),
        aggregateViewNode.getInput().getRowType().getFieldCount() + offset));

因此，本文我们就来看一看这个ImmutableBitSet一些核心的操作函数。

使用示例

这里我们介绍一个ImmutableBitSet使用的例子，后续相关的函数操作说明会使用这个例子：

ImmutableBitSet.Builder builder = ImmutableBitSet.builder();
builder.set(3);
builder.set(5);
builder.set(67);
builder.set(70);
builder.set(129);
ImmutableBitSet bitSet = builder.build();
System.out.println(bitSet.cardinality());
System.out.println(bitSet.nth(1));

如上所示，通常先构造一个Builder，然后set指定的bit位，最终构造生成ImmutableBitSet。下面来看一下主要的成员和关键函数。

主要成员

ImmutableBitSet中比较重要的几个成员变量如下所示：

private static final int ADDRESS_BITS_PER_WORD = 6;
private static final int BITS_PER_WORD = 1 << ADDRESS_BITS_PER_WORD;

private static final long[] EMPTY_LONGS = new long[0];

private static final ImmutableBitSet EMPTY = new ImmutableBitSet(EMPTY_LONGS);
  
private final long[] words;

这些变量含义如下：

这里最主要的变量就是words，这是一个long类型的数组，实际存储设置的bit位；
Long是64位，使用6个bit就可以表示，因此ADDRESS_BITS_PER_WORD设置为6，用来进行一些二进制的左移和右移的运算；
BITS_PER_WORD表示将1左移6位，结果就是64。当words数组中存在多个成员的时候，取指定long成员对应的bit位时会用到，后面会详细介绍，这里不再赘述。

Builder相关函数

要使用ImmutableBitSet，我们首先需要构造一个ImmutableBitSet.Builder，因此先看一下这个Builder的关键函数。

set(int i)

set函数是最常用的，将指定位的bit设置为true，表示该位上面有值，对应的函数代码如下所示：

public Builder set(int bit) {
  if (words == null) {
    throw new IllegalArgumentException("can only use builder once");
  }
  int wordIndex = wordIndex(bit);
  if (wordIndex >= words.length) {
    words = Arrays.copyOf(words, wordIndex + 1);
  }
  words[wordIndex] |= 1L << bit;
  return this;
}

private static int wordIndex(int bitIndex) {
  return bitIndex >> ADDRESS_BITS_PER_WORD;
}

通过这个函数，就可以往Builder里面设置对应的bit位，上述代码主要处理逻辑如下：

通过传入的bit位，获取在words中的索引，即数组的下标，通过wordIndex函数来实现。这个函数只有一行代码，即将传入的bit右移6位。其效果是：0～63会返回0，表示存储在words中的第一个long成员；64～127返回1，表示存储在words中的第二个long成员，以此类推；
如果word索引大于当前数组的长度，则会扩容，数组长度+1；
将1L左移bit位，并与指定index的数组成员，即words[wordIndex]进行或操作，并更新到当前数组成员，然后返回。

这里我们结合上述的例子来看一下：

//3 >> 6 = 0，表示wordIndex是0，存储在第一个long成员；1L << 3，得到的结果是0000 1000
builder.set(3);
//5 >> 6 = 0，表示wordIndex是0，存储在第一个long成员；1L << 5，得到的结果是0010 0000
builder.set(5);
//67 >> 6 = 1，表示wordIndex是1，存储在第二个long成员；1L << 67，得到的结果是0000 1000
builder.set(67);
//70 >> 6 = 1，表示wordIndex是1，存储在第二个long成员；1L << 70，得到的结果是0100 1000
builder.set(70);
//129 >> 6 = 2，表示wordIndex是2，存储在第三个long成员；1L << 129，得到的结果是0000 0010
builder.set(129);

这里我们只展示了二进制的最后8位，前面56位都是0，因此省略。此时，words包含3个long成员，分别是：

第一个long成员的值就是0000 1000 ｜ 0010 0000 = 0010 1000，十进制是40；
第二个long成员的值就是0000 1000 ｜ 0100 0000 = 0100 1000，十进制是72；
第三个long成员的值就是0000 0010，十进制是2；

我们通过debug可以进行验证：

需要注意的是，在对1L进行左移操作时，由于long是64位的，因此我们需要先对bit进行64取余操作，因此上面大于64的左移操作实际效果如下：

1L << 67 => 1L << (67 % 64) => 1L << 3 => 0000 1000
1L << 70 => 1L << (70 % 64) => 1L << 6 => 0100 0000
1L << 129 => 1L << (129 % 64) => 1L << 1 => 0000 0010

build()

Builder填充完成之后，就可以通过build()方法来生成ImmutableBitSet对象了。这个函数比较简单，就是将words成员作为参数来构造ImmutableBitSet，同时将words成员本身置空。也就是说，Builder只能调用一次build()，这点需要注意。

addAll(ImmutableBitSet bitSet)

set方法只能设置单个bit位，如果想要批量设置的话，可以调用addAll()函数。这个函数有三个重载方法，这里我们看一下与ImmutableBitSet相关的：

public Builder addAll(ImmutableBitSet bitSet) {
  for (Integer bit : bitSet) {
    set(bit);
  }
  return this;
}

该方法比较简单，就是通过ImmutableBitSet的iterator()方法来获取每一个bit，然后set到当前的Builder中。关于iterator()处理逻辑，我们下面再介绍。

ImmutableBitSet相关函数

看完Builder常用的几个函数后，我们再来看一下ImmutableBitSet相关的函数。

cardinality()

这个函数会返回当前ImmutableBitSet包含的所有bit位，相关代码如下所示：

public int cardinality() {
  return countBits(words);
}

private static int countBits(long[] words) {
  int sum = 0;
  for (long word : words) {
    sum += Long.bitCount(word);
  }
  return sum;
}

可以看到，实际是通过countBits函数来进行统计的。该函数会遍历words数组所有的成员，依次统计每个成员包含的bit位，然后累计求和。在上述例子中，对应的cardinality就是5，表示一共有5个bit位被设置为true，前面两个成员分别包含2个为true的bit位，第三个成员包含1个。

nth(int)

该函数是Calcite中也会经常用到的。该函数的功能是取第n位对应的值，代码如下所示：

public int nth(int n) {
  int start = 0;
  for (long word : words) {
    final int bitCount = Long.bitCount(word);
    if (n < bitCount) {
      while (word != 0) {
        if ((word & 1) == 1) {
          if (n == 0) {
            return start;
          }
          --n;
        }
        word >>= 1;
        ++start;
      }
    }
    start += 64;
    n -= bitCount;
  }
  throw new IndexOutOfBoundsException("index out of range: " + n);
}

代码对应的主要处理逻辑如下：

遍历words数组，通过每个成员包含的bit位数量，来确定第n个bit位于哪个数组成员；
遍历的同时，更新每个数组成员对应的起始值start，第一个成员是0，第二个成员是64，第三个成员是128，依次类推；
定位到对应的数组成员之后，循环处理这个long，每次右移1位，然后start加1；
右移1位之后，判断long的最后一位是否为1；是的话，则表示该bit位被set过，然后再判断是否为第n个bit，是的话，则直接返回start。否则继续右移long，然后更新start。

我们结合上述示例来看一下这个处理过程。n=3时，整个流程如下所示：

iterator()

上面我们提到，Builder.addAll()函数主要就是通过ImmutableBitSet的迭代器iterator()来实现批量set的，这里我们就来看一下ImmutableBitSet是如何实现的。当我们使用for循环取ImmutableBitSet的成员时，会自动调用迭代器的next()函数不断获取下一个成员，next()函数实现：

//ImmutableBitSet#iterator()
@Override public Integer next() {
  int prev = i;
  i = nextSetBit(i + 1);
  return prev;
}

可以看到，这里主要是通过nextSetBit()函数来取下一个bit位的，下面我们看下这个函数的实现逻辑。

nextSetBit(int fromIndex)

这个函数是用来取给定bit位以及其之后的，下一个被set的bit位（包含fromIndex对应的bit位本身），函数代码如下所示：

public int nextSetBit(int fromIndex) {
  if (fromIndex < 0) {
    throw new IndexOutOfBoundsException("fromIndex < 0: " + fromIndex);
  }
  int u = wordIndex(fromIndex);
  if (u >= words.length) {
    return -1;
  }
  long word = words[u] & (WORD_MASK << fromIndex);
   while (true) {
    if (word != 0) {
      return (u * BITS_PER_WORD) + Long.numberOfTrailingZeros(word);
    }
    if (++u == words.length) {
      return -1;
    }
    word = words[u];
  }
}

函数代码主要的处理逻辑如下：

获取传入index所在的数组成员下标，即位于第几个数组成员中；
将word mask右移index位，例如index为3，则结果是1111 1000（省略mask前面的1），然后将该结果与对应的数组成员进行与操作，得到一个新的word。这行代码的目的是为了屏蔽index前面的成员，进行后续的比较。例如数组成员是1111 1111，与操作的结果则是1111 1000。那么word最后三个1（对应十进制的1、2、4）都被屏蔽；
如果word为空，表示传入的index，不在当前这个word中，对应的bit没有被set，此时更新word为下一个数组成员。由于下一个数组成员的起始index是必定大于传入index的，因此word不需要进行上面的左移和与操作；
word不为空，先获取trailing zero数量，然后加上实际的offset即可（数组成员每往后移动一个，index起始offset就要加64）。

我们同样使用上述的例子来看一下。fromIndex设置为6时，整个流程如下所示：

get(int bitIndex)

该函数是用来判断bit set的指定bit位是否被set为true，函数代码如下：

public boolean get(int bitIndex) {
  if (bitIndex < 0) {
    throw new IndexOutOfBoundsException("bitIndex < 0: " + bitIndex);
  }
  int wordIndex = wordIndex(bitIndex);
  return (wordIndex < words.length)
      && ((words[wordIndex] & (1L << bitIndex)) != 0);
}

核心的代码逻辑就是words[wordIndex] & (1L << bitIndex)，如果为0，表示对应bit没有被set，即为0。这里我们仍然结合上述例子来看：

//假设bitIndex为69，此时位于words的第二个数组成员，范围为64～127
1L << 69 => 1L << (75 % 69) => 1L << 6 => 0100 0000
0100 1000 & 0100 0000 = 0100 0000
//假设bitIndex为75，此时位于words的第二个数组成员，范围为64～127
1L << 75 => 1L << (75 % 64) => 1L << 11 => 1000 0000 0000
0000 0100 1000 & 1000 0000 0000 = 0

因此，get(69)返回true，对应上述示例的set(70)；get(75)则返回false。

小结

本文主要介绍了Builder和ImmutableBitSet中的一些关键函数的处理逻辑，并结合示例展示了具体的用法。ImmutableBitSet还有不少其他的函数，其中大多都是基于上述的函数进行实现或者本身实现比较简单，这里不再一一介绍，有兴趣的同学可以自行阅读源码。
关于Calcite为什么要开发一个这样的bit set，笔者猜测应该是跟逻辑执行计划有关。使用ImmutableBitSet可以通过与操作快速判断字段对应的index是否在bit set中，效率比较高。后续阅读到planner优化相关的代码时，再做分享。