surrealdb-core 3.1.2

//! KNN scan operator for ANN index-backed vector search.
//!
//! This operator performs approximate nearest-neighbor search using an ANN
//! index. It retrieves the top-K records closest to a query vector, ordered by
//! distance (nearest first).

use std::collections::HashSet;
use std::sync::Arc;

use reblessive::TreeStack;

use super::common::fetch_and_filter_records_batch;
use super::pipeline::{ScanPipeline, build_field_state};
use super::resolved::ResolvedTableContext;
use crate::catalog::Index;
use crate::err::Error;
use crate::exec::index::access_path::IndexRef;
use crate::exec::permission::{
	PhysicalPermission, convert_permission_to_physical_runtime, should_check_perms,
	validate_record_user_access,
};
use crate::exec::{
	AccessMode, CardinalityHint, ContextLevel, ExecOperator, ExecutionContext, FlowResult,
	OperatorMetrics, PhysicalExpr, ValueBatch, ValueBatchStream, monitor_stream,
};
use crate::expr::{Cond, ControlFlow, ControlFlowExt};
use crate::iam::Action;
use crate::kvs::CachePolicy;
use crate::val::Number;

/// KNN scan operator using an ANN index.
///
/// Executes an approximate nearest-neighbor search against an ANN index
/// and returns the top-K matching records ordered by distance.
#[derive(Debug)]
pub struct KnnScan {
	/// Reference to the ANN index definition
	pub index_ref: IndexRef,
	/// The query vector to search for nearest neighbors of
	pub vector: Vec<Number>,
	/// Number of nearest neighbors to return
	pub k: u32,
	/// ANN search expansion factor
	pub ef: u32,
	/// Table name for record fetching
	pub table_name: crate::val::TableName,
	/// Optional VERSION timestamp for time-travel queries.
	pub(crate) version: Option<Arc<dyn PhysicalExpr>>,
	/// Plan-time resolved table context. When present, `execute()` skips
	/// runtime table def + permission lookup.
	pub(crate) resolved: Option<ResolvedTableContext>,
	/// Per-operator runtime metrics for EXPLAIN ANALYZE.
	pub(crate) metrics: Arc<OperatorMetrics>,
	/// KNN distance context, shared with IndexFunctionExec for vector::distance::knn().
	pub(crate) knn_context: Option<Arc<crate::exec::function::KnnContext>>,
	/// Residual WHERE condition (non-KNN predicates) to push down into ANN
	/// search. When present, the ANN search will only consider candidates that
	/// satisfy this condition, preventing non-matching rows from consuming
	/// top-K slots.
	///
	/// SECURITY: the cond is evaluated against raw stored records inside the
	/// ANN search (`idx/trees/{hnsw,diskann}/filter.rs::is_record_truthy`),
	/// before any SELECT pipeline filtering. The permission gate that keeps
	/// this safe lives at that chokepoint, which applies the table's SELECT
	/// permission to each candidate BEFORE invoking the cond — so hidden
	/// rows are skipped pre-cond and a record user cannot use the cond to
	/// probe their field values. Preserve that ordering when touching
	/// `is_record_truthy`; see the SECURITY note there for the threat model.
	pub(crate) residual_cond: Option<Cond>,
	/// Projection-aware field set for computed-field materialization.
	/// Outer `None` = sub-operator mode (parent handles fields).
	/// `Some(None)` = all fields, `Some(Some(set))` = specific fields.
	pub(crate) needed_fields: Option<Option<HashSet<String>>>,
}

impl KnnScan {
	#[allow(clippy::too_many_arguments)]
	pub(crate) fn new(
		index_ref: IndexRef,
		vector: Vec<Number>,
		k: u32,
		ef: u32,
		table_name: crate::val::TableName,
		version: Option<Arc<dyn PhysicalExpr>>,
		knn_context: Option<Arc<crate::exec::function::KnnContext>>,
		residual_cond: Option<Cond>,
		needed_fields: Option<Option<HashSet<String>>>,
	) -> Self {
		Self {
			index_ref,
			vector,
			k,
			ef,
			table_name,
			version,
			resolved: None,
			metrics: Arc::new(OperatorMetrics::new()),
			knn_context,
			residual_cond,
			needed_fields,
		}
	}

	/// Set the plan-time resolved table context.
	pub(crate) fn with_resolved(mut self, resolved: ResolvedTableContext) -> Self {
		self.resolved = Some(resolved);
		self
	}
}
impl ExecOperator for KnnScan {
	fn name(&self) -> &'static str {
		"KnnScan"
	}

	fn attrs(&self) -> Vec<(String, String)> {
		vec![
			("index".to_string(), self.index_ref.name.to_string()),
			("k".to_string(), self.k.to_string()),
			("ef".to_string(), self.ef.to_string()),
			("dimension".to_string(), self.vector.len().to_string()),
		]
	}

	fn required_context(&self) -> ContextLevel {
		ContextLevel::Database
	}

	fn access_mode(&self) -> AccessMode {
		AccessMode::ReadOnly
	}

	fn cardinality_hint(&self) -> CardinalityHint {
		CardinalityHint::Bounded(self.k as usize)
	}

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

	fn execute(&self, ctx: &ExecutionContext) -> FlowResult<ValueBatchStream> {
		let db_ctx = ctx.database()?.clone();

		// Validate record user has access to this namespace/database
		validate_record_user_access(&db_ctx)?;

		// Check if we need to enforce permissions
		let check_perms = should_check_perms(&db_ctx, Action::View)?;

		// Clone for the async block
		let index_ref = self.index_ref.clone();
		let vector = self.vector.clone();
		let k = self.k;
		let ef = self.ef;
		let table_name = self.table_name.clone();
		let version_expr = self.version.clone();
		let knn_context = self.knn_context.clone();
		let residual_cond = self.residual_cond.clone();
		let resolved = self.resolved.clone();
		let needed_fields = self.needed_fields.clone();
		let ctx = ctx.clone();

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

			// Evaluate VERSION expression
			let version: Option<u64> = match &version_expr {
				Some(expr) => {
					let eval_ctx = crate::exec::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(),
			};

			// Get the FrozenContext from the root context
			let root = ctx.root();
			let frozen_ctx = &root.ctx;

			// Resolve table permissions and table_id: plan-time fast path or runtime fallback
			let (select_permission, table_id) = if let Some(ref res) = resolved {
				let perm = res.select_permission(check_perms);
				(perm, res.table_def.table_id)
			} else {
				let table_def = db_ctx
					.get_table_def(&table_name, version)
					.await
					.context("Failed to get table")?;

				let table_def = match table_def {
					Some(def) => def,
					None => {
						Err(ControlFlow::Err(anyhow::Error::new(Error::TbNotFound {
							name: table_name.clone(),
						})))?;
						unreachable!()
					}
				};

				let select_permission = if check_perms {
					convert_permission_to_physical_runtime(
						&table_def.permissions.select,
						ctx.ctx(),
					)
					.await
					.context("Failed to convert permission")?
				} else {
					PhysicalPermission::Allow
				};
				(select_permission, table_def.table_id)
			};

			// Early exit if denied
			if matches!(select_permission, PhysicalPermission::Deny) {
				return;
			}

			// Resolve field state for computed fields and field-level
			// permissions. When needed_fields is None (sub-operator mode),
			// the parent operator handles field processing.
			let field_state = match &needed_fields {
				Some(nf) => {
					if let Some(ref res) = resolved {
						res.field_state_for_projection(nf.as_ref())
					} else {
						build_field_state(
							&ctx, &table_name, check_perms, nf.as_ref(),
						).await?
					}
				}
				None => super::pipeline::FieldState::empty(),
			};

			// Get the ANN parameters from the index definition
			let index_def = index_ref.definition();
			let knn_results = match &index_def.index {
				Index::Hnsw(hnsw_params) => {
					// Obtain the shared HNSW index
					let hnsw_index = frozen_ctx
						.get_index_stores()
						.get_index_hnsw(
							ns.namespace_id,
							db.database_id,
							frozen_ctx,
							table_id,
							index_def,
							hnsw_params,
						)
						.await
						.context("Failed to get HNSW index")?;

					// Ensure the HNSW index state is current
					hnsw_index
						.check_state(frozen_ctx)
						.await
						.context("Failed to check HNSW index state")?;

					let cond_filter = match (residual_cond.clone(), ctx.options()) {
						(Some(cond), Some(opt)) => Some((opt, Arc::new(cond))),
						_ => None,
					};

					let mut stack = TreeStack::new();
					stack
						.enter(|stk| {
							let hnsw_index = &hnsw_index;
							let vector = &vector;
							async move {
								hnsw_index
									.knn_search(
										frozen_ctx,
										stk,
										vector,
										k as usize,
										ef as usize,
										cond_filter,
									)
									.await
							}
						})
						.finish()
						.await
						.context("HNSW KNN search failed")?
				}
				#[cfg(not(target_family = "wasm"))]
				Index::DiskAnn(diskann_params) => {
					let diskann_index = frozen_ctx
						.get_index_stores()
						.get_index_diskann(
							ns.namespace_id,
							db.database_id,
							table_id,
							index_def,
							diskann_params,
						)
						.await
						.context("Failed to get DiskANN index")?;

					diskann_index.check_state().await.context("Failed to check DiskANN index state")?;

					let cond_filter = match (residual_cond.clone(), ctx.options()) {
						(Some(cond), Some(opt)) => Some((opt, Arc::new(cond))),
						_ => None,
					};

					let mut stack = TreeStack::new();
					stack
						.enter(|stk| {
							let diskann_index = &diskann_index;
							let vector = &vector;
							async move {
								diskann_index
									.knn_search(
										frozen_ctx,
										stk,
										vector,
										k as usize,
										ef as usize,
										cond_filter,
									)
									.await
							}
						})
						.finish()
						.await
						.context("DiskANN KNN search failed")?
				}
				_ => {
					Err(ControlFlow::Err(anyhow::anyhow!(
						"Index '{}' is not an ANN index",
						index_def.name
					)))?;
					unreachable!()
				}
			};

			let mut rids = Vec::with_capacity(knn_results.len());
			// Populate KNN distance context (if present) before yielding records.
			// This makes distances available to vector::distance::knn() during
			// downstream projection evaluation.
			if let Some(ref knn_ctx) = knn_context {
				for (rid, distance, _) in &knn_results {
					knn_ctx.insert(rid.as_ref().clone(), Number::Float(*distance)).await;
					rids.push(rid.as_ref().clone());
				}
			} else {
				for (rid, _, _) in &knn_results {
					rids.push(rid.as_ref().clone());
				}
			}

			// Table-level permissions are handled by fetch_and_filter_records_batch.
			// The pipeline handles computed fields and field-level permissions.
			let mut pipeline = ScanPipeline::new(
				PhysicalPermission::Allow,
				None,
				field_state,
				check_perms,
				None,
				0,
			);

			// Batch-fetch all records and apply permission filtering
			let mut values = fetch_and_filter_records_batch(
				&ctx,
				&txn,
				ns.namespace_id,
				db.database_id,
				&rids,
				&select_permission,
				check_perms,
				version,
				CachePolicy::ReadWrite,
			).await?;

			pipeline.process_batch(&mut values, &ctx).await?;

			if !values.is_empty() {
				yield ValueBatch { values };
			}
		};

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