fix(rule): preserve hidden distance for sort

Keep _distance in an inner projection when ORDER BY uses a vector\ndistance expression that is not part of the final select list.\n\nThis fixes split-provider execution for queries like SELECT id ORDER\nBY l2_distance(vector, ARRAY[...]) LIMIT k while preserving the final\noutput schema. Add an execution test for the direct ORDER BY shape to\ncover the production case.

Author: anoop-narang

src/rule.rs

@@ -18,9 +18,10 @@
 //
 // Replacement:
 //
-//   Sort(fetch=k)                                             ← kept (sort order)
-//     Projection([col(a), col(b), col("_distance").alias("dist")])
-//       USearchNode                                           ← executes ANN
+//   Projection([final output cols])
+//     Sort(fetch=k)
+//       Projection([final output cols + optional hidden _distance])
+//         USearchNode
 
 use std::collections::HashMap;
 use std::sync::Arc;
@@ -151,25 +152,33 @@ impl USearchRule {
             node: Arc::new(node) as Arc<dyn UserDefinedLogicalNode>,
         }));
 
-        // Build Projection over USearchNode matching the original output schema.
+        // Build the final user-visible projection over USearchNode output.
         let dist_alias_str = dist_alias.as_deref().unwrap_or("_distance");
-        let new_proj_exprs = if proj_exprs_slice.is_empty() {
+        let final_proj_exprs = if proj_exprs_slice.is_empty() {
             passthrough_projection(&vsn_df_schema, &table_ref)
         } else {
             remap_projections(proj_exprs_slice, dist_alias_str, &table_ref)
         };
-        let new_proj = Projection::try_new(new_proj_exprs, node_plan).ok()?;
-
-        // Keep the Sort node so DataFusion handles ordering by _distance / dist.
-        // USearch returns results in arbitrary (internal) order when the underlying
-        // data is fetched from the TableProvider.
-        Some(LogicalPlan::Sort(
-            datafusion::logical_expr::logical_plan::Sort {
-                expr: sort.expr.clone(),
-                input: Arc::new(LogicalPlan::Projection(new_proj)),
-                fetch: sort.fetch,
-            },
-        ))
+        let remapped_sort_exprs = remap_sort_exprs(&sort.expr, dist_alias.as_deref());
+        let needs_hidden_distance = remapped_sort_exprs.iter().any(|e| {
+            matches!(&e.expr, Expr::Column(c) if c.relation.is_none() && c.name == "_distance")
+        }) && !projection_exposes_name(&final_proj_exprs, "_distance");
+
+        let mut sort_input_exprs = final_proj_exprs.clone();
+        if needs_hidden_distance {
+            sort_input_exprs.push(col("_distance"));
+        }
+
+        let sort_input = Projection::try_new(sort_input_exprs, node_plan).ok()?;
+        let sorted = LogicalPlan::Sort(datafusion::logical_expr::logical_plan::Sort {
+            expr: remapped_sort_exprs,
+            input: Arc::new(LogicalPlan::Projection(sort_input)),
+            fetch: sort.fetch,
+        });
+
+        let outer_proj_exprs = build_outer_projection(&final_proj_exprs);
+        let outer_proj = Projection::try_new(outer_proj_exprs, Arc::new(sorted)).ok()?;
+        Some(LogicalPlan::Projection(outer_proj))
     }
 }
 
@@ -283,7 +292,11 @@ fn dist_type_matches_metric(dist_type: &DistanceType, metric: MetricKind) -> boo
 }
 
 fn is_distance_expr(expr: &Expr) -> bool {
-    matches!(expr, Expr::ScalarFunction(sf) if is_dist_udf_name(sf.func.name()))
+    let inner = match expr {
+        Expr::Alias(a) => a.expr.as_ref(),
+        other => other,
+    };
+    matches!(inner, Expr::ScalarFunction(sf) if is_dist_udf_name(sf.func.name()))
 }
 
 fn try_extract_distance(expr: &Expr) -> Option<(String, String, Vec<f64>)> {
@@ -322,6 +335,45 @@ fn remap_projections(
         .collect()
 }
 
+fn remap_sort_exprs(
+    sort_exprs: &[datafusion::logical_expr::SortExpr],
+    dist_alias_name: Option<&str>,
+) -> Vec<datafusion::logical_expr::SortExpr> {
+    sort_exprs
+        .iter()
+        .map(|sort_expr| {
+            let remapped_expr = match &sort_expr.expr {
+                Expr::Column(c) if Some(c.name.as_str()) == dist_alias_name => col(c.name.as_str()),
+                expr if is_distance_expr(expr) => col("_distance"),
+                other => other.clone(),
+            };
+            datafusion::logical_expr::SortExpr {
+                expr: remapped_expr,
+                asc: sort_expr.asc,
+                nulls_first: sort_expr.nulls_first,
+            }
+        })
+        .collect()
+}
+
+fn projection_exposes_name(exprs: &[Expr], name: &str) -> bool {
+    exprs.iter().any(|expr| match expr {
+        Expr::Alias(a) => a.name == name,
+        Expr::Column(c) => c.name == name,
+        _ => false,
+    })
+}
+
+fn build_outer_projection(exprs: &[Expr]) -> Vec<Expr> {
+    exprs.iter()
+        .filter_map(|expr| match expr {
+            Expr::Alias(a) => Some(col(a.name.as_str())),
+            Expr::Column(c) => Some(Expr::Column(c.clone())),
+            _ => None,
+        })
+        .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

tests/execution.rs

@@ -530,6 +530,17 @@ async fn exec_split_provider_select_specific_columns() {
     assert_eq!(ids.len(), 2, "expected 2 results; got {ids:?}");
 }
 
+/// SELECT specific columns without projecting the distance expression.
+/// This is the production shape behind `vector_distance(...)`.
+#[tokio::test]
+async fn exec_split_provider_order_by_udf_direct() {
+    let ctx = make_split_provider_ctx("items::vector").await;
+    let sql = format!("SELECT id FROM items ORDER BY l2_distance(vector, {Q}) ASC LIMIT 2");
+    let ids = collect_ids(&ctx, &sql).await;
+    assert_eq!(ids[0], 1, "closest must be row 1\nids: {ids:?}");
+    assert_eq!(ids.len(), 2, "expected 2 results; got {ids:?}");
+}
+
 /// SELECT * with distance UDF — should fall back to UDF brute-force
 /// (since vector column is not in lookup provider schema).
 #[tokio::test]