UBI goes distrib

Pardon, I barely understand what's going on there.

solr/core/src/java/org/apache/solr/handler/component/UBIComponent.java

@@ -237,27 +237,44 @@ public void process(ResponseBuilder rb) throws IOException {
 
   @Override
   public int distributedProcess(ResponseBuilder rb) throws IOException {
-
     SolrParams params = rb.req.getParams();
     if (!params.getBool(COMPONENT_NAME, false)) {
       return ResponseBuilder.STAGE_DONE;
     }
 
-    if (rb.stage != ResponseBuilder.STAGE_GET_FIELDS) {
-      return ResponseBuilder.STAGE_DONE;
+    if (rb.stage < ResponseBuilder.STAGE_GET_FIELDS) {
+      return ResponseBuilder.STAGE_GET_FIELDS;
     }
 
-    doStuff(rb);
+    if (rb.stage == ResponseBuilder.STAGE_GET_FIELDS) {
+      doDistribStuff(rb);
+      return ResponseBuilder.STAGE_DONE;
+    }
 
     return ResponseBuilder.STAGE_DONE;
   }
 
   public void doStuff(ResponseBuilder rb) throws IOException {
+    UBIQuery ubiQuery = getUbiQuery(rb);
+    if (ubiQuery == null) return;
+
+    ResultContext rc = (ResultContext) rb.rsp.getResponse();
+    DocList docs = rc.getDocList();
+    // DocList docs = rb.getResults().docList;
+
+    String docIds = extractDocIds(docs, rb.req.getSearcher());
+
+    ubiQuery.setDocIds(docIds);
 
-    // not sure why but sometimes we get it twoice...  how can a response have the
+    addUserBehaviorInsightsToResponse(ubiQuery, rb);
+    recordQuery(ubiQuery);
+  }
+
+  private static UBIQuery getUbiQuery(ResponseBuilder rb) {
+    // not sure why but sometimes we get it tw(o)ice...  how can a response have the
     // the same component run twice?
     if (rb.rsp.getValues().get("ubi") != null) {
-      return;
+      return null;
     }
     SolrParams params = rb.req.getParams();
 
@@ -270,9 +287,9 @@ public void doStuff(ResponseBuilder rb) throws IOException {
     ubiQuery.setApplication(params.get(APPLICATION));
     if (ubiQuery.getApplication() == null) {
       ubiQuery.setApplication(
-          rb.isDistrib
-              ? rb.req.getCloudDescriptor().getCollectionName()
-              : searcher.getCore().getName());
+              rb.isDistrib
+                      ? rb.req.getCloudDescriptor().getCollectionName()
+                      : searcher.getCore().getName());
     }
 
     String queryAttributes = params.get(QUERY_ATTRIBUTES);
@@ -292,12 +309,19 @@ public void doStuff(ResponseBuilder rb) throws IOException {
         }
       }
     }
+    return ubiQuery;
+  }
+
+  public void doDistribStuff(ResponseBuilder rb) throws IOException {
+
+    // not sure why but sometimes we get it tw(o)ice...  how can a response have the
+    // the same component run twice?
+    UBIQuery ubiQuery = getUbiQuery(rb);
+    if (ubiQuery == null) return;
 
-    ResultContext rc = (ResultContext) rb.rsp.getResponse();
-    DocList docs = rc.getDocList();
-    // DocList docs = rb.getResults().docList;
 
-    String docIds = extractDocIds(docs, searcher);
+    //String docIds = extractDocIds(docs, searcher);
+    String docIds =String.join(",", rb.resultIds.keySet().stream().map(Object::toString).toList());
     ubiQuery.setDocIds(docIds);
 
     addUserBehaviorInsightsToResponse(ubiQuery, rb);

solr/core/src/test/org/apache/solr/handler/component/UBIComponentDistrQueriesTest.java

@@ -0,0 +1,110 @@
+/*
+ * 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 com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.commons.io.input.ReversedLinesFileReader;
+import org.apache.lucene.tests.util.LuceneTestCase;
+import org.apache.solr.client.solrj.io.SolrClientCache;
+import org.apache.solr.client.solrj.io.Tuple;
+import org.apache.solr.client.solrj.io.stream.*;
+import org.apache.solr.client.solrj.io.stream.expr.StreamExpression;
+import org.apache.solr.client.solrj.io.stream.expr.StreamExpressionParser;
+import org.apache.solr.client.solrj.io.stream.expr.StreamFactory;
+import org.apache.solr.client.solrj.request.CollectionAdminRequest;
+import org.apache.solr.client.solrj.request.UpdateRequest;
+import org.apache.solr.client.solrj.response.QueryResponse;
+import org.apache.solr.cloud.AbstractDistribZkTestBase;
+import org.apache.solr.cloud.SolrCloudTestCase;
+import org.apache.solr.common.SolrInputDocument;
+import org.apache.solr.common.params.MapSolrParams;
+import org.apache.solr.common.params.SolrParams;
+import org.apache.solr.core.SolrCore;
+import org.apache.solr.embedded.JettySolrRunner;
+import org.apache.solr.handler.LoggingStream;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Path;
+import java.time.Instant;
+import java.util.*;
+
+@LuceneTestCase.SuppressCodecs({"Lucene3x", "Lucene40", "Lucene41", "Lucene42", "Lucene45"})
+public class UBIComponentDistrQueriesTest extends SolrCloudTestCase {
+
+  private static final String COLLECTIONORALIAS = "collection1";
+  private static final int TIMEOUT = DEFAULT_TIMEOUT;
+  private static final String id = "id";
+
+  private static boolean useAlias;
+
+  @BeforeClass
+  public static void setupCluster() throws Exception {
+    configureCluster(4)
+        .addConfig(
+            "conf", TEST_PATH().resolve("configsets").resolve("ubi-enabled").resolve("conf"))
+        .configure();
+
+    String collection;
+    useAlias = random().nextBoolean();
+    if (useAlias) {
+      collection = COLLECTIONORALIAS + "_collection";
+    } else {
+      collection = COLLECTIONORALIAS;
+    }
+
+    CollectionAdminRequest.createCollection(collection, "conf", 2, 1)
+        .process(cluster.getSolrClient());
+
+    cluster.waitForActiveCollection(collection, 2, 2);
+
+    AbstractDistribZkTestBase.waitForRecoveriesToFinish(
+        collection, cluster.getZkStateReader(), false, true, TIMEOUT);
+    if (useAlias) {
+      CollectionAdminRequest.createAlias(COLLECTIONORALIAS, collection)
+          .process(cluster.getSolrClient());
+    }
+
+    // -------------------
+
+    CollectionAdminRequest.createCollection("ubi_queries", "_default", 1, 1)
+            .process(cluster.getSolrClient());
+
+    cluster.waitForActiveCollection("ubi_queries", 1, 1);
+
+    AbstractDistribZkTestBase.waitForRecoveriesToFinish(
+            "ubi_queries", cluster.getZkStateReader(), false, true, TIMEOUT);
+  }
+
+  @Before
+  public void cleanIndex() throws Exception {
+    new UpdateRequest().deleteByQuery("*:*").commit(cluster.getSolrClient(), COLLECTIONORALIAS);
+  }
+
+  @Test
+  public void testUBIQueryStream() throws Exception {
+    cluster.getSolrClient(COLLECTIONORALIAS).add(List.of(new SolrInputDocument("id", "1", "subject", "aa"),
+            new SolrInputDocument("id", "two", "subject", "aa"),
+            new SolrInputDocument("id", "3", "subject", "aa")));
+    cluster.getSolrClient(COLLECTIONORALIAS).commit(true, true);
+    QueryResponse queryResponse = cluster.getSolrClient(COLLECTIONORALIAS).query(new MapSolrParams(Map.of("q", "aa", "rows", "2", "ubi", "true")));
+    System.out.println(queryResponse);
+  }
+}