From ee0a6b45ca50b94bf1a68032687442919d19ecb9 Mon Sep 17 00:00:00 2001
From: Joel Bernstein <jbernste@apache.org>
Date: Wed, 11 Dec 2019 20:23:20 -0500
Subject: [PATCH] SOLR-14060: Allow the random Streaming Expression to create
 the x-axis

---
 .../client/solrj/io/stream/RandomStream.java  | 31 +++++++++++++++++++
 .../solrj/io/stream/StreamExpressionTest.java |  8 +++++
 2 files changed, 39 insertions(+)

diff --git a/solr/solrj/src/java/org/apache/solr/client/solrj/io/stream/RandomStream.java b/solr/solrj/src/java/org/apache/solr/client/solrj/io/stream/RandomStream.java
index 8f065000d65..aca0e3d4af8 100644
--- a/solr/solrj/src/java/org/apache/solr/client/solrj/io/stream/RandomStream.java
+++ b/solr/solrj/src/java/org/apache/solr/client/solrj/io/stream/RandomStream.java
@@ -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();
diff --git a/solr/solrj/src/test/org/apache/solr/client/solrj/io/stream/StreamExpressionTest.java b/solr/solrj/src/test/org/apache/solr/client/solrj/io/stream/StreamExpressionTest.java
index 7bb0584f8bb..a722f042af4 100644
--- a/solr/solrj/src/test/org/apache/solr/client/solrj/io/stream/StreamExpressionTest.java
+++ b/solr/solrj/src/test/org/apache/solr/client/solrj/io/stream/StreamExpressionTest.java
@@ -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();
     }