Allow pipeline parameter binding for k in lance_knn and pg_knn (#76)

lance_knn and pg_knn rejected non-literal k at planning time, so pipeline
`{k}` parameters (which surface as NULL during schema inference) could not
be used — unlike sqlite_knn which already had a NULL fallback. Hoist the
parsing into a shared knn_utils::extract_k helper that handles the NULL
inference placeholder and the MAX_KNN_K bound check, and adopt it in all
three KNN table functions. Update the lance/pg demo pipelines and their
READMEs to take k (and keep limit where it was already a parameter).

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: BtXin

crates/skardi/src/sources/providers/knn_utils.rs

@@ -9,12 +9,48 @@ pub const MAX_KNN_K: usize = 500;
 use arrow::array::{
     Array, BinaryArray, FixedSizeListArray, Float32Array, Float64Array, ListArray, StringArray,
 };
+use datafusion::common::{ScalarValue, plan_err};
 use datafusion::error::{DataFusionError, Result as DFResult};
 use datafusion::execution::TaskContext;
+use datafusion::logical_expr::Expr;
 use datafusion::physical_plan::{ExecutionPlan, execute_stream};
 use futures::StreamExt;
 use std::sync::Arc;
 
+/// Extract the `k` argument for a KNN table function from a planning-time expression.
+///
+/// Accepts positive integer literals (`Int32`, `Int64`, `UInt64`) and enforces the
+/// global [`MAX_KNN_K`] upper bound. During pipeline schema inference, pipeline
+/// `{param}` placeholders are converted to SQL `?` which reach the table function
+/// as `ScalarValue::Null` — in that case we return a dummy value of `1` so inference
+/// succeeds. The real integer is substituted textually at request time before
+/// re-planning, so the dummy value never runs.
+///
+/// `fn_name` is used as the error-message prefix (e.g. `"sqlite_knn"`).
+pub fn extract_k(expr: &Expr, fn_name: &str) -> DFResult<usize> {
+    let k = match expr {
+        Expr::Literal(ScalarValue::Int64(Some(v @ 1..)), _) => *v as usize,
+        Expr::Literal(ScalarValue::Int64(Some(v)), _) => {
+            return plan_err!("{fn_name}: k must be a positive integer, got {v}");
+        }
+        Expr::Literal(ScalarValue::Int32(Some(v @ 1..)), _) => *v as usize,
+        Expr::Literal(ScalarValue::Int32(Some(v)), _) => {
+            return plan_err!("{fn_name}: k must be a positive integer, got {v}");
+        }
+        Expr::Literal(ScalarValue::UInt64(Some(v)), _) if *v > 0 => *v as usize,
+        // Pipeline `{k}` parameter arrives as `?` → Null during schema inference.
+        // Return a dummy value; the real k is substituted textually at execute time.
+        Expr::Literal(ScalarValue::Null, _) => 1,
+        _ => {
+            return plan_err!("{fn_name}: k must be a positive integer literal");
+        }
+    };
+    if k > MAX_KNN_K {
+        return plan_err!("{fn_name}: k must be between 1 and {MAX_KNN_K}, got {k}");
+    }
+    Ok(k)
+}
+
 /// Execute a child plan and extract the first row's first column as `Vec<f32>`.
 ///
 /// Returns `None` if the plan produces no rows (e.g. a subquery filter matched nothing).

crates/skardi/src/sources/providers/lance/knn_table_function.rs

@@ -29,7 +29,7 @@ use std::sync::Arc;
 
 use super::knn_exec::LanceKnnExec;
 use super::utils::expr_to_lance_sql;
-use crate::sources::providers::knn_utils::MAX_KNN_K;
+use crate::sources::providers::knn_utils::extract_k;
 use crate::sources::providers::{DatasetEntry, DatasetRegistry};
 
 /// Table function that creates KNN search on Lance tables
@@ -56,13 +56,7 @@ impl TableFunctionImpl for LanceKnnTableFunction {
         // Extract string arguments
         let table_name = extract_string(&exprs[0], "table_name")?;
         let vector_column = extract_string(&exprs[1], "vector_column")?;
-        let k = {
-            let k = extract_int(&exprs[3], "k")?;
-            if k == 0 || k > MAX_KNN_K {
-                return plan_err!("lance_knn: k must be between 1 and {MAX_KNN_K}, got {k}");
-            }
-            k
-        };
+        let k = extract_k(&exprs[3], "lance_knn")?;
         let filter = if exprs.len() == 5 {
             Some(extract_string(&exprs[4], "filter")?)
         } else {
@@ -244,15 +238,6 @@ fn extract_string(expr: &Expr, name: &str) -> DFResult<String> {
     }
 }
 
-fn extract_int(expr: &Expr, name: &str) -> DFResult<usize> {
-    match expr {
-        Expr::Literal(ScalarValue::Int64(Some(n)), _) => Ok(*n as usize),
-        Expr::Literal(ScalarValue::Int32(Some(n)), _) => Ok(*n as usize),
-        Expr::Literal(ScalarValue::UInt64(Some(n)), _) => Ok(*n as usize),
-        _ => plan_err!("lance_knn: {} must be an integer literal", name),
-    }
-}
-
 fn try_extract_vector(expr: &Expr) -> DFResult<Option<ArrayRef>> {
     match expr {
         Expr::Literal(ScalarValue::List(arr), _) => {

crates/skardi/src/sources/providers/sqlite/knn_table_function.rs

@@ -36,7 +36,7 @@ use tokio_rusqlite::Connection;
 
 use super::knn_exec::SqliteKnnExec;
 use super::{expr_to_sqlite_sql, extract_string};
-use crate::sources::providers::knn_utils::MAX_KNN_K;
+use crate::sources::providers::knn_utils::extract_k;
 use crate::sources::providers::{DatasetEntry, DatasetRegistry};
 
 /// Entry stored in the registry for each registered SQLite table.
@@ -76,7 +76,7 @@ impl TableFunctionImpl for SqliteKnnTableFunction {
 
         let table_name = extract_string(&exprs[0], "table")?;
         let vector_col = extract_string(&exprs[1], "vector_col")?;
-        let k = extract_k(&exprs[3])?;
+        let k = extract_k(&exprs[3], "sqlite_knn")?;
 
         // Try to extract a literal vector.
         let literal_vector = extract_vector(&exprs[2]).ok();
@@ -259,27 +259,6 @@ pub fn register_sqlite_knn_udtf(ctx: &SessionContext, registry: DatasetRegistry)
 
 // ─── Argument extraction helpers ─────────────────────────────────────────────
 
-fn extract_k(expr: &Expr) -> DFResult<usize> {
-    let k = match expr {
-        Expr::Literal(ScalarValue::Int64(Some(v @ 1..)), _) => Ok(*v as usize),
-        Expr::Literal(ScalarValue::Int64(Some(v)), _) => {
-            plan_err!("sqlite_knn: k must be a positive integer, got {}", v)
-        }
-        Expr::Literal(ScalarValue::Int32(Some(v @ 1..)), _) => Ok(*v as usize),
-        Expr::Literal(ScalarValue::Int32(Some(v)), _) => {
-            plan_err!("sqlite_knn: k must be a positive integer, got {}", v)
-        }
-        Expr::Literal(ScalarValue::UInt64(Some(v)), _) => Ok(*v as usize),
-        // NULL placeholder during schema inference
-        Expr::Literal(ScalarValue::Null, _) => Ok(1),
-        _ => plan_err!("sqlite_knn: k must be a positive integer literal"),
-    }?;
-    if k > MAX_KNN_K {
-        return plan_err!("sqlite_knn: k must be between 1 and {MAX_KNN_K}, got {k}");
-    }
-    Ok(k)
-}
-
 fn extract_vector(expr: &Expr) -> DFResult<Vec<f32>> {
     let values: Arc<dyn arrow::array::Array> = match expr {
         Expr::Literal(ScalarValue::List(arr), _) => {

crates/skardi/src/sources/providers/sqlx/pg/knn_table_function.rs

@@ -41,7 +41,7 @@ use std::sync::Arc;
 
 use super::knn_exec::{DistanceMetric, PgKnnExec, PgVectorFetchExec};
 use super::utils::expr_to_pg_sql;
-use crate::sources::providers::knn_utils::MAX_KNN_K;
+use crate::sources::providers::knn_utils::extract_k;
 use crate::sources::providers::{DatasetEntry, DatasetRegistry};
 
 /// Entry stored in the registry for each registered Postgres table.
@@ -89,7 +89,7 @@ impl TableFunctionImpl for PgKnnTableFunction {
                 .map_err(datafusion::error::DataFusionError::Plan)?
         };
 
-        let k = extract_k(&exprs[4])?;
+        let k = extract_k(&exprs[4], "pg_knn")?;
 
         let inline_filter = if exprs.len() == 6 {
             extract_string(&exprs[5], "filter").ok()
@@ -467,19 +467,6 @@ fn pg_type_to_arrow(
 
 // ─── Argument extraction helpers ─────────────────────────────────────────────
 
-fn extract_k(expr: &Expr) -> DFResult<usize> {
-    let k = match expr {
-        Expr::Literal(ScalarValue::Int64(Some(n)), _) => Ok(*n as usize),
-        Expr::Literal(ScalarValue::Int32(Some(n)), _) => Ok(*n as usize),
-        Expr::Literal(ScalarValue::UInt64(Some(n)), _) => Ok(*n as usize),
-        _ => plan_err!("pg_knn: k must be a positive integer literal"),
-    }?;
-    if k == 0 || k > MAX_KNN_K {
-        return plan_err!("pg_knn: k must be between 1 and {MAX_KNN_K}, got {k}");
-    }
-    Ok(k)
-}
-
 fn extract_string(expr: &Expr, name: &str) -> DFResult<String> {
     match expr {
         Expr::Literal(ScalarValue::Utf8(Some(s)), _) => Ok(s.clone()),

demo/embeddings/candle/README.md

@@ -99,11 +99,12 @@ cargo run --bin skardi-server --features candle -- \
 curl -X POST http://localhost:8080/semantic-search/execute \
   -H "Content-Type: application/json" \
   -d '{
-    "query": "how does similarity search work in vector databases?"
+    "query": "how does similarity search work in vector databases?",
+    "k": 10
   }' | jq .
 ```
 
-**Response** (truncated — returns up to 10 results):
+**Response** (truncated — returns up to `k` results):
 ```json
 {
       "id": 1,
@@ -173,24 +174,25 @@ curl -X POST http://localhost:8080/semantic-search/execute \
 # Retrieval-Augmented Generation
 curl -X POST http://localhost:8080/semantic-search/execute \
   -H "Content-Type: application/json" \
-  -d '{"query": "how to ground LLM responses with retrieved documents"}' | jq .
+  -d '{"query": "how to ground LLM responses with retrieved documents", "k": 10}' | jq .
 
 # Arrow / columnar formats
 curl -X POST http://localhost:8080/semantic-search/execute \
   -H "Content-Type: application/json" \
-  -d '{"query": "columnar data formats for analytics"}' | jq .
+  -d '{"query": "columnar data formats for analytics", "k": 5}' | jq .
 
 # Model quantization
 curl -X POST http://localhost:8080/semantic-search/execute \
   -H "Content-Type: application/json" \
-  -d '{"query": "running models on CPU without a GPU"}' | jq .
+  -d '{"query": "running models on CPU without a GPU", "k": 5}' | jq .
 ```
 
 ## Pipeline Parameters
 
 | Parameter | Type | Required | Description |
 |---|---|---|---|
 | `query` | string | Yes | Free-text search query |
+| `k` | integer | Yes | Number of nearest neighbours to return |
 
 ## Directory Layout
 

demo/embeddings/candle/pipelines/pipeline_semantic_search.yaml

@@ -7,13 +7,16 @@ metadata:
     the nearest documents by cosine similarity.
   author: Skardi Demo
 
+# Parameters:
+#   {query} - Natural language search query (embedded on the fly)
+#   {k}     - Number of nearest neighbours to return
+
 query: |
   SELECT id, title, content, _distance
   FROM lance_knn(
     'doc_embeddings',
     'embedding',
     candle('models/generated/bge-small-en-v1.5', {query}),
-    10
+    {k}
   )
   ORDER BY _distance
-  LIMIT 10

demo/embeddings/gguf/README.md

@@ -104,7 +104,8 @@ cargo run --bin skardi-server --features gguf -- \
 curl -X POST http://localhost:8080/semantic-search-gguf/execute \
   -H "Content-Type: application/json" \
   -d '{
-    "query": "how does similarity search work in vector databases?"
+    "query": "how does similarity search work in vector databases?",
+    "k": 10
   }' | jq .
 ```
 
@@ -186,24 +187,25 @@ curl -X POST http://localhost:8080/semantic-search-gguf/execute \
 # Retrieval-Augmented Generation
 curl -X POST http://localhost:8080/semantic-search-gguf/execute \
   -H "Content-Type: application/json" \
-  -d '{"query": "how to ground LLM responses with retrieved documents"}' | jq .
+  -d '{"query": "how to ground LLM responses with retrieved documents", "k": 10}' | jq .
 
 # GGUF and quantisation
 curl -X POST http://localhost:8080/semantic-search-gguf/execute \
   -H "Content-Type: application/json" \
-  -d '{"query": "running quantised models on CPU without a GPU"}' | jq .
+  -d '{"query": "running quantised models on CPU without a GPU", "k": 5}' | jq .
 
 # Arrow / columnar formats
 curl -X POST http://localhost:8080/semantic-search-gguf/execute \
   -H "Content-Type: application/json" \
-  -d '{"query": "columnar data formats for analytics"}' | jq .
+  -d '{"query": "columnar data formats for analytics", "k": 5}' | jq .
 ```
 
 ## Pipeline Parameters
 
 | Parameter | Type | Required | Description |
 |---|---|---|---|
 | `query` | string | Yes | Free-text search query |
+| `k` | integer | Yes | Number of nearest neighbours to return |
 
 ## Directory Layout
 

demo/embeddings/gguf/pipelines/pipeline_semantic_search_gguf.yaml

@@ -8,13 +8,16 @@ metadata:
     the nearest documents by cosine similarity.
   author: Skardi Demo
 
+# Parameters:
+#   {query} - Natural language search query (embedded on the fly)
+#   {k}     - Number of nearest neighbours to return
+
 query: |
   SELECT id, title, content, _distance
   FROM lance_knn(
     'doc_embeddings_gguf',
     'embedding',
     gguf('models/generated/embeddinggemma-300m', {query}),
-    10
+    {k}
   )
   ORDER BY _distance
-  LIMIT 10

demo/embeddings/remote/README.md

@@ -45,7 +45,7 @@ cargo run --bin skardi-server --features remote-embed -- \
 ```bash
 curl -s "http://localhost:8080/semantic-search-remote/execute" \
   -H 'Content-Type: application/json' \
-  -d '{"query": "how does semantic search work?"}' | jq .
+  -d '{"query": "how does semantic search work?", "k": 10}' | jq .
 ```
 
 **Response**:

demo/embeddings/remote/pipelines/pipeline_semantic_search_remote.yaml

@@ -7,13 +7,16 @@ metadata:
     OpenAI API; lance_knn() finds the nearest documents by cosine similarity.
   author: Skardi Demo
 
+# Parameters:
+#   {query} - Natural language search query (embedded on the fly)
+#   {k}     - Number of nearest neighbours to return
+
 query: |
   SELECT id, title, content, _distance
   FROM lance_knn(
     'doc_embeddings_openai',
     'embedding',
     remote_embed('openai', 'text-embedding-3-small', {query}),
-    10
+    {k}
   )
   ORDER BY _distance
-  LIMIT 10