GH-4920 SPARQLConnection.size() method should not fetch every stateme…

…nt in the repository. Just send a count query instead.

Signed-off-by: Jerven Bolleman <jerven.bolleman@sib.swiss>

core/repository/sparql/src/main/java/org/eclipse/rdf4j/repository/sparql/SPARQLConnection.java

@@ -17,7 +17,9 @@
 import java.io.InputStream;
 import java.io.Reader;
 import java.net.URL;
+import java.util.Arrays;
 import java.util.Objects;
+import java.util.stream.Collectors;
 
 import org.apache.http.client.HttpClient;
 import org.eclipse.rdf4j.common.iteration.CloseableIteration;
@@ -80,6 +82,8 @@
  */
 public class SPARQLConnection extends AbstractRepositoryConnection implements HttpClientDependent {
 
+	private static final String COUNT_EVERYTHING = "SELECT (COUNT(*) AS ?count) WHERE { ?s ?p ?o }";
+
 	private static final String EVERYTHING = "CONSTRUCT { ?s ?p ?o } WHERE { ?s ?p ?o }";
 
 	private static final String EVERYTHING_WITH_GRAPH = "SELECT * WHERE {  ?s ?p ?o . OPTIONAL { GRAPH ?ctx { ?s ?p ?o } } }";
@@ -288,14 +292,40 @@ public boolean isEmpty() throws RepositoryException {
 
 	@Override
 	public long size(Resource... contexts) throws RepositoryException {
-		try (RepositoryResult<Statement> stmts = getStatements(null, null, null, true, contexts)) {
-			long i = 0;
-			while (stmts.hasNext()) {
-				stmts.next();
-				i++;
+		String query = sizeAsTupleQuery(contexts);
+		TupleQuery tq = prepareTupleQuery(SPARQL, query);
+		try (TupleQueryResult res = tq.evaluate()) {
+			if (res.hasNext()) {
+
+				Value value = res.next().getBinding("count").getValue();
+				if (value instanceof Literal) {
+					return ((Literal) value).longValue();
+				} else {
+					return 0;
+				}
+			}
+		} catch (QueryEvaluationException e) {
+			throw new RepositoryException(e);
+		}
+		return 0;
+	}
+
+	String sizeAsTupleQuery(Resource... contexts) {
+		String query = COUNT_EVERYTHING;
+		if (contexts != null && isQuadMode()) {
+			if (contexts.length == 1 && contexts[0].isIRI()) {
+				query = "SELECT (COUNT(*) AS ?count) WHERE { GRAPH <" + ((IRI) contexts[0]).stringValue()
+						+ "> { ?s ?p ?o}}";
+			} else if (contexts.length > 0) {
+				String graphs = Arrays.stream(contexts)
+						.filter(Resource::isIRI)
+						.map(Resource::stringValue)
+						.map(s -> "FROM <" + s + ">")
+						.collect(Collectors.joining(" "));
+				query = "SELECT (COUNT(*) AS ?count) " + graphs + "WHERE { ?s ?p ?o}";
 			}
-			return i;
 		}
+		return query;
 	}
 
 	@Override

core/repository/sparql/src/test/java/org/eclipse/rdf4j/repository/sparql/SPARQLConnectionTest.java

@@ -12,13 +12,19 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.eclipse.rdf4j.model.util.Values.iri;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.lang.ref.WeakReference;
 
 import org.eclipse.rdf4j.http.client.SPARQLProtocolSession;
 import org.eclipse.rdf4j.model.IRI;
@@ -27,10 +33,18 @@
 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.impl.MapBindingSet;
+import org.eclipse.rdf4j.query.impl.SimpleBinding;
+import org.eclipse.rdf4j.query.impl.TupleQueryResultBuilder;
+import org.eclipse.rdf4j.query.parser.ParsedQuery;
+import org.eclipse.rdf4j.query.parser.sparql.SPARQLParser;
+import org.eclipse.rdf4j.query.parser.sparql.SPARQLParserFactory;
 import org.eclipse.rdf4j.rio.ParserConfig;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.invocation.InvocationOnMock;
 
 public class SPARQLConnectionTest {
 
@@ -100,6 +114,26 @@ public void testAddSingleContextHandling() throws Exception {
 		assertThat(sparqlUpdate).containsPattern(expectedAddPattern).containsPattern(expectedRemovePattern);
 	}
 
+	@Test
+	public void testSizeQuery() throws Exception {
+
+		String sizeAsTupleQuery = subject.sizeAsTupleQuery();
+		ParsedQuery query = new SPARQLParserFactory().getParser().parseQuery(sizeAsTupleQuery, "http://example.org/");
+		assertNotNull(query);
+
+		sizeAsTupleQuery = subject.sizeAsTupleQuery(vf.createIRI("urn:g1"));
+		query = new SPARQLParserFactory().getParser().parseQuery(sizeAsTupleQuery, "http://example.org/");
+		assertNotNull(query);
+
+		sizeAsTupleQuery = subject.sizeAsTupleQuery(vf.createIRI("urn:g1"), vf.createIRI("urn:g2"));
+		query = new SPARQLParserFactory().getParser().parseQuery(sizeAsTupleQuery, "http://example.org/");
+		assertNotNull(query);
+
+		sizeAsTupleQuery = subject.sizeAsTupleQuery(vf.createIRI("urn:g1"), vf.createBNode());
+		query = new SPARQLParserFactory().getParser().parseQuery(sizeAsTupleQuery, "http://example.org/");
+		assertNotNull(query);
+	}
+
 	@Test
 	public void testAddMultipleContextHandling() throws Exception {
 		ArgumentCaptor<String> sparqlUpdateCaptor = ArgumentCaptor.forClass(String.class);