Cassandra Vector Store initial impl follow up

- add concurrency to store.add(..) (bc embeddingClient is slow)
- CassandraVectorStoreAutoConfiguration uses CassandraAutoConfiguration
- driver profiles for production stability+performance,
- small cleanups and naming fixes,
- main doc tidy-up
- astradb compatibility (protocol V4)
– don't create embeddings again for documents that already have them
  similar to https://github.com/spring-projects/spring-ai/pull/413

vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraFilterExpressionConverter.java

@@ -79,7 +79,7 @@ final class CassandraFilterExpressionConverter extends AbstractFilterExpressionC
 			// TODO SAI supports collections
 			// reach out to mck@apache.org if you'd like these implemented
 			// case CONTAINS -> context.append(" CONTAINS ");
-			// case CONTAINS_KEY -> context.append(" CONTAINS KEY ");
+			// case CONTAINS_KEY -> context.append(" CONTAINS_KEY ");
 			default -> throw new UnsupportedOperationException(
 					String.format("Expression type %s not yet implemented. Patches welcome.", type));
 		}

vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java

@@ -17,16 +17,20 @@ package org.springframework.ai.vectorstore;
 
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 
 import com.datastax.oss.driver.api.core.cql.BoundStatement;
 import com.datastax.oss.driver.api.core.cql.BoundStatementBuilder;
 import com.datastax.oss.driver.api.core.cql.PreparedStatement;
 import com.datastax.oss.driver.api.core.cql.Row;
+import com.datastax.oss.driver.api.core.cql.SimpleStatement;
 import com.datastax.oss.driver.api.core.data.CqlVector;
 import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata;
 import com.datastax.oss.driver.api.querybuilder.QueryBuilder;
@@ -40,6 +44,7 @@ import org.slf4j.LoggerFactory;
 
 import org.springframework.ai.document.Document;
 import org.springframework.ai.embedding.EmbeddingClient;
+import org.springframework.ai.vectorstore.CassandraVectorStoreConfig;
 import org.springframework.ai.vectorstore.CassandraVectorStoreConfig.SchemaColumn;
 import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
 import org.springframework.beans.factory.InitializingBean;
@@ -53,13 +58,13 @@ import org.springframework.beans.factory.InitializingBean;
  * fields in the documents to be stored alongside the vector and content data.
  *
  * This class requires a CassandraVectorStoreConfig configuration object for
- * initialization, which includes settings like connection details, index name, field
+ * initialization, which includes settings like connection details, index name, column
  * names, etc. It also requires an EmbeddingClient to convert documents into embeddings
  * before storing them.
  *
  * A schema matching the configuration is automatically created if it doesn't exist.
  * Missing columns and indexes in existing tables will also be automatically created.
- * Disable this with the disallowSchemaCreation.
+ * Disable this with the CassandraVectorStoreConfig#disallowSchemaChanges().
  *
  * This class is designed to work with brand new tables that it creates for you, or on top
  * of existing Cassandra tables. The latter is appropriate when wanting to keep data in
@@ -69,9 +74,20 @@ import org.springframework.beans.factory.InitializingBean;
  * Instances of this class are not dynamic against server-side schema changes. If you
  * change the schema server-side you need a new CassandraVectorStore instance.
  *
+ * When adding documents with the method {@link #add(List<Document>)} it first calls
+ * embeddingClient to create the embeddings. This is slow. Configure
+ * {@link CassandraVectorStoreConfig.Builder#withFixedThreadPoolExecutorSize(int)}
+ * accordingly to improve performance so embeddings are created and the documents are
+ * added concurrently. The default concurrency is 16
+ * ({@link CassandraVectorStoreConfig#DEFAULT_ADD_CONCURRENCY}). Remote transformers
+ * probably want higher concurrency, and local transformers may need lower concurrency.
+ * This concurrency limit does not need to be higher than the max parallel calls made to
+ * the {@link #add(List<Document>)} method multiplied by the list size. This setting can
+ * also serve as a protecting throttle against your embedding model.
+ *
  * @author Mick Semb Wever
  * @see VectorStore
- * @see CassandraVectorStoreConfig
+ * @see org.springframework.ai.vectorstore.CassandraVectorStoreConfig
  * @see EmbeddingClient
  * @since 1.0.0
  */
@@ -87,10 +103,14 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean
 
 	}
 
-	private static final String QUERY_FORMAT = "select %s,%s,%s%s from %s.%s ? order by %s ann of ? limit ?";
-
 	public static final String SIMILARITY_FIELD_NAME = "similarity_score";
 
+	public static final String DRIVER_PROFILE_UPDATES = "spring-ai-updates";
+
+	public static final String DRIVER_PROFILE_SEARCH = "spring-ai-search";
+
+	private static final String QUERY_FORMAT = "select %s,%s,%s%s from %s.%s ? order by %s ann of ? limit ?";
+
 	private static final Logger logger = LoggerFactory.getLogger(CassandraVectorStore.class);
 
 	private final CassandraVectorStoreConfig conf;
@@ -99,7 +119,7 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean
 
 	private final FilterExpressionConverter filterExpressionConverter;
 
-	private final Map<Set<String>, PreparedStatement> addStmts = new HashMap<>();
+	private final ConcurrentMap<Set<String>, PreparedStatement> addStmts = new ConcurrentHashMap<>();
 
 	private final PreparedStatement deleteStmt;
 
@@ -133,30 +153,39 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean
 
 	@Override
 	public void add(List<Document> documents) {
-		CompletableFuture[] futures = new CompletableFuture[documents.size()];
-		short i = 0;
+		var futures = new CompletableFuture[documents.size()];
+
+		int i = 0;
 		for (Document d : documents) {
-			List<Object> primaryKeyValues = this.conf.documentIdTranslator.apply(d.getId());
-			var embedding = this.embeddingClient.embed(d).stream().map(Double::floatValue).toList();
+			futures[i++] = CompletableFuture.runAsync(() -> {
+				List<Object> primaryKeyValues = this.conf.documentIdTranslator.apply(d.getId());
 
-			BoundStatementBuilder builder = prepareAddStatement(d.getMetadata().keySet()).boundStatementBuilder();
-			for (int k = 0; k < primaryKeyValues.size(); ++k) {
-				SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);
-				builder = builder.set(keyColumn.name(), primaryKeyValues.get(k), keyColumn.javaType());
-			}
+				var embedding = (null != d.getEmbedding() && !d.getEmbedding().isEmpty() ? d.getEmbedding()
+						: this.embeddingClient.embed(d))
+					.stream()
+					.map(Double::floatValue)
+					.toList();
 
-			builder = builder.setString(this.conf.schema.content(), d.getContent())
-				.setVector(this.conf.schema.embedding(), CqlVector.newInstance(embedding), Float.class);
+				BoundStatementBuilder builder = prepareAddStatement(d.getMetadata().keySet()).boundStatementBuilder();
+				for (int k = 0; k < primaryKeyValues.size(); ++k) {
+					SchemaColumn keyColumn = this.conf.getPrimaryKeyColumn(k);
+					builder = builder.set(keyColumn.name(), primaryKeyValues.get(k), keyColumn.javaType());
+				}
 
-			for (var metadataColumn : this.conf.schema.metadataColumns()
-				.stream()
-				.filter((mc) -> d.getMetadata().containsKey(mc.name()))
-				.toList()) {
+				builder = builder.setString(this.conf.schema.content(), d.getContent())
+					.setVector(this.conf.schema.embedding(), CqlVector.newInstance(embedding), Float.class);
 
-				builder = builder.set(metadataColumn.name(), d.getMetadata().get(metadataColumn.name()),
-						metadataColumn.javaType());
-			}
-			futures[i++] = this.conf.session.executeAsync(builder.build()).toCompletableFuture();
+				for (var metadataColumn : this.conf.schema.metadataColumns()
+					.stream()
+					.filter((mc) -> d.getMetadata().containsKey(mc.name()))
+					.toList()) {
+
+					builder = builder.set(metadataColumn.name(), d.getMetadata().get(metadataColumn.name()),
+							metadataColumn.javaType());
+				}
+				BoundStatement s = builder.build().setExecutionProfileName(DRIVER_PROFILE_UPDATES);
+				this.conf.session.execute(s);
+			}, this.conf.executor);
 		}
 		CompletableFuture.allOf(futures).join();
 	}
@@ -164,7 +193,7 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean
 	@Override
 	public Optional<Boolean> delete(List<String> idList) {
 		CompletableFuture[] futures = new CompletableFuture[idList.size()];
-		short i = 0;
+		int i = 0;
 		for (String id : idList) {
 			List<Object> primaryKeyValues = this.conf.documentIdTranslator.apply(id);
 			BoundStatement s = this.deleteStmt.bind(primaryKeyValues.toArray());
@@ -191,8 +220,9 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean
 		String query = String.format(this.similarityStmt, cqlVector, whereClause, cqlVector, request.getTopK());
 		List<Document> documents = new ArrayList<>();
 		logger.trace("Executing {}", query);
+		SimpleStatement s = SimpleStatement.newInstance(query).setExecutionProfileName(DRIVER_PROFILE_SEARCH);
 
-		for (Row row : this.conf.session.execute(query)) {
+		for (Row row : this.conf.session.execute(s)) {
 			float score = row.getFloat(0);
 			if (score < request.getSimilarityThreshold()) {
 				break;
@@ -248,7 +278,16 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean
 	}
 
 	private PreparedStatement prepareAddStatement(Set<String> metadataFields) {
-		if (!this.addStmts.containsKey(metadataFields)) {
+
+		// metadata fields that are not configured as metadata columns are not added
+		Set<String> fieldsThatAreColumns = new HashSet<>(this.conf.schema.metadataColumns()
+			.stream()
+			.map((mc) -> mc.name())
+			.filter((mc) -> metadataFields.contains(mc))
+			.toList());
+
+		return this.addStmts.computeIfAbsent(fieldsThatAreColumns, (fields) -> {
+
 			RegularInsert stmt = null;
 			InsertInto stmtStart = QueryBuilder.insertInto(this.conf.schema.keyspace(), this.conf.schema.table());
 
@@ -262,17 +301,11 @@ public final class CassandraVectorStore implements VectorStore, InitializingBean
 			stmt = stmt.value(this.conf.schema.content(), QueryBuilder.bindMarker(this.conf.schema.content()))
 				.value(this.conf.schema.embedding(), QueryBuilder.bindMarker(this.conf.schema.embedding()));
 
-			for (String metadataField : this.conf.schema.metadataColumns()
-				.stream()
-				.map((mc) -> mc.name())
-				.filter((mc) -> metadataFields.contains(mc))
-				.toList()) {
-
+			for (String metadataField : fields) {
 				stmt = stmt.value(metadataField, QueryBuilder.bindMarker(metadataField));
 			}
-			this.addStmts.putIfAbsent(metadataFields, this.conf.session.prepare(stmt.build()));
-		}
-		return this.addStmts.get(metadataFields);
+			return this.conf.session.prepare(stmt.build());
+		});
 	}
 
 	private String similaritySearchStatement() {

vector-stores/spring-ai-cassandra/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStoreConfig.java

@@ -22,6 +22,8 @@ import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
 import java.util.function.Function;
 import java.util.stream.Stream;
 
@@ -48,15 +50,17 @@ import org.slf4j.LoggerFactory;
 /**
  * Configuration for the Cassandra vector store.
  *
- * All metadata fields configured to the store will be fetched and added to all queried
+ * All metadata columns configured to the store will be fetched and added to all queried
  * documents.
  *
- * If you wish to metadata search against a field its 'searchable' argument must be true.
+ * To filter expression search against a metadata column configure it with
+ * SchemaColumnTags.INDEXED
  *
  * The Cassandra Java Driver is configured via the application.conf resource found in the
  * classpath. See
  * https://github.com/apache/cassandra-java-driver/tree/4.x/manual/core/configuration
  *
+ * @author Mick Semb Wever
  * @since 1.0.0
  */
 public final class CassandraVectorStoreConfig implements AutoCloseable {
@@ -73,6 +77,8 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
 
 	public static final String DEFAULT_EMBEDDING_COLUMN_NAME = "embedding";
 
+	public static final int DEFAULT_ADD_CONCURRENCY = 16;
+
 	private static final Logger logger = LoggerFactory.getLogger(CassandraVectorStore.class);
 
 	record Schema(String keyspace, String table, List<SchemaColumn> partitionKeys, List<SchemaColumn> clusteringKeys,
@@ -127,6 +133,8 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
 
 	final PrimaryKeyTranslator primaryKeyTranslator;
 
+	final Executor executor;
+
 	private final boolean closeSessionOnClose;
 
 	private CassandraVectorStoreConfig(Builder builder) {
@@ -139,6 +147,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
 		this.disallowSchemaChanges = builder.disallowSchemaCreation;
 		this.documentIdTranslator = builder.documentIdTranslator;
 		this.primaryKeyTranslator = builder.primaryKeyTranslator;
+		this.executor = Executors.newFixedThreadPool(builder.fixedThreadPoolExecutorSize);
 	}
 
 	public static Builder builder() {
@@ -187,6 +196,8 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
 
 		private boolean disallowSchemaCreation = false;
 
+		private int fixedThreadPoolExecutorSize = DEFAULT_ADD_CONCURRENCY;
+
 		private DocumentIdTranslator documentIdTranslator = (String id) -> List.of(id);
 
 		private PrimaryKeyTranslator primaryKeyTranslator = (List<Object> primaryKeyColumns) -> {
@@ -261,20 +272,27 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
 			return this;
 		}
 
-		public Builder addMetadataColumn(SchemaColumn... fields) {
+		public Builder addMetadataColumns(SchemaColumn... columns) {
 			Builder builder = this;
-			for (SchemaColumn f : fields) {
+			for (SchemaColumn f : columns) {
 				builder = builder.addMetadataColumn(f);
 			}
 			return builder;
 		}
 
-		public Builder addMetadataColumn(SchemaColumn field) {
+		public Builder addMetadataColumns(List<SchemaColumn> columns) {
+			Builder builder = this;
+			this.metadataColumns.addAll(columns);
+			return builder;
+		}
 
-			Preconditions.checkArgument(this.metadataColumns.stream().noneMatch((sc) -> sc.name().equals(field.name())),
-					"A metadata field with name %s has already been added", field.name());
+		public Builder addMetadataColumn(SchemaColumn column) {
 
-			this.metadataColumns.add(field);
+			Preconditions.checkArgument(
+					this.metadataColumns.stream().noneMatch((sc) -> sc.name().equals(column.name())),
+					"A metadata column with name %s has already been added", column.name());
+
+			this.metadataColumns.add(column);
 			return this;
 		}
 
@@ -283,6 +301,18 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
 			return this;
 		}
 
+		/**
+		 * Executor to use when adding documents. The hotspot is the call to the
+		 * embeddingClient. For remote transformers you probably want a higher value to
+		 * utilize network. For local transformers you probably want a lower value to
+		 * avoid saturation.
+		 **/
+		public Builder withFixedThreadPoolExecutorSize(int threads) {
+			Preconditions.checkArgument(0 < threads);
+			this.fixedThreadPoolExecutorSize = threads;
+			return this;
+		}
+
 		public Builder withDocumentIdTranslator(DocumentIdTranslator documentIdTranslator) {
 			this.documentIdTranslator = documentIdTranslator;
 			return this;
@@ -480,7 +510,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
 			if (column.isPresent()) {
 
 				Preconditions.checkArgument(column.get().getType().equals(metadata.type()),
-						"Cannot change type on metadata field %s from %s to %s", metadata.name(),
+						"Cannot change type on metadata column %s from %s to %s", metadata.name(),
 						column.get().getType(), metadata.type());
 			}
 			else {

vector-stores/spring-ai-cassandra/src/main/resources/application.conf

@@ -0,0 +1,24 @@
+# Reference configuration for the DataStax Java driver for Apache Cassandra®
+#  see https://github.com/apache/cassandra-java-driver/tree/4.x/manual/core/configuration
+#
+# 
+# when using spring-boot autoconfigure this will not be used
+#  instead CassandraVectorStoreAutoConfiguration.driverConfigLoaderBuilderCustomizer() is used
+datastax-java-driver {
+  profiles {
+    spring-ai-updates {
+      basic.request {
+        consistency = LOCAL_QUORUM
+        timeout = 1 seconds
+        default-idempotence = true
+      }
+    }
+    spring-ai-search {
+        basic.request {
+          consistency = LOCAL_ONE
+          timeout = 10 seconds
+          default-idempotence = true
+      }
+    }
+  }
+}

vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraRichSchemaVectorStoreIT.java

@@ -17,23 +17,29 @@ package org.springframework.ai.vectorstore;
 
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import org.junit.jupiter.api.Assertions;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ThreadLocalRandom;
 
 import com.datastax.oss.driver.api.core.CqlSession;
 import com.datastax.oss.driver.api.core.CqlSessionBuilder;
 import com.datastax.oss.driver.api.core.servererrors.InvalidQueryException;
 import com.datastax.oss.driver.api.core.servererrors.SyntaxError;
 import com.datastax.oss.driver.api.core.type.DataTypes;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.testcontainers.containers.CassandraContainer;
 import org.testcontainers.junit.jupiter.Container;
 import org.testcontainers.junit.jupiter.Testcontainers;
+import org.testcontainers.shaded.org.apache.commons.lang3.RandomStringUtils;
 import org.testcontainers.utility.DockerImageName;
 
 import org.springframework.ai.document.Document;
@@ -174,6 +180,51 @@ class CassandraRichSchemaVectorStoreIT {
 		});
 	}
 
+	@Test
+	void addAndSearchPoormansBench() {
+		// todo – replace with JMH (parameters: nThreads, rounds, runs, docsPerAdd)
+		int nThreads = CassandraVectorStoreConfig.DEFAULT_ADD_CONCURRENCY;
+		int runs = 10; // 100;
+		int docsPerAdd = 12; // 128;
+		int rounds = 3;
+
+		contextRunner.run(context -> {
+
+			try (CassandraVectorStore store = new CassandraVectorStore(
+					storeBuilder(context, List.of()).withFixedThreadPoolExecutorSize(nThreads).build(),
+					context.getBean(EmbeddingClient.class))) {
+
+				var executor = Executors.newFixedThreadPool((int) (nThreads * 1.2));
+				for (int k = 0; k < rounds; ++k) {
+					long start = System.nanoTime();
+					var futures = new CompletableFuture[runs];
+					for (int j = 0; j < runs; ++j) {
+						futures[j] = CompletableFuture.runAsync(() -> {
+							List<Document> documents = new ArrayList<>();
+							for (int i = docsPerAdd; i >= 0; --i) {
+
+								documents.add(new Document(
+										RandomStringUtils.randomAlphanumeric(4) + "§¶"
+												+ ThreadLocalRandom.current().nextInt(1, 10),
+										RandomStringUtils.randomAlphanumeric(1024), Map.of("revision",
+												ThreadLocalRandom.current().nextInt(1, 100000), "id", 1000)));
+							}
+							store.add(documents);
+
+							var results = store.similaritySearch(
+									SearchRequest.query(RandomStringUtils.randomAlphanumeric(20)).withTopK(10));
+
+							assertThat(results).hasSize(10);
+						}, executor);
+					}
+					CompletableFuture.allOf(futures).join();
+					long time = System.nanoTime() - start;
+					logger.info("add+search took an average of {} ms", Duration.ofNanos(time / runs).toMillis());
+				}
+			}
+		});
+	}
+
 	@Test
 	void searchWithPartitionFilter() throws InterruptedException {
 		contextRunner.run(context -> {
@@ -456,22 +507,37 @@ class CassandraRichSchemaVectorStoreIT {
 	}
 
 	private StoreWrapper<CassandraVectorStore, CassandraVectorStoreConfig> createStore(ApplicationContext context,
-			List<SchemaColumn> extraMetadataFields, boolean disallowSchemaCreation, boolean dropKeyspaceFirst)
+			List<SchemaColumn> columnOverrides, boolean disallowSchemaCreation, boolean dropKeyspaceFirst)
 			throws IOException {
 
-		Optional<SchemaColumn> wikiOverride = extraMetadataFields.stream()
+		CassandraVectorStoreConfig.Builder builder = storeBuilder(context, columnOverrides);
+		if (disallowSchemaCreation) {
+			builder = builder.disallowSchemaChanges();
+		}
+
+		CassandraVectorStoreConfig conf = builder.build();
+		if (dropKeyspaceFirst) {
+			conf.dropKeyspace();
+		}
+		return new StoreWrapper(new CassandraVectorStore(conf, context.getBean(EmbeddingClient.class)), conf);
+	}
+
+	static CassandraVectorStoreConfig.Builder storeBuilder(ApplicationContext context,
+			List<SchemaColumn> columnOverrides) throws IOException {
+
+		Optional<SchemaColumn> wikiOverride = columnOverrides.stream()
 			.filter((f) -> "wiki".equals(f.name()))
 			.findFirst();
 
-		Optional<SchemaColumn> langOverride = extraMetadataFields.stream()
+		Optional<SchemaColumn> langOverride = columnOverrides.stream()
 			.filter((f) -> "language".equals(f.name()))
 			.findFirst();
 
-		Optional<SchemaColumn> titleOverride = extraMetadataFields.stream()
+		Optional<SchemaColumn> titleOverride = columnOverrides.stream()
 			.filter((f) -> "title".equals(f.name()))
 			.findFirst();
 
-		Optional<SchemaColumn> chunkNoOverride = extraMetadataFields.stream()
+		Optional<SchemaColumn> chunkNoOverride = columnOverrides.stream()
 			.filter((f) -> "chunk_no".equals(f.name()))
 			.findFirst();
 
@@ -493,7 +559,7 @@ class CassandraRichSchemaVectorStoreIT {
 			.withEmbeddingColumnName("all_minilm_l6_v2_embedding")
 			.withIndexName("all_minilm_l6_v2_ann")
 
-			.addMetadataColumn(new SchemaColumn("revision", DataTypes.INT),
+			.addMetadataColumns(new SchemaColumn("revision", DataTypes.INT),
 					new SchemaColumn("id", DataTypes.INT, CassandraVectorStoreConfig.SchemaColumnTags.INDEXED))
 
 			// this store uses '§¶' as a deliminator in the document id between db columns
@@ -511,21 +577,7 @@ class CassandraRichSchemaVectorStoreIT {
 				return List.of("simplewiki", "en", title, chunk_no);
 			});
 
-		for (SchemaColumn cf : extraMetadataFields) {
-			if (!partitionKeys.contains(cf) && !clusteringKeys.contains(cf)) {
-				builder = builder.addMetadataColumn(cf);
-			}
-		}
-
-		if (disallowSchemaCreation) {
-			builder = builder.disallowSchemaChanges();
-		}
-
-		CassandraVectorStoreConfig conf = builder.build();
-		if (dropKeyspaceFirst) {
-			conf.dropKeyspace();
-		}
-		return new StoreWrapper(new CassandraVectorStore(conf, context.getBean(EmbeddingClient.class)), conf);
+		return builder;
 	}
 
 	private void executeCqlFile(ApplicationContext context, String filename) throws IOException {

vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/CassandraVectorStoreIT.java

@@ -345,8 +345,9 @@ class CassandraVectorStoreIT {
 		public CassandraVectorStore store(CqlSession cqlSession, EmbeddingClient embeddingClient) {
 
 			CassandraVectorStoreConfig conf = storeBuilder(cqlSession)
-				.addMetadataColumn(new SchemaColumn("meta1", DataTypes.TEXT), new SchemaColumn("meta2", DataTypes.TEXT),
-						new SchemaColumn("country", DataTypes.TEXT), new SchemaColumn("year", DataTypes.SMALLINT))
+				.addMetadataColumns(new SchemaColumn("meta1", DataTypes.TEXT),
+						new SchemaColumn("meta2", DataTypes.TEXT), new SchemaColumn("country", DataTypes.TEXT),
+						new SchemaColumn("year", DataTypes.SMALLINT))
 				.build();
 
 			conf.dropKeyspace();
@@ -378,7 +379,7 @@ class CassandraVectorStoreIT {
 	private CassandraVectorStore createTestStore(ApplicationContext context, SchemaColumn... metadataFields) {
 
 		CassandraVectorStoreConfig.Builder builder = storeBuilder(context.getBean(CqlSession.class))
-			.addMetadataColumn(metadataFields);
+			.addMetadataColumns(metadataFields);
 
 		CassandraVectorStoreConfig conf = builder.build();
 		conf.dropKeyspace();

vector-stores/spring-ai-cassandra/src/test/java/org/springframework/ai/vectorstore/WikiVectorStoreExample.java

@@ -81,25 +81,30 @@ class WikiVectorStoreExample {
 		@Bean
 		public CassandraVectorStore store(CqlSession cqlSession, EmbeddingClient embeddingClient) {
 
+			List<SchemaColumn> partitionColumns = List.of(new SchemaColumn("wiki", DataTypes.TEXT),
+					new SchemaColumn("language", DataTypes.TEXT), new SchemaColumn("title", DataTypes.TEXT));
+
+			List<SchemaColumn> clusteringColumns = List.of(new SchemaColumn("chunk_no", DataTypes.INT),
+					new SchemaColumn("bert_embedding_no", DataTypes.INT));
+
+			List<SchemaColumn> extraColumns = List.of(new SchemaColumn("revision", DataTypes.INT),
+					new SchemaColumn("id", DataTypes.INT));
+
 			CassandraVectorStoreConfig conf = CassandraVectorStoreConfig.builder()
 				.withCqlSession(cqlSession)
 				.withKeyspaceName("wikidata")
 				.withTableName("articles")
-
-				.withPartitionKeys(List.of(new SchemaColumn("wiki", DataTypes.TEXT),
-						new SchemaColumn("language", DataTypes.TEXT), new SchemaColumn("title", DataTypes.TEXT)))
-
-				.withClusteringKeys(List.of(new SchemaColumn("chunk_no", DataTypes.INT),
-						new SchemaColumn("bert_embedding_no", DataTypes.INT)))
-
+				.withPartitionKeys(partitionColumns)
+				.withClusteringKeys(clusteringColumns)
 				.withContentColumnName("body")
 				.withEmbeddingColumnName("all_minilm_l6_v2_embedding")
 				.withIndexName("all_minilm_l6_v2_ann")
 				.disallowSchemaChanges()
-
-				.addMetadataColumn(new SchemaColumn("revision", DataTypes.INT), new SchemaColumn("id", DataTypes.INT))
+				.addMetadataColumns(extraColumns)
 
 				.withPrimaryKeyTranslator((List<Object> primaryKeys) -> {
+					// the deliminator used to join fields together into the document's id
+					// is arbitary, here "§¶" is used
 					if (primaryKeys.isEmpty()) {
 						return "test§¶0";
 					}