GH-5030: fix FedXDataset support for queries with FROM clauses

FedX supports interpret the Dataset being defined on the query: if it is
of type FedXDataset, the federation is reduced to the members defined by
FedXDataset#getEndpoints().

This is a nice means to evaluate a query on a subset of the federation
and is generally working fine.

When the query contains a FROM clause, however, the externally passed
Dataset is wrapped into a FallbackDataset (which internally keeps the
FedXDataset as primary one.

This change now inspects the FallbackDataset and unwraps it to reduce
the members. Note that for execution still the "FallbackDataset" is
used.

core/query/src/main/java/org/eclipse/rdf4j/query/impl/FallbackDataset.java

@@ -113,4 +113,12 @@ private void appendURI(StringBuilder sb, IRI uri) {
 		}
 	}
 
+	public Dataset getPrimary() {
+		return primary;
+	}
+
+	public Dataset getFallback() {
+		return fallback;
+	}
+
 }

tools/federation/src/main/java/org/eclipse/rdf4j/federated/evaluation/FederationEvalStrategy.java

@@ -121,6 +121,7 @@
 import org.eclipse.rdf4j.query.algebra.helpers.TupleExprs;
 import org.eclipse.rdf4j.query.algebra.helpers.collectors.VarNameCollector;
 import org.eclipse.rdf4j.query.impl.EmptyBindingSet;
+import org.eclipse.rdf4j.query.impl.FallbackDataset;
 import org.eclipse.rdf4j.repository.RepositoryException;
 import org.eclipse.rdf4j.repository.sparql.federation.CollectionIteration;
 import org.eclipse.rdf4j.repository.sparql.federation.RepositoryFederatedService;
@@ -179,19 +180,8 @@ public TupleExpr optimize(TupleExpr expr, EvaluationStatistics evaluationStatist
 
 		FederationEvaluationStatistics stats = (FederationEvaluationStatistics) evaluationStatistics;
 		QueryInfo queryInfo = stats.getQueryInfo();
-		Dataset dataset = stats.getDataset();
 
-		FederationContext federationContext = queryInfo.getFederationContext();
-		List<Endpoint> members;
-		if (dataset instanceof FedXDataset) {
-			// run the query against a selected set of endpoints
-			FedXDataset ds = (FedXDataset) dataset;
-			members = federationContext.getEndpointManager().getEndpoints(ds.getEndpoints());
-		} else {
-			// evaluate against entire federation
-			FedX fed = federationContext.getFederation();
-			members = fed.getMembers();
-		}
+		List<Endpoint> members = getMembersFromContext(stats);
 
 		// Clone the tuple expression to allow for more aggressive optimizations
 		TupleExpr clone = expr.clone();
@@ -270,6 +260,34 @@ public TupleExpr optimize(TupleExpr expr, EvaluationStatistics evaluationStatist
 		return query;
 	}
 
+	/**
+	 * Returns the federation members that are active in the current federation. By default it is all federation
+	 * members. If the passed {@link Dataset} is a {@link FedXDataset}, the defined {@link FedXDataset#getEndpoints()}
+	 * are used.
+	 *
+	 * @param evaluationStatisticss to keep the current query context
+	 * @return
+	 */
+	protected List<Endpoint> getMembersFromContext(FederationEvaluationStatistics evaluationStatisticss) {
+		Dataset queryDataset = evaluationStatisticss.getDataset();
+
+		// if the query uses a FROM clause, the external dataset
+		// is wrapped as primary dataset in FallbackDataset
+		if (queryDataset instanceof FallbackDataset) {
+			queryDataset = ((FallbackDataset) queryDataset).getPrimary();
+		}
+
+		if (queryDataset instanceof FedXDataset) {
+			// run the query against a selected set of endpoints
+			FedXDataset ds = (FedXDataset) queryDataset;
+			return federationContext.getEndpointManager().getEndpoints(ds.getEndpoints());
+		} else {
+			// evaluate against entire federation
+			FedX fed = federationContext.getFederation();
+			return fed.getMembers();
+		}
+	}
+
 	/**
 	 * Perform source selection for all statements of the query. As a result of this method all statement nodes are
 	 * annotated with their relevant sources.

tools/federation/src/test/java/org/eclipse/rdf4j/federated/BasicTests.java

@@ -24,6 +24,8 @@
 import org.eclipse.rdf4j.model.Statement;
 import org.eclipse.rdf4j.model.util.Values;
 import org.eclipse.rdf4j.model.vocabulary.DCTERMS;
+import org.eclipse.rdf4j.model.vocabulary.FOAF;
+import org.eclipse.rdf4j.model.vocabulary.RDF;
 import org.eclipse.rdf4j.model.vocabulary.RDFS;
 import org.eclipse.rdf4j.query.AbstractTupleQueryResultHandler;
 import org.eclipse.rdf4j.query.BindingSet;
@@ -35,6 +37,7 @@
 import org.eclipse.rdf4j.query.TupleQueryResult;
 import org.eclipse.rdf4j.query.TupleQueryResultHandlerException;
 import org.eclipse.rdf4j.query.impl.SimpleDataset;
+import org.eclipse.rdf4j.repository.Repository;
 import org.eclipse.rdf4j.repository.RepositoryConnection;
 import org.eclipse.rdf4j.repository.util.Repositories;
 import org.junit.jupiter.api.Assertions;
@@ -575,4 +578,83 @@ public void test_EscapingQuotedLiteral() throws Exception {
 		}
 
 	}
+
+	@Test
+	public void test_reduceFederation() throws Exception {
+
+		List<Endpoint> endpoints = prepareTest(
+				Arrays.asList("/tests/basic/data_emptyStore.ttl", "/tests/basic/data_emptyStore.ttl"));
+
+		Repository repo1 = getRepository(1);
+		Repository repo2 = getRepository(2);
+
+		String repo1Id = endpoints.get(0).getId();
+
+		IRI graph1 = Values.iri("http://example.org/graph1");
+		IRI graph2 = Values.iri("http://example.org/graph2");
+
+		try (RepositoryConnection con = repo1.getConnection()) {
+			con.add(Values.iri("http://example.org/repo1/p1"), RDF.TYPE, FOAF.PERSON, graph1);
+			con.add(Values.iri("http://example.org/repo2/p2"), RDF.TYPE, FOAF.PERSON, graph2);
+		}
+
+		try (RepositoryConnection con = repo2.getConnection()) {
+			con.add(Values.iri("http://example.org/repo2/p3"), RDF.TYPE, FOAF.PERSON, graph1);
+		}
+
+		Repository fedxRepo = fedxRule.getRepository();
+
+		// 1: regular federation
+		try (RepositoryConnection con = fedxRepo.getConnection()) {
+			TupleQuery tupleQuery = con.prepareTupleQuery(
+					"PREFIX foaf: <http://xmlns.com/foaf/0.1/> "
+							+ "SELECT * WHERE { "
+							+ "   ?subClass a foaf:Person. "
+							+ " } "
+			);
+
+			// expect from both repos
+			Assertions.assertEquals(3, QueryResults.asSet(tupleQuery.evaluate()).size());
+
+			// now we scope it additional to graph1
+			tupleQuery = con.prepareTupleQuery(
+					"PREFIX foaf: <http://xmlns.com/foaf/0.1/> "
+							+ "SELECT * FROM <http://example.org/graph1> WHERE { "
+							+ "   ?subClass a foaf:Person. "
+							+ " } ");
+
+			// expect results defined in graph1 (1 in repo1, 2 from repo2)
+			Assertions.assertEquals(2, QueryResults.asSet(tupleQuery.evaluate()).size());
+		}
+
+		// 2: reduce to federation member 1 id
+		// 2a: additionall restrict to named graph
+		FedXDataset fedXDataset = new FedXDataset(new SimpleDataset());
+		fedXDataset.addEndpoint(repo1Id);
+
+		try (RepositoryConnection con = fedxRepo.getConnection()) {
+			TupleQuery tupleQuery = con.prepareTupleQuery(
+					"PREFIX foaf: <http://xmlns.com/foaf/0.1/> "
+							+ "SELECT * WHERE { "
+							+ "   ?subClass a foaf:Person. "
+							+ " } "
+			);
+			tupleQuery.setDataset(fedXDataset);
+
+			// expect result from repo 1
+			Assertions.assertEquals(2, QueryResults.asSet(tupleQuery.evaluate()).size());
+
+			// now we scope it additional to graph1
+			tupleQuery = con.prepareTupleQuery(
+					"PREFIX foaf: <http://xmlns.com/foaf/0.1/> "
+							+ "SELECT * FROM <http://example.org/graph1> WHERE { "
+							+ "   ?subClass a foaf:Person. "
+							+ " } ");
+			tupleQuery.setDataset(fedXDataset);
+
+			// expect result from graph1 from repo1
+			Assertions.assertEquals(1, QueryResults.asSet(tupleQuery.evaluate()).size());
+		}
+
+	}
 }