From 9f660ec3079b3ae049e81062075837bf5b6da3ca Mon Sep 17 00:00:00 2001 From: Sumit Soreng Date: Tue, 25 Aug 2026 00:05:35 +0530 Subject: [PATCH 1/2] [ISSUE-851] implemented VectorStore and VectorStoreMetadata interfaces with validation and persistence --- .../ai/index/vectorstore/VectorStore.java | 71 +++++++ .../vectorstore/VectorStoreMetadata.java | 190 ++++++++++++++++++ .../vectorstore/VectorStoreMetadataTest.java | 128 ++++++++++++ 3 files changed, 389 insertions(+) create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStore.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadata.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadataTest.java diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStore.java new file mode 100644 index 000000000..175e8c412 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStore.java @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.ai.index.vectorstore; + +import java.util.List; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.geaflow.ai.index.vector.IVector; + +public interface VectorStore { + + /** + * Get the metadata of this vector store. + * + * @return The vector store metadata. + */ + VectorStoreMetadata getMetadata(); + + /** + * Initialize the vector store. + * Implementations should validate the expected metadata against the stored metadata. + * + * @param expectedMetadata The expected metadata for validation. + */ + void init(VectorStoreMetadata expectedMetadata); + + /** + * Add a vector to the store with a given ID. + * + * @param id The identifier for the vector. + * @param vector The vector to store. + */ + void add(String id, IVector vector); + + /** + * Delete a vector by ID. + * + * @param id The identifier for the vector. + */ + void delete(String id); + + /** + * Search the top-K closest vectors to the query vector. + * + * @param queryVector The query vector. + * @param topK The number of results to return. + * @return A list of pairs containing the vector ID and distance score. + */ + List> search(IVector queryVector, int topK); + + /** + * Close the vector store and release any resources. + */ + void close(); +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadata.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadata.java new file mode 100644 index 000000000..9335312c9 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadata.java @@ -0,0 +1,190 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.ai.index.vectorstore; + +import com.google.gson.Gson; +import com.google.gson.GsonBuilder; +import com.google.gson.annotations.SerializedName; +import java.io.BufferedReader; +import java.io.BufferedWriter; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Objects; + +public class VectorStoreMetadata { + + @SerializedName("model_name") + private String modelName; + + @SerializedName("dimension") + private int dimension; + + @SerializedName("distance") + private String distance; + + @SerializedName("index_version") + private String indexVersion; + + @SerializedName("created_at") + private long createdAt; + + @SerializedName("format_version") + private String formatVersion; + + public VectorStoreMetadata() { + } + + public VectorStoreMetadata(String modelName, int dimension, String distance, + String indexVersion, long createdAt, String formatVersion) { + this.modelName = modelName; + this.dimension = dimension; + this.distance = distance; + this.indexVersion = indexVersion; + this.createdAt = createdAt; + this.formatVersion = formatVersion; + } + + public String getModelName() { + return modelName; + } + + public void setModelName(String modelName) { + this.modelName = modelName; + } + + public int getDimension() { + return dimension; + } + + public void setDimension(int dimension) { + this.dimension = dimension; + } + + public String getDistance() { + return distance; + } + + public void setDistance(String distance) { + this.distance = distance; + } + + public String getIndexVersion() { + return indexVersion; + } + + public void setIndexVersion(String indexVersion) { + this.indexVersion = indexVersion; + } + + public long getCreatedAt() { + return createdAt; + } + + public void setCreatedAt(long createdAt) { + this.createdAt = createdAt; + } + + public String getFormatVersion() { + return formatVersion; + } + + public void setFormatVersion(String formatVersion) { + this.formatVersion = formatVersion; + } + + public void validate(VectorStoreMetadata expected) { + if (expected == null) { + return; + } + if (this.dimension != expected.dimension) { + throw new IllegalArgumentException(String.format( + "Dimension mismatch. Expected %d, but got %d", expected.dimension, this.dimension)); + } + if (!Objects.equals(this.modelName, expected.modelName)) { + throw new IllegalArgumentException(String.format( + "Model mismatch. Expected %s, but got %s", expected.modelName, this.modelName)); + } + } + + public void validateComplete() { + if (modelName == null || modelName.isEmpty()) { + throw new IllegalArgumentException("Missing metadata: model_name"); + } + if (dimension <= 0) { + throw new IllegalArgumentException("Missing or invalid metadata: dimension"); + } + if (distance == null || distance.isEmpty()) { + throw new IllegalArgumentException("Missing metadata: distance"); + } + if (indexVersion == null || indexVersion.isEmpty()) { + throw new IllegalArgumentException("Missing metadata: index_version"); + } + if (formatVersion == null || formatVersion.isEmpty()) { + throw new IllegalArgumentException("Missing metadata: format_version"); + } + } + + public static VectorStoreMetadata load(Path metadataPath) throws IOException { + if (!Files.exists(metadataPath)) { + throw new IOException("Metadata file does not exist: " + metadataPath.toAbsolutePath()); + } + try (BufferedReader reader = Files.newBufferedReader(metadataPath, StandardCharsets.UTF_8)) { + Gson gson = new Gson(); + VectorStoreMetadata metadata = gson.fromJson(reader, VectorStoreMetadata.class); + if (metadata == null) { + throw new IOException("Failed to parse metadata from file"); + } + metadata.validateComplete(); + return metadata; + } + } + + public void save(Path metadataPath) throws IOException { + validateComplete(); + try (BufferedWriter writer = Files.newBufferedWriter(metadataPath, StandardCharsets.UTF_8)) { + Gson gson = new GsonBuilder().setPrettyPrinting().create(); + gson.toJson(this, writer); + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + VectorStoreMetadata that = (VectorStoreMetadata) o; + return dimension == that.dimension + && createdAt == that.createdAt + && Objects.equals(modelName, that.modelName) + && Objects.equals(distance, that.distance) + && Objects.equals(indexVersion, that.indexVersion) + && Objects.equals(formatVersion, that.formatVersion); + } + + @Override + public int hashCode() { + return Objects.hash(modelName, dimension, distance, indexVersion, createdAt, formatVersion); + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadataTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadataTest.java new file mode 100644 index 000000000..9f310cea1 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadataTest.java @@ -0,0 +1,128 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.geaflow.ai.index.vectorstore; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +public class VectorStoreMetadataTest { + + private Path tempFile; + + @BeforeEach + public void setUp() throws IOException { + tempFile = Files.createTempFile("metadata_test", ".json"); + } + + @AfterEach + public void tearDown() throws IOException { + Files.deleteIfExists(tempFile); + } + + @Test + public void testPersistAndLoadMetadata() throws IOException { + VectorStoreMetadata metadata = new VectorStoreMetadata( + "test_model", + 128, + "COSINE", + "1.0", + System.currentTimeMillis(), + "v1" + ); + metadata.save(tempFile); + + VectorStoreMetadata loaded = VectorStoreMetadata.load(tempFile); + assertEquals(metadata, loaded); + assertEquals("test_model", loaded.getModelName()); + assertEquals(128, loaded.getDimension()); + } + + @Test + public void testMissingMetadataValidationFails() { + VectorStoreMetadata metadata = new VectorStoreMetadata(); + assertThrows(IllegalArgumentException.class, metadata::validateComplete); + + metadata.setModelName("test_model"); + assertThrows(IllegalArgumentException.class, metadata::validateComplete); + + metadata.setDimension(128); + assertThrows(IllegalArgumentException.class, metadata::validateComplete); + + metadata.setDistance("L2"); + assertThrows(IllegalArgumentException.class, metadata::validateComplete); + + metadata.setIndexVersion("1.0"); + assertThrows(IllegalArgumentException.class, metadata::validateComplete); + + metadata.setFormatVersion("v1"); + // Should not throw now + metadata.validateComplete(); + } + + @Test + public void testDimensionMismatchFails() { + VectorStoreMetadata expected = new VectorStoreMetadata("test_model", 128, "L2", "1.0", 0, "v1"); + VectorStoreMetadata actual = new VectorStoreMetadata("test_model", 256, "L2", "1.0", 0, "v1"); + + IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, () -> { + actual.validate(expected); + }); + assertEquals("Dimension mismatch. Expected 128, but got 256", ex.getMessage()); + } + + @Test + public void testModelMismatchFails() { + VectorStoreMetadata expected = new VectorStoreMetadata("expected_model", 128, "L2", "1.0", 0, "v1"); + VectorStoreMetadata actual = new VectorStoreMetadata("actual_model", 128, "L2", "1.0", 0, "v1"); + + IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, () -> { + actual.validate(expected); + }); + assertEquals("Model mismatch. Expected expected_model, but got actual_model", ex.getMessage()); + } + + @Test + public void testLoadInvalidJsonFails() throws IOException { + Files.write(tempFile, "{ \"dimension\": 128, \"model_name\": \"test\" ".getBytes(StandardCharsets.UTF_8)); // Invalid JSON + + assertThrows(com.google.gson.JsonSyntaxException.class, () -> { + VectorStoreMetadata.load(tempFile); + }); + } + + @Test + public void testLoadMissingFieldsFails() throws IOException { + // Missing distance, index_version, format_version + Files.write(tempFile, "{ \"dimension\": 128, \"model_name\": \"test\" }".getBytes(StandardCharsets.UTF_8)); + + assertThrows(IllegalArgumentException.class, () -> { + VectorStoreMetadata.load(tempFile); + }); + } +} From 90a6333d6db291be998603f080cae8fe2025d003 Mon Sep 17 00:00:00 2001 From: Sumit Soreng Date: Fri, 28 Aug 2026 12:34:19 +0530 Subject: [PATCH 2/2] feat(ai): implemented decoupled VectorStore SPI architecture --- .../ai/index/vectorstore/DistanceMetric.java | 24 ++ .../ai/index/vectorstore/DistanceUtils.java | 70 +++++ .../vectorstore/InMemoryVectorStore.java | 119 ++++++++ .../index/vectorstore/LocalVectorStore.java | 256 ++++++++++++++++++ .../ai/index/vectorstore/VectorHit.java | 42 +++ .../ai/index/vectorstore/VectorQuery.java | 51 ++++ .../ai/index/vectorstore/VectorRecord.java | 69 +++++ .../ai/index/vectorstore/VectorStore.java | 71 +++-- .../vectorstore/VectorStoreException.java | 52 ++++ .../vectorstore/VectorStoreMetadata.java | 160 +++-------- .../vectorstore/VectorStoreMetadataTest.java | 128 --------- .../ai/index/vectorstore/VectorStoreTest.java | 111 ++++++++ 12 files changed, 864 insertions(+), 289 deletions(-) create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/DistanceMetric.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/DistanceUtils.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/InMemoryVectorStore.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/LocalVectorStore.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorHit.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorQuery.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorRecord.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreException.java delete mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadataTest.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreTest.java diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/DistanceMetric.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/DistanceMetric.java new file mode 100644 index 000000000..db5445c06 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/DistanceMetric.java @@ -0,0 +1,24 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.geaflow.ai.index.vectorstore; + +public enum DistanceMetric { + COSINE, + L2, + DOT_PRODUCT +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/DistanceUtils.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/DistanceUtils.java new file mode 100644 index 000000000..6750490b0 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/DistanceUtils.java @@ -0,0 +1,70 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.geaflow.ai.index.vectorstore; + +public class DistanceUtils { + + public static double compute(double[] a, double[] b, DistanceMetric metric) { + if (a.length != b.length) { + throw new VectorStoreException(VectorStoreException.ErrorCode.DIMENSION_MISMATCH, + "Vector lengths differ: " + a.length + " vs " + b.length); + } + switch (metric) { + case COSINE: + return cosineSimilarity(a, b); + case L2: + return 1.0 / (1.0 + euclideanDistance(a, b)); + case DOT_PRODUCT: + return dotProduct(a, b); + default: + throw new IllegalArgumentException("Unknown metric: " + metric); + } + } + + private static double cosineSimilarity(double[] a, double[] b) { + double dot = 0.0; + double normA = 0.0; + double normB = 0.0; + for (int i = 0; i < a.length; i++) { + dot += a[i] * b[i]; + normA += a[i] * a[i]; + normB += b[i] * b[i]; + } + if (normA == 0.0 || normB == 0.0) { + return 0.0; + } + return dot / (Math.sqrt(normA) * Math.sqrt(normB)); + } + + private static double euclideanDistance(double[] a, double[] b) { + double sum = 0.0; + for (int i = 0; i < a.length; i++) { + double diff = a[i] - b[i]; + sum += diff * diff; + } + return Math.sqrt(sum); + } + + private static double dotProduct(double[] a, double[] b) { + double dot = 0.0; + for (int i = 0; i < a.length; i++) { + dot += a[i] * b[i]; + } + return dot; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/InMemoryVectorStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/InMemoryVectorStore.java new file mode 100644 index 000000000..e0ee63729 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/InMemoryVectorStore.java @@ -0,0 +1,119 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.geaflow.ai.index.vectorstore; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.PriorityQueue; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; + +public class InMemoryVectorStore implements VectorStore { + private final VectorStoreMetadata metadata; + private final ConcurrentHashMap store = new ConcurrentHashMap<>(); + private final Set deletedIds = ConcurrentHashMap.newKeySet(); + + public InMemoryVectorStore(VectorStoreMetadata metadata) { + this.metadata = Objects.requireNonNull(metadata); + } + + @Override + public void upsert(VectorRecord record) { + if (record.getEmbedding().length != metadata.getDimension()) { + throw new VectorStoreException(VectorStoreException.ErrorCode.DIMENSION_MISMATCH, + "Expected dimension " + metadata.getDimension() + ", got " + record.getEmbedding().length); + } + store.put(record.getVectorId(), record); + deletedIds.remove(record.getVectorId()); + } + + @Override + public void upsertBatch(List records) { + for (VectorRecord record : records) { + upsert(record); + } + } + + @Override + public List search(VectorQuery query) { + if (query.getFilterMetadata().containsKey("model_name")) { + String filterModel = query.getFilterMetadata().get("model_name"); + if (!Objects.equals(filterModel, metadata.getModelName())) { + throw new VectorStoreException(VectorStoreException.ErrorCode.MODEL_MISMATCH, + "Expected model " + metadata.getModelName() + ", got " + filterModel); + } + } + + PriorityQueue hits = new PriorityQueue<>(query.getTopK(), Comparator.comparingDouble(VectorHit::getScore)); + for (VectorRecord record : store.values()) { + if (!deletedIds.contains(record.getVectorId())) { + boolean match = true; + for (Map.Entry entry : query.getFilterMetadata().entrySet()) { + if (entry.getKey().equals("model_name")) { + continue; + } + if (!Objects.equals(record.getMetadata().get(entry.getKey()), entry.getValue()) && !Objects.equals(record.getSourceType(), entry.getValue())) { + match = false; + break; + } + } + if (match) { + double score = DistanceUtils.compute(query.getQueryVector(), record.getEmbedding(), metadata.getDistance()); + + if (hits.size() < query.getTopK()) { + hits.add(new VectorHit(record.getVectorId(), score, record)); + } else if (score > hits.peek().getScore()) { + hits.poll(); + hits.add(new VectorHit(record.getVectorId(), score, record)); + } + } + } + } + + List topKHits = new ArrayList<>(); + while (!hits.isEmpty()) { + topKHits.add(hits.poll()); + } + + Collections.reverse(topKHits); + return topKHits; + } + + @Override + public void markDeleted(String vectorId) { + if (!store.containsKey(vectorId)) { + throw new VectorStoreException(VectorStoreException.ErrorCode.RECORD_NOT_FOUND, + "Record not found: " + vectorId); + } + deletedIds.add(vectorId); + } + + @Override + public VectorStoreMetadata getMetadata() { + return metadata; + } + + @Override + public void close() { + + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/LocalVectorStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/LocalVectorStore.java new file mode 100644 index 000000000..73904303e --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/LocalVectorStore.java @@ -0,0 +1,256 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.geaflow.ai.index.vectorstore; + +import com.google.gson.Gson; +import com.google.gson.JsonObject; +import com.google.gson.JsonParser; +import java.io.BufferedReader; +import java.io.BufferedWriter; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.PriorityQueue; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.zip.CRC32; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class LocalVectorStore implements VectorStore { + private static final Logger LOGGER = LoggerFactory.getLogger(LocalVectorStore.class); + + private final VectorStoreMetadata metadata; + private final Path jsonlPath; + private final Path quarantinePath; + private final Gson gson = new Gson(); + private final ConcurrentHashMap store = new ConcurrentHashMap<>(); + private final Set deletedIds = ConcurrentHashMap.newKeySet(); + private BufferedWriter writer; + + public LocalVectorStore(VectorStoreMetadata metadata, Path jsonlPath) { + this.metadata = Objects.requireNonNull(metadata); + this.jsonlPath = Objects.requireNonNull(jsonlPath); + this.quarantinePath = jsonlPath.resolveSibling(jsonlPath.getFileName() + ".quarantine"); + init(); + } + + private void init() { + if (Files.exists(jsonlPath)) { + try (BufferedReader reader = Files.newBufferedReader(jsonlPath, StandardCharsets.UTF_8)) { + String line; + while ((line = reader.readLine()) != null) { + processLine(line); + } + } catch (IOException e) { + throw new VectorStoreException(VectorStoreException.ErrorCode.PERSISTENCE_ERROR, "Failed to read store file", e); + } + } + try { + this.writer = Files.newBufferedWriter(jsonlPath, StandardCharsets.UTF_8, + StandardOpenOption.CREATE, StandardOpenOption.APPEND); + } catch (IOException e) { + throw new VectorStoreException(VectorStoreException.ErrorCode.PERSISTENCE_ERROR, "Failed to open store file for writing", e); + } + } + + private void processLine(String line) { + try { + JsonObject json = new JsonParser().parse(line).getAsJsonObject(); + String expectedChecksum = json.get("__checksum").getAsString(); + + JsonObject dataForChecksum = new JsonParser().parse(line).getAsJsonObject(); + dataForChecksum.remove("__checksum"); + String actualChecksum = computeChecksum(dataForChecksum.toString()); + + if (!expectedChecksum.equals(actualChecksum)) { + quarantineLine(line); + return; + } + + String vectorId = json.get("vectorId").getAsString(); + if (json.has("__deleted") && json.get("__deleted").getAsBoolean()) { + store.remove(vectorId); + deletedIds.add(vectorId); + } else { + VectorRecord record = gson.fromJson(dataForChecksum, VectorRecord.class); + store.put(vectorId, record); + deletedIds.remove(vectorId); + } + } catch (Exception e) { + quarantineLine(line); + } + } + + private void quarantineLine(String line) { + try { + Files.write(quarantinePath, (line + System.lineSeparator()).getBytes(StandardCharsets.UTF_8), + StandardOpenOption.CREATE, StandardOpenOption.APPEND); + } catch (IOException e) { + LOGGER.warn("Ignoring quarantine IOException"); + } + } + + private String computeChecksum(String data) { + CRC32 crc32 = new CRC32(); + crc32.update(data.getBytes(StandardCharsets.UTF_8)); + return Long.toHexString(crc32.getValue()); + } + + private void appendLine(JsonObject json) { + String dataStr = json.toString(); + String checksum = computeChecksum(dataStr); + json.addProperty("__checksum", checksum); + try { + writer.write(json.toString()); + writer.newLine(); + } catch (IOException e) { + throw new VectorStoreException(VectorStoreException.ErrorCode.PERSISTENCE_ERROR, "Failed to write record", e); + } + } + + @Override + public void upsert(VectorRecord record) { + if (record.getEmbedding().length != metadata.getDimension()) { + throw new VectorStoreException(VectorStoreException.ErrorCode.DIMENSION_MISMATCH, + "Expected dimension " + metadata.getDimension() + ", got " + record.getEmbedding().length); + } + JsonObject json = gson.toJsonTree(record).getAsJsonObject(); + appendLine(json); + + store.put(record.getVectorId(), record); + deletedIds.remove(record.getVectorId()); + + try { + writer.flush(); + } catch (IOException e) { + throw new VectorStoreException(VectorStoreException.ErrorCode.PERSISTENCE_ERROR, "Failed to flush record", e); + } + } + + @Override + public void upsertBatch(List records) { + for (VectorRecord record : records) { + if (record.getEmbedding().length != metadata.getDimension()) { + throw new VectorStoreException(VectorStoreException.ErrorCode.DIMENSION_MISMATCH, + "Expected dimension " + metadata.getDimension() + ", got " + record.getEmbedding().length); + } + JsonObject json = gson.toJsonTree(record).getAsJsonObject(); + appendLine(json); + + store.put(record.getVectorId(), record); + deletedIds.remove(record.getVectorId()); + } + try { + writer.flush(); + } catch (IOException e) { + throw new VectorStoreException(VectorStoreException.ErrorCode.PERSISTENCE_ERROR, "Failed to flush batch", e); + } + } + + @Override + public List search(VectorQuery query) { + if (query.getFilterMetadata().containsKey("model_name")) { + String filterModel = query.getFilterMetadata().get("model_name"); + if (!Objects.equals(filterModel, metadata.getModelName())) { + throw new VectorStoreException(VectorStoreException.ErrorCode.MODEL_MISMATCH, + "Expected model " + metadata.getModelName() + ", got " + filterModel); + } + } + + PriorityQueue hits = new PriorityQueue<>(query.getTopK(), Comparator.comparingDouble(VectorHit::getScore)); + + for (VectorRecord record : store.values()) { + if (!deletedIds.contains(record.getVectorId())) { + boolean match = true; + for (Map.Entry entry : query.getFilterMetadata().entrySet()) { + if (entry.getKey().equals("model_name")) { + continue; + } + if (!Objects.equals(record.getMetadata().get(entry.getKey()), entry.getValue()) && !Objects.equals(record.getSourceType(), entry.getValue())) { + match = false; + break; + } + } + if (match) { + double score = DistanceUtils.compute(query.getQueryVector(), record.getEmbedding(), metadata.getDistance()); + if (hits.size() < query.getTopK()) { + hits.add(new VectorHit(record.getVectorId(), score, record)); + } else if (!hits.isEmpty() && score > hits.peek().getScore()) { + hits.poll(); + hits.add(new VectorHit(record.getVectorId(), score, record)); + } + } + } + } + + List topKHits = new ArrayList<>(); + while (!hits.isEmpty()) { + topKHits.add(hits.poll()); + } + + Collections.reverse(topKHits); + return topKHits; + } + + @Override + public void markDeleted(String vectorId) { + if (!store.containsKey(vectorId)) { + throw new VectorStoreException(VectorStoreException.ErrorCode.RECORD_NOT_FOUND, + "Record not found: " + vectorId); + } + JsonObject json = new JsonObject(); + json.addProperty("vectorId", vectorId); + json.addProperty("__deleted", true); + appendLine(json); + + try { + writer.flush(); + } catch (IOException e) { + throw new VectorStoreException(VectorStoreException.ErrorCode.PERSISTENCE_ERROR, "Failed to flush deletion", e); + } + + store.remove(vectorId); + deletedIds.add(vectorId); + } + + @Override + public VectorStoreMetadata getMetadata() { + return metadata; + } + + @Override + public void close() { + try { + if (writer != null) { + writer.close(); + } + } catch (IOException e) { + LOGGER.warn("Failed to close VectorStore cleanly", e); + } + + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorHit.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorHit.java new file mode 100644 index 000000000..80232d31f --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorHit.java @@ -0,0 +1,42 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.geaflow.ai.index.vectorstore; + +public class VectorHit { + private final String vectorId; + private final double score; + private final VectorRecord record; + + public VectorHit(String vectorId, double score, VectorRecord record) { + this.vectorId = vectorId; + this.score = score; + this.record = record; + } + + public String getVectorId() { + return vectorId; + } + + public double getScore() { + return score; + } + + public VectorRecord getRecord() { + return record; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorQuery.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorQuery.java new file mode 100644 index 000000000..5ca3fd58e --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorQuery.java @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.geaflow.ai.index.vectorstore; + +import java.util.Collections; +import java.util.Map; + +public class VectorQuery { + private final double[] queryVector; + private final int topK; + private final Map filterMetadata; + + public VectorQuery(double[] queryVector, int topK, Map filterMetadata) { + if (queryVector == null || queryVector.length == 0) { + throw new IllegalArgumentException("queryVector cannot be null or empty"); + } + if (topK <= 0) { + throw new IllegalArgumentException("topK must be greater than 0"); + } + this.queryVector = queryVector; + this.topK = topK; + this.filterMetadata = filterMetadata == null ? Collections.emptyMap() : filterMetadata; + } + + public double[] getQueryVector() { + return queryVector; + } + + public int getTopK() { + return topK; + } + + public Map getFilterMetadata() { + return filterMetadata; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorRecord.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorRecord.java new file mode 100644 index 000000000..a7969fac1 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorRecord.java @@ -0,0 +1,69 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.geaflow.ai.index.vectorstore; + +import java.util.Collections; +import java.util.Map; + +public class VectorRecord { + private final String vectorId; + private final double[] embedding; + private final String sourceType; + private final String sourceId; + private final Map metadata; + + public VectorRecord(String vectorId, double[] embedding, String sourceType, String sourceId, Map metadata) { + if (vectorId == null || vectorId.isEmpty()) { + throw new IllegalArgumentException("vectorId cannot be null or empty"); + } + if (embedding == null || embedding.length == 0) { + throw new IllegalArgumentException("embedding cannot be null or empty"); + } + if (sourceType == null || sourceType.isEmpty()) { + throw new IllegalArgumentException("sourceType cannot be null or empty"); + } + if (sourceId == null || sourceId.isEmpty()) { + throw new IllegalArgumentException("sourceId cannot be null or empty"); + } + this.vectorId = vectorId; + this.embedding = embedding; + this.sourceType = sourceType; + this.sourceId = sourceId; + this.metadata = metadata == null ? Collections.emptyMap() : metadata; + } + + public String getVectorId() { + return vectorId; + } + + public double[] getEmbedding() { + return embedding; + } + + public String getSourceType() { + return sourceType; + } + + public String getSourceId() { + return sourceId; + } + + public Map getMetadata() { + return metadata; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStore.java index 175e8c412..730b6e2ae 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStore.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStore.java @@ -1,71 +1,66 @@ /* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * http://www.apache.org/licenses/LICENSE-2.0 + * http://www.apache.org/licenses/LICENSE-2.0 * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.geaflow.ai.index.vectorstore; import java.util.List; -import org.apache.commons.lang3.tuple.Pair; -import org.apache.geaflow.ai.index.vector.IVector; public interface VectorStore { /** - * Get the metadata of this vector store. + * Insert or update a single vector record. + * Upsert is idempotent: same vectorId overwrites the previous record. * - * @return The vector store metadata. + * @throws VectorStoreException if record.embedding.length != metadata.dimension (DIMENSION_MISMATCH) + * @throws VectorStoreException if required metadata fields are missing (METADATA_INCOMPLETE) */ - VectorStoreMetadata getMetadata(); + void upsert(VectorRecord record); /** - * Initialize the vector store. - * Implementations should validate the expected metadata against the stored metadata. - * - * @param expectedMetadata The expected metadata for validation. + * Batch insert/update. Atomic per record; partial failure does not + * roll back already-committed records. */ - void init(VectorStoreMetadata expectedMetadata); + void upsertBatch(List records); /** - * Add a vector to the store with a given ID. + * Search for nearest vectors using the configured distance metric. + * Results are ordered by score descending (best match first). * - * @param id The identifier for the vector. - * @param vector The vector to store. + * @return empty list if no results match, never null + * @throws VectorStoreException if filterMetadata specifies + * a model_name different from the store's model_name (MODEL_MISMATCH) */ - void add(String id, IVector vector); + List search(VectorQuery query); /** - * Delete a vector by ID. + * Soft-delete a vector by id. The record remains on disk but is + * excluded from future search results. * - * @param id The identifier for the vector. + * @throws VectorStoreException if vectorId does not exist (RECORD_NOT_FOUND) */ - void delete(String id); + void markDeleted(String vectorId); /** - * Search the top-K closest vectors to the query vector. - * - * @param queryVector The query vector. - * @param topK The number of results to return. - * @return A list of pairs containing the vector ID and distance score. + * Returns metadata about this store instance. */ - List> search(IVector queryVector, int topK); + VectorStoreMetadata getMetadata(); /** - * Close the vector store and release any resources. + * Release resources. Implementations must flush pending writes before returning. */ void close(); } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreException.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreException.java new file mode 100644 index 000000000..a8881f00e --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreException.java @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.geaflow.ai.index.vectorstore; + +public class VectorStoreException extends RuntimeException { + public enum ErrorCode { + DIMENSION_MISMATCH, + METADATA_INCOMPLETE, + MODEL_MISMATCH, + FORMAT_VERSION_UNKNOWN, + RECORD_NOT_FOUND, + PERSISTENCE_ERROR + } + + private final ErrorCode errorCode; + private final String detail; + + public VectorStoreException(ErrorCode errorCode, String detail) { + super(errorCode.name() + ": " + detail); + this.errorCode = errorCode; + this.detail = detail; + } + + public VectorStoreException(ErrorCode errorCode, String detail, Throwable cause) { + super(errorCode.name() + ": " + detail, cause); + this.errorCode = errorCode; + this.detail = detail; + } + + public ErrorCode getErrorCode() { + return errorCode; + } + + public String getDetail() { + return detail; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadata.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadata.java index 9335312c9..dd20dd748 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadata.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadata.java @@ -1,60 +1,53 @@ /* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at * - * http://www.apache.org/licenses/LICENSE-2.0 + * http://www.apache.org/licenses/LICENSE-2.0 * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. */ package org.apache.geaflow.ai.index.vectorstore; -import com.google.gson.Gson; -import com.google.gson.GsonBuilder; import com.google.gson.annotations.SerializedName; -import java.io.BufferedReader; -import java.io.BufferedWriter; -import java.io.IOException; -import java.nio.charset.StandardCharsets; -import java.nio.file.Files; -import java.nio.file.Path; import java.util.Objects; public class VectorStoreMetadata { - @SerializedName("model_name") - private String modelName; - + private final String modelName; @SerializedName("dimension") - private int dimension; - + private final int dimension; @SerializedName("distance") - private String distance; - + private final DistanceMetric distance; @SerializedName("index_version") - private String indexVersion; - + private final String indexVersion; @SerializedName("created_at") - private long createdAt; - + private final long createdAt; @SerializedName("format_version") - private String formatVersion; + private final int formatVersion; - public VectorStoreMetadata() { - } - - public VectorStoreMetadata(String modelName, int dimension, String distance, - String indexVersion, long createdAt, String formatVersion) { + public VectorStoreMetadata(String modelName, int dimension, DistanceMetric distance, + String indexVersion, long createdAt, int formatVersion) { + if (modelName == null || modelName.isEmpty()) { + throw new IllegalArgumentException("modelName cannot be null or empty"); + } + if (dimension <= 0) { + throw new IllegalArgumentException("dimension must be greater than 0"); + } + if (distance == null) { + throw new IllegalArgumentException("distance cannot be null"); + } + if (indexVersion == null || indexVersion.isEmpty()) { + throw new IllegalArgumentException("indexVersion cannot be null or empty"); + } this.modelName = modelName; this.dimension = dimension; this.distance = distance; @@ -67,105 +60,26 @@ public String getModelName() { return modelName; } - public void setModelName(String modelName) { - this.modelName = modelName; - } - public int getDimension() { return dimension; } - public void setDimension(int dimension) { - this.dimension = dimension; - } - - public String getDistance() { + public DistanceMetric getDistance() { return distance; } - public void setDistance(String distance) { - this.distance = distance; - } - public String getIndexVersion() { return indexVersion; } - public void setIndexVersion(String indexVersion) { - this.indexVersion = indexVersion; - } - public long getCreatedAt() { return createdAt; } - public void setCreatedAt(long createdAt) { - this.createdAt = createdAt; - } - - public String getFormatVersion() { + public int getFormatVersion() { return formatVersion; } - public void setFormatVersion(String formatVersion) { - this.formatVersion = formatVersion; - } - - public void validate(VectorStoreMetadata expected) { - if (expected == null) { - return; - } - if (this.dimension != expected.dimension) { - throw new IllegalArgumentException(String.format( - "Dimension mismatch. Expected %d, but got %d", expected.dimension, this.dimension)); - } - if (!Objects.equals(this.modelName, expected.modelName)) { - throw new IllegalArgumentException(String.format( - "Model mismatch. Expected %s, but got %s", expected.modelName, this.modelName)); - } - } - - public void validateComplete() { - if (modelName == null || modelName.isEmpty()) { - throw new IllegalArgumentException("Missing metadata: model_name"); - } - if (dimension <= 0) { - throw new IllegalArgumentException("Missing or invalid metadata: dimension"); - } - if (distance == null || distance.isEmpty()) { - throw new IllegalArgumentException("Missing metadata: distance"); - } - if (indexVersion == null || indexVersion.isEmpty()) { - throw new IllegalArgumentException("Missing metadata: index_version"); - } - if (formatVersion == null || formatVersion.isEmpty()) { - throw new IllegalArgumentException("Missing metadata: format_version"); - } - } - - public static VectorStoreMetadata load(Path metadataPath) throws IOException { - if (!Files.exists(metadataPath)) { - throw new IOException("Metadata file does not exist: " + metadataPath.toAbsolutePath()); - } - try (BufferedReader reader = Files.newBufferedReader(metadataPath, StandardCharsets.UTF_8)) { - Gson gson = new Gson(); - VectorStoreMetadata metadata = gson.fromJson(reader, VectorStoreMetadata.class); - if (metadata == null) { - throw new IOException("Failed to parse metadata from file"); - } - metadata.validateComplete(); - return metadata; - } - } - - public void save(Path metadataPath) throws IOException { - validateComplete(); - try (BufferedWriter writer = Files.newBufferedWriter(metadataPath, StandardCharsets.UTF_8)) { - Gson gson = new GsonBuilder().setPrettyPrinting().create(); - gson.toJson(this, writer); - } - } - @Override public boolean equals(Object o) { if (this == o) { @@ -177,10 +91,10 @@ public boolean equals(Object o) { VectorStoreMetadata that = (VectorStoreMetadata) o; return dimension == that.dimension && createdAt == that.createdAt + && formatVersion == that.formatVersion && Objects.equals(modelName, that.modelName) - && Objects.equals(distance, that.distance) - && Objects.equals(indexVersion, that.indexVersion) - && Objects.equals(formatVersion, that.formatVersion); + && distance == that.distance + && Objects.equals(indexVersion, that.indexVersion); } @Override diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadataTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadataTest.java deleted file mode 100644 index 9f310cea1..000000000 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadataTest.java +++ /dev/null @@ -1,128 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -package org.apache.geaflow.ai.index.vectorstore; - -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -import java.io.IOException; -import java.nio.charset.StandardCharsets; -import java.nio.file.Files; -import java.nio.file.Path; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; - -public class VectorStoreMetadataTest { - - private Path tempFile; - - @BeforeEach - public void setUp() throws IOException { - tempFile = Files.createTempFile("metadata_test", ".json"); - } - - @AfterEach - public void tearDown() throws IOException { - Files.deleteIfExists(tempFile); - } - - @Test - public void testPersistAndLoadMetadata() throws IOException { - VectorStoreMetadata metadata = new VectorStoreMetadata( - "test_model", - 128, - "COSINE", - "1.0", - System.currentTimeMillis(), - "v1" - ); - metadata.save(tempFile); - - VectorStoreMetadata loaded = VectorStoreMetadata.load(tempFile); - assertEquals(metadata, loaded); - assertEquals("test_model", loaded.getModelName()); - assertEquals(128, loaded.getDimension()); - } - - @Test - public void testMissingMetadataValidationFails() { - VectorStoreMetadata metadata = new VectorStoreMetadata(); - assertThrows(IllegalArgumentException.class, metadata::validateComplete); - - metadata.setModelName("test_model"); - assertThrows(IllegalArgumentException.class, metadata::validateComplete); - - metadata.setDimension(128); - assertThrows(IllegalArgumentException.class, metadata::validateComplete); - - metadata.setDistance("L2"); - assertThrows(IllegalArgumentException.class, metadata::validateComplete); - - metadata.setIndexVersion("1.0"); - assertThrows(IllegalArgumentException.class, metadata::validateComplete); - - metadata.setFormatVersion("v1"); - // Should not throw now - metadata.validateComplete(); - } - - @Test - public void testDimensionMismatchFails() { - VectorStoreMetadata expected = new VectorStoreMetadata("test_model", 128, "L2", "1.0", 0, "v1"); - VectorStoreMetadata actual = new VectorStoreMetadata("test_model", 256, "L2", "1.0", 0, "v1"); - - IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, () -> { - actual.validate(expected); - }); - assertEquals("Dimension mismatch. Expected 128, but got 256", ex.getMessage()); - } - - @Test - public void testModelMismatchFails() { - VectorStoreMetadata expected = new VectorStoreMetadata("expected_model", 128, "L2", "1.0", 0, "v1"); - VectorStoreMetadata actual = new VectorStoreMetadata("actual_model", 128, "L2", "1.0", 0, "v1"); - - IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, () -> { - actual.validate(expected); - }); - assertEquals("Model mismatch. Expected expected_model, but got actual_model", ex.getMessage()); - } - - @Test - public void testLoadInvalidJsonFails() throws IOException { - Files.write(tempFile, "{ \"dimension\": 128, \"model_name\": \"test\" ".getBytes(StandardCharsets.UTF_8)); // Invalid JSON - - assertThrows(com.google.gson.JsonSyntaxException.class, () -> { - VectorStoreMetadata.load(tempFile); - }); - } - - @Test - public void testLoadMissingFieldsFails() throws IOException { - // Missing distance, index_version, format_version - Files.write(tempFile, "{ \"dimension\": 128, \"model_name\": \"test\" }".getBytes(StandardCharsets.UTF_8)); - - assertThrows(IllegalArgumentException.class, () -> { - VectorStoreMetadata.load(tempFile); - }); - } -} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreTest.java new file mode 100644 index 000000000..d34e5db4e --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreTest.java @@ -0,0 +1,111 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.geaflow.ai.index.vectorstore; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Collections; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +public class VectorStoreTest { + + private Path tempFile; + private VectorStoreMetadata metadata; + + @BeforeEach + public void setUp() throws IOException { + tempFile = Files.createTempFile("vector_store_test", ".jsonl"); + metadata = new VectorStoreMetadata("test_model", 128, DistanceMetric.COSINE, "1.0", System.currentTimeMillis(), 1); + } + + @AfterEach + public void tearDown() throws IOException { + Files.deleteIfExists(tempFile); + Files.deleteIfExists(tempFile.resolveSibling(tempFile.getFileName() + ".quarantine")); + } + + @Test + public void testMetadataValidation() { + assertThrows(IllegalArgumentException.class, () -> { + new VectorStoreMetadata("", 128, DistanceMetric.COSINE, "1.0", 0, 1); + }); + + assertThrows(IllegalArgumentException.class, () -> { + new VectorStoreMetadata("test", 0, DistanceMetric.COSINE, "1.0", 0, 1); + }); + + assertThrows(IllegalArgumentException.class, () -> { + new VectorStoreMetadata("test", 128, null, "1.0", 0, 1); + }); + } + + @Test + public void testInMemoryVectorStore() { + VectorStore store = new InMemoryVectorStore(metadata); + testVectorStore(store); + } + + @Test + public void testLocalVectorStore() { + VectorStore store = new LocalVectorStore(metadata, tempFile); + testVectorStore(store); + } + + private void testVectorStore(VectorStore store) { + double[] embedding = new double[128]; + embedding[0] = 1.0; + + VectorRecord record = new VectorRecord("v1", embedding, "chunk", "c1", Collections.emptyMap()); + store.upsert(record); + + VectorQuery query = new VectorQuery(embedding, 10, Collections.emptyMap()); + List hits = store.search(query); + + assertEquals(1, hits.size()); + assertEquals("v1", hits.get(0).getVectorId()); + + // Dimension mismatch + double[] badEmbedding = new double[64]; + VectorRecord badRecord = new VectorRecord("v2", badEmbedding, "chunk", "c2", Collections.emptyMap()); + + assertThrows(VectorStoreException.class, () -> { + store.upsert(badRecord); + }); + + // Model mismatch + VectorQuery badQuery = new VectorQuery(embedding, 10, Collections.singletonMap("model_name", "wrong_model")); + assertThrows(VectorStoreException.class, () -> { + store.search(badQuery); + }); + + // Delete + store.markDeleted("v1"); + hits = store.search(query); + assertEquals(0, hits.size()); + + store.close(); + } +}