diff --git a/changelog/unreleased/SOLR-18466.yml b/changelog/unreleased/SOLR-18466.yml new file mode 100644 index 00000000000..1873e5702a4 --- /dev/null +++ b/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 diff --git a/solr/solrj-streaming/src/java/org/apache/solr/client/solrj/io/stream/CloudSolrStream.java b/solr/solrj-streaming/src/java/org/apache/solr/client/solrj/io/stream/CloudSolrStream.java index 5f16e7348b2..c4d6e4cb271 100644 --- a/solr/solrj-streaming/src/java/org/apache/solr/client/solrj/io/stream/CloudSolrStream.java +++ b/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 getSlices( String collectionName, CloudSolrClient cloudSolrClient, boolean checkAlias) throws IOException { + return getSlices(collectionName, cloudSolrClient, checkAlias, new ModifiableSolrParams()); + } + + public static List getSlices( + String collectionName, + CloudSolrClient cloudSolrClient, + boolean checkAlias, + SolrParams solrParams) + throws IOException { Stream allCollections = Arrays.stream(collectionName.split(",")); @@ -385,14 +396,15 @@ public static List 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 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); @@ -400,6 +412,14 @@ public static List getSlices( return slices; } + private static Stream 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. diff --git a/solr/solrj-streaming/src/java/org/apache/solr/client/solrj/io/stream/TupleStream.java b/solr/solrj-streaming/src/java/org/apache/solr/client/solrj/io/stream/TupleStream.java index 3dc3ceb6f78..64c35d54ba3 100644 --- a/solr/solrj-streaming/src/java/org/apache/solr/client/solrj/io/stream/TupleStream.java +++ b/solr/solrj-streaming/src/java/org/apache/solr/client/solrj/io/stream/TupleStream.java @@ -158,7 +158,6 @@ static List getReplicas( try { CloudSolrClient cloudSolrClient = solrClientCache.getCloudSolrClient(solrConnection); ClusterState clusterState = cloudSolrClient.getClusterStateProvider().getClusterState(); - List slices = CloudSolrStream.getSlices(collection, cloudSolrClient, true); Set liveNodes = clusterState.getLiveNodes(); RequestReplicaListTransformerGenerator requestReplicaListTransformerGenerator; @@ -176,6 +175,8 @@ static List getReplicas( } solrParams.add(requestParams); + List slices = CloudSolrStream.getSlices(collection, cloudSolrClient, true, solrParams); + ReplicaListTransformer replicaListTransformer = requestReplicaListTransformerGenerator.getReplicaListTransformer(solrParams); diff --git a/solr/solrj-streaming/src/test/org/apache/solr/client/solrj/io/stream/StreamExpressionTest.java b/solr/solrj-streaming/src/test/org/apache/solr/client/solrj/io/stream/StreamExpressionTest.java index 11bfe63ae0a..c5690130efc 100644 --- a/solr/solrj-streaming/src/test/org/apache/solr/client/solrj/io/stream/StreamExpressionTest.java +++ b/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 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 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 expected = + docCollection.getRouter().getSearchSlices("user1!", routeParams, docCollection).stream() + .map(Slice::getName) + .collect(Collectors.toSet()); + Set 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 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 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 {