Merge pull request #18 from tembo-io/return_cols

Flexible return columns

Author: ChuckHend

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "vectorize"
-version = "0.2.0"
+version = "0.3.0"
 edition = "2021"
 publish = false
 

README.md

@@ -57,17 +57,17 @@ Finally, search.
 ```sql
 SELECT * FROM vectorize.search(
     job_name => 'product_search',
-    return_col => 'product_name',
     query => 'accessories for mobile devices',
-    api_key => 'my-openai-key"',
+    api_key => 'my-openai-key',
+    return_columns => ARRAY['product_id', 'product_name'],
     num_results => 3
 );
 ```
 
 ```text
-                                          search_results                                          
---------------------------------------------------------------------------------------------------
- {"value": "Phone Charger", "column": "product_name", "similarity_score": 0.8530797672121025}
- {"value": "Tablet Holder", "column": "product_name", "similarity_score": 0.8284493388477342}
- {"value": "Bluetooth Speaker", "column": "product_name", "similarity_score": 0.8255034841826178}
+                                         search_results                                         
+------------------------------------------------------------------------------------------------
+ {"product_id": 13, "product_name": "Phone Charger", "similarity_score": 0.8564774308489237}
+ {"product_id": 24, "product_name": "Tablet Holder", "similarity_score": 0.8295404213393001}
+ {"product_id": 4, "product_name": "Bluetooth Speaker", "similarity_score": 0.8248579643539758}
 ```

Trunk.toml

@@ -6,7 +6,7 @@ description = "The simplest implementation of LLM-backed vector search on Postgr
 homepage = "https://github.com/tembo-io/pg_vectorize"
 documentation = "https://github.com/tembo-io/pg_vectorize"
 categories = ["orchestration", "machine_learning"]
-version = "0.2.0"
+version = "0.3.0"
 
 [build]
 postgres_version = "15"

sql/vectorize--0.2.0--0.3.0.sql

@@ -0,0 +1,33 @@
+DROP FUNCTION vectorize."table";
+
+CREATE  FUNCTION vectorize."table"(
+        "table" TEXT, /* &str */
+        "columns" TEXT[], /* alloc::vec::Vec<alloc::string::String> */
+        "job_name" TEXT, /* core::option::Option<alloc::string::String> */
+        "args" json, /* pgrx::datum::json::Json */
+        "primary_key" TEXT, /* alloc::string::String */
+        "schema" TEXT DEFAULT 'public', /* alloc::string::String */
+        "update_col" TEXT DEFAULT 'last_updated_at', /* alloc::string::String */
+        "transformer" vectorize.Transformer DEFAULT 'openai', /* vectorize::types::Transformer */
+        "search_alg" vectorize.SimilarityAlg DEFAULT 'pgv_cosine_similarity', /* vectorize::types::SimilarityAlg */
+        "table_method" vectorize.TableMethod DEFAULT 'append', /* vectorize::init::TableMethod */
+        "schedule" TEXT DEFAULT '* * * * *' /* alloc::string::String */
+) RETURNS TEXT /* core::result::Result<alloc::string::String, anyhow::Error> */
+LANGUAGE c /* Rust */
+
+AS 'MODULE_PATHNAME', 'table_wrapper';
+
+DROP FUNCTION vectorize."search";
+
+CREATE  FUNCTION vectorize."search"(
+        "job_name" TEXT, /* &str */
+        "query" TEXT, /* &str */
+        "api_key" TEXT, /* &str */
+        "return_columns" TEXT[] DEFAULT ARRAY['*']::text[], /* alloc::vec::Vec<alloc::string::String> */
+        "num_results" INT DEFAULT 10 /* i32 */
+) RETURNS TABLE (
+        "search_results" jsonb  /* pgrx::datum::json::JsonB */
+)
+STRICT
+LANGUAGE c /* Rust */
+AS 'MODULE_PATHNAME', 'search_wrapper';

src/api.rs

@@ -116,9 +116,9 @@ fn table(
 #[pg_extern]
 fn search(
     job_name: &str,
-    return_col: &str,
     query: &str,
     api_key: &str,
+    return_columns: default!(Vec<String>, "ARRAY['*']::text[]"),
     num_results: default!(i32, 10),
 ) -> Result<TableIterator<'static, (name!(search_results, pgrx::JsonB),)>, spi::Error> {
     // note: this is not the most performant implementation
@@ -159,7 +159,7 @@ fn search(
         job_name,
         &schema,
         &table,
-        return_col,
+        &return_columns,
         num_results,
         &embeddings[0],
     )?;

src/search.rs

@@ -4,43 +4,36 @@ pub fn cosine_similarity_search(
     project: &str,
     schema: &str,
     table: &str,
-    return_col: &str,
+    return_columns: &[String],
     num_results: i32,
     embeddings: &[f64],
 ) -> Result<Vec<(pgrx::JsonB,)>, spi::Error> {
     let emb = serde_json::to_string(&embeddings).expect("failed to serialize embeddings");
     let query = format!(
         "
-    SELECT 
-        1 - ({project}_embeddings <=> '{emb}'::vector) AS cosine_similarity,
-        *
+    SELECT to_jsonb(t)
+    as results FROM (
+        SELECT 
+        1 - ({project}_embeddings <=> '{emb}'::vector) AS similarity_score,
+        {cols}
     FROM {schema}.{table}
     WHERE {project}_updated_at is NOT NULL
-    ORDER BY cosine_similarity DESC
-    LIMIT {num_results};
-    "
+    ORDER BY similarity_score DESC
+    LIMIT {num_results}
+    ) t
+    ",
+        cols = return_columns.join(", "),
     );
     log!("query: {}", query);
     Spi::connect(|client| {
         let mut results: Vec<(pgrx::JsonB,)> = Vec::new();
         let tup_table = client.select(&query, None, None)?;
-
         for row in tup_table {
-            let v = row[return_col]
-                .value::<String>()
-                .expect("failed to get value");
-            let score = row["cosine_similarity"]
-                .value::<f64>()
-                .expect("failed to get value");
-
-            let r = serde_json::json!({
-                "column": return_col,
-                "value": v,
-                "similarity_score": score
-            });
-            results.push((pgrx::JsonB(r),));
+            match row["results"].value()? {
+                Some(r) => results.push((r,)),
+                None => error!("failed to get results"),
+            }
         }
-
         Ok(results)
     })
 }