GH-5686 feat: add lmdb page-cardinality estimator package and regressions

Author: hmottestad

core/sail/lmdb/src/main/java/org/eclipse/rdf4j/sail/lmdb/estimate/LmdbBtreeRangeCounter.java

@@ -0,0 +1,295 @@
+/*******************************************************************************
+ * Copyright (c) 2026 Eclipse RDF4J contributors.
+ *
+ * All rights reserved. This program and the accompanying materials
+ * are made available under the terms of the Eclipse Distribution License v1.0
+ * which accompanies this distribution, and is available at
+ * http://www.eclipse.org/org/documents/edl-v10.php.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ *******************************************************************************/
+// Some portions generated by Codex
+package org.eclipse.rdf4j.sail.lmdb.estimate;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+
+import org.eclipse.rdf4j.sail.lmdb.util.GroupMatcher;
+
+final class LmdbBtreeRangeCounter {
+
+	private final LmdbDataFile dataFile;
+	private final LmdbMeta meta;
+
+	LmdbBtreeRangeCounter(LmdbDataFile dataFile, LmdbMeta meta) {
+		this.dataFile = dataFile;
+		this.meta = meta;
+	}
+
+	RangeCountResult countRange(LmdbDb db, byte[] minKey, int minKeyLength, byte[] maxKey, int maxKeyLength,
+			GroupMatcher matcher) throws IOException {
+		RangeCountResult result = new RangeCountResult();
+		if (db.isEmpty()) {
+			return result;
+		}
+
+		SeekCursor cursor = seek(db, minKey, minKeyLength, result);
+		if (cursor == null) {
+			return result;
+		}
+
+		while (true) {
+			while (cursor.leafIndex < cursor.leafPage.numKeys) {
+				LmdbNode node = cursor.leafPage.node(cursor.leafIndex);
+				int cmpMax = LmdbKeyComparator.compare(cursor.leafPage.buffer, node.keyOffset, node.keySize, maxKey,
+						maxKeyLength);
+				if (cmpMax > 0) {
+					return result;
+				}
+				if (matcher == null || matches(node, cursor.leafPage.buffer, matcher)) {
+					result.entries += countNodeEntries(node, cursor.leafPage, result);
+				}
+				cursor.leafIndex++;
+			}
+
+			if (!advanceToNextLeaf(cursor, result)) {
+				return result;
+			}
+		}
+	}
+
+	byte[] findValueByExactKey(LmdbDb db, byte[] key, int keyLength, RangeCountResult ioStats) throws IOException {
+		if (db.isEmpty()) {
+			return null;
+		}
+		SeekCursor cursor = seek(db, key, keyLength, ioStats);
+		if (cursor == null || cursor.leafIndex >= cursor.leafPage.numKeys) {
+			return null;
+		}
+
+		LmdbNode node = cursor.leafPage.node(cursor.leafIndex);
+		int cmp = LmdbKeyComparator.compare(cursor.leafPage.buffer, node.keyOffset, node.keySize, key, keyLength);
+		if (cmp != 0) {
+			return null;
+		}
+
+		byte[] value = new byte[node.valueSize];
+		ByteBuffer duplicate = cursor.leafPage.buffer.duplicate();
+		duplicate.position(node.valueOffset);
+		duplicate.get(value, 0, node.valueSize);
+		return value;
+	}
+
+	private SeekCursor seek(LmdbDb db, byte[] searchKey, int searchKeyLength, RangeCountResult stats)
+			throws IOException {
+		List<BranchFrame> branchPath = new ArrayList<>();
+		LmdbPage page = dataFile.readPage(db.rootPgno, meta);
+		if (page.isBranch()) {
+			stats.branchPagesRead++;
+		}
+		if (page.isLeaf() || page.isLeaf2()) {
+			stats.leafPagesRead++;
+		}
+
+		while (page.isBranch()) {
+			if (page.numKeys == 0) {
+				throw new IOException("Corrupt branch page with zero keys: " + page.expectedPgno);
+			}
+			SearchResult search = findFirstGreaterOrEqual(page, searchKey, searchKeyLength, false);
+			int childIndex;
+			if (search.index >= page.numKeys) {
+				childIndex = page.numKeys - 1;
+			} else {
+				childIndex = search.index;
+				if (!search.exact) {
+					childIndex--;
+				}
+			}
+			if (childIndex < 0 || childIndex >= page.numKeys) {
+				throw new IOException("Corrupt branch descent index " + childIndex + " for page " + page.expectedPgno);
+			}
+			LmdbNode branchNode = page.node(childIndex);
+			branchPath.add(new BranchFrame(page, childIndex));
+			page = dataFile.readPage(branchNode.branchPgno, meta);
+			if (page.isBranch()) {
+				stats.branchPagesRead++;
+			}
+			if (page.isLeaf() || page.isLeaf2()) {
+				stats.leafPagesRead++;
+			}
+		}
+
+		if (!page.isLeaf() && !page.isLeaf2()) {
+			throw new IOException("Expected leaf page, found flags=" + page.flags + " on page " + page.expectedPgno);
+		}
+
+		SearchResult leafSearch = findFirstGreaterOrEqual(page, searchKey, searchKeyLength, true);
+		SeekCursor cursor = new SeekCursor(page, leafSearch.index, branchPath);
+		if (cursor.leafIndex >= cursor.leafPage.numKeys && !advanceToNextLeaf(cursor, stats)) {
+			return null;
+		}
+		return cursor;
+	}
+
+	private boolean advanceToNextLeaf(SeekCursor cursor, RangeCountResult stats) throws IOException {
+		while (!cursor.branchPath.isEmpty()) {
+			BranchFrame last = cursor.branchPath.get(cursor.branchPath.size() - 1);
+			int nextChild = last.childIndex + 1;
+			if (nextChild < last.page.numKeys) {
+				last.childIndex = nextChild;
+				LmdbNode nextNode = last.page.node(nextChild);
+				LmdbPage page = dataFile.readPage(nextNode.branchPgno, meta);
+				if (page.isBranch()) {
+					stats.branchPagesRead++;
+				}
+				if (page.isLeaf() || page.isLeaf2()) {
+					stats.leafPagesRead++;
+				}
+
+				while (page.isBranch()) {
+					if (page.numKeys == 0) {
+						throw new IOException("Corrupt branch page with zero keys: " + page.expectedPgno);
+					}
+					cursor.branchPath.add(new BranchFrame(page, 0));
+					LmdbNode firstNode = page.node(0);
+					page = dataFile.readPage(firstNode.branchPgno, meta);
+					if (page.isBranch()) {
+						stats.branchPagesRead++;
+					}
+					if (page.isLeaf() || page.isLeaf2()) {
+						stats.leafPagesRead++;
+					}
+				}
+
+				if (!page.isLeaf() && !page.isLeaf2()) {
+					throw new IOException("Expected leaf page while advancing, found flags=" + page.flags);
+				}
+				cursor.leafPage = page;
+				cursor.leafIndex = 0;
+				return true;
+			}
+			cursor.branchPath.remove(cursor.branchPath.size() - 1);
+		}
+		return false;
+	}
+
+	private SearchResult findFirstGreaterOrEqual(LmdbPage page, byte[] key, int keyLength, boolean leafSearch)
+			throws IOException {
+		if (page.numKeys == 0) {
+			return new SearchResult(0, false);
+		}
+
+		int low;
+		int high = page.numKeys - 1;
+		if (leafSearch || page.isLeaf()) {
+			low = 0;
+		} else {
+			low = 1;
+		}
+		int index = 0;
+		int rc = -1;
+
+		while (low <= high) {
+			index = (low + high) >>> 1;
+			LmdbNode node = page.node(index);
+			rc = LmdbKeyComparator.compare(key, keyLength, page.buffer, node.keyOffset, node.keySize);
+			if (rc == 0) {
+				return new SearchResult(index, true);
+			}
+			if (rc > 0) {
+				low = index + 1;
+			} else {
+				high = index - 1;
+			}
+		}
+
+		if (rc > 0) {
+			index++;
+		}
+		if (low > high) {
+			index = low;
+		}
+		return new SearchResult(index, false);
+	}
+
+	private boolean matches(LmdbNode node, ByteBuffer pageBuffer, GroupMatcher matcher) {
+		ByteBuffer keySlice = pageBuffer.duplicate();
+		keySlice.order(pageBuffer.order());
+		keySlice.position(node.keyOffset);
+		keySlice.limit(node.keyOffset + node.keySize);
+		ByteBuffer keyView = keySlice.slice();
+		keyView.order(pageBuffer.order());
+		return matcher.matches(keyView);
+	}
+
+	private long countNodeEntries(LmdbNode node, LmdbPage page, RangeCountResult stats) throws IOException {
+		if ((node.nodeFlags & LmdbFormat.F_SUBDATA) != 0 && node.valueSize >= LmdbFormat.META_DB_SIZE) {
+			ByteBuffer dup = page.buffer.duplicate();
+			dup.order(page.buffer.order());
+			dup.position(node.valueOffset);
+			LmdbDb subDb = LmdbDb.parse(dup, node.valueOffset);
+			return subDb.entries;
+		}
+		if ((node.nodeFlags & LmdbFormat.F_DUPDATA) != 0 && node.valueSize >= LmdbFormat.PAGE_HEADER_SIZE) {
+			return countSubPageEntries(page.buffer, node.valueOffset, node.valueSize);
+		}
+		if ((node.nodeFlags & LmdbFormat.F_BIGDATA) != 0 && node.valueSize >= Long.BYTES) {
+			long overflowPgno = page.buffer.getLong(node.valueOffset);
+			LmdbPage overflowPage = dataFile.readPage(overflowPgno, meta);
+			stats.overflowPagesRead += Math.max(overflowPage.overflowPages, 1);
+		}
+		return 1;
+	}
+
+	private long countSubPageEntries(ByteBuffer buffer, int offset, int length) {
+		if (offset + length > buffer.limit()) {
+			return 0;
+		}
+		int flags = LmdbFormat.unsignedShort(buffer, offset + 10);
+		int lower = LmdbFormat.unsignedShort(buffer, offset + 12);
+		if ((flags & LmdbFormat.P_LEAF2) != 0) {
+			int keySize = LmdbFormat.unsignedShort(buffer, offset + 8);
+			if (keySize <= 0) {
+				return 0;
+			}
+			int bytes = Math.max(lower - LmdbFormat.PAGE_HEADER_SIZE, 0);
+			return bytes / keySize;
+		}
+		return Math.max(LmdbFormat.numKeys(lower), 0);
+	}
+
+	private static final class SearchResult {
+		final int index;
+		final boolean exact;
+
+		SearchResult(int index, boolean exact) {
+			this.index = index;
+			this.exact = exact;
+		}
+	}
+
+	private static final class BranchFrame {
+		final LmdbPage page;
+		int childIndex;
+
+		BranchFrame(LmdbPage page, int childIndex) {
+			this.page = page;
+			this.childIndex = childIndex;
+		}
+	}
+
+	private static final class SeekCursor {
+		LmdbPage leafPage;
+		int leafIndex;
+		final List<BranchFrame> branchPath;
+
+		SeekCursor(LmdbPage leafPage, int leafIndex, List<BranchFrame> branchPath) {
+			this.leafPage = leafPage;
+			this.leafIndex = leafIndex;
+			this.branchPath = branchPath;
+		}
+	}
+
+}