diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/ann/LinearNearestNeighborIndex.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/ann/LinearNearestNeighborIndex.java new file mode 100644 index 000000000..ecac74213 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/ann/LinearNearestNeighborIndex.java @@ -0,0 +1,94 @@ +/* + * 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.ann; + +import java.util.ArrayList; +import java.util.List; +import org.apache.geaflow.ai.index.vector.IVector; + +/** + * Exact linear scan over the added entries — the default {@link NearestNeighborIndex} adapter for + * Phase 1 and the reference implementation an ANN adapter is measured against. + * + *

Every search scores every entry through {@code entryVector.match(query)} and sorts by the + * interface contract: score descending, ties broken by insertion order. The sort is stable over + * the insertion sequence, which is what makes equal-scored neighbors come out in the order they + * were added. Results are a fresh list; mutating it does not affect the index. + * + *

This adapter makes no performance claim: it is the correctness baseline. Swap-in of an ANN + * implementation is a change behind the same interface, not an optimization of this class. + * + * @param id type of the indexed entries + */ +public class LinearNearestNeighborIndex implements NearestNeighborIndex { + + private final List> entries = new ArrayList<>(); + + @Override + public void add(T id, IVector vector) { + if (id == null || vector == null) { + throw new IllegalArgumentException("id and vector must not be null"); + } + entries.add(new Entry<>(id, vector)); + } + + @Override + public List> search(IVector query, int topK) { + if (query == null) { + throw new IllegalArgumentException("query must not be null"); + } + if (topK <= 0) { + throw new IllegalArgumentException("topK must be positive, got " + topK); + } + double[] scores = new double[entries.size()]; + Integer[] order = new Integer[entries.size()]; + for (int i = 0; i < entries.size(); i++) { + scores[i] = entries.get(i).vector.match(query); + order[i] = i; + } + // Tie-break key: the insertion ordinal travels alongside the score so equal-scored + // neighbors keep the order they were added in, independent of sort stability. + java.util.Arrays.sort(order, (a, b) -> { + int byScore = Double.compare(scores[b], scores[a]); + return byScore != 0 ? byScore : Integer.compare(a, b); + }); + List> result = new ArrayList<>(Math.min(topK, order.length)); + for (int i = 0; i < Math.min(topK, order.length); i++) { + Entry entry = entries.get(order[i]); + result.add(new Neighbor<>(entry.id, scores[order[i]])); + } + return result; + } + + @Override + public int size() { + return entries.size(); + } + + private static final class Entry { + private final T id; + private final IVector vector; + + private Entry(T id, IVector vector) { + this.id = id; + this.vector = vector; + } + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/ann/NearestNeighborIndex.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/ann/NearestNeighborIndex.java new file mode 100644 index 000000000..6e00d6053 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/ann/NearestNeighborIndex.java @@ -0,0 +1,76 @@ +/* + * 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.ann; + +import java.util.List; +import org.apache.geaflow.ai.index.vector.IVector; + +/** + * Nearest-neighbor search over {@link IVector} entries, keyed by an opaque id. + * + *

Phase 1 ships this contract with the exact {@link LinearNearestNeighborIndex} only — vector + * search today is a linear TopN over the candidate set, and this interface exists so an ANN + * implementation (HNSW and friends) can slot in behind the same calls later without touching the + * retrieval path. No performance claim is made or implied by implementing this interface. + * + *

Contract. Implementations are not required to be thread-safe. {@link #add} appends an + * entry; adding the same id twice keeps both entries (append-only semantics, like most ANN + * structures). {@link #search}: + * + *

+ * + *

Score direction is fixed as {@code entryVector.match(query)} — never the other way around — + * because some {@link IVector} implementations are not symmetric in their arguments. + * + *

An exact adapter must score every entry with {@link IVector#match} and honor the full + * ordering above. An approximate adapter may miss neighbors and coarsen scores, but must still + * respect the result shape: at most {@code topK} entries, no id appearing twice unless it was + * added twice, and a non-increasing score order within the returned page. + * + * @param id type of the indexed entries + */ +public interface NearestNeighborIndex { + + /** + * Appends one entry to the index. + * + * @param id caller-defined id returned in search results + * @param vector representation the entry is searched by + */ + void add(T id, IVector vector); + + /** + * Returns up to {@code topK} nearest entries for the query, ordered by the contract above. + * + * @param query vector to search against + * @param topK maximum number of results, must be positive + * @return ranked neighbors, highest score first + */ + List> search(IVector query, int topK); + + /** Number of entries currently held (an id added twice counts twice). */ + int size(); +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/ann/Neighbor.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/ann/Neighbor.java new file mode 100644 index 000000000..55a85e0e4 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/ann/Neighbor.java @@ -0,0 +1,67 @@ +/* + * 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.ann; + +import java.util.Objects; + +/** + * One search result of a {@link NearestNeighborIndex}: the caller-supplied id plus its similarity + * score against the query, as produced by {@link org.apache.geaflow.ai.index.vector.IVector#match}. + */ +public final class Neighbor { + + private final T id; + private final double score; + + public Neighbor(T id, double score) { + this.id = Objects.requireNonNull(id); + this.score = score; + } + + public T getId() { + return id; + } + + public double getScore() { + return score; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + Neighbor other = (Neighbor) o; + return Double.compare(score, other.score) == 0 && id.equals(other.id); + } + + @Override + public int hashCode() { + return Objects.hash(id, score); + } + + @Override + public String toString() { + return "Neighbor{" + "id=" + id + ", score=" + score + '}'; + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/ann/LinearNearestNeighborIndexTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/ann/LinearNearestNeighborIndexTest.java new file mode 100644 index 000000000..f34430122 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/ann/LinearNearestNeighborIndexTest.java @@ -0,0 +1,58 @@ +/* + * 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.ann; + +import org.apache.geaflow.ai.index.vector.EmbeddingVector; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** + * The linear adapter runs the full {@link NearestNeighborIndexContractTest} suite — including the + * exactness tests — plus its own argument validation. This is the adapter a future ANN + * implementation is measured against: same tests, same orderings, no approximations. + */ +public class LinearNearestNeighborIndexTest extends NearestNeighborIndexContractTest { + + @Override + protected NearestNeighborIndex newIndex() { + return new LinearNearestNeighborIndex<>(); + } + + @Test + public void rejectsNullArguments() { + LinearNearestNeighborIndex index = new LinearNearestNeighborIndex<>(); + Assertions.assertThrows(IllegalArgumentException.class, + () -> index.add(null, new EmbeddingVector(new double[]{1, 0}))); + Assertions.assertThrows(IllegalArgumentException.class, + () -> index.add("a", null)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> index.search(null, 1)); + } + + /** The returned page is a fresh list: caller-side mutation cannot corrupt the index. */ + @Test + public void returnedPageIsIndependentOfTheIndex() { + LinearNearestNeighborIndex index = new LinearNearestNeighborIndex<>(); + index.add("a", new EmbeddingVector(new double[]{1, 0})); + index.add("b", new EmbeddingVector(new double[]{0, 1})); + index.search(new EmbeddingVector(new double[]{1, 0}), 2).clear(); + Assertions.assertEquals(2, index.search(new EmbeddingVector(new double[]{1, 0}), 2).size()); + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/ann/NearestNeighborIndexContractTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/ann/NearestNeighborIndexContractTest.java new file mode 100644 index 000000000..b39ff0c57 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/index/ann/NearestNeighborIndexContractTest.java @@ -0,0 +1,175 @@ +/* + * 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.ann; + +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.geaflow.ai.index.vector.EmbeddingVector; +import org.apache.geaflow.ai.index.vector.IVector; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** + * Contract every {@link NearestNeighborIndex} must hold, written against a fresh index supplied + * by {@link #newIndex()} so any future adapter runs the identical suite — that is the issue #854 + * acceptance criterion "the linear adapter passes the same contract that an ANN adapter will use". + * + *

The tests in {@link #resultShape()} hold for any adapter, exact or approximate. The tests in + * {@link #rankingIsExact()} assert the exact ordering — score from {@link IVector#match}, + * descending, ties on insertion order — and are required only of adapters that claim exactness. + * An ANN implementation should extend this class, keep the shape tests, and override + * {@link #rankingIsExact()} (or exclude it) while its recall is approximate; the distinction is + * deliberate and documented in the interface. + */ +public abstract class NearestNeighborIndexContractTest { + + /** Fresh, empty index of the implementation under test. */ + protected abstract NearestNeighborIndex newIndex(); + + private static EmbeddingVector vec(double... components) { + return new EmbeddingVector(components); + } + + /** + * Shape rules that must hold for every adapter: topK bound, result-page score monotonicity, + * empty-index behavior, argument validation, append-only size accounting. + */ + @Test + public void resultShape() { + NearestNeighborIndex index = newIndex(); + Assertions.assertEquals(0, index.size()); + Assertions.assertTrue(index.search(vec(1, 0), 5).isEmpty(), + "an empty index returns no neighbors"); + + index.add("a", vec(1, 0)); + index.add("b", vec(0, 1)); + index.add("a", vec(1, 1)); + Assertions.assertEquals(3, index.size(), + "an id added twice is two entries, per the append-only contract"); + + List> page = index.search(vec(1, 0), 2); + Assertions.assertEquals(2, page.size(), "at most topK neighbors are returned"); + Assertions.assertTrue(page.get(0).getScore() >= page.get(1).getScore(), + "the returned page is ordered by non-increasing score"); + + Assertions.assertEquals(3, index.search(vec(1, 0), 10).size(), + "topK beyond the index size returns everything"); + Assertions.assertThrows(IllegalArgumentException.class, + () -> index.search(vec(1, 0), 0), "topK must be positive"); + Assertions.assertThrows(IllegalArgumentException.class, + () -> index.search(vec(1, 0), -1), "topK must be positive"); + } + + /** + * The exactness rules: scores equal {@code entryVector.match(query)} verbatim, the global + * order is score descending with ties on insertion order, and repeated searches over + * unchanged data are bit-for-bit identical. + */ + @Test + public void rankingIsExact() { + NearestNeighborIndex index = newIndex(); + // Cosine similarity against (1, 0): exactly 1, exactly 0, and 1/sqrt(2) in that order — + // three distinct scores with hand-checkable values. + index.add("east", vec(1, 0)); + index.add("north", vec(0, 1)); + index.add("diagonal", vec(1, 1)); + // Zero vector: cosine is undefined and EmbeddingVector.match returns 0.0 for it. + index.add("origin", vec(0, 0)); + + List> ranked = index.search(vec(1, 0), 4); + Assertions.assertEquals(4, ranked.size()); + Assertions.assertAll( + () -> Assertions.assertEquals("east", ranked.get(0).getId()), + () -> Assertions.assertEquals(1.0, ranked.get(0).getScore(), 1e-12), + () -> Assertions.assertEquals("diagonal", ranked.get(1).getId()), + () -> Assertions.assertEquals(1.0 / Math.sqrt(2), ranked.get(1).getScore(), 1e-12), + () -> Assertions.assertEquals("north", ranked.get(2).getId()), + () -> Assertions.assertEquals(0.0, ranked.get(2).getScore(), 1e-12), + () -> Assertions.assertEquals("origin", ranked.get(3).getId()), + () -> Assertions.assertEquals(0.0, ranked.get(3).getScore(), 1e-12)); + + List> again = index.search(vec(1, 0), 4); + Assertions.assertEquals(ranked, again, + "search is deterministic: identical query and data give identical results"); + } + + /** + * The tie-break rule the issue asks to be documented and pinned: equal scores come out in + * insertion order. "third" and "twin" are collinear with the query (identical cosine of 1), + * and the three orthogonal axes all score exactly 0 — each group must rank in insertion + * order, and no member of the zero group may jump ahead of the collinear group. + */ + @Test + public void tiesBreakOnInsertionOrder() { + NearestNeighborIndex index = newIndex(); + index.add("second", vec(0, 1, 0, 0)); + index.add("fourth", vec(0, 0, 0, 1)); + index.add("first", vec(1, 0, 0, 0)); + index.add("third", vec(0, 0, 1, 0)); + index.add("twin", vec(0, 0, 2, 0)); + + List> ranked = index.search(vec(0, 0, 1, 0), 5); + Assertions.assertEquals( + java.util.Arrays.asList("third", "twin", "second", "fourth", "first"), + ids(ranked), + "collinear entries first (insertion order among themselves), " + + "then the zero-scored ones, again in insertion order"); + Assertions.assertEquals(ranked.get(0).getScore(), ranked.get(1).getScore(), 0.0, + "collinear vectors produce exactly equal scores"); + } + + /** + * The page window slides by score: raising topK never reorders what a smaller topK returned — + * the first k entries of the topK=n page equal the topK=k page. + */ + @Test + public void largerTopKExtendsWithoutReordering() { + NearestNeighborIndex index = newIndex(); + for (int i = 0; i < 8; i++) { + index.add("v" + i, vec(i, 8 - i)); + } + List> top3 = index.search(vec(1, 1), 3); + List> top8 = index.search(vec(1, 1), 8); + Assertions.assertEquals(ids(top3), ids(top8).subList(0, 3), + "a larger page must extend, not reorder"); + } + + /** No id may appear more often in a page than it was added — append-only accounting. */ + @Test + public void pagesDoNotDuplicateEntries() { + NearestNeighborIndex index = newIndex(); + index.add("dup", vec(1, 0)); + index.add("dup", vec(1, 0)); + index.add("other", vec(1, 0)); + List> page = index.search(vec(1, 0), 3); + Set uniqueIds = new HashSet<>(ids(page)); + Assertions.assertEquals(2, uniqueIds.size(), "ids collapse to distinct values"); + Assertions.assertEquals(3, page.size(), "both dup entries and other are returned"); + } + + private static List ids(List> page) { + List ids = new java.util.ArrayList<>(page.size()); + for (Neighbor neighbor : page) { + ids.add(neighbor.getId()); + } + return ids; + } +}