Merge pull request #5 from dev-diaries41/feat/query

feat: support retrieval using start and end indices

Author: dev-diaries41

core/src/main/java/com/fpf/smartscansdk/core/ml/embeddings/EmbeddingTypes.kt

@@ -22,7 +22,7 @@ interface IEmbeddingStore {
     val exists: Boolean
     suspend fun add(newEmbeddings: List<Embedding>)
     suspend fun remove(ids: List<Long>)
-    suspend fun getAll(): List<Embedding> // getAll used instead of get to make clear that loading full index in memory is required
+    suspend fun get(): List<Embedding>
     fun clear()
 }
 

extensions/src/main/java/com/fpf/smartscansdk/extensions/embeddings/FileEmbeddingRetriever.kt

@@ -8,26 +8,43 @@ import com.fpf.smartscansdk.core.ml.embeddings.getTopN
 class FileEmbeddingRetriever(
     private val store: FileEmbeddingStore
 ): IRetriever {
+
+    private var cachedIds: List<Long>? = null
+
     override suspend fun query(
         embedding: FloatArray,
         topK: Int,
         threshold: Float
     ): List<Embedding> {
 
-        val storedEmbeddings = store.getAll()
+        cachedIds = null // clear on new search
 
-        if (storedEmbeddings.isEmpty()) {
-            return emptyList()
-        }
+        val storedEmbeddings = store.get()
+
+        if (storedEmbeddings.isEmpty()) return emptyList()
 
         val similarities = getSimilarities(embedding, storedEmbeddings.map { it.embeddings })
-        val results = getTopN(similarities, topK, threshold)
+        val resultIndices = getTopN(similarities, topK, threshold)
 
-        if (results.isEmpty()) {
-            return emptyList()
+        if (resultIndices.isEmpty()) return emptyList()
+
+        val idsToCache = mutableListOf<Long>()
+        val results = resultIndices.map{idx ->
+            idsToCache.add( storedEmbeddings[idx].id)
+            storedEmbeddings[idx]
         }
+        cachedIds = idsToCache
+        return results
+    }
+
+    suspend fun query(start: Int, end: Int): List<Embedding> {
+        val ids = cachedIds ?: return emptyList()
+        val s = start.coerceAtLeast(0)
+        val e = end.coerceAtMost(ids.size)
+        if (s >= e) return emptyList()
 
-        return results.map{idx -> storedEmbeddings[idx]}
+        val batch = store.get(ids.subList(s, e))
+        return batch
     }
 }
 

extensions/src/main/java/com/fpf/smartscansdk/extensions/embeddings/FileEmbeddingStore.kt

@@ -78,7 +78,7 @@ class FileEmbeddingStore(
     }
 
     // This explicitly makes clear the design constraints that requires the full index to be loaded in memory
-    override suspend fun getAll(): List<Embedding> = withContext(Dispatchers.IO){
+    override suspend fun get(): List<Embedding> = withContext(Dispatchers.IO){
         cache?.let { return@withContext it.values.toList() };
 
         FileInputStream(file).channel.use { ch ->
@@ -170,7 +170,7 @@ class FileEmbeddingStore(
         try {
             val map = cache ?: run {
                 // Load all embeddings into the map if cache is empty
-                val all = getAll()
+                val all = get()
                 LinkedHashMap(all.associateBy { it.id })
             }
 
@@ -188,6 +188,26 @@ class FileEmbeddingStore(
         }
     }
 
+    suspend fun get(ids: List<Long>): List<Embedding> = withContext(Dispatchers.IO) {
+        val map = cache ?: run {
+            val all = get()
+            LinkedHashMap(all.associateBy { it.id })
+        }
+        val embeddings = mutableListOf<Embedding>()
+
+        for (id in ids) {
+            map.get(id)?.let { embeddings.add(it) }
+        }
+        embeddings
+    }
+
+    suspend fun get(id: Long): Embedding? = withContext(Dispatchers.IO) {
+        val map = cache ?: run {
+            val all = get()
+            LinkedHashMap(all.associateBy { it.id })
+        }
+        map.get(id)
+    }
 
     override fun clear(){
         cache = null

extensions/src/test/kotlin/com/fpf/smartscansdk/extensions/embeddings/FileEmbeddingRetriever.kt

@@ -0,0 +1,55 @@
+package com.fpf.smartscansdk.extensions.embeddings
+
+import com.fpf.smartscansdk.core.ml.embeddings.Embedding
+import kotlinx.coroutines.test.runTest
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.TestInstance
+import org.junit.jupiter.api.io.TempDir
+import org.junit.jupiter.api.Test
+import java.io.File
+
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+class FileEmbeddingRetrieverTest {
+
+    @TempDir
+    lateinit var tempDir: File
+
+    private val embeddingLength = 4
+
+    private fun embedding(id: Long, date: Long, values: FloatArray) =
+        Embedding(id, date, values)
+
+    private fun createStore(fileName: String = "embeddings.bin") =
+        FileEmbeddingStore(tempDir, fileName, embeddingLength)
+
+    @Test
+    fun `query batch retrieval with start and end works`() = runTest {
+        val store = createStore()
+        val retriever = FileEmbeddingRetriever(store)
+
+        val embeddings = listOf(
+            embedding(1, 100, floatArrayOf(0.1f, 0.2f, 0.3f, 0.4f)),
+            embedding(2, 200, floatArrayOf(0.5f, 0.6f, 0.7f, 0.8f)),
+            embedding(3, 300, floatArrayOf(0.9f, 1.0f, 1.1f, 1.2f))
+        )
+        store.add(embeddings)
+
+        // trigger initial query to populate cachedIds
+        retriever.query(floatArrayOf(0.1f, 0.2f, 0.3f, 0.4f), topK = 3, threshold = 0f)
+
+        // fetch first two cached embeddings (order-agnostic)
+        val batch1 = retriever.query(0, 2)
+        assertEquals(2, batch1.size)
+        assertTrue(batch1.map { it.id }.all { it in listOf(1L, 2L, 3L) })
+
+        // fetch last cached embedding
+        val batch2 = retriever.query(2, 3)
+        assertEquals(1, batch2.size)
+        assertTrue(batch2[0].id in listOf(1L, 2L, 3L))
+
+        // out-of-bounds requests return empty
+        val batch3 = retriever.query(3, 5)
+        assertEquals(0, batch3.size)
+    }
+}

extensions/src/test/kotlin/com/fpf/smartscansdk/extensions/embeddings/FileEmbeddingStoreTest.kt

@@ -51,7 +51,7 @@ class FileEmbeddingStoreTest {
         )
 
         store.add(embeddings)
-        val loaded = store.getAll()
+        val loaded = store.get()
 
         Assertions.assertEquals(2, loaded.size)
         Assertions.assertEquals(embeddings[0].id, loaded[0].id)
@@ -72,7 +72,7 @@ class FileEmbeddingStoreTest {
         store.add(first)
         store.add(second)
 
-        val all = store.getAll()
+        val all = store.get()
         Assertions.assertEquals(2, all.size)
         Assertions.assertEquals(1, all[0].id)
         Assertions.assertEquals(2, all[1].id)
@@ -90,7 +90,7 @@ class FileEmbeddingStoreTest {
         store.add(embeddings)
         store.remove(listOf(2L))
 
-        val remaining = store.getAll()
+        val remaining = store.get()
         Assertions.assertEquals(2, remaining.size)
         Assertions.assertFalse(remaining.any { it.id == 2L })
     }
@@ -101,7 +101,7 @@ class FileEmbeddingStoreTest {
         val embeddings = listOf(embedding(1, 100, FloatArray(embeddingLength) { 0.1f }))
         store.add(embeddings)
 
-        val firstLoad = store.getAll()
+        val firstLoad = store.get()
         if(store.useCache){
             Assertions.assertTrue(store.isCached)
         }else{
@@ -111,7 +111,7 @@ class FileEmbeddingStoreTest {
         store.clear()
         Assertions.assertFalse(store.isCached)
 
-        val secondLoad = store.getAll()
+        val secondLoad = store.get()
         assertTrue(firstLoad.zip(secondLoad).all { (a, b) ->
             a.id == b.id && a.date == b.date && a.embeddings.contentEquals(b.embeddings)
         })
@@ -123,13 +123,13 @@ class FileEmbeddingStoreTest {
         val embeddings = listOf(embedding(1, 100, FloatArray(embeddingLength) { 0.1f }))
         store.add(embeddings)
 
-        val firstLoad = store.getAll()
+        val firstLoad = store.get()
         Assertions.assertFalse(store.isCached)
 
         store.clear()
         Assertions.assertFalse(store.isCached)
 
-        val secondLoad = store.getAll()
+        val secondLoad = store.get()
         assertTrue(firstLoad.zip(secondLoad).all { (a, b) ->
             a.id == b.id && a.date == b.date && a.embeddings.contentEquals(b.embeddings)
         })