Merge pull request #178 from milderhc/postgres-vector-index

Add Postgres vector index support

Author: Milder Hernandez

.github/_typos.toml

@@ -29,6 +29,7 @@ ans = "ans"               # Short for answers
 arange = "arange"         # Method in Python numpy package
 prompty = "prompty"       # prompty is a format name.
 ist = "ist"               # German language
+Prelease = "Prelease"     # Prelease is a format name.
 
 [default.extend-identifiers]
 ags = "ags" # Azure Graph Service

.github/workflows/java-build.yml

@@ -56,7 +56,7 @@ jobs:
         run: ./mvnw -B -Pbug-check -Pcompile-jdk${{ matrix.java-versions }} test --file pom.xml
 
       # Uploads test artifacts for each JDK version
-      - uses: actions/upload-artifact@v2
+      - uses: actions/upload-artifact@v4
         if: always()
         with:
           name: test_output_sk_jdk${{ matrix.java-versions }}u

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

@@ -21,6 +21,8 @@ public class Hotel {
     @JsonProperty("summaryEmbedding")
     @VectorStoreRecordVectorAttribute(dimensions = 3)
     private final List<Float> descriptionEmbedding;
+    @VectorStoreRecordVectorAttribute(dimensions = 3, indexKind = "hnsw", distanceFunction = "cosine")
+    private final List<Float> additionalEmbedding;
     @VectorStoreRecordDataAttribute
     private double rating;
 
@@ -41,6 +43,7 @@ public Hotel(
         this.code = code;
         this.description = description;
         this.descriptionEmbedding = descriptionEmbedding;
+        this.additionalEmbedding = descriptionEmbedding;
         this.rating = rating;
     }
 
@@ -63,6 +66,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

@@ -36,7 +36,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

@@ -19,6 +19,7 @@
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.SQLException;
+import java.sql.Statement;
 import java.time.OffsetDateTime;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -150,6 +151,29 @@ private String getColumnNamesAndTypesForVectorFields(
             .collect(Collectors.joining(", "));
     }
 
+    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.
      *
@@ -158,29 +182,44 @@ private String getColumnNamesAndTypesForVectorFields(
      * @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,
         VectorStoreRecordDefinition recordDefinition) {
 
-        String createStorageTable = "CREATE TABLE IF NOT EXISTS "
-            + getCollectionTableName(collectionName) + " ("
-            + getKeyColumnName(recordDefinition.getKeyField()) + " VARCHAR(255) PRIMARY KEY, "
-            + getColumnNamesAndTypes(new ArrayList<>(recordDefinition.getDataFields()),
-                supportedDataTypes)
-            + ", "
-            + getColumnNamesAndTypesForVectorFields(recordDefinition.getVectorFields())
-            + ");";
-
-        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) + " ("
+                + getKeyColumnName(recordDefinition.getKeyField()) + " VARCHAR(255) PRIMARY KEY, "
+                + getColumnNamesAndTypes(new ArrayList<>(recordDefinition.getDataFields()),
+                    supportedDataTypes)
+                + ", "
+                + getColumnNamesAndTypesForVectorFields(recordDefinition.getVectorFields())
+                + ");";
+
+            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

@@ -37,7 +37,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()) {