GH-5686 perf: add lmdb estimator-driven join optimization

Author: hmottestad

core/queryalgebra/evaluation/src/main/java/org/eclipse/rdf4j/query/algebra/evaluation/optimizer/QueryJoinOptimizer.java

@@ -222,6 +222,24 @@
 			<plugin>
 				<artifactId>maven-assembly-plugin</artifactId>
 			</plugin>
+			<plugin>
+				<groupId>org.apache.maven.plugins</groupId>
+				<artifactId>maven-compiler-plugin</artifactId>
+				<executions>
+					<execution>
+						<id>default-testCompile</id>
+						<configuration>
+							<annotationProcessorPaths>
+								<path>
+									<groupId>org.openjdk.jmh</groupId>
+									<artifactId>jmh-generator-annprocess</artifactId>
+									<version>${jmhVersion}</version>
+								</path>
+							</annotationProcessorPaths>
+						</configuration>
+					</execution>
+				</executions>
+			</plugin>
 		</plugins>
 	</build>
 </project>

core/queryalgebra/evaluation/src/test/java/org/eclipse/rdf4j/query/algebra/evaluation/impl/QueryJoinOptimizerTest.java

@@ -8,9 +8,12 @@
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *******************************************************************************/
+// Some portions generated by Codex
 package org.eclipse.rdf4j.sail.lmdb;
 
 import java.io.IOException;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 import org.eclipse.rdf4j.model.IRI;
 import org.eclipse.rdf4j.model.Resource;
@@ -28,14 +31,19 @@
 class LmdbEvaluationStatistics extends EvaluationStatistics {
 
 	private static final Logger log = LoggerFactory.getLogger(LmdbEvaluationStatistics.class);
+	private static final int SHARED_CACHE_MAX_ENTRIES = 262_144;
+	private static final Map<SharedCardinalityKey, Double> sharedCardinalityCache = new ConcurrentHashMap<>();
 
 	private final ValueStore valueStore;
 
 	private final TripleStore tripleStore;
+	private final int tripleStoreIdentity;
+	private final Map<CardinalityKey, Double> cardinalityCache = new ConcurrentHashMap<>();
 
 	public LmdbEvaluationStatistics(ValueStore valueStore, TripleStore tripleStore) {
 		this.valueStore = valueStore;
 		this.tripleStore = tripleStore;
+		this.tripleStoreIdentity = System.identityHashCode(tripleStore);
 	}
 
 	@Override
@@ -112,6 +120,116 @@ private double cardinality(Resource subj, IRI pred, Value obj, Resource context)
 			}
 		}
 
-		return tripleStore.cardinality(subjID, predID, objID, contextID);
+		CardinalityKey key = new CardinalityKey(subjID, predID, objID, contextID);
+		Double cached = cardinalityCache.get(key);
+		if (cached != null) {
+			return cached;
+		}
+
+		long revisionId = valueStore.getRevision().getRevisionId();
+		SharedCardinalityKey sharedKey = new SharedCardinalityKey(tripleStoreIdentity, revisionId, key);
+		Double sharedCached = sharedCardinalityCache.get(sharedKey);
+		if (sharedCached != null) {
+			cardinalityCache.put(key, sharedCached);
+			return sharedCached;
+		}
+
+		double cardinality = tripleStore.cardinality(subjID, predID, objID, contextID);
+		cardinalityCache.put(key, cardinality);
+		cacheSharedCardinality(sharedKey, cardinality);
+		return cardinality;
+	}
+
+	private static void cacheSharedCardinality(SharedCardinalityKey key, double cardinality) {
+		if (sharedCardinalityCache.size() >= SHARED_CACHE_MAX_ENTRIES) {
+			sharedCardinalityCache.clear();
+		}
+		sharedCardinalityCache.put(key, cardinality);
+	}
+
+	private static final class CardinalityKey {
+		private final long subjID;
+		private final long predID;
+		private final long objID;
+		private final long contextID;
+
+		private CardinalityKey(long subjID, long predID, long objID, long contextID) {
+			this.subjID = subjID;
+			this.predID = predID;
+			this.objID = objID;
+			this.contextID = contextID;
+		}
+
+		@Override
+		public boolean equals(Object obj) {
+			if (this == obj) {
+				return true;
+			}
+			if (!(obj instanceof CardinalityKey)) {
+				return false;
+			}
+			CardinalityKey other = (CardinalityKey) obj;
+			return subjID == other.subjID
+					&& predID == other.predID
+					&& objID == other.objID
+					&& contextID == other.contextID;
+		}
+
+		@Override
+		public int hashCode() {
+			int result = Long.hashCode(subjID);
+			result = 31 * result + Long.hashCode(predID);
+			result = 31 * result + Long.hashCode(objID);
+			result = 31 * result + Long.hashCode(contextID);
+			return result;
+		}
+	}
+
+	private static final class SharedCardinalityKey {
+		private final int tripleStoreIdentity;
+		private final long revisionId;
+		private final long subjID;
+		private final long predID;
+		private final long objID;
+		private final long contextID;
+		private final int hashCode;
+
+		private SharedCardinalityKey(int tripleStoreIdentity, long revisionId, CardinalityKey cardinalityKey) {
+			this.tripleStoreIdentity = tripleStoreIdentity;
+			this.revisionId = revisionId;
+			this.subjID = cardinalityKey.subjID;
+			this.predID = cardinalityKey.predID;
+			this.objID = cardinalityKey.objID;
+			this.contextID = cardinalityKey.contextID;
+			int hash = Integer.hashCode(tripleStoreIdentity);
+			hash = 31 * hash + Long.hashCode(revisionId);
+			hash = 31 * hash + Long.hashCode(subjID);
+			hash = 31 * hash + Long.hashCode(predID);
+			hash = 31 * hash + Long.hashCode(objID);
+			hash = 31 * hash + Long.hashCode(contextID);
+			this.hashCode = hash;
+		}
+
+		@Override
+		public boolean equals(Object obj) {
+			if (this == obj) {
+				return true;
+			}
+			if (!(obj instanceof SharedCardinalityKey)) {
+				return false;
+			}
+			SharedCardinalityKey other = (SharedCardinalityKey) obj;
+			return tripleStoreIdentity == other.tripleStoreIdentity
+					&& revisionId == other.revisionId
+					&& subjID == other.subjID
+					&& predID == other.predID
+					&& objID == other.objID
+					&& contextID == other.contextID;
+		}
+
+		@Override
+		public int hashCode() {
+			return hashCode;
+		}
 	}
 }

core/sail/lmdb/pom.xml

@@ -8,6 +8,7 @@
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *******************************************************************************/
+// Some portions generated by Codex
 package org.eclipse.rdf4j.sail.lmdb;
 
 import static org.eclipse.rdf4j.sail.lmdb.LmdbUtil.E;
@@ -52,6 +53,7 @@
 import static org.lwjgl.util.lmdb.LMDB.mdb_txn_abort;
 import static org.lwjgl.util.lmdb.LMDB.mdb_txn_begin;
 import static org.lwjgl.util.lmdb.LMDB.mdb_txn_commit;
+import static org.lwjgl.util.lmdb.LMDB.mdb_txn_id;
 
 import java.io.Closeable;
 import java.io.File;
@@ -85,6 +87,7 @@
 import org.eclipse.rdf4j.sail.lmdb.TxnRecordCache.Record;
 import org.eclipse.rdf4j.sail.lmdb.TxnRecordCache.RecordCacheIterator;
 import org.eclipse.rdf4j.sail.lmdb.config.LmdbStoreConfig;
+import org.eclipse.rdf4j.sail.lmdb.estimate.LmdbPageCardinalityEstimator;
 import org.eclipse.rdf4j.sail.lmdb.util.GroupMatcher;
 import org.eclipse.rdf4j.sail.lmdb.util.IndexKeyWriters;
 import org.lwjgl.PointerBuffer;
@@ -152,6 +155,7 @@ class TripleStore implements Closeable {
 	 * The directory that is used to store the index files.
 	 */
 	private final File dir;
+	private final File dataMdbFile;
 	/**
 	 * Object containing meta-data for the triple store.
 	 */
@@ -166,9 +170,11 @@ class TripleStore implements Closeable {
 	private final int contextsDbi;
 	private int pageSize;
 	private final boolean autoGrow;
+	private final boolean usePageEstimator;
 	private long mapSize;
 	private long writeTxn;
 	private final TxnManager txnManager;
+	private final LmdbPageCardinalityEstimator pageEstimator;
 
 	private TxnRecordCache recordCache = null;
 
@@ -195,8 +201,10 @@ public int compareRegion(ByteBuffer array1, int startIdx1, ByteBuffer array2, in
 
 	TripleStore(File dir, LmdbStoreConfig config, ValueStore valueStore) throws IOException, SailException {
 		this.dir = dir;
+		this.dataMdbFile = new File(dir, "data.mdb");
 		boolean forceSync = config.getForceSync();
 		this.autoGrow = config.getAutoGrow();
+		this.usePageEstimator = config.getPageCardinalityEstimator();
 		this.valueStore = valueStore;
 
 		// create directory if it not exists
@@ -229,6 +237,7 @@ public int compareRegion(ByteBuffer array1, int startIdx1, ByteBuffer array2, in
 		});
 
 		txnManager = new TxnManager(env, Mode.RESET);
+		pageEstimator = usePageEstimator ? new LmdbPageCardinalityEstimator(dataMdbFile) : null;
 
 		File propFile = new File(this.dir, PROPERTIES_FILE);
 		String indexSpecStr = config.getTripleIndexes();
@@ -462,6 +471,14 @@ public void close() throws IOException {
 			endTransaction(false);
 
 			List<Throwable> caughtExceptions = new ArrayList<>();
+			if (pageEstimator != null) {
+				try {
+					pageEstimator.close();
+				} catch (Throwable e) {
+					logger.warn("Failed to close page estimator", e);
+					caughtExceptions.add(e);
+				}
+			}
 			for (TripleIndex index : indexes) {
 				try {
 					index.close();
@@ -674,6 +691,64 @@ protected void filterUsedIds(Collection<Long> ids) throws IOException {
 
 	protected double cardinality(long subj, long pred, long obj, long context) throws IOException {
 		TripleIndex index = getBestIndex(subj, pred, obj, context);
+		if (!usePageEstimator) {
+			return cardinalityUsingSamplingEstimator(index, subj, pred, obj, context);
+		}
+
+		try {
+			return cardinalityUsingPageEstimator(index, subj, pred, obj, context);
+		} catch (IOException | RuntimeException e) {
+			logger.warn("Page-walk cardinality estimator failed for index {}, falling back to sampling",
+					new String(index.getFieldSeq()), e);
+			return cardinalityUsingSamplingEstimator(index, subj, pred, obj, context);
+		}
+	}
+
+	private double cardinalityUsingPageEstimator(TripleIndex index, long subj, long pred, long obj, long context)
+			throws IOException {
+		LmdbPageCardinalityEstimator estimator = pageEstimator;
+		if (estimator == null) {
+			return cardinalityUsingSamplingEstimator(index, subj, pred, obj, context);
+		}
+		int relevantParts = index.getPatternScore(subj, pred, obj, context);
+		final String explicitDbName = new String(index.getFieldSeq());
+		final String inferredDbName = explicitDbName + "-inf";
+
+		return txnManager.doWith((stack, txn) -> {
+			long txnId = mdb_txn_id(txn);
+			if (relevantParts == 0) {
+				long explicitEntries = estimator.totalEntries(txnId, explicitDbName);
+				long inferredEntries = estimator.totalEntries(txnId, inferredDbName);
+				return (double) (explicitEntries + inferredEntries);
+			}
+
+			ByteBuffer minKeyBuffer = ByteBuffer.allocate(MAX_KEY_LENGTH);
+			index.getMinKey(minKeyBuffer, subj, pred, obj, context);
+			minKeyBuffer.flip();
+			byte[] minKey = toArray(minKeyBuffer);
+
+			ByteBuffer maxKeyBuffer = ByteBuffer.allocate(MAX_KEY_LENGTH);
+			index.getMaxKey(maxKeyBuffer, subj, pred, obj, context);
+			maxKeyBuffer.flip();
+			byte[] maxKey = toArray(maxKeyBuffer);
+
+			GroupMatcher matcher = index.createMatcher(subj, pred, obj, context);
+			long explicitCount = estimator.estimateEntries(txnId, explicitDbName, minKey, minKey.length, maxKey,
+					maxKey.length, matcher);
+			long inferredCount = estimator.estimateEntries(txnId, inferredDbName, minKey, minKey.length, maxKey,
+					maxKey.length, matcher);
+			return (double) (explicitCount + inferredCount);
+		});
+	}
+
+	private static byte[] toArray(ByteBuffer buffer) {
+		byte[] data = new byte[buffer.remaining()];
+		buffer.get(data);
+		return data;
+	}
+
+	private double cardinalityUsingSamplingEstimator(TripleIndex index, long subj, long pred, long obj, long context)
+			throws IOException {
 
 		int relevantParts = index.getPatternScore(subj, pred, obj, context);
 		if (relevantParts == 0) {