blur-query/src/main/java/org/apache/blur/lucene/search/FacetExecutor.java

/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.blur.lucene.search;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLongArray;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

import org.apache.blur.log.Log;
import org.apache.blur.log.LogFactory;
import org.apache.blur.trace.Trace;
import org.apache.blur.trace.Tracer;
import org.apache.lucene.index.AtomicReader;
import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.util.OpenBitSet;

public class FacetExecutor {

  private static final Log LOG = LogFactory.getLog(FacetExecutor.class);

  static Comparator<Entry<Object, Info>> COMPARATOR = new Comparator<Entry<Object, Info>>() {
    @Override
    public int compare(Entry<Object, Info> o1, Entry<Object, Info> o2) {
      AtomicReader r1 = o1.getValue()._reader;
      AtomicReader r2 = o2.getValue()._reader;
      return r2.maxDoc() - r1.maxDoc();
    }
  };

  static class Finished extends IOException {
    private static final long serialVersionUID = 1L;
  }

  static class SimpleCollector extends Collector {

    int _hits;
    final OpenBitSet _bitSet;
    Scorer _scorer;

    SimpleCollector(OpenBitSet bitSet) {
      _bitSet = bitSet;
    }

    @Override
    public void collect(int doc) throws IOException {
      if (_bitSet.fastGet(doc)) {
        _hits++;
      } else {
        int nextSetBit = _bitSet.nextSetBit(doc);
        if (nextSetBit < 0) {
          throw new Finished();
        } else {
          int advance = _scorer.advance(nextSetBit);
          if (advance == DocIdSetIterator.NO_MORE_DOCS) {
            throw new Finished();
          }
          if (_bitSet.fastGet(advance)) {
            _hits++;
          }
        }
      }
    }

    @Override
    public void setScorer(Scorer scorer) throws IOException {
      _scorer = scorer;
    }

    @Override
    public void setNextReader(AtomicReaderContext context) throws IOException {

    }

    @Override
    public boolean acceptsDocsOutOfOrder() {
      return false;
    }

  }

  static class SimpleCollectorExitEarly extends SimpleCollector {

    private long _currentCount;
    private final long _min;

    SimpleCollectorExitEarly(OpenBitSet bitSet, long currentCount, long min) {
      super(bitSet);
      _currentCount = currentCount;
      _min = min;
    }

    @Override
    public void collect(int doc) throws IOException {
      super.collect(doc);
      _currentCount++;
      if (_currentCount >= _min) {
        throw new Finished();
      }
    }

  }

  static class Info {
    final OpenBitSet _bitSet;
    final Scorer[] _scorers;
    final AtomicReader _reader;
    final String _readerStr;
    final int _maxDoc;
    final Lock[] _locks;

    @Override
    public String toString() {
      return "Info scorers length [" + _scorers.length + "] reader [" + _reader + "]";
    }

    Info(AtomicReaderContext context, Scorer[] scorers, Lock[] locks) {
      AtomicReader reader = context.reader();
      _bitSet = new OpenBitSet(reader.maxDoc());
      _scorers = scorers;
      _reader = reader;
      _readerStr = _reader.toString();
      _maxDoc = _reader.maxDoc();
      _locks = locks;
    }

    void process(AtomicLongArray counts, long[] minimumsBeforeReturning, AtomicBoolean running) throws IOException {
      if (minimumsBeforeReturning == null) {
        Tracer trace = Trace.trace("processing facet - segment", Trace.param("reader", _readerStr),
            Trace.param("maxDoc", _maxDoc), Trace.param("minimums", "NONE"), Trace.param("scorers", _scorers.length));
        try {
          for (int i = 0; i < _scorers.length && running.get(); i++) {
            SimpleCollector col = new SimpleCollector(_bitSet);
            runFacet(counts, col, i);
          }
        } finally {
          trace.done();
        }
      } else {
        BlockingQueue<Integer> ids = new ArrayBlockingQueue<Integer>(_scorers.length + 1);
        try {
          populate(ids);
          while (!ids.isEmpty() && running.get()) {
            int id = ids.take();
            Lock lock = _locks[id];
            if (lock.tryLock()) {
              try {
                long min = minimumsBeforeReturning[id];
                long currentCount = counts.get(id);
                if (currentCount < min) {
                  LOG.debug("Running facet, current count [{0}] min [{1}]", currentCount, min);
                  SimpleCollectorExitEarly col = new SimpleCollectorExitEarly(_bitSet, currentCount, min);
                  runFacet(counts, col, id);
                }
              } finally {
                lock.unlock();
              }
            } else {
              ids.put(id);
            }
          }
        } catch (Exception e) {
          throw new IOException(e);
        }
      }
    }

    private void populate(BlockingQueue<Integer> ids) throws InterruptedException {
      for (int i = 0; i < _scorers.length; i++) {
        ids.put(i);
      }
    }

    private void runFacet(AtomicLongArray counts, SimpleCollector col, int i) throws IOException {
      Scorer scorer = _scorers[i];
      if (scorer != null) {
        Tracer traceInner = Trace.trace("processing facet - segment - scorer", Trace.param("scorer", scorer),
            Trace.param("scorer.cost", scorer.cost()));
        try {
          // new ExitScorer(scorer).score(col);
          scorer.score(col);
        } catch (Finished e) {
          // Do nothing, exiting early.
        } finally {
          traceInner.done();
        }
        int hits = col._hits;
        LOG.debug("Facet [{0}] result [{1}]", i, hits);
        counts.addAndGet(i, hits);
      }
      col._hits = 0;
    }
  }

  private final Map<Object, Info> _infoMap = new ConcurrentHashMap<Object, FacetExecutor.Info>();
  private final int _length;
  private final AtomicLongArray _counts;
  private final long[] _minimumsBeforeReturning;
  private final Lock[] _locks;
  private final AtomicBoolean _running;
  private boolean _processed;

  public FacetExecutor(int length) {
    this(length, null, new AtomicLongArray(length));
  }

  public FacetExecutor(int length, long[] minimumsBeforeReturning, AtomicLongArray counts) {
    this(length, minimumsBeforeReturning, counts, new AtomicBoolean(true));
  }

  public FacetExecutor(int length, long[] minimumsBeforeReturning) {
    this(length, minimumsBeforeReturning, new AtomicLongArray(length));
  }

  public FacetExecutor(int length, long[] minimumsBeforeReturning, AtomicLongArray counts, AtomicBoolean running) {
    _length = length;
    _counts = counts;
    _minimumsBeforeReturning = minimumsBeforeReturning;
    _locks = new Lock[_length];
    for (int i = 0; i < _length; i++) {
      _locks[i] = new ReentrantReadWriteLock().writeLock();
    }
    _running = running;
  }

  public void addScorers(AtomicReaderContext context, Scorer[] scorers) throws IOException {
    if (scorers.length != _length) {
      throw new IOException("Scorer length is not correct expecting [" + _length + "] actual [" + scorers.length + "]");
    }
    Object key = getKey(context);
    Info info = _infoMap.get(key);
    if (info == null) {
      info = new Info(context, scorers, _locks);
      _infoMap.put(key, info);
    } else {
      AtomicReader reader = context.reader();
      LOG.warn("Info about reader context [{0}] already created, existing Info [{1}] current reader [{2}].", context,
          info, reader);
    }
  }

  public boolean scorersAlreadyAdded(AtomicReaderContext context) {
    Object key = getKey(context);
    return _infoMap.containsKey(key);
  }

  private Object getKey(AtomicReaderContext context) {
    return context.reader().getCoreCacheKey();
  }

  public OpenBitSet getBitSet(AtomicReaderContext context) throws IOException {
    Info info = _infoMap.get(getKey(context));
    if (info == null) {
      throw new IOException("Info object missing.");
    }
    return info._bitSet;
  }

  public int length() {
    return _length;
  }

  public long get(int i) throws IOException {
    if (!_processed) {
      throw new IOException("Has not been processed.");
    }
    return _counts.get(i);
  }

  public void processFacets(ExecutorService executor) throws IOException {
    if (!_processed) {
      Tracer trace = Trace.trace("processing facets");
      try {
        processInternal(executor);
      } finally {
        trace.done();
      }
      _processed = true;
    }
  }

  private void processInternal(ExecutorService executor) throws IOException {
    List<Entry<Object, Info>> entries = new ArrayList<Entry<Object, Info>>(_infoMap.entrySet());
    Collections.sort(entries, COMPARATOR);
    if (executor == null) {
      for (Entry<Object, Info> e : entries) {
        if (_running.get()) {
          e.getValue().process(_counts, _minimumsBeforeReturning, _running);
        }
      }
    } else {
      final AtomicInteger finished = new AtomicInteger();
      for (Entry<Object, Info> e : entries) {
        if (_running.get()) {
          final Entry<Object, Info> entry = e;
          executor.submit(new Runnable() {
            @Override
            public void run() {
              try {
                entry.getValue().process(_counts, _minimumsBeforeReturning, _running);
              } catch (Throwable e) {
                LOG.error("Unknown error", e);
              } finally {
                finished.incrementAndGet();
              }
            }
          });
        }
      }

      while (finished.get() < entries.size()) {
        synchronized (this) {
          try {
            wait(1);
          } catch (InterruptedException e) {
            throw new IOException(e);
          }
        }
      }
    }
  }
}