fix(rule): keep index for k-NN that returns metadata, fall back when vector is projected (#25)

A k-NN query that projects the indexed vector column (SELECT *, or
SELECT id, embedding) crashed with a post-optimizer schema-mismatch when an
index was present: the passthrough branch built its output from the index
node's columns (addressing key + non-vector columns), which can't include the
vector, so the rewritten plan's schema differed from the original and
DataFusion's invariant check aborted the query.

The fix is output-aware. The rule now also anchors on a Projection sitting
directly over a passthrough k-NN Sort and drives the rewrite from that outer
projection's columns — the query's real output:

  - vector NOT in output (e.g. SELECT id ... ORDER BY l2_distance(emb, ...),
    the common "nearest ids" query) -> every output column is producible from
    the node, so the index is still used.
  - vector IN output (SELECT *, SELECT id, embedding) -> the rewrite can't
    produce it, so the rule declines and the query falls back to exact
    brute-force search (correct, like the existing DESC / metric-mismatch
    fallbacks) instead of crashing.

This keeps the metadata-only k-NN path on the index (no regression) while
fixing the crash. A code comment records the rejected alternative (have
USearchExec reconstruct the vector via index.get) and why: it would make the
index a second source of returned vectors that must byte-match the source,
which breaks under F16 quantization.

Regression tests model production (lookup schema excludes the vector column,
which the existing tests' provider included, masking the bug). README documents
the fallback.

Fixes #508

Author: anoop-narang

README.md

@@ -4,7 +4,7 @@ A DataFusion extension that integrates [USearch](https://github.com/unum-cloud/u
 
 Queries matching the `ORDER BY distance_fn(col, query) LIMIT k` pattern are **transparently rewritten** by an optimizer rule into a native USearch index call — no query rewrite needed from the caller. `WHERE` clause filters are handled adaptively: high-selectivity filters use USearch's in-graph predicate API; low-selectivity filters bypass HNSW entirely and scan the data directly.
 
-**DataFusion:** 52.2   **USearch:** 2.24
+**DataFusion:** 53   **USearch:** 2.24
 
 ---
 
@@ -230,20 +230,33 @@ tests/
 
 ### Optimizer rewrite
 
-The rule (`rule.rs`) matches two logical plan shapes:
+The rule (`rule.rs`) matches the `Sort(fetch=k)` over a `TableScan`, with an
+optional `Projection` between them and an optional `Filter` directly above the
+scan:
 
 ```
 Sort(fetch=k, ORDER BY dist ASC)
-  Projection([..., distance_fn(col, lit) AS dist, ...])
-    TableScan(name)
-
-Sort(fetch=k, ORDER BY dist ASC)
-  Projection([..., distance_fn(col, lit) AS dist, ...])
-    Filter(predicate)
+  [ Projection([..., distance_fn(col, lit) AS dist, ...]) ]   ← optional
+    [ Filter(predicate) ]                                     ← optional
       TableScan(name)
 ```
 
-Preconditions: sort is `ASC`, distance UDF matches index metric, table is registered, query vector is a literal. When the rule fires, it replaces the inner nodes with a `USearchNode` leaf carrying: table name, vector column, query vector, k, distance type, and absorbed filter predicates. The `Sort` node is preserved above for final ordering.
+DataFusion omits the `Projection` for `SELECT *` (and for any SELECT whose
+columns come straight from the scan), so the `Sort` can sit directly on the
+`TableScan`.
+
+Preconditions: sort is `ASC`, distance UDF matches index metric, table is
+registered, query vector is a literal. When the rule fires, it replaces the inner
+nodes with a `USearchNode` leaf carrying: table name, vector column, query
+vector, k, distance type, and absorbed filter predicates. The `Sort` node is
+preserved above for final ordering.
+
+**Schema preservation:** an optimizer rule must not change the plan's output
+schema. The `USearchNode` produces only what the `lookup_provider` can fetch by
+key (addressing key + non-vector columns) plus `_distance` — it cannot produce
+the indexed vector column. If the matched `Sort`'s output would include the
+vector column (e.g. `SELECT *`), the rule declines and the query falls back to
+exact execution rather than emitting a schema-incompatible plan.
 
 Physical planning (`planner.rs`) translates `USearchNode` into `USearchExec`, a physical plan node that executes the actual search.
 
@@ -305,6 +318,7 @@ Tests cover optimizer rule matching/rejection, end-to-end execution through both
 
 | Limitation | Notes |
 |---|---|
+| Projecting the indexed vector column | A k-NN query whose output includes the vector column itself (e.g. `SELECT *`, or `SELECT id, vector`) falls back to exact execution. The `lookup_provider` does not store the vector column (see [registration](#register-providers-and-set-up-the-sessioncontext)), so the rewrite cannot reproduce it. Project the metadata columns and the distance instead. |
 | Stacked `Filter` nodes | Only one `Filter -> TableScan` layer is absorbed. `Filter -> Filter -> TableScan` falls back to exact execution. DataFusion typically combines multiple WHERE conditions into a single Filter, so this rarely occurs. |
 | Runtime query vectors | The query vector must be a compile-time literal (`ARRAY[0.1, ...]`). Column references or subquery results are not rewritten. Use `vector_search_vector` for explicit ANN queries. |
 | `ef_search` per-query | `expansion_search` is global to the index instance. Per-query adjustment is not supported. |

src/rule.rs

@@ -48,21 +48,66 @@ impl USearchRule {
     }
 
     fn try_match(&self, plan: &LogicalPlan) -> Option<LogicalPlan> {
+        match plan {
+            // Anchor on the Sort itself. The projection (if any) sits *below* the
+            // Sort and supplies its output columns; SELECT * omits it entirely.
+            LogicalPlan::Sort(sort) => {
+                let (proj_exprs_slice, after_sort): (&[Expr], &LogicalPlan) =
+                    match sort.input.as_ref() {
+                        LogicalPlan::Projection(p) => (p.expr.as_slice(), p.input.as_ref()),
+                        other => (&[], other),
+                    };
+                self.build_rewrite(sort, proj_exprs_slice, after_sort)
+            }
+
+            // Output-aware passthrough. When a Projection sits directly over a
+            // k-NN Sort that rests on the scan (no projection between them), drive
+            // the rewrite with the OUTER projection's columns — i.e. the query's
+            // real output. The rewrite can only produce the index node's columns
+            // (addressing key + non-vector columns + _distance), never the indexed
+            // vector itself. Routing the output columns through `build_rewrite`
+            // lets it fire when they're all producible (e.g. `SELECT id … ORDER BY
+            // l2_distance(emb, …)`) and decline — falling back to exact search —
+            // when the output needs the vector (`SELECT *`, `SELECT id, emb`),
+            // rather than emitting a schema the consumer can't satisfy (issue #508).
+            //
+            // ALTERNATIVE (not taken): teach USearchExec to reconstruct the vector
+            // column for the result keys via `index.get(key)`, so even
+            // vector-returning queries stay on the index. Rejected to keep a single
+            // source of truth for returned vectors — the index would otherwise be a
+            // second source that must byte-match the parquet (breaks under F16
+            // quantization, and relies on USearch never transforming stored vectors).
+            // See the README "Limitations" entry and runtimedb issue #508.
+            LogicalPlan::Projection(outer) => {
+                let LogicalPlan::Sort(sort) = outer.input.as_ref() else {
+                    return None;
+                };
+                // Only the passthrough shape; the remap shape (projection *below*
+                // the Sort) is handled when we visit the Sort above.
+                if !matches!(
+                    sort.input.as_ref(),
+                    LogicalPlan::TableScan(_) | LogicalPlan::Filter(_)
+                ) {
+                    return None;
+                }
+                self.build_rewrite(sort, &outer.expr, sort.input.as_ref())
+            }
+
+            _ => None,
+        }
+    }
+
+    fn build_rewrite(
+        &self,
+        sort: &datafusion::logical_expr::logical_plan::Sort,
+        proj_exprs_slice: &[Expr],
+        after_sort: &LogicalPlan,
+    ) -> Option<LogicalPlan> {
         use datafusion::logical_expr::logical_plan::TableScan;
 
         // Require Sort with embedded fetch limit.
-        let sort = match plan {
-            LogicalPlan::Sort(s) => s,
-            _ => return None,
-        };
         let k = sort.fetch?;
 
-        // Projection is optional — DataFusion 51 omits it for SELECT * queries.
-        let (proj_exprs_slice, after_sort): (&[Expr], &LogicalPlan) = match sort.input.as_ref() {
-            LogicalPlan::Projection(p) => (p.expr.as_slice(), p.input.as_ref()),
-            other => (&[], other),
-        };
-
         // Accept TableScan directly, or Filter(TableScan) for WHERE clauses.
         // Deeper nesting (Filter→Filter→…) is not absorbed — the rule does
         // not fire and DataFusion falls back to exact execution.
@@ -155,7 +200,14 @@ impl USearchRule {
         // Build the final user-visible projection over USearchNode output.
         let dist_alias_str = dist_alias.as_deref().unwrap_or("_distance");
         let final_proj_exprs = if proj_exprs_slice.is_empty() {
-            passthrough_projection(&vsn_df_schema, &table_ref)
+            // No explicit Projection node (e.g. SELECT *, or a SELECT whose
+            // columns come straight from the scan, so the Sort sits directly on
+            // the TableScan). The rewrite must reproduce the original output
+            // columns; if any isn't producible from the node — the indexed
+            // vector column is never stored in the fetch path — bail so the
+            // query falls back to exact brute-force search, like the other
+            // unsupported shapes (DESC, metric mismatch, stacked filters).
+            passthrough_projection(after_sort.schema().as_ref(), &vsn_df_schema, &table_ref)?
         } else {
             remap_projections(proj_exprs_slice, dist_alias_str, &table_ref)
         };
@@ -375,21 +427,39 @@ fn build_outer_projection(exprs: &[Expr]) -> Vec<Expr> {
         .collect()
 }
 
-/// Build a passthrough Projection for SELECT * queries (no original Projection node).
-/// Projects only the original table columns (not `_distance`) so the output schema
-/// matches the original Sort schema. The Sort re-evaluates the distance UDF expression
-/// on the k result rows returned by USearchExec (O(k × dim), negligible for small k).
-fn passthrough_projection(schema: &DFSchema, table_ref: &TableReference) -> Vec<Expr> {
-    schema
+/// Build a passthrough Projection for queries with no explicit Projection node
+/// (e.g. `SELECT *`, or a SELECT whose columns come straight from the scan so the
+/// Sort sits directly on the TableScan).
+///
+/// The projection must reproduce the *original* output columns (`original_schema`,
+/// the Sort's input). The `USearchNode` can only produce the columns in
+/// `node_schema` — the fetch path's addressing key + non-vector columns +
+/// `_distance`; the indexed vector column is never stored there (see
+/// `PointLookupProvider`). If the original output needs a column the node can't
+/// produce (the vector column), return `None` so the rule declines to rewrite and
+/// the query falls back to exact brute-force search. The Sort re-evaluates the
+/// distance UDF on the k result rows returned by USearchExec (O(k × dim)).
+fn passthrough_projection(
+    original_schema: &DFSchema,
+    node_schema: &DFSchema,
+    table_ref: &TableReference,
+) -> Option<Vec<Expr>> {
+    original_schema
         .inner()
         .fields()
         .iter()
-        .filter(|f| f.name() != "_distance")
         .map(|f| {
-            Expr::Column(datafusion::common::Column::new(
-                Some(table_ref.clone()),
-                f.name().as_str(),
-            ))
+            let producible = node_schema
+                .inner()
+                .fields()
+                .iter()
+                .any(|nf| nf.name() == f.name());
+            producible.then(|| {
+                Expr::Column(datafusion::common::Column::new(
+                    Some(table_ref.clone()),
+                    f.name().as_str(),
+                ))
+            })
         })
         .collect()
 }