diff --git a/README.md b/README.md index 1d69984f3..0d1e92abe 100644 --- a/README.md +++ b/README.md @@ -432,13 +432,75 @@ the `expectedVersion` option is used. ### `db.getEstimatedKeyCount(): number` Retrieves the estimated number of keys in the database. This is an alias for -`db.getDBIntProperty('rocksdb.estimate-num-keys')`. +`db.getDBIntProperty('rocksdb.estimate-num-keys')`; use `estimateCount()` for range support and a +confidence indicator. ```typescript const estimated = db.getEstimatedKeyCount(); console.log(estimated); ``` +### `db.estimateCount(options?: CountEstimateOptions): CountEstimate` + +Estimates the number of keys in the database, or within a key range, returning +`{ count, confidence }`. Unlike `getKeysCount()`, this never iterates: the estimate is derived +from RocksDB statistics (memtable stats plus approximate SST sizes converted through the entry +density of the SSTs overlapping the range), so its cost scales with the number of SSTs overlapping +the range rather than the number of keys. Reading cold table properties can do I/O through the +table cache, so narrow ranges are preferable. A start-only range is computed as the +whole-database estimate minus the complement, so it does the work of the range _below_ `start`. +Accuracy improves with range size. Resolution is bounded by SST data-block granularity, so a range +narrower than a block is unreliable in either direction: it may over-report or report 0 for present +keys, and its low `confidence` is the signal. Recently deleted or overwritten entries may be counted +until compaction. +Estimates always reflect committed state; writes pending in a transaction are not included. Set +`reverse: true` to use `getRange()`'s reverse convention (`start` is the upper bound and `end` is +the lower bound). An inverted range (`start` ≥ `end`) returns +`{ count: 0, confidence: 1 }`. + +`confidence` is a heuristic 0–1 indicator of how trustworthy `count` is — exactly 1 only when the +count is exact. It is derived from the estimate's resolution (data-block/memtable-sampling +granularity relative to the count), the tombstone fraction of the overlapping SSTs, and — for +start-only ranges — the error compounded by complement subtraction. Treat it as an ordering +signal (e.g. when to trust an estimate for query planning vs fall back to a heuristic), not a +statistical bound. + +```typescript +const { count, confidence } = db.estimateCount({ start: 'a', end: 'z' }); +``` + +### `db.createCountEstimator(options?: CountEstimatorOptions): CountEstimator` + +Creates an estimator that progressively refines a range count estimate while the range is being +iterated — useful for reporting a total alongside a page of results without scanning the full +range. Before any traversal, `estimate()` returns the pure statistical estimate (same as +`estimateCount(range)`). As the caller reports progress with `advance(lastKey, count)` (e.g. once +per page), `estimate()` returns the exact traversed count plus a statistical estimate of the +remainder, calibrated by the observed ratio of actual-to-estimated entries over the portion +already traversed — so the count converges toward the exact total. `confidence` is the +exactness-weighted blend of the traversed portion and the remainder's confidence, so it approaches 1 +as the exact portion grows, although a checkpoint may decrease when calibration makes a large +correction. +Each checkpoint reads committed state, so a traversal performed against a transaction snapshot may +be calibrated against data committed after that snapshot. +When traversal completes, call `finish()` and `estimate()` returns the exact count with +confidence 1. Reverse ranges follow `getRange`: set `start` to the upper bound and `end` to the +lower bound, then set `reverse: true`. The caller owns the progress contract: cursors must move +monotonically through the range and each entry must be reported exactly once. + +```typescript +const range = { start: 'a', end: 'z' }; +const estimator = db.createCountEstimator(range); +let lastKey; +let pageSize = 0; +for (const { key } of db.getRange({ ...range, limit: 25 })) { + lastKey = key; + pageSize++; +} +estimator.advance(lastKey, pageSize); +const { count, confidence } = estimator.estimate(); +``` + ### `db.getKeys(options?: IteratorOptions): ExtendedIterable` Retrieves all keys within a range. diff --git a/src/binding/database/database.cpp b/src/binding/database/database.cpp index 5beb92f72..76d02bacc 100644 --- a/src/binding/database/database.cpp +++ b/src/binding/database/database.cpp @@ -922,6 +922,205 @@ napi_value Database::GetCount(napi_env env, napi_callback_info info) { return result; } +struct RangeEstimate { + double count = 0; + double memtableCount = 0; + double sstCount = 0; + double entriesPerBlock = 0; + double liveFraction = 1; + bool degraded = false; +}; + +/** + * Estimates the number of live keys in `[start, end)` from RocksDB statistics + * alone — no iteration: + * + * - memtable portion: `GetApproximateMemTableStats` returns an entry count + * directly (it counts all memtable entries, including tombstones and + * overwrites, so it can over-report a recently-deleted range). + * - SST portion: the approximate file bytes covered by the range + * (`GetApproximateSizes`) converted to entries using the live-entry density + * of only the SSTs overlapping the range (`GetPropertiesOfTablesInRange`: + * `(num_entries - num_deletions) / file bytes`). Using range-local table + * properties keeps the density honest when entry sizes vary across the + * keyspace, and needs no cache/invalidation. + * + * Overlapping versions of a key in multiple levels are counted once per + * level, so the estimate skews high on heavily-overwritten ranges until + * compaction; resolution is bounded by SST data-block granularity, so tiny + * ranges can over-report or report zero for present keys. + */ +static RangeEstimate estimateRangeCount(rocksdb::DB* db, rocksdb::ColumnFamilyHandle* cf, const rocksdb::Slice& start, const rocksdb::Slice& end) { + rocksdb::Range range(start, end); + RangeEstimate result; + + uint64_t memtableCount = 0; + uint64_t memtableSize = 0; + db->GetApproximateMemTableStats(cf, range, &memtableCount, &memtableSize); + result.memtableCount = static_cast(memtableCount); + result.count = result.memtableCount; + + rocksdb::SizeApproximationOptions sizeOptions; + sizeOptions.include_memtables = false; + sizeOptions.files_size_error_margin = 0.1; + uint64_t sstBytes = 0; + rocksdb::Status status = db->GetApproximateSizes(sizeOptions, cf, &range, 1, &sstBytes); + if (!status.ok()) { + result.degraded = true; + return result; + } + + rocksdb::TablePropertiesCollection props; + status = db->GetPropertiesOfTablesInRange(cf, &range, 1, &props); + uint64_t entries = 0; + uint64_t deletions = 0; + uint64_t fileBytes = 0; + uint64_t dataBlocks = 0; + if (status.ok()) { + for (const auto& prop : props) { + if (!prop.second) { + result.degraded = true; + continue; + } + const rocksdb::TableProperties& p = *prop.second; + entries += p.num_entries; + deletions += p.num_deletions; + // Approximate the on-disk file size covered by table properties; + // GetApproximateSizes offsets span data + index + filter blocks, + // so the density denominator must too. + fileBytes += p.data_size + p.index_size + p.filter_size; + dataBlocks += p.num_data_blocks; + } + } else { + result.degraded = true; + } + if (entries <= deletions || fileBytes == 0) { + // A nonzero byte estimate without density leaves the SST portion unknown. + result.degraded = result.degraded || sstBytes != 0; + return result; + } + if (sstBytes == 0) { + result.degraded = true; + return result; + } + + double density = static_cast(entries - deletions) / static_cast(fileBytes); + result.sstCount = static_cast(sstBytes) * density; + result.count += result.sstCount; + result.entriesPerBlock = dataBlocks > 0 + ? static_cast(entries - deletions) / static_cast(dataBlocks) + : 0; + result.liveFraction = static_cast(entries - deletions) / static_cast(entries); + return result; +} + +/** + * Heuristic [0, 1] trust indicator for a range estimate — 1 only when exact. + * Combines the estimate's resolution (SST portion is quantized to data-block + * granularity, memtable counts to skip-list sampling granularity) with the + * tombstone fraction of the overlapping SSTs (a proxy for overwrite/delete + * skew the estimate cannot see). + */ +static double estimateConfidence(const RangeEstimate& est) { + if (est.degraded) { + return 0.1; + } + if (est.count <= 0) { + return 0.95; + } + double sstResolution = std::max(est.entriesPerBlock, 1.0); + double memtableResolution = 8; + double resolution = (est.sstCount * sstResolution + est.memtableCount * memtableResolution) / est.count; + double granularity = est.count / (est.count + resolution); + return granularity * (0.5 + 0.5 * est.liveFraction); +} + +/** + * Estimates the number of keys within a range without iterating, returning + * `{ count, confidence }`. Both keys are optional buffers; an open-ended side + * is handled by subtracting the complementary range from the + * whole-column-family `estimate-num-keys` (an empty slice is the *smallest* + * key, so it must never be passed as an upper bound). + * + * @example + * ```typescript + * const db = NativeDatabase.open('path/to/db'); + * const { count, confidence } = db.estimateCount(startBuffer, endBuffer); + * ``` + */ +napi_value Database::EstimateCount(napi_env env, napi_callback_info info) { + NAPI_METHOD_ARGV(2); + UNWRAP_DB_HANDLE_AND_OPEN(); + ACQUIRE_OPERATIONS_LOCK(); + + rocksdb::DB* db = (*dbHandle)->descriptor->db.get(); + rocksdb::ColumnFamilyHandle* cf = (*dbHandle)->getColumnFamilyHandle(); + + // N-API may return a null data pointer for a zero-length buffer. + void* startData = nullptr; + size_t startLength = 0; + napi_valuetype startType; + NAPI_STATUS_THROWS(::napi_typeof(env, argv[0], &startType)); + bool hasStart = startType != napi_undefined && startType != napi_null; + if (hasStart) { + NAPI_STATUS_THROWS(::napi_get_buffer_info(env, argv[0], &startData, &startLength)); + } + + void* endData = nullptr; + size_t endLength = 0; + napi_valuetype endType; + NAPI_STATUS_THROWS(::napi_typeof(env, argv[1], &endType)); + bool hasEnd = endType != napi_undefined && endType != napi_null; + if (hasEnd) { + NAPI_STATUS_THROWS(::napi_get_buffer_info(env, argv[1], &endData, &endLength)); + } + + rocksdb::Slice startSlice(startLength ? static_cast(startData) : "", startLength); + rocksdb::Slice endSlice(endLength ? static_cast(endData) : "", endLength); + + double estimate = 0; + double confidence = 0; + if (!hasEnd) { + uint64_t totalKeys = 0; + bool totalOk = db->GetIntProperty(cf, rocksdb::DB::Properties::kEstimateNumKeys, &totalKeys); + double total = static_cast(totalKeys); + if (!totalOk) { + estimate = 0; + confidence = 0; + } else if (!hasStart || startLength == 0 || totalKeys == 0) { + estimate = total; + confidence = totalKeys == 0 ? 0.95 : 0.9; + } else { + // No upper bound: estimate [start, ∞) as total minus [min, start). + RangeEstimate complement = estimateRangeCount(db, cf, rocksdb::Slice(), startSlice); + estimate = std::max(0.0, total - complement.count); + double share = estimate / std::max(estimate + complement.count, 1.0); + confidence = std::min(0.9, estimateConfidence(complement)) * share; + } + } else if (endLength == 0 || startSlice.compare(endSlice) >= 0) { + // Empty end bound (below every key) or inverted/empty range: + // GetApproximateSizes would underflow (end offset minus start offset + // in uint64). Comparator is always bytewise (db_descriptor.cpp), so + // Slice::compare matches key order. Empty by construction, so exact. + estimate = 0; + confidence = 1.0; + } else { + RangeEstimate rangeEstimate = estimateRangeCount(db, cf, startSlice, endSlice); + estimate = rangeEstimate.count; + confidence = estimateConfidence(rangeEstimate); + } + + napi_value result; + NAPI_STATUS_THROWS(::napi_create_object(env, &result)); + napi_value countValue; + NAPI_STATUS_THROWS(::napi_create_double(env, std::round(estimate), &countValue)); + NAPI_STATUS_THROWS(::napi_set_named_property(env, result, "count", countValue)); + napi_value confidenceValue; + NAPI_STATUS_THROWS(::napi_create_double(env, confidence, &confidenceValue)); + NAPI_STATUS_THROWS(::napi_set_named_property(env, result, "confidence", confidenceValue)); + return result; +} + napi_value Database::GetMonotonicTimestamp(napi_env env, napi_callback_info info) { NAPI_METHOD(); UNWRAP_DB_HANDLE_AND_OPEN(); @@ -1987,6 +2186,7 @@ void Database::Init(napi_env env, napi_value exports) { { "destroy", nullptr, Destroy, nullptr, nullptr, nullptr, napi_default, nullptr }, { "drop", nullptr, Drop, nullptr, nullptr, nullptr, napi_default, nullptr }, { "dropSync", nullptr, DropSync, nullptr, nullptr, nullptr, napi_default, nullptr }, + { "estimateCount", nullptr, EstimateCount, nullptr, nullptr, nullptr, napi_default, nullptr }, { "flush", nullptr, Flush, nullptr, nullptr, nullptr, napi_default, nullptr }, { "flushSync", nullptr, FlushSync, nullptr, nullptr, nullptr, napi_default, nullptr }, { "get", nullptr, Get, nullptr, nullptr, nullptr, napi_default, nullptr }, diff --git a/src/binding/database/database.h b/src/binding/database/database.h index a1272c906..1c6a11d97 100644 --- a/src/binding/database/database.h +++ b/src/binding/database/database.h @@ -285,6 +285,7 @@ struct Database final { static napi_value FlushSync(napi_env env, napi_callback_info info); static napi_value Get(napi_env env, napi_callback_info info); static napi_value GetCompression(napi_env env, napi_callback_info info); + static napi_value EstimateCount(napi_env env, napi_callback_info info); static napi_value GetCount(napi_env env, napi_callback_info info); static napi_value GetDBIntProperty(napi_env env, napi_callback_info info); static napi_value GetDBProperty(napi_env env, napi_callback_info info); diff --git a/src/count-estimator.ts b/src/count-estimator.ts new file mode 100644 index 000000000..29d04c4e7 --- /dev/null +++ b/src/count-estimator.ts @@ -0,0 +1,172 @@ +import type { CountEstimate, CountEstimateOptions } from './dbi.ts'; +import type { Key } from './encoding.ts'; +import type { Store } from './store.ts'; + +export interface CountEstimatorOptions extends CountEstimateOptions {} + +/** + * Below this many traversed entries the observed density is too noisy to + * calibrate with, so `estimate()` returns the raw statistical estimate. + */ +const CALIBRATION_MIN_TRAVERSED = 16; +const CALIBRATION_MIN_CONFIDENCE = 0.8; + +/** + * Bounds on how far the observed-vs-estimated ratio may scale the remainder; + * the statistical estimate can be arbitrarily wrong at data-block granularity + * and an unclamped ratio would let one bad sample dominate. + */ +const CALIBRATION_MAX = 8; + +/** + * Progressively refines a range key-count estimate as an iterator traverses + * the range. The estimate starts as the pure statistical estimate + * (`estimateCount`) and converges toward the exact count: the traversed + * portion is exact, and the statistical estimate of the untraversed remainder + * is calibrated by the observed ratio of actual-to-estimated entries over the + * portion already traversed. + * + * The estimator never touches the iterator itself — the caller reports + * progress with `advance(lastKey, count)` whenever it wants a checkpoint + * (e.g. once per page), then reads `estimate()`. The caller owns the + * progress contract: cursors must move monotonically through the range and + * each entry must be reported once (a re-reported page inflates the count + * undetectably). When traversal completes, call `finish()` so `estimate()` + * returns the exact total instead of adding a block-granular remainder. + * + * Each `estimate()` checkpoint queries RocksDB statistics for the two range + * segments (cost scales with the SSTs overlapping them, not with key count); + * results are memoized per checkpoint, so repeated reads between `advance()` + * calls are free. + */ +export class CountEstimator { + #store: Store; + #start: Key | Uint8Array | undefined; + #end: Key | Uint8Array | undefined; + #exclusiveStart: boolean; + #inclusiveEnd: boolean; + #reverse: boolean; + #cursor: Key | Uint8Array | undefined; + #traversed = 0; + #finished = false; + #memoized: CountEstimate | undefined; + + constructor(store: Store, options?: CountEstimatorOptions) { + this.#store = store; + this.#reverse = options?.reverse ?? false; + this.#start = this.#reverse ? options?.end : options?.start; + this.#end = this.#reverse ? options?.start : options?.end; + this.#exclusiveStart = options?.exclusiveStart ?? this.#reverse; + this.#inclusiveEnd = options?.inclusiveEnd ?? this.#reverse; + } + + /** + * The number of entries reported traversed so far. + */ + get traversed(): number { + return this.#traversed; + } + + /** + * Records that iteration has advanced through `count` more entries, ending + * at `lastKey`. Pass `count: 0` when an empty page has no last key. + */ + advance(lastKey: Key | Uint8Array | undefined, count = 1): void { + if (lastKey === undefined) { + if (count !== 0) { + throw new Error('CountEstimator.advance requires lastKey when count is nonzero'); + } + return; + } + this.#cursor = lastKey; + this.#traversed += count; + this.#memoized = undefined; + } + + /** + * Marks traversal of the range as complete: `estimate()` becomes the exact + * traversed count. + */ + finish(): void { + this.#finished = true; + } + + /** + * Estimates the total number of entries in the full range: the exact + * traversed count plus a calibrated statistical estimate of the remainder. + * `confidence` is the exactness-weighted blend of the traversed portion + * (exact) and the remainder's statistical confidence, so it approaches 1 + * as the exact portion grows (though a checkpoint can decrease when + * calibration makes a large correction), and is exactly 1 after `finish()`. + */ + estimate(): CountEstimate { + if (this.#finished) { + return { count: this.#traversed, confidence: 1 }; + } + if (this.#memoized !== undefined) { + return { ...this.#memoized }; + } + if (this.#cursor === undefined) { + this.#memoized = this.#store.estimateCount({ + start: this.#start, + end: this.#end, + exclusiveStart: this.#exclusiveStart, + inclusiveEnd: this.#inclusiveEnd, + }); + return { ...this.#memoized }; + } + + // The cursor entry itself belongs to the traversed side, so the + // remainder excludes it in both directions (an inclusive lower bound + // forward would count it twice and block convergence). The range's own + // bound flags stay with their original edge of the full range. + const traversedRange = this.#reverse + ? { start: this.#cursor, end: this.#end, inclusiveEnd: this.#inclusiveEnd } + : { + start: this.#start, + exclusiveStart: this.#exclusiveStart, + end: this.#cursor, + inclusiveEnd: true, + }; + const remainingRange = this.#reverse + ? { start: this.#start, exclusiveStart: this.#exclusiveStart, end: this.#cursor } + : { + start: this.#cursor, + exclusiveStart: true, + end: this.#end, + inclusiveEnd: this.#inclusiveEnd, + }; + + const remainingEstimate = this.#store.estimateCount(remainingRange); + let remaining = remainingEstimate.count; + let calibrationConfidence = 1; + if (this.#traversed >= CALIBRATION_MIN_TRAVERSED) { + const traversedEstimate = this.#store.estimateCount(traversedRange); + if ( + traversedEstimate.count >= CALIBRATION_MIN_TRAVERSED && + traversedEstimate.confidence >= CALIBRATION_MIN_CONFIDENCE + ) { + const calibration = Math.min( + CALIBRATION_MAX, + Math.max(1 / CALIBRATION_MAX, this.#traversed / traversedEstimate.count) + ); + remaining *= calibration; + calibrationConfidence = Math.min(calibration, 1 / calibration); + } + } + + const count = Math.round(this.#traversed + remaining); + // Cap below 1: only finish() (or an exact-by-construction range) may + // claim exactness, even when rounding makes the remainder vanish. + const confidence = + count > 0 + ? Math.min( + 0.999, + (this.#traversed + calibrationConfidence * remainingEstimate.confidence * remaining) / + count + ) + : Math.min(0.999, calibrationConfidence * remainingEstimate.confidence); + this.#memoized = { count, confidence }; + return { ...this.#memoized }; + } +} diff --git a/src/database.ts b/src/database.ts index 61a218e6d..e5a536560 100644 --- a/src/database.ts +++ b/src/database.ts @@ -1,6 +1,12 @@ import type { BackupStreamOptions } from './backup-stream.ts'; import type { BackupOptions } from './backup.ts'; -import { DBI, type DBITransactional } from './dbi.ts'; +import { CountEstimator, type CountEstimatorOptions } from './count-estimator.ts'; +import { + DBI, + type CountEstimate, + type CountEstimateOptions, + type DBITransactional, +} from './dbi.ts'; import type { BufferWithDataView, Encoder, EncoderFunction, Key } from './encoding.ts'; import { addGlobalListener, @@ -455,7 +461,9 @@ export class RocksDatabase extends DBI { } /** - * Retrieves the estimated number of keys in the database. + * Retrieves the estimated number of keys in the database. This is an alias + * for `db.estimateCount().count`; use `estimateCount()` for range support + * and a confidence indicator. * * @example * ```typescript @@ -468,6 +476,60 @@ export class RocksDatabase extends DBI { return this.getDBIntProperty('rocksdb.estimate-num-keys') ?? 0; } + /** + * Estimates the number of keys in the database, or within a key range, + * returning `{ count, confidence }`. Unlike `getKeysCount()`, this never + * iterates: the estimate is derived from RocksDB statistics (memtable + * stats plus approximate SST sizes converted through the entry density of + * the SSTs overlapping the range), so its cost scales with the number of + * SSTs overlapping the range rather than the number of keys — though + * reading table properties for cold files can do I/O through the table + * cache, and a start-only range does the work of its complement below + * `start`. Accuracy improves with range size. Resolution is bounded by SST + * data-block granularity, so tiny ranges may over-report or report zero for + * present keys; low `confidence` is the signal. Recently deleted or + * overwritten entries may be counted until compaction. An inverted range + * (`start` >= `end`) returns `{ count: 0, confidence: 1 }`. + * + * `confidence` is a heuristic 0–1 trust indicator (1 only when exact) — + * see `CountEstimate`. Estimates always reflect committed state; writes + * pending in a transaction are not included. + * + * @example + * ```typescript + * const db = RocksDatabase.open('/path/to/database'); + * const { count, confidence } = db.estimateCount({ start: 'a', end: 'z' }); + * ``` + */ + estimateCount(options?: CountEstimateOptions): CountEstimate { + return this.store.estimateCount(options); + } + + /** + * Creates a `CountEstimator` for progressively refining a range count + * estimate while iterating the range: report progress with + * `advance(lastKey, count)` (e.g. once per page of results) and + * `estimate()` returns the exact traversed count plus a calibrated + * estimate of the remainder, converging toward the exact total. Call + * `finish()` when traversal completes to make `estimate()` exact. + * + * @example + * ```typescript + * const estimator = db.createCountEstimator({ start: 'a', end: 'z' }); + * let lastKey; + * let pageSize = 0; + * for (const { key } of db.getRange({ start: 'a', end: 'z', limit: 25 })) { + * lastKey = key; + * pageSize++; + * } + * estimator.advance(lastKey, pageSize); + * const { count, confidence } = estimator.estimate(); + * ``` + */ + createCountEstimator(options?: CountEstimatorOptions): CountEstimator { + return new CountEstimator(this.store, options); + } + /** * Returns the current timestamp as a monotonically increasing timestamp in * milliseconds represented as a decimal number. diff --git a/src/dbi.ts b/src/dbi.ts index dd72742fe..ad9ffaaf8 100644 --- a/src/dbi.ts +++ b/src/dbi.ts @@ -62,6 +62,23 @@ export interface RocksDBOptions { tailing?: boolean; } +export interface CountEstimate { + /** + * The estimated number of keys. + */ + count: number; + + /** + * Heuristic 0–1 indicator of how trustworthy `count` is; exactly 1 only + * when the count is exact. Derived from the estimate's resolution + * (SST data-block / memtable sampling granularity relative to the count), + * the tombstone fraction of the overlapping SSTs, and — for open-ended + * starts — the error compounded by complement subtraction. A heuristic + * ordering signal, not a statistical bound. + */ + confidence: number; +} + export interface RangeOptions extends RocksDBOptions { /** * The range end key, otherwise known as the "upper bound". Defaults to @@ -88,6 +105,14 @@ export interface RangeOptions extends RocksDBOptions { start?: Key | Uint8Array; } +export interface CountEstimateOptions extends RangeOptions { + /** + * Interpret `start` as the upper bound and `end` as the lower bound, as + * `getRange()` does for reverse iteration. Defaults to `false`. + */ + reverse?: boolean; +} + export interface IteratorOptions extends RangeOptions { // decoder?: (value: any) => any, diff --git a/src/index.ts b/src/index.ts index 598e792f8..8c6692e22 100644 --- a/src/index.ts +++ b/src/index.ts @@ -14,8 +14,9 @@ export { type RocksDBStat, type RocksDBStats, } from './database.ts'; +export { CountEstimator, type CountEstimatorOptions } from './count-estimator.ts'; export { DBIterator } from './dbi-iterator.ts'; -export { DBI, type IteratorOptions } from './dbi.ts'; +export { DBI, type CountEstimate, type CountEstimateOptions, type IteratorOptions } from './dbi.ts'; export type { Key } from './encoding.ts'; export type * from './stats.ts'; export { diff --git a/src/load-binding.ts b/src/load-binding.ts index f3e61f0ca..7a14200a3 100644 --- a/src/load-binding.ts +++ b/src/load-binding.ts @@ -333,6 +333,7 @@ export type NativeDatabase = { txnId?: number, expectedVersion?: number ): number; + estimateCount(startKey?: Buffer, endKey?: Buffer): { count: number; confidence: number }; getCompression(): { algorithm: string; level?: number }; getCount(options?: RangeOptions, txnId?: number): number; getDBIntProperty(propertyName: string): number | undefined; diff --git a/src/store.ts b/src/store.ts index c578f4951..24a2ce415 100644 --- a/src/store.ts +++ b/src/store.ts @@ -1,7 +1,13 @@ import { type BackupStreamOptions, backupToStream } from './backup-stream.ts'; import { assertBackupDirOutsideDatabase, type BackupOptions } from './backup.ts'; import { DBIterator, type DBIteratorValue } from './dbi-iterator.ts'; -import type { DBITransactional, IteratorOptions, RangeOptions } from './dbi.ts'; +import type { + CountEstimate, + CountEstimateOptions, + DBITransactional, + IteratorOptions, + RangeOptions, +} from './dbi.ts'; import { type BufferWithDataView, createFixedBuffer, @@ -846,6 +852,36 @@ export class Store { return result; } + /** + * Estimates the number of keys in a range from RocksDB statistics + * (memtable stats + approximate SST sizes with range-local entry density) + * without iterating. Estimates always reflect committed state, so there is + * no transactional variant. + */ + estimateCount(options?: CountEstimateOptions): CountEstimate { + let startBuffer: Buffer | undefined; + let endBuffer: Buffer | undefined; + const reverse = options?.reverse ?? false; + const start = reverse ? options?.end : options?.start; + const end = reverse ? options?.start : options?.end; + const exclusiveStart = options?.exclusiveStart ?? reverse; + const inclusiveEnd = options?.inclusiveEnd ?? reverse; + + if (start !== undefined) { + const encodedStart = this.encodeKey(start); + // A zero byte appended to a key is its bytewise successor, turning an + // inclusive bound into an exclusive one (and vice versa for the end). + startBuffer = copyEncodedKey(encodedStart, exclusiveStart); + } + + if (end !== undefined) { + const encodedEnd = this.encodeKey(end); + endBuffer = copyEncodedKey(encodedEnd, inclusiveEnd); + } + + return this.db.estimateCount(startBuffer, endBuffer); + } + getCount(context: StoreContext, options?: StoreRangeOptions): number { options = { ...options }; @@ -1235,6 +1271,21 @@ export class Store { } } +/** + * Copies an encoded key out of the shared key buffer (which the next + * `encodeKey` call would clobber), optionally appending a zero byte to + * produce the key's bytewise successor. + */ +function copyEncodedKey(encoded: BufferWithDataView, appendSuccessorByte: boolean): Buffer { + const length = encoded.end - encoded.start; + const copy = Buffer.allocUnsafe(length + (appendSuccessorByte ? 1 : 0)); + copy.set(encoded.subarray(encoded.start, encoded.end)); + if (appendSuccessorByte) { + copy[length] = 0; + } + return copy; +} + /** * Ensure that they key has been copied into our shared buffer, and return the ending position * @param keyBuffer diff --git a/test/estimate-count.test.ts b/test/estimate-count.test.ts new file mode 100644 index 000000000..6b331bf97 --- /dev/null +++ b/test/estimate-count.test.ts @@ -0,0 +1,335 @@ +import { dbRunner } from './lib/util.ts'; +import { describe, expect, it, vi } from 'vitest'; + +/** + * Estimates are statistical (block-granular SST approximation + memtable + * skip-list approximation), so assertions use a tolerance factor rather than + * exact bounds. Uniform fixed-size entries keep the real accuracy well inside + * these bounds; the factor only guards against wild regressions. + */ +function expectWithin(estimate: number, exact: number, factor: number) { + expect(estimate).toBeGreaterThanOrEqual(exact / factor); + expect(estimate).toBeLessThanOrEqual(exact * factor); +} + +function expectConfidence(confidence: number, min = 0, max = 1) { + expect(confidence).toBeGreaterThanOrEqual(min); + expect(confidence).toBeLessThanOrEqual(max); +} + +const KEY = (i: number) => `key-${String(i).padStart(6, '0')}`; + +describe('estimateCount', () => { + it('should estimate counts for ranges over flushed data', () => + dbRunner(async ({ db }) => { + const N = 20000; + for (let i = 0; i < N; i++) { + await db.put(KEY(i), `value-${i}-${'x'.repeat(50)}`); + } + await db.flush(); + + // full range, both open-ended and bounded + expectWithin(db.getEstimatedKeyCount(), N, 2); + const full = db.estimateCount({ start: KEY(0), end: KEY(N) }); + expectWithin(full.count, N, 2); + // a large uniform range should be high-confidence + expectConfidence(full.confidence, 0.5, 1); + + const singleKey = db.estimateCount({ + start: KEY(N / 2), + end: KEY(N / 2), + inclusiveEnd: true, + }); + expect(singleKey.count === 1 || singleKey.confidence <= 0.5).toBe(true); + + // half range [25%, 75%) + const half = db.estimateCount({ start: KEY(N / 4), end: KEY((3 * N) / 4) }); + expectWithin(half.count, N / 2, 2); + + // open-ended: start only and end only + expectWithin(db.estimateCount({ start: KEY(N / 2) }).count, N / 2, 2); + expectWithin(db.estimateCount({ end: KEY(N / 2) }).count, N / 2, 2); + + // a range past all data should estimate near zero relative to N + const empty = db.estimateCount({ start: 'z', end: 'zz' }); + expect(empty.count).toBeLessThan(N / 20); + expectConfidence(empty.confidence); + })); + + it('should estimate counts for data still in the memtable', () => + dbRunner(async ({ db }) => { + const N = 10000; + for (let i = 0; i < N; i++) { + await db.put(KEY(i), `value-${i}`); + } + // no flush: everything is in the memtable + expectWithin(db.estimateCount({ start: KEY(0), end: KEY(N) }).count, N, 2); + expectWithin(db.estimateCount({ start: KEY(N / 4), end: KEY((3 * N) / 4) }).count, N / 2, 2); + })); + + it('should estimate counts spanning memtable and SST data', () => + dbRunner(async ({ db }) => { + const N = 10000; + for (let i = 0; i < N; i++) { + await db.put(KEY(i), `value-${i}`); + } + await db.flush(); + for (let i = N; i < 2 * N; i++) { + await db.put(KEY(i), `value-${i}`); + } + expectWithin(db.estimateCount({ start: KEY(0), end: KEY(2 * N) }).count, 2 * N, 2); + })); + + it('should return a confident 0 for an empty database', () => + dbRunner(async ({ db }) => { + expect(db.getEstimatedKeyCount()).toBe(0); + // even a zero from estimate-num-keys is an estimate, not exact + expect(db.estimateCount()).toEqual({ count: 0, confidence: 0.95 }); + const range = db.estimateCount({ start: 'a', end: 'z' }); + expect(range.count).toBe(0); + expectConfidence(range.confidence, 0.9, 0.99); + })); + + it('should return an exact 0 for an inverted range', () => + dbRunner(async ({ db }) => { + for (let i = 0; i < 5000; i++) { + await db.put(KEY(i), `value-${i}`); + } + await db.flush(); + // inverted/empty by construction: exact, so confidence is 1 + expect(db.estimateCount({ start: KEY(4000), end: KEY(1000) })).toEqual({ + count: 0, + confidence: 1, + }); + expect(db.estimateCount({ start: KEY(1000), end: KEY(1000) })).toEqual({ + count: 0, + confidence: 1, + }); + })); + + it('should treat zero-length native bounds safely', () => + dbRunner(async ({ db }) => { + for (let i = 0; i < 5000; i++) { + await db.put(KEY(i), `value-${i}`); + } + await db.flush(); + // encodeKey rejects zero-length keys, so exercise the native surface + // directly: an empty end bound sorts below every key (empty range), + // an empty start bound is the minimum key (no-op lower bound) + const native = (db as any).store.db; + expect(native.estimateCount(undefined, Buffer.alloc(0)).count).toBe(0); + expect(native.estimateCount(Buffer.alloc(0), undefined).count).toBe( + db.getEstimatedKeyCount() + ); + })); + + it('should not count uncommitted transaction writes', () => + dbRunner(async ({ db }) => { + const N = 10000; + for (let i = 0; i < N; i++) { + await db.put(KEY(i), `value-${i}`); + } + await db.transaction(async (txn) => { + for (let i = N; i < 2 * N; i++) { + txn.putSync(KEY(i), `value-${i}`); + } + const estimate = db.estimateCount({ start: KEY(0), end: KEY(2 * N) }); + expect(estimate.count).toBeLessThan(N * 1.5); + }); + })); + + it('should scale estimates with range width', () => + dbRunner(async ({ db }) => { + const N = 20000; + for (let i = 0; i < N; i++) { + await db.put(KEY(i), `value-${i}-${'x'.repeat(30)}`); + } + await db.flush(); + + const tenth = db.estimateCount({ start: KEY(0), end: KEY(N / 10) }).count; + const half = db.estimateCount({ start: KEY(0), end: KEY(N / 2) }).count; + const full = db.estimateCount({ start: KEY(0), end: KEY(N) }).count; + expect(tenth).toBeLessThan(half); + expect(half).toBeLessThan(full); + })); + + it('should use range-local density for varied value sizes', () => + dbRunner(async ({ db }) => { + const N = 3000; + for (let i = 0; i < N; i++) { + await db.put(`small-${KEY(i)}`, 'x'.repeat(20)); + } + await db.flush(); + for (let i = 0; i < N; i++) { + await db.put(`large-${KEY(i)}`, 'x'.repeat(1000)); + } + await db.flush(); + + expectWithin(db.estimateCount({ start: 'small-', end: 'small.' }).count, N, 2); + expectWithin(db.estimateCount({ start: 'large-', end: 'large.' }).count, N, 2); + })); +}); + +describe('CountEstimator', () => { + it('should refine the estimate as iteration advances', () => + dbRunner(async ({ db }) => { + const N = 20000; + for (let i = 0; i < N; i++) { + await db.put(KEY(i), `value-${i}-${'x'.repeat(50)}`); + } + await db.flush(); + + const range = { start: KEY(0), end: KEY(N) }; + const estimator = db.createCountEstimator(range); + + // before any traversal: the pure statistical estimate + const initial = estimator.estimate(); + expectWithin(initial.count, N, 2); + expectConfidence(initial.confidence); + + // walk the first half, checkpointing once per "page" + let traversed = 0; + let lastKey: unknown; + for (const key of db.getKeys(range)) { + lastKey = key; + if (++traversed % 1000 === 0) { + estimator.advance(lastKey as string, 1000); + } + if (traversed >= N / 2) { + break; + } + } + expect(estimator.traversed).toBe(N / 2); + + // With half the range traversed exactly, the estimate must be at + // least the traversed count, within a tighter overall bound, and at + // least as confident as the exact fraction of the estimate. + const refined = estimator.estimate(); + expect(refined.count).toBeGreaterThanOrEqual(N / 2); + expectWithin(refined.count, N, 1.6); + expectConfidence(refined.confidence, estimator.traversed / refined.count, 1); + })); + + it('should support reverse iteration', () => + dbRunner(async ({ db }) => { + const N = 10000; + for (let i = 0; i < N; i++) { + await db.put(KEY(i), `value-${i}`); + } + await db.flush(); + + const range = { start: KEY(N), end: KEY(0), reverse: true }; + const estimator = db.createCountEstimator(range); + expectWithin(db.estimateCount(range).count, N - 1, 2); + let upperBound = range.start; + let inclusiveEnd = true; + for (;;) { + const page = Array.from( + db.getKeys({ ...range, start: upperBound, inclusiveEnd, limit: 500 }) + ); + if (page.length === 0) { + break; + } + upperBound = page[page.length - 1] as string; + inclusiveEnd = false; + estimator.advance(upperBound, page.length); + const checkpoint = estimator.estimate(); + expect(checkpoint.count).toBeGreaterThanOrEqual(estimator.traversed); + expectWithin(checkpoint.count, N - 1, 2); + } + expect(estimator.traversed).toBe(N - 1); + estimator.finish(); + expect(estimator.estimate()).toEqual({ count: N - 1, confidence: 1 }); + })); + + it('should memoize before traversal and ignore an empty page', () => + dbRunner(async ({ db }) => { + for (let i = 0; i < 1000; i++) { + await db.put(KEY(i), `value-${i}`); + } + await db.flush(); + + const estimateCount = vi.spyOn(db.store, 'estimateCount'); + const estimator = db.createCountEstimator({ start: KEY(0), end: KEY(1000) }); + const initial = estimator.estimate(); + const expected = { ...initial }; + expect(estimator.estimate()).toEqual(initial); + expect(estimateCount).toHaveBeenCalledTimes(1); + initial.count = 0; + expect(estimator.estimate()).toEqual(expected); + + expect(() => estimator.advance(undefined, 1)).toThrow( + 'CountEstimator.advance requires lastKey when count is nonzero' + ); + expect(estimator.traversed).toBe(0); + expect(estimator.estimate()).toEqual(expected); + expect(estimateCount).toHaveBeenCalledTimes(1); + })); + + it('should not calibrate from a block-granular partial page', () => + dbRunner(async ({ db }) => { + const N = 20000; + for (let i = 0; i < N; i++) { + await db.put(KEY(i), `value-${i}`); + } + await db.flush(); + + const estimator = db.createCountEstimator({ start: KEY(0), end: KEY(N) }); + estimator.advance(KEY(24), 25); + expectWithin(estimator.estimate().count, N, 2); + })); + + it('should converge to near-exact as traversal completes', () => + dbRunner(async ({ db }) => { + const N = 5000; + for (let i = 0; i < N; i++) { + await db.put(KEY(i), `value-${i}`); + } + await db.flush(); + + const estimator = db.createCountEstimator({ start: KEY(0), end: KEY(N) }); + estimator.advance(KEY(N - 1), N); + // remainder is (KEY(N-1), KEY(N)) — essentially empty, though the + // estimate of it is block-granular, so allow a small overshoot + const final = estimator.estimate(); + expect(final.count).toBeGreaterThanOrEqual(N); + expect(final.count).toBeLessThan(N * 1.25); + expectConfidence(final.confidence, 0.7, 1); + + estimator.finish(); + expect(estimator.estimate()).toEqual({ count: N, confidence: 1 }); + })); + + it('should report the exact total for a paginated loop driven to completion', () => + dbRunner(async ({ db }) => { + const N = 5000; + const PAGE = 250; + for (let i = 0; i < N; i++) { + await db.put(KEY(i), `value-${i}`); + } + await db.flush(); + + const range = { start: KEY(0), end: KEY(N) }; + const estimator = db.createCountEstimator(range); + let pageStart: string | undefined; + let exclusiveStart = false; + for (;;) { + const page = Array.from( + db.getKeys({ ...range, start: pageStart ?? range.start, exclusiveStart, limit: PAGE }) + ); + if (page.length === 0) { + break; + } + pageStart = page[page.length - 1] as string; + exclusiveStart = true; + estimator.advance(pageStart, page.length); + const checkpoint = estimator.estimate(); + expect(checkpoint.count).toBeGreaterThanOrEqual(estimator.traversed); + // only finish() may claim exactness + expect(checkpoint.confidence).toBeLessThan(1); + } + expect(estimator.traversed).toBe(N); + estimator.finish(); + expect(estimator.estimate()).toEqual({ count: N, confidence: 1 }); + })); +});