CompactMultiValueFacetHandler.java

package com.browseengine.bobo.facets.impl;

import it.unimi.dsi.fastutil.floats.FloatArrayList;
import it.unimi.dsi.fastutil.floats.FloatList;
import it.unimi.dsi.fastutil.ints.IntArrayList;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Properties;

import org.apache.lucene.index.AtomicReader;
import org.apache.lucene.index.DocsEnum;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.Terms;
import org.apache.lucene.index.TermsEnum;
import org.apache.lucene.search.Explanation;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.util.BytesRef;

import com.browseengine.bobo.api.BoboSegmentReader;
import com.browseengine.bobo.api.BrowseFacet;
import com.browseengine.bobo.api.BrowseSelection;
import com.browseengine.bobo.api.FacetIterator;
import com.browseengine.bobo.api.FacetSpec;
import com.browseengine.bobo.facets.FacetCountCollector;
import com.browseengine.bobo.facets.FacetCountCollectorSource;
import com.browseengine.bobo.facets.FacetHandler;
import com.browseengine.bobo.facets.data.FacetDataCache;
import com.browseengine.bobo.facets.data.TermListFactory;
import com.browseengine.bobo.facets.data.TermStringList;
import com.browseengine.bobo.facets.data.TermValueList;
import com.browseengine.bobo.facets.filter.CompactMultiValueFacetFilter;
import com.browseengine.bobo.facets.filter.EmptyFilter;
import com.browseengine.bobo.facets.filter.RandomAccessAndFilter;
import com.browseengine.bobo.facets.filter.RandomAccessFilter;
import com.browseengine.bobo.facets.filter.RandomAccessNotFilter;
import com.browseengine.bobo.query.scoring.BoboDocScorer;
import com.browseengine.bobo.query.scoring.FacetScoreable;
import com.browseengine.bobo.query.scoring.FacetTermScoringFunctionFactory;
import com.browseengine.bobo.sort.DocComparator;
import com.browseengine.bobo.sort.DocComparatorSource;
import com.browseengine.bobo.util.BigIntArray;
import com.browseengine.bobo.util.BigSegmentedArray;
import com.browseengine.bobo.util.StringArrayComparator;

public class CompactMultiValueFacetHandler extends FacetHandler<FacetDataCache<?>> implements
    FacetScoreable {
  private static final int MAX_VAL_COUNT = 32;
  private final TermListFactory<?> _termListFactory;
  private final String _indexFieldName;

  public CompactMultiValueFacetHandler(String name, String indexFieldName,
      TermListFactory<?> termListFactory) {
    super(name);
    _indexFieldName = indexFieldName;
    _termListFactory = termListFactory;
  }

  public CompactMultiValueFacetHandler(String name, TermListFactory<?> termListFactory) {
    this(name, name, termListFactory);
  }

  public CompactMultiValueFacetHandler(String name, String indexFieldName) {
    this(name, indexFieldName, null);
  }

  public CompactMultiValueFacetHandler(String name) {
    this(name, name, null);
  }

  @Override
  public DocComparatorSource getDocComparatorSource() {
    return new CompactMultiFacetDocComparatorSource(this);
  }

  @Override
  public RandomAccessFilter buildRandomAccessFilter(String value, Properties prop)
      throws IOException {
    return new CompactMultiValueFacetFilter(this, value);
  }

  @Override
  public RandomAccessFilter buildRandomAccessAndFilter(String[] vals, Properties prop)
      throws IOException {
    ArrayList<RandomAccessFilter> filterList = new ArrayList<RandomAccessFilter>(vals.length);

    for (String val : vals) {
      RandomAccessFilter f = buildRandomAccessFilter(val, prop);
      if (f != null) {
        filterList.add(f);
      } else {
        return EmptyFilter.getInstance();
      }
    }
    if (filterList.size() == 1) return filterList.get(0);
    return new RandomAccessAndFilter(filterList);
  }

  @Override
  public RandomAccessFilter buildRandomAccessOrFilter(String[] vals, Properties prop, boolean isNot)
      throws IOException {
    RandomAccessFilter filter = null;

    if (vals.length > 0) {
      filter = new CompactMultiValueFacetFilter(this, vals);
    } else {
      filter = EmptyFilter.getInstance();
    }
    if (isNot) {
      filter = new RandomAccessNotFilter(filter);
    }
    return filter;
  }

  private static int countBits(int val) {
    int c = 0;
    for (c = 0; val > 0; c++) {
      val &= val - 1;
    }
    return c;
  }

  @Override
  public int getNumItems(BoboSegmentReader reader, int id) {
    FacetDataCache<?> dataCache = getFacetData(reader);
    if (dataCache == null) return 0;
    int encoded = dataCache.orderArray.get(id);
    return countBits(encoded);
  }

  @Override
  public String[] getFieldValues(BoboSegmentReader reader, int id) {
    FacetDataCache<?> dataCache = getFacetData(reader);
    if (dataCache == null) return new String[0];
    int encoded = dataCache.orderArray.get(id);
    if (encoded == 0) {
      return new String[] { "" };
    } else {
      int count = 1;
      ArrayList<String> valList = new ArrayList<String>(MAX_VAL_COUNT);

      while (encoded != 0) {
        if ((encoded & 0x00000001) != 0x0) {
          valList.add(dataCache.valArray.get(count));
        }
        count++;
        encoded >>>= 1;
      }
      return valList.toArray(new String[valList.size()]);
    }
  }

  @Override
  public Object[] getRawFieldValues(BoboSegmentReader reader, int id) {
    FacetDataCache<?> dataCache = getFacetData(reader);
    if (dataCache == null) return new String[0];
    int encoded = dataCache.orderArray.get(id);
    if (encoded == 0) {
      return new Object[0];
    } else {
      int count = 1;
      ArrayList<Object> valList = new ArrayList<Object>(MAX_VAL_COUNT);

      while (encoded != 0) {
        if ((encoded & 0x00000001) != 0x0) {
          valList.add(dataCache.valArray.getRawValue(count));
        }
        count++;
        encoded >>>= 1;
      }
      return valList.toArray(new Object[valList.size()]);
    }
  }

  @Override
  public FacetCountCollectorSource getFacetCountCollectorSource(final BrowseSelection sel,
      final FacetSpec ospec) {
    return new FacetCountCollectorSource() {

      @Override
      public FacetCountCollector getFacetCountCollector(BoboSegmentReader reader, int docBase) {
        final FacetDataCache<?> dataCache = CompactMultiValueFacetHandler.this.getFacetData(reader);
        return new CompactMultiValueFacetCountCollector(_name, sel, dataCache, docBase, ospec);
      }
    };
  }

  @SuppressWarnings({ "unchecked", "rawtypes" })
  @Override
  public FacetDataCache<?> load(BoboSegmentReader reader) throws IOException {
    int maxDoc = reader.maxDoc();

    BigIntArray order = new BigIntArray(maxDoc);

    TermValueList<?> mterms = _termListFactory == null ? new TermStringList() : _termListFactory
        .createTermList();

    IntArrayList minIDList = new IntArrayList();
    IntArrayList maxIDList = new IntArrayList();
    IntArrayList freqList = new IntArrayList();

    int t = 0; // current term number
    mterms.add(null);
    minIDList.add(-1);
    maxIDList.add(-1);
    freqList.add(0);
    t++;
    Terms terms = reader.terms(_indexFieldName);
    if (terms != null) {
      TermsEnum termsEnum = terms.iterator(null);
      BytesRef text;
      while ((text = termsEnum.next()) != null) {
        // store term text
        // we expect that there is at most one term per document
        if (t > MAX_VAL_COUNT) {
          throw new IOException("maximum number of value cannot exceed: " + MAX_VAL_COUNT);
        }
        String val = text.utf8ToString();
        mterms.add(val);
        int bit = (0x00000001 << (t - 1));
        Term term = new Term(_indexFieldName, val);
        DocsEnum docsEnum = reader.termDocsEnum(term);
        // freqList.add(termEnum.docFreq()); // removed because the df doesn't take into account the
        // num of deletedDocs
        int df = 0;
        int minID = -1;
        int maxID = -1;
        int docID = -1;
        while ((docID = docsEnum.nextDoc()) != DocsEnum.NO_MORE_DOCS) {
          df++;
          order.add(docID, order.get(docID) | bit);
          minID = docID;
          while (docsEnum.nextDoc() != DocsEnum.NO_MORE_DOCS) {
            docID = docsEnum.docID();
            df++;
            order.add(docID, order.get(docID) | bit);
          }
          maxID = docID;
        }
        freqList.add(df);
        minIDList.add(minID);
        maxIDList.add(maxID);
        t++;
      }
    }

    mterms.seal();

    return new FacetDataCache(order, mterms, freqList.toIntArray(), minIDList.toIntArray(),
        maxIDList.toIntArray(), TermCountSize.large);
  }

  private static class CompactMultiFacetDocComparatorSource extends DocComparatorSource {
    private final CompactMultiValueFacetHandler _facetHandler;

    private CompactMultiFacetDocComparatorSource(CompactMultiValueFacetHandler facetHandler) {
      _facetHandler = facetHandler;
    }

    @Override
    public DocComparator getComparator(final AtomicReader reader, int docbase) throws IOException {
      if (!(reader instanceof BoboSegmentReader)) throw new IllegalStateException(
          "reader must be instance of " + BoboSegmentReader.class);
      final FacetDataCache<?> dataCache = _facetHandler.getFacetData((BoboSegmentReader) reader);
      return new DocComparator() {
        @Override
        public int compare(ScoreDoc doc1, ScoreDoc doc2) {
          int encoded1 = dataCache.orderArray.get(doc1.doc);
          int encoded2 = dataCache.orderArray.get(doc2.doc);
          return encoded1 - encoded2;
        }

        @Override
        public Comparable<?> value(ScoreDoc doc) {
          return new StringArrayComparator(_facetHandler.getFieldValues((BoboSegmentReader) reader,
            doc.doc));
        }
      };
    }
  }

  @Override
  public BoboDocScorer getDocScorer(BoboSegmentReader reader,
      FacetTermScoringFunctionFactory scoringFunctionFactory, Map<String, Float> boostMap) {
    FacetDataCache<?> dataCache = getFacetData(reader);
    float[] boostList = BoboDocScorer.buildBoostList(dataCache.valArray, boostMap);
    return new CompactMultiValueDocScorer(dataCache, scoringFunctionFactory, boostList);
  }

  private static final class CompactMultiValueDocScorer extends BoboDocScorer {
    private final FacetDataCache<?> _dataCache;

    CompactMultiValueDocScorer(FacetDataCache<?> dataCache,
        FacetTermScoringFunctionFactory scoreFunctionFactory, float[] boostList) {
      super(scoreFunctionFactory.getFacetTermScoringFunction(dataCache.valArray.size(),
        dataCache.orderArray.size()), boostList);
      _dataCache = dataCache;
    }

    @Override
    public Explanation explain(int doc) {
      int encoded = _dataCache.orderArray.get(doc);

      int count = 1;
      FloatList scoreList = new FloatArrayList(_dataCache.valArray.size());
      ArrayList<Explanation> explList = new ArrayList<Explanation>(scoreList.size());
      while (encoded != 0) {
        if ((encoded & 0x00000001) != 0x0) {
          int idx = count - 1;
          scoreList.add(_function.score(_dataCache.freqs[idx], _boostList[idx]));
          explList.add(_function.explain(_dataCache.freqs[idx], _boostList[idx]));
        }
        count++;
        encoded >>>= 1;
      }
      Explanation topLevel = _function.explain(scoreList.toFloatArray());
      for (Explanation sub : explList) {
        topLevel.addDetail(sub);
      }
      return topLevel;
    }

    @Override
    public final float score(int docid) {
      _function.clearScores();
      int encoded = _dataCache.orderArray.get(docid);

      int count = 1;

      while (encoded != 0) {
        int idx = count - 1;
        if ((encoded & 0x00000001) != 0x0) {
          _function.scoreAndCollect(_dataCache.freqs[idx], _boostList[idx]);
        }
        count++;
        encoded >>>= 1;
      }
      return _function.getCurrentScore();
    }

  }

  private static final class CompactMultiValueFacetCountCollector extends
      DefaultFacetCountCollector {
    private final BigSegmentedArray _array;
    private final int[] _combinationCount = new int[16 * 8];
    private int _noValCount = 0;
    private boolean _aggregated = false;

    CompactMultiValueFacetCountCollector(String name, BrowseSelection sel,
        FacetDataCache<?> dataCache, int docBase, FacetSpec ospec) {
      super(name, dataCache, docBase, sel, ospec);
      _array = _dataCache.orderArray;
    }

    @Override
    public final void collectAll() {
      _count = BigIntArray.fromArray(_dataCache.freqs);
      _aggregated = true;
    }

    @Override
    public final void collect(int docid) {
      int encoded = _array.get(docid);
      if (encoded == 0) {
        _noValCount++;
      } else {
        int offset = 0;
        while (true) {
          _combinationCount[(encoded & 0x0F) + offset]++;
          encoded = (encoded >>> 4);
          if (encoded == 0) break;
          offset += 16;
        }
      }
    }

    @Override
    public BrowseFacet getFacet(String value) {
      if (!_aggregated) aggregateCounts();
      return super.getFacet(value);
    }

    @Override
    public int getFacetHitsCount(Object value) {
      if (!_aggregated) aggregateCounts();
      return super.getFacetHitsCount(value);
    }

    @Override
    public BigSegmentedArray getCountDistribution() {
      if (!_aggregated) aggregateCounts();
      return _count;
    }

    @Override
    public List<BrowseFacet> getFacets() {
      if (!_aggregated) aggregateCounts();
      return super.getFacets();
    }

    private void aggregateCounts() {
      _count.add(0, _noValCount);

      for (int i = 1; i < _combinationCount.length; i++) {
        int count = _combinationCount[i];
        if (count > 0) {
          int offset = (i >> 4) * 4;
          int encoded = (i & 0x0F);
          int index = 1;
          while (encoded != 0) {
            if ((encoded & 0x00000001) != 0x0) {
              int idx = index + offset;
              _count.add(idx, _count.get(idx) + count);
            }
            index++;
            encoded >>>= 1;
          }
        }
      }
      _aggregated = true;
    }

    @Override
    public FacetIterator iterator() {
      if (!_aggregated) aggregateCounts();
      return super.iterator();
    }
  }
}