Bind stable vector handles during planning

Author: anoop-narang

src/planner.rs

@@ -140,18 +140,8 @@ impl ExtensionPlanner for USearchExecPlanner {
             None => return Ok(None),
         };
 
-        // Async: ensure the index is loaded (downloads on cache miss).
-        self.registry.ensure_loaded(&node.table_name).await?;
-
-        let registered = match self.registry.resolve(&node.table_name) {
-            Some(r) => r,
-            None => {
-                return Err(DataFusionError::Execution(format!(
-                    "USearchExecPlanner: table '{}' not available after ensure_loaded",
-                    node.table_name
-                )));
-            }
-        };
+        // Async: bind a query-stable registry entry during planning.
+        let registered = self.registry.prepare(&node.table_name).await?;
 
         let exec_props = session_state.execution_props();
 
@@ -223,17 +213,22 @@ impl ExtensionPlanner for USearchExecPlanner {
             None
         };
 
+        let schema = registered.schema.clone();
+        let key_col = registered.key_col.clone();
+        let scalar_kind = registered.scalar_kind;
+        let brute_force_threshold = registered.config.brute_force_selectivity_threshold;
+
         Ok(Some(Arc::new(USearchExec::new(SearchParams {
             table_name: node.table_name.clone(),
-            registry: self.registry.clone(),
+            registered,
             query_vec: node.query_vec_f64(),
             k: node.k,
             distance_type: node.distance_type.clone(),
             has_filters: !node.filters.is_empty(),
-            schema: registered.schema.clone(),
-            key_col: registered.key_col.clone(),
-            scalar_kind: registered.scalar_kind,
-            brute_force_threshold: registered.config.brute_force_selectivity_threshold,
+            schema,
+            key_col,
+            scalar_kind,
+            brute_force_threshold,
             provider_scan,
         }))))
     }
@@ -242,10 +237,10 @@ impl ExtensionPlanner for USearchExecPlanner {
 // ── Search parameters ─────────────────────────────────────────────────────────
 
 /// All parameters needed to run a USearch query, cloned cheaply into execute().
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 struct SearchParams {
     table_name: String,
-    registry: Arc<dyn VectorIndexResolver>,
+    registered: Arc<crate::registry::RegisteredTable>,
     query_vec: Vec<f64>,
     k: usize,
     distance_type: DistanceType,
@@ -261,6 +256,24 @@ struct SearchParams {
     provider_scan: Option<Arc<dyn ExecutionPlan>>,
 }
 
+impl fmt::Debug for SearchParams {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("SearchParams")
+            .field("table_name", &self.table_name)
+            .field("k", &self.k)
+            .field("has_filters", &self.has_filters)
+            .field("schema", &self.schema)
+            .field("key_col", &self.key_col)
+            .field("scalar_kind", &self.scalar_kind)
+            .field("brute_force_threshold", &self.brute_force_threshold)
+            .field(
+                "provider_scan",
+                &self.provider_scan.as_ref().map(|_| "Some(..)"),
+            )
+            .finish_non_exhaustive()
+    }
+}
+
 // ── Physical execution node ───────────────────────────────────────────────────
 
 /// Leaf execution plan that defers all I/O to execute() time.
@@ -361,16 +374,7 @@ async fn usearch_execute(
     params: SearchParams,
     task_ctx: Arc<TaskContext>,
 ) -> Result<Vec<RecordBatch>> {
-    // Re-fetch at execute time and reload on cache miss so eviction between
-    // plan and execute is handled correctly for async-backed resolvers.
-    params.registry.ensure_loaded(&params.table_name).await?;
-
-    let registered = params.registry.resolve(&params.table_name).ok_or_else(|| {
-        DataFusionError::Execution(format!(
-            "USearchExec: table '{}' not available after ensure_loaded at execute time",
-            params.table_name
-        ))
-    })?;
+    let registered = params.registered.clone();
 
     if !params.has_filters {
         // ── Unfiltered path ───────────────────────────────────────────────

src/registry.rs

@@ -230,30 +230,34 @@ pub struct VectorIndexMeta {
 
 /// Trait for resolving vector index entries by name.
 ///
-/// Split into two operations to accommodate DataFusion's mixed sync/async pipeline:
+/// DataFusion's optimizer is synchronous but physical planning is async, so the
+/// resolver exposes two levels:
 /// - `peek()` — sync, cheap: used by the optimizer rule to decide if ANN rewrite applies
-/// - `resolve()` — sync: returns fully loaded entry from cache (for executor hot path)
-/// - `ensure_loaded()` — async: loads the index on cache miss (for planner/executor)
+/// - `resolve()` — sync, cache-only: used by synchronous callers such as the UDTF
+/// - `prepare()` — async: returns a query-stable loaded entry for planning/execution
 ///
-/// The built-in [`USearchRegistry`] implements all three as direct hashmap lookups
-/// (always cached). Production systems can implement catalog-backed loading in
-/// `ensure_loaded`.
+/// The important contract is that `prepare()` returns the exact loaded
+/// [`RegisteredTable`] the query should execute against. This avoids a second
+/// lookup during execution and removes the planner/execute handoff gap where a
+/// newer generation could replace the registry entry mid-query.
 #[async_trait::async_trait]
 pub trait VectorIndexResolver: Send + Sync + std::fmt::Debug {
     /// Sync, cheap: check if a vector index exists and return its metadata.
     /// Used by the optimizer rule to decide if the ANN rewrite should apply.
     /// Must NOT do I/O — only check local cache or lightweight state.
     fn peek(&self, name: &str) -> Option<VectorIndexMeta>;
 
-    /// Sync: return a fully loaded entry from cache.
-    /// Returns `None` on cache miss — the caller should call `ensure_loaded` first.
+    /// Sync, cache-only: return a fully loaded entry from local cache.
+    /// Returns `None` on cache miss.
     fn resolve(&self, name: &str) -> Option<Arc<RegisteredTable>>;
 
-    /// Async: ensure the index is loaded and available in cache.
-    /// On cache miss or stale entry, downloads files and loads the index.
-    /// On cache hit with current generation, returns immediately.
-    /// Called from the async physical planner before building the execution plan.
-    async fn ensure_loaded(&self, name: &str) -> datafusion::common::Result<()>;
+    /// Async: return a query-stable loaded entry.
+    ///
+    /// On cache miss or stale entry, implementations may download files and
+    /// load the current generation. On success, the returned [`Arc`] must stay
+    /// valid for the lifetime of the query even if a newer generation is loaded
+    /// concurrently afterwards.
+    async fn prepare(&self, name: &str) -> datafusion::common::Result<Arc<RegisteredTable>>;
 }
 
 // ── USearchRegistry ───────────────────────────────────────────────────────────
@@ -438,9 +442,10 @@ impl VectorIndexResolver for USearchRegistry {
         self.get(name)
     }
 
-    async fn ensure_loaded(&self, _name: &str) -> datafusion::common::Result<()> {
-        // USearchRegistry is always fully cached — nothing to load.
-        Ok(())
+    async fn prepare(&self, name: &str) -> datafusion::common::Result<Arc<RegisteredTable>> {
+        self.get(name).ok_or_else(|| {
+            DataFusionError::Execution(format!("USearchRegistry: table '{name}' is not loaded"))
+        })
     }
 }
 

src/udtf.rs

@@ -44,8 +44,8 @@ use crate::registry::VectorIndexResolver;
 ///
 /// This entry point is synchronous. For async-backed [`VectorIndexResolver`]
 /// implementations, it only works when the target index is already loaded in
-/// the local cache. `vector_usearch()` does not call `ensure_loaded()` and
-/// cannot trigger async index loads.
+/// the local cache. `vector_usearch()` does not call `prepare()` and cannot
+/// trigger async index loads.
 pub struct USearchUDTF {
     registry: Arc<dyn VectorIndexResolver>,
 }