Skip to content

Commit 0da98f9

Browse files
committed
feat: Neo4j vector index method support
This commit introduces support for specifying vector index methods (HNSW) for Neo4j targets. - Modified `src/base/spec.rs` to derive `Eq` for `VectorIndexMethod`. - Modified `src/ops/targets/neo4j.rs` to: - Allow `VectorIndexMethod` to be passed to `IndexDef::from_vector_index_def`. - Store `VectorIndexMethod` in `IndexDef::VectorIndex`. - Implement error handling for unsupported `VectorIndexMethod` (IVFFlat). - Update `SetupComponentOperator::describe_state` to display the method. - Update `SetupComponentOperator::create` to include HNSW parameters in the Cypher query.
1 parent 0dc1a48 commit 0da98f9

File tree

2 files changed

+34
-6
lines changed

2 files changed

+34
-6
lines changed

src/base/spec.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -384,7 +384,7 @@ impl fmt::Display for VectorSimilarityMetric {
384384
}
385385
}
386386

387-
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
387+
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
388388
#[serde(tag = "kind")]
389389
pub enum VectorIndexMethod {
390390
Hnsw {

src/ops/targets/neo4j.rs

Lines changed: 33 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -568,9 +568,6 @@ impl SetupState {
568568
.map(|f| (f.name.as_str(), &f.value_type.typ))
569569
.collect::<HashMap<_, _>>();
570570
for index_def in index_options.vector_indexes.iter() {
571-
if index_def.method.is_some() {
572-
api_bail!("Vector index method is not configurable for Neo4j yet");
573-
}
574571
sub_components.push(ComponentState {
575572
object_label: schema.elem_type.clone(),
576573
index_def: IndexDef::from_vector_index_def(
@@ -583,6 +580,7 @@ impl SetupState {
583580
index_def.field_name
584581
)
585582
})?,
583+
index_def.method.clone(),
586584
)?,
587585
});
588586
}
@@ -644,14 +642,20 @@ enum IndexDef {
644642
field_name: String,
645643
metric: spec::VectorSimilarityMetric,
646644
vector_size: usize,
645+
method: Option<spec::VectorIndexMethod>,
647646
},
648647
}
649648

650649
impl IndexDef {
651650
fn from_vector_index_def(
652651
index_def: &spec::VectorIndexDef,
653652
field_typ: &schema::ValueType,
653+
method: Option<spec::VectorIndexMethod>,
654654
) -> Result<Self> {
655+
let method = index_def.method.clone();
656+
if let Some(spec::VectorIndexMethod::IvfFlat { .. }) = method {
657+
api_bail!("IVFFlat vector index method is not supported for Neo4j");
658+
}
655659
Ok(Self::VectorIndex {
656660
field_name: index_def.field_name.clone(),
657661
vector_size: (match field_typ {
@@ -664,6 +668,7 @@ impl IndexDef {
664668
api_error!("Vector index field must be a vector with fixed dimension")
665669
})?,
666670
metric: index_def.metric,
671+
method,
667672
})
668673
}
669674
}
@@ -723,9 +728,14 @@ impl components::SetupOperator for SetupComponentOperator {
723728
field_name,
724729
metric,
725730
vector_size,
731+
method,
726732
} => {
733+
let method_str = method
734+
.as_ref()
735+
.map(|m| format!(", method: {}", m))
736+
.unwrap_or_default();
727737
format!(
728-
"{key_desc} ON {label} (field_name: {field_name}, vector_size: {vector_size}, metric: {metric})",
738+
"{key_desc} ON {label} (field_name: {field_name}, vector_size: {vector_size}, metric: {metric}{method_str})",
729739
)
730740
}
731741
}
@@ -752,14 +762,32 @@ impl components::SetupOperator for SetupComponentOperator {
752762
field_name,
753763
metric,
754764
vector_size,
765+
method,
755766
} => {
767+
let method_config =
768+
if let Some(spec::VectorIndexMethod::Hnsw { m, ef_construction }) = method {
769+
let mut parts = vec![];
770+
if let Some(m_val) = m {
771+
parts.push(format!("`hnsw.m`: {}", m_val));
772+
}
773+
if let Some(ef_val) = ef_construction {
774+
parts.push(format!("`hnsw.ef_construction`: {}", ef_val));
775+
}
776+
if parts.is_empty() {
777+
"".to_string()
778+
} else {
779+
format!(", {}", parts.join(", "))
780+
}
781+
} else {
782+
"".to_string()
783+
};
756784
formatdoc! {"
757785
CREATE VECTOR INDEX {name} IF NOT EXISTS
758786
FOR {matcher} ON {qualifier}.{field_name}
759787
OPTIONS {{
760788
indexConfig: {{
761789
`vector.dimensions`: {vector_size},
762-
`vector.similarity_function`: '{metric}'
790+
`vector.similarity_function`: '{metric}'{method_config}
763791
}}
764792
}}",
765793
name = key.name,

0 commit comments

Comments
 (0)