SortCollector.java

package com.browseengine.bobo.sort;

import java.io.IOException;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.Set;

import javax.management.MBeanServer;
import javax.management.ObjectName;
import javax.management.StandardMBean;

import org.apache.log4j.Logger;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.SortField;

import com.browseengine.bobo.api.BoboCustomSortField;
import com.browseengine.bobo.api.BoboSegmentReader;
import com.browseengine.bobo.api.Browsable;
import com.browseengine.bobo.api.BrowseHit;
import com.browseengine.bobo.api.FacetAccessible;
import com.browseengine.bobo.facets.FacetHandler;
import com.browseengine.bobo.jmx.JMXUtil;
import com.browseengine.bobo.sort.DocComparatorSource.DocIdDocComparatorSource;
import com.browseengine.bobo.sort.DocComparatorSource.RelevanceDocComparatorSource;
import com.browseengine.bobo.util.MemoryManager;
import com.browseengine.bobo.util.MemoryManagerAdminMBean;

public abstract class SortCollector extends Collector {
  private static final Logger logger = Logger.getLogger(SortCollector.class);

  protected static MemoryManager<int[]> intarraymgr = new MemoryManager<int[]>(
      new MemoryManager.Initializer<int[]>() {
        @Override
        public void init(int[] buf) {
          Arrays.fill(buf, 0);
        }

        @Override
        public int[] newInstance(int size) {
          return new int[size];
        }

        @Override
        public int size(int[] buf) {
          assert buf != null;
          return buf.length;
        }

      });

  protected static MemoryManager<float[]> floatarraymgr = new MemoryManager<float[]>(
      new MemoryManager.Initializer<float[]>() {
        @Override
        public void init(float[] buf) {
          Arrays.fill(buf, 0);
        }

        @Override
        public float[] newInstance(int size) {
          return new float[size];
        }

        @Override
        public int size(float[] buf) {
          assert buf != null;
          return buf.length;
        }

      });

  static {
    try {
      // register memory manager mbean
      MBeanServer mbeanServer = java.lang.management.ManagementFactory.getPlatformMBeanServer();
      ObjectName mbeanName = new ObjectName(JMXUtil.JMX_DOMAIN, "name",
          "SortCollectorImpl-MemoryManager-Int");
      StandardMBean mbean = new StandardMBean(intarraymgr.getAdminMBean(),
          MemoryManagerAdminMBean.class);
      mbeanServer.registerMBean(mbean, mbeanName);

      mbeanName = new ObjectName(JMXUtil.JMX_DOMAIN, "name",
          "SortCollectorImpl-MemoryManager-Float");
      mbean = new StandardMBean(floatarraymgr.getAdminMBean(), MemoryManagerAdminMBean.class);
      mbeanServer.registerMBean(mbean, mbeanName);
    } catch (Exception e) {
      logger.error(e.getMessage(), e);
    }
  }

  public static class CollectorContext {
    public BoboSegmentReader reader;
    public int base;
    public DocComparator comparator;
    public int length;

    public CollectorContext(BoboSegmentReader reader, int base, DocComparator comparator) {
      this.reader = reader;
      this.base = base;
      this.comparator = comparator;
    }
  }

  public FacetHandler<?> groupBy = null; // Point to the first element of groupByMulti to avoid
                                         // array lookups.
  public FacetHandler<?>[] groupByMulti = null;

  public LinkedList<CollectorContext> contextList;
  public LinkedList<int[]> docidarraylist;
  public LinkedList<float[]> scorearraylist;

  public static int BLOCK_SIZE = 4096;

  protected Collector _collector = null;
  protected final SortField[] _sortFields;
  protected final boolean _fetchStoredFields;
  protected boolean _closed = false;

  protected SortCollector(SortField[] sortFields, boolean fetchStoredFields) {
    _sortFields = sortFields;
    _fetchStoredFields = fetchStoredFields;
  }

  abstract public BrowseHit[] topDocs() throws IOException;

  abstract public int getTotalHits();

  abstract public int getTotalGroups();

  abstract public FacetAccessible[] getGroupAccessibles();

  private static DocComparatorSource getNonFacetComparatorSource(SortField sf) {
    String fieldname = sf.getField();

    SortField.Type type = sf.getType();

    switch (type) {
    case INT:
      return new DocComparatorSource.IntDocComparatorSource(fieldname);

    case FLOAT:
      return new DocComparatorSource.FloatDocComparatorSource(fieldname);

    case LONG:
      return new DocComparatorSource.LongDocComparatorSource(fieldname);

    case DOUBLE:
      return new DocComparatorSource.DoubleDocComparatorSource(fieldname);

    case BYTE:
      return new DocComparatorSource.ByteDocComparatorSource(fieldname);

    case SHORT:
      return new DocComparatorSource.ShortDocComparatorSource(fieldname);

    case STRING:
      return new DocComparatorSource.StringOrdComparatorSource(fieldname);

    case STRING_VAL:
      return new DocComparatorSource.StringValComparatorSource(fieldname);

    case CUSTOM:
      throw new IllegalArgumentException("lucene custom sort no longer supported: " + fieldname);

    default:
      throw new IllegalStateException("Illegal sort type: " + type + ", for field: " + fieldname);
    }
  }

  private static DocComparatorSource getComparatorSource(Browsable browser, SortField sf) {
    DocComparatorSource compSource = null;
    if (SortField.FIELD_DOC.equals(sf)) {
      compSource = new DocIdDocComparatorSource();
    } else if (SortField.FIELD_SCORE.equals(sf) || sf.getType() == SortField.Type.SCORE) {
      // we want to do reverse sorting regardless for relevance
      compSource = new ReverseDocComparatorSource(new RelevanceDocComparatorSource());
    } else if (sf instanceof BoboCustomSortField) {
      BoboCustomSortField custField = (BoboCustomSortField) sf;
      DocComparatorSource src = custField.getCustomComparatorSource();
      assert src != null;
      compSource = src;
    } else {
      Set<String> facetNames = browser.getFacetNames();
      String sortName = sf.getField();
      if (facetNames.contains(sortName)) {
        FacetHandler<?> handler = browser.getFacetHandler(sortName);
        assert handler != null;
        compSource = handler.getDocComparatorSource();
      } else { // default lucene field
        logger.info("doing default lucene sort for: " + sf);
        compSource = getNonFacetComparatorSource(sf);
      }
    }
    boolean reverse = sf.getReverse();
    if (reverse) {
      compSource = new ReverseDocComparatorSource(compSource);
    }
    compSource.setReverse(reverse);
    return compSource;
  }

  private static SortField convert(Browsable browser, SortField sort) {
    String field = sort.getField();
    FacetHandler<?> facetHandler = browser.getFacetHandler(field);
    if (facetHandler != null) {
      BoboCustomSortField sortField = new BoboCustomSortField(field, sort.getReverse(),
          facetHandler.getDocComparatorSource());
      return sortField;
    } else {
      return sort;
    }
  }

  public static SortCollector buildSortCollector(Browsable browser, Query q, SortField[] sort,
      int offset, int count, boolean fetchStoredFields, Set<String> termVectorsToFetch,
      String[] groupBy, int maxPerGroup, boolean collectDocIdCache) {
    if (sort == null || sort.length == 0) {
      if (q != null && !(q instanceof MatchAllDocsQuery)) {
        sort = new SortField[] { SortField.FIELD_SCORE };
      } else {
        sort = new SortField[] { SortField.FIELD_DOC };
      }
    }

    boolean doScoring = false;
    for (SortField sf : sort) {
      if (sf.getType() == SortField.Type.SCORE) {
        doScoring = true;
        break;
      }
    }

    DocComparatorSource compSource;
    if (sort.length == 1) {
      SortField sf = convert(browser, sort[0]);
      compSource = getComparatorSource(browser, sf);
    } else {
      DocComparatorSource[] compSources = new DocComparatorSource[sort.length];
      for (int i = 0; i < sort.length; ++i) {
        compSources[i] = getComparatorSource(browser, convert(browser, sort[i]));
      }
      compSource = new MultiDocIdComparatorSource(compSources);
    }
    return new SortCollectorImpl(compSource, sort, browser, offset, count, doScoring,
        fetchStoredFields, termVectorsToFetch, groupBy, maxPerGroup, collectDocIdCache);
  }

  public SortCollector setCollector(Collector collector) {
    _collector = collector;
    return this;
  }

  public Collector getCollector() {
    return _collector;
  }

  public void close() {
    if (!_closed) {
      _closed = true;
      if (docidarraylist != null) {
        while (!docidarraylist.isEmpty()) {
          intarraymgr.release(docidarraylist.poll());
        }
      }
      if (scorearraylist != null) {
        while (!scorearraylist.isEmpty()) {
          floatarraymgr.release(scorearraylist.poll());
        }
      }
    }
  }
}