SchemaObject changes to support multiple vector configs

src/main/java/io/stargate/sgv2/jsonapi/api/v1/CollectionResource.java

@@ -198,15 +198,15 @@ public Uni<RestResponse<CommandResult>> postCommand(
                 // TODO: refactor this code to be cleaner so it assigns on one line
                 EmbeddingProvider embeddingProvider = null;
                 final VectorConfig.VectorizeConfig vectorizeConfig =
-                    schemaObject.vectorConfig().vectorizeConfig();
+                    schemaObject.vectorConfigs().get(0).vectorizeConfig();
                 if (vectorizeConfig != null) {
                   embeddingProvider =
                       embeddingProviderFactory.getConfiguration(
                           dataApiRequestInfo.getTenantId(),
                           dataApiRequestInfo.getCassandraToken(),
                           vectorizeConfig.provider(),
                           vectorizeConfig.modelName(),
-                          schemaObject.vectorConfig().vectorSize(),
+                          schemaObject.vectorConfigs().get(0).vectorSize(),
                           vectorizeConfig.parameters(),
                           vectorizeConfig.authentication(),
                           command.getClass().getSimpleName());

src/main/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/DatabaseSchemaObject.java

@@ -1,5 +1,7 @@
 package io.stargate.sgv2.jsonapi.service.cqldriver.executor;
 
+import java.util.List;
+
 public class DatabaseSchemaObject extends SchemaObject {
 
   public static final SchemaObjectType TYPE = SchemaObjectType.DATABASE;
@@ -9,8 +11,8 @@ public DatabaseSchemaObject() {
   }
 
   @Override
-  public VectorConfig vectorConfig() {
-    return VectorConfig.notEnabledVectorConfig();
+  public List<VectorConfig> vectorConfigs() {
+    return List.of(VectorConfig.notEnabledVectorConfig());
   }
 
   @Override

src/main/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/KeyspaceSchemaObject.java

@@ -1,6 +1,7 @@
 package io.stargate.sgv2.jsonapi.service.cqldriver.executor;
 
 import io.stargate.sgv2.jsonapi.service.schema.collections.CollectionSchemaObject;
+import java.util.List;
 
 public class KeyspaceSchemaObject extends SchemaObject {
 
@@ -39,8 +40,8 @@ public static KeyspaceSchemaObject fromSchemaObject(TableSchemaObject table) {
   }
 
   @Override
-  public VectorConfig vectorConfig() {
-    return VectorConfig.notEnabledVectorConfig();
+  public List<VectorConfig> vectorConfigs() {
+    return List.of(VectorConfig.notEnabledVectorConfig());
   }
 
   @Override

src/main/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/NamespaceCache.java

@@ -104,7 +104,7 @@ private Uni<SchemaObject> loadSchemaObject(
               }
 
               // 04-Sep-2024, tatu: Used to check that API Tables enabled; no longer checked here
-              return new TableSchemaObject(table);
+              return TableSchemaObject.getTableSettings(table, objectMapper);
             });
   }
 

src/main/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/SchemaObject.java

@@ -1,5 +1,7 @@
 package io.stargate.sgv2.jsonapi.service.cqldriver.executor;
 
+import java.util.List;
+
 /** A Collection or Table the command works on */
 public abstract class SchemaObject {
 
@@ -29,12 +31,13 @@ public SchemaObjectName name() {
   }
 
   /**
-   * Subclasses must always return an instance of VectorConfig, if there is no vector config they
-   * should return VectorConfig.notEnabledVectorConfig()
+   * Subclasses must always return List of VectorConfig, if there is no vector config they should
+   * return VectorConfig.notEnabledVectorConfig(). This needs to be a list to support tables with
+   * multiple vector columns.
    *
    * @return
    */
-  public abstract VectorConfig vectorConfig();
+  public abstract List<VectorConfig> vectorConfigs();
 
   /**
    * Call to get an instance of the appropriate {@link IndexUsage} for this schema object

src/main/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/TableSchemaObject.java

@@ -1,22 +1,98 @@
 package io.stargate.sgv2.jsonapi.service.cqldriver.executor;
 
+import com.datastax.oss.driver.api.core.CqlIdentifier;
+import com.datastax.oss.driver.api.core.data.ByteUtils;
+import com.datastax.oss.driver.api.core.metadata.schema.ColumnMetadata;
+import com.datastax.oss.driver.api.core.metadata.schema.IndexMetadata;
 import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata;
+import com.datastax.oss.driver.api.core.type.VectorType;
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import io.stargate.sgv2.jsonapi.service.schema.SimilarityFunction;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 
 public class TableSchemaObject extends TableBasedSchemaObject {
 
   public static final SchemaObjectType TYPE = SchemaObjectType.TABLE;
 
-  public TableSchemaObject(TableMetadata tableMetadata) {
+  private final List<VectorConfig> vectorConfigs;
+
+  private TableSchemaObject(TableMetadata tableMetadata, List<VectorConfig> vectorConfigs) {
     super(TYPE, tableMetadata);
+    this.vectorConfigs = vectorConfigs;
   }
 
   @Override
-  public VectorConfig vectorConfig() {
-    return VectorConfig.notEnabledVectorConfig();
+  public List<VectorConfig> vectorConfigs() {
+    return vectorConfigs;
   }
 
   @Override
   public IndexUsage newIndexUsage() {
     return IndexUsage.NO_OP;
   }
+
+  /**
+   * Get table schema object from table metadata
+   *
+   * @param tableMetadata
+   * @param objectMapper
+   * @return
+   */
+  public static TableSchemaObject getTableSettings(
+      TableMetadata tableMetadata, ObjectMapper objectMapper) {
+    Map<String, String> extensions =
+        (Map<String, String>)
+            tableMetadata.getOptions().get(CqlIdentifier.fromInternal("extensions"));
+    String vectorize = extensions != null ? extensions.get("vectorize") : null;
+    Map<String, VectorConfig.VectorizeConfig> resultMap = new HashMap<>();
+    if (vectorize != null) {
+      String vectorizeJson = new String(ByteUtils.fromHexString(vectorize).array());
+      // Define the TypeReference for Map<String, VectorConfig.VectorizeConfig>
+      TypeReference<Map<String, VectorConfig.VectorizeConfig>> typeRef =
+          new TypeReference<Map<String, VectorConfig.VectorizeConfig>>() {};
+
+      // Convert JSON string to Map
+      try {
+        resultMap = objectMapper.readValue(vectorizeJson, typeRef);
+      } catch (JsonProcessingException e) {
+        throw new RuntimeException(e);
+      }
+    }
+
+    List<VectorConfig> vectorConfigs = new ArrayList<>();
+    for (Map.Entry<CqlIdentifier, ColumnMetadata> column : tableMetadata.getColumns().entrySet()) {
+      if (column.getValue().getType() instanceof VectorType vectorType) {
+        final Optional<IndexMetadata> index = tableMetadata.getIndex(column.getKey());
+        SimilarityFunction similarityFunction = SimilarityFunction.COSINE;
+        if (index.isPresent()) {
+          final IndexMetadata indexMetadata = index.get();
+          final Map<String, String> indexOptions = indexMetadata.getOptions();
+          final String similarityFunctionValue = indexOptions.get("similarity_function");
+          if (similarityFunctionValue != null) {
+            similarityFunction = SimilarityFunction.fromString(similarityFunctionValue);
+          }
+        }
+        int dimension = vectorType.getDimensions();
+        VectorConfig vectorConfig =
+            new VectorConfig(
+                true,
+                column.getKey().asInternal(),
+                dimension,
+                similarityFunction,
+                resultMap.get(column.getKey().asInternal()));
+        vectorConfigs.add(vectorConfig);
+      }
+    }
+    if (vectorConfigs.isEmpty()) {
+      vectorConfigs.add(VectorConfig.notEnabledVectorConfig());
+    }
+    return new TableSchemaObject(tableMetadata, Collections.unmodifiableList(vectorConfigs));
+  }
 }

src/main/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/VectorConfig.java

@@ -2,6 +2,7 @@
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import io.stargate.sgv2.jsonapi.config.constants.DocumentConstants;
 import io.stargate.sgv2.jsonapi.service.schema.SimilarityFunction;
 import java.util.Map;
 
@@ -15,27 +16,56 @@
  */
 public record VectorConfig(
     boolean vectorEnabled,
+    String fieldName,
     int vectorSize,
     SimilarityFunction similarityFunction,
     VectorizeConfig vectorizeConfig) {
 
   // TODO: this is an immutable record, this can be singleton
   // TODO: Remove the use of NULL for the objects like vectorizeConfig
   public static VectorConfig notEnabledVectorConfig() {
-    return new VectorConfig(false, -1, null, null);
+    return new VectorConfig(false, null, -1, null, null);
   }
 
-  // convert a vector jsonNode from table comment to vectorConfig
+  // convert a vector jsonNode from table comment to vectorConfig, used for collection
   public static VectorConfig fromJson(JsonNode jsonNode, ObjectMapper objectMapper) {
     // dimension, similarityFunction, must exist
     int dimension = jsonNode.get("dimension").asInt();
     SimilarityFunction similarityFunction =
         SimilarityFunction.fromString(jsonNode.get("metric").asText());
 
+    return fromJson(
+        DocumentConstants.Fields.VECTOR_EMBEDDING_TEXT_FIELD,
+        dimension,
+        similarityFunction,
+        jsonNode,
+        objectMapper);
+  }
+
+  // convert a vector jsonNode from table extension to vectorConfig, used for tables
+  public static VectorConfig fromJson(
+      String fieldName,
+      int dimension,
+      SimilarityFunction similarityFunction,
+      JsonNode jsonNode,
+      ObjectMapper objectMapper) {
     VectorizeConfig vectorizeConfig = null;
     // construct vectorizeConfig
     JsonNode vectorizeServiceNode = jsonNode.get("service");
     if (vectorizeServiceNode != null) {
+      vectorizeConfig = VectorizeConfig.fromJson(vectorizeServiceNode, objectMapper);
+    }
+    return new VectorConfig(true, fieldName, dimension, similarityFunction, vectorizeConfig);
+  }
+
+  public record VectorizeConfig(
+      String provider,
+      String modelName,
+      Map<String, String> authentication,
+      Map<String, Object> parameters) {
+
+    protected static VectorizeConfig fromJson(
+        JsonNode vectorizeServiceNode, ObjectMapper objectMapper) {
       // provider, modelName, must exist
       String provider = vectorizeServiceNode.get("provider").asText();
       String modelName = vectorizeServiceNode.get("modelName").asText();
@@ -51,17 +81,8 @@ public static VectorConfig fromJson(JsonNode jsonNode, ObjectMapper objectMapper
           vectorizeServiceParameterNode == null
               ? null
               : objectMapper.convertValue(vectorizeServiceParameterNode, Map.class);
-      vectorizeConfig =
-          new VectorizeConfig(
-              provider, modelName, vectorizeServiceAuthentication, vectorizeServiceParameter);
+      return new VectorizeConfig(
+          provider, modelName, vectorizeServiceAuthentication, vectorizeServiceParameter);
     }
-
-    return new VectorConfig(true, dimension, similarityFunction, vectorizeConfig);
   }
-
-  public record VectorizeConfig(
-      String provider,
-      String modelName,
-      Map<String, String> authentication,
-      Map<String, Object> parameters) {}
 }

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/DataVectorizer.java

@@ -14,6 +14,7 @@
 import io.stargate.sgv2.jsonapi.exception.ErrorCodeV1;
 import io.stargate.sgv2.jsonapi.exception.JsonApiException;
 import io.stargate.sgv2.jsonapi.service.cqldriver.executor.SchemaObject;
+import io.stargate.sgv2.jsonapi.service.cqldriver.executor.VectorConfig;
 import io.stargate.sgv2.jsonapi.service.embedding.operation.EmbeddingProvider;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -30,6 +31,7 @@ public class DataVectorizer {
   private final JsonNodeFactory nodeFactory;
   private final EmbeddingCredentials embeddingCredentials;
   private final SchemaObject schemaObject;
+  private final VectorConfig vectorConfig;
 
   /**
    * Constructor
@@ -49,6 +51,9 @@ public DataVectorizer(
     this.nodeFactory = nodeFactory;
     this.embeddingCredentials = embeddingCredentials;
     this.schemaObject = schemaObject;
+    // This is getting element at 0 since only one vector is stored in the schema.
+    // This logic needs to be changed to handle multiple vectors columns for tables,
+    vectorConfig = schemaObject.vectorConfigs().get(0);
   }
 
   /**
@@ -114,7 +119,7 @@ public Uni<Boolean> vectorize(List<JsonNode> documents) {
                   if (vectorData.size() != vectorizeTexts.size()) {
                     throw EMBEDDING_PROVIDER_UNEXPECTED_RESPONSE.toApiException(
                         "Embedding provider '%s' didn't return the expected number of embeddings. Expect: '%d'. Actual: '%d'",
-                        schemaObject.vectorConfig().vectorizeConfig().provider(),
+                        vectorConfig.vectorizeConfig().provider(),
                         vectorizeTexts.size(),
                         vectorData.size());
                   }
@@ -125,11 +130,11 @@ public Uni<Boolean> vectorize(List<JsonNode> documents) {
                     JsonNode document = documents.get(position);
                     float[] vector = vectorData.get(vectorPosition);
                     // check if all vectors have the expected size
-                    if (vector.length != schemaObject.vectorConfig().vectorSize()) {
+                    if (vector.length != vectorConfig.vectorSize()) {
                       throw EMBEDDING_PROVIDER_UNEXPECTED_RESPONSE.toApiException(
                           "Embedding provider '%s' did not return expected embedding length. Expect: '%d'. Actual: '%d'",
-                          schemaObject.vectorConfig().vectorizeConfig().provider(),
-                          schemaObject.vectorConfig().vectorSize(),
+                          vectorConfig.vectorizeConfig().provider(),
+                          vectorConfig.vectorSize(),
                           vector.length);
                     }
                     final ArrayNode arrayNode = nodeFactory.arrayNode(vector.length);
@@ -170,11 +175,11 @@ public Uni<float[]> vectorize(String vectorizeContent) {
             vectorData -> {
               float[] vector = vectorData.get(0);
               // check if vector have the expected size
-              if (vector.length != schemaObject.vectorConfig().vectorSize()) {
+              if (vector.length != vectorConfig.vectorSize()) {
                 throw EMBEDDING_PROVIDER_UNEXPECTED_RESPONSE.toApiException(
                     "Embedding provider '%s' did not return expected embedding length. Expect: '%d'. Actual: '%d'",
-                    schemaObject.vectorConfig().vectorizeConfig().provider(),
-                    schemaObject.vectorConfig().vectorSize(),
+                    vectorConfig.vectorizeConfig().provider(),
+                    vectorConfig.vectorSize(),
                     vector.length);
               }
               return vector;
@@ -212,11 +217,11 @@ public Uni<Boolean> vectorize(SortClause sortClause) {
                 vectorData -> {
                   float[] vector = vectorData.get(0);
                   // check if vector have the expected size
-                  if (vector.length != schemaObject.vectorConfig().vectorSize()) {
+                  if (vector.length != vectorConfig.vectorSize()) {
                     throw EMBEDDING_PROVIDER_UNEXPECTED_RESPONSE.toApiException(
                         "Embedding provider '%s' did not return expected embedding length. Expect: '%d'. Actual: '%d'",
-                        schemaObject.vectorConfig().vectorizeConfig().provider(),
-                        schemaObject.vectorConfig().vectorSize(),
+                        vectorConfig.vectorizeConfig().provider(),
+                        vectorConfig.vectorSize(),
                         vector.length);
                   }
                   sortExpressions.clear();

src/main/java/io/stargate/sgv2/jsonapi/service/processor/MeteredCommandProcessor.java

@@ -22,12 +22,14 @@
 import io.stargate.sgv2.jsonapi.config.CommandLevelLoggingConfig;
 import io.stargate.sgv2.jsonapi.config.constants.DocumentConstants;
 import io.stargate.sgv2.jsonapi.service.cqldriver.executor.SchemaObject;
+import io.stargate.sgv2.jsonapi.service.cqldriver.executor.VectorConfig;
 import jakarta.enterprise.context.ApplicationScoped;
 import jakarta.enterprise.inject.Produces;
 import jakarta.inject.Inject;
 import jakarta.inject.Singleton;
 import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -243,8 +245,12 @@ private <T extends SchemaObject> Tags getCustomTags(
               result.errors().get(0).fieldsForMetricsTag().getOrDefault("errorCode", UNKNOWN_VALUE);
       errorCodeTag = Tag.of(jsonApiMetricsConfig.errorCode(), errorCode);
     }
+    final Optional<VectorConfig> first =
+        commandContext.schemaObject().vectorConfigs().stream()
+            .filter(a -> a.vectorEnabled())
+            .findFirst();
     Tag vectorEnabled =
-        commandContext.schemaObject().vectorConfig().vectorEnabled()
+        first.isPresent()
             ? Tag.of(jsonApiMetricsConfig.vectorEnabled(), "true")
             : Tag.of(jsonApiMetricsConfig.vectorEnabled(), "false");
     JsonApiMetricsConfig.SortType sortType = getVectorTypeTag(command);