surrealdb-core 3.2.3

//! IndexCountScan operator - optimized COUNT() using index count metadata.
//!
//! When a query is `SELECT count() FROM <table> WHERE <cond> GROUP ALL` and a
//! COUNT index exists whose stored condition matches the WHERE clause exactly,
//! this operator replaces the full Scan -> Filter -> Aggregate pipeline.
//!
//! Instead of deserializing and filtering every record, it sums the delta
//! entries stored in `IndexCountKey` for the matching COUNT index.  This is
//! O(index entries) with no record I/O.
//!
//! The planner emits this operator (via `is_indexed_count_eligible`) when:
//! - Fields are count-all-only
//! - GROUP ALL is present
//! - A WHERE clause is present
//! - No SPLIT, ORDER BY, FETCH, or OMIT clauses
//! - A single table source
//!
//! At execution time the operator:
//! 1. Resolves the table and looks up its indexes.
//! 2. Finds a `Index::Count(cond)` whose condition equals the query's WHERE.
//! 3. Scans `IndexCountKey` deltas to compute the total count.
//! 4. Falls back to a full scan + filter + count if no matching COUNT index is found or permissions
//!    are conditional.

use std::sync::Arc;

use tracing::instrument;

use crate::catalog::{DatabaseId, Index, NamespaceId, Permission};
use crate::err::Error;
use crate::exec::index::access_path::{BTreeAccess, IndexRef};
use crate::exec::permission::{
	PhysicalPermission, convert_permission_to_physical_runtime, should_check_perms,
	validate_record_user_access,
};
use crate::exec::{
	AccessMode, CardinalityHint, ContextLevel, EvalContext, ExecOperator, ExecutionContext,
	FlowResult, OperatorMetrics, PhysicalExpr, ValueBatch, ValueBatchStream, monitor_stream,
};
use crate::expr::cond::Cond;
use crate::expr::{ControlFlow, ControlFlowExt};
use crate::iam::Action;
use crate::key::index::iu::IndexCountKey;
use crate::key::record;
use crate::kvs::KVValue;
use crate::val::{Number, Object, TableName, Value};

/// Optimized operator for `SELECT count() FROM <table> WHERE <cond> GROUP ALL`
/// when a matching COUNT index exists.
///
/// Falls back to B-tree index key counting (when a covering B-tree index is
/// available) or full scan + filter + count if no index can service the query.
#[derive(Debug, Clone)]
pub struct IndexCountScan {
	/// Expression that evaluates to the table name.
	pub(crate) source: Arc<dyn PhysicalExpr>,
	/// The physical expression for the WHERE predicate (used for fallback).
	pub(crate) predicate: Arc<dyn PhysicalExpr>,
	/// The AST-level WHERE condition for exact matching against COUNT index
	/// conditions.
	pub(crate) condition: Cond,
	/// Optional VERSION expression for time-travel queries.
	pub(crate) version: Option<Arc<dyn PhysicalExpr>>,
	/// Output field names for the count result (one per SELECT field).
	/// For `SELECT count() as c FROM t WHERE ... GROUP ALL` this would be `["c"]`.
	/// For `SELECT count() FROM t WHERE ... GROUP ALL` this would be `["count"]`.
	pub(crate) field_names: Vec<String>,
	/// Optional B-tree index access path for key-only counting when no
	/// matching COUNT index exists.  The planner resolves this from the
	/// same index analysis it performs for regular queries.
	pub(crate) btree_access: Option<(IndexRef, BTreeAccess)>,
	/// Per-operator runtime metrics for EXPLAIN ANALYZE.
	pub(crate) metrics: Arc<OperatorMetrics>,
}

impl IndexCountScan {
	pub(crate) fn new(
		source: Arc<dyn PhysicalExpr>,
		predicate: Arc<dyn PhysicalExpr>,
		condition: Cond,
		version: Option<Arc<dyn PhysicalExpr>>,
		field_names: Vec<String>,
	) -> Self {
		debug_assert!(!field_names.is_empty(), "IndexCountScan requires at least one field name");
		Self {
			source,
			predicate,
			condition,
			version,
			field_names,
			btree_access: None,
			metrics: Arc::new(OperatorMetrics::new()),
		}
	}

	/// Set the B-tree index access path for key-only counting.
	pub(crate) fn with_btree_access(mut self, access: Option<(IndexRef, BTreeAccess)>) -> Self {
		self.btree_access = access;
		self
	}
}
impl ExecOperator for IndexCountScan {
	fn name(&self) -> &'static str {
		"IndexCountScan"
	}

	fn attrs(&self) -> Vec<(String, String)> {
		vec![
			("source".to_string(), self.source.to_sql()),
			("condition".to_string(), self.predicate.to_sql()),
		]
	}

	fn required_context(&self) -> ContextLevel {
		// IndexCountScan needs database context, combined with expression contexts
		self.source
			.required_context()
			.max(self.predicate.required_context())
			.max(ContextLevel::Database)
	}

	fn metrics(&self) -> Option<&OperatorMetrics> {
		Some(&self.metrics)
	}

	fn expressions(&self) -> Vec<(&str, &Arc<dyn PhysicalExpr>)> {
		vec![("source", &self.source), ("predicate", &self.predicate)]
	}

	fn access_mode(&self) -> AccessMode {
		self.source.access_mode().combine(self.predicate.access_mode())
	}

	fn cardinality_hint(&self) -> CardinalityHint {
		CardinalityHint::AtMostOne
	}

	#[instrument(name = "IndexCountScan::execute", level = "trace", skip_all)]
	fn execute(&self, ctx: &ExecutionContext) -> FlowResult<ValueBatchStream> {
		let db_ctx = ctx.database()?.clone();
		validate_record_user_access(&db_ctx)?;
		let check_perms = should_check_perms(&db_ctx, Action::View)?;

		let source_expr = Arc::clone(&self.source);
		let predicate_expr = Arc::clone(&self.predicate);
		let condition = self.condition.clone();
		let version = self.version.clone();
		let field_names = self.field_names.clone();
		let btree_access = self.btree_access.clone();
		let ctx = ctx.clone();

		let stream = async_stream::try_stream! {
			let db_ctx = ctx.database().context("IndexCountScan requires database context")?;
			let txn = ctx.txn();
			let ns = Arc::clone(&db_ctx.ns_ctx.ns);
			let db = Arc::clone(&db_ctx.db);

			// Evaluate VERSION expression to a timestamp
			let version: Option<u64> = match &version {
				Some(expr) => {
					let eval_ctx = EvalContext::from_exec_ctx(&ctx);
					let v = expr.evaluate(eval_ctx).await?;
					Some(
						v.cast_to::<crate::val::Datetime>()
							.map_err(|e| anyhow::anyhow!("{e}"))?
							.to_version_stamp(txn.timestamp_impl().as_ref())?,
					)
				}
				None => ctx.version_stamp(),
			};

			// Evaluate source expression to get the table name.
			let eval_ctx = EvalContext::from_exec_ctx(&ctx);
			let table_value = source_expr.evaluate(eval_ctx).await?;

			let table_name = match table_value {
				Value::Table(t) => t,
				_ => {
					Err(ControlFlow::Err(anyhow::anyhow!(
						"IndexCountScan received a non-table source"
					)))?;
					unreachable!()
				}
			};

			// Verify table exists.
			let table_def = db_ctx
				.get_table_def(&table_name, version)
				.await
				.context("Failed to get table")?;

			if table_def.is_none() {
				Err(ControlFlow::Err(anyhow::Error::new(Error::TbNotFound {
					name: table_name.clone(),
				})))?;
			}

			// Resolve SELECT permission.
			let select_permission = if check_perms {
				let catalog_perm = match &table_def {
					Some(def) => def.permissions.select.clone(),
					None => Permission::None,
				};
				convert_permission_to_physical_runtime(&catalog_perm, ctx.ctx())
					.await
					.context("Failed to convert permission")?
			} else {
				PhysicalPermission::Allow
			};

			match select_permission {
				PhysicalPermission::Deny => {
					// Table is invisible.
					return;
				}
				PhysicalPermission::Conditional(_) => {
					// Per-record permissions: fall back to full scan + filter + count.
				let count = count_with_filter_fallback(
					&ctx,
					ns.namespace_id,
					db.database_id,
					&table_name,
					version,
					&select_permission,
					&predicate_expr,
				)
				.await?;
				yield make_count_batch(count, &field_names);
				return;
				}
				PhysicalPermission::Allow => {
					// Proceed to look for a matching COUNT index.
				}
			}

			// Look up all indexes for the table (using the execution-level cache).
			let indexes = db_ctx
				.get_table_indexes(&table_name, version)
				.await
				.context("Failed to fetch table indexes")?;

			let matching_index = indexes.iter().find(|ix| {
				if let Index::Count(ref idx_cond) = ix.index {
					// The COUNT index condition must exactly match the WHERE clause.
					idx_cond.as_ref() == Some(&condition)
				} else {
					false
				}
			});

			if let Some(ix_def) = matching_index {
				// Fast path: sum delta counts from the COUNT index.
				let count = sum_index_count_deltas(
					&ctx,
					&txn,
					ns.namespace_id,
					db.database_id,
					&table_name,
					ix_def.index_id,
				)
				.await?;
				yield make_count_batch(count, &field_names);
			} else if let Some((ref ix_ref, ref access)) = btree_access {
				// Medium path: count entries by iterating B-tree index
				// keys only — no record value deserialization.
				let count = count_btree_index_keys(
					&ctx,
					&txn,
					ns.namespace_id,
					db.database_id,
					ix_ref,
					access,
				)
				.await?;
				yield make_count_batch(count, &field_names);
			} else {
				// No matching COUNT index found: fall back to full scan + filter + count.
				let perm = PhysicalPermission::Allow;
				let count = count_with_filter_fallback(
					&ctx,
					ns.namespace_id,
					db.database_id,
					&table_name,
					version,
					&perm,
					&predicate_expr,
				)
				.await?;
				yield make_count_batch(count, &field_names);
			}
		};

		Ok(monitor_stream(Box::pin(stream), "IndexCountScan", &self.metrics))
	}
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Build the single-row batch that the Aggregate operator would normally
/// produce for `SELECT count() … GROUP ALL`.
///
/// Each entry in `field_names` becomes a key in the output object, all
/// mapping to the same count value. For example:
/// - `SELECT count() FROM t WHERE … GROUP ALL`      -> `{ "count": N }`
/// - `SELECT count() AS c FROM t WHERE … GROUP ALL`  -> `{ "c": N }`
/// - `SELECT count() AS a, count() AS b …`           -> `{ "a": N, "b": N }`
fn make_count_batch(count: usize, field_names: &[String]) -> ValueBatch {
	let mut obj = Object::default();
	let count_val = Value::Number(Number::Int(count as i64));
	for name in field_names {
		obj.insert(name.clone(), count_val.clone());
	}
	ValueBatch {
		values: vec![Value::Object(obj)],
	}
}

/// Sum the delta entries in `IndexCountKey` for a given COUNT index.
pub(crate) async fn sum_index_count_deltas(
	ctx: &ExecutionContext,
	txn: &crate::kvs::Transaction,
	ns: NamespaceId,
	db: DatabaseId,
	tb: &TableName,
	ix: crate::catalog::IndexId,
) -> Result<usize, ControlFlow> {
	let range =
		IndexCountKey::range(ns, db, tb, ix).context("Failed to compute index count key range")?;
	let mut cursor = txn
		.open_keys_cursor(range, crate::idx::planner::ScanDirection::Forward, 0, None)
		.await
		.context("Failed to open index-count cursor")?;
	let mut count: i64 = 0;
	loop {
		if ctx.cancellation().is_cancelled() {
			return Err(ControlFlow::Err(anyhow::anyhow!(Error::QueryCancelled)));
		}
		let batch = cursor
			.next_batch(crate::kvs::NORMAL_BATCH_SIZE)
			.await
			.context("Failed to scan index count keys")?;
		if batch.is_empty() {
			break;
		}
		for key in &batch {
			let iu = IndexCountKey::decode_key(key).context("Failed to decode index count key")?;
			if iu.pos {
				count += iu.count as i64;
			} else {
				count -= iu.count as i64;
			}
		}
	}
	Ok(count.max(0) as usize)
}

/// Fallback: scan all records, apply the predicate, and count matches.
///
/// Used when no matching COUNT index exists or when per-record permissions
/// require row-level evaluation.
async fn count_with_filter_fallback(
	ctx: &ExecutionContext,
	ns_id: NamespaceId,
	db_id: DatabaseId,
	table_name: &TableName,
	version: Option<u64>,
	permission: &PhysicalPermission,
	predicate: &Arc<dyn PhysicalExpr>,
) -> Result<usize, ControlFlow> {
	use crate::exec::permission::PhysicalPermission;

	let txn = ctx.txn();
	let beg = record::prefix(ns_id, db_id, table_name)?;
	let end = record::suffix(ns_id, db_id, table_name)?;

	let mut cursor = txn
		.open_vals_cursor(beg..end, crate::idx::planner::ScanDirection::Forward, 0, version)
		.await
		.context("Failed to open scan cursor")?;
	let mut count = 0usize;
	loop {
		if ctx.cancellation().is_cancelled() {
			return Err(ControlFlow::Err(anyhow::anyhow!(Error::QueryCancelled)));
		}
		let batch = cursor
			.next_batch(crate::kvs::NORMAL_BATCH_SIZE)
			.await
			.context("Failed to scan record")?;
		if batch.is_empty() {
			break;
		}
		for (key, val) in &batch {
			let decoded_key = crate::key::record::RecordKey::decode_key(key)
				.context("Failed to decode record key")?;
			let rid_val = crate::val::RecordId {
				table: decoded_key.tb.into_owned(),
				key: decoded_key.id,
			};
			let record = crate::catalog::Record::kv_decode_value(val, rid_val)
				.context("Failed to deserialize record")?;
			let value = record.data;

			// Check per-record permission first.
			let perm_allowed = match permission {
				PhysicalPermission::Allow => true,
				PhysicalPermission::Deny => false,
				PhysicalPermission::Conditional(expr) => {
					let eval_ctx = EvalContext::from_exec_ctx(ctx).with_value(&value);
					expr.evaluate(eval_ctx).await.map(|v| v.is_truthy()).map_err(|e| {
						ControlFlow::Err(anyhow::anyhow!("Failed to check permission: {e}"))
					})?
				}
			};
			if !perm_allowed {
				continue;
			}

			// Apply the WHERE predicate.
			let eval_ctx = EvalContext::from_exec_ctx(ctx).with_value(&value);
			let matches =
				predicate.evaluate(eval_ctx).await.map(|v| v.is_truthy()).map_err(|e| {
					ControlFlow::Err(anyhow::anyhow!("Failed to evaluate predicate: {e}"))
				})?;
			if matches {
				count += 1;
			}
		}
	}

	Ok(count)
}

/// Count matching records by iterating B-tree index keys only.
///
/// This is much faster than the full-scan fallback because it avoids
/// reading and deserializing record values.  Each index entry corresponds
/// to exactly one matching record, so we simply count entries in the
/// appropriate key range.
async fn count_btree_index_keys(
	ctx: &ExecutionContext,
	txn: &crate::kvs::Transaction,
	ns_id: NamespaceId,
	db_id: DatabaseId,
	index_ref: &IndexRef,
	access: &BTreeAccess,
) -> Result<usize, ControlFlow> {
	use crate::exec::index::iterator::btree::{
		CompoundEqualIterator, CompoundRangeIterator, IndexEqualIterator, IndexRangeIterator,
		UniqueEqualIterator, UniqueRangeIterator,
	};
	use crate::idx::planner::ScanDirection;

	let ix = index_ref.definition();
	let is_unique = index_ref.is_unique();
	let mut count = 0usize;

	match (access, is_unique) {
		(BTreeAccess::Equality(value), true) => {
			let mut iter = UniqueEqualIterator::new(ns_id, db_id, ix, value)
				.context("Failed to create unique equal iterator")?;
			loop {
				if ctx.cancellation().is_cancelled() {
					return Err(ControlFlow::Err(anyhow::anyhow!(Error::QueryCancelled)));
				}
				let rids = iter.next_batch(txn).await.context("Failed to iterate index")?;
				if rids.is_empty() {
					break;
				}
				count += rids.len();
			}
		}
		(BTreeAccess::Equality(value), false) => {
			// Non-unique equality: iterate all matching entries.
			let mut iter = IndexEqualIterator::new(ns_id, db_id, ix, value)
				.context("Failed to create index equal iterator")?;
			loop {
				if ctx.cancellation().is_cancelled() {
					return Err(ControlFlow::Err(anyhow::anyhow!(Error::QueryCancelled)));
				}
				let rids = iter.next_batch(txn).await.context("Failed to iterate index")?;
				if rids.is_empty() {
					break;
				}
				count += rids.len();
			}
		}
		(
			BTreeAccess::Range {
				from,
				to,
			},
			true,
		) => {
			let mut iter = UniqueRangeIterator::new(
				ns_id,
				db_id,
				ix,
				from.as_ref(),
				to.as_ref(),
				ScanDirection::Forward,
			)
			.context("Failed to create unique range iterator")?;
			loop {
				if ctx.cancellation().is_cancelled() {
					return Err(ControlFlow::Err(anyhow::anyhow!(Error::QueryCancelled)));
				}
				let rids = iter.next_batch(txn).await.context("Failed to iterate index")?;
				if rids.is_empty() {
					break;
				}
				count += rids.len();
			}
		}
		(
			BTreeAccess::Range {
				from,
				to,
			},
			false,
		) => {
			let mut iter = IndexRangeIterator::new(
				ns_id,
				db_id,
				ix,
				from.as_ref(),
				to.as_ref(),
				ScanDirection::Forward,
			)
			.context("Failed to create index range iterator")?;
			loop {
				if ctx.cancellation().is_cancelled() {
					return Err(ControlFlow::Err(anyhow::anyhow!(Error::QueryCancelled)));
				}
				let rids = iter.next_batch(txn).await.context("Failed to iterate index")?;
				if rids.is_empty() {
					break;
				}
				count += rids.len();
			}
		}
		(
			BTreeAccess::Compound {
				prefix,
				range: Some(range),
			},
			_,
		) => {
			let mut iter =
				CompoundRangeIterator::new(ns_id, db_id, ix, prefix, range, ScanDirection::Forward)
					.context("Failed to create compound range iterator")?;
			loop {
				if ctx.cancellation().is_cancelled() {
					return Err(ControlFlow::Err(anyhow::anyhow!(Error::QueryCancelled)));
				}
				let rids = iter.next_batch(txn, 1000).await.context("Failed to iterate index")?;
				if rids.is_empty() {
					break;
				}
				count += rids.len();
			}
		}
		(
			BTreeAccess::Compound {
				prefix,
				range: None,
			},
			_,
		) => {
			let mut iter =
				CompoundEqualIterator::new(ns_id, db_id, ix, prefix, None, ScanDirection::Forward)
					.context("Failed to create compound equal iterator")?;
			loop {
				if ctx.cancellation().is_cancelled() {
					return Err(ControlFlow::Err(anyhow::anyhow!(Error::QueryCancelled)));
				}
				let rids = iter.next_batch(txn, 1000).await.context("Failed to iterate index")?;
				if rids.is_empty() {
					break;
				}
				count += rids.len();
			}
		}
		// FullText and Knn are not supported for counting.
		_ => {
			return Err(ControlFlow::Err(anyhow::anyhow!(
				"Unsupported BTreeAccess type for index key counting"
			)));
		}
	}

	Ok(count)
}