Add Postgres vector index support

Author: Milder Hernandez Cagua

api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/Hotel.java

@@ -17,6 +17,8 @@ public class Hotel {
     private final String description;
     @VectorStoreRecordVectorAttribute(dimensions = 3)
     private final List<Float> descriptionEmbedding;
+    @VectorStoreRecordVectorAttribute(dimensions = 3, indexKind = "hnsw", distanceFunction = "cosine")
+    private final List<Float> additionalEmbedding;
     @VectorStoreRecordDataAttribute
     private double rating;
 
@@ -30,6 +32,7 @@ public Hotel(String id, String name, int code, String description, List<Float> d
         this.code = code;
         this.description = description;
         this.descriptionEmbedding = descriptionEmbedding;
+        this.additionalEmbedding = descriptionEmbedding;
         this.rating = rating;
     }
 
@@ -52,6 +55,9 @@ public String getDescription() {
     public List<Float> getDescriptionEmbedding() {
         return descriptionEmbedding;
     }
+    public List<Float> getAdditionalEmbedding() {
+        return additionalEmbedding;
+    }
 
     public double getRating() {
         return rating;

semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreCollectionCreateMapping.java

@@ -34,7 +34,7 @@ private static VectorSearchAlgorithmMetric getAlgorithmMetric(
         }
 
         switch (vectorField.getDistanceFunction()) {
-            case COSINE_SIMILARITY:
+            case COSINE:
                 return VectorSearchAlgorithmMetric.COSINE;
             case DOT_PRODUCT:
                 return VectorSearchAlgorithmMetric.DOT_PRODUCT;

semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/postgres/PostgreSQLVectorDistanceFunction.java

@@ -0,0 +1,42 @@
+// Copyright (c) Microsoft. All rights reserved.
+package com.microsoft.semantickernel.connectors.data.postgres;
+
+import com.microsoft.semantickernel.data.recorddefinition.DistanceFunction;
+
+public enum PostgreSQLVectorDistanceFunction {
+    L2("vector_l2_ops", "<->"), COSINE("vector_cosine_ops", "<=>"), INNER_PRODUCT("vector_ip_ops",
+        "<#>");
+
+    private final String value;
+    private final String operator;
+
+    PostgreSQLVectorDistanceFunction(String value, String operator) {
+        this.value = value;
+        this.operator = operator;
+    }
+
+    public String getValue() {
+        return value;
+    }
+
+    public String getOperator() {
+        return operator;
+    }
+
+    public static PostgreSQLVectorDistanceFunction fromDistanceFunction(DistanceFunction function) {
+        if (function == null) {
+            return null;
+        }
+
+        switch (function) {
+            case EUCLIDEAN:
+                return L2;
+            case COSINE:
+                return COSINE;
+            case DOT_PRODUCT:
+                return INNER_PRODUCT;
+            default:
+                throw new IllegalArgumentException("Unsupported distance function: " + function);
+        }
+    }
+}

semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/postgres/PostgreSQLVectorIndexKind.java

@@ -0,0 +1,33 @@
+// Copyright (c) Microsoft. All rights reserved.
+package com.microsoft.semantickernel.connectors.data.postgres;
+
+import com.microsoft.semantickernel.data.recorddefinition.IndexKind;
+
+public enum PostgreSQLVectorIndexKind {
+    HNSW("hnsw"), IVFFLAT("ivfflat");
+
+    private final String value;
+
+    PostgreSQLVectorIndexKind(String value) {
+        this.value = value;
+    }
+
+    public String getValue() {
+        return value;
+    }
+
+    public static PostgreSQLVectorIndexKind fromIndexKind(IndexKind indexKind) {
+        if (indexKind == null) {
+            return null;
+        }
+
+        switch (indexKind) {
+            case HNSW:
+                return HNSW;
+            case FLAT:
+                return IVFFLAT;
+            default:
+                throw new IllegalArgumentException("Unsupported index kind: " + indexKind);
+        }
+    }
+}

semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/postgres/PostgreSQLVectorStoreQueryProvider.java

@@ -5,6 +5,8 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreDefaultQueryProvider;
 import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreQueryProvider;
+import com.microsoft.semantickernel.data.recorddefinition.DistanceFunction;
+import com.microsoft.semantickernel.data.recorddefinition.IndexKind;
 import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition;
 import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordField;
 import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordKeyField;
@@ -19,6 +21,7 @@
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Statement;
 import java.time.OffsetDateTime;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -155,6 +158,29 @@ private String getColumnNamesAndTypesForVectorFields(List<VectorStoreRecordVecto
         return columnNames.toString();
     }
 
+    private String createIndexForVectorField(String collectionName,
+        VectorStoreRecordVectorField vectorField) {
+        PostgreSQLVectorIndexKind indexKind = PostgreSQLVectorIndexKind
+            .fromIndexKind(vectorField.getIndexKind());
+        PostgreSQLVectorDistanceFunction distanceFunction = PostgreSQLVectorDistanceFunction
+            .fromDistanceFunction(vectorField.getDistanceFunction());
+
+        // If indexKind is not specified, no index is created
+        // and pgvector performs exact nearest neighbor search.
+        if (indexKind == null) {
+            return null;
+        }
+        if (distanceFunction == null) {
+            throw new SKException(
+                "Distance function is required for vector field: " + vectorField.getName());
+        }
+
+        return "CREATE INDEX IF NOT EXISTS " + getCollectionTableName(collectionName) + "_index"
+            + " ON " + getCollectionTableName(collectionName)
+            + " USING " + indexKind.getValue()
+            + " (" + vectorField.getName() + " " + distanceFunction.getValue() + ");";
+    }
+
     /**
      * Creates a collection.
      *
@@ -164,29 +190,44 @@ private String getColumnNamesAndTypesForVectorFields(List<VectorStoreRecordVecto
      * @throws SKException if an error occurs while creating the collection
      */
     @Override
-    @SuppressFBWarnings("SQL_PREPARED_STATEMENT_GENERATED_FROM_NONCONSTANT_STRING") // SQL query is generated dynamically with valid identifiers
+    @SuppressFBWarnings(value = {
+            "SQL_NONCONSTANT_STRING_PASSED_TO_EXECUTE",
+            "SQL_PREPARED_STATEMENT_GENERATED_FROM_NONCONSTANT_STRING"
+    }) // SQL query is generated dynamically with valid identifiers
     public void createCollection(String collectionName, Class<?> recordClass,
         VectorStoreRecordDefinition recordDefinition) {
         Field keyDeclaredField = recordDefinition.getKeyDeclaredField(recordClass);
         List<Field> dataDeclaredFields = recordDefinition.getDataDeclaredFields(recordClass);
 
-        String createStorageTable = "CREATE TABLE IF NOT EXISTS "
-            + getCollectionTableName(collectionName)
-            + " (" + keyDeclaredField.getName() + " VARCHAR(255) PRIMARY KEY, "
-            + getColumnNamesAndTypes(dataDeclaredFields, supportedDataTypes) + ", "
-            + getColumnNamesAndTypesForVectorFields(recordDefinition.getVectorFields(), recordClass)
-            + ");";
-
-        String insertCollectionQuery = "INSERT INTO " + validateSQLidentifier(collectionsTable)
-            + " (collectionId) VALUES (?)";
+        List<VectorStoreRecordVectorField> vectorFields = recordDefinition.getVectorFields();
 
         try (Connection connection = dataSource.getConnection();
-            PreparedStatement createTable = connection.prepareStatement(createStorageTable)) {
-            createTable.execute();
+            Statement createTableAndIndexes = connection.createStatement()) {
+
+            String createStorageTable = "CREATE TABLE IF NOT EXISTS "
+                + getCollectionTableName(collectionName)
+                + " (" + keyDeclaredField.getName() + " VARCHAR(255) PRIMARY KEY, "
+                + getColumnNamesAndTypes(dataDeclaredFields, supportedDataTypes) + ", "
+                + getColumnNamesAndTypesForVectorFields(vectorFields, recordClass)
+                + ");";
+
+            createTableAndIndexes.addBatch(createStorageTable);
+            for (VectorStoreRecordVectorField vectorField : vectorFields) {
+                String createVectorIndex = createIndexForVectorField(collectionName, vectorField);
+
+                if (createVectorIndex != null) {
+                    createTableAndIndexes.addBatch(createVectorIndex);
+                }
+            }
+
+            createTableAndIndexes.executeBatch();
         } catch (SQLException e) {
             throw new SKException("Failed to create collection", e);
         }
 
+        String insertCollectionQuery = "INSERT INTO " + validateSQLidentifier(collectionsTable)
+            + " (collectionId) VALUES (?)";
+
         try (Connection connection = dataSource.getConnection();
             PreparedStatement insert = connection.prepareStatement(insertCollectionQuery)) {
             insert.setObject(1, collectionName);

semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStoreCollectionCreateMapping.java

@@ -35,7 +35,7 @@ private static String getAlgorithmMetric(
         }
 
         switch (vectorField.getDistanceFunction()) {
-            case COSINE_SIMILARITY:
+            case COSINE:
                 return RedisVectorDistanceMetric.COSINE;
             case DOT_PRODUCT:
                 return RedisVectorDistanceMetric.DOT_PRODUCT;

semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/recorddefinition/DistanceFunction.java

@@ -2,7 +2,18 @@
 package com.microsoft.semantickernel.data.recorddefinition;
 
 public enum DistanceFunction {
-    COSINE_SIMILARITY("cosineSimilarity"), DOT_PRODUCT("dotProduct"), EUCLIDEAN("euclidean");
+    /**
+     * Cosine (angular) similarity function.
+     */
+    COSINE("cosine"),
+    /**
+     * Dot product between two vectors.
+     */
+    DOT_PRODUCT("dotProduct"),
+    /**
+     * Euclidean distance function. Also known as L2 norm.
+     */
+    EUCLIDEAN("euclidean");
 
     private final String value;
 
@@ -23,7 +34,7 @@ public String getValue() {
      */
     public static DistanceFunction fromString(String text) {
         if (text == null || text.isEmpty()) {
-            return COSINE_SIMILARITY;
+            return null;
         }
 
         for (DistanceFunction b : DistanceFunction.values()) {

semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/recorddefinition/IndexKind.java

@@ -23,7 +23,7 @@ public String getValue() {
      */
     public static IndexKind fromString(String text) {
         if (text == null || text.isEmpty()) {
-            return FLAT;
+            return null;
         }
 
         for (IndexKind b : IndexKind.values()) {