es的聚合查询会涉及到很多概念,比如fielddata,DocValue,也会引出很多问题,比如聚合查询导致的内存溢出。在没有真正了解聚合查询的情况下,我们往往对这些概念,问题都是云山雾绕的。本文我们分析一下ES聚合查询的源码,理清楚聚合查询的流程。穿越层层迷雾来认清聚合的本质。

聚合查询的入口

es的聚合查询的入口代码如下:

public void execute(SearchContext searchContext) throws QueryPhaseExecutionException {
        aggregationPhase.preProcess(searchContext);  <1>
        boolean rescore = execute(searchContext, searchContext.searcher());<2>
        aggregationPhase.execute(searchContext);<3>
        }
    }

<1>为聚合查询做准备
<2>根据条件进行查询,获取查询结果
<3>对查询结果进行聚合

<3>才是聚合的真正入口,但是要想真正理解ES的聚合,我们必须了解<1><2>。因为<1>中提供了聚合查询必要的采集器(collector), 正排索引。<2>为聚合查询提供了数据基础,即<3>是在<2>中采集出来的数据的基础上进行的。下面以这3步为大纲,分析es的聚合查询源码

聚合前的准备

聚合前所需要做的准备主要就是一件事:构建采集器aggregators。
aggregators是一个由aggregator组成的列表,aggregator包装着聚合的实现逻辑,因为es拥有多种聚合方式,所以也就有多种不同实现逻辑的aggregator。在查询阶段中,es会调用aggregator中的逻辑去采集数据。

aggregator的构建

AggregatorFactories factories = context.aggregations().factories();   <1>
aggregators = factories.createTopLevelAggregators(aggregationContext); <2>

<1>从上下文中获取aggregator的工厂
<2>工厂生产出aggregators

值得一提的是<2>步骤表面是只是生产了aggregator。实际上还偷偷干了一件重要的事情:加载Doc Value 。
Doc Value 和FieldData是es的正排索引,它对提升es聚合查询的性能起着至关重要的作用。因此,有必要探究一下它的加载逻辑。

DocValue的加载

DocValue的加载的源码位置:AggregatorFactories:createTopLevelAggregators -> AggregatorFactory:create -> ValuesSourceAggregatorFactory:createInternal

public Aggregator createInternal(AggregationContext context, Aggregator parent, boolean collectsFromSingleBucket,
            List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData) throws IOException {
        VS vs = config.toValuesSource(context.getQueryShardContext()); <1>
        return doCreateInternal(vs, context, parent, collectsFromSingleBucket, pipelineAggregators, metaData); <2>
    }

<1> 从config中获取value source,vs中包含了DocValue
<2>fielddata 作为vs参数传入该方法中。

protected Aggregator doCreateInternal(ValuesSource valuesSource, Aggregator parent, boolean collectsFromSingleBucket,
            List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData) throws IOException {
 
       ...
                ValuesSource.Bytes.WithOrdinals valueSourceWithOrdinals = (ValuesSource.Bytes.WithOrdinals) valuesSource;
                IndexSearcher indexSearcher = context.searcher();
                maxOrd = valueSourceWithOrdinals.globalMaxOrd(indexSearcher);<1>
                ratio = maxOrd / ((double) indexSearcher.getIndexReader().numDocs());<2>
           ...
    }

<1>从vs中加载docValue,它首先会尝试去本地缓存中找,如果本地缓存中没有DocValue的话,就从磁盘文件中读取,着就是传入indexSearcher的目的。获取到docvalue之后就可以获得maxOrd,它表示这个字段中term的总数。
<2>ratio=词项总数/文档总数。radio越小聚合出来结果的bucket数量就越小。根据ratio的值我们应该选择适合的聚合模式以优化聚合查询的性能。

加载出来的docValue真正排上用场是在执行查询的过程中。

在查询过程中采集数据

这个步骤的入口代码位于QueryPhase:execute方法中,这个方法很长,内容很多,但是我们只关注它与聚合部分的联系,因此我们只需要看到其中的一行代码

searcher.search(query, collector);

这行代码是es正式开始查询的入口,它直接调用的lucene的查询接口,query参数包含了查询条件,collector则是封装了aggregator,它携带了聚合的逻辑,我们称collector为采集器。

protected void search(List<LeafReaderContext> leaves, Weight weight, Collector collector)
      throws IOException {
    for (LeafReaderContext ctx : leaves) { // search each subreader
      final LeafCollector leafCollector;
      try {
        leafCollector = collector.getLeafCollector(ctx);
      } catch (CollectionTerminatedException e) {
        continue;
      }
      BulkScorer scorer = weight.bulkScorer(ctx);<1>
      if (scorer != null) {
        try {
          scorer.score(leafCollector, ctx.reader().getLiveDocs());<1>
        } catch (CollectionTerminatedException e) {
          // collection was terminated prematurely
          // continue with the following leaf
        }
      }
    }
  }

这段代码的逻辑一目了然。我们知道es中一个索引包含多个分片,一个分片包含多个段,这里的leaves就是段的集合。代码中遍历每个段,去查询段中符合查询条件的文档,给文档打分,用采集器收集匹配查询文档的聚合指标数据。
<1>找出匹配条件的文档集合
<2>遍历匹配的文档集合,用采集器采集指标数据。
我们只关注跟聚合相关的<2>

for (int doc = iterator.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = iterator.nextDoc()) {
          if (acceptDocs == null || acceptDocs.get(doc)) {
            collector.collect(doc);
          }
        }
public void collect(int doc) throws IOException {
      final LeafCollector[] collectors = this.collectors;
      int numCollectors = this.numCollectors;
      for (int i = 0; i < numCollectors; ) {
        final LeafCollector collector = collectors[i];
        try {
          collector.collect(doc);
          ++i;
        } catch (CollectionTerminatedException e) {
          removeCollector(i);
          numCollectors = this.numCollectors;
          if (numCollectors == 0) {
            throw new CollectionTerminatedException();
          }
        }
      }
    }

遍历匹配的文档集合,用采集器采集这个文档的指标

public void collect(int doc, long bucket) throws IOException {
                    assert bucket == 0;
                    final int ord = singleValues.getOrd(doc);<1>
                    if (ord >= 0) {
                        collectGlobalOrd(doc, ord, sub);<2>
                    }
                }

<1>用正排索引DocValue寻找指定文档对应的词项,这里就是前面加载的DocValue排上用场的地方了
<2>更新词项对应的指标

private void collectGlobalOrd(int doc, long globalOrd, LeafBucketCollector sub) throws IOException {
      
          collectExistingBucket(sub, doc, globalOrd);
    }
public final void collectExistingBucket(LeafBucketCollector subCollector, int doc, long bucketOrd) throws IOException {
        docCounts.increment(bucketOrd, 1);
    }

docCounts可以理解为一个Map,以词项作为key,以词项对应的文档数量作为value。这里说的词项实际上是一个数字bucketOrd,它是词项在全局的唯一标志。
到这里查询阶段数据的采集完成,docCounts就是在查询阶段为聚合准备的数据。聚合中的bucket就是从docCounts的基础上构建出来的。

ES request断路器对docCounts的内存限制

docCounts的大小取决于词项的数量,我们假设如果聚合请求涉及到的词项非常庞大,那么docCounts占用的内存空间也会非常庞大,这是不是有OOM的风险呢?所幸ES对此早已有了对策,那就是通过request 断路器来限制docCounts的大小。request 断路器的作用就是防止每个请求(比如聚合查询请求)的数据结构占用的内存超出一定的量。

private IntArray docCounts;

    public BucketsAggregator(String name, AggregatorFactories factories, SearchContext context, Aggregator parent,
            List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData) throws IOException {
        super(name, factories, context, parent, pipelineAggregators, metaData);
        bigArrays = context.bigArrays();
        docCounts = bigArrays.newIntArray(1, true);
    }

public IntArray newIntArray(long size, boolean clearOnResize) {
        if (size > INT_PAGE_SIZE) {
            // when allocating big arrays, we want to first ensure we have the capacity by
            // checking with the circuit breaker before attempting to allocate
            adjustBreaker(BigIntArray.estimateRamBytes(size), false);
            return new BigIntArray(size, this, clearOnResize);
        } 
    }

以上代码可以看到,docCounts的类型是IntArray 。而在创建一个IntArray 对象的时候,会调用adjustBreaker方法预估,加上这个intArray之后占用的内存会不会达到request 断路器定义的limit,如果超过limit就会抛出异常终止查询。这就是断路器对内存的保护。

数据聚合

进过前两个步骤,我们已经获取到了用于聚合的基础数据,现在我们可以开始聚合了。
聚合的关键代码:

aggregations.add(aggregator.buildAggregation(0));
public InternalAggregation buildAggregation(long owningBucketOrdinal) throws IOException {
       
		final int size;
        BucketPriorityQueue<OrdBucket> ordered = new BucketPriorityQueue<>(size, order.comparator(this));<1>
        OrdBucket spare = new OrdBucket(-1, 0, null, showTermDocCountError, 0);<2>
        for (long globalTermOrd = 0; globalTermOrd < valueCount; ++globalTermOrd) {<3>
           
            final long bucketOrd = getBucketOrd(globalTermOrd);
            final int bucketDocCount = bucketOrd < 0 ? 0 : bucketDocCount(bucketOrd);
            if (bucketCountThresholds.getMinDocCount() > 0 && bucketDocCount == 0) {
                continue;
            }
            otherDocCount += bucketDocCount;
            spare.globalOrd = globalTermOrd;
            spare.bucketOrd = bucketOrd;
            spare.docCount = bucketDocCount;
            if (bucketCountThresholds.getShardMinDocCount() <= spare.docCount) {
                spare = ordered.insertWithOverflow(spare);
                if (spare == null) {
                    spare = new OrdBucket(-1, 0, null, showTermDocCountError, 0);
                }
            }
        }

        // Get the top buckets
        final StringTerms.Bucket[] list = new StringTerms.Bucket[ordered.size()];<4>
        long survivingBucketOrds[] = new long[ordered.size()];
        for (int i = ordered.size() - 1; i >= 0; --i) {
            final OrdBucket bucket = ordered.pop();
            survivingBucketOrds[i] = bucket.bucketOrd;
            BytesRef scratch = new BytesRef();
            copy(lookupGlobalOrd.apply(bucket.globalOrd), scratch);
            list[i] = new StringTerms.Bucket(scratch, bucket.docCount, null, showTermDocCountError, 0, format);
            list[i].bucketOrd = bucket.bucketOrd;
            otherDocCount -= list[i].docCount;
        }
        //replay any deferred collections
        runDeferredCollections(survivingBucketOrds);

        //Now build the aggs
        for (int i = 0; i < list.length; i++) {
            StringTerms.Bucket bucket = list[i];
            bucket.aggregations = bucket.docCount == 0 ? bucketEmptyAggregations() : bucketAggregations(bucket.bucketOrd);
            bucket.docCountError = 0;
        }

        return new StringTerms(name, order, bucketCountThresholds.getRequiredSize(), bucketCountThresholds.getMinDocCount(),
                pipelineAggregators(), metaData(), format, bucketCountThresholds.getShardSize(), showTermDocCountError,
                otherDocCount, Arrays.asList(list), 0);<5>
    }

<1>创建一个bucket队列ordered ,存放bucket
<2>构建一个空的bucket对象spare
<3>构建所有bucket并且添加到ordered。其中spare.globalOrd这个词项在全局的序号,spare.bucketOrd是这个词项在段中的序号,spare.docCount这个词项拥有的文档数,这些信息都是从DocValue和docCounts中获取的。
<4>创建bucket列表list,将ordered中的bucket放入list中
<5>最后用StringTerms对象包装list返回聚合结果。

这里需要提醒的是,bucket列表list是存储在内存中的,如果这个list中bucket的数量太过庞大,比如达到了几千万甚至上亿的数据量,很可能会引发esOOM的惨案。事实上着中情况在es的运维过程中时有发生,一些不同了解es聚合原理的业余操作者,动不动就在上亿数据量的索引上对时间戳,主键这种唯一标志的字段做聚合查询,导致生成上亿个bucket最终出发es OOM。最后还抱怨不好用。对于这种问题,目前主要的规避方法是:1、培训es操作者,杜绝提交这种不合理的查询请求;2、在es上层做一层网关或者代理,过拒绝这种恶意请求。
不过我发现es6.2.0中新推出了一个search.max_buckets的配置,如果查询产生的buckets数量超过配置的数量,就能终止查询,防止es OOM。

看到这里各位看官可能会疑惑,前面不是说过es 的request断路器可以保护内存的吗?为什么阻止不了bucket列表的内存溢出。这里我们需要知道,request断路器监控的只是查询过程中产生的docCounts占用内存的大小,并没有监控聚合阶段bucket列表占用的内存。千万不要错误的以为request断路器会监控聚合查询过程中所有数据结构占用的内存!