feat(kuzu): support kuzu as storage target (#561)

* feat(kuzu): support setup

* style: rearrange `kuzu.rs` into sections

* refactor(kuzu): expose single `register()` method to seal internal types

* refactor: use standalone methods to build setup query

* feat(kuzu): support kuzu as target

Author: badmonster0

docs/docs/core/settings.mdx

@@ -6,7 +6,7 @@ description: Provide settings for CocoIndex, e.g. database connection, app names
 import Tabs from '@theme/Tabs';
 import TabItem from '@theme/TabItem';
 
-# CocoIndex Settings
+# CocoIndex Setting
 
 Certain settings need to be provided for CocoIndex to work, e.g. database connections, app namespace, etc.
 

examples/docs_to_knowledge_graph/main.py

@@ -26,23 +26,41 @@ class Relationship:
     object: str
 
 
+neo4j_conn_spec = cocoindex.add_auth_entry(
+    "Neo4jConnection",
+    cocoindex.storages.Neo4jConnection(
+        uri="bolt://localhost:7687",
+        user="neo4j",
+        password="cocoindex",
+    ),
+)
+kuzu_conn_spec = cocoindex.add_auth_entry(
+    "KuzuConnection",
+    cocoindex.storages.KuzuConnection(
+        api_server_url="http://localhost:8123",
+    ),
+)
+
+# Use Neo4j as the graph database
+GraphDbSpec = cocoindex.storages.Neo4j
+GraphDbConnection = cocoindex.storages.Neo4jConnection
+GraphDbDeclaration = cocoindex.storages.Neo4jDeclaration
+conn_spec = neo4j_conn_spec
+
+# Use Kuzu as the graph database
+#  GraphDbSpec = cocoindex.storages.Kuzu
+#  GraphDbConnection = cocoindex.storages.KuzuConnection
+#  GraphDbDeclaration = cocoindex.storages.KuzuDeclaration
+#  conn_spec = kuzu_conn_spec
+
+
 @cocoindex.flow_def(name="DocsToKG")
 def docs_to_kg_flow(
     flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope
-):
+) -> None:
     """
     Define an example flow that extracts relationship from files and build knowledge graph.
     """
-    # configure neo4j connection
-    conn_spec = cocoindex.add_auth_entry(
-        "Neo4jConnection",
-        cocoindex.storages.Neo4jConnection(
-            uri="bolt://localhost:7687",
-            user="neo4j",
-            password="cocoindex",
-        ),
-    )
-
     data_scope["documents"] = flow_builder.add_source(
         cocoindex.sources.LocalFile(
             path="../../docs/docs/core", included_patterns=["*.md", "*.mdx"]
@@ -112,22 +130,22 @@ def docs_to_kg_flow(
     # export to neo4j
     document_node.export(
         "document_node",
-        cocoindex.storages.Neo4j(
+        GraphDbSpec(
             connection=conn_spec, mapping=cocoindex.storages.Nodes(label="Document")
         ),
         primary_key_fields=["filename"],
     )
     # Declare reference Node to reference entity node in a relationship
     flow_builder.declare(
-        cocoindex.storages.Neo4jDeclaration(
+        GraphDbDeclaration(
             connection=conn_spec,
             nodes_label="Entity",
             primary_key_fields=["value"],
         )
     )
     entity_relationship.export(
         "entity_relationship",
-        cocoindex.storages.Neo4j(
+        GraphDbSpec(
             connection=conn_spec,
             mapping=cocoindex.storages.Relationships(
                 rel_type="RELATIONSHIP",
@@ -153,7 +171,7 @@ def docs_to_kg_flow(
     )
     entity_mention.export(
         "entity_mention",
-        cocoindex.storages.Neo4j(
+        GraphDbSpec(
             connection=conn_spec,
             mapping=cocoindex.storages.Relationships(
                 rel_type="MENTION",

examples/product_recommendation/main.py

@@ -77,10 +77,38 @@ def extract_product_info(product: cocoindex.Json, filename: str) -> ProductInfo:
     )
 
 
+neo4j_conn_spec = cocoindex.add_auth_entry(
+    "Neo4jConnection",
+    cocoindex.storages.Neo4jConnection(
+        uri="bolt://localhost:7687",
+        user="neo4j",
+        password="cocoindex",
+    ),
+)
+kuzu_conn_spec = cocoindex.add_auth_entry(
+    "KuzuConnection",
+    cocoindex.storages.KuzuConnection(
+        api_server_url="http://localhost:8123",
+    ),
+)
+
+# Use Neo4j as the graph database
+GraphDbSpec = cocoindex.storages.Neo4j
+GraphDbConnection = cocoindex.storages.Neo4jConnection
+GraphDbDeclaration = cocoindex.storages.Neo4jDeclaration
+conn_spec = neo4j_conn_spec
+
+# Use Kuzu as the graph database
+#  GraphDbSpec = cocoindex.storages.Kuzu
+#  GraphDbConnection = cocoindex.storages.KuzuConnection
+#  GraphDbDeclaration = cocoindex.storages.KuzuDeclaration
+#  conn_spec = kuzu_conn_spec
+
+
 @cocoindex.flow_def(name="StoreProduct")
 def store_product_flow(
     flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope
-):
+) -> None:
     """
     Define an example flow that extracts triples from files and build knowledge graph.
     """
@@ -122,25 +150,16 @@ def store_product_flow(
                 taxonomy=t["name"],
             )
 
-    conn_spec = cocoindex.add_auth_entry(
-        "Neo4jConnection",
-        cocoindex.storages.Neo4jConnection(
-            uri="bolt://localhost:7687",
-            user="neo4j",
-            password="cocoindex",
-        ),
-    )
-
     product_node.export(
         "product_node",
-        cocoindex.storages.Neo4j(
+        GraphDbSpec(
             connection=conn_spec, mapping=cocoindex.storages.Nodes(label="Product")
         ),
         primary_key_fields=["id"],
     )
 
     flow_builder.declare(
-        cocoindex.storages.Neo4jDeclaration(
+        GraphDbDeclaration(
             connection=conn_spec,
             nodes_label="Taxonomy",
             primary_key_fields=["value"],
@@ -149,7 +168,7 @@ def store_product_flow(
 
     product_taxonomy.export(
         "product_taxonomy",
-        cocoindex.storages.Neo4j(
+        GraphDbSpec(
             connection=conn_spec,
             mapping=cocoindex.storages.Relationships(
                 rel_type="PRODUCT_TAXONOMY",
@@ -175,7 +194,7 @@ def store_product_flow(
     )
     product_complementary_taxonomy.export(
         "product_complementary_taxonomy",
-        cocoindex.storages.Neo4j(
+        GraphDbSpec(
             connection=conn_spec,
             mapping=cocoindex.storages.Relationships(
                 rel_type="PRODUCT_COMPLEMENTARY_TAXONOMY",

python/cocoindex/storages.py

@@ -25,16 +25,6 @@ class Qdrant(op.StorageSpec):
     api_key: str | None = None
 
 
-@dataclass
-class Neo4jConnection:
-    """Connection spec for Neo4j."""
-
-    uri: str
-    user: str
-    password: str
-    db: str | None = None
-
-
 @dataclass
 class TargetFieldMapping:
     """Mapping for a graph element (node or relationship) field."""
@@ -88,6 +78,16 @@ class Relationships:
 NodeReferenceMapping = NodeFromFields
 
 
+@dataclass
+class Neo4jConnection:
+    """Connection spec for Neo4j."""
+
+    uri: str
+    user: str
+    password: str
+    db: str | None = None
+
+
 class Neo4j(op.StorageSpec):
     """Graph storage powered by Neo4j."""
 
@@ -103,3 +103,26 @@ class Neo4jDeclaration(op.DeclarationSpec):
     nodes_label: str
     primary_key_fields: Sequence[str]
     vector_indexes: Sequence[index.VectorIndexDef] = ()
+
+
+@dataclass
+class KuzuConnection:
+    """Connection spec for Kuzu."""
+
+    api_server_url: str
+
+
+class Kuzu(op.StorageSpec):
+    """Graph storage powered by Kuzu."""
+
+    connection: AuthEntryReference[KuzuConnection]
+    mapping: Nodes | Relationships
+
+
+class KuzuDeclaration(op.DeclarationSpec):
+    """Declarations for Kuzu."""
+
+    kind = "Kuzu"
+    connection: AuthEntryReference[KuzuConnection]
+    nodes_label: str
+    primary_key_fields: Sequence[str]

src/base/schema.rs

@@ -86,7 +86,7 @@ impl std::fmt::Display for BasicValueType {
     }
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
 pub struct StructSchema {
     pub fields: Arc<Vec<FieldSchema>>,
 
@@ -138,7 +138,7 @@ impl std::fmt::Display for TableKind {
     }
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 pub struct TableSchema {
     pub kind: TableKind,
     pub row: StructSchema,
@@ -191,7 +191,7 @@ impl TableSchema {
     }
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 #[serde(tag = "kind")]
 pub enum ValueType {
     Struct(StructSchema),
@@ -222,7 +222,7 @@ impl ValueType {
     }
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 pub struct EnrichedValueType<DataType = ValueType> {
     #[serde(rename = "type")]
     pub typ: DataType,
@@ -295,7 +295,7 @@ impl std::fmt::Display for ValueType {
     }
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 pub struct FieldSchema<DataType = ValueType> {
     /// ID is used to identify the field in the schema.
     pub name: FieldName,
@@ -342,7 +342,7 @@ impl std::fmt::Display for FieldSchema {
 pub struct CollectorSchema {
     pub fields: Vec<FieldSchema>,
     /// If specified, the collector will have an automatically generated UUID field with the given index.
-    pub auto_uuid_field_idx: Option<u32>,
+    pub auto_uuid_field_idx: Option<usize>,
 }
 
 impl std::fmt::Display for CollectorSchema {

src/base/value.rs

@@ -1,5 +1,6 @@
 use super::schema::*;
 use crate::base::duration::parse_duration;
+use crate::prelude::invariance_violation;
 use crate::{api_bail, api_error};
 use anyhow::Result;
 use base64::prelude::*;
@@ -175,6 +176,26 @@ impl std::fmt::Display for KeyValue {
 }
 
 impl KeyValue {
+    pub fn from_json(value: serde_json::Value, fields_schema: &[FieldSchema]) -> Result<Self> {
+        let value = if fields_schema.len() == 1 {
+            Value::from_json(value, &fields_schema[0].value_type.typ)?
+        } else {
+            let field_values: FieldValues = FieldValues::from_json(value, fields_schema)?;
+            Value::Struct(field_values)
+        };
+        Ok(value.as_key()?)
+    }
+
+    pub fn from_values<'a>(values: impl ExactSizeIterator<Item = &'a Value>) -> Result<Self> {
+        let key = if values.len() == 1 {
+            let mut values = values;
+            values.next().ok_or_else(invariance_violation)?.as_key()?
+        } else {
+            KeyValue::Struct(values.map(|v| v.as_key()).collect::<Result<Vec<_>>>()?)
+        };
+        Ok(key)
+    }
+
     pub fn fields_iter(&self, num_fields: usize) -> Result<impl Iterator<Item = &KeyValue>> {
         let slice = if num_fields == 1 {
             std::slice::from_ref(self)

src/builder/analyzer.rs

@@ -987,7 +987,6 @@ impl AnalyzerContext<'_> {
                                     .fields
                                     .iter()
                                     .position(|field| &field.name == f)
-                                    .map(|idx| idx as u32)
                                     .ok_or_else(|| anyhow!("field not found: {}", f))
                             })
                             .collect::<Result<Vec<_>>>()?;
@@ -1007,7 +1006,7 @@ impl AnalyzerContext<'_> {
                         let mut value_fields_schema: Vec<FieldSchema> = vec![];
                         let mut value_fields_idx = vec![];
                         for (idx, field) in collector_schema.fields.iter().enumerate() {
-                            if !pk_fields_idx.contains(&(idx as u32)) {
+                            if !pk_fields_idx.contains(&idx) {
                                 value_fields_schema.push(field.clone());
                                 value_fields_idx.push(idx as u32);
                             }

src/builder/plan.rs

@@ -94,7 +94,7 @@ pub struct AnalyzedCollectOp {
 }
 
 pub enum AnalyzedPrimaryKeyDef {
-    Fields(Vec<u32>),
+    Fields(Vec<usize>),
 }
 
 pub struct AnalyzedExportOp {

src/execution/row_indexer.rs

@@ -24,20 +24,11 @@ pub fn extract_primary_key(
     primary_key_def: &AnalyzedPrimaryKeyDef,
     record: &FieldValues,
 ) -> Result<KeyValue> {
-    let key = match primary_key_def {
+    match primary_key_def {
         AnalyzedPrimaryKeyDef::Fields(fields) => {
-            if fields.len() == 1 {
-                record.fields[fields[0] as usize].as_key()?
-            } else {
-                let mut key_values = Vec::with_capacity(fields.len());
-                for field in fields.iter() {
-                    key_values.push(record.fields[*field as usize].as_key()?);
-                }
-                KeyValue::Struct(key_values)
-            }
+            KeyValue::from_values(fields.iter().map(|field| &record.fields[*field as usize]))
         }
-    };
-    Ok(key)
+    }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]

src/ops/registration.rs

@@ -6,6 +6,8 @@ use anyhow::Result;
 use std::sync::{Arc, LazyLock, RwLock, RwLockReadGuard};
 
 fn register_executor_factories(registry: &mut ExecutorFactoryRegistry) -> Result<()> {
+    let reqwest_client = reqwest::Client::new();
+
     sources::local_file::Factory.register(registry)?;
     sources::google_drive::Factory.register(registry)?;
     sources::amazon_s3::Factory.register(registry)?;
@@ -16,6 +18,7 @@ fn register_executor_factories(registry: &mut ExecutorFactoryRegistry) -> Result
 
     storages::postgres::Factory::default().register(registry)?;
     Arc::new(storages::qdrant::Factory::default()).register(registry)?;
+    storages::kuzu::register(registry, reqwest_client)?;
 
     storages::neo4j::Factory::new().register(registry)?;