Skip to content

Commit 77f0e19

Browse files
clunclun
committed
Update Stargate
1 parent 420f536 commit 77f0e19

File tree

5 files changed

+234
-2
lines changed

5 files changed

+234
-2
lines changed

‎astra-db-client/src/main/java/com/dtsx/astra/sdk/cassio/ClusteredMetadataVectorCassandraTable.java‎

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
package com.dtsx.astra.sdk.cassio;
2+
3+
import com.datastax.oss.driver.api.core.CqlSession;
4+
import com.datastax.oss.driver.api.core.cql.PreparedStatement;
5+
import com.datastax.oss.driver.api.core.cql.Row;
6+
import lombok.NonNull;
7+
import lombok.extern.slf4j.Slf4j;
8+
9+
/**
10+
* Table representing persistence for LangChain operations.
11+
* - parition key: partitionId
12+
* - clustering key: rowId
13+
* - column: value
14+
*/
15+
@Slf4j
16+
public class ClusteredMetadataVectorCassandraTable
17+
extends AbstractCassandraTable<ClusteredMetadataVectorCassandraTable.Record> {
18+
19+
/**
20+
* Dimension of the vector in use
21+
*/
22+
private final int vectorDimension;
23+
24+
/**
25+
* Similarity Metric, Vector is indexed with this metric.
26+
*/
27+
private final SimilarityMetric similarityMetric;
28+
29+
/**
30+
* Prepared statements
31+
*/
32+
private PreparedStatement findPartitionStatement;
33+
private PreparedStatement deletePartitionStatement;
34+
private PreparedStatement deleteRowStatement;
35+
private PreparedStatement insertRowStatement;
36+
private PreparedStatement findRowStatement;
37+
38+
/**
39+
* Constructor with mandatory parameters.
40+
*
41+
* @param session cassandra session
42+
* @param keyspaceName keyspace name
43+
* @param tableName table name
44+
* @param vectorDimension vector dimension
45+
* @param metric similarity metric
46+
*/
47+
public ClusteredMetadataVectorCassandraTable(
48+
@NonNull CqlSession session,
49+
@NonNull String keyspaceName,
50+
@NonNull String tableName,
51+
@NonNull Integer vectorDimension,
52+
@NonNull SimilarityMetric metric) {
53+
super(session, keyspaceName, tableName);
54+
this.vectorDimension = vectorDimension;
55+
this.similarityMetric = metric;
56+
}
57+
58+
/**
59+
* Builder class for creating instances of {@link ClusteredMetadataVectorCassandraTable}.
60+
* This class follows the builder pattern to allow setting various parameters
61+
* before creating an instance of {@link ClusteredMetadataVectorCassandraTable}.
62+
*/
63+
public static class Builder {
64+
private CqlSession session;
65+
private String keyspaceName;
66+
private String tableName;
67+
private Integer vectorDimension;
68+
private SimilarityMetric metric = SimilarityMetric.COS;
69+
70+
/**
71+
* Sets the CqlSession.
72+
*
73+
* @param session The CqlSession to be used by the ClusteredMetadataVectorCassandraTable.
74+
* @return The current Builder instance for chaining.
75+
*/
76+
public Builder withSession(CqlSession session) {
77+
this.session = session;
78+
return this;
79+
}
80+
81+
/**
82+
* Sets the keyspace name.
83+
*
84+
* @param keyspaceName The name of the keyspace to be used.
85+
* @return The current Builder instance for chaining.
86+
*/
87+
public Builder withKeyspaceName(String keyspaceName) {
88+
this.keyspaceName = keyspaceName;
89+
return this;
90+
}
91+
92+
/**
93+
* Sets the table name.
94+
*
95+
* @param tableName The name of the table to be used.
96+
* @return The current Builder instance for chaining.
97+
*/
98+
public Builder withTableName(String tableName) {
99+
this.tableName = tableName;
100+
return this;
101+
}
102+
103+
/**
104+
* Sets the vector dimension.
105+
*
106+
* @param vectorDimension The vector dimension to be used.
107+
* @return The current Builder instance for chaining.
108+
*/
109+
public Builder withVectorDimension(Integer vectorDimension) {
110+
this.vectorDimension = vectorDimension;
111+
return this;
112+
}
113+
114+
/**
115+
* Sets the similarity metric.
116+
*
117+
* @param metric The SimilarityMetric to be used.
118+
* @return The current Builder instance for chaining.
119+
*/
120+
public Builder withMetric(SimilarityMetric metric) {
121+
this.metric = metric;
122+
return this;
123+
}
124+
125+
/**
126+
* Creates a new instance of ClusteredMetadataVectorCassandraTable with the current builder settings.
127+
*
128+
* @return A new instance of ClusteredMetadataVectorCassandraTable.
129+
*/
130+
public ClusteredMetadataVectorCassandraTable build() {
131+
return new ClusteredMetadataVectorCassandraTable(session, keyspaceName, tableName, vectorDimension, metric);
132+
}
133+
134+
/**
135+
* Default constructor for Builder.
136+
*/
137+
public Builder() {}
138+
}
139+
140+
/**
141+
* Builder for the class.
142+
*
143+
* @return
144+
* builder for the class
145+
*/
146+
public static ClusteredMetadataVectorCassandraTable.Builder builder() {
147+
return new Builder();
148+
}
149+
150+
/**
151+
* Prepare statements on first request.
152+
*/
153+
private synchronized void prepareStatements() {
154+
if (findPartitionStatement == null) {
155+
findPartitionStatement = cqlSession.prepare(
156+
"select * from " + keyspaceName + "." + tableName
157+
+ " where " + PARTITION_ID + " = ? ");
158+
deletePartitionStatement = cqlSession.prepare(
159+
"delete from " + keyspaceName + "." + tableName
160+
+ " where " + PARTITION_ID + " = ? ");
161+
findRowStatement = cqlSession.prepare(
162+
"select * from " + keyspaceName + "." + tableName
163+
+ " where " + PARTITION_ID + " = ? "
164+
+ " and " + ROW_ID + " = ? ");
165+
deleteRowStatement = cqlSession.prepare(
166+
"delete from " + keyspaceName + "." + tableName
167+
+ " where " + PARTITION_ID + " = ? "
168+
+ " and " + ROW_ID + " = ? ");
169+
insertRowStatement = cqlSession.prepare(
170+
"insert into " + keyspaceName + "." + tableName
171+
+ " (" + PARTITION_ID + ", " + ROW_ID + ", " + BODY_BLOB + ") "
172+
+ " values (?, ?, ?)");
173+
}
174+
}
175+
176+
@Override
177+
public void create() {
178+
// Create Table
179+
cqlSession.execute("CREATE TABLE IF NOT EXISTS " + tableName + " (" +
180+
PARTITION_ID + "uuid," +
181+
ROW_ID + "text, " +
182+
ATTRIBUTES_BLOB + " text, " +
183+
BODY_BLOB + " text, " +
184+
METADATA_S + " map<text, text>, " +
185+
VECTOR + " vector<float, " + vectorDimension + ">, " +
186+
"PRIMARY KEY ((" + PARTITION_ID + "), " + ROW_ID + ")) " +
187+
"WITH CLUSTERING ORDER BY (" + ROW_ID + " DESC)");
188+
cqlSession.execute(
189+
"CREATE CUSTOM INDEX IF NOT EXISTS idx_vector_" + tableName
190+
+ " ON " + tableName + " (" + VECTOR + ") "
191+
+ "USING 'org.apache.cassandra.index.sai.StorageAttachedIndex' "
192+
+ "WITH OPTIONS = { 'similarity_function': '" + similarityMetric.getOption() + "'};");
193+
log.info("+ Index '{}' has been created (if needed).", "idx_vector_" + tableName);
194+
// Create Metadata Index
195+
cqlSession.execute(
196+
"CREATE CUSTOM INDEX IF NOT EXISTS eidx_metadata_s_" + tableName
197+
+ " ON " + tableName + " (ENTRIES(" + METADATA_S + ")) "
198+
+ "USING 'org.apache.cassandra.index.sai.StorageAttachedIndex' ");
199+
log.info("+ Index '{}' has been created (if needed).", "eidx_metadata_s_" + tableName);
200+
201+
}
202+
203+
@Override
204+
public void put(Record row) {
205+
206+
}
207+
208+
@Override
209+
public Record mapRow(Row row) {
210+
return null;
211+
}
212+
213+
public static class Record {}
214+
215+
}

‎astra-db-client/src/test/java/com/dtsx/astra/sdk/AstraDBIntegrationsTest.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,4 +219,21 @@ public void shouldInsertWithSimpleCollectionJsonMapping() {
219219
);
220220
}
221221

222+
@Test
223+
@Order(9)
224+
@DisplayName("09. Upsert with CollectionClient")
225+
public void shouldUpsertDocument() {
226+
if (astraDb == null) {
227+
astraDb = astraDbClient.database(TEST_DBNAME);
228+
}
229+
CollectionClient collection = astraDb.collection(TEST_COLLECTION_NAME);
230+
Assertions.assertNotNull(collection);
231+
JsonDocument doc = new JsonDocument()
232+
.id("4")
233+
.put("name", "Coded Cleats Copy")
234+
.put("description", "ChatGPT integrated sneakers that talk to you");
235+
collection.upsert(doc);
236+
collection.upsert(doc);
237+
}
238+
222239
}

‎astra-db-client/src/test/java/com/dtsx/astra/sdk/vector/AstraDBQuickStart.java‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,6 @@ public static void main(String[] args) {
3838
.put("description", "Vision Vector Frame - A deep learning display that controls your mood")
3939
.vector(new float[]{0.1f, 0.05f, 0.08f, 0.3f, 0.6f})
4040
));
41-
4241
// Search
4342
List<JsonResult> resultsSet =
4443
testCollection.similaritySearch(new float[]{0.15f, 0.1f, 0.1f, 0.35f, 0.55f},10);

‎astra-db-client/src/test/java/com/dtsx/astra/sdk/vector/AstraVectorQuickStart.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ public void quickStartTest() {
3535
// 2. Create a store (delete if exist)
3636
AstraDB astraDB = astraDBClient.database(databaseName);
3737

38+
3839
// 3. Insert data in the store
3940
astraDB.deleteCollection(collectionName);
4041
CollectionClient collection = astraDB.createCollection(collectionName, 14);

‎pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@
3838
<lombok.version>1.18.30</lombok.version>
3939

4040
<!-- Stargate -->
41-
<stargate-sdk.version>2.1.4-SNAPSHOT</stargate-sdk.version>
41+
<stargate-sdk.version>2.1.4</stargate-sdk.version>
4242
<stargate-grpc.version>2.0.17</stargate-grpc.version>
4343
<grpc-netty.version>1.56.1</grpc-netty.version>
4444

0 commit comments

Comments
 (0)