SOLR-18466: Support _route_ in the streaming expression search function

changelog/unreleased/SOLR-18466.yml

@@ -0,0 +1,32 @@
+# (DELETE ALL COMMENTS UP HERE AFTER FILLING THIS IN
+
+# See https://github.com/apache/solr/blob/main/dev-docs/changelog.adoc
+
+# If the change is minor, don't bother adding a changelog entry.
+# For `other` type entries, the threshold to bother with a changelog entry should be even higher.
+
+# title:
+#  * The audience is end-users and administrators, not committers.
+#  * Be short and focused on the user impact.  Multiple sentences is fine!
+#  * For technical/geeky details, prefer the commit message instead of changelog.
+#  * Reference JIRA issues like `SOLR-12345`, or if no JIRA but have a GitHub PR then `PR#12345`.
+
+# type:
+#  `added` for new features/improvements, opt-in by the user typically documented in the ref guide
+#  `changed` for improvements; not opt-in
+#  `fixed` for improvements that are deemed to have fixed buggy behavior
+#  `deprecated` for marking things deprecated
+#  `removed` for code removed
+#  `dependency_update` for updates to dependencies
+#  `other` for anything else, like large/significant refactorings, build changes,
+#    test infrastructure, or documentation.
+#    Most such changes are too small/minor to bother with a changelog entry.
+
+title: Support _route_ in the streaming expression function to prune shard fan-out
+
+type: Improvement
+authors:
+  - name: Mathew Skaria
+links:
+  - name: SOLR-18466
+    url: https://issues.apache.org/jira/browse/SOLR-18466

solr/solrj-streaming/src/java/org/apache/solr/client/solrj/io/stream/CloudSolrStream.java

@@ -51,9 +51,11 @@
 import org.apache.solr.client.solrj.io.stream.expr.StreamExpressionValue;
 import org.apache.solr.client.solrj.io.stream.expr.StreamFactory;
 import org.apache.solr.common.cloud.ClusterState;
+import org.apache.solr.common.cloud.DocCollection;
 import org.apache.solr.common.cloud.Replica;
 import org.apache.solr.common.cloud.Slice;
 import org.apache.solr.common.params.ModifiableSolrParams;
+import org.apache.solr.common.params.ShardParams;
 import org.apache.solr.common.params.SolrParams;
 import org.apache.solr.common.util.URLUtil;
 
@@ -375,6 +377,15 @@ private StreamComparator parseComp(String sort, String fl) throws IOException {
   public static List<Slice> getSlices(
       String collectionName, CloudSolrClient cloudSolrClient, boolean checkAlias)
       throws IOException {
+    return getSlices(collectionName, cloudSolrClient, checkAlias, new ModifiableSolrParams());
+  }
+
+  public static List<Slice> getSlices(
+      String collectionName,
+      CloudSolrClient cloudSolrClient,
+      boolean checkAlias,
+      SolrParams solrParams)
+      throws IOException {
 
     Stream<String> allCollections = Arrays.stream(collectionName.split(","));
 
@@ -385,21 +396,30 @@ public static List<Slice> getSlices(
           allCollections.flatMap(
               col -> cloudSolrClient.getClusterStateProvider().resolveAlias(col).stream());
     }
-
+    // Check for _route_ param
+    final String routeKeys = solrParams.get(ShardParams._ROUTE_);
     // Lookup all actives slices for these collections
     ClusterState clusterState = cloudSolrClient.getClusterState();
     List<Slice> slices =
         allCollections
             .map(c -> clusterState.getCollectionOrNull(c, true))
             .filter(Objects::nonNull)
-            .flatMap(docCol -> docCol.getActiveSlices().stream())
+            .flatMap(docCol -> sliceResolution(docCol, routeKeys, solrParams))
             .toList();
     if (slices.isEmpty()) {
       throw new IOException("Slices not found for " + collectionName);
     }
     return slices;
   }
 
+  private static Stream<Slice> sliceResolution(
+      DocCollection docCol, String routeKeys, SolrParams solrParams) {
+    if (routeKeys == null || routeKeys.isEmpty()) {
+      return docCol.getActiveSlices().stream();
+    }
+    return docCol.getRouter().getSearchSlices(routeKeys, solrParams, docCol).stream();
+  }
+
   protected void constructStreams() throws IOException {
     final ModifiableSolrParams mParams = adjustParams(new ModifiableSolrParams(params));
     mParams.set(DISTRIB, "false"); // We are the aggregator.

solr/solrj-streaming/src/java/org/apache/solr/client/solrj/io/stream/TupleStream.java

@@ -158,7 +158,6 @@ static List<Replica> getReplicas(
     try {
       CloudSolrClient cloudSolrClient = solrClientCache.getCloudSolrClient(solrConnection);
       ClusterState clusterState = cloudSolrClient.getClusterStateProvider().getClusterState();
-      List<Slice> slices = CloudSolrStream.getSlices(collection, cloudSolrClient, true);
       Set<String> liveNodes = clusterState.getLiveNodes();
 
       RequestReplicaListTransformerGenerator requestReplicaListTransformerGenerator;
@@ -176,6 +175,8 @@ static List<Replica> getReplicas(
       }
       solrParams.add(requestParams);
 
+      List<Slice> slices = CloudSolrStream.getSlices(collection, cloudSolrClient, true, solrParams);
+
       ReplicaListTransformer replicaListTransformer =
           requestReplicaListTransformerGenerator.getReplicaListTransformer(solrParams);
 

solr/solrj-streaming/src/test/org/apache/solr/client/solrj/io/stream/StreamExpressionTest.java

@@ -32,6 +32,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import java.util.zip.GZIPOutputStream;
@@ -58,7 +59,12 @@
 import org.apache.solr.client.solrj.request.CollectionAdminRequest;
 import org.apache.solr.client.solrj.request.UpdateRequest;
 import org.apache.solr.cloud.SolrCloudTestCase;
+import org.apache.solr.common.SolrException;
+import org.apache.solr.common.cloud.CompositeIdRouter;
+import org.apache.solr.common.cloud.DocCollection;
+import org.apache.solr.common.cloud.Slice;
 import org.apache.solr.common.params.ModifiableSolrParams;
+import org.apache.solr.common.params.ShardParams;
 import org.apache.solr.common.util.URLUtil;
 import org.apache.solr.core.CoreDescriptor;
 import org.apache.solr.embedded.JettySolrRunner;
@@ -428,6 +434,223 @@ public void testSqlStream() throws Exception {
     }
   }
 
+  @Test
+  public void testCloudSolrStreamWithRouteParam() throws Exception {
+    final String routeCollection = "routeCollection";
+    CollectionAdminRequest.createCollectionWithImplicitRouter(routeCollection, "conf", "a,b,c", 1)
+        .process(cluster.getSolrClient());
+    cluster.waitForActiveCollection(routeCollection, 3, 3);
+
+    StreamContext streamContext = new StreamContext();
+    SolrClientCache solrClientCache = new SolrClientCache();
+    streamContext.setSolrClientCache(solrClientCache);
+    try {
+      new UpdateRequest()
+          .add(id, "0", "a_i", "0")
+          .add(id, "1", "a_i", "1")
+          .add(id, "2", "a_i", "2")
+          .withRoute("a")
+          .process(cluster.getSolrClient(), routeCollection);
+      new UpdateRequest()
+          .add(id, "3", "a_i", "3")
+          .add(id, "4", "a_i", "4")
+          .withRoute("b")
+          .process(cluster.getSolrClient(), routeCollection);
+      new UpdateRequest()
+          .add(id, "5", "a_i", "5")
+          .withRoute("c")
+          .commit(cluster.getSolrClient(), routeCollection);
+
+      CloudSolrClient cloudSolrClient = cluster.getSolrClient();
+
+      assertEquals(Set.of("a"), routedSliceNames(routeCollection, cloudSolrClient, "a"));
+      assertEquals(Set.of("a", "c"), routedSliceNames(routeCollection, cloudSolrClient, "a,c"));
+      assertEquals(
+          3,
+          CloudSolrStream.getSlices(
+                  routeCollection, cloudSolrClient, true, new ModifiableSolrParams())
+              .size());
+
+      StreamFactory factory =
+          new StreamFactory()
+              .withDefaultSolrConnection(solrConnection)
+              .withCollectionUseThisConnection(routeCollection, solrConnection);
+
+      assertRouteQuery(routeCollection, factory, streamContext, "a", 1, new int[] {0, 1, 2});
+      assertRouteQuery(routeCollection, factory, streamContext, "a,c", 2, new int[] {0, 1, 2, 5});
+      assertRouteQuery(
+          routeCollection, factory, streamContext, null, 3, new int[] {0, 1, 2, 3, 4, 5});
+    } finally {
+      CollectionAdminRequest.deleteCollection(routeCollection).process(cluster.getSolrClient());
+      solrClientCache.close();
+    }
+  }
+
+  private Set<String> routedSliceNames(
+      String collection, CloudSolrClient cloudSolrClient, String routeKey) throws IOException {
+    ModifiableSolrParams params = new ModifiableSolrParams();
+    params.add(ShardParams._ROUTE_, routeKey);
+    return CloudSolrStream.getSlices(collection, cloudSolrClient, true, params).stream()
+        .map(Slice::getName)
+        .collect(Collectors.toSet());
+  }
+
+  private void assertRouteQuery(
+      String collection,
+      StreamFactory factory,
+      StreamContext streamContext,
+      String routeKey,
+      int expectedStreamCount,
+      int[] expectedIds)
+      throws Exception {
+    CloudSolrStream stream = routeStream(collection, factory, streamContext, routeKey);
+    List<Tuple> tuples = getTuples(stream);
+    assertEquals(expectedStreamCount, stream.children().size());
+    assertEquals(expectedIds.length, tuples.size());
+    assertOrderOf(tuples, "a_i", expectedIds);
+  }
+
+  private CloudSolrStream routeStream(
+      String collection, StreamFactory factory, StreamContext streamContext, String routeKey)
+      throws IOException {
+    String routeClause = routeKey == null ? "" : ", _route_=\"" + routeKey + "\"";
+    StreamExpression expression =
+        StreamExpressionParser.parse(
+            "search("
+                + collection
+                + ", q=*:*, fl=\"id,a_i\", sort=\"a_i asc\""
+                + routeClause
+                + ")");
+    CloudSolrStream stream = new CloudSolrStream(expression, factory);
+    stream.setStreamContext(streamContext);
+    return stream;
+  }
+
+  @Test
+  public void testCloudSolrStreamWithCompositeIdRoute() throws Exception {
+    new UpdateRequest()
+        .add(id, "user1!0", "a_i", "0")
+        .add(id, "user1!1", "a_i", "1")
+        .add(id, "user1!2", "a_i", "2")
+        .add(id, "user2!3", "a_i", "3")
+        .add(id, "user2!4", "a_i", "4")
+        .commit(cluster.getSolrClient(), COLLECTIONORALIAS);
+
+    CloudSolrClient cloudSolrClient = cluster.getSolrClient();
+    String realCollection =
+        cloudSolrClient.getClusterStateProvider().resolveAlias(COLLECTIONORALIAS).get(0);
+    DocCollection docCollection = cloudSolrClient.getClusterState().getCollection(realCollection);
+    assertEquals(
+        "This test requires the compositeId router",
+        CompositeIdRouter.NAME,
+        docCollection.getRouter().getName());
+
+    ModifiableSolrParams routeParams = new ModifiableSolrParams();
+    routeParams.add(ShardParams._ROUTE_, "user1!");
+    Set<String> expected =
+        docCollection.getRouter().getSearchSlices("user1!", routeParams, docCollection).stream()
+            .map(Slice::getName)
+            .collect(Collectors.toSet());
+    Set<String> routedSliceNames =
+        CloudSolrStream.getSlices(COLLECTIONORALIAS, cloudSolrClient, true, routeParams).stream()
+            .map(Slice::getName)
+            .collect(Collectors.toSet());
+    assertEquals(
+        "getSlices should match the compositeId router's slice selection",
+        expected,
+        routedSliceNames);
+
+    StreamFactory factory =
+        new StreamFactory()
+            .withDefaultSolrConnection(solrConnection)
+            .withCollectionUseThisConnection(COLLECTIONORALIAS, solrConnection);
+    StreamContext streamContext = new StreamContext();
+    SolrClientCache solrClientCache = new SolrClientCache();
+    streamContext.setSolrClientCache(solrClientCache);
+    try {
+      StreamExpression expression =
+          StreamExpressionParser.parse(
+              "search("
+                  + COLLECTIONORALIAS
+                  + ", q=*:*, fq=\"id:user1!*\", fl=\"id,a_i\", sort=\"a_i asc\", _route_=\"user1!\")");
+      CloudSolrStream stream = new CloudSolrStream(expression, factory);
+      stream.setStreamContext(streamContext);
+      List<Tuple> tuples = getTuples(stream);
+      assertEquals(
+          "search() should open one SolrStream per routed slice",
+          routedSliceNames.size(),
+          stream.children().size());
+      assertEquals(3, tuples.size());
+      assertOrderOf(tuples, "a_i", 0, 1, 2);
+    } finally {
+      solrClientCache.close();
+    }
+  }
+
+  @Test
+  public void testCloudSolrStreamWithEmptyRoute() throws Exception {
+    new UpdateRequest()
+        .add(id, "0", "a_i", "0")
+        .add(id, "1", "a_i", "1")
+        .add(id, "2", "a_i", "2")
+        .commit(cluster.getSolrClient(), COLLECTIONORALIAS);
+
+    CloudSolrClient cloudSolrClient = cluster.getSolrClient();
+    String realCollection =
+        cloudSolrClient.getClusterStateProvider().resolveAlias(COLLECTIONORALIAS).get(0);
+    DocCollection docCollection = cloudSolrClient.getClusterState().getCollection(realCollection);
+    int activeSliceCount = docCollection.getActiveSlices().size();
+
+    ModifiableSolrParams emptyRouteParams = new ModifiableSolrParams();
+    emptyRouteParams.add(ShardParams._ROUTE_, "");
+    List<Slice> slices =
+        CloudSolrStream.getSlices(COLLECTIONORALIAS, cloudSolrClient, true, emptyRouteParams);
+    assertEquals(
+        "Empty _route_ should fall back to all active slices", activeSliceCount, slices.size());
+  }
+
+  @Test
+  public void testCloudSolrStreamWithInvalidRoute() throws Exception {
+    final String routeCollection = "invalidRouteCollection";
+    CollectionAdminRequest.createCollectionWithImplicitRouter(routeCollection, "conf", "a,b", 1)
+        .process(cluster.getSolrClient());
+    cluster.waitForActiveCollection(routeCollection, 2, 2);
+
+    StreamContext streamContext = new StreamContext();
+    SolrClientCache solrClientCache = new SolrClientCache();
+    streamContext.setSolrClientCache(solrClientCache);
+    try {
+      new UpdateRequest()
+          .add(id, "0", "a_i", "0")
+          .withRoute("a")
+          .commit(cluster.getSolrClient(), routeCollection);
+
+      CloudSolrClient cloudSolrClient = cluster.getSolrClient();
+      ModifiableSolrParams badRouteParams = new ModifiableSolrParams();
+      badRouteParams.add(ShardParams._ROUTE_, "doesNotExist");
+
+      SolrException ex =
+          expectThrows(
+              SolrException.class,
+              () ->
+                  CloudSolrStream.getSlices(
+                      routeCollection, cloudSolrClient, true, badRouteParams));
+      assertTrue(
+          "Error should mention the missing shard", ex.getMessage().contains("doesNotExist"));
+
+      StreamFactory factory =
+          new StreamFactory()
+              .withDefaultSolrConnection(solrConnection)
+              .withCollectionUseThisConnection(routeCollection, solrConnection);
+      expectThrows(
+          IOException.class,
+          () -> getTuples(routeStream(routeCollection, factory, streamContext, "doesNotExist")));
+    } finally {
+      CollectionAdminRequest.deleteCollection(routeCollection).process(cluster.getSolrClient());
+      solrClientCache.close();
+    }
+  }
+
   @Test
   public void testCloudSolrStreamWithZkHost() throws Exception {