From 40fdf563ed39a9e4f00127de1b3b5351bcf6da0b Mon Sep 17 00:00:00 2001
From: Daniel Rammer <hamersaw@protonmail.com>
Date: Wed, 16 Sep 2026 21:16:59 -0500
Subject: [PATCH] feat(knn): let a caller rebuild a batch KNN node with a new k
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A batch `KNNVectorDistanceExec` bounds every query's candidates by its own
`k` and carries no enclosing top-k `SortExec`, so a plan rewriter that wants
to widen the candidate set has nothing to move — `try_new_batch` is
`pub(crate)` and `retain_vector` is private, so it can't rebuild the node
either.

`with_k` covers that: `k` feeds only the execute-time cut, never the schema
or plan properties, so everything else carries over unchanged.

Sophon's WAL union needs it to over-fetch the base arm on the unindexed
batch shape, the way it already does for `ANNIvfBatchExec`.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 rust/lance/src/io/exec/knn.rs | 77 +++++++++++++++++++++++++++++++++++
 1 file changed, 77 insertions(+)

diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs
index 3f10653cf32..d2ed1877fc0 100644
--- a/rust/lance/src/io/exec/knn.rs
+++ b/rust/lance/src/io/exec/knn.rs
@@ -553,6 +553,38 @@ impl KNNVectorDistanceExec {
         })
     }
 
+    /// Rebuild this node with a different per-query `k`.
+    ///
+    /// A batch node bounds every query's candidates by its own `k` and carries no
+    /// enclosing top-k `SortExec`, so a caller that rewrites the plan to widen the
+    /// candidate set has nothing else to move. `k` feeds only the execute-time cut,
+    /// never the schema or plan properties, so the rest of the node carries over.
+    ///
+    /// Returns an error for a zero `k` on a batch node, matching [`Self::try_new_batch`].
+    pub fn with_k(&self, k: usize) -> Result<Self> {
+        if self.is_batch && k == 0 {
+            return Err(Error::invalid_input(
+                "k must be positive for batch KNN".to_string(),
+            ));
+        }
+        Ok(Self {
+            input: self.input.clone(),
+            query: self.query.clone(),
+            is_batch: self.is_batch,
+            query_count: self.query_count,
+            k,
+            lower_bound: self.lower_bound,
+            upper_bound: self.upper_bound,
+            column: self.column.clone(),
+            distance_type: self.distance_type,
+            retain_vector: self.retain_vector,
+            input_schema: self.input_schema.clone(),
+            output_schema: self.output_schema.clone(),
+            properties: self.properties.clone(),
+            metrics: ExecutionPlanMetricsSet::new(),
+        })
+    }
+
     fn take_vector_row(vectors: &dyn Array, row_index: u32) -> DataFusionResult<ArrayRef> {
         let indices = UInt32Array::from_iter([Some(row_index)]);
         arrow_select::take::take(vectors, &indices, None)
@@ -4393,6 +4425,51 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_batch_with_k_rebuilds_the_cut_and_keeps_the_schema() {
+        let schema = Arc::new(ArrowSchema::new(vec![
+            ArrowField::new("i", DataType::Int32, true),
+            ArrowField::new(
+                "vec",
+                DataType::FixedSizeList(
+                    Arc::new(ArrowField::new("item", DataType::Float32, true)),
+                    4,
+                ),
+                true,
+            ),
+            ROW_ID_FIELD.clone(),
+        ]));
+        let batch = RecordBatch::new_empty(schema);
+        let input: Arc<dyn ExecutionPlan> = Arc::new(TestingExec::new(vec![batch]));
+        let query = Arc::new(Float32Array::from(vec![
+            0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0,
+        ])) as ArrayRef;
+        let plan = KNNVectorDistanceExec::try_new_batch(
+            input,
+            "vec",
+            query,
+            KnnBatchParams {
+                is_batch: true,
+                query_count: 2,
+                k: 2,
+                lower_bound: None,
+                upper_bound: None,
+                distance_type: DistanceType::L2,
+                retain_vector: false,
+            },
+        )
+        .unwrap();
+
+        let widened = plan.with_k(7).unwrap();
+        assert_eq!(widened.k, 7);
+        assert_eq!(widened.query_count, 2);
+        assert!(widened.is_batch);
+        assert_eq!(widened.schema(), plan.schema());
+        assert_eq!(plan.k, 2, "the original node is left alone");
+
+        assert!(plan.with_k(0).is_err(), "batch k must stay positive");
+    }
+
     #[test]
     fn test_batch_partition_statistics_aligns_with_output_schema() {
         let schema = Arc::new(ArrowSchema::new(vec![