From 18c9215bc1356335849ea969262eb1f1047c96e9 Mon Sep 17 00:00:00 2001 From: Joel Bernstein Date: Wed, 14 May 2014 20:12:12 +0000 Subject: [PATCH] SOLR-5973: Pluggable Ranking Collectors and Merge Strategies git-svn-id: https://svn.apache.org/repos/asf/lucene/dev/trunk@1594698 13f79535-47bb-0310-9956-ffa450edef68 --- solr/CHANGES.txt | 3 + .../solr/handler/component/MergeStrategy.java | 76 ++ .../handler/component/QueryComponent.java | 38 +- .../handler/component/ResponseBuilder.java | 23 +- .../solr/handler/component/ShardDoc.java | 2 +- .../org/apache/solr/search/RankQuery.java | 29 + .../apache/solr/search/SolrIndexSearcher.java | 8 +- .../conf/solrconfig-plugcollector.xml | 579 +++++++++++++ .../apache/solr/search/MergeStrategyTest.java | 181 ++++ .../org/apache/solr/search/RankQueryTest.java | 115 +++ .../solr/search/TestRankQueryPlugin.java | 816 ++++++++++++++++++ 11 files changed, 1866 insertions(+), 4 deletions(-) create mode 100644 solr/core/src/java/org/apache/solr/handler/component/MergeStrategy.java create mode 100644 solr/core/src/java/org/apache/solr/search/RankQuery.java create mode 100644 solr/core/src/test-files/solr/collection1/conf/solrconfig-plugcollector.xml create mode 100644 solr/core/src/test/org/apache/solr/search/MergeStrategyTest.java create mode 100644 solr/core/src/test/org/apache/solr/search/RankQueryTest.java create mode 100644 solr/core/src/test/org/apache/solr/search/TestRankQueryPlugin.java diff --git a/solr/CHANGES.txt b/solr/CHANGES.txt index cca3c367d27..1963944ee4c 100644 --- a/solr/CHANGES.txt +++ b/solr/CHANGES.txt @@ -113,6 +113,9 @@ New Features * SOLR-6043: Add ability to set http headers in solr response (Tomás Fernández Löbbe via Ryan Ernst) +* SOLR-5973: Pluggable Ranking Collectors and Merge Strategies + (Joel Bernstein) + Bug Fixes ---------------------- diff --git a/solr/core/src/java/org/apache/solr/handler/component/MergeStrategy.java b/solr/core/src/java/org/apache/solr/handler/component/MergeStrategy.java new file mode 100644 index 00000000000..0ff19bd523e --- /dev/null +++ b/solr/core/src/java/org/apache/solr/handler/component/MergeStrategy.java @@ -0,0 +1,76 @@ +/* +* 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.handler.component; + +import org.apache.solr.search.SolrIndexSearcher; + +import java.util.Comparator; +import java.io.IOException; + +/** + * The MergeStrategy class defines custom merge logic for distributed searches. + **/ + + +public interface MergeStrategy { + + /** + * merge defines the merging behaving of results that are collected from the + * shards during a distributed search. + * + **/ + + public void merge(ResponseBuilder rb, ShardRequest sreq); + + /** + * mergesIds must return true if the merge method merges document ids from the shards. + * If it merges other output from the shards it must return false. + * */ + + public boolean mergesIds(); + + + /** + * handlesMergeFields must return true if the MergeStrategy + * implements a custom handleMergeFields(ResponseBuilder rb, SolrIndexSearch searcher) + * */ + + public boolean handlesMergeFields(); + + + /** + * Implement handleMergeFields(ResponseBuilder rb, SolrIndexSearch searcher) if + * your merge strategy needs more complex data then the sort fields provide. + * */ + + public void handleMergeFields(ResponseBuilder rb, SolrIndexSearcher searcher) throws IOException; + + /** + * Defines the order that the mergeStrategies are applied. Lower costs are applied first. + * */ + public int getCost(); + + public static final Comparator MERGE_COMP = new Comparator() { + public int compare(Object o1, Object o2) { + MergeStrategy m1 = (MergeStrategy)o1; + MergeStrategy m2 = (MergeStrategy)o2; + return m1.getCost()-m2.getCost(); + } + }; + +} \ No newline at end of file diff --git a/solr/core/src/java/org/apache/solr/handler/component/QueryComponent.java b/solr/core/src/java/org/apache/solr/handler/component/QueryComponent.java index 3be4f6437e5..ac2689b9d17 100644 --- a/solr/core/src/java/org/apache/solr/handler/component/QueryComponent.java +++ b/solr/core/src/java/org/apache/solr/handler/component/QueryComponent.java @@ -77,6 +77,7 @@ import org.apache.solr.search.ReturnFields; import org.apache.solr.search.SolrIndexSearcher; import org.apache.solr.search.SolrReturnFields; import org.apache.solr.search.SortSpec; +import org.apache.solr.search.RankQuery; import org.apache.solr.search.SyntaxError; import org.apache.solr.search.grouping.CommandHandler; import org.apache.solr.search.grouping.GroupingSpecification; @@ -98,6 +99,8 @@ import org.apache.solr.search.grouping.endresulttransformer.GroupedEndResultTran import org.apache.solr.search.grouping.endresulttransformer.MainEndResultTransformer; import org.apache.solr.search.grouping.endresulttransformer.SimpleEndResultTransformer; import org.apache.solr.util.SolrPluginUtils; +import java.util.Collections; +import java.util.Comparator; /** * TODO! @@ -147,6 +150,17 @@ public class QueryComponent extends SearchComponent // normalize a null query to a query that matches nothing q = new BooleanQuery(); } + + if(q instanceof RankQuery) { + MergeStrategy mergeStrategy = ((RankQuery)q).getMergeStrategy(); + if(mergeStrategy != null) { + rb.addMergeStrategy(mergeStrategy); + if(mergeStrategy.handlesMergeFields()) { + rb.mergeFieldHandler = mergeStrategy; + } + } + } + rb.setQuery( q ); rb.setSortSpec( parser.getSort(true) ); rb.setQparser(parser); @@ -473,7 +487,13 @@ public class QueryComponent extends SearchComponent rb.getNextCursorMark().getSerializedTotem()); } } - doFieldSortValues(rb, searcher); + + if(rb.mergeFieldHandler != null) { + rb.mergeFieldHandler.handleMergeFields(rb, searcher); + } else { + doFieldSortValues(rb, searcher); + } + doPrefetch(rb); } @@ -821,6 +841,22 @@ public class QueryComponent extends SearchComponent private void mergeIds(ResponseBuilder rb, ShardRequest sreq) { + List mergeStrategies = rb.getMergeStrategies(); + if(mergeStrategies != null) { + Collections.sort(mergeStrategies, MergeStrategy.MERGE_COMP); + boolean idsMerged = false; + for(MergeStrategy mergeStrategy : mergeStrategies) { + mergeStrategy.merge(rb, sreq); + if(mergeStrategy.mergesIds()) { + idsMerged = true; + } + } + + if(idsMerged) { + return; //ids were merged above so return. + } + } + SortSpec ss = rb.getSortSpec(); Sort sort = ss.getSort(); diff --git a/solr/core/src/java/org/apache/solr/handler/component/ResponseBuilder.java b/solr/core/src/java/org/apache/solr/handler/component/ResponseBuilder.java index 5fd9fa248b2..c30ac710c5e 100644 --- a/solr/core/src/java/org/apache/solr/handler/component/ResponseBuilder.java +++ b/solr/core/src/java/org/apache/solr/handler/component/ResponseBuilder.java @@ -40,6 +40,7 @@ import org.apache.solr.search.grouping.distributed.command.QueryCommandResult; import java.util.Collection; import java.util.HashMap; import java.util.List; +import java.util.ArrayList; import java.util.Map; import java.util.Set; @@ -58,6 +59,7 @@ public class ResponseBuilder public boolean doExpand; public boolean doStats; public boolean doTerms; + public MergeStrategy mergeFieldHandler; private boolean needDocList = false; private boolean needDocSet = false; @@ -74,6 +76,9 @@ public class ResponseBuilder private CursorMark cursorMark; private CursorMark nextCursorMark; + private List mergeStrategies; + + private DocListAndSet results = null; private NamedList debugInfo = null; private RTimer timer = null; @@ -230,7 +235,23 @@ public class ResponseBuilder debugResults = dbg; debugTrack = dbg; } - + + public void addMergeStrategy(MergeStrategy mergeStrategy) { + if(mergeStrategies == null) { + mergeStrategies = new ArrayList(); + } + + mergeStrategies.add(mergeStrategy); + } + + public List getMergeStrategies() { + return this.mergeStrategies; + } + + public void setResponseDocs(SolrDocumentList _responseDocs) { + this._responseDocs = _responseDocs; + } + public boolean isDebugTrack() { return debugTrack; } diff --git a/solr/core/src/java/org/apache/solr/handler/component/ShardDoc.java b/solr/core/src/java/org/apache/solr/handler/component/ShardDoc.java index ec1e33f20f2..97b831b3d7f 100644 --- a/solr/core/src/java/org/apache/solr/handler/component/ShardDoc.java +++ b/solr/core/src/java/org/apache/solr/handler/component/ShardDoc.java @@ -35,7 +35,7 @@ public class ShardDoc extends FieldDoc { public String shard; public String shardAddress; // TODO - int orderInShard; + public int orderInShard; // the position of this doc within the shard... this can be used // to short-circuit comparisons if the shard is equal, and can // also be used to break ties within the same shard. diff --git a/solr/core/src/java/org/apache/solr/search/RankQuery.java b/solr/core/src/java/org/apache/solr/search/RankQuery.java new file mode 100644 index 00000000000..da8c00a1387 --- /dev/null +++ b/solr/core/src/java/org/apache/solr/search/RankQuery.java @@ -0,0 +1,29 @@ +/* +* 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.search; + +import org.apache.lucene.search.TopDocsCollector; +import org.apache.lucene.search.Query; +import org.apache.solr.handler.component.MergeStrategy; + +public abstract class RankQuery extends Query { + + public abstract TopDocsCollector getTopDocsCollector(int len, SolrIndexSearcher.QueryCommand cmd); + public abstract MergeStrategy getMergeStrategy(); + +} \ No newline at end of file diff --git a/solr/core/src/java/org/apache/solr/search/SolrIndexSearcher.java b/solr/core/src/java/org/apache/solr/search/SolrIndexSearcher.java index 176c0df733f..1361bc710c5 100644 --- a/solr/core/src/java/org/apache/solr/search/SolrIndexSearcher.java +++ b/solr/core/src/java/org/apache/solr/search/SolrIndexSearcher.java @@ -1482,7 +1482,13 @@ public class SolrIndexSearcher extends IndexSearcher implements Closeable,SolrIn * TopDocsCollector to use. */ private TopDocsCollector buildTopDocsCollector(int len, QueryCommand cmd) throws IOException { - + + Query q = cmd.getQuery(); + if(q instanceof RankQuery) { + RankQuery rq = (RankQuery)q; + return rq.getTopDocsCollector(len, cmd); + } + if (null == cmd.getSort()) { assert null == cmd.getCursorMark() : "have cursor but no sort"; return TopScoreDocCollector.create(len, true); diff --git a/solr/core/src/test-files/solr/collection1/conf/solrconfig-plugcollector.xml b/solr/core/src/test-files/solr/collection1/conf/solrconfig-plugcollector.xml new file mode 100644 index 00000000000..f1431c61b48 --- /dev/null +++ b/solr/core/src/test-files/solr/collection1/conf/solrconfig-plugcollector.xml @@ -0,0 +1,579 @@ + + + + + + + + + + + + ${solr.data.dir:} + + + + 1000000 + 2000000 + 3000000 + 4000000 + + + ${tests.luceneMatchVersion:LUCENE_CURRENT} + + + + + + + + + + + + + ${solr.ulog.dir:} + + + + ${solr.commitwithin.softcommit:true} + + + + + + + 1024 + + + + + + + + + + + + true + + + + + + 10 + + + + + + + + + + + + + + + + + + + + + + + + + true + + + + + true + + + + + + dismax + *:* + 0.01 + + text^0.5 features_t^1.0 subject^1.4 title_stemmed^2.0 + + + text^0.2 features_t^1.1 subject^1.4 title_stemmed^2.0 title^1.5 + + + ord(weight)^0.5 recip(rord(iind),1,1000,1000)^0.3 + + + 3<-1 5<-2 6<90% + + 100 + + + + + + + + + + + 4 + true + text,name,subject,title,whitetok + + + + + + + 4 + true + text,name,subject,title,whitetok + + + + + + + + lowerpunctfilt + + + default + lowerfilt + spellchecker1 + false + + + direct + DirectSolrSpellChecker + lowerfilt + 3 + + + wordbreak + solr.WordBreakSolrSpellChecker + lowerfilt + true + true + 10 + + + multipleFields + lowerfilt1and2 + spellcheckerMultipleFields + false + + + + jarowinkler + lowerfilt + + org.apache.lucene.search.spell.JaroWinklerDistance + spellchecker2 + + + + solr.FileBasedSpellChecker + external + spellings.txt + UTF-8 + spellchecker3 + + + + freq + lowerfilt + spellcheckerFreq + + freq + false + + + fqcn + lowerfilt + spellcheckerFQCN + org.apache.solr.spelling.SampleComparator + false + + + perDict + org.apache.solr.handler.component.DummyCustomParamSpellChecker + lowerfilt + + + + + + + + termsComp + + + + + + + + + false + + false + + 1 + + + spellcheck + + + + + direct + false + false + 1 + + + spellcheck + + + + + default + wordbreak + 20 + + + spellcheck + + + + + direct + wordbreak + 20 + + + spellcheck + + + + + dismax + lowerfilt1^1 + + + spellcheck + + + + + + + + + + + + + + + tvComponent + + + + + + string + elevate.xml + + + + + + explicit + + + elevate + + + + + + + + + + + + + 100 + + + + + + 70 + + + + + + + ]]> + ]]> + + + + + + + + + + + + + 10 + .,!? + + + + + + WORD + en + US + + + + + + + + + + max-age=30, public + + + + + + + explicit + true + + + + + solr + solrconfig.xml schema.xml admin-extra.html + + + + prefix-${solr.test.sys.prop2}-suffix + + + + + + false + true + v_t,t_field + org.apache.solr.update.processor.TextProfileSignature + + + + + + false + false + id + + org.apache.solr.update.processor.Lookup3Signature + + + + + + + true + non_indexed_signature_sS + false + v_t,t_field + org.apache.solr.update.processor.TextProfileSignature + + + + + + + uniq + uniq2 + uniq3 + + + + + + + + + regex_dup_A_s + x + x_x + + + + regex_dup_B_s + x + x_x + + + + + + + + regex_dup_A_s + x + x_x + + + regex_dup_B_s + x + x_x + + + + + + diff --git a/solr/core/src/test/org/apache/solr/search/MergeStrategyTest.java b/solr/core/src/test/org/apache/solr/search/MergeStrategyTest.java new file mode 100644 index 00000000000..68b8c9e8a7d --- /dev/null +++ b/solr/core/src/test/org/apache/solr/search/MergeStrategyTest.java @@ -0,0 +1,181 @@ +package org.apache.solr.search; + +/* + * 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. + */ + +import org.apache.solr.BaseDistributedSearchTestCase; +import org.apache.solr.client.solrj.response.QueryResponse; +import org.apache.solr.common.SolrDocument; +import org.apache.solr.common.SolrDocumentList; +import org.apache.solr.common.params.CommonParams; +import org.apache.solr.common.params.ModifiableSolrParams; +import org.apache.solr.common.util.NamedList; +import org.apache.solr.handler.component.MergeStrategy; +import org.apache.solr.handler.component.ResponseBuilder; +import org.apache.solr.handler.component.ShardRequest; +import org.junit.BeforeClass; + +import java.util.Arrays; + +/** + * Test for QueryComponent's distributed querying + * + * @see org.apache.solr.handler.component.QueryComponent + */ +public class MergeStrategyTest extends BaseDistributedSearchTestCase { + + public MergeStrategyTest() { + fixShardCount = true; + shardCount = 3; + stress = 0; + } + + @BeforeClass + public static void setUpBeforeClass() throws Exception { + initCore("solrconfig-plugcollector.xml", "schema15.xml"); + } + + @Override + public void doTest() throws Exception { + del("*:*"); + + index_specific(0,"id","1", "sort_i", "5"); + index_specific(0,"id","2", "sort_i", "50"); + index_specific(1,"id","5", "sort_i", "4"); + index_specific(1,"id","6", "sort_i", "10"); + index_specific(0,"id","7", "sort_i", "1"); + index_specific(1,"id","8", "sort_i", "2"); + index_specific(2,"id","9", "sort_i", "1000"); + index_specific(2,"id","10", "sort_i", "1500"); + index_specific(2,"id","11", "sort_i", "1300"); + index_specific(1,"id","12", "sort_i", "15"); + index_specific(1,"id","13", "sort_i", "16"); + + commit(); + + handle.put("explain", SKIPVAL); + handle.put("QTime", SKIPVAL); + handle.put("timestamp", SKIPVAL); + handle.put("score", SKIPVAL); + handle.put("wt", SKIP); + handle.put("distrib", SKIP); + handle.put("shards.qt", SKIP); + handle.put("shards", SKIP); + handle.put("q", SKIP); + handle.put("maxScore", SKIPVAL); + handle.put("_version_", SKIP); + + //Test mergeStrategy that uses score + query("q", "{!rank q=$qq}", "qq", "*:*", "rows","12", "sort", "sort_i asc", "fl","*,score"); + + //Test without mergeStrategy + query("q", "*:*", "rows","12", "sort", "sort_i asc"); + + //Test mergeStrategy1 that uses a sort field. + query("q", "{!rank mergeStrategy=1 q=$qq}", "qq", "*:*", "rows","12", "sort", "sort_i asc"); + + ModifiableSolrParams params = new ModifiableSolrParams(); + params.add("qq", "*:*"); + params.add("rows", "12"); + params.add("q", "{!rank q=$qq}"); + params.add("sort", "sort_i asc"); + params.add("fl","*,score"); + setDistributedParams(params); + QueryResponse rsp = queryServer(params); + assertOrder(rsp,"10","11","9","2","13","12","6","1","5","8","7"); + + params = new ModifiableSolrParams(); + params.add("q", "*:*"); + params.add("rows", "12"); + params.add("sort", "sort_i asc"); + params.add("fl","*,score"); + setDistributedParams(params); + rsp = queryServer(params); + assertOrder(rsp,"7","8","5","1","6","12","13","2","9","11","10"); + + MergeStrategy m1 = new MergeStrategy() { + @Override + public void merge(ResponseBuilder rb, ShardRequest sreq) { + } + + public boolean mergesIds() { + return true; + } + + public boolean handlesMergeFields() { return false;} + public void handleMergeFields(ResponseBuilder rb, SolrIndexSearcher searcher) {} + + @Override + public int getCost() { + return 1; + } + }; + + MergeStrategy m2 = new MergeStrategy() { + @Override + public void merge(ResponseBuilder rb, ShardRequest sreq) { + } + + public boolean mergesIds() { + return true; + } + + public boolean handlesMergeFields() { return false;} + public void handleMergeFields(ResponseBuilder rb, SolrIndexSearcher searcher) {} + + @Override + public int getCost() { + return 100; + } + }; + + MergeStrategy m3 = new MergeStrategy() { + @Override + public void merge(ResponseBuilder rb, ShardRequest sreq) { + } + + public boolean mergesIds() { + return false; + } + + public boolean handlesMergeFields() { return false;} + public void handleMergeFields(ResponseBuilder rb, SolrIndexSearcher searcher) {} + + @Override + public int getCost() { + return 50; + } + }; + + MergeStrategy[] merges = {m1,m2,m3}; + Arrays.sort(merges, MergeStrategy.MERGE_COMP); + assert(merges[0].getCost() == 1); + assert(merges[1].getCost() == 50); + assert(merges[2].getCost() == 100); + } + + private void assertOrder(QueryResponse rsp, String ... docs) throws Exception { + SolrDocumentList list = rsp.getResults(); + for(int i=0; i uniqueDoc = new HashMap<>(); + + + NamedList shardInfo = null; + if(rb.req.getParams().getBool(ShardParams.SHARDS_INFO, false)) { + shardInfo = new SimpleOrderedMap<>(); + rb.rsp.getValues().add(ShardParams.SHARDS_INFO,shardInfo); + } + + IndexSchema schema = rb.req.getSchema(); + SchemaField uniqueKeyField = schema.getUniqueKeyField(); + + long numFound = 0; + Float maxScore=null; + boolean partialResults = false; + List shardDocs = new ArrayList(); + + for (ShardResponse srsp : sreq.responses) { + SolrDocumentList docs = null; + + if(shardInfo!=null) { + SimpleOrderedMap nl = new SimpleOrderedMap<>(); + + if (srsp.getException() != null) { + Throwable t = srsp.getException(); + if(t instanceof SolrServerException) { + t = ((SolrServerException)t).getCause(); + } + nl.add("error", t.toString() ); + StringWriter trace = new StringWriter(); + t.printStackTrace(new PrintWriter(trace)); + nl.add("trace", trace.toString() ); + if (srsp.getShardAddress() != null) { + nl.add("shardAddress", srsp.getShardAddress()); + } + } + else { + docs = (SolrDocumentList)srsp.getSolrResponse().getResponse().get("response"); + nl.add("numFound", docs.getNumFound()); + nl.add("maxScore", docs.getMaxScore()); + nl.add("shardAddress", srsp.getShardAddress()); + } + if(srsp.getSolrResponse()!=null) { + nl.add("time", srsp.getSolrResponse().getElapsedTime()); + } + + shardInfo.add(srsp.getShard(), nl); + } + // now that we've added the shard info, let's only proceed if we have no error. + if (srsp.getException() != null) { + partialResults = true; + continue; + } + + if (docs == null) { // could have been initialized in the shards info block above + docs = (SolrDocumentList)srsp.getSolrResponse().getResponse().get("response"); + } + + NamedList responseHeader = (NamedList)srsp.getSolrResponse().getResponse().get("responseHeader"); + if (responseHeader != null && Boolean.TRUE.equals(responseHeader.get("partialResults"))) { + partialResults = true; + } + + // calculate global maxScore and numDocsFound + if (docs.getMaxScore() != null) { + maxScore = maxScore==null ? docs.getMaxScore() : Math.max(maxScore, docs.getMaxScore()); + } + numFound += docs.getNumFound(); + + + for (int i=0; i= 0) { + // TODO: remove previous from priority queue + // continue; + // } + } + + ShardDoc shardDoc = new ShardDoc(); + shardDoc.id = id; + shardDoc.shard = srsp.getShard(); + shardDoc.orderInShard = i; + Object scoreObj = doc.getFieldValue("score"); + if (scoreObj != null) { + if (scoreObj instanceof String) { + shardDoc.score = Float.parseFloat((String)scoreObj); + } else { + shardDoc.score = (Float)scoreObj; + } + } + shardDocs.add(shardDoc); + } // end for-each-doc-in-response + } // end for-each-response + + Collections.sort(shardDocs, new Comparator() { + @Override + public int compare(ShardDoc o1, ShardDoc o2) { + if(o1.score < o2.score) { + return 1; + } else if (o1.score > o2.score) { + return -1; + } else { + return 0; //To change body of implemented methods use File | Settings | File Templates. + } + } + }); + + int resultSize = shardDocs.size(); + + Map resultIds = new HashMap<>(); + for (int i=0; i sortVals = new NamedList<>(); // order is important for the sort fields + IndexReaderContext topReaderContext = searcher.getTopReaderContext(); + List leaves = topReaderContext.leaves(); + AtomicReaderContext currentLeaf = null; + if (leaves.size()==1) { + // if there is a single segment, use that subReader and avoid looking up each time + currentLeaf = leaves.get(0); + leaves=null; + } + + DocList docList = rb.getResults().docList; + + // sort ids from lowest to highest so we can access them in order + int nDocs = docList.size(); + final long[] sortedIds = new long[nDocs]; + final float[] scores = new float[nDocs]; // doc scores, parallel to sortedIds + DocList docs = rb.getResults().docList; + DocIterator it = docs.iterator(); + for (int i=0; i schemaFields = sortSpec.getSchemaFields(); + + for (int fld = 0; fld < schemaFields.size(); fld++) { + SchemaField schemaField = schemaFields.get(fld); + FieldType ft = null == schemaField? null : schemaField.getType(); + SortField sortField = sortFields[fld]; + + SortField.Type type = sortField.getType(); + // :TODO: would be simpler to always serialize every position of SortField[] + if (type==SortField.Type.SCORE || type==SortField.Type.DOC) continue; + + FieldComparator comparator = null; + Object[] vals = new Object[nDocs]; + + int lastIdx = -1; + int idx = 0; + + for (int i = 0; i < sortedIds.length; ++i) { + long idAndPos = sortedIds[i]; + float score = scores[i]; + int doc = (int)(idAndPos >>> 32); + int position = (int)idAndPos; + + if (leaves != null) { + idx = ReaderUtil.subIndex(doc, leaves); + currentLeaf = leaves.get(idx); + if (idx != lastIdx) { + // we switched segments. invalidate comparator. + comparator = null; + } + } + + if (comparator == null) { + comparator = sortField.getComparator(1,0); + comparator = comparator.setNextReader(currentLeaf); + } + + doc -= currentLeaf.docBase; // adjust for what segment this is in + comparator.setScorer(new FakeScorer(doc, score)); + comparator.copy(0, doc); + Object val = comparator.value(0); + if (null != ft) val = ft.marshalSortValue(val); + vals[position] = val; + } + + sortVals.add(sortField.getField(), vals); + } + + rsp.add("merge_values", sortVals); + } + } + + private class FakeScorer extends Scorer { + final int docid; + final float score; + + FakeScorer(int docid, float score) { + super(null); + this.docid = docid; + this.score = score; + } + + @Override + public int docID() { + return docid; + } + + @Override + public float score() throws IOException { + return score; + } + + @Override + public int freq() throws IOException { + throw new UnsupportedOperationException(); + } + + @Override + public int nextDoc() throws IOException { + throw new UnsupportedOperationException(); + } + + @Override + public int advance(int target) throws IOException { + throw new UnsupportedOperationException(); + } + + @Override + public long cost() { + return 1; + } + + @Override + public Weight getWeight() { + throw new UnsupportedOperationException(); + } + + @Override + public Collection getChildren() { + throw new UnsupportedOperationException(); + } + } + + public void merge(ResponseBuilder rb, ShardRequest sreq) { + + // id to shard mapping, to eliminate any accidental dups + HashMap uniqueDoc = new HashMap<>(); + + + NamedList shardInfo = null; + if(rb.req.getParams().getBool(ShardParams.SHARDS_INFO, false)) { + shardInfo = new SimpleOrderedMap<>(); + rb.rsp.getValues().add(ShardParams.SHARDS_INFO,shardInfo); + } + + IndexSchema schema = rb.req.getSchema(); + SchemaField uniqueKeyField = schema.getUniqueKeyField(); + + long numFound = 0; + Float maxScore=null; + boolean partialResults = false; + List shardDocs = new ArrayList(); + + for (ShardResponse srsp : sreq.responses) { + SolrDocumentList docs = null; + + if(shardInfo!=null) { + SimpleOrderedMap nl = new SimpleOrderedMap<>(); + + if (srsp.getException() != null) { + Throwable t = srsp.getException(); + if(t instanceof SolrServerException) { + t = ((SolrServerException)t).getCause(); + } + nl.add("error", t.toString() ); + StringWriter trace = new StringWriter(); + t.printStackTrace(new PrintWriter(trace)); + nl.add("trace", trace.toString() ); + if (srsp.getShardAddress() != null) { + nl.add("shardAddress", srsp.getShardAddress()); + } + } + else { + docs = (SolrDocumentList)srsp.getSolrResponse().getResponse().get("response"); + nl.add("numFound", docs.getNumFound()); + nl.add("maxScore", docs.getMaxScore()); + nl.add("shardAddress", srsp.getShardAddress()); + } + if(srsp.getSolrResponse()!=null) { + nl.add("time", srsp.getSolrResponse().getElapsedTime()); + } + + shardInfo.add(srsp.getShard(), nl); + } + // now that we've added the shard info, let's only proceed if we have no error. + if (srsp.getException() != null) { + partialResults = true; + continue; + } + + if (docs == null) { // could have been initialized in the shards info block above + docs = (SolrDocumentList)srsp.getSolrResponse().getResponse().get("response"); + } + + NamedList responseHeader = (NamedList)srsp.getSolrResponse().getResponse().get("responseHeader"); + if (responseHeader != null && Boolean.TRUE.equals(responseHeader.get("partialResults"))) { + partialResults = true; + } + + // calculate global maxScore and numDocsFound + if (docs.getMaxScore() != null) { + maxScore = maxScore==null ? docs.getMaxScore() : Math.max(maxScore, docs.getMaxScore()); + } + numFound += docs.getNumFound(); + + SortSpec ss = rb.getSortSpec(); + Sort sort = ss.getSort(); + + NamedList sortFieldValues = (NamedList)(srsp.getSolrResponse().getResponse().get("merge_values")); + NamedList unmarshalledSortFieldValues = unmarshalSortValues(ss, sortFieldValues, schema); + List lst = (List)unmarshalledSortFieldValues.getVal(0); + + for (int i=0; i= 0) { + // TODO: remove previous from priority queue + // continue; + // } + } + + ShardDoc shardDoc = new ShardDoc(); + shardDoc.id = id; + shardDoc.shard = srsp.getShard(); + shardDoc.orderInShard = i; + Object scoreObj = lst.get(i); + if (scoreObj != null) { + shardDoc.score = ((Integer)scoreObj).floatValue(); + } + shardDocs.add(shardDoc); + } // end for-each-doc-in-response + } // end for-each-response + + Collections.sort(shardDocs, new Comparator() { + @Override + public int compare(ShardDoc o1, ShardDoc o2) { + if(o1.score < o2.score) { + return 1; + } else if (o1.score > o2.score) { + return -1; + } else { + return 0; //To change body of implemented methods use File | Settings | File Templates. + } + } + }); + + int resultSize = shardDocs.size(); + + Map resultIds = new HashMap<>(); + for (int i=0; i schemaFields = sortSpec.getSchemaFields(); + SortField[] sortFields = sortSpec.getSort().getSort(); + + int marshalledFieldNum = 0; + for (int sortFieldNum = 0; sortFieldNum < sortFields.length; sortFieldNum++) { + final SortField sortField = sortFields[sortFieldNum]; + final SortField.Type type = sortField.getType(); + + // :TODO: would be simpler to always serialize every position of SortField[] + if (type==SortField.Type.SCORE || type==SortField.Type.DOC) continue; + + final String sortFieldName = sortField.getField(); + final String valueFieldName = sortFieldValues.getName(marshalledFieldNum); + assert sortFieldName.equals(valueFieldName) + : "sortFieldValues name key does not match expected SortField.getField"; + + List sortVals = (List)sortFieldValues.getVal(marshalledFieldNum); + + final SchemaField schemaField = schemaFields.get(sortFieldNum); + if (null == schemaField) { + unmarshalledSortValsPerField.add(sortField.getField(), sortVals); + } else { + FieldType fieldType = schemaField.getType(); + List unmarshalledSortVals = new ArrayList(); + for (Object sortVal : sortVals) { + unmarshalledSortVals.add(fieldType.unmarshalSortValue(sortVal)); + } + unmarshalledSortValsPerField.add(sortField.getField(), unmarshalledSortVals); + } + marshalledFieldNum++; + } + return unmarshalledSortValsPerField; + } + } + + + class TestCollector extends TopDocsCollector { + + private List list = new ArrayList(); + private FieldCache.Ints values; + private int base; + + public TestCollector(PriorityQueue pq) { + super(pq); + } + + public boolean acceptsDocsOutOfOrder() { + return false; + } + + public void doSetNextReader(AtomicReaderContext context) throws IOException { + values = FieldCache.DEFAULT.getInts(context.reader(), "sort_i", false); + base = context.docBase; + } + + public void collect(int doc) { + list.add(new ScoreDoc(doc+base, (float)values.get(doc))); + } + + public int topDocsSize() { + return list.size(); + } + + public TopDocs topDocs() { + Collections.sort(list, new Comparator() { + public int compare(Object o1, Object o2) { + ScoreDoc s1 = (ScoreDoc) o1; + ScoreDoc s2 = (ScoreDoc) o2; + if (s1.score == s2.score) { + return 0; + } else if (s1.score < s2.score) { + return 1; + } else { + return -1; + } + } + }); + ScoreDoc[] scoreDocs = list.toArray(new ScoreDoc[list.size()]); + return new TopDocs(list.size(), scoreDocs, 0.0f); + } + + public TopDocs topDocs(int start, int len) { + return topDocs(); + } + + public int getTotalHits() { + return list.size(); + } + } + + class TestCollector1 extends TopDocsCollector { + + private List list = new ArrayList(); + private int base; + private Scorer scorer; + + public TestCollector1(PriorityQueue pq) { + super(pq); + } + + public boolean acceptsDocsOutOfOrder() { + return false; + } + + public void doSetNextReader(AtomicReaderContext context) throws IOException { + base = context.docBase; + } + + public void setScorer(Scorer scorer) { + this.scorer = scorer; + } + + public void collect(int doc) throws IOException { + list.add(new ScoreDoc(doc+base, scorer.score())); + } + + public int topDocsSize() { + return list.size(); + } + + public TopDocs topDocs() { + Collections.sort(list, new Comparator() { + public int compare(Object o1, Object o2) { + ScoreDoc s1 = (ScoreDoc) o1; + ScoreDoc s2 = (ScoreDoc) o2; + if (s1.score == s2.score) { + return 0; + } else if (s1.score > s2.score) { + return 1; + } else { + return -1; + } + } + }); + ScoreDoc[] scoreDocs = list.toArray(new ScoreDoc[list.size()]); + return new TopDocs(list.size(), scoreDocs, 0.0f); + } + + public TopDocs topDocs(int start, int len) { + return topDocs(); + } + + public int getTotalHits() { + return list.size(); + } + } + + + + +}