GH-5030: fix FedXDataset support for queries with FROM clauses (#5031)

Author: hmottestad

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());
+		}
+
+	}
 }