feat(native-composite-key): support multiple key parts in KTable

Author: georgeh0

examples/postgres_source/main.py

@@ -97,8 +97,8 @@ def postgres_product_indexing_flow(
     with data_scope["products"].row() as product:
         product["full_description"] = flow_builder.transform(
             make_full_description,
-            product["_key"]["product_category"],
-            product["_key"]["product_name"],
+            product["product_category"],
+            product["product_name"],
             product["description"],
         )
         product["total_value"] = flow_builder.transform(
@@ -112,8 +112,8 @@ def postgres_product_indexing_flow(
             )
         )
         indexed_product.collect(
-            product_category=product["_key"]["product_category"],
-            product_name=product["_key"]["product_name"],
+            product_category=product["product_category"],
+            product_name=product["product_name"],
             description=product["description"],
             price=product["price"],
             amount=product["amount"],

python/cocoindex/convert.py

@@ -14,7 +14,6 @@
 import numpy as np
 
 from .typing import (
-    KEY_FIELD_NAME,
     TABLE_TYPES,
     AnalyzedAnyType,
     AnalyzedBasicType,
@@ -96,14 +95,24 @@ def encode_struct_list(value: Any) -> Any:
                 f"Value type for dict is required to be a struct (e.g. dataclass or NamedTuple), got {variant.value_type}. "
                 f"If you want a free-formed dict, use `cocoindex.Json` instead."
             )
+        value_encoder = make_engine_value_encoder(value_type_info)
 
-        key_encoder = make_engine_value_encoder(analyze_type_info(variant.key_type))
-        value_encoder = make_engine_value_encoder(analyze_type_info(variant.value_type))
+        key_type_info = analyze_type_info(variant.key_type)
+        key_encoder = make_engine_value_encoder(key_type_info)
+        if isinstance(key_type_info.variant, AnalyzedBasicType):
+
+            def encode_row(k: Any, v: Any) -> Any:
+                return [key_encoder(k)] + value_encoder(v)
+
+        else:
+
+            def encode_row(k: Any, v: Any) -> Any:
+                return key_encoder(k) + value_encoder(v)
 
         def encode_struct_dict(value: Any) -> Any:
             if not value:
                 return []
-            return [[key_encoder(k)] + value_encoder(v) for k, v in value.items()]
+            return [encode_row(k, v) for k, v in value.items()]
 
         return encode_struct_dict
 
@@ -234,25 +243,47 @@ def decode(value: Any) -> Any | None:
                         f"declared `{dst_type_info.core_type}`, a dict type expected"
                     )
 
-                key_field_schema = engine_fields_schema[0]
-                field_path.append(f".{key_field_schema.get('name', KEY_FIELD_NAME)}")
-                key_decoder = make_engine_value_decoder(
-                    field_path,
-                    key_field_schema["type"],
-                    analyze_type_info(key_type),
-                    for_key=True,
-                )
-                field_path.pop()
+                num_key_parts = src_type.get("num_key_parts", 1)
+                key_type_info = analyze_type_info(key_type)
+                key_decoder: Callable[..., Any] | None = None
+                if (
+                    isinstance(
+                        key_type_info.variant, (AnalyzedBasicType, AnalyzedAnyType)
+                    )
+                    and num_key_parts == 1
+                ):
+                    single_key_decoder = make_engine_value_decoder(
+                        field_path,
+                        engine_fields_schema[0]["type"],
+                        key_type_info,
+                        for_key=True,
+                    )
+
+                    def key_decoder(value: list[Any]) -> Any:
+                        return single_key_decoder(value[0])
+
+                else:
+                    key_decoder = make_engine_struct_decoder(
+                        field_path,
+                        engine_fields_schema[0:num_key_parts],
+                        key_type_info,
+                        for_key=True,
+                    )
                 value_decoder = make_engine_struct_decoder(
                     field_path,
-                    engine_fields_schema[1:],
+                    engine_fields_schema[num_key_parts:],
                     analyze_type_info(value_type),
                 )
 
                 def decode(value: Any) -> Any | None:
                     if value is None:
                         return None
-                    return {key_decoder(v[0]): value_decoder(v[1:]) for v in value}
+                    return {
+                        key_decoder(v[0:num_key_parts]): value_decoder(
+                            v[num_key_parts:]
+                        )
+                        for v in value
+                    }
 
         return decode
 

python/cocoindex/typing.py

@@ -330,35 +330,50 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
 
 def _encode_struct_schema(
     struct_type: type, key_type: type | None = None
-) -> dict[str, Any]:
+) -> tuple[dict[str, Any], int | None]:
     fields = []
 
-    def add_field(name: str, t: Any) -> None:
+    def add_field(name: str, analyzed_type: AnalyzedTypeInfo) -> None:
         try:
-            type_info = encode_enriched_type_info(analyze_type_info(t))
+            type_info = encode_enriched_type_info(analyzed_type)
         except ValueError as e:
             e.add_note(
                 f"Failed to encode annotation for field - "
-                f"{struct_type.__name__}.{name}: {t}"
+                f"{struct_type.__name__}.{name}: {analyzed_type.core_type}"
             )
             raise
         type_info["name"] = name
         fields.append(type_info)
 
+    def add_fields_from_struct(struct_type: type) -> None:
+        if dataclasses.is_dataclass(struct_type):
+            for field in dataclasses.fields(struct_type):
+                add_field(field.name, analyze_type_info(field.type))
+        elif is_namedtuple_type(struct_type):
+            for name, field_type in struct_type.__annotations__.items():
+                add_field(name, analyze_type_info(field_type))
+        else:
+            raise ValueError(f"Unsupported struct type: {struct_type}")
+
+    result: dict[str, Any] = {}
+    num_key_parts = None
     if key_type is not None:
-        add_field(KEY_FIELD_NAME, key_type)
+        key_type_info = analyze_type_info(key_type)
+        if isinstance(key_type_info.variant, AnalyzedBasicType):
+            add_field(KEY_FIELD_NAME, key_type_info)
+            num_key_parts = 1
+        elif isinstance(key_type_info.variant, AnalyzedStructType):
+            add_fields_from_struct(key_type_info.variant.struct_type)
+            num_key_parts = len(fields)
+        else:
+            raise ValueError(f"Unsupported key type: {key_type}")
 
-    if dataclasses.is_dataclass(struct_type):
-        for field in dataclasses.fields(struct_type):
-            add_field(field.name, field.type)
-    elif is_namedtuple_type(struct_type):
-        for name, field_type in struct_type.__annotations__.items():
-            add_field(name, field_type)
+    add_fields_from_struct(struct_type)
 
-    result: dict[str, Any] = {"fields": fields}
+    result["fields"] = fields
     if doc := inspect.getdoc(struct_type):
         result["description"] = doc
-    return result
+    return result, num_key_parts
 
 
 def _encode_type(type_info: AnalyzedTypeInfo) -> dict[str, Any]:
@@ -374,7 +389,7 @@ def _encode_type(type_info: AnalyzedTypeInfo) -> dict[str, Any]:
         return {"kind": variant.kind}
 
     if isinstance(variant, AnalyzedStructType):
-        encoded_type = _encode_struct_schema(variant.struct_type)
+        encoded_type, _ = _encode_struct_schema(variant.struct_type)
         encoded_type["kind"] = "Struct"
         return encoded_type
 
@@ -384,10 +399,8 @@ def _encode_type(type_info: AnalyzedTypeInfo) -> dict[str, Any]:
         if isinstance(elem_type_info.variant, AnalyzedStructType):
             if variant.vector_info is not None:
                 raise ValueError("LTable type must not have a vector info")
-            return {
-                "kind": "LTable",
-                "row": _encode_struct_schema(elem_type_info.variant.struct_type),
-            }
+            row_type, _ = _encode_struct_schema(elem_type_info.variant.struct_type)
+            return {"kind": "LTable", "row": row_type}
         else:
             vector_info = variant.vector_info
             return {
@@ -402,12 +415,14 @@ def _encode_type(type_info: AnalyzedTypeInfo) -> dict[str, Any]:
             raise ValueError(
                 f"KTable value must have a Struct type, got {value_type_info.core_type}"
             )
+        row_type, num_key_parts = _encode_struct_schema(
+            value_type_info.variant.struct_type,
+            variant.key_type,
+        )
         return {
             "kind": "KTable",
-            "row": _encode_struct_schema(
-                value_type_info.variant.struct_type,
-                variant.key_type,
-            ),
+            "row": row_type,
+            "num_key_parts": num_key_parts,
         }
 
     if isinstance(variant, AnalyzedUnionType):

src/base/schema.rs

@@ -136,52 +136,63 @@ impl std::fmt::Display for StructSchema {
 }
 
 #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(tag = "kind")]
 #[allow(clippy::enum_variant_names)]
 pub enum TableKind {
     /// An table with unordered rows, without key.
     UTable,
-    /// A table's first field is the key.
+    /// A table's first field is the key. The value is number of fields serving as the key
     #[serde(alias = "Table")]
-    KTable,
+    KTable {
+        // Omit the field if num_key_parts is 1 for backward compatibility.
+        #[serde(default = "default_num_key_parts", skip_serializing_if = "is_one")]
+        num_key_parts: usize,
+    },
+
     /// A table whose rows orders are preserved.
     #[serde(alias = "List")]
     LTable,
 }
 
+fn default_num_key_parts() -> usize {
+    1
+}
+
+fn is_one(value: &usize) -> bool {
+    *value == 1
+}
+
 impl std::fmt::Display for TableKind {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             TableKind::UTable => write!(f, "Table"),
-            TableKind::KTable => write!(f, "KTable"),
+            TableKind::KTable { num_key_parts } => write!(f, "KTable({num_key_parts})"),
             TableKind::LTable => write!(f, "LTable"),
         }
     }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 pub struct TableSchema {
+    #[serde(flatten)]
     pub kind: TableKind,
+
     pub row: StructSchema,
 }
 
+impl std::fmt::Display for TableSchema {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}({})", self.kind, self.row)
+    }
+}
+
 impl TableSchema {
-    pub fn has_key(&self) -> bool {
-        match self.kind {
-            TableKind::KTable => true,
-            TableKind::UTable | TableKind::LTable => false,
-        }
+    pub fn new(kind: TableKind, row: StructSchema) -> Self {
+        Self { kind, row }
     }
 
-    pub fn key_type(&self) -> Option<&EnrichedValueType> {
-        match self.kind {
-            TableKind::KTable => self
-                .row
-                .fields
-                .first()
-                .as_ref()
-                .map(|field| &field.value_type),
-            TableKind::UTable | TableKind::LTable => None,
-        }
+    pub fn has_key(&self) -> bool {
+        !self.key_schema().is_empty()
     }
 
     pub fn without_attrs(&self) -> Self {
@@ -190,23 +201,11 @@ impl TableSchema {
             row: self.row.without_attrs(),
         }
     }
-}
-
-impl std::fmt::Display for TableSchema {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}({})", self.kind, self.row)
-    }
-}
-
-impl TableSchema {
-    pub fn new(kind: TableKind, row: StructSchema) -> Self {
-        Self { kind, row }
-    }
 
-    pub fn key_field(&self) -> Option<&FieldSchema> {
+    pub fn key_schema(&self) -> &[FieldSchema] {
         match self.kind {
-            TableKind::KTable => Some(self.row.fields.first().unwrap()),
-            TableKind::UTable | TableKind::LTable => None,
+            TableKind::KTable { num_key_parts: n } => &self.row.fields[..n],
+            TableKind::UTable | TableKind::LTable => &[],
         }
     }
 }
@@ -224,11 +223,11 @@ pub enum ValueType {
 }
 
 impl ValueType {
-    pub fn key_type(&self) -> Option<&EnrichedValueType> {
+    pub fn key_schema(&self) -> &[FieldSchema] {
         match self {
-            ValueType::Basic(_) => None,
-            ValueType::Struct(_) => None,
-            ValueType::Table(c) => c.key_type(),
+            ValueType::Basic(_) => &[],
+            ValueType::Struct(_) => &[],
+            ValueType::Table(c) => c.key_schema(),
         }
     }