Merge branch 'georgia-tech-db:feature/remote_indexes' into feature/ze…

…ro-vector-opt
georgia-tech-db · Apr 12, 2024 · ec90e20 · ec90e20
2 parents 02c404f + 6327b14
commit ec90e20
Show file tree

Hide file tree

Showing 9 changed files with 268 additions and 25 deletions.
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -52,7 +52,7 @@ jobs:
       - run: psql test -c 'alter database test set enable_seqscan = off'
 
       # setup the database for testing
-      - run: make installcheck REGRESS="pinecone_crud pinecone_medium_create pinecone_zero_vector_insert" REGRESS_OPTS="--dbname=test --inputdir=./test --use-existing"
+      - run: make installcheck REGRESS="pinecone_crud pinecone_medium_create pinecone_zero_vector_insert pinecone_build_after_insert pinecone_invalid_config" REGRESS_OPTS="--dbname=test --inputdir=./test --use-existing"
       - if: ${{ failure() }}
         run: cat regression.diffs
   # mac:

diff --git a/comparison.html b/comparison.html
@@ -0,0 +1,92 @@
+<div class="relative overflow-hidden shadow-md rounded-lg">
+    <table class="table-fixed w-full text-left">
+        <thead class="uppercase bg-[#4666a7] text-[#ffffff]" style="background-color: #4666a7; color: #ffffff;">
+            <tr>
+                <td class="py-1 border text-center  p-4" contenteditable="true">
+                    <br>
+                </td>
+                <td class="py-1 border text-center  p-4" contenteditable="true">PGVECTOR</td>
+                <td class="py-1 border text-center  p-4" contenteditable="true">PINECONE</td>
+                <td class="py-1 border text-center  p-4" contenteditable="true">PGVECTOR-REMOTE</td>
+            </tr>
+        </thead>
+        <tbody class="bg-white text-gray-500 bg-[#FFFFFF] text-[#0f2e6c]" style="background-color: #FFFFFF; color: #0f2e6c;">
+            <tr class="py-5">
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">SQL Interface
+                    <br>
+                </td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true"><code>✔</code>
+
+                </td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">
+                    <br>
+                </td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true"><code>✔</code>
+
+                </td>
+            </tr>
+            <tr class="py-5">
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">Transaction Support
+                    <br>
+                </td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true"><code>✔</code>
+
+                </td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true"></td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true"><code>✔</code>
+
+                </td>
+            </tr>
+            <tr class="py-5">
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">10+ Client Languages
+                    <br>
+                </td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true"><code>✔</code>
+
+                </td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true"></td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true"><code>✔</code>
+
+                </td>
+            </tr>
+            <tr class="py-5">
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">Data Export
+                    <br>
+                </td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true"><code>✔</code>
+
+                </td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">
+                    <br>
+                </td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true"><code>✔</code>
+
+                </td>
+            </tr>
+            <tr class="py-5">
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">Latency*</td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">300ms</td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">4ms</td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">7ms</td>
+            </tr>
+            <tr class="py-5">
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">Throughput*</td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">3QPS</td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">300QPS</td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">300QPS</td>
+            </tr>
+            <tr class="py-5">
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true">Metadata Filtering
+                    <br>
+                </td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true"></td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true"><code>✔</code>
+
+                </td>
+                <td class="py-5 border text-center font-bold p-4" contenteditable="true"><code>✔</code>
+
+                </td>
+            </tr>
+        </tbody>
+    </table>
+</div>
diff --git a/src/pinecone/pinecone.h b/src/pinecone/pinecone.h
@@ -200,7 +200,7 @@ IndexBulkDeleteResult *pinecone_bulkdelete(IndexVacuumInfo *info, IndexBulkDelet
 IndexBulkDeleteResult *no_vacuumcleanup(IndexVacuumInfo *info, IndexBulkDeleteResult *stats);
 
 // validate
-void pinecone_spec_validator(const char *spec);
+void pinecone_spec_validator(const PineconeOptions *opts);
 void pinecone_host_validator(const char *spec);
 void validate_api_key(void);
 bool validate_vector_nonzero(Vector* vector);

diff --git a/src/pinecone/pinecone_build.c b/src/pinecone/pinecone_build.c
@@ -53,12 +53,18 @@ IndexBuildResult *pinecone_build(Relation heap, Relation index, IndexInfo *index
     PineconeOptions *opts = (PineconeOptions *) index->rd_options;
     IndexBuildResult *result = palloc(sizeof(IndexBuildResult));
     VectorMetric metric = get_opclass_metric(index);
-    cJSON* spec_json = cJSON_Parse(GET_STRING_RELOPTION(opts, spec));
-    int dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
-    char* pinecone_index_name = get_pinecone_index_name(index);
-    char* host = GET_STRING_RELOPTION(opts, host);
+
+    cJSON* spec_json;
+    char* pinecone_index_name;
+    char* host;
+    int dimensions;
     cJSON* describe_index_response;
 
+    pinecone_spec_validator(opts);
+    spec_json = cJSON_Parse(GET_STRING_RELOPTION(opts, spec));
+    dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
+    pinecone_index_name = get_pinecone_index_name(index);
+    host = GET_STRING_RELOPTION(opts, host);
     validate_api_key();
 
     // if the host is specified, check that it is empty
@@ -93,13 +99,20 @@ IndexBuildResult *pinecone_build(Relation heap, Relation index, IndexInfo *index
     } else {
         cJSON* index_stats_response;
         InsertBaseTable(heap, index, indexInfo, host, result);
-        // wait for the remote index to finish processing the vectors
-        // i.e. describe stats is equal to result->index_tuples
-        index_stats_response = pinecone_get_index_stats(pinecone_api_key, host);
-        while (cJSON_GetObjectItemCaseSensitive(index_stats_response, "totalVectorCount")->valueint < result->index_tuples) {
-            sleep(1);
-            index_stats_response = pinecone_get_index_stats(pinecone_api_key, host);
-        }
+
+        #ifdef PINECONE_MOCK
+            if (!pinecone_use_mock_response) {
+        #endif
+                // wait for the remote index to finish processing the vectors
+                // i.e. describe stats is equal to result->index_tuples
+                index_stats_response = pinecone_get_index_stats(pinecone_api_key, host);
+                while (cJSON_GetObjectItemCaseSensitive(index_stats_response, "totalVectorCount")->valueint < result->index_tuples) {
+                    sleep(1);
+                    index_stats_response = pinecone_get_index_stats(pinecone_api_key, host);
+                }
+        #ifdef PINECONE_MOCK
+            }
+        #endif
     }
     return result;
 }
@@ -125,6 +138,7 @@ char* CreatePineconeIndexAndWait(Relation index, cJSON* spec_json, VectorMetric
         sleep(1);
         create_response = describe_index(pinecone_api_key, pinecone_index_name);
     }
+    sleep(1); // TODO: ping the host with get_index_stats instead of pinging pinecone.io with describe_index
     return host;
 }
 

diff --git a/src/pinecone/pinecone_validate.c b/src/pinecone/pinecone_validate.c
@@ -23,16 +23,15 @@ bool validate_vector_nonzero(Vector* vector) {
 }
 
 
-void pinecone_spec_validator(const char *spec)
+void pinecone_spec_validator(const PineconeOptions *opts)
 {
-    bool empty = strcmp(spec, "") == 0;
-    if (empty || cJSON_Parse(spec) == NULL)
+    if (opts == NULL || cJSON_Parse(GET_STRING_RELOPTION(opts, spec)) == NULL || strcmp(GET_STRING_RELOPTION(opts, spec), "") == 0)
     {
         ereport(ERROR,
                 (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
-                (empty ? errmsg("Spec cannot be empty") : errmsg("Invalid spec: %s", spec)),
-                errhint("Spec should be a valid JSON object e.g. WITH (spec='{\"serverless\":{\"cloud\":\"aws\",\"region\":\"us-west-2\"}}').\n \
-                        Refer to https://docs.pinecone.io/reference/create_index")));
+                 errmsg("Invalid spec"),
+                 errhint("Spec should be a valid JSON object e.g. WITH (spec='{\"serverless\":{\"cloud\":\"aws\",\"region\":\"us-west-2\"}}').\n \
+                         Refer to https://docs.pinecone.io/reference/create_index")));
     }
 }
 

diff --git a/test/expected/pinecone_invalid_config.out b/test/expected/pinecone_invalid_config.out
@@ -1,4 +1,5 @@
 SET enable_seqscan = off;
+SET client_min_messages = 'notice';
 ALTER SYSTEM RESET pinecone.api_key;
 SELECT pg_reload_conf();
  pg_reload_conf 
@@ -10,13 +11,15 @@ CREATE TABLE t (val vector(3));
 CREATE INDEX i2 ON t USING pinecone (val) WITH (spec = '{"serverless":{"cloud":"aws","region":"us-west-2"}}');
 ERROR:  Pinecone API key not set
 HINT:  Set the pinecone API key using the pinecone.api_key GUC. E.g. ALTER SYSTEM SET pinecone.api_key TO 'your-api-key'
-ALTER SYSTEM SET pinecone.api_key = '5b2c1031-ba58-4acc-a634-9f943d68822c';
+ALTER SYSTEM SET pinecone.api_key = 'fake-key';
 SELECT pg_reload_conf();
  pg_reload_conf 
 ----------------
  t
 (1 row)
 
 CREATE INDEX i2 ON t USING pinecone (val);
-ERROR:  Spec cannot be empty
-DROP TABLE t;
+ERROR:  Invalid spec
+HINT:  Spec should be a valid JSON object e.g. WITH (spec='{"serverless":{"cloud":"aws","region":"us-west-2"}}').
+                          Refer to https://docs.pinecone.io/reference/create_index
+DROP TABLE t;
diff --git a/test/expected/pinecone_metrics.out b/test/expected/pinecone_metrics.out
@@ -1,6 +1,86 @@
+delete from pinecone_mock;
+SET client_min_messages = 'notice';
 SET enable_seqscan = off;
+-- flush each vector individually
+SET pinecone.vectors_per_request = 1;
+SET pinecone.requests_per_batch = 1;
+-- CREATE TABLE
 CREATE TABLE t (val vector(3));
+-- mock create index
+INSERT INTO pinecone_mock (url_prefix, method, response)
+VALUES ('https://api.pinecone.io/indexes', 'POST', $${
+        "name": "invalid",
+        "metric": "euclidean",
+        "dimension": 3,
+        "status": {
+                "ready": true,
+                "state": "Ready"
+        },
+        "host": "fakehost",
+        "spec": {
+                "serverless": {
+                        "cloud": "aws",
+                        "region": "us-west-2"
+                }
+        }
+}$$);
+-- mock describe index stats
+INSERT INTO pinecone_mock (url_prefix, method, response)
+VALUES ('https://fakehost/describe_index_stats', 'GET', '{"namespaces":{},"dimension":3,"indexFullness":0,"totalVectorCount":0}');
+-- create index
 CREATE INDEX i2 ON t USING pinecone (val vector_l2_ops) WITH (spec = '{"serverless":{"cloud":"aws","region":"us-west-2"}}');
+-- mock upsert
+INSERT INTO pinecone_mock (url_prefix, method, response) VALUES ('https://fakehost/vectors/upsert', 'POST', '{"upsertedCount":1}');
+INSERT INTO t (val) VALUES ('[1,0,0]');
+INSERT INTO t (val) VALUES ('[1,0,1]');
+INSERT INTO t (val) VALUES ('[1,1,0]');
+EXPLAIN SELECT val,val<->'[1,1,1]' as dist FROM t ORDER BY val <-> '[1, 1, 1]';
+                           QUERY PLAN                           
+----------------------------------------------------------------
+ Index Scan using i2 on t  (cost=0.00..7.41 rows=1360 width=40)
+   Order By: (val <-> '[1,1,1]'::vector)
+(2 rows)
+
+EXPLAIN SELECT val,val<=>'[1,1,1]' as dist FROM t ORDER BY val <=> '[1, 1, 1]';
+                                  QUERY PLAN                                   
+-------------------------------------------------------------------------------
+ Sort  (cost=10000000097.78..10000000101.18 rows=1360 width=40)
+   Sort Key: ((val <=> '[1,1,1]'::vector))
+   ->  Seq Scan on t  (cost=10000000000.00..10000000027.00 rows=1360 width=40)
+(3 rows)
+
+DROP INDEX i2;
+INSERT INTO pinecone_mock (url_prefix, method, response)
+VALUES ('https://api.pinecone.io/indexes', 'POST', $${
+        "name": "invalid",
+        "metric": "cosine",
+        "dimension": 3,
+        "status": {
+                "ready": true,
+                "state": "Ready"
+        },
+        "host": "fakehost",
+        "spec": {
+                "serverless": {
+                        "cloud": "aws",
+                        "region": "us-west-2"
+                }
+        }
+}$$);
 CREATE INDEX i3 ON t USING pinecone (val vector_cosine_ops) WITH (spec = '{"serverless":{"cloud":"aws","region":"us-west-2"}}');
-CREATE INDEX i1 ON t USING pinecone (val vector_ip_ops) WITH (spec = '{"serverless":{"cloud":"aws","region":"us-west-2"}}');
+EXPLAIN SELECT val,val<=>'[1,1,1]' as dist FROM t ORDER BY val <=> '[1, 1, 1]';
+                         QUERY PLAN                          
+-------------------------------------------------------------
+ Index Scan using i3 on t  (cost=0.00..4.02 rows=3 width=40)
+   Order By: (val <=> '[1,1,1]'::vector)
+(2 rows)
+
+EXPLAIN SELECT val,val<->'[1,1,1]' as dist FROM t ORDER BY val <-> '[1, 1, 1]';
+                                 QUERY PLAN                                 
+----------------------------------------------------------------------------
+ Sort  (cost=10000000001.06..10000000001.07 rows=3 width=40)
+   Sort Key: ((val <-> '[1,1,1]'::vector))
+   ->  Seq Scan on t  (cost=10000000000.00..10000000001.04 rows=3 width=40)
+(3 rows)
+
 DROP TABLE t;
diff --git a/test/sql/pinecone_invalid_config.sql b/test/sql/pinecone_invalid_config.sql
@@ -1,9 +1,10 @@
 SET enable_seqscan = off;
+SET client_min_messages = 'notice';
 ALTER SYSTEM RESET pinecone.api_key;
 SELECT pg_reload_conf();
 CREATE TABLE t (val vector(3));
 CREATE INDEX i2 ON t USING pinecone (val) WITH (spec = '{"serverless":{"cloud":"aws","region":"us-west-2"}}');
-ALTER SYSTEM SET pinecone.api_key = '5b2c1031-ba58-4acc-a634-9f943d68822c';
+ALTER SYSTEM SET pinecone.api_key = 'fake-key';
 SELECT pg_reload_conf();
 CREATE INDEX i2 ON t USING pinecone (val);
 DROP TABLE t;
diff --git a/test/sql/pinecone_metrics.sql b/test/sql/pinecone_metrics.sql
@@ -1,6 +1,60 @@
+delete from pinecone_mock;
+SET client_min_messages = 'notice';
 SET enable_seqscan = off;
+-- flush each vector individually
+SET pinecone.vectors_per_request = 1;
+SET pinecone.requests_per_batch = 1;
+-- CREATE TABLE
 CREATE TABLE t (val vector(3));
+-- mock create index
+INSERT INTO pinecone_mock (url_prefix, method, response)
+VALUES ('https://api.pinecone.io/indexes', 'POST', $${
+        "name": "invalid",
+        "metric": "euclidean",
+        "dimension": 3,
+        "status": {
+                "ready": true,
+                "state": "Ready"
+        },
+        "host": "fakehost",
+        "spec": {
+                "serverless": {
+                        "cloud": "aws",
+                        "region": "us-west-2"
+                }
+        }
+}$$);
+-- mock describe index stats
+INSERT INTO pinecone_mock (url_prefix, method, response)
+VALUES ('https://fakehost/describe_index_stats', 'GET', '{"namespaces":{},"dimension":3,"indexFullness":0,"totalVectorCount":0}');
+-- create index
 CREATE INDEX i2 ON t USING pinecone (val vector_l2_ops) WITH (spec = '{"serverless":{"cloud":"aws","region":"us-west-2"}}');
+-- mock upsert
+INSERT INTO pinecone_mock (url_prefix, method, response) VALUES ('https://fakehost/vectors/upsert', 'POST', '{"upsertedCount":1}');
+INSERT INTO t (val) VALUES ('[1,0,0]');
+INSERT INTO t (val) VALUES ('[1,0,1]');
+INSERT INTO t (val) VALUES ('[1,1,0]');
+EXPLAIN SELECT val,val<->'[1,1,1]' as dist FROM t ORDER BY val <-> '[1, 1, 1]';
+EXPLAIN SELECT val,val<=>'[1,1,1]' as dist FROM t ORDER BY val <=> '[1, 1, 1]';
+DROP INDEX i2;
+INSERT INTO pinecone_mock (url_prefix, method, response)
+VALUES ('https://api.pinecone.io/indexes', 'POST', $${
+        "name": "invalid",
+        "metric": "cosine",
+        "dimension": 3,
+        "status": {
+                "ready": true,
+                "state": "Ready"
+        },
+        "host": "fakehost",
+        "spec": {
+                "serverless": {
+                        "cloud": "aws",
+                        "region": "us-west-2"
+                }
+        }
+}$$);
 CREATE INDEX i3 ON t USING pinecone (val vector_cosine_ops) WITH (spec = '{"serverless":{"cloud":"aws","region":"us-west-2"}}');
-CREATE INDEX i1 ON t USING pinecone (val vector_ip_ops) WITH (spec = '{"serverless":{"cloud":"aws","region":"us-west-2"}}');
+EXPLAIN SELECT val,val<=>'[1,1,1]' as dist FROM t ORDER BY val <=> '[1, 1, 1]';
+EXPLAIN SELECT val,val<->'[1,1,1]' as dist FROM t ORDER BY val <-> '[1, 1, 1]';
 DROP TABLE t;