Support pg vector as pg_knn for postgres source (#57)

* # This is a combination of 3 commits.
# This is the 1st commit message:

support pg vector

support pg vector example

update example

hard-code k

fix CLI

remove fallback case

fix fallback

remove hnsw

demo yaml

fix k format

property

support subquery

update shared logic

use subquery

fix sub query path

# This is the commit message #2:

support multi type of query

# This is the commit message #3:

add tests

* support pg vector

support pg vector example

update example

hard-code k

fix CLI

remove fallback case

fix fallback

remove hnsw

demo yaml

fix k format

property

support subquery

update shared logic

use subquery

fix sub query path

support multi type of query

add tests

test updates

test updates

simplify and add tests

cleanup

limit k

update pipelines

bug fix

add conversion

no cast

* address comments

* stick to 64

* revert lance

* Revert "revert lance"

This reverts commit 75b8c6a.

Author: reluctanthero104

crates/cli/src/main.rs

@@ -18,9 +18,10 @@ use object_store::http::HttpBuilder;
 use serde::Deserialize;
 use skardi::sources::providers::lance::fts_table_function::register_lance_fts_udtf;
 use skardi::sources::providers::lance::knn_table_function::register_lance_knn_udtf;
+use skardi::sources::providers::sqlx::register_pg_knn_udtf;
 use skardi::sources::providers::{
-    iceberg::register_iceberg_table, lance::register_lance_table, mongo::register_mongo_tables,
-    mysql::register_mysql_tables, sqlite::register_sqlite_tables,
+    DatasetRegistry, iceberg::register_iceberg_table, lance::register_lance_table,
+    mongo::register_mongo_tables, mysql::register_mysql_tables, sqlite::register_sqlite_tables,
     sqlx::postgres::register_postgres_tables,
 };
 use std::collections::HashMap;
@@ -29,9 +30,6 @@ use std::path::{Path, PathBuf};
 use std::sync::{Arc, RwLock};
 use url::Url;
 
-/// Shared registry mapping table names to Lance datasets, used by the `lance_knn` UDTF.
-type DatasetRegistry = Arc<RwLock<HashMap<String, Arc<Dataset>>>>;
-
 #[derive(Parser)]
 #[command(name = "skardi")]
 #[command(about = "CLI tool for managing Skardi pipelines and data sources", long_about = None)]
@@ -278,7 +276,10 @@ impl UrlTableFactory for SkardiUrlTableFactory {
                 .unwrap_or(url)
                 .to_string();
             if let Ok(mut reg) = self.dataset_registry.write() {
-                reg.insert(table_name, Arc::clone(&dataset_arc));
+                reg.insert(
+                    table_name,
+                    skardi::sources::providers::DatasetEntry::Lance(Arc::clone(&dataset_arc)),
+                );
             }
 
             let provider: Arc<dyn TableProvider> = dataset_arc;
@@ -291,7 +292,7 @@ impl UrlTableFactory for SkardiUrlTableFactory {
 }
 
 /// Create a new SessionContext with custom URL table support (built-in files + Lance)
-/// and the `lance_knn` UDTF registered. Returns the context and the shared dataset registry.
+/// and the `lance_knn` / `pg_knn` UDTFs registered.
 fn new_session_context() -> (SessionContext, DatasetRegistry) {
     let dataset_registry: DatasetRegistry = Arc::new(RwLock::new(HashMap::new()));
     let session_store = SessionStore::new();
@@ -317,9 +318,10 @@ fn new_session_context() -> (SessionContext, DatasetRegistry) {
 
     factory.session_store().with_state(ctx.state_weak_ref());
 
-    // Register the lance_knn and lance_fts table functions
+    // Register the lance_knn, lance_fts, and pg_knn table functions, all sharing one registry
     register_lance_knn_udtf(&ctx, Arc::clone(&dataset_registry));
     register_lance_fts_udtf(&ctx, Arc::clone(&dataset_registry));
+    register_pg_knn_udtf(&ctx, Arc::clone(&dataset_registry));
 
     (ctx, dataset_registry)
 }
@@ -503,6 +505,7 @@ async fn register_source(
                 conn_str,
                 source.options.as_ref(),
                 false,
+                Some(dataset_registry),
             )
             .await
             .with_context(|| format!("Failed to register Postgres '{}'", source.name))?;

crates/server/src/config.rs

@@ -774,12 +774,14 @@ async fn register_data_source(
             );
 
             // Register PostgreSQL table using the sqlx-based provider
+            let pg_knn_registry = optimizer_registry.map(|r| r.pg_knn_pools());
             skardi::sources::providers::sqlx::postgres::register_postgres_tables(
                 session_ctx,
                 &source.name,
                 connection_string,
                 source.options.as_ref(),
                 source.access_mode.is_read_write(),
+                pg_knn_registry.as_ref(),
             )
             .await
             .map_err(|e| {

crates/server/src/optimizer_registry.rs

@@ -12,6 +12,8 @@ use anyhow::Result;
 use datafusion::prelude::SessionContext;
 use lance::dataset::Dataset;
 use skardi::sources::providers::lance::{register_lance_fts_udtf, register_lance_knn_udtf};
+use skardi::sources::providers::sqlx::register_pg_knn_udtf;
+use skardi::sources::providers::{DatasetEntry, DatasetRegistry};
 use std::collections::{HashMap, HashSet};
 use std::sync::{Arc, RwLock};
 
@@ -24,16 +26,17 @@ use crate::config::{DataSource, DataSourceType};
 /// 2. **Dataset Store** - Maintains references to datasets needed by table functions
 /// 3. **Lifecycle Coordinator** - Ensures datasets are available when functions need them
 pub struct OptimizerRegistry {
-    /// Lance datasets indexed by table name
-    /// Used by lance_knn table function to access datasets
-    lance_datasets: Arc<RwLock<HashMap<String, Arc<Dataset>>>>,
+    /// Unified dataset registry indexed by table name.
+    /// Stores both Lance datasets and Postgres entries, used by
+    /// lance_knn, lance_fts, and pg_knn table functions.
+    dataset_registry: DatasetRegistry,
 }
 
 impl OptimizerRegistry {
     /// Create a new registry
     pub fn new() -> Self {
         Self {
-            lance_datasets: Arc::new(RwLock::new(HashMap::new())),
+            dataset_registry: Arc::new(RwLock::new(HashMap::new())),
         }
     }
 
@@ -81,9 +84,9 @@ impl OptimizerRegistry {
             self.register_lance_functions(ctx)?;
         }
 
-        // Future: Register Postgres-specific UDFs
+        // Register Postgres-specific table functions
         if source_types.contains(&DataSourceType::Postgres) {
-            self.register_postgres_udfs(ctx)?;
+            self.register_postgres_functions(ctx)?;
         }
 
         Ok(())
@@ -95,46 +98,54 @@ impl OptimizerRegistry {
     fn register_lance_functions(&self, ctx: &mut SessionContext) -> Result<()> {
         tracing::info!("Registering Lance table functions");
 
-        // Register lance_knn table function
-        register_lance_knn_udtf(ctx, self.lance_datasets());
+        register_lance_knn_udtf(ctx, self.datasets());
         tracing::info!("✓ Registered lance_knn table function");
 
-        // Register lance_fts table function
-        register_lance_fts_udtf(ctx, self.lance_datasets());
+        register_lance_fts_udtf(ctx, self.datasets());
         tracing::info!("✓ Registered lance_fts table function");
 
         Ok(())
     }
 
-    /// Register Postgres-specific UDFs (placeholder for future)
-    fn register_postgres_udfs(&self, _ctx: &mut SessionContext) -> Result<()> {
-        tracing::debug!("Postgres UDFs not yet implemented");
+    /// Register Postgres-specific table functions.
+    fn register_postgres_functions(&self, ctx: &mut SessionContext) -> Result<()> {
+        tracing::info!("Registering pg_knn table function");
+        register_pg_knn_udtf(ctx, self.datasets());
+        tracing::info!("✓ Registered pg_knn table function");
         Ok(())
     }
 
-    // === Lance Dataset Management ===
+    // === Dataset Management ===
 
     /// Store a Lance dataset in the registry
-    ///
-    /// Called when registering a Lance table with DataFusion.
-    /// The lance_knn table function uses this to look up datasets by name.
     pub fn register_lance_dataset(&self, table_name: &str, dataset: Arc<Dataset>) {
-        let mut datasets = self.lance_datasets.write().unwrap();
-        datasets.insert(table_name.to_string(), dataset);
+        let mut reg = self.dataset_registry.write().unwrap();
+        reg.insert(table_name.to_string(), DatasetEntry::Lance(dataset));
         tracing::debug!("Registered Lance dataset '{}' in registry", table_name);
     }
 
     /// Get a Lance dataset by table name
     pub fn get_lance_dataset(&self, table_name: &str) -> Option<Arc<Dataset>> {
-        let datasets = self.lance_datasets.read().unwrap();
-        datasets.get(table_name).cloned()
+        let reg = self.dataset_registry.read().unwrap();
+        reg.get(table_name).and_then(|e| match e {
+            DatasetEntry::Lance(ds) => Some(Arc::clone(ds)),
+            _ => None,
+        })
     }
 
-    /// Get a clone of the Lance datasets map
-    ///
-    /// Returns an `Arc<RwLock<>>` that can be shared with table functions
-    pub fn lance_datasets(&self) -> Arc<RwLock<HashMap<String, Arc<Dataset>>>> {
-        Arc::clone(&self.lance_datasets)
+    /// Get a clone of the unified dataset registry to share with table functions.
+    pub fn datasets(&self) -> DatasetRegistry {
+        Arc::clone(&self.dataset_registry)
+    }
+
+    /// Alias for `datasets()` — used when passing to `register_postgres_tables`.
+    pub fn pg_knn_pools(&self) -> DatasetRegistry {
+        self.datasets()
+    }
+
+    /// Alias for `datasets()` — used when passing to `register_lance_table`.
+    pub fn lance_datasets(&self) -> DatasetRegistry {
+        self.datasets()
     }
 }