SOLR-14060: Allow the random Streaming Expression to create the x-axis

This commit is contained in:
Joel Bernstein 2019-12-11 20:23:20 -05:00
parent f894bd019e
commit ee0a6b45ca
2 changed files with 39 additions and 0 deletions

View File

@ -47,6 +47,7 @@ import org.apache.solr.client.solrj.request.QueryRequest;
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 static org.apache.solr.common.params.CommonParams.SORT;
@ -65,6 +66,8 @@ public class RandomStream extends TupleStream implements Expressible {
protected transient SolrClientCache cache;
protected transient CloudSolrClient cloudSolrClient;
private Iterator<SolrDocument> documentIterator;
private int x;
private boolean outputX;
public RandomStream() {
// Used by the RandomFacade
@ -98,6 +101,7 @@ public class RandomStream extends TupleStream implements Expressible {
}
// zkHost, optional - if not provided then will look into factory list to get
String zkHost = null;
if(null == zkHostExpression){
@ -121,6 +125,25 @@ public class RandomStream extends TupleStream implements Expressible {
this.zkHost = zkHost;
this.props = props;
this.collection = collection;
if(props.containsKey(CommonParams.FL)) {
String fl = props.get(CommonParams.FL);
if(fl != null) {
if(fl.equals("*")) {
outputX = true;
} else {
String[] fields = fl.split(",");
for (String f : fields) {
if (f.trim().equals("x")) {
outputX = true;
}
}
}
} else {
outputX = true;
}
} else {
outputX = true;
}
}
@Override
@ -213,9 +236,17 @@ public class RandomStream extends TupleStream implements Expressible {
if(documentIterator.hasNext()) {
Map map = new HashMap();
SolrDocument doc = documentIterator.next();
// Put the generated x-axis first. If there really is an x field it will overwrite it.
if(outputX) {
map.put("x", x++);
}
for(Entry<String, Object> entry : doc.entrySet()) {
map.put(entry.getKey(), entry.getValue());
}
return new Tuple(map);
} else {
Map fields = new HashMap();

View File

@ -662,6 +662,9 @@ public class StreamExpressionTest extends SolrCloudTestCase {
SolrStream solrStream = new SolrStream(jetty.getBaseUrl().toString() + "/collection1", sParams);
List<Tuple> tuples4 = getTuples(solrStream);
assert (tuples4.size() == 1);
//Assert no x-axis
assertNull(tuples4.get(0).get("x"));
sParams = new ModifiableSolrParams(StreamingTest.mapParams(CommonParams.QT, "/stream"));
sParams.add("expr", "random(" + COLLECTIONORALIAS + ")");
@ -674,6 +677,11 @@ public class StreamExpressionTest extends SolrCloudTestCase {
assert(fields.containsKey("a_f"));
assert(fields.containsKey("a_i"));
assert(fields.containsKey("a_s"));
//Assert the x-axis:
for(int i=0; i<tuples4.size(); i++) {
assertEquals(tuples4.get(i).getLong("x").longValue(), i);
}
} finally {
cache.close();
}