<commit_before>/*
 * 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.solr.response;

import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.index.AtomicReader;
import org.apache.lucene.index.MultiDocValues;
import org.apache.lucene.index.SortedDocValues;
import org.apache.lucene.index.NumericDocValues;
import org.apache.lucene.index.SortedSetDocValues;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.SortField;
import org.apache.lucene.util.ArrayUtil;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.CharsRef;
import org.apache.lucene.util.CharsRefBuilder;
import org.apache.lucene.util.LongValues;
import org.apache.lucene.util.FixedBitSet;
import org.apache.solr.common.params.SolrParams;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.request.SolrRequestInfo;
import org.apache.solr.schema.IndexSchema;
import org.apache.solr.schema.SchemaField;
import org.apache.solr.schema.FieldType;
import org.apache.solr.schema.TrieFloatField;
import org.apache.solr.schema.TrieDoubleField;
import org.apache.solr.schema.TrieIntField;
import org.apache.solr.schema.TrieLongField;
import org.apache.solr.schema.StrField;
import org.apache.solr.search.SolrIndexSearcher;
import org.apache.solr.search.SortSpec;
import org.apache.solr.search.SyntaxError;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.Writer;
import java.io.PrintWriter;
import java.net.SocketException;
import java.util.List;


public class SortingResponseWriter implements QueryResponseWriter {

  private final static Logger logger = LoggerFactory.getLogger(SortingResponseWriter.class);

  public void init(NamedList args) {
    /* NOOP */
  }

  public String getContentType(SolrQueryRequest req, SolrQueryResponse res) {
    return "application/json";
  }

  public void write(Writer writer, SolrQueryRequest req, SolrQueryResponse res) throws IOException {
    Exception e1 = res.getException();
    if(e1 != null) {
      e1.printStackTrace(new PrintWriter(writer));
      return;
    }
    SolrRequestInfo info = SolrRequestInfo.getRequestInfo();
    SortSpec sortSpec = info.getResponseBuilder().getSortSpec();

    if(sortSpec == null) {
      throw new IOException(new SyntaxError("No sort criteria was provided."));
    }

    SolrIndexSearcher searcher = req.getSearcher();
    Sort sort = searcher.weightSort(sortSpec.getSort());

    if(sort == null) {
      throw new IOException(new SyntaxError("No sort criteria was provided."));
    }

    if(sort.needsScores()) {
      throw new IOException(new SyntaxError("Scoring is not currently supported with xsort."));
    }

    FixedBitSet[] sets = (FixedBitSet[])req.getContext().get("export");
    Integer th = (Integer)req.getContext().get("totalHits");

    if(sets == null) {
      throw new IOException(new SyntaxError("xport RankQuery is required for xsort: rq={!xport}"));
    }

    int totalHits = th.intValue();
    SolrParams params = req.getParams();
    String fl = params.get("fl");

    if(fl == null) {
      throw new IOException(new SyntaxError("export field list (fl) must be specified."));
    }

    String[] fields = fl.split(",");

    for(int i=0;i<fields.length; i++) {
      if(fl.trim().equals("score")) {
        throw new IOException(new SyntaxError("Scoring is not currently supported with xsort."));
      }
    }

    FieldWriter[] fieldWriters = getFieldWriters(fields, req.getSearcher());
    writer.write("{\"numFound\":"+totalHits+", \"docs\":[");

    //Write the data.
    List<AtomicReaderContext> leaves = req.getSearcher().getTopReaderContext().leaves();
    SortDoc sortDoc = getSortDoc(req.getSearcher(), sort.getSort());
    int count = 0;
    int queueSize = 30000;
    SortQueue queue = new SortQueue(queueSize, sortDoc);
    SortDoc[] outDocs = new SortDoc[queueSize];

    long total = 0;

    while(count < totalHits) {
      //long begin = System.nanoTime();
      boolean commaNeeded = false;
      queue.reset();
      SortDoc top = queue.top();
      for(int i=0; i<leaves.size(); i++) {
        sortDoc.setNextReader(leaves.get(i));
        DocIdSetIterator it = sets[i].iterator();
        int docId = -1;
        while((docId = it.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
          sortDoc.setValues(docId);
          if(top.lessThan(sortDoc)) {
            top.setValues(sortDoc);
            top = queue.updateTop();
          }
        }
      }

      int outDocsIndex = -1;

      for(int i=0; i<queueSize; i++) {
        SortDoc s = queue.pop();
        if(s.docId > -1) {
          outDocs[++outDocsIndex] = s;
        }
      }

     //long end = System.nanoTime();

      count += (outDocsIndex+1);

      try {
        for(int i=outDocsIndex; i>=0; --i) {
          SortDoc s = outDocs[i];
          if(commaNeeded){writer.write(',');}
          writer.write('{');
          writeDoc(s, leaves, fieldWriters, sets, writer);
          writer.write('}');
          commaNeeded = true;
          s.reset();
        }
      } catch(Throwable e) {
        Throwable ex = e;
        while(ex != null) {
          String m = ex.getMessage();
          if(m != null && m.contains("Broken pipe")) {
            logger.info("Early client disconnect during export");
            return;
          }
          ex = ex.getCause();
        }

        if(e instanceof IOException) {
          throw ((IOException)e);
        } else {
          throw new IOException(e);
        }
      }
    }

    //System.out.println("Sort Time 2:"+Long.toString(total/1000000));
    writer.write("]}");
    writer.flush();
  }




  protected void writeDoc(SortDoc sortDoc,
                          List<AtomicReaderContext> leaves,
                          FieldWriter[] fieldWriters,
                          FixedBitSet[] sets,<commit_msg> SOLR-5244: Exporting Full Sorted Result Sets<commit_after>/*
 * 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.solr.response;

import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.index.AtomicReader;
import org.apache.lucene.index.MultiDocValues;
import org.apache.lucene.index.SortedDocValues;
import org.apache.lucene.index.NumericDocValues;
import org.apache.lucene.index.SortedSetDocValues;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.SortField;
import org.apache.lucene.util.ArrayUtil;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.CharsRef;
import org.apache.lucene.util.CharsRefBuilder;
import org.apache.lucene.util.LongValues;
import org.apache.lucene.util.FixedBitSet;
import org.apache.solr.common.params.SolrParams;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.request.SolrRequestInfo;
import org.apache.solr.schema.IndexSchema;
import org.apache.solr.schema.SchemaField;
import org.apache.solr.schema.FieldType;
import org.apache.solr.schema.TrieFloatField;
import org.apache.solr.schema.TrieDoubleField;
import org.apache.solr.schema.TrieIntField;
import org.apache.solr.schema.TrieLongField;
import org.apache.solr.schema.StrField;
import org.apache.solr.search.SolrIndexSearcher;
import org.apache.solr.search.SortSpec;
import org.apache.solr.search.SyntaxError;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.Writer;
import java.io.PrintWriter;
import java.net.SocketException;
import java.util.List;


public class SortingResponseWriter implements QueryResponseWriter {

  private final static Logger logger = LoggerFactory.getLogger(SortingResponseWriter.class);

  public void init(NamedList args) {
    /* NOOP */
  }

  public String getContentType(SolrQueryRequest req, SolrQueryResponse res) {
    return "application/json";
  }

  public void write(Writer writer, SolrQueryRequest req, SolrQueryResponse res) throws IOException {
    Exception e1 = res.getException();
    if(e1 != null) {
      if(!(e1 instanceof IgnoreException)) {
        e1.printStackTrace(new PrintWriter(writer));
      }
      return;
    }
    SolrRequestInfo info = SolrRequestInfo.getRequestInfo();
    SortSpec sortSpec = info.getResponseBuilder().getSortSpec();

    if(sortSpec == null) {
      throw new IOException(new SyntaxError("No sort criteria was provided."));
    }

    SolrIndexSearcher searcher = req.getSearcher();
    Sort sort = searcher.weightSort(sortSpec.getSort());

    if(sort == null) {
      throw new IOException(new SyntaxError("No sort criteria was provided."));
    }

    if(sort.needsScores()) {
      throw new IOException(new SyntaxError("Scoring is not currently supported with xsort."));
    }

    FixedBitSet[] sets = (FixedBitSet[])req.getContext().get("export");
    Integer th = (Integer)req.getContext().get("totalHits");

    if(sets == null) {
      throw new IOException(new SyntaxError("xport RankQuery is required for xsort: rq={!xport}"));
    }

    int totalHits = th.intValue();
    SolrParams params = req.getParams();
    String fl = params.get("fl");

    if(fl == null) {
      throw new IOException(new SyntaxError("export field list (fl) must be specified."));
    }

    String[] fields = fl.split(",");

    for(int i=0;i<fields.length; i++) {
      if(fl.trim().equals("score")) {
        throw new IOException(new SyntaxError("Scoring is not currently supported with xsort."));
      }
    }

    FieldWriter[] fieldWriters = getFieldWriters(fields, req.getSearcher());
    writer.write("{\"numFound\":"+totalHits+", \"docs\":[");

    //Write the data.
    List<AtomicReaderContext> leaves = req.getSearcher().getTopReaderContext().leaves();
    SortDoc sortDoc = getSortDoc(req.getSearcher(), sort.getSort());
    int count = 0;
    int queueSize = 30000;
    SortQueue queue = new SortQueue(queueSize, sortDoc);
    SortDoc[] outDocs = new SortDoc[queueSize];

    long total = 0;

    while(count < totalHits) {
      //long begin = System.nanoTime();
      boolean commaNeeded = false;
      queue.reset();
      SortDoc top = queue.top();
      for(int i=0; i<leaves.size(); i++) {
        sortDoc.setNextReader(leaves.get(i));
        DocIdSetIterator it = sets[i].iterator();
        int docId = -1;
        while((docId = it.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
          sortDoc.setValues(docId);
          if(top.lessThan(sortDoc)) {
            top.setValues(sortDoc);
            top = queue.updateTop();
          }
        }
      }

      int outDocsIndex = -1;

      for(int i=0; i<queueSize; i++) {
        SortDoc s = queue.pop();
        if(s.docId > -1) {
          outDocs[++outDocsIndex] = s;
        }
      }

     //long end = System.nanoTime();

      count += (outDocsIndex+1);

      try {
        for(int i=outDocsIndex; i>=0; --i) {
          SortDoc s = outDocs[i];
          if(commaNeeded){writer.write(',');}
          writer.write('{');
          writeDoc(s, leaves, fieldWriters, sets, writer);
          writer.write('}');
          commaNeeded = true;
          s.reset();
        }
      } catch(Throwable e) {
        Throwable ex = e;
        while(ex != null) {
          String m = ex.getMessage();
          if(m != null && m.contains("Broken pipe")) {
            throw new IgnoreException();
          }
          ex = ex.getCause();
        }

        if(e instanceof IOException) {
          throw ((IOException)e);
        } else {
          throw new IOException(e);
        }
      }
    }

    //System.out.println("Sort Time 2:"+Long.toString(total/1000000));
    writer.write("]}");
    writer.flush();
  }

  public static class IgnoreException extends IOException {
    public void printStackTrace(PrintWriter pw) {
      pw.print("Early Client Disconnect");

    }

    public String getMessage() {
      return "Early Client Disconnect";
    }
  }


  protected void writeDoc(SortDoc sortDoc,
                          List<AtomicReaderContext> leaves,
                          FieldWriter[] fieldWriters,
                          FixedBitSet[] sets,
{'fraction_non_alphanumeric': 0.12098711509063115, 'fraction_numerical': 0.005022930770910679, 'mean_word_length': 3.089907710628163, 'pattern_counts': {'":': 4, '<': 15, '<?xml version=': 0, '>': 11, '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': '680088', 'n_tokens_mistral': 4232, 'n_tokens_neox': 4103, 'n_words': 983}