-
Notifications
You must be signed in to change notification settings - Fork 7
feat: add vector search support #87
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| @@ -0,0 +1,153 @@ | ||||||||||||||||||||||||||||||
| /** | ||||||||||||||||||||||||||||||
| * CollectionManagerVector.kt | ||||||||||||||||||||||||||||||
| * | ||||||||||||||||||||||||||||||
| * Vector index support for Android/Kotlin implementation. | ||||||||||||||||||||||||||||||
| * This code should be integrated into CollectionManager.kt | ||||||||||||||||||||||||||||||
| */ | ||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||
| package com.cblreactnative | ||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||
| import com.couchbase.lite.Collection | ||||||||||||||||||||||||||||||
| import com.couchbase.lite.VectorIndexConfiguration | ||||||||||||||||||||||||||||||
| import com.couchbase.lite.VectorEncoding | ||||||||||||||||||||||||||||||
| import com.couchbase.lite.DistanceMetric | ||||||||||||||||||||||||||||||
| import com.couchbase.lite.ScalarQuantizerType | ||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||
| /** | ||||||||||||||||||||||||||||||
| * Creates a vector index on the specified collection. | ||||||||||||||||||||||||||||||
| * | ||||||||||||||||||||||||||||||
| * This method handles the "vector" case in the createIndex switch/when statement. | ||||||||||||||||||||||||||||||
| * | ||||||||||||||||||||||||||||||
| * @param indexName Name for the new index | ||||||||||||||||||||||||||||||
| * @param indexConfig Map containing vector index configuration | ||||||||||||||||||||||||||||||
| * @param collection The collection to create the index on | ||||||||||||||||||||||||||||||
| * @throws Exception if index creation fails | ||||||||||||||||||||||||||||||
| */ | ||||||||||||||||||||||||||||||
| fun createVectorIndex( | ||||||||||||||||||||||||||||||
| indexName: String, | ||||||||||||||||||||||||||||||
| indexConfig: Map<String, Any>, | ||||||||||||||||||||||||||||||
| collection: Collection | ||||||||||||||||||||||||||||||
| ) { | ||||||||||||||||||||||||||||||
| // Extract required parameters | ||||||||||||||||||||||||||||||
| val expression = indexConfig["expression"] as? String | ||||||||||||||||||||||||||||||
| ?: throw Exception("Vector index requires 'expression' parameter") | ||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||
| val dimensions = (indexConfig["dimensions"] as? Number)?.toLong() | ||||||||||||||||||||||||||||||
| ?: throw Exception("Vector index requires 'dimensions' parameter") | ||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||
| val centroids = (indexConfig["centroids"] as? Number)?.toLong() | ||||||||||||||||||||||||||||||
| ?: throw Exception("Vector index requires 'centroids' parameter") | ||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||
| // Create vector index configuration | ||||||||||||||||||||||||||||||
| val vectorConfig = VectorIndexConfiguration(expression, dimensions, centroids) | ||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||
| // Set distance metric | ||||||||||||||||||||||||||||||
| val metricStr = indexConfig["metric"] as? String | ||||||||||||||||||||||||||||||
| if (metricStr != null) { | ||||||||||||||||||||||||||||||
| vectorConfig.metric = when (metricStr) { | ||||||||||||||||||||||||||||||
| "cosine" -> DistanceMetric.COSINE | ||||||||||||||||||||||||||||||
| "euclidean" -> DistanceMetric.EUCLIDEAN | ||||||||||||||||||||||||||||||
| "euclideanSquared" -> DistanceMetric.EUCLIDEAN_SQUARED | ||||||||||||||||||||||||||||||
| "dot" -> DistanceMetric.DOT | ||||||||||||||||||||||||||||||
| else -> DistanceMetric.EUCLIDEAN_SQUARED // Default | ||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||
|
Comment on lines
+47
to
+53
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The
Suggested change
|
||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||
| // Set encoding | ||||||||||||||||||||||||||||||
| val encodingConfig = indexConfig["encoding"] as? Map<String, Any> | ||||||||||||||||||||||||||||||
| if (encodingConfig != null) { | ||||||||||||||||||||||||||||||
| val encodingType = encodingConfig["type"] as? String | ||||||||||||||||||||||||||||||
| vectorConfig.encoding = when (encodingType) { | ||||||||||||||||||||||||||||||
| "none" -> VectorEncoding.none() | ||||||||||||||||||||||||||||||
| "SQ" -> { | ||||||||||||||||||||||||||||||
| // Scalar Quantizer - determine type from config or default to SQ8 | ||||||||||||||||||||||||||||||
| val sqType = encodingConfig["sqType"] as? String | ||||||||||||||||||||||||||||||
| when (sqType) { | ||||||||||||||||||||||||||||||
| "SQ4" -> VectorEncoding.scalarQuantizer(ScalarQuantizerType.SQ4) | ||||||||||||||||||||||||||||||
| "SQ6" -> VectorEncoding.scalarQuantizer(ScalarQuantizerType.SQ6) | ||||||||||||||||||||||||||||||
| "SQ8" -> VectorEncoding.scalarQuantizer(ScalarQuantizerType.SQ8) | ||||||||||||||||||||||||||||||
| else -> VectorEncoding.scalarQuantizer(ScalarQuantizerType.SQ8) | ||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||
| "PQ" -> { | ||||||||||||||||||||||||||||||
| // Product Quantizer | ||||||||||||||||||||||||||||||
| val subquantizers = (encodingConfig["subquantizers"] as? Number)?.toLong() ?: 0 | ||||||||||||||||||||||||||||||
| val bits = (encodingConfig["bits"] as? Number)?.toLong() ?: 8 | ||||||||||||||||||||||||||||||
| VectorEncoding.productQuantizer(subquantizers, bits) | ||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||
| else -> VectorEncoding.none() | ||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||
|
Comment on lines
+60
to
+79
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This block has a few areas for improvement:
Here is a suggested refactoring that addresses these points. vectorConfig.encoding = when (encodingType) {
"none", null -> VectorEncoding.none()
"SQ" -> {
// Scalar Quantizer - determine type from config or default to SQ8
val sqType = encodingConfig["sqType"] as? String
when (sqType) {
"SQ4" -> VectorEncoding.scalarQuantizer(ScalarQuantizerType.SQ4)
"SQ6" -> VectorEncoding.scalarQuantizer(ScalarQuantizerType.SQ6)
"SQ8", null -> VectorEncoding.scalarQuantizer(ScalarQuantizerType.SQ8)
else -> throw IllegalArgumentException("Unknown Scalar Quantizer type: $sqType")
}
}
"PQ" -> {
// Product Quantizer
val subquantizers = (encodingConfig["subquantizers"] as? Number)?.toLong()
?: throw IllegalArgumentException("Product Quantizer encoding requires 'subquantizers' parameter")
if (subquantizers <= 0) {
throw IllegalArgumentException("'subquantizers' must be a positive number")
}
val bits = (encodingConfig["bits"] as? Number)?.toLong() ?: 8
VectorEncoding.productQuantizer(subquantizers, bits)
}
else -> throw IllegalArgumentException("Unknown vector encoding type: $encodingType")
} |
||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||
| // Set optional training parameters | ||||||||||||||||||||||||||||||
| val minTrainingSize = (indexConfig["minTrainingSize"] as? Number)?.toLong() | ||||||||||||||||||||||||||||||
| if (minTrainingSize != null && minTrainingSize > 0) { | ||||||||||||||||||||||||||||||
| vectorConfig.minTrainingSize = minTrainingSize | ||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||
| val maxTrainingSize = (indexConfig["maxTrainingSize"] as? Number)?.toLong() | ||||||||||||||||||||||||||||||
| if (maxTrainingSize != null && maxTrainingSize > 0) { | ||||||||||||||||||||||||||||||
| vectorConfig.maxTrainingSize = maxTrainingSize | ||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||
| // Set number of probes (affects search accuracy vs speed) | ||||||||||||||||||||||||||||||
| val numProbes = (indexConfig["numProbes"] as? Number)?.toLong() | ||||||||||||||||||||||||||||||
| if (numProbes != null && numProbes > 0) { | ||||||||||||||||||||||||||||||
| vectorConfig.numProbes = numProbes | ||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||
| // Set lazy indexing (index built on first query) | ||||||||||||||||||||||||||||||
| val isLazy = indexConfig["isLazy"] as? Boolean | ||||||||||||||||||||||||||||||
| if (isLazy != null) { | ||||||||||||||||||||||||||||||
| vectorConfig.isLazy = isLazy | ||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||
| // Create the index | ||||||||||||||||||||||||||||||
| collection.createIndex(indexName, vectorConfig) | ||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||
| /* | ||||||||||||||||||||||||||||||
| * INTEGRATION INSTRUCTIONS | ||||||||||||||||||||||||||||||
| * ======================== | ||||||||||||||||||||||||||||||
| * | ||||||||||||||||||||||||||||||
| * To integrate this into CollectionManager.kt: | ||||||||||||||||||||||||||||||
| * | ||||||||||||||||||||||||||||||
| * 1. Add the imports at the top of the file: | ||||||||||||||||||||||||||||||
| * | ||||||||||||||||||||||||||||||
| * import com.couchbase.lite.VectorIndexConfiguration | ||||||||||||||||||||||||||||||
| * import com.couchbase.lite.VectorEncoding | ||||||||||||||||||||||||||||||
| * import com.couchbase.lite.DistanceMetric | ||||||||||||||||||||||||||||||
| * import com.couchbase.lite.ScalarQuantizerType | ||||||||||||||||||||||||||||||
| * | ||||||||||||||||||||||||||||||
| * 2. Modify the createIndex function signature to accept optional config: | ||||||||||||||||||||||||||||||
| * | ||||||||||||||||||||||||||||||
| * fun createIndex( | ||||||||||||||||||||||||||||||
| * indexName: String, | ||||||||||||||||||||||||||||||
| * indexType: String, | ||||||||||||||||||||||||||||||
| * items: List<List<Any>>, | ||||||||||||||||||||||||||||||
| * indexConfig: Map<String, Any>? = null, // ADD THIS | ||||||||||||||||||||||||||||||
| * collectionName: String, | ||||||||||||||||||||||||||||||
| * scopeName: String, | ||||||||||||||||||||||||||||||
| * databaseName: String | ||||||||||||||||||||||||||||||
| * ) | ||||||||||||||||||||||||||||||
| * | ||||||||||||||||||||||||||||||
| * 3. Add the vector case to the when statement: | ||||||||||||||||||||||||||||||
| * | ||||||||||||||||||||||||||||||
| * when (indexType) { | ||||||||||||||||||||||||||||||
| * "value" -> { | ||||||||||||||||||||||||||||||
| * // existing code... | ||||||||||||||||||||||||||||||
| * } | ||||||||||||||||||||||||||||||
| * "full-text" -> { | ||||||||||||||||||||||||||||||
| * // existing code... | ||||||||||||||||||||||||||||||
| * } | ||||||||||||||||||||||||||||||
| * "vector" -> { | ||||||||||||||||||||||||||||||
| * val config = indexConfig | ||||||||||||||||||||||||||||||
| * ?: throw Exception("Vector index requires config") | ||||||||||||||||||||||||||||||
| * createVectorIndex(indexName, config, collection) | ||||||||||||||||||||||||||||||
| * } | ||||||||||||||||||||||||||||||
| * else -> throw Exception("Unknown index type: $indexType") | ||||||||||||||||||||||||||||||
| * } | ||||||||||||||||||||||||||||||
| * | ||||||||||||||||||||||||||||||
| * 4. Update the React Native bridge module (CblReactnativeModule.kt) to | ||||||||||||||||||||||||||||||
| * extract and pass the full config map for vector indexes. | ||||||||||||||||||||||||||||||
| */ | ||||||||||||||||||||||||||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
For better error handling and clarity, it's recommended to throw a more specific exception like
IllegalArgumentExceptioninstead of the genericExceptionwhen a required parameter is missing. This helps callers to better understand and handle the error. Please apply this to all three required parameter checks (expression,dimensions,centroids) and update the function's KDoc@throwstag accordingly.