<commit_before>/*
 This software was produced for the U. S. Government
 under Contract No. W15P7T-11-C-F600, and is
 subject to the Rights in Noncommercial Computer Software
 and Noncommercial Computer Software Documentation
 Clause 252.227-7014 (JUN 1995)

 Copyright 2013 The MITRE Corporation. All Rights Reserved.

 Licensed 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.mitre.solr.tagger;

import com.google.common.io.CharStreams;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.index.ReaderUtil;
import org.apache.lucene.queries.function.FunctionValues;
import org.apache.lucene.queries.function.ValueSource;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Sort;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.IntsRef;
import org.apache.lucene.util.OpenBitSet;
import org.apache.solr.common.SolrException;
import org.apache.solr.common.params.CommonParams;
import org.apache.solr.common.params.SolrParams;
import org.apache.solr.common.util.ContentStream;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.handler.RequestHandlerBase;
import org.apache.solr.request.LocalSolrQueryRequest;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.response.SolrQueryResponse;
import org.apache.solr.schema.SchemaField;
import org.apache.solr.search.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.io.Reader;
import java.io.StringReader;
import java.util.*;

/**
 * @author <NAME> - <EMAIL>
 */
public class TaggerRequestHandler extends RequestHandlerBase {

  private static final String OVERLAPS = "overlaps";
  private final Logger log = LoggerFactory.getLogger(getClass());

  /** Request parameter. */
  public static final String TAGS_LIMIT = "tagsLimit";
  /** Request parameter. */
  public static final String SUB_TAGS = "subTags";//deprecated
  private static final String MATCH_TEXT = "matchText";

  private TaggerFstCorpus _corpus;//synchronized access

  @Override
  public void handleRequestBody(SolrQueryRequest req, SolrQueryResponse rsp) throws Exception {
    boolean build = req.getParams().getBool("build", false);
    TaggerFstCorpus corpus = getCorpus(build, req, rsp);
    if (build)//just build; that's it.
      return;

    final TagClusterReducer tagClusterReducer;
    String overlaps = req.getParams().get(OVERLAPS);
    if (overlaps == null) {//deprecated; should always be specified
      if (req.getParams().getBool(SUB_TAGS, false))//deprecated
        tagClusterReducer = TagClusterReducer.ALL;
      else
        tagClusterReducer = TagClusterReducer.NO_SUB;
    } else if (overlaps.equals("ALL")) {
      tagClusterReducer = TagClusterReducer.ALL;
    } else if (overlaps.equals("NO_SUB")) {
      tagClusterReducer = TagClusterReducer.NO_SUB;
    } else if (overlaps.equals("LONGEST_DOMINANT_RIGHT")) {
      tagClusterReducer = TagClusterReducer.LONGEST_DOMINANT_RIGHT;
    } else {
      throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
          "unknown tag overlap mode: "+overlaps);
    }

    final int rows = req.getParams().getInt(CommonParams.ROWS, 10000);
    final int tagsLimit = req.getParams().getInt(TAGS_LIMIT, 1000);
    final boolean addMatchText = req.getParams().getBool(MATCH_TEXT, false);
    final String indexedField = corpus.getIndexedField();
    final SchemaField idSchemaField = req.getSchema().getUniqueKeyField();

    //Get posted data
    Reader reader = null;
    Iterable<ContentStream> streams = req.getContentStreams();
    if (streams != null) {
      Iterator<ContentStream> iter = streams.iterator();
      if (iter.hasNext()) {
        reader = iter.next().getReader();
      }
      if (iter.hasNext()) {
        throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
            getClass().getSimpleName()+" does not support multiple ContentStreams");
      }
    }
    if (reader == null) {
      throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
          getClass().getSimpleName()+" requires text to be POSTed to it");
    }
    final String bufferedInput;
    if (addMatchText) {
      //read the input fully into a String buffer we'll use later to get
      // the match text, then replace the input with a reader wrapping the buffer.
      bufferedInput = CharStreams.toString(reader);//(closes reader)
      reader = new StringReader(bufferedInput);
    } else {
      bufferedInput = null;//not used
    }

    final SolrIndexSearcher searcher = req.getSearcher();
    final OpenBitSet matchDocIds = new OpenBitSet(searcher.maxDoc());
    final List tags = new ArrayList(2000);

    try {
      Analyzer analyzer = req.getSchema().getField(indexedField).getType().getAnalyzer();

      new Tagger(corpus, analyzer, reader, tagClusterReducer) {
        @SuppressWarnings("unchecked")
        @Override
        protected void tagCallback(int startOffset, int endOffset, long docIdsKey) {
          if (tags.size() >= tagsLimit)
            return;
          NamedList tag = new NamedList();
          tag.add("startOffset", startOffset);
          tag.add("endOffset", endOffset);
          if (addMatchText)
            tag.add("matchText", bufferedInput.substring(startOffset,
                endOffset));
          //below caches, and also flags matchDocIds
          tag.add("ids", lookupSchemaDocIds(docIdsKey));
          tags.add(tag);
        }

        Map<Long,List> docIdsListCache = new HashMap<Long, List>(2000);

        ValueSourceAccessor uniqueKeyCache = new ValueSourceAccessor(searcher,
            idSchemaField.getType().getValueSource(idSchemaField, null));

        @SuppressWarnings("unchecked")
        private List lookupSchemaDocIds(long docIdsKey) {
          List schemaDocIds = docIdsListCache.get(docIdsKey);
          if (schemaDocIds != null)
            return schemaDocIds;
          IntsRef docIds = lookupDocIds(docIdsKey);
          //translate lucene docIds to schema ids
          schemaDocIds = new ArrayList(docIds.length);
          for (int i = docIds.offset; i < docIds.offset + docIds.length; i++) {
            int docId = docIds.ints[i];
            matchDocIds.set(docId);//also, flip docid in bitset
            schemaDocIds.add(uniqueKeyCache.objectVal(docId));//translates here
          }
          docIdsListCache.put(docIdsKey, schemaDocIds);
          return schemaDocIds;
        }

      }.process();
    } finally {
      reader.close();
    }
    rsp.add("tagsCount",tags.size());
    rsp.add("tags", tags);

    ReturnFields returnFields = new SolrReturnFields( req );
    rsp.setReturnFields( returnFields );

    DocList docs = searcher.getDocList(
        new MatchAllDocsQuery(),
        new BitDocSet(matchDocIds),
        Sort.INDEXORDER, 0, rows);
    rsp.add("matchingDocs",docs);
  }

  /** Gets the corpus if it's ready and not stale, otherwise initializes it. */
  private synchronized TaggerFstCorpus getCorpus(boolean forceBuild,
      SolrQueryRequest req, SolrQueryResponse rsp) throws IOException {
    long indexVersion = req.getSearcher().getIndexReader().getVersion();
    if (_corpus != null && indexVersion != _corpus.getIndexVersion())
      forceBuild = true;
    if (forceBuild || _corpus == null) {
      rsp.setHttpCaching(false);
      _corpus = null;//help GC
      _corpus = initCorpus(req, forceBuild);
    }
    return _corpus;
  }

  private TaggerFstCorpus initCorpus(SolrQueryRequest req, boolean forceRebuild) throws IOException {
    SolrParams params = SolrParams.wrapDefaults(req.getParams(), SolrParams.toSolrParams(getInitArgs()));
    //--load params
    String indexedField = params.get("indexedField");
    if (indexedField == null)
      throw new RuntimeException("required param 'indexedField'");
    String storedField = params.get("storedField", indexedField);
    boolean partialMatches = params.getBool("partialMatches", false);
    String taggerCacheFile = params.get("cacheFile");
    File cacheFile = (taggerCacheFile == null ? null : new File(req.getCore().getDataDir(), taggerCacheFile));<commit_msg>Instead of "searching" for the matched docs, directly create the Solr DocList because the tagger knows exactly which docs matched.<commit_after>/*
 This software was produced for the U. S. Government
 under Contract No. W15P7T-11-C-F600, and is
 subject to the Rights in Noncommercial Computer Software
 and Noncommercial Computer Software Documentation
 Clause 252.227-7014 (JUN 1995)

 Copyright 2013 The MITRE Corporation. All Rights Reserved.

 Licensed 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.mitre.solr.tagger;

import com.google.common.io.CharStreams;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.index.ReaderUtil;
import org.apache.lucene.queries.function.FunctionValues;
import org.apache.lucene.queries.function.ValueSource;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.IntsRef;
import org.apache.lucene.util.OpenBitSet;
import org.apache.solr.common.SolrException;
import org.apache.solr.common.params.CommonParams;
import org.apache.solr.common.params.SolrParams;
import org.apache.solr.common.util.ContentStream;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.handler.RequestHandlerBase;
import org.apache.solr.request.LocalSolrQueryRequest;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.response.SolrQueryResponse;
import org.apache.solr.schema.SchemaField;
import org.apache.solr.search.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.io.Reader;
import java.io.StringReader;
import java.util.*;

/**
 * @author <NAME> - <EMAIL>
 */
public class TaggerRequestHandler extends RequestHandlerBase {

  private static final String OVERLAPS = "overlaps";
  private final Logger log = LoggerFactory.getLogger(getClass());

  /** Request parameter. */
  public static final String TAGS_LIMIT = "tagsLimit";
  /** Request parameter. */
  public static final String SUB_TAGS = "subTags";//deprecated
  private static final String MATCH_TEXT = "matchText";

  private TaggerFstCorpus _corpus;//synchronized access

  @Override
  public void handleRequestBody(SolrQueryRequest req, SolrQueryResponse rsp) throws Exception {
    boolean build = req.getParams().getBool("build", false);
    TaggerFstCorpus corpus = getCorpus(build, req, rsp);
    if (build)//just build; that's it.
      return;

    final TagClusterReducer tagClusterReducer;
    String overlaps = req.getParams().get(OVERLAPS);
    if (overlaps == null) {//deprecated; should always be specified
      if (req.getParams().getBool(SUB_TAGS, false))//deprecated
        tagClusterReducer = TagClusterReducer.ALL;
      else
        tagClusterReducer = TagClusterReducer.NO_SUB;
    } else if (overlaps.equals("ALL")) {
      tagClusterReducer = TagClusterReducer.ALL;
    } else if (overlaps.equals("NO_SUB")) {
      tagClusterReducer = TagClusterReducer.NO_SUB;
    } else if (overlaps.equals("LONGEST_DOMINANT_RIGHT")) {
      tagClusterReducer = TagClusterReducer.LONGEST_DOMINANT_RIGHT;
    } else {
      throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
          "unknown tag overlap mode: "+overlaps);
    }

    final int rows = req.getParams().getInt(CommonParams.ROWS, 10000);
    final int tagsLimit = req.getParams().getInt(TAGS_LIMIT, 1000);
    final boolean addMatchText = req.getParams().getBool(MATCH_TEXT, false);
    final String indexedField = corpus.getIndexedField();
    final SchemaField idSchemaField = req.getSchema().getUniqueKeyField();

    //Get posted data
    Reader reader = null;
    Iterable<ContentStream> streams = req.getContentStreams();
    if (streams != null) {
      Iterator<ContentStream> iter = streams.iterator();
      if (iter.hasNext()) {
        reader = iter.next().getReader();
      }
      if (iter.hasNext()) {
        throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
            getClass().getSimpleName()+" does not support multiple ContentStreams");
      }
    }
    if (reader == null) {
      throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
          getClass().getSimpleName()+" requires text to be POSTed to it");
    }
    final String bufferedInput;
    if (addMatchText) {
      //read the input fully into a String buffer we'll use later to get
      // the match text, then replace the input with a reader wrapping the buffer.
      bufferedInput = CharStreams.toString(reader);//(closes reader)
      reader = new StringReader(bufferedInput);
    } else {
      bufferedInput = null;//not used
    }

    final SolrIndexSearcher searcher = req.getSearcher();
    final OpenBitSet matchDocIdsBS = new OpenBitSet(searcher.maxDoc());
    final List tags = new ArrayList(2000);

    try {
      Analyzer analyzer = req.getSchema().getField(indexedField).getType().getAnalyzer();

      new Tagger(corpus, analyzer, reader, tagClusterReducer) {
        @SuppressWarnings("unchecked")
        @Override
        protected void tagCallback(int startOffset, int endOffset, long docIdsKey) {
          if (tags.size() >= tagsLimit)
            return;
          NamedList tag = new NamedList();
          tag.add("startOffset", startOffset);
          tag.add("endOffset", endOffset);
          if (addMatchText)
            tag.add("matchText", bufferedInput.substring(startOffset,
                endOffset));
          //below caches, and also flags matchDocIdsBS
          tag.add("ids", lookupSchemaDocIds(docIdsKey));
          tags.add(tag);
        }

        Map<Long,List> docIdsListCache = new HashMap<Long, List>(2000);

        ValueSourceAccessor uniqueKeyCache = new ValueSourceAccessor(searcher,
            idSchemaField.getType().getValueSource(idSchemaField, null));

        @SuppressWarnings("unchecked")
        private List lookupSchemaDocIds(long docIdsKey) {
          List schemaDocIds = docIdsListCache.get(docIdsKey);
          if (schemaDocIds != null)
            return schemaDocIds;
          IntsRef docIds = lookupDocIds(docIdsKey);
          //translate lucene docIds to schema ids
          schemaDocIds = new ArrayList(docIds.length);
          for (int i = docIds.offset; i < docIds.offset + docIds.length; i++) {
            int docId = docIds.ints[i];
            matchDocIdsBS.set(docId);//also, flip docid in bitset
            schemaDocIds.add(uniqueKeyCache.objectVal(docId));//translates here
          }
          docIdsListCache.put(docIdsKey, schemaDocIds);
          return schemaDocIds;
        }

      }.process();
    } finally {
      reader.close();
    }
    rsp.add("tagsCount",tags.size());
    rsp.add("tags", tags);

    ReturnFields returnFields = new SolrReturnFields( req );
    rsp.setReturnFields( returnFields );

    //Now we must supply a Solr DocList and add it to the response.
    //  Typically this is gotten via a SolrIndexSearcher.search(), but in this case we
    //  know exactly what documents to return, the order doesn't matter nor does
    //  scoring.
    //  Ideally an implementation of DocList could be directly implemented off
    //  of a BitSet, but there are way too many methods to implement for a minor
    //  payoff.
    int matchDocs = (int) matchDocIdsBS.cardinality();
    int[] docIds = new int[ Math.min(rows, matchDocs) ];
    DocIdSetIterator docIdIter = matchDocIdsBS.iterator();
    for (int i = 0; i < docIds.length; i++) {
      docIds[i] = docIdIter.nextDoc();
    }
    DocList docs = new DocSlice(0, docIds.length, docIds, null, matchDocs, 1f);
    rsp.add("matchingDocs", docs);
  }

  /** Gets the corpus if it's ready and not stale, otherwise initializes it. */
  private synchronized TaggerFstCorpus getCorpus(boolean forceBuild,
      SolrQueryRequest req, SolrQueryResponse rsp) throws IOException {
    long indexVersion = req.getSearcher().getIndexReader().getVersion();
    if (_corpus != null && indexVersion != _corpus.getIndexVersion())
      forceBuild = true;
    if (forceBuild || _corpus == null) {
      rsp.setHttpCaching(false);
      _corpus = null;//help GC
      _corpus = initCorpus(req, forceBuild);
    }
    return _corpus;
  }

  private TaggerFstCorpus initCorpus(SolrQueryRequest req, boolean forceRebuild) throws IOException {
    SolrParams params = SolrParams.wrapDefaults(req.getParams(), SolrParams.toSolrParams(getInitArgs()));
    //--load params
    String indexedField = params.get("indexedField");
    if (indexedField == null)
      throw new RuntimeException("required param 'indexedField'");
    String storedField = params.get("storedField", indexedField);
    boolean partialMatches = params.getBool("partialMatches", false);
    String taggerCacheFile = params.get("cacheFile");
    File cacheFile = (taggerCacheFile == null ? null : new File(req.getCore().getDataDir(), taggerCacheFile));
{'fraction_non_alphanumeric': 0.10328354930814847, 'fraction_numerical': 0.0056007028332967276, 'mean_word_length': 3.5418952618453865, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6887841', 'n_tokens_mistral': 5459, 'n_tokens_neox': 5238, 'n_words': 1376}