GH-5685 initial implementation

Author: hmottestad

core/sail/lmdb/src/main/java/org/eclipse/rdf4j/sail/lmdb/TripleStore.java

@@ -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) {