feat: allow configuring Postgres vector index method and options

Author: banrovegrie

docs/docs/core/flow_def.mdx

@@ -313,6 +313,7 @@ Types of the fields must be key types. See [Key Types](data_types#key-types) for
 
     *   `field_name`: the field to create vector index.
     *   `metric`: the similarity metric to use.
+    *   `method`: the index algorithm and optional tuning parameters. Defaults to HNSW. Use `cocoindex.HnswVectorIndexMethod()` or `cocoindex.IvfFlatVectorIndexMethod()` to customize the method and its parameters.
 
 #### Similarity Metrics
 

docs/docs/examples/examples/simple_vector_index.md

@@ -105,6 +105,16 @@ doc_embeddings.export(
 CocoIndex supports other vector databases as well, with 1-line switch.
 <DocumentationButton url="https://cocoindex.io/docs/ops/targets" text="Targets" />
 
+Need IVFFlat or custom HNSW parameters? Pass a method, for example:
+
+```python
+cocoindex.VectorIndexDef(
+    field_name="embedding",
+    metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY,
+    method=cocoindex.IvfFlatVectorIndexMethod(lists=200),
+)
+```
+
 ## Query the index
 
 ### Define a shared flow for both indexing and querying

python/cocoindex/__init__.py

@@ -21,7 +21,13 @@
 from .flow import update_all_flows_async, setup_all_flows, drop_all_flows
 from .lib import settings, init, start_server, stop
 from .llm import LlmSpec, LlmApiType
-from .index import VectorSimilarityMetric, VectorIndexDef, IndexOptions
+from .index import (
+    VectorSimilarityMetric,
+    VectorIndexDef,
+    IndexOptions,
+    HnswVectorIndexMethod,
+    IvfFlatVectorIndexMethod,
+)
 from .setting import DatabaseConnectionSpec, Settings, ServerSettings
 from .setting import get_app_namespace
 from .query_handler import QueryHandlerResultFields, QueryInfo, QueryOutput
@@ -82,6 +88,8 @@
     "VectorSimilarityMetric",
     "VectorIndexDef",
     "IndexOptions",
+    "HnswVectorIndexMethod",
+    "IvfFlatVectorIndexMethod",
     # Settings
     "DatabaseConnectionSpec",
     "Settings",

python/cocoindex/index.py

@@ -1,6 +1,6 @@
 from enum import Enum
-from dataclasses import dataclass
-from typing import Sequence
+from dataclasses import dataclass, field
+from typing import Sequence, Union
 
 
 class VectorSimilarityMetric(Enum):
@@ -9,6 +9,26 @@ class VectorSimilarityMetric(Enum):
     INNER_PRODUCT = "InnerProduct"
 
 
+@dataclass
+class HnswVectorIndexMethod:
+    """HNSW vector index parameters."""
+
+    type: str = field(init=False, default="hnsw")
+    m: int | None = None
+    ef_construction: int | None = None
+
+
+@dataclass
+class IvfFlatVectorIndexMethod:
+    """IVFFlat vector index parameters."""
+
+    type: str = field(init=False, default="ivfflat")
+    lists: int | None = None
+
+
+VectorIndexMethod = Union[HnswVectorIndexMethod, IvfFlatVectorIndexMethod]
+
+
 @dataclass
 class VectorIndexDef:
     """
@@ -17,6 +37,7 @@ class VectorIndexDef:
 
     field_name: str
     metric: VectorSimilarityMetric
+    method: VectorIndexMethod = field(default_factory=HnswVectorIndexMethod)
 
 
 @dataclass

src/base/spec.rs

@@ -384,15 +384,86 @@ impl fmt::Display for VectorSimilarityMetric {
     }
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum VectorIndexMethod {
+    Hnsw {
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        m: Option<u32>,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        ef_construction: Option<u32>,
+    },
+    IvfFlat {
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        lists: Option<u32>,
+    },
+}
+
+impl Default for VectorIndexMethod {
+    fn default() -> Self {
+        Self::Hnsw {
+            m: None,
+            ef_construction: None,
+        }
+    }
+}
+
+impl VectorIndexMethod {
+    pub fn kind(&self) -> &'static str {
+        match self {
+            Self::Hnsw { .. } => "hnsw",
+            Self::IvfFlat { .. } => "ivfflat",
+        }
+    }
+
+    pub fn is_default(&self) -> bool {
+        matches!(self, Self::Hnsw { m: None, ef_construction: None })
+    }
+}
+
+impl fmt::Display for VectorIndexMethod {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Hnsw { m, ef_construction } => {
+                let mut parts = Vec::new();
+                if let Some(m) = m {
+                    parts.push(format!("m={}", m));
+                }
+                if let Some(ef) = ef_construction {
+                    parts.push(format!("ef_construction={}", ef));
+                }
+                if parts.is_empty() {
+                    write!(f, "hnsw")
+                } else {
+                    write!(f, "hnsw({})", parts.join(","))
+                }
+            }
+            Self::IvfFlat { lists } => {
+                if let Some(lists) = lists {
+                    write!(f, "ivfflat(lists={lists})")
+                } else {
+                    write!(f, "ivfflat")
+                }
+            }
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct VectorIndexDef {
     pub field_name: FieldName,
     pub metric: VectorSimilarityMetric,
+    #[serde(default)]
+    pub method: VectorIndexMethod,
 }
 
 impl fmt::Display for VectorIndexDef {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}:{}", self.field_name, self.metric)
+        if self.method.is_default() {
+            write!(f, "{}:{}", self.field_name, self.metric)
+        } else {
+            write!(f, "{}:{}:{}", self.field_name, self.metric, self.method)
+        }
     }
 }
 

src/ops/targets/postgres.rs

@@ -461,21 +461,51 @@ fn to_vector_similarity_metric_sql(metric: VectorSimilarityMetric) -> &'static s
 }
 
 fn to_index_spec_sql(index_spec: &VectorIndexDef) -> Cow<'static, str> {
+    let method = match &index_spec.method {
+        spec::VectorIndexMethod::Hnsw { .. } => "hnsw",
+        spec::VectorIndexMethod::IvfFlat { .. } => "ivfflat",
+    };
+    let options = match &index_spec.method {
+        spec::VectorIndexMethod::Hnsw { m, ef_construction } => {
+            let mut opts = Vec::new();
+            if let Some(m) = m {
+                opts.push(format!("m = {}", m));
+            }
+            if let Some(ef) = ef_construction {
+                opts.push(format!("ef_construction = {}", ef));
+            }
+            opts
+        }
+        spec::VectorIndexMethod::IvfFlat { lists } => lists
+            .map(|lists| vec![format!("lists = {}", lists)])
+            .unwrap_or_default(),
+    };
+    let with_clause = if options.is_empty() {
+        String::new()
+    } else {
+        format!(" WITH ({})", options.join(", "))
+    };
     format!(
-        "USING hnsw ({} {})",
+        "USING {method} ({} {}){}",
         index_spec.field_name,
-        to_vector_similarity_metric_sql(index_spec.metric)
+        to_vector_similarity_metric_sql(index_spec.metric),
+        with_clause
     )
     .into()
 }
 
 fn to_vector_index_name(table_name: &str, vector_index_def: &spec::VectorIndexDef) -> String {
-    format!(
+    let mut name = format!(
         "{}__{}__{}",
         table_name,
         vector_index_def.field_name,
         to_vector_similarity_metric_sql(vector_index_def.metric)
-    )
+    );
+    if !vector_index_def.method.is_default() {
+        name.push_str("__");
+        name.push_str(vector_index_def.method.kind());
+    }
+    name
 }
 
 fn describe_index_spec(index_name: &str, index_spec: &VectorIndexDef) -> String {