Merge branch 'main' into qdrant

Author: Anush008

Cargo.toml

@@ -25,7 +25,7 @@ log = "0.4.26"
 regex = "1.11.1"
 serde = { version = "1.0.219", features = ["derive"] }
 serde_json = "1.0.140"
-sqlx = { version = "0.8.3", features = ["chrono", "postgres", "runtime-tokio"] }
+sqlx = { version = "0.8.3", features = ["chrono", "postgres", "runtime-tokio", "uuid"] }
 tokio = { version = "1.44.1", features = [
     "macros",
     "rt-multi-thread",
@@ -90,4 +90,4 @@ http-body-util = "0.1.3"
 yaml-rust2 = "0.10.0"
 urlencoding = "2.1.3"
 qdrant-client = "1.13.0"
-
+uuid = { version = "1.16.0", features = ["serde", "v4", "v8"] }

docs/docs/core/data_types.mdx

@@ -24,6 +24,7 @@ This is the list of all basic types supported by CocoIndex:
 | float32 | `cocoindex.typing.Float32` |`float` | 
 | float64 |  `cocoindex.typing.Float64` |`float` |
 | range | `cocoindex.typing.Range`  | `tuple[int, int]` |
+| uuid | `uuid.UUId` | `uuid.UUID` |
 | vector[*type*, *N*?] |`Annotated[list[type], cocoindex.typing.Vector(dim=N)]` | `list[type]` | 
 | json | `cocoindex.typing.Json` | Any type convertible to JSON by `json` package | 
 
@@ -73,6 +74,7 @@ Currently, the following types are supported as types for key fields:
 - `bool`
 - `int64`
 - `range`
+- `uuid`
 - Struct with all fields being key types
 
 ### Vector Type

examples/pdf_embedding/main.py

@@ -55,13 +55,14 @@ def pdf_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoinde
 
         with doc["chunks"].row() as chunk:
             chunk["embedding"] = chunk["text"].call(text_to_embedding)
-            doc_embeddings.collect(filename=doc["filename"], location=chunk["location"],
+            doc_embeddings.collect(id=cocoindex.GeneratedField.UUID,
+                                   filename=doc["filename"], location=chunk["location"],
                                    text=chunk["text"], embedding=chunk["embedding"])
 
     doc_embeddings.export(
         "doc_embeddings",
         cocoindex.storages.Qdrant(qdrant_url="http://localhost:6333", collection_name="cocoindex"),
-        primary_key_fields=["filename", "location"],
+        primary_key_fields=["id"],
         vector_index=[("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)])
 
 query_handler = cocoindex.query.SimpleSemanticsQueryHandler(

python/cocoindex/__init__.py

@@ -2,7 +2,8 @@
 Cocoindex is a framework for building and running indexing pipelines.
 """
 from . import flow, functions, query, sources, storages, cli
-from .flow import FlowBuilder, DataScope, DataSlice, Flow, flow_def, EvaluateAndDumpOptions
+from .flow import FlowBuilder, DataScope, DataSlice, Flow, flow_def
+from .flow import EvaluateAndDumpOptions, GeneratedField
 from .llm import LlmSpec, LlmApiType
 from .vector import VectorSimilarityMetric
 from .lib import *

python/cocoindex/convert.py

@@ -3,6 +3,7 @@
 """
 import dataclasses
 import inspect
+import uuid
 
 from typing import Any, Callable
 from .typing import analyze_type_info, COLLECTION_TYPES
@@ -13,6 +14,8 @@ def to_engine_value(value: Any) -> Any:
         return [to_engine_value(getattr(value, f.name)) for f in dataclasses.fields(value)]
     if isinstance(value, (list, tuple)):
         return [to_engine_value(v) for v in value]
+    if isinstance(value, uuid.UUID):
+        return value.bytes
     return value
 
 def make_engine_value_converter(
@@ -62,6 +65,9 @@ def make_engine_value_converter(
         field_path.pop()
         return lambda value: [elem_converter(v) for v in value] if value is not None else None
 
+    if src_type_kind == 'Uuid':
+        return lambda value: uuid.UUID(bytes=value)
+
     return lambda value: value
 
 def _make_engine_struct_value_converter(

python/cocoindex/flow.py

@@ -238,6 +238,12 @@ def add_collector(self, name: str | None = None) -> DataCollector:
             )
         )
 
+class GeneratedField(Enum):
+    """
+    A generated field is automatically set by the engine.
+    """
+    UUID = "Uuid"
+
 class DataCollector:
     """A data collector is used to collect data into a collector."""
     _flow_builder_state: _FlowBuilderState
@@ -248,12 +254,25 @@ def __init__(self, flow_builder_state: _FlowBuilderState,
         self._flow_builder_state = flow_builder_state
         self._engine_data_collector = data_collector
 
-    def collect(self, **kwargs: DataSlice):
+    def collect(self, **kwargs: DataSlice | GeneratedField):
         """
         Collect data into the collector.
         """
+        regular_kwargs = []
+        auto_uuid_field = None
+        for k, v in kwargs.items():
+            if isinstance(v, GeneratedField):
+                if v == GeneratedField.UUID:
+                    if auto_uuid_field is not None:
+                        raise ValueError("Only one generated UUID field is allowed")
+                    auto_uuid_field = k
+                else:
+                    raise ValueError(f"Unexpected generated field: {v}")
+            else:
+                regular_kwargs.append((k, _data_slice_state(v).engine_data_slice))
+
         self._flow_builder_state.engine_flow_builder.collect(
-            self._engine_data_collector, [(k, _data_slice_state(v).engine_data_slice) for k, v in kwargs.items()])
+            self._engine_data_collector, regular_kwargs, auto_uuid_field)
 
     def export(self, name: str, target_spec: op.StorageSpec, /, *,
               primary_key_fields: Sequence[str] | None = None,

python/cocoindex/typing.py

@@ -3,6 +3,7 @@
 import dataclasses
 import types
 import inspect
+import uuid
 from typing import Annotated, NamedTuple, Any, TypeVar, TYPE_CHECKING, overload
 
 class Vector(NamedTuple):
@@ -130,6 +131,8 @@ def analyze_type_info(t) -> AnalyzedTypeInfo:
             kind = 'Int64'
         elif t is float:
             kind = 'Float64'
+        elif t is uuid.UUID:
+            kind = 'Uuid'
         else:
             raise ValueError(f"type unsupported yet: {t}")
 

src/base/json_schema.rs

@@ -3,12 +3,19 @@ use schemars::schema::{
     ArrayValidation, InstanceType, Metadata, ObjectValidation, Schema, SchemaObject, SingleOrVec,
 };
 
+pub struct ToJsonSchemaOptions {
+    /// If true, mark all fields as required.
+    /// Use union type (with `null`) for optional fields instead.
+    /// Models like OpenAI will reject the schema if a field is not required.
+    pub fields_always_required: bool,
+}
+
 pub trait ToJsonSchema {
-    fn to_json_schema(&self) -> SchemaObject;
+    fn to_json_schema(&self, options: &ToJsonSchemaOptions) -> SchemaObject;
 }
 
 impl ToJsonSchema for schema::BasicValueType {
-    fn to_json_schema(&self) -> SchemaObject {
+    fn to_json_schema(&self, options: &ToJsonSchemaOptions) -> SchemaObject {
         let mut schema = SchemaObject::default();
         match self {
             schema::BasicValueType::Str => {
@@ -26,9 +33,6 @@ impl ToJsonSchema for schema::BasicValueType {
             schema::BasicValueType::Float32 | schema::BasicValueType::Float64 => {
                 schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::Number)));
             }
-            schema::BasicValueType::Json => {
-                // Can be any value. No type constraint.
-            }
             schema::BasicValueType::Range => {
                 schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::Array)));
                 schema.array = Some(Box::new(ArrayValidation {
@@ -51,11 +55,18 @@ impl ToJsonSchema for schema::BasicValueType {
                     .description =
                     Some("A range, start pos (inclusive), end pos (exclusive).".to_string());
             }
+            schema::BasicValueType::Uuid => {
+                schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::String)));
+                schema.format = Some("uuid".to_string());
+            }
+            schema::BasicValueType::Json => {
+                // Can be any value. No type constraint.
+            }
             schema::BasicValueType::Vector(s) => {
                 schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::Array)));
                 schema.array = Some(Box::new(ArrayValidation {
                     items: Some(SingleOrVec::Single(Box::new(
-                        s.element_type.to_json_schema().into(),
+                        s.element_type.to_json_schema(options).into(),
                     ))),
                     min_items: s.dimension.and_then(|d| u32::try_from(d).ok()),
                     max_items: s.dimension.and_then(|d| u32::try_from(d).ok()),
@@ -68,7 +79,7 @@ impl ToJsonSchema for schema::BasicValueType {
 }
 
 impl ToJsonSchema for schema::StructSchema {
-    fn to_json_schema(&self) -> SchemaObject {
+    fn to_json_schema(&self, options: &ToJsonSchemaOptions) -> SchemaObject {
         SchemaObject {
             metadata: Some(Box::new(Metadata {
                 description: self.description.as_ref().map(|s| s.to_string()),
@@ -79,12 +90,25 @@ impl ToJsonSchema for schema::StructSchema {
                 properties: self
                     .fields
                     .iter()
-                    .map(|f| (f.name.to_string(), f.value_type.to_json_schema().into()))
+                    .map(|f| {
+                        let mut schema = f.value_type.to_json_schema(options);
+                        if options.fields_always_required && f.value_type.nullable {
+                            if let Some(instance_type) = &mut schema.instance_type {
+                                let mut types = match instance_type {
+                                    SingleOrVec::Single(t) => vec![**t],
+                                    SingleOrVec::Vec(t) => std::mem::take(t),
+                                };
+                                types.push(InstanceType::Null);
+                                *instance_type = SingleOrVec::Vec(types);
+                            }
+                        }
+                        (f.name.to_string(), schema.into())
+                    })
                     .collect(),
                 required: self
                     .fields
                     .iter()
-                    .filter(|&f| (!f.value_type.nullable))
+                    .filter(|&f| (options.fields_always_required || !f.value_type.nullable))
                     .map(|f| f.name.to_string())
                     .collect(),
                 additional_properties: Some(Schema::Bool(false).into()),
@@ -96,14 +120,16 @@ impl ToJsonSchema for schema::StructSchema {
 }
 
 impl ToJsonSchema for schema::ValueType {
-    fn to_json_schema(&self) -> SchemaObject {
+    fn to_json_schema(&self, options: &ToJsonSchemaOptions) -> SchemaObject {
         match self {
-            schema::ValueType::Basic(b) => b.to_json_schema(),
-            schema::ValueType::Struct(s) => s.to_json_schema(),
+            schema::ValueType::Basic(b) => b.to_json_schema(options),
+            schema::ValueType::Struct(s) => s.to_json_schema(options),
             schema::ValueType::Collection(c) => SchemaObject {
                 instance_type: Some(SingleOrVec::Single(Box::new(InstanceType::Array))),
                 array: Some(Box::new(ArrayValidation {
-                    items: Some(SingleOrVec::Single(Box::new(c.row.to_json_schema().into()))),
+                    items: Some(SingleOrVec::Single(Box::new(
+                        c.row.to_json_schema(options).into(),
+                    ))),
                     ..Default::default()
                 })),
                 ..Default::default()
@@ -113,7 +139,7 @@ impl ToJsonSchema for schema::ValueType {
 }
 
 impl ToJsonSchema for schema::EnrichedValueType {
-    fn to_json_schema(&self) -> SchemaObject {
-        self.typ.to_json_schema()
+    fn to_json_schema(&self, options: &ToJsonSchemaOptions) -> SchemaObject {
+        self.typ.to_json_schema(options)
     }
 }

src/base/schema.rs

@@ -35,6 +35,9 @@ pub enum BasicValueType {
     /// A range, with a start offset and a length.
     Range,
 
+    /// A UUID.
+    Uuid,
+
     /// A JSON value.
     Json,
 
@@ -52,6 +55,7 @@ impl std::fmt::Display for BasicValueType {
             BasicValueType::Float32 => write!(f, "float32"),
             BasicValueType::Float64 => write!(f, "float64"),
             BasicValueType::Range => write!(f, "range"),
+            BasicValueType::Uuid => write!(f, "uuid"),
             BasicValueType::Json => write!(f, "json"),
             BasicValueType::Vector(s) => write!(
                 f,
@@ -121,7 +125,7 @@ pub struct CollectionSchema {
     pub row: StructSchema,
 
     #[serde(default = "Vec::new", skip_serializing_if = "Vec::is_empty")]
-    pub collectors: Vec<NamedSpec<StructSchema>>,
+    pub collectors: Vec<NamedSpec<Arc<CollectorSchema>>>,
 }
 
 impl CollectionSchema {
@@ -153,7 +157,7 @@ impl CollectionSchema {
                 .iter()
                 .map(|c| NamedSpec {
                     name: c.name.clone(),
-                    spec: c.spec.without_attrs(),
+                    spec: Arc::from(c.spec.without_attrs()),
                 })
                 .collect(),
         }
@@ -335,13 +339,65 @@ impl std::fmt::Display for FieldSchema {
     }
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+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<usize>,
+}
+
+impl std::fmt::Display for CollectorSchema {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "Collector(")?;
+        for (i, field) in self.fields.iter().enumerate() {
+            if i > 0 {
+                write!(f, ", ")?;
+            }
+            write!(f, "{}", field)?;
+        }
+        write!(f, ")")
+    }
+}
+
+impl CollectorSchema {
+    pub fn from_fields(fields: Vec<FieldSchema>, auto_uuid_field: Option<FieldName>) -> Self {
+        let mut fields = fields;
+        let auto_uuid_field_idx = if let Some(auto_uuid_field) = auto_uuid_field {
+            fields.insert(
+                0,
+                FieldSchema::new(
+                    auto_uuid_field,
+                    EnrichedValueType {
+                        typ: ValueType::Basic(BasicValueType::Uuid),
+                        nullable: false,
+                        attrs: Default::default(),
+                    },
+                ),
+            );
+            Some(0)
+        } else {
+            None
+        };
+        Self {
+            fields,
+            auto_uuid_field_idx,
+        }
+    }
+    pub fn without_attrs(&self) -> Self {
+        Self {
+            fields: self.fields.iter().map(|f| f.without_attrs()).collect(),
+            auto_uuid_field_idx: self.auto_uuid_field_idx,
+        }
+    }
+}
+
 /// Top-level schema for a flow instance.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct DataSchema {
     pub schema: StructSchema,
 
     #[serde(default = "Vec::new", skip_serializing_if = "Vec::is_empty")]
-    pub collectors: Vec<NamedSpec<StructSchema>>,
+    pub collectors: Vec<NamedSpec<Arc<CollectorSchema>>>,
 }
 
 impl Deref for DataSchema {

src/base/spec.rs

@@ -181,9 +181,15 @@ pub struct ForEachOpSpec {
 /// Emit data to a given collector at the given scope.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct CollectOpSpec {
+    /// Field values to be collected.
     pub input: StructMapping,
+    /// Scope for the collector.
     pub scope_name: ScopeName,
+    /// Name of the collector.
     pub collector_name: FieldName,
+    /// If specified, the collector will have an automatically generated UUID field with the given name.
+    /// The uuid will remain stable when collected input values remain unchanged.
+    pub auto_uuid_field: Option<FieldName>,
 }
 
 #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]