PG-Vector support in Babelfish (#2232)

With GenAI capturing the imagination and Vector Similarity search being one of the foundational features in this space, we MUST leverage what is being offered by pgvector and make it available for T-SQL applications. With this commit, users can install the vector extension in SYS schema and start using not only the datatype but also all of the other features available in pgvector extension [https://github.com/pgvector/pgvector]. NOTE: Except for Creation of IVFFLAT and HNSW indexes, rest all queries must follow TSQL syntax and semantics. TASK: BABEL-4687 Signed-off-by: Kushaal Shroff [email protected]
babelfish-for-postgresql · Jan 29, 2024 · 5d7b982 · 5d7b982
1 parent 1ea412a
commit 5d7b982
Show file tree

Hide file tree

Showing 29 changed files with 5,345 additions and 21 deletions.
diff --git a/.github/composite-actions/build-vector-extension/action.yml b/.github/composite-actions/build-vector-extension/action.yml
@@ -0,0 +1,22 @@
+name: 'Build vector Extension'
+
+inputs:
+  install_dir:
+    description: 'Engine install directory'
+    required: no
+    default: psql
+
+runs:
+  using: "composite"
+  steps:
+    - name: Build vector Extension
+      run: |
+        cd ..
+        export VECTOR_VERSION="0.5.1"
+        sudo apt-get install wget
+        wget https://github.com/pgvector/pgvector/archive/refs/tags/v${VECTOR_VERSION}.tar.gz
+        tar -xvzf v${VECTOR_VERSION}.tar.gz
+        cd pgvector-${VECTOR_VERSION}/
+        make USE_PGXS=1 PG_CONFIG=~/${{ inputs.install_dir }}/bin/pg_config
+        sudo make USE_PGXS=1 PG_CONFIG=~/${{ inputs.install_dir }}/bin/pg_config install
+      shell: bash
diff --git a/.github/composite-actions/setup-base-version/action.yml b/.github/composite-actions/setup-base-version/action.yml
@@ -63,6 +63,18 @@ runs:
         sudo make USE_PGXS=1 PG_CONFIG=~/${{ inputs.install_dir }}/bin/pg_config install
       shell: bash
 
+    - name: Build vector Extension
+      run: |
+        cd ..
+        export VECTOR_VERSION="0.5.1"
+        sudo apt-get install wget
+        wget https://github.com/pgvector/pgvector/archive/refs/tags/v${VECTOR_VERSION}.tar.gz 
+        tar -xvzf v${VECTOR_VERSION}.tar.gz
+        cd pgvector-${VECTOR_VERSION}/
+        make USE_PGXS=1 PG_CONFIG=~/${{ inputs.install_dir }}/bin/pg_config
+        sudo make USE_PGXS=1 PG_CONFIG=~/${{ inputs.install_dir }}/bin/pg_config install
+      shell: bash
+
     - name: Build PostGIS Extension
       run: |
         cd ..

diff --git a/.github/composite-actions/setup-new-version/action.yml b/.github/composite-actions/setup-new-version/action.yml
@@ -44,9 +44,16 @@ runs:
       with:
         install_dir: ${{ inputs.pg_new_dir }}
 
+    - name: Build vector Extension
+      id: build-vector-extension
+      if: always() && steps.build-tds_fdw-extension.outcome == 'success'
+      uses: ./.github/composite-actions/build-vector-extension
+      with:
+        install_dir: ${{ inputs.pg_new_dir }}
+
     - name: Build PostGIS Extension
       id: build-postgis-extension
-      if: always() && steps.build-tds_fdw-extension.outcome == 'success'
+      if: always() && steps.build-vector-extension.outcome == 'success'
       uses: ./.github/composite-actions/build-postgis-extension
       with:
         install_dir: ${{ inputs.pg_new_dir }}

diff --git a/.github/workflows/code-coverage.yml b/.github/workflows/code-coverage.yml
@@ -44,9 +44,14 @@ jobs:
         if: always() &&  steps.build-extensions.outcome == 'success'
         uses: ./.github/composite-actions/build-tds_fdw-extension
 
+      - name: Build vector Extension
+        id: build-vector-extension
+        if: always() &&  steps.build-tds_fdw-extension.outcome == 'success'
+        uses: ./.github/composite-actions/build-vector-extension
+
       - name: Build PostGIS Extension
         id: build-postgis-extension
-        if: always() &&  steps.build-tds_fdw-extension.outcome == 'success'
+        if: always() &&  steps.build-vector-extension.outcome == 'success'
         uses: ./.github/composite-actions/build-postgis-extension
 
       - name: Install Extensions

diff --git a/.github/workflows/jdbc-tests-with-non-default-server-collation.yml b/.github/workflows/jdbc-tests-with-non-default-server-collation.yml
@@ -37,9 +37,14 @@ jobs:
         if: always() && steps.build-extensions.outcome == 'success'
         uses: ./.github/composite-actions/build-tds_fdw-extension
 
+      - name: Build vector Extension
+        id: build-vector-extension
+        if: always() &&  steps.build-tds_fdw-extension.outcome == 'success'
+        uses: ./.github/composite-actions/build-vector-extension
+
       - name: Build PostGIS Extension
         id: build-postgis-extension
-        if: always() && steps.build-tds_fdw-extension.outcome == 'success'
+        if: always() && steps.build-vector-extension.outcome == 'success'
         uses: ./.github/composite-actions/build-postgis-extension
 
       - name: Install Extensions

diff --git a/.github/workflows/jdbc-tests-with-parallel-query.yml b/.github/workflows/jdbc-tests-with-parallel-query.yml
@@ -35,9 +35,14 @@ jobs:
         if: always() && steps.build-extensions.outcome == 'success'
         uses: ./.github/composite-actions/build-tds_fdw-extension
 
+      - name: Build vector Extension
+        id: build-vector-extension
+        if: always() &&  steps.build-tds_fdw-extension.outcome == 'success'
+        uses: ./.github/composite-actions/build-vector-extension
+
       - name: Build PostGIS Extension
         id: build-postgis-extension
-        if: always() && steps.build-tds_fdw-extension.outcome == 'success'
+        if: always() &&  steps.build-vector-extension.outcome == 'success'
         uses: ./.github/composite-actions/build-postgis-extension
 
       - name: Install Extensions

diff --git a/.github/workflows/jdbc-tests.yml b/.github/workflows/jdbc-tests.yml
@@ -38,9 +38,14 @@ jobs:
         if: always() && steps.build-extensions.outcome == 'success'
         uses: ./.github/composite-actions/build-tds_fdw-extension
 
+      - name: Build vector Extension
+        id: build-vector-extension
+        if: always() &&  steps.build-tds_fdw-extension.outcome == 'success'
+        uses: ./.github/composite-actions/build-vector-extension
+
       - name: Build PostGIS Extension
         id: build-postgis-extension
-        if: always() && steps.build-tds_fdw-extension.outcome == 'success'
+        if: always() &&  steps.build-vector-extension.outcome == 'success'
         uses: ./.github/composite-actions/build-postgis-extension
 
       - name: Install Extensions

diff --git a/.github/workflows/major-version-upgrade.yml b/.github/workflows/major-version-upgrade.yml
@@ -154,9 +154,16 @@ jobs:
         with:
           install_dir: ${{env.NEW_INSTALL_DIR}}
 
+      - name: Build vector Extension
+        id: build-vector-extension
+        if: always() && steps.build-tds_fdw-extension.outcome == 'success'
+        uses: ./.github/composite-actions/build-vector-extension
+        with:
+          install_dir: ${{env.NEW_INSTALL_DIR}}
+
       - name: Build PostGIS Extension
         id: build-postgis-extension
-        if: always() && steps.build-tds_fdw-extension.outcome == 'success'
+        if: always() && steps.build-vector-extension.outcome == 'success'
         uses: ./.github/composite-actions/build-postgis-extension
         with:
           install_dir: ${{env.NEW_INSTALL_DIR}}

diff --git a/.github/workflows/minor-version-upgrade.yml b/.github/workflows/minor-version-upgrade.yml
@@ -92,6 +92,18 @@ jobs:
           sudo make USE_PGXS=1 PG_CONFIG=~/psql/bin/pg_config install
         shell: bash
 
+      - name: Build vector Extension
+        run: |
+          cd ..
+          export VECTOR_VERSION="0.5.1"
+          sudo apt-get install wget
+          wget https://github.com/pgvector/pgvector/archive/refs/tags/v${VECTOR_VERSION}.tar.gz
+          tar -xvzf v${VECTOR_VERSION}.tar.gz
+          cd pgvector-${VECTOR_VERSION}/
+          make USE_PGXS=1 PG_CONFIG=~/psql/bin/pg_config
+          sudo make USE_PGXS=1 PG_CONFIG=~/psql/bin/pg_config install
+        shell: bash
+
       - name: Install extensions
         id: install-extensions-older
         if: always() && steps.build-tds_fdw-extension.outcome == 'success'

diff --git a/contrib/babelfishpg_common/src/typecode.c b/contrib/babelfishpg_common/src/typecode.c
@@ -52,6 +52,7 @@ type_info_t type_infos[TOTAL_TYPECODE_COUNT] =
 	{0, 1, "sysname", "sysname", 5, 31, 5},
 	{0, 1, "rowversion", "timestamp", 8, 32, 3},
 	{0, 1, "timestamp", "timestamp", 8, 33, 3},
+	{0, 1, "vector", "vector", 9, 34, 3},
 	/*
 	 * Geospatial types cannot be stored in SQL variant so setting sqlvariant header size to 1
 	 */

diff --git a/contrib/babelfishpg_common/src/typecode.h b/contrib/babelfishpg_common/src/typecode.h
@@ -45,7 +45,7 @@
 #define FIXEDDECIMAL_MULTIPLIER 10000LL
 #endif
 
-#define TOTAL_TYPECODE_COUNT 35
+#define TOTAL_TYPECODE_COUNT 36
 
 struct Node;
 

diff --git a/contrib/babelfishpg_tds/babelfishpg_tds--1.0.0.sql b/contrib/babelfishpg_tds/babelfishpg_tds--1.0.0.sql
@@ -47,4 +47,4 @@ LANGUAGE SQL;
 CREATE FUNCTION sys.disable_injected_fault_all()
 RETURNS text
 AS $$ SELECT sys.inject_fault('all', 0) $$
-LANGUAGE SQL;
+LANGUAGE SQL;
diff --git a/contrib/babelfishpg_tds/src/backend/tds/tds_data_map.c b/contrib/babelfishpg_tds/src/backend/tds/tds_data_map.c
@@ -188,6 +188,7 @@ TdsIoFunctionRawData TdsIoFunctionRawData_data[] =
 	/* Mapping TDS listener sender to basic Postgres datatypes. */
 	{"pg_catalog", "oid", TDS_TYPE_INTEGER, 4, 1, TDS_SEND_INTEGER, TDS_RECV_INVALID},
 	{"pg_catalog", "sql_identifier", TDS_TYPE_VARCHAR, -1, 1, TDS_SEND_VARCHAR, TDS_RECV_INVALID},
+	{"sys", "vector", TDS_TYPE_VARCHAR, -1, 2, TDS_SEND_VARCHAR, TDS_RECV_VARCHAR},
 	{"pg_catalog", "name", TDS_TYPE_VARCHAR, -1, 1, TDS_SEND_VARCHAR, TDS_RECV_INVALID},
 	{"pg_catalog", "character_data", TDS_TYPE_VARCHAR, -1, 2, TDS_SEND_VARCHAR, TDS_RECV_INVALID},
 	{"pg_catalog", "bool", TDS_TYPE_BIT, 1, 1, TDS_SEND_BIT, TDS_RECV_INVALID},
@@ -221,7 +222,9 @@ TdsIoFunctionRawData TdsIoFunctionRawData_data[] =
 	{"pg_catalog", "tid", TDS_TYPE_VARCHAR, -1, 2, TDS_SEND_VARCHAR, TDS_RECV_INVALID},
 	{"pg_catalog", "inet", TDS_TYPE_VARCHAR, -1, 2, TDS_SEND_VARCHAR, TDS_RECV_INVALID},
 	{"pg_catalog", "interval", TDS_TYPE_VARCHAR, -1, 2, TDS_SEND_VARCHAR, TDS_RECV_INVALID},
-	{"pg_catalog", "bytea", TDS_TYPE_VARBINARY, -1, 2, TDS_SEND_VARBINARY, TDS_RECV_INVALID}
+	{"pg_catalog", "bytea", TDS_TYPE_VARBINARY, -1, 2, TDS_SEND_VARBINARY, TDS_RECV_INVALID},
+	{"pg_catalog", "record", TDS_TYPE_VARCHAR, -1, 2, TDS_SEND_VARCHAR, TDS_RECV_INVALID},
+	{"pg_catalog", "_record", TDS_TYPE_VARCHAR, -1, 2, TDS_SEND_VARCHAR, TDS_RECV_INVALID}
 };
 
 size_t		TdsIoFunctionRawData_datasize = lengthof(TdsIoFunctionRawData_data);
diff --git a/contrib/babelfishpg_tsql/antlr/TSqlLexer.g4 b/contrib/babelfishpg_tsql/antlr/TSqlLexer.g4
@@ -441,6 +441,7 @@ HIGH:                                            H I G H;
 HINT:                                            H I N T;
 HISTORY_RETENTION_PERIOD:                        H I S T O R Y  UNDERLINE  R E T E N T I O N  UNDERLINE  P E R I O D;
 HISTORY_TABLE:                                   H I S T O R Y  UNDERLINE  T A B L E;
+HNSW:                                            H N S W;
 HOLDLOCK:                                        H O L D L O C K;
 HONOR_BROKER_PRIORITY:                           H O N O R  UNDERLINE  B R O K E R  UNDERLINE  P R I O R I T Y;
 HOUR:                                            H O U R;
@@ -484,6 +485,7 @@ IS:                                              I S;
 ISDESCENDANTOF:                                  I S D E S C E N D A N T O F;
 ISNULL:                                          I S N U L L;
 ISOLATION:                                       I S O L A T I O N;
+IVFFLAT:                                         I V F F L A T;
 JOB:                                             J O B;
 JOIN:                                            J O I N;
 JSON:                                            J S O N;
@@ -1070,6 +1072,9 @@ VAR:                                             V A R;
 VARBINARY_KEYWORD:                               V A R B I N A R Y;
 VARP:                                            V A R P;
 VARYING:                                         V A R Y I N G;
+VECTOR_COSINE_OPS:                               V E C T O R UNDERLINE C O S I N E UNDERLINE O P S;
+VECTOR_IP_OPS:                                   V E C T O R UNDERLINE I P UNDERLINE O P S;
+VECTOR_L2_OPS:                                   V E C T O R UNDERLINE L '2' UNDERLINE O P S;
 VERBOSELOGGING:                                  V E R B O S E L O G G I N G;
 VERSION:                                         V E R S I O N;
 VIEW:                                            V I E W;
@@ -1173,6 +1178,9 @@ MOD_ASSIGN:          '%=';
 AND_ASSIGN:          '&=';
 XOR_ASSIGN:          '^=';
 OR_ASSIGN:           '|=';
+VECTOR_L2:           '<->';
+VECTOR_COSINE:       '<=>';
+VECTOR_IP:           '<#>';
 
 DOT:                 '.';
 UNDERLINE:           '_';

diff --git a/contrib/babelfishpg_tsql/antlr/TSqlParser.g4 b/contrib/babelfishpg_tsql/antlr/TSqlParser.g4
@@ -1822,9 +1822,21 @@ create_index
     (WHERE where=search_condition)?
     with_index_options?
     (ON storage_partition_clause)?
+ SEMI?
+    |CREATE UNIQUE? clustered? COLUMNSTORE? INDEX id ON table_name (USING vector_index_method)? (LR_BRACKET column_name_list_with_order_for_vector RR_BRACKET)?
+    (INCLUDE LR_BRACKET column_name_list RR_BRACKET )?
+    (WHERE where=search_condition)?
+    with_index_options?
+    (ON storage_partition_clause)?
  SEMI?
     ;
 
+/* We introduce specific index methods so as to avoid PG syntax leaks. */
+vector_index_method
+    : HNSW
+    | IVFFLAT
+    ;
+
 alter_index
     : ALTER INDEX (id | ALL) ON table_name alter_index_options
  SEMI?
@@ -3574,8 +3586,9 @@ xml_common_directives
     : COMMA (BINARY_KEYWORD BASE64 | TYPE | ROOT ( LR_BRACKET char_string RR_BRACKET )?)
     ;
 
-order_by_expression
+order_by_expression 
     : order_by=expression (ascending=ASC | descending=DESC)?
+    | order_by=expression vector_operator expression (ascending=ASC | descending=DESC)? 
     ;
 
 group_by_item
@@ -5187,6 +5200,11 @@ column_name_list_with_order
     : simple_column_name (ASC | DESC)? (COMMA simple_column_name (ASC | DESC)?)*
     ;
 
+/* We introduce specific index methods so as to avoid PG syntax leaks. */
+column_name_list_with_order_for_vector
+    : simple_column_name (ASC | DESC)? (VECTOR_COSINE_OPS | VECTOR_IP_OPS | VECTOR_L2_OPS)? (COMMA simple_column_name (ASC | DESC)? (VECTOR_COSINE_OPS | VECTOR_IP_OPS | VECTOR_L2_OPS)?)*
+    ;
+
 //For some reason, sql server allows any number of prefixes:  Here, h is the column: a.b.c.d.e.f.g.h
 insert_column_name_list
     : col+=insert_column_id (COMMA col+=insert_column_id)*
@@ -5227,7 +5245,11 @@ local_id
 // https://msdn.microsoft.com/en-us/library/ms188074.aspx
 // Spaces are allowed for comparison operators.
 comparison_operator
-    : EQUAL | GREATER | LESS | LESS EQUAL | GREATER EQUAL | LESS GREATER | EXCLAMATION EQUAL | EXCLAMATION GREATER | EXCLAMATION LESS | MULT_ASSIGN | EQUAL_STAR_OJ
+    : EQUAL | GREATER | LESS | LESS EQUAL | GREATER EQUAL | LESS GREATER | EXCLAMATION EQUAL | EXCLAMATION GREATER | EXCLAMATION LESS | MULT_ASSIGN | EQUAL_STAR_OJ | vector_operator
+    ;
+
+vector_operator
+    : VECTOR_COSINE | VECTOR_IP | VECTOR_L2
     ;
 
 assignment_operator

diff --git a/contrib/babelfishpg_tsql/sql/babelfishpg_tsql.sql b/contrib/babelfishpg_tsql/sql/babelfishpg_tsql.sql
@@ -639,6 +639,7 @@ BEGIN
 		WHEN 'sql_variant' THEN tds_id = 98;
 		WHEN 'datetimeoffset' THEN tds_id = 43;
 		WHEN 'timestamp' THEN tds_id = 173;
+		WHEN 'vector' THEN tds_id = 167; -- Same as varchar 
 		WHEN 'geometry' THEN tds_id = 240;
 		WHEN 'geography' THEN tds_id = 240;
 		ELSE tds_id = 0;

diff --git a/contrib/babelfishpg_tsql/sql/information_schema_tsql.sql b/contrib/babelfishpg_tsql/sql/information_schema_tsql.sql
@@ -57,7 +57,7 @@ $$SELECT
 		THEN 1073741823
 		WHEN type = 'sysname'
 		THEN 128
-		WHEN type IN ('xml', 'geometry', 'geography')
+		WHEN type IN ('xml', 'vector', 'geometry', 'geography')
 		THEN -1
 		WHEN type = 'sql_variant'
 		THEN 0

diff --git a/contrib/babelfishpg_tsql/sql/sys_views.sql b/contrib/babelfishpg_tsql/sql/sys_views.sql
@@ -353,6 +353,7 @@ BEGIN
 			ELSIF typemod <= 7 THEN max_length = 5;
 			END IF;
 		WHEN 'timestamp' THEN max_length = 8;
+		WHEN 'vector' THEN max_length = -1; -- dummy as varchar max
 		ELSE max_length = typelen;
 		END CASE;
 		RETURN max_length;

diff --git a/contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--3.4.0--3.5.0.sql b/contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--3.4.0--3.5.0.sql
@@ -2211,6 +2211,7 @@ AS SELECT
 WHERE FALSE;
 GRANT SELECT ON sys.availability_groups TO PUBLIC;
 
+
 -- BABELFISH_SCHEMA_PERMISSIONS
 CREATE TABLE IF NOT EXISTS sys.babelfish_schema_permissions (
   dbid smallint NOT NULL,
@@ -2703,6 +2704,7 @@ BEGIN
 			ELSIF typemod <= 7 THEN max_length = 5;
 			END IF;
 		WHEN 'timestamp' THEN max_length = 8;
+    WHEN 'vector' THEN max_length = -1; -- dummy as varchar max
 		ELSE max_length = typelen;
 		END CASE;
 		RETURN max_length;
@@ -2761,7 +2763,7 @@ $$SELECT
 		THEN 1073741823
 		WHEN type = 'sysname'
 		THEN 128
-		WHEN type IN ('xml', 'geometry', 'geography')
+		WHEN type IN ('xml', 'vector', 'geometry', 'geography')
 		THEN -1
 		WHEN type = 'sql_variant'
 		THEN 0
@@ -2810,6 +2812,7 @@ BEGIN
 		WHEN 'sql_variant' THEN tds_id = 98;
 		WHEN 'datetimeoffset' THEN tds_id = 43;
 		WHEN 'timestamp' THEN tds_id = 173;
+		WHEN 'vector' THEN tds_id = 167; -- Same as varchar 
 		WHEN 'geometry' THEN tds_id = 240;
 		WHEN 'geography' THEN tds_id = 240;
 		ELSE tds_id = 0;

diff --git a/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-epilogue.y.c b/contrib/babelfishpg_tsql/src/backend_parser/gram-tsql-epilogue.y.c
@@ -1886,8 +1886,8 @@ tsql_index_nulls_order(List *indexParams, const char *accessMethod)
 		if (indexElem->nulls_ordering != SORTBY_NULLS_DEFAULT)
 			continue;
 
-		/* GIN indexes don't support NULLS FIRST/LAST options */
-		if (strcmp(accessMethod, "gin") == 0)
+		/* GIN, HNSW and IVFFLAT indexes don't support NULLS FIRST/LAST options */
+		if (strcmp(accessMethod, "gin") == 0 || strcmp(accessMethod, "hnsw") == 0 || strcmp(accessMethod, "ivfflat") == 0)
 			return;
 
 		switch (indexElem->ordering)

diff --git a/contrib/babelfishpg_tsql/src/datatype_info.h b/contrib/babelfishpg_tsql/src/datatype_info.h
@@ -4,7 +4,7 @@
 #define NULLVAL		PG_INT32_MIN
 #define NULLVAL_STR	"NULL"
 
-#define DATATYPE_INFO_TABLE_ROWS 39
+#define DATATYPE_INFO_TABLE_ROWS 40
 
 typedef struct DatatypeInfo
 {
@@ -488,6 +488,19 @@ static const DatatypeInfo datatype_info_table[DATATYPE_INFO_TABLE_ROWS] = {
 		0, 0, 9, 3, NULLVAL, NULLVAL, 22, 16, 111,
 		"smalldatetime"
 	},
+	{
+		/* Dummy Entry for vector datatype since its not a TSQL datatype. */
+		"vector",
+		NULLVAL, NULLVAL, NULLVAL, NULLVAL,
+		0,
+		"'",
+		"'",
+		NULLVAL_STR,
+		1, 0, 0, NULLVAL, 0, NULLVAL,
+		"vector",
+		NULLVAL, NULLVAL, NULLVAL, NULLVAL, NULLVAL, NULLVAL, NULLVAL, NULLVAL, NULLVAL,
+		"vector"
+    },
 	{
 		"geometry",
 		-4, -4, -151, -151,