From 9f660ec3079b3ae049e81062075837bf5b6da3ca Mon Sep 17 00:00:00 2001 From: Sumit Soreng Date: Tue, 25 Aug 2026 00:05:35 +0530 Subject: [PATCH] [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); + }); + } +}