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 new file mode 100644 index 000000000..730b6e2ae --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStore.java @@ -0,0 +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 + * + * 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; + +public interface VectorStore { + + /** + * Insert or update a single vector record. + * Upsert is idempotent: same vectorId overwrites the previous record. + * + * @throws VectorStoreException if record.embedding.length != metadata.dimension (DIMENSION_MISMATCH) + * @throws VectorStoreException if required metadata fields are missing (METADATA_INCOMPLETE) + */ + void upsert(VectorRecord record); + + /** + * Batch insert/update. Atomic per record; partial failure does not + * roll back already-committed records. + */ + void upsertBatch(List records); + + /** + * Search for nearest vectors using the configured distance metric. + * Results are ordered by score descending (best match first). + * + * @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) + */ + List search(VectorQuery query); + + /** + * Soft-delete a vector by id. The record remains on disk but is + * excluded from future search results. + * + * @throws VectorStoreException if vectorId does not exist (RECORD_NOT_FOUND) + */ + void markDeleted(String vectorId); + + /** + * Returns metadata about this store instance. + */ + VectorStoreMetadata getMetadata(); + + /** + * 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 new file mode 100644 index 000000000..dd20dd748 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/vectorstore/VectorStoreMetadata.java @@ -0,0 +1,104 @@ +/* + * 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.annotations.SerializedName; +import java.util.Objects; + +public class VectorStoreMetadata { + @SerializedName("model_name") + private final String modelName; + @SerializedName("dimension") + private final int dimension; + @SerializedName("distance") + private final DistanceMetric distance; + @SerializedName("index_version") + private final String indexVersion; + @SerializedName("created_at") + private final long createdAt; + @SerializedName("format_version") + private final int 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; + this.indexVersion = indexVersion; + this.createdAt = createdAt; + this.formatVersion = formatVersion; + } + + public String getModelName() { + return modelName; + } + + public int getDimension() { + return dimension; + } + + public DistanceMetric getDistance() { + return distance; + } + + public String getIndexVersion() { + return indexVersion; + } + + public long getCreatedAt() { + return createdAt; + } + + public int getFormatVersion() { + return formatVersion; + } + + @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 + && formatVersion == that.formatVersion + && Objects.equals(modelName, that.modelName) + && distance == that.distance + && Objects.equals(indexVersion, that.indexVersion); + } + + @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/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(); + } +}