Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Avoid calling CREATE KEYSPACE|TABLE IF NOT EXISTS … unnecessarily

Browse Source 
This also is a fix for AstraDB, which throws an exception if you try to `CREATE KEYSPACE IF NOT EXISTS …`
And use more compatible `USING 'StorageAttachedIndex'` index creation syntax.
This commit is contained in:
mck
2024-04-29 15:40:33 +02:00
committed by Christian Tzolov
parent 3c40268d80
commit ac1a0d67e9
1 changed files with 41 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -441,7 +441,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
{
SimpleStatement indexStmt = SchemaBuilder.createIndex(this.schema.index)
.ifNotExists()
.custom("SAI")
.custom("StorageAttachedIndex")
.onTable(this.schema.keyspace, this.schema.table)
.andColumn(this.schema.embedding)
.build();
@@ -457,7 +457,7 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
SimpleStatement indexStmt = SchemaBuilder.createIndex(String.format("%s_idx", metadata.name()))
.ifNotExists()
.custom("SAI")
.custom("StorageAttachedIndex")
.onTable(this.schema.keyspace, this.schema.table)
.andColumn(metadata.name())
.build();
@@ -468,39 +468,41 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
}
private void ensureTableExists(int vectorDimension) {
if (this.session.getMetadata().getKeyspace(this.schema.keyspace).get().getTable(this.schema.table).isEmpty()) {
CreateTable createTable = null;
CreateTable createTable = null;
CreateTableStart createTableStart = SchemaBuilder.createTable(this.schema.keyspace, this.schema.table)
.ifNotExists();
CreateTableStart createTableStart = SchemaBuilder.createTable(this.schema.keyspace, this.schema.table)
.ifNotExists();
for (SchemaColumn partitionKey : this.schema.partitionKeys) {
createTable = (null != createTable ? createTable : createTableStart).withPartitionKey(partitionKey.name,
partitionKey.type);
for (SchemaColumn partitionKey : this.schema.partitionKeys) {
createTable = (null != createTable ? createTable : createTableStart).withPartitionKey(partitionKey.name,
partitionKey.type);
}
for (SchemaColumn clusteringKey : this.schema.clusteringKeys) {
createTable = createTable.withClusteringColumn(clusteringKey.name, clusteringKey.type);
}
createTable = createTable.withColumn(this.schema.content, DataTypes.TEXT);
for (SchemaColumn metadata : this.schema.metadataColumns) {
createTable = createTable.withColumn(metadata.name(), metadata.type());
}
// https://datastax-oss.atlassian.net/browse/JAVA-3118
// .withColumn(config.embedding, new DefaultVectorType(DataTypes.FLOAT,
// vectorDimension));
StringBuilder tableStmt = new StringBuilder(createTable.asCql());
tableStmt.setLength(tableStmt.length() - 1);
tableStmt.append(',')
.append(this.schema.embedding)
.append(" vector<float,")
.append(vectorDimension)
.append(">)");
logger.debug("Executing {}", tableStmt.toString());
this.session.execute(tableStmt.toString());
}
for (SchemaColumn clusteringKey : this.schema.clusteringKeys) {
createTable = createTable.withClusteringColumn(clusteringKey.name, clusteringKey.type);
}
createTable = createTable.withColumn(this.schema.content, DataTypes.TEXT);
for (SchemaColumn metadata : this.schema.metadataColumns) {
createTable = createTable.withColumn(metadata.name(), metadata.type());
}
// https://datastax-oss.atlassian.net/browse/JAVA-3118
// .withColumn(config.embedding, new DefaultVectorType(DataTypes.FLOAT,
// vectorDimension));
StringBuilder tableStmt = new StringBuilder(createTable.asCql());
tableStmt.setLength(tableStmt.length() - 1);
tableStmt.append(',')
.append(this.schema.embedding)
.append(" vector<float,")
.append(vectorDimension)
.append(">)");
logger.debug("Executing {}", tableStmt.toString());
this.session.execute(tableStmt.toString());
}
private void ensureTableColumnsExist(int vectorDimension) {
@@ -563,14 +565,15 @@ public final class CassandraVectorStoreConfig implements AutoCloseable {
}
private void ensureKeyspaceExists() {
if (this.session.getMetadata().getKeyspace(this.schema.keyspace).isEmpty()) {
SimpleStatement keyspaceStmt = SchemaBuilder.createKeyspace(this.schema.keyspace)
.ifNotExists()
.withSimpleStrategy(1)
.build();
SimpleStatement keyspaceStmt = SchemaBuilder.createKeyspace(this.schema.keyspace)
.ifNotExists()
.withSimpleStrategy(1)
.build();
logger.debug("Executing {}", keyspaceStmt.getQuery());
this.session.execute(keyspaceStmt);
logger.debug("Executing {}", keyspaceStmt.getQuery());
this.session.execute(keyspaceStmt);
}
}
}