From 62b2879592daaef224ca45bfe69ccf80fa8ea70e Mon Sep 17 00:00:00 2001 From: zhhanz Date: Tue, 14 Jul 2026 15:23:23 -0700 Subject: [PATCH 01/15] Regroup Rust SDK modules --- .../file_summarization/rust/src/main.rs | 2 +- dev/agent-skills/target-connector/SKILL.md | 2 + examples/rust/amazon_s3_embedding/Cargo.toml | 2 +- examples/rust/amazon_s3_embedding/README.md | 2 +- examples/rust/amazon_s3_embedding/src/main.rs | 6 +-- examples/rust/audio_to_text/README.md | 2 +- examples/rust/audio_to_text/src/main.rs | 4 +- examples/rust/code_embedding/Cargo.toml | 2 +- examples/rust/code_embedding/README.md | 2 +- examples/rust/code_embedding/src/main.rs | 2 +- .../rust/code_embedding_lancedb/README.md | 6 +-- .../rust/code_embedding_lancedb/src/main.rs | 4 +- .../rust/conversation_to_knowledge/README.md | 2 +- .../conversation_to_knowledge/src/clients.rs | 2 +- .../conversation_to_knowledge/src/main.rs | 1 - .../conversation_to_knowledge/src/pipeline.rs | 2 +- examples/rust/csv_to_iggy/README.md | 4 +- examples/rust/csv_to_iggy/src/main.rs | 6 +-- examples/rust/csv_to_kafka/src/main.rs | 6 +-- examples/rust/files_to_sqlite/README.md | 4 +- examples/rust/files_to_sqlite/src/main.rs | 4 +- examples/rust/files_transform/src/main.rs | 2 +- .../rust/gdrive_text_embedding/Cargo.toml | 2 +- .../rust/gdrive_text_embedding/src/main.rs | 4 +- examples/rust/hn_trending_topics/README.md | 2 +- examples/rust/hn_trending_topics/src/main.rs | 2 +- examples/rust/image_search/src/main.rs | 2 +- examples/rust/image_search_colpali/README.md | 2 +- .../rust/image_search_colpali/src/main.rs | 4 +- examples/rust/kafka_consume/README.md | 2 +- examples/rust/kafka_consume/src/main.rs | 3 +- .../meeting_notes_graph_falkordb/README.md | 2 +- .../meeting_notes_graph_falkordb/src/main.rs | 4 +- .../rust/meeting_notes_graph_neo4j/README.md | 2 +- .../meeting_notes_graph_neo4j/src/main.rs | 4 +- .../multi_codebase_summarization/src/main.rs | 2 +- .../oci_object_storage_embedding/Cargo.toml | 2 +- .../oci_object_storage_embedding/README.md | 2 +- .../oci_object_storage_embedding/src/main.rs | 6 +-- examples/rust/paper_metadata/Cargo.toml | 2 +- examples/rust/paper_metadata/README.md | 2 +- examples/rust/paper_metadata/src/main.rs | 6 +-- examples/rust/pdf_embedding/Cargo.toml | 2 +- examples/rust/pdf_embedding/README.md | 2 +- examples/rust/pdf_embedding/src/main.rs | 2 +- examples/rust/pdf_to_markdown/README.md | 2 +- examples/rust/pdf_to_markdown/src/main.rs | 2 +- examples/rust/postgres_source/Cargo.toml | 2 +- examples/rust/postgres_source/src/main.rs | 2 +- examples/rust/text_embedding/Cargo.toml | 2 +- examples/rust/text_embedding/README.md | 2 +- examples/rust/text_embedding/src/main.rs | 4 +- .../rust/text_embedding_lancedb/README.md | 10 ++-- .../rust/text_embedding_lancedb/src/main.rs | 4 +- examples/rust/text_embedding_qdrant/README.md | 10 ++-- .../rust/text_embedding_qdrant/src/main.rs | 4 +- .../rust/text_embedding_turbopuffer/README.md | 10 ++-- .../text_embedding_turbopuffer/src/main.rs | 4 +- rust/core/src/engine/progress_display.rs | 2 +- rust/sdk/SHOWCASE.md | 5 +- rust/sdk/cocoindex/benches/sdk_microbench.rs | 2 +- .../src/{ => connectors}/amazon_s3.rs | 8 ++-- .../cocoindex/src/{ => connectors}/doris.rs | 0 .../src/{ => connectors}/falkordb.rs | 0 .../cocoindex/src/{ => connectors}/gdrive.rs | 2 +- .../cocoindex/src/{ => connectors}/iggy.rs | 2 +- .../cocoindex/src/{ => connectors}/kafka.rs | 0 .../cocoindex/src/{ => connectors}/lancedb.rs | 0 rust/sdk/cocoindex/src/connectors/mod.rs | 32 +++++++++++++ .../cocoindex/src/{ => connectors}/neo4j.rs | 0 .../{ => connectors}/oci_object_storage.rs | 12 ++--- .../src/{ => connectors}/postgres.rs | 0 .../cocoindex/src/{ => connectors}/qdrant.rs | 0 .../cocoindex/src/{ => connectors}/sqlite.rs | 0 .../src/{ => connectors}/surrealdb.rs | 0 .../src/{ => connectors}/turbopuffer.rs | 0 .../cocoindex/src/{ => connectors}/valkey.rs | 0 rust/sdk/cocoindex/src/lib.rs | 46 +------------------ rust/sdk/cocoindex/src/memo.rs | 10 ++-- rust/sdk/cocoindex/src/ops/api.rs | 2 +- rust/sdk/cocoindex/src/prelude.rs | 18 ++++---- .../sdk/cocoindex/src/{ => resources}/file.rs | 0 rust/sdk/cocoindex/src/{ => resources}/fs.rs | 8 ++-- rust/sdk/cocoindex/src/{ => resources}/id.rs | 0 rust/sdk/cocoindex/src/resources/mod.rs | 3 ++ rust/sdk/cocoindex/src/row_schema.rs | 2 +- rust/sdk/cocoindex/tests/amazon_s3_source.rs | 8 ++-- rust/sdk/cocoindex/tests/doris_target.rs | 4 +- rust/sdk/cocoindex/tests/fs_live.rs | 4 +- rust/sdk/cocoindex/tests/gdrive_source.rs | 6 ++- .../sdk/cocoindex/tests/graph_vector_index.rs | 4 +- rust/sdk/cocoindex/tests/iggy_source.rs | 2 +- rust/sdk/cocoindex/tests/iggy_target.rs | 2 +- rust/sdk/cocoindex/tests/kafka_source.rs | 3 +- rust/sdk/cocoindex/tests/kafka_target.rs | 3 +- rust/sdk/cocoindex/tests/lancedb_target.rs | 2 +- rust/sdk/cocoindex/tests/oci_live.rs | 5 +- rust/sdk/cocoindex/tests/pipeline.rs | 16 +++---- rust/sdk/cocoindex/tests/postgres_source.rs | 3 +- rust/sdk/cocoindex/tests/postgres_target.rs | 3 +- rust/sdk/cocoindex/tests/qdrant_target.rs | 2 +- rust/sdk/cocoindex/tests/schema_from_row.rs | 10 ++-- rust/sdk/cocoindex/tests/sqlite_target.rs | 3 +- rust/sdk/cocoindex/tests/surrealdb_target.rs | 2 +- .../sdk/cocoindex/tests/turbopuffer_target.rs | 4 +- rust/sdk/cocoindex/tests/valkey_target.rs | 4 +- 106 files changed, 219 insertions(+), 215 deletions(-) rename rust/sdk/cocoindex/src/{ => connectors}/amazon_s3.rs (99%) rename rust/sdk/cocoindex/src/{ => connectors}/doris.rs (100%) rename rust/sdk/cocoindex/src/{ => connectors}/falkordb.rs (100%) rename rust/sdk/cocoindex/src/{ => connectors}/gdrive.rs (99%) rename rust/sdk/cocoindex/src/{ => connectors}/iggy.rs (99%) rename rust/sdk/cocoindex/src/{ => connectors}/kafka.rs (100%) rename rust/sdk/cocoindex/src/{ => connectors}/lancedb.rs (100%) create mode 100644 rust/sdk/cocoindex/src/connectors/mod.rs rename rust/sdk/cocoindex/src/{ => connectors}/neo4j.rs (100%) rename rust/sdk/cocoindex/src/{ => connectors}/oci_object_storage.rs (99%) rename rust/sdk/cocoindex/src/{ => connectors}/postgres.rs (100%) rename rust/sdk/cocoindex/src/{ => connectors}/qdrant.rs (100%) rename rust/sdk/cocoindex/src/{ => connectors}/sqlite.rs (100%) rename rust/sdk/cocoindex/src/{ => connectors}/surrealdb.rs (100%) rename rust/sdk/cocoindex/src/{ => connectors}/turbopuffer.rs (100%) rename rust/sdk/cocoindex/src/{ => connectors}/valkey.rs (100%) rename rust/sdk/cocoindex/src/{ => resources}/file.rs (100%) rename rust/sdk/cocoindex/src/{ => resources}/fs.rs (99%) rename rust/sdk/cocoindex/src/{ => resources}/id.rs (100%) diff --git a/benchmarks/file_summarization/rust/src/main.rs b/benchmarks/file_summarization/rust/src/main.rs index e3df64c5b..62fcdf69b 100644 --- a/benchmarks/file_summarization/rust/src/main.rs +++ b/benchmarks/file_summarization/rust/src/main.rs @@ -5,7 +5,7 @@ use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; -use cocoindex::fs::{FileEntry, walk}; +use cocoindex::resources::fs::{FileEntry, walk}; use cocoindex::prelude::*; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; diff --git a/dev/agent-skills/target-connector/SKILL.md b/dev/agent-skills/target-connector/SKILL.md index 1d8c43ae4..852c7f485 100644 --- a/dev/agent-skills/target-connector/SKILL.md +++ b/dev/agent-skills/target-connector/SKILL.md @@ -354,6 +354,8 @@ def test_vector_support(connector_with_vec: tuple[Connection, Path]) -> None: **Reference implementations:** - `python/tests/connectors/test_sqlite_target.py` - SQLite tests with vector support +- Rust SDK target connectors live under `rust/sdk/cocoindex/src/connectors/`; see + `sqlite.rs`, `postgres.rs`, and `doris.rs` there for relational connector patterns. ## Attachment Providers diff --git a/examples/rust/amazon_s3_embedding/Cargo.toml b/examples/rust/amazon_s3_embedding/Cargo.toml index c38345027..cfec7e2e7 100644 --- a/examples/rust/amazon_s3_embedding/Cargo.toml +++ b/examples/rust/amazon_s3_embedding/Cargo.toml @@ -18,7 +18,7 @@ cocoindex = { path = "../../../rust/sdk/cocoindex", features = [ ] } tokio = { version = "1", features = ["full"] } serde = { version = "1", features = ["derive"] } -# Query path (pgvector similarity); target writes go through cocoindex::postgres. +# Query path (pgvector similarity); target writes go through cocoindex::connectors::postgres. sqlx = { version = "0.8", default-features = false, features = [ "runtime-tokio", "tls-rustls", diff --git a/examples/rust/amazon_s3_embedding/README.md b/examples/rust/amazon_s3_embedding/README.md index 02276b7c5..a6d8990ca 100644 --- a/examples/rust/amazon_s3_embedding/README.md +++ b/examples/rust/amazon_s3_embedding/README.md @@ -10,7 +10,7 @@ serves similarity search. | Concern | Python | Rust (this example) | | ---------------- | ---------------------------------------- | ---------------------------------------------------- | -| Source | `amazon_s3.list_objects` (aiobotocore) | `cocoindex::amazon_s3::list_objects` (aws-sdk-s3) | +| Source | `amazon_s3.list_objects` (aiobotocore) | `cocoindex::connectors::amazon_s3::list_objects` (aws-sdk-s3) | | Per-file compute | `@coco.fn(memo=True) process_file` | `#[cocoindex::function(memo)] process_file` | | Chunking | `RecursiveSplitter` (markdown, 2000/500) | `cocoindex_ops_text` `RecursiveChunker` (markdown) | | Embeddings | `sentence-transformers/all-MiniLM-L6-v2` | `fastembed` `AllMiniLML6V2` (same model, 384-dim) | diff --git a/examples/rust/amazon_s3_embedding/src/main.rs b/examples/rust/amazon_s3_embedding/src/main.rs index 218a4710d..6dd1f47d3 100644 --- a/examples/rust/amazon_s3_embedding/src/main.rs +++ b/examples/rust/amazon_s3_embedding/src/main.rs @@ -14,11 +14,11 @@ use std::path::PathBuf; use std::sync::{Arc, LazyLock}; -use cocoindex::amazon_s3::{self, ListOptions, S3Client, S3File}; -use cocoindex::file::PatternFilePathMatcher; +use cocoindex::connectors::amazon_s3::{self, ListOptions, S3Client, S3File}; +use cocoindex::resources::file::PatternFilePathMatcher; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; -use cocoindex::postgres; +use cocoindex::connectors::postgres; use cocoindex::prelude::*; use sqlx::Row; use sqlx::postgres::{PgPool, PgPoolOptions}; diff --git a/examples/rust/audio_to_text/README.md b/examples/rust/audio_to_text/README.md index 02d629025..5337795f3 100644 --- a/examples/rust/audio_to_text/README.md +++ b/examples/rust/audio_to_text/README.md @@ -10,7 +10,7 @@ available transcriptions. | Concern | Python | Rust (this example) | | ---------------- | ---------------------------------------- | ---------------------------------------------------- | -| Source | `localfs.walk_dir` | `cocoindex::fs::walk` | +| Source | `localfs.walk_dir` | `cocoindex::resources::fs::walk` | | Per-file compute | `@coco.fn(memo=True) process_file` | `#[cocoindex::function(memo)] transcribe` | | Transcription | `LiteLLMTranscriber("whisper-1")` | OpenAI `/v1/audio/transcriptions` (`whisper-1`) | | Target | `postgres.mount_table_target` | `postgres::mount_table_target` | diff --git a/examples/rust/audio_to_text/src/main.rs b/examples/rust/audio_to_text/src/main.rs index 936046cb3..b777fd673 100644 --- a/examples/rust/audio_to_text/src/main.rs +++ b/examples/rust/audio_to_text/src/main.rs @@ -9,7 +9,7 @@ //! cargo run -- [AUDIO_DIR] # default AUDIO_DIR = ./audio_files //! //! Parallels the Python example: -//! - source : `cocoindex::fs::walk` (cf. `localfs.walk_dir`) +//! - source : `cocoindex::resources::fs::walk` (cf. `localfs.walk_dir`) //! - per-file work : `#[cocoindex::function(memo)]` (cf. `@coco.fn(memo=True)`) //! - transcription : `cocoindex::ops::api::ApiTranscriber` (cf. `LiteLLMTranscriber("whisper-1")`) //! - target : `postgres::TableTarget` (cf. `postgres.mount_table_target`) @@ -18,7 +18,7 @@ use std::path::PathBuf; use std::sync::LazyLock; use cocoindex::ops::api::ApiTranscriber; -use cocoindex::postgres; +use cocoindex::connectors::postgres; use cocoindex::prelude::*; const TABLE: &str = "audio_transcriptions"; diff --git a/examples/rust/code_embedding/Cargo.toml b/examples/rust/code_embedding/Cargo.toml index 8f670668f..6016b799a 100644 --- a/examples/rust/code_embedding/Cargo.toml +++ b/examples/rust/code_embedding/Cargo.toml @@ -16,7 +16,7 @@ cocoindex = { path = "../../../rust/sdk/cocoindex", features = [ ] } tokio = { version = "1", features = ["full"] } serde = { version = "1", features = ["derive"] } -# Postgres query path. Target writes go through cocoindex::postgres. +# Postgres query path. Target writes go through cocoindex::connectors::postgres. sqlx = { version = "0.8", default-features = false, features = [ "runtime-tokio", "tls-rustls", diff --git a/examples/rust/code_embedding/README.md b/examples/rust/code_embedding/README.md index 95f2b0c1a..85ffab7a1 100644 --- a/examples/rust/code_embedding/README.md +++ b/examples/rust/code_embedding/README.md @@ -14,7 +14,7 @@ Pipeline: **walk → detect language → tree-sitter chunk → embed → store i | Chunking | `RecursiveSplitter` | `cocoindex_ops_text::split::RecursiveChunker` | | Embeddings | `SentenceTransformerEmbedder` (all-MiniLM-L6-v2) | `fastembed` `AllMiniLML6V2` — **the same model**, local ONNX | | Embedder change-detection | `ContextKey(..., detect_change=True)` | `ContextKey::new_with_state("embedder", \|e\| e.model_name)` | -| Vector store | `postgres.TableTarget` + `declare_vector_index` | `cocoindex::postgres` `TableTarget` + `declare_vector_index` | +| Vector store | `postgres.TableTarget` + `declare_vector_index` | `cocoindex::connectors::postgres` `TableTarget` + `declare_vector_index` | | Stable row ids | `IdGenerator.next_id(chunk.text)` | `IdGenerator::next_id(ctx, chunk_text)` | | Query | pgvector `<=>` | pgvector `<=>` | diff --git a/examples/rust/code_embedding/src/main.rs b/examples/rust/code_embedding/src/main.rs index d16c1d28b..a9a0fed76 100644 --- a/examples/rust/code_embedding/src/main.rs +++ b/examples/rust/code_embedding/src/main.rs @@ -19,7 +19,7 @@ use std::sync::LazyLock; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter, detect_code_language}; -use cocoindex::postgres; +use cocoindex::connectors::postgres; use cocoindex::prelude::*; use serde::{Deserialize, Serialize}; use sqlx::Row; diff --git a/examples/rust/code_embedding_lancedb/README.md b/examples/rust/code_embedding_lancedb/README.md index bf85c0b8b..603cb50eb 100644 --- a/examples/rust/code_embedding_lancedb/README.md +++ b/examples/rust/code_embedding_lancedb/README.md @@ -6,18 +6,18 @@ Walks a source tree, detects each file's language, chunks it (tree-sitter-aware) embeds the chunks, and stores them in **LanceDB** — then serves vector search. Same pipeline as [`code_embedding`](../code_embedding), but the target is the -native `cocoindex::lancedb` connector instead of Postgres/pgvector. +native `cocoindex::connectors::lancedb` connector instead of Postgres/pgvector. ## Parallel to the Python example | Concern | Python | Rust (this example) | | ---------------- | ---------------------------------------- | -------------------------------------------------- | -| Source | `localfs.walk_dir` | `cocoindex::fs::walk` | +| Source | `localfs.walk_dir` | `cocoindex::resources::fs::walk` | | Per-file compute | `@coco.fn(memo=True) process_file` | `#[cocoindex::function(memo)] process_file` | | Language detect | `detect_code_language` | `cocoindex_ops_text::prog_langs::detect_language` | | Chunking | `RecursiveSplitter` (1000/300/300) | `cocoindex_ops_text` `RecursiveChunker` (1000/300/300) | | Embeddings | `sentence-transformers/all-MiniLM-L6-v2` | `fastembed` `AllMiniLML6V2` (same model, 384-dim) | -| Target | `lancedb.mount_table_target` | `cocoindex::lancedb::mount_table_target` | +| Target | `lancedb.mount_table_target` | `cocoindex::connectors::lancedb::mount_table_target` | Incrementality: unchanged files are memo-skipped; chunks of a removed/edited file are reconciled away by the managed LanceDB `TableTarget`. diff --git a/examples/rust/code_embedding_lancedb/src/main.rs b/examples/rust/code_embedding_lancedb/src/main.rs index 7c760a37a..19f36bb97 100644 --- a/examples/rust/code_embedding_lancedb/src/main.rs +++ b/examples/rust/code_embedding_lancedb/src/main.rs @@ -7,13 +7,13 @@ //! cargo run -- query "your query" # LanceDB vector search //! //! Same pipeline as the `code_embedding` example, but the target is the native -//! `cocoindex::lancedb` connector instead of Postgres/pgvector (parallels the +//! `cocoindex::connectors::lancedb` connector instead of Postgres/pgvector (parallels the //! Python example's use of `cocoindex.connectors.lancedb`). use std::path::PathBuf; use std::sync::LazyLock; -use cocoindex::lancedb::{self, ColumnDef, ColumnType, LanceDatabase, TableSchema}; +use cocoindex::connectors::lancedb::{self, ColumnDef, ColumnType, LanceDatabase, TableSchema}; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter, detect_code_language}; use cocoindex::prelude::*; diff --git a/examples/rust/conversation_to_knowledge/README.md b/examples/rust/conversation_to_knowledge/README.md index a205cc87b..2ff7a3017 100644 --- a/examples/rust/conversation_to_knowledge/README.md +++ b/examples/rust/conversation_to_knowledge/README.md @@ -17,7 +17,7 @@ sessions, statements, persons, techs, orgs, and the relationships between them. | LLM extraction (2 passes) | `instructor` + `litellm` | `reqwest` → OpenAI JSON mode | | Stable ids | `IdGenerator` | `cocoindex::IdGenerator` | | Entity resolution | `ops.entity_resolution` (faiss + LLM) | `cocoindex::entity_resolution` + `fastembed` Snowflake embeddings + LLM pair resolver | -| Graph store | `surrealdb` connector (`TableTarget`/`RelationTarget`) | `cocoindex::surrealdb` targets over the native `surrealdb` crate | +| Graph store | `surrealdb` connector (`TableTarget`/`RelationTarget`) | `cocoindex::connectors::surrealdb` targets over the native `surrealdb` crate | | Embedder change-detection | `ContextKey(..., detect_change=True)` | `ContextKey::new_with_state(...)` | ### Design notes / where it differs diff --git a/examples/rust/conversation_to_knowledge/src/clients.rs b/examples/rust/conversation_to_knowledge/src/clients.rs index 17656972a..c2c520776 100644 --- a/examples/rust/conversation_to_knowledge/src/clients.rs +++ b/examples/rust/conversation_to_knowledge/src/clients.rs @@ -5,7 +5,7 @@ use std::path::PathBuf; use std::sync::{Arc, LazyLock}; use cocoindex::prelude::*; -pub use cocoindex::surrealdb::Graph; +pub use cocoindex::connectors::surrealdb::Graph; use fastembed::{ EmbeddingModel, InitOptions, InitOptionsUserDefined, Pooling, TextEmbedding, TokenizerFiles, UserDefinedEmbeddingModel, read_file_to_bytes, diff --git a/examples/rust/conversation_to_knowledge/src/main.rs b/examples/rust/conversation_to_knowledge/src/main.rs index 8b5ce4848..55290f07a 100644 --- a/examples/rust/conversation_to_knowledge/src/main.rs +++ b/examples/rust/conversation_to_knowledge/src/main.rs @@ -28,7 +28,6 @@ use std::collections::HashMap; use std::path::PathBuf; use cocoindex::prelude::*; -use cocoindex::walk; use serde::Deserialize; use clients::{Embedder, Graph, LlmClient}; diff --git a/examples/rust/conversation_to_knowledge/src/pipeline.rs b/examples/rust/conversation_to_knowledge/src/pipeline.rs index 609ece0b8..59f4b3caa 100644 --- a/examples/rust/conversation_to_knowledge/src/pipeline.rs +++ b/examples/rust/conversation_to_knowledge/src/pipeline.rs @@ -5,7 +5,7 @@ use std::collections::{HashMap, HashSet}; use async_trait::async_trait; use cocoindex::entity_resolution; use cocoindex::prelude::*; -use cocoindex::surrealdb; +use cocoindex::connectors::surrealdb; use serde::Deserialize; use crate::clients::{EMBEDDER, Embedder, GRAPH, LLM, LlmClient, RESOLVER_LLM}; diff --git a/examples/rust/csv_to_iggy/README.md b/examples/rust/csv_to_iggy/README.md index 1d3662ef2..13ee7d1d2 100644 --- a/examples/rust/csv_to_iggy/README.md +++ b/examples/rust/csv_to_iggy/README.md @@ -3,7 +3,7 @@ Rust analogue of the Python [`csv_to_kafka`](../../csv_to_kafka) example, targeting **Apache Iggy** instead of Kafka. (Python ships an `iggy` connector with the same shape as its `kafka` one; this mirrors the Rust [`csv_to_kafka`](../csv_to_kafka) -example onto `cocoindex::iggy`.) +example onto `cocoindex::connectors::iggy`.) Reads local CSV files, converts each row to a JSON object (header row as keys), and publishes one Iggy message per row via CocoIndex's declarative @@ -14,7 +14,7 @@ and publishes one Iggy message per row via CocoIndex's declarative | Concern | Python (`csv_to_kafka` / `iggy`) | Rust (this example) | | ---------------- | ---------------------------------------- | --------------------------------------------------- | | Per-file compute | `@coco.fn(memo=True)` | `#[cocoindex::function(memo)] process_csv` | -| Target | `mount_kafka_topic_target` / `mount_iggy_topic_target` | `cocoindex::iggy::mount_iggy_topic_target` | +| Target | `mount_kafka_topic_target` / `mount_iggy_topic_target` | `cocoindex::connectors::iggy::mount_iggy_topic_target` | | Declare a message| `target.declare_target_state(key, value)`| `target.declare_message(key, value)` | Incrementality (two layers): unchanged CSV files are memo-skipped, and the diff --git a/examples/rust/csv_to_iggy/src/main.rs b/examples/rust/csv_to_iggy/src/main.rs index d803c8482..d4ae33147 100644 --- a/examples/rust/csv_to_iggy/src/main.rs +++ b/examples/rust/csv_to_iggy/src/main.rs @@ -18,8 +18,8 @@ //! Stream/topic are user-managed (CocoIndex never creates/drops them); `index` //! creates them up front as a convenience. -use cocoindex::fs::FileEntry; -use cocoindex::iggy::{self, IggyProducer}; +use cocoindex::resources::fs::FileEntry; +use cocoindex::connectors::iggy::{self, IggyProducer}; use cocoindex::prelude::*; use iggy::prelude::{ CompressionAlgorithm, Consumer, Identifier, IggyExpiry, MaxTopicSize, MessageClient, @@ -120,7 +120,7 @@ async fn index(producer: IggyProducer, stream: String, topic: String) -> Result< let target = iggy::mount_iggy_topic_target(&ctx, &producer, stream, topic, options)?; - let files = cocoindex::fs::walk_items("./data", &["**/*.csv"])?; + let files = cocoindex::resources::fs::walk_items("./data", &["**/*.csv"])?; println!("found {} CSV file(s)", files.len()); let per_file = mount_each!(files, |file| process_csv(ctx, file)).await?; diff --git a/examples/rust/csv_to_kafka/src/main.rs b/examples/rust/csv_to_kafka/src/main.rs index 50880edae..b01d0d97d 100644 --- a/examples/rust/csv_to_kafka/src/main.rs +++ b/examples/rust/csv_to_kafka/src/main.rs @@ -22,8 +22,8 @@ //! Note: unlike the Python example (`live=True` continuous watch), this runs a //! single pass per `index` invocation — the Rust SDK's `fs::walk` is one-shot. -use cocoindex::fs::FileEntry; -use cocoindex::kafka::{self, KafkaProducer}; +use cocoindex::resources::fs::FileEntry; +use cocoindex::connectors::kafka::{self, KafkaProducer}; use cocoindex::prelude::*; use rskafka::client::ClientBuilder; use rskafka::client::partition::UnknownTopicHandling; @@ -90,7 +90,7 @@ async fn index(producer: KafkaProducer, topic: String) -> Result<()> { kafka::KafkaTopicOptions::default(), )?; - let files = cocoindex::fs::walk_items("./data", &["**/*.csv"])?; + let files = cocoindex::resources::fs::walk_items("./data", &["**/*.csv"])?; println!("found {} CSV file(s)", files.len()); let per_file = mount_each!(files, |file| process_csv(ctx, file)).await?; diff --git a/examples/rust/files_to_sqlite/README.md b/examples/rust/files_to_sqlite/README.md index 466e552a0..bdb7c5c3c 100644 --- a/examples/rust/files_to_sqlite/README.md +++ b/examples/rust/files_to_sqlite/README.md @@ -1,7 +1,7 @@ # Files → SQLite (Rust) A self-contained example of the Rust SDK's embedded **SQLite table target** -(`cocoindex::sqlite`). No external server is needed — SQLite is embedded. +(`cocoindex::connectors::sqlite`). No external server is needed — SQLite is embedded. It walks `**/*.md` / `**/*.txt` files under a source directory, computes a small per-file summary (word count + first line), and writes one row per file into a @@ -14,7 +14,7 @@ rows for files that disappear. | Connection | `sqlite::Database::connect(path)` | | Schema | `sqlite::TableSchema` + `sqlite::ColumnDef` | | Target | `sqlite::mount_table_target` + `TableTarget::declare_row` | -| Source | `cocoindex::fs::walk` + `Ctx::mount_each` | +| Source | `cocoindex::resources::fs::walk` + `Ctx::mount_each` | ## Run diff --git a/examples/rust/files_to_sqlite/src/main.rs b/examples/rust/files_to_sqlite/src/main.rs index 47cd64798..add015d15 100644 --- a/examples/rust/files_to_sqlite/src/main.rs +++ b/examples/rust/files_to_sqlite/src/main.rs @@ -16,7 +16,7 @@ use std::path::PathBuf; use std::sync::LazyLock; use cocoindex::prelude::*; -use cocoindex::sqlite; +use cocoindex::connectors::sqlite; use serde::{Deserialize, Serialize}; use sqlx::Row as _; @@ -85,7 +85,7 @@ async fn index(source_dir: PathBuf, db_path: String) -> Result<()> { async move { let table = sqlite::mount_table_target(&ctx, &DB, TABLE, files_schema()?).await?; - let files = cocoindex::fs::walk_items(&source_dir, &["**/*.md", "**/*.txt"])?; + let files = cocoindex::resources::fs::walk_items(&source_dir, &["**/*.md", "**/*.txt"])?; mount_each!(files, |file| process_file(ctx, file, table)).await?; Ok(()) } diff --git a/examples/rust/files_transform/src/main.rs b/examples/rust/files_transform/src/main.rs index 3978d7a79..23be5a022 100644 --- a/examples/rust/files_transform/src/main.rs +++ b/examples/rust/files_transform/src/main.rs @@ -85,7 +85,7 @@ async fn main() -> Result<()> { // `(key, file)` pairs for `mount_each!`. (Output names join the // relative-path components with `__`, so nested files don't // collide — see `output_name_for`.) - let files = cocoindex::fs::walk_items(&source_dir, &["**/*.md"])?; + let files = cocoindex::resources::fs::walk_items(&source_dir, &["**/*.md"])?; // One processing component per file. `ctx` is the parent; the // macro substitutes each child scope's `&Ctx`, and fingerprints diff --git a/examples/rust/gdrive_text_embedding/Cargo.toml b/examples/rust/gdrive_text_embedding/Cargo.toml index 561c2de6c..7863fa6b4 100644 --- a/examples/rust/gdrive_text_embedding/Cargo.toml +++ b/examples/rust/gdrive_text_embedding/Cargo.toml @@ -18,7 +18,7 @@ cocoindex = { path = "../../../rust/sdk/cocoindex", features = [ ] } tokio = { version = "1", features = ["full"] } serde = { version = "1", features = ["derive"] } -# Query path (pgvector similarity); target writes go through cocoindex::postgres. +# Query path (pgvector similarity); target writes go through cocoindex::connectors::postgres. sqlx = { version = "0.8", default-features = false, features = [ "runtime-tokio", "tls-rustls", diff --git a/examples/rust/gdrive_text_embedding/src/main.rs b/examples/rust/gdrive_text_embedding/src/main.rs index d1584c8fb..f0053dce0 100644 --- a/examples/rust/gdrive_text_embedding/src/main.rs +++ b/examples/rust/gdrive_text_embedding/src/main.rs @@ -7,10 +7,10 @@ use std::path::PathBuf; use std::sync::LazyLock; -use cocoindex::gdrive::{DriveFile, GoogleDriveClient, GoogleDriveSource}; +use cocoindex::connectors::gdrive::{DriveFile, GoogleDriveClient, GoogleDriveSource}; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; -use cocoindex::postgres; +use cocoindex::connectors::postgres; use cocoindex::prelude::*; use serde::{Deserialize, Serialize}; use sqlx::Row; diff --git a/examples/rust/hn_trending_topics/README.md b/examples/rust/hn_trending_topics/README.md index 1c429cf2c..75fb47993 100644 --- a/examples/rust/hn_trending_topics/README.md +++ b/examples/rust/hn_trending_topics/README.md @@ -11,7 +11,7 @@ Scrapes recent HackerNews threads + comments via the public [Algolia HN API](htt | Fetch threads/comments | `aiohttp` + Algolia HN API | `reqwest` + Algolia HN API (no key) | | Per-thread incremental skip | component memo | `#[cocoindex::function(memo)]` on `process_thread` | | Topic extraction | `litellm` (`gemini-2.5-flash`) | `reqwest` → OpenAI JSON (`gpt-4o-mini`) | -| Store | `postgres.TableTarget` (`hn_messages`, `hn_topics`) | `cocoindex::postgres` `TableTarget` (same two tables) | +| Store | `postgres.TableTarget` (`hn_messages`, `hn_topics`) | `cocoindex::connectors::postgres` `TableTarget` (same two tables) | | Trending / search | SQL | SQL (same scoring: thread mention = 5, comment = 1) | **Differences:** Python defaults to Gemini; this uses OpenAI (`OPENAI_API_KEY`). Target writes are declarative in both SDKs; the incremental win is the per-thread memo plus target-state reconciliation. diff --git a/examples/rust/hn_trending_topics/src/main.rs b/examples/rust/hn_trending_topics/src/main.rs index 6549ce2b8..53274c8ba 100644 --- a/examples/rust/hn_trending_topics/src/main.rs +++ b/examples/rust/hn_trending_topics/src/main.rs @@ -19,7 +19,7 @@ use std::sync::LazyLock; -use cocoindex::postgres; +use cocoindex::connectors::postgres; use cocoindex::prelude::*; use serde::{Deserialize, Serialize}; use sqlx::Row; diff --git a/examples/rust/image_search/src/main.rs b/examples/rust/image_search/src/main.rs index f178ae757..7ad304b71 100644 --- a/examples/rust/image_search/src/main.rs +++ b/examples/rust/image_search/src/main.rs @@ -24,7 +24,7 @@ use std::sync::LazyLock; use cocoindex::ops::image::ImageEmbedder; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::prelude::*; -use cocoindex::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; +use cocoindex::connectors::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; use serde_json::json; /// CLIP ViT-B/32 vision tower (images) and text tower (queries). Both output diff --git a/examples/rust/image_search_colpali/README.md b/examples/rust/image_search_colpali/README.md index af3565a5c..5ed923bca 100644 --- a/examples/rust/image_search_colpali/README.md +++ b/examples/rust/image_search_colpali/README.md @@ -3,7 +3,7 @@ Rust port of the Python [`image_search_colpali`](../../image_search_colpali) example. Embeds each image with **ColPali** into a *list* of vectors (late-interaction), indexes them in a Qdrant **MAX_SIM multi-vector** collection -(`cocoindex::qdrant::CollectionSchema::multivector` + +(`cocoindex::connectors::qdrant::CollectionSchema::multivector` + `declare_multivector_point`), and searches with `qdrant::multivector_search`. ColPali has no pure-Rust model, so inference is offloaded to a small external diff --git a/examples/rust/image_search_colpali/src/main.rs b/examples/rust/image_search_colpali/src/main.rs index d6d2a3b67..d98db4ae7 100644 --- a/examples/rust/image_search_colpali/src/main.rs +++ b/examples/rust/image_search_colpali/src/main.rs @@ -19,7 +19,7 @@ //! where `D` is the per-vector dimension (128 for `vidore/colpali-v1.2`). A //! reference Python server is in the README. Everything else — the incremental //! pipeline and the Qdrant MAX_SIM multi-vector collection — is native Rust via -//! `cocoindex::qdrant`. +//! `cocoindex::connectors::qdrant`. //! //! Build note: `qdrant-client` compiles protobufs, so `protoc` is required. @@ -27,7 +27,7 @@ use std::path::PathBuf; use std::sync::LazyLock; use cocoindex::prelude::*; -use cocoindex::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; +use cocoindex::connectors::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; use serde::Deserialize; use serde_json::json; diff --git a/examples/rust/kafka_consume/README.md b/examples/rust/kafka_consume/README.md index 10a4e1f9a..708845398 100644 --- a/examples/rust/kafka_consume/README.md +++ b/examples/rust/kafka_consume/README.md @@ -1,7 +1,7 @@ # Kafka consume (Rust) A self-contained example of the Rust SDK's native Kafka **source** -(`cocoindex::kafka::topic_as_map` + `Ctx::mount_each_live`) — the analogue of +(`cocoindex::connectors::kafka::topic_as_map` + `Ctx::mount_each_live`) — the analogue of Python's `topic_as_map`. Pairs with the [`csv_to_kafka`](../csv_to_kafka) target example: produce with that, consume with this. diff --git a/examples/rust/kafka_consume/src/main.rs b/examples/rust/kafka_consume/src/main.rs index 2680dc74e..1c0165e77 100644 --- a/examples/rust/kafka_consume/src/main.rs +++ b/examples/rust/kafka_consume/src/main.rs @@ -13,7 +13,8 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; -use cocoindex::{App, UpdateOptions, kafka}; +use cocoindex::connectors::kafka; +use cocoindex::{App, UpdateOptions}; #[tokio::main] async fn main() -> cocoindex::Result<()> { diff --git a/examples/rust/meeting_notes_graph_falkordb/README.md b/examples/rust/meeting_notes_graph_falkordb/README.md index fe07e72f4..60cd7c893 100644 --- a/examples/rust/meeting_notes_graph_falkordb/README.md +++ b/examples/rust/meeting_notes_graph_falkordb/README.md @@ -12,7 +12,7 @@ It builds the same graph shape in FalkorDB: For local end-to-end testing this Rust example uses deterministic Markdown parsing over `input/*.md` instead of Google Drive + LLM extraction. The target -side is native CocoIndex Rust: `cocoindex::falkordb` table and relation targets +side is native CocoIndex Rust: `cocoindex::connectors::falkordb` table and relation targets with target-state reconciliation. ## Run diff --git a/examples/rust/meeting_notes_graph_falkordb/src/main.rs b/examples/rust/meeting_notes_graph_falkordb/src/main.rs index 37a4d7efd..6208691f9 100644 --- a/examples/rust/meeting_notes_graph_falkordb/src/main.rs +++ b/examples/rust/meeting_notes_graph_falkordb/src/main.rs @@ -2,8 +2,8 @@ use std::collections::{BTreeMap, BTreeSet}; use std::path::PathBuf; use chrono::NaiveDate; -use cocoindex::falkordb::{self, ColumnDef, TableSchema}; -use cocoindex::fs; +use cocoindex::connectors::falkordb::{self, ColumnDef, TableSchema}; +use cocoindex::resources::fs; use cocoindex::prelude::*; use serde::{Deserialize, Serialize}; diff --git a/examples/rust/meeting_notes_graph_neo4j/README.md b/examples/rust/meeting_notes_graph_neo4j/README.md index c0e941485..269eb7012 100644 --- a/examples/rust/meeting_notes_graph_neo4j/README.md +++ b/examples/rust/meeting_notes_graph_neo4j/README.md @@ -12,7 +12,7 @@ It builds the same graph shape in Neo4j: For local end-to-end testing this Rust example uses deterministic Markdown parsing over `input/*.md` instead of Google Drive + LLM extraction. The target -side is native CocoIndex Rust: `cocoindex::neo4j` table and relation targets +side is native CocoIndex Rust: `cocoindex::connectors::neo4j` table and relation targets with target-state reconciliation. ## Run diff --git a/examples/rust/meeting_notes_graph_neo4j/src/main.rs b/examples/rust/meeting_notes_graph_neo4j/src/main.rs index d32e0537d..f7addcfbc 100644 --- a/examples/rust/meeting_notes_graph_neo4j/src/main.rs +++ b/examples/rust/meeting_notes_graph_neo4j/src/main.rs @@ -2,8 +2,8 @@ use std::collections::{BTreeMap, BTreeSet}; use std::path::PathBuf; use chrono::NaiveDate; -use cocoindex::fs; -use cocoindex::neo4j::{self, ColumnDef, TableSchema}; +use cocoindex::resources::fs; +use cocoindex::connectors::neo4j::{self, ColumnDef, TableSchema}; use cocoindex::prelude::*; use serde::{Deserialize, Serialize}; diff --git a/examples/rust/multi_codebase_summarization/src/main.rs b/examples/rust/multi_codebase_summarization/src/main.rs index 881fba970..6c45ec615 100644 --- a/examples/rust/multi_codebase_summarization/src/main.rs +++ b/examples/rust/multi_codebase_summarization/src/main.rs @@ -377,7 +377,7 @@ async fn main() -> Result<()> { let project_dir = entry.path(); // Match both root-level and nested Python files. - let files = cocoindex::fs::walk(&project_dir, &["*.py", "**/*.py"])?; + let files = cocoindex::resources::fs::walk(&project_dir, &["*.py", "**/*.py"])?; let files: Vec<_> = files .into_iter() diff --git a/examples/rust/oci_object_storage_embedding/Cargo.toml b/examples/rust/oci_object_storage_embedding/Cargo.toml index a5ddbacb0..759ad66dd 100644 --- a/examples/rust/oci_object_storage_embedding/Cargo.toml +++ b/examples/rust/oci_object_storage_embedding/Cargo.toml @@ -18,7 +18,7 @@ cocoindex = { path = "../../../rust/sdk/cocoindex", features = [ ] } tokio = { version = "1", features = ["full"] } serde = { version = "1", features = ["derive"] } -# Query path (pgvector similarity); target writes go through cocoindex::postgres. +# Query path (pgvector similarity); target writes go through cocoindex::connectors::postgres. sqlx = { version = "0.8", default-features = false, features = [ "runtime-tokio", "tls-rustls", diff --git a/examples/rust/oci_object_storage_embedding/README.md b/examples/rust/oci_object_storage_embedding/README.md index 4952da314..2717811f3 100644 --- a/examples/rust/oci_object_storage_embedding/README.md +++ b/examples/rust/oci_object_storage_embedding/README.md @@ -11,7 +11,7 @@ Postgres/pgvector — then serves similarity search. | Concern | Python | Rust (this example) | | ---------------- | ----------------------------------------------- | --------------------------------------------------------- | -| Source | `oci_object_storage.list_objects` (`oci` SDK) | `cocoindex::oci_object_storage::list_objects` (native REST)| +| Source | `oci_object_storage.list_objects` (`oci` SDK) | `cocoindex::connectors::oci_object_storage::list_objects` (native REST)| | Auth | `oci.config.from_file` + `ObjectStorageClient` | `OciClient::connect` reads `~/.oci/config`, signs requests | | Per-file compute | `@coco.fn(memo=True) process_file` | `#[cocoindex::function(memo)] process_file` | | Chunking | `RecursiveSplitter` (markdown, 2000/500) | `RecursiveSplitter` (markdown, 2000/500) | diff --git a/examples/rust/oci_object_storage_embedding/src/main.rs b/examples/rust/oci_object_storage_embedding/src/main.rs index 67f343823..4f8c876c1 100644 --- a/examples/rust/oci_object_storage_embedding/src/main.rs +++ b/examples/rust/oci_object_storage_embedding/src/main.rs @@ -15,11 +15,11 @@ use std::path::PathBuf; use std::sync::{Arc, LazyLock}; -use cocoindex::file::PatternFilePathMatcher; -use cocoindex::oci_object_storage::{self, ListOptions, OciClient, OciFile}; +use cocoindex::resources::file::PatternFilePathMatcher; +use cocoindex::connectors::oci_object_storage::{self, ListOptions, OciClient, OciFile}; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; -use cocoindex::postgres; +use cocoindex::connectors::postgres; use cocoindex::prelude::*; use sqlx::Row; use sqlx::postgres::{PgPool, PgPoolOptions}; diff --git a/examples/rust/paper_metadata/Cargo.toml b/examples/rust/paper_metadata/Cargo.toml index e03ab11e8..91b84c255 100644 --- a/examples/rust/paper_metadata/Cargo.toml +++ b/examples/rust/paper_metadata/Cargo.toml @@ -18,7 +18,7 @@ serde_json = "1" reqwest = { version = "0.12", features = ["json"] } # PDF parsing: page count + first-page text extraction. lopdf = "0.40" -# Query path (pgvector similarity); target writes go through cocoindex::postgres. +# Query path (pgvector similarity); target writes go through cocoindex::connectors::postgres. sqlx = { version = "0.8", default-features = false, features = [ "runtime-tokio", "tls-rustls", diff --git a/examples/rust/paper_metadata/README.md b/examples/rust/paper_metadata/README.md index ef6db5efe..caa1c14eb 100644 --- a/examples/rust/paper_metadata/README.md +++ b/examples/rust/paper_metadata/README.md @@ -10,7 +10,7 @@ everything across three Postgres tables — then serves similarity search. | Concern | Python | Rust (this example) | | ---------------- | -------------------------------------------- | --------------------------------------------------------- | -| Source | `localfs.walk_dir` (`**/*.pdf`, live) | `cocoindex::fs::walk` (`**/*.pdf`) | +| Source | `localfs.walk_dir` (`**/*.pdf`, live) | `cocoindex::resources::fs::walk` (`**/*.pdf`) | | Per-file compute | `@coco.fn(memo=True) process_file` | `#[cocoindex::function(memo)] process_file` | | PDF parsing | `pypdf` (first-page text + page count) | `lopdf` (first-page text + page count) | | LLM extraction | `openai` chat completions (`gpt-4o`, JSON) | OpenAI chat completions REST (`gpt-4o`, JSON mode) | diff --git a/examples/rust/paper_metadata/src/main.rs b/examples/rust/paper_metadata/src/main.rs index c604e4b11..3d3b2c79e 100644 --- a/examples/rust/paper_metadata/src/main.rs +++ b/examples/rust/paper_metadata/src/main.rs @@ -8,7 +8,7 @@ //! cargo run -- query "your query" # pgvector similarity search (no index) //! //! Parallels the Python example: -//! - source : `cocoindex::fs::walk` (cf. `localfs.walk_dir`) +//! - source : `cocoindex::resources::fs::walk` (cf. `localfs.walk_dir`) //! - per-file compute : `#[cocoindex::function(memo)]` (cf. `@coco.fn(memo=True)`) //! - PDF parsing : `lopdf` (cf. `pypdf` first-page text + page count) //! - LLM extraction : OpenAI chat completions, JSON mode (cf. `openai` client) @@ -28,10 +28,10 @@ use std::path::PathBuf; use std::sync::LazyLock; -use cocoindex::id::UuidGenerator; +use cocoindex::resources::id::UuidGenerator; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{CustomLanguageConfig, RecursiveChunkConfig, RecursiveSplitter}; -use cocoindex::postgres; +use cocoindex::connectors::postgres; use cocoindex::prelude::*; use serde::de::DeserializeOwned; use sqlx::Row; diff --git a/examples/rust/pdf_embedding/Cargo.toml b/examples/rust/pdf_embedding/Cargo.toml index f4d030d5e..f8ad0f170 100644 --- a/examples/rust/pdf_embedding/Cargo.toml +++ b/examples/rust/pdf_embedding/Cargo.toml @@ -18,7 +18,7 @@ tokio = { version = "1", features = ["full"] } serde = { version = "1", features = ["derive"] } # PDF text extraction (Rust-native stand-in for Python's docling converter). lopdf = "0.40" -# Query path (pgvector similarity); target writes go through cocoindex::postgres. +# Query path (pgvector similarity); target writes go through cocoindex::connectors::postgres. sqlx = { version = "0.8", default-features = false, features = [ "runtime-tokio", "tls-rustls", diff --git a/examples/rust/pdf_embedding/README.md b/examples/rust/pdf_embedding/README.md index 5b7f032e3..6da420f09 100644 --- a/examples/rust/pdf_embedding/README.md +++ b/examples/rust/pdf_embedding/README.md @@ -9,7 +9,7 @@ them in Postgres/pgvector — then serves similarity search. | Concern | Python | Rust (this example) | | ---------------- | ---------------------------------------- | -------------------------------------------------- | -| Source | `localfs.walk_dir` (`**/*.pdf`) | `cocoindex::fs::walk` (`**/*.pdf`) | +| Source | `localfs.walk_dir` (`**/*.pdf`) | `cocoindex::resources::fs::walk` (`**/*.pdf`) | | PDF → text | `docling` (PDF → Markdown, ML pipeline) | `lopdf` text extraction | | Per-file compute | `@coco.fn(memo=True) process_file` | `#[cocoindex::function(memo)] process_file` | | Chunking | `RecursiveSplitter` (markdown, 2000/500) | `cocoindex_ops_text` `RecursiveChunker` (markdown, 2000/500) | diff --git a/examples/rust/pdf_embedding/src/main.rs b/examples/rust/pdf_embedding/src/main.rs index 03b6165d5..58772d172 100644 --- a/examples/rust/pdf_embedding/src/main.rs +++ b/examples/rust/pdf_embedding/src/main.rs @@ -18,7 +18,7 @@ use std::sync::LazyLock; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; -use cocoindex::postgres; +use cocoindex::connectors::postgres; use cocoindex::prelude::*; use sqlx::Row; use sqlx::postgres::{PgPool, PgPoolOptions}; diff --git a/examples/rust/pdf_to_markdown/README.md b/examples/rust/pdf_to_markdown/README.md index 03df1a485..9e15a0b9c 100644 --- a/examples/rust/pdf_to_markdown/README.md +++ b/examples/rust/pdf_to_markdown/README.md @@ -11,7 +11,7 @@ automatically. | Concern | Python | Rust (this example) | | ---------------- | ---------------------------------------- | ----------------------------------------- | -| Source | `localfs.walk_dir` (`**/*.pdf`) | `cocoindex::fs::walk` (`**/*.pdf`) | +| Source | `localfs.walk_dir` (`**/*.pdf`) | `cocoindex::resources::fs::walk` (`**/*.pdf`) | | PDF → markdown | `docling` (PDF → Markdown, ML pipeline) | `lopdf` text extraction | | Per-file compute | `@coco.fn(memo=True) process_file` | `#[cocoindex::function(memo)] convert_pdf` | | Output | `localfs.declare_file` (`.md`) | `DirTarget::declare_file` (`.md`) | diff --git a/examples/rust/pdf_to_markdown/src/main.rs b/examples/rust/pdf_to_markdown/src/main.rs index f1dde28b5..6f6e2d520 100644 --- a/examples/rust/pdf_to_markdown/src/main.rs +++ b/examples/rust/pdf_to_markdown/src/main.rs @@ -75,7 +75,7 @@ async fn main() -> Result<()> { let output_dir = output_dir.clone(); async move { let target = DirTarget::mount(&ctx, &output_dir)?; - let files = cocoindex::fs::walk_items(&source_dir, &["**/*.pdf"])?; + let files = cocoindex::resources::fs::walk_items(&source_dir, &["**/*.pdf"])?; println!( "converting {} PDF(s) from {}", files.len(), diff --git a/examples/rust/postgres_source/Cargo.toml b/examples/rust/postgres_source/Cargo.toml index 5a1f978e5..df8bf1165 100644 --- a/examples/rust/postgres_source/Cargo.toml +++ b/examples/rust/postgres_source/Cargo.toml @@ -9,7 +9,7 @@ edition = "2024" cocoindex = { path = "../../../rust/sdk/cocoindex", features = ["postgres", "fastembed"] } tokio = { version = "1", features = ["full"] } serde = { version = "1", features = ["derive"] } -# Query path (pgvector similarity); target writes go through cocoindex::postgres. +# Query path (pgvector similarity); target writes go through cocoindex::connectors::postgres. sqlx = { version = "0.8", default-features = false, features = [ "runtime-tokio", "tls-rustls", diff --git a/examples/rust/postgres_source/src/main.rs b/examples/rust/postgres_source/src/main.rs index cf7b4c04f..d23db9b70 100644 --- a/examples/rust/postgres_source/src/main.rs +++ b/examples/rust/postgres_source/src/main.rs @@ -20,7 +20,7 @@ use std::sync::LazyLock; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; -use cocoindex::postgres; +use cocoindex::connectors::postgres; use cocoindex::prelude::*; use sqlx::Row; use sqlx::postgres::{PgPool, PgPoolOptions}; diff --git a/examples/rust/text_embedding/Cargo.toml b/examples/rust/text_embedding/Cargo.toml index 72fd227a7..1269f4aea 100644 --- a/examples/rust/text_embedding/Cargo.toml +++ b/examples/rust/text_embedding/Cargo.toml @@ -16,7 +16,7 @@ cocoindex = { path = "../../../rust/sdk/cocoindex", features = [ ] } tokio = { version = "1", features = ["full"] } serde = { version = "1", features = ["derive"] } -# Query path (pgvector similarity); target writes go through cocoindex::postgres. +# Query path (pgvector similarity); target writes go through cocoindex::connectors::postgres. sqlx = { version = "0.8", default-features = false, features = [ "runtime-tokio", "tls-rustls", diff --git a/examples/rust/text_embedding/README.md b/examples/rust/text_embedding/README.md index 438f14b2c..ad8ac72bc 100644 --- a/examples/rust/text_embedding/README.md +++ b/examples/rust/text_embedding/README.md @@ -9,7 +9,7 @@ and stores them in Postgres/pgvector — then serves similarity search. | Concern | Python | Rust (this example) | | ---------------- | ---------------------------------------- | -------------------------------------------------- | -| Source | `localfs.walk_dir` | `cocoindex::fs::walk` | +| Source | `localfs.walk_dir` | `cocoindex::resources::fs::walk` | | Per-file compute | `@coco.fn(memo=True) process_file` | `#[cocoindex::function(memo)] process_file` | | Chunking | `RecursiveSplitter` (markdown) | `cocoindex_ops_text` `RecursiveChunker` (markdown) | | Embeddings | `sentence-transformers/all-MiniLM-L6-v2` | `fastembed` `AllMiniLML6V2` (same model, 384-dim) | diff --git a/examples/rust/text_embedding/src/main.rs b/examples/rust/text_embedding/src/main.rs index 3bf7d9a73..0a5f19611 100644 --- a/examples/rust/text_embedding/src/main.rs +++ b/examples/rust/text_embedding/src/main.rs @@ -6,7 +6,7 @@ //! cargo run -- query "your query" # pgvector similarity search //! //! Parallels the Python example: -//! - source : `cocoindex::fs::walk` (cf. `localfs.walk_dir`) +//! - source : `cocoindex::resources::fs::walk` (cf. `localfs.walk_dir`) //! - per-file compute : `#[cocoindex::function(memo)]` (cf. `@coco.fn(memo=True)`) //! - chunking : `cocoindex::ops::text::RecursiveSplitter` (cf. `RecursiveSplitter`) //! - embeddings : `cocoindex::ops::sentence_transformers` all-MiniLM-L6-v2 @@ -17,7 +17,7 @@ use std::sync::LazyLock; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; -use cocoindex::postgres; +use cocoindex::connectors::postgres; use cocoindex::prelude::*; use sqlx::Row; use sqlx::postgres::{PgPool, PgPoolOptions}; diff --git a/examples/rust/text_embedding_lancedb/README.md b/examples/rust/text_embedding_lancedb/README.md index a7fcae114..13ffd8039 100644 --- a/examples/rust/text_embedding_lancedb/README.md +++ b/examples/rust/text_embedding_lancedb/README.md @@ -3,20 +3,20 @@ Rust port of the Python [`text_embedding_lancedb`](../../text_embedding_lancedb) example. Same pipeline as [`text_embedding`](../text_embedding), but the vector store is **LanceDB** (an embedded, file-based vector database) via the native -`cocoindex::lancedb` connector instead of Postgres/pgvector. +`cocoindex::connectors::lancedb` connector instead of Postgres/pgvector. ## Parallel to the Python example | Concern | Python | Rust (this example) | | ---------------- | ---------------------------------------- | -------------------------------------------------- | -| Source | `localfs.walk_dir` | `cocoindex::fs::walk` | +| Source | `localfs.walk_dir` | `cocoindex::resources::fs::walk` | | Per-file compute | `@coco.fn(memo=True) process_file` | `#[cocoindex::function(memo)] process_file` | | Chunking | `RecursiveSplitter` (markdown) | `cocoindex_ops_text` `RecursiveChunker` (markdown) | | Embeddings | `sentence-transformers/all-MiniLM-L6-v2` | `fastembed` `AllMiniLML6V2` (same model, 384-dim) | -| Target | `lancedb.TableTarget` | `cocoindex::lancedb::LanceTableTarget` | -| Vector search | `table.search(vec)` | `cocoindex::lancedb::vector_search` (cosine) | +| Target | `lancedb.TableTarget` | `cocoindex::connectors::lancedb::LanceTableTarget` | +| Vector search | `table.search(vec)` | `cocoindex::connectors::lancedb::vector_search` (cosine) | -The `cocoindex::lancedb` connector is a declarative two-level managed target +The `cocoindex::connectors::lancedb` connector is a declarative two-level managed target (table → rows), mirroring `postgres`: it creates the table to match the schema, upserts changed rows, skips unchanged ones (fingerprint tracking), and deletes rows that are no longer declared. It's built on the native Rust `lancedb` crate + diff --git a/examples/rust/text_embedding_lancedb/src/main.rs b/examples/rust/text_embedding_lancedb/src/main.rs index da4adc3a3..0c784799f 100644 --- a/examples/rust/text_embedding_lancedb/src/main.rs +++ b/examples/rust/text_embedding_lancedb/src/main.rs @@ -7,13 +7,13 @@ //! cargo run -- query "your query" # LanceDB vector search //! //! Same pipeline as the `text_embedding` example, but the target is the native -//! `cocoindex::lancedb` connector instead of Postgres/pgvector. Parallels the +//! `cocoindex::connectors::lancedb` connector instead of Postgres/pgvector. Parallels the //! Python example's use of `cocoindex.connectors.lancedb`. use std::path::PathBuf; use std::sync::LazyLock; -use cocoindex::lancedb::{self, ColumnDef, ColumnType, LanceDatabase, TableSchema}; +use cocoindex::connectors::lancedb::{self, ColumnDef, ColumnType, LanceDatabase, TableSchema}; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; use cocoindex::prelude::*; diff --git a/examples/rust/text_embedding_qdrant/README.md b/examples/rust/text_embedding_qdrant/README.md index e97ce63a1..badd48d9e 100644 --- a/examples/rust/text_embedding_qdrant/README.md +++ b/examples/rust/text_embedding_qdrant/README.md @@ -3,19 +3,19 @@ Rust port of the Python [`text_embedding_qdrant`](../../text_embedding_qdrant) example. Same pipeline as [`text_embedding`](../text_embedding), but the vector store is -**Qdrant** via the native `cocoindex::qdrant` collection target. +**Qdrant** via the native `cocoindex::connectors::qdrant` collection target. ## Parallel to the Python example | Concern | Python | Rust (this example) | | ---------------- | ---------------------------------------- | -------------------------------------------------- | -| Source | `localfs.walk_dir` | `cocoindex::fs::walk` | +| Source | `localfs.walk_dir` | `cocoindex::resources::fs::walk` | | Chunking | `RecursiveSplitter` (markdown) | `cocoindex_ops_text` `RecursiveChunker` (markdown) | | Embeddings | `sentence-transformers/all-MiniLM-L6-v2` | `fastembed` `AllMiniLML6V2` (same model, 384-dim) | -| Target | `qdrant.CollectionTarget` | `cocoindex::qdrant::CollectionTarget` | -| Search | `client.query_points(...)` | `cocoindex::qdrant::vector_search` (cosine score) | +| Target | `qdrant.CollectionTarget` | `cocoindex::connectors::qdrant::CollectionTarget` | +| Search | `client.query_points(...)` | `cocoindex::connectors::qdrant::vector_search` (cosine score) | -The `cocoindex::qdrant` connector is a declarative two-level **managed target** +The `cocoindex::connectors::qdrant` connector is a declarative two-level **managed target** (collection → points) built on CocoIndex's public target-state facade: it creates the collection to match the vector schema, upserts changed points, skips unchanged ones (fingerprint tracking), deletes orphaned points, and recreates diff --git a/examples/rust/text_embedding_qdrant/src/main.rs b/examples/rust/text_embedding_qdrant/src/main.rs index 530f5dba8..97c05e935 100644 --- a/examples/rust/text_embedding_qdrant/src/main.rs +++ b/examples/rust/text_embedding_qdrant/src/main.rs @@ -7,7 +7,7 @@ //! cargo run -- query "your query" # Qdrant vector search //! //! Same pipeline as `text_embedding`, but the target is the native -//! `cocoindex::qdrant` collection connector (built on the public target-state +//! `cocoindex::connectors::qdrant` collection connector (built on the public target-state //! facade). Parallels the Python example's `cocoindex.connectors.qdrant`. //! //! Build note: the `qdrant-client` crate compiles protobufs, so a `protoc` @@ -19,7 +19,7 @@ use std::sync::LazyLock; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; use cocoindex::prelude::*; -use cocoindex::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; +use cocoindex::connectors::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; use serde_json::json; const EMBED_MODEL: &str = "sentence-transformers/all-MiniLM-L6-v2"; diff --git a/examples/rust/text_embedding_turbopuffer/README.md b/examples/rust/text_embedding_turbopuffer/README.md index c593cdb99..91d568680 100644 --- a/examples/rust/text_embedding_turbopuffer/README.md +++ b/examples/rust/text_embedding_turbopuffer/README.md @@ -3,19 +3,19 @@ Rust port of the Python [`text_embedding_turbopuffer`](../../text_embedding_turbopuffer) example. Same pipeline as [`text_embedding`](../text_embedding), but the vector store is **Turbopuffer** (a hosted vector database) via the native -`cocoindex::turbopuffer` namespace target. +`cocoindex::connectors::turbopuffer` namespace target. ## Parallel to the Python example | Concern | Python | Rust (this example) | | ---------------- | ---------------------------------------- | -------------------------------------------------- | -| Source | `localfs.walk_dir` | `cocoindex::fs::walk` | +| Source | `localfs.walk_dir` | `cocoindex::resources::fs::walk` | | Chunking | `RecursiveSplitter` (markdown) | `cocoindex_ops_text` `RecursiveChunker` (markdown) | | Embeddings | `sentence-transformers/all-MiniLM-L6-v2` | `fastembed` `AllMiniLML6V2` (same model, 384-dim) | -| Target | `turbopuffer.NamespaceTarget` | `cocoindex::turbopuffer::NamespaceTarget` | -| Search | `ns.query(rank_by=("vector","ANN",...))` | `cocoindex::turbopuffer::vector_search` | +| Target | `turbopuffer.NamespaceTarget` | `cocoindex::connectors::turbopuffer::NamespaceTarget` | +| Search | `ns.query(rank_by=("vector","ANN",...))` | `cocoindex::connectors::turbopuffer::vector_search` | -The `cocoindex::turbopuffer` connector is a declarative two-level **managed +The `cocoindex::connectors::turbopuffer` connector is a declarative two-level **managed target** (namespace → rows) built on CocoIndex's public target-state facade: it upserts changed rows, skips unchanged ones (fingerprint tracking), deletes orphaned rows, and clears the namespace if the vector schema changes. Turbopuffer diff --git a/examples/rust/text_embedding_turbopuffer/src/main.rs b/examples/rust/text_embedding_turbopuffer/src/main.rs index 70ee0ebb2..8751595f2 100644 --- a/examples/rust/text_embedding_turbopuffer/src/main.rs +++ b/examples/rust/text_embedding_turbopuffer/src/main.rs @@ -8,7 +8,7 @@ //! cargo run -- query "your query" # Turbopuffer vector search //! //! Same pipeline as `text_embedding`, but the target is the native -//! `cocoindex::turbopuffer` namespace connector (built on the public +//! `cocoindex::connectors::turbopuffer` namespace connector (built on the public //! target-state facade). Parallels `cocoindex.connectors.turbopuffer`. //! //! Config (env, e.g. via `.env`): @@ -22,7 +22,7 @@ use std::sync::LazyLock; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; use cocoindex::prelude::*; -use cocoindex::turbopuffer::{self, DistanceMetric, NamespaceSchema, TurbopufferConnection}; +use cocoindex::connectors::turbopuffer::{self, DistanceMetric, NamespaceSchema, TurbopufferConnection}; use serde_json::json; const EMBED_MODEL: &str = "sentence-transformers/all-MiniLM-L6-v2"; diff --git a/rust/core/src/engine/progress_display.rs b/rust/core/src/engine/progress_display.rs index 09d21d607..9be877a14 100644 --- a/rust/core/src/engine/progress_display.rs +++ b/rust/core/src/engine/progress_display.rs @@ -427,7 +427,7 @@ async fn show_progress_pty( tokio::select! { result = handle.changed() => { let version = result?; - if version >= TERMINATED_VERSION { + if version == TERMINATED_VERSION { break; } } diff --git a/rust/sdk/SHOWCASE.md b/rust/sdk/SHOWCASE.md index 8f9f361cc..a3e47e58e 100644 --- a/rust/sdk/SHOWCASE.md +++ b/rust/sdk/SHOWCASE.md @@ -77,7 +77,7 @@ impl FileEntry { ``` ```rust -let files = cocoindex::fs::walk(&dir, &["**/*.rs", "**/*.py"])?; +let files = cocoindex::resources::fs::walk(&dir, &["**/*.rs", "**/*.py"])?; ``` ### `RunStats` @@ -281,7 +281,8 @@ async fn main() -> cocoindex::Result<()> { for entry in std::fs::read_dir(&root_dir)? { let project_name = entry?.file_name().to_string_lossy().to_string(); let project_dir = entry?.path(); - let files = cocoindex::fs::walk(&project_dir, &["*.py", "**/*.py"])?; + let files = + cocoindex::resources::fs::walk(&project_dir, &["*.py", "**/*.py"])?; let file_infos = ctx .mount_each( diff --git a/rust/sdk/cocoindex/benches/sdk_microbench.rs b/rust/sdk/cocoindex/benches/sdk_microbench.rs index c47f8449a..35611dabf 100644 --- a/rust/sdk/cocoindex/benches/sdk_microbench.rs +++ b/rust/sdk/cocoindex/benches/sdk_microbench.rs @@ -5,7 +5,7 @@ use std::time::{Duration, SystemTime}; use cocoindex::memo::{ finish_key_fingerprinter, key_bytes_result, new_key_fingerprinter, write_key_fingerprint_part, }; -use cocoindex::walk; +use cocoindex::resources::fs::walk; use cocoindex_utils::fingerprint::Fingerprint; use criterion::{BenchmarkId, Criterion, SamplingMode, black_box, criterion_group, criterion_main}; use rustc_hash::FxHashSet; diff --git a/rust/sdk/cocoindex/src/amazon_s3.rs b/rust/sdk/cocoindex/src/connectors/amazon_s3.rs similarity index 99% rename from rust/sdk/cocoindex/src/amazon_s3.rs rename to rust/sdk/cocoindex/src/connectors/amazon_s3.rs index 97bfbc14f..045cca161 100644 --- a/rust/sdk/cocoindex/src/amazon_s3.rs +++ b/rust/sdk/cocoindex/src/connectors/amazon_s3.rs @@ -15,7 +15,7 @@ //! //! Like the Google Drive source, [`S3File`] serializes only stable metadata for //! memo keys. The clone-cheap client and read cache are skipped by serde, while -//! the public type still implements the shared async [`crate::file::FileLike`] +//! the public type still implements the shared async [`crate::resources::file::FileLike`] //! trait. use std::path::PathBuf; @@ -34,7 +34,7 @@ use serde::{Deserialize, Serialize}; pub use aws_sdk_s3; use crate::error::{Error, Result}; -use crate::file::{ +use crate::resources::file::{ FileContentCache, FileLike, FileMetadata, FilePath, FilePathMatcher, FileSourceItem, MatchAllFilePathMatcher, decode_bytes, }; @@ -523,8 +523,8 @@ fn relative_key(prefix: &str, key: &str) -> Option { #[cfg(test)] mod tests { use super::*; - use crate::file::FileLike; - use crate::fs::PatternFilePathMatcher; + use crate::resources::file::FileLike; + use crate::resources::fs::PatternFilePathMatcher; #[test] fn parse_uri_ok_and_errors() { diff --git a/rust/sdk/cocoindex/src/doris.rs b/rust/sdk/cocoindex/src/connectors/doris.rs similarity index 100% rename from rust/sdk/cocoindex/src/doris.rs rename to rust/sdk/cocoindex/src/connectors/doris.rs diff --git a/rust/sdk/cocoindex/src/falkordb.rs b/rust/sdk/cocoindex/src/connectors/falkordb.rs similarity index 100% rename from rust/sdk/cocoindex/src/falkordb.rs rename to rust/sdk/cocoindex/src/connectors/falkordb.rs diff --git a/rust/sdk/cocoindex/src/gdrive.rs b/rust/sdk/cocoindex/src/connectors/gdrive.rs similarity index 99% rename from rust/sdk/cocoindex/src/gdrive.rs rename to rust/sdk/cocoindex/src/connectors/gdrive.rs index b9dbd89e8..6b4832c9d 100644 --- a/rust/sdk/cocoindex/src/gdrive.rs +++ b/rust/sdk/cocoindex/src/connectors/gdrive.rs @@ -26,7 +26,7 @@ use serde::{Deserialize, Serialize}; use sha2::Sha256; use crate::error::{Error, Result}; -use crate::file::{ +use crate::resources::file::{ FileContentCache, FileLike, FileMetadata, FilePath, FileSourceItem, decode_bytes, }; diff --git a/rust/sdk/cocoindex/src/iggy.rs b/rust/sdk/cocoindex/src/connectors/iggy.rs similarity index 99% rename from rust/sdk/cocoindex/src/iggy.rs rename to rust/sdk/cocoindex/src/connectors/iggy.rs index 7b1b0af7b..a34290087 100644 --- a/rust/sdk/cocoindex/src/iggy.rs +++ b/rust/sdk/cocoindex/src/connectors/iggy.rs @@ -33,7 +33,7 @@ use serde::{Deserialize, Serialize}; /// Re-export of the upstream [`iggy`] crate prelude. Streams and topics are /// user-managed, so callers use this to create/manage them and to poll messages /// back — without having to depend on the `iggy` crate directly. e.g. -/// `cocoindex::iggy::prelude::{StreamClient, TopicClient}`. +/// `cocoindex::connectors::iggy::prelude::{StreamClient, TopicClient}`. pub use ::iggy::prelude; use crate::ctx::Ctx; diff --git a/rust/sdk/cocoindex/src/kafka.rs b/rust/sdk/cocoindex/src/connectors/kafka.rs similarity index 100% rename from rust/sdk/cocoindex/src/kafka.rs rename to rust/sdk/cocoindex/src/connectors/kafka.rs diff --git a/rust/sdk/cocoindex/src/lancedb.rs b/rust/sdk/cocoindex/src/connectors/lancedb.rs similarity index 100% rename from rust/sdk/cocoindex/src/lancedb.rs rename to rust/sdk/cocoindex/src/connectors/lancedb.rs diff --git a/rust/sdk/cocoindex/src/connectors/mod.rs b/rust/sdk/cocoindex/src/connectors/mod.rs new file mode 100644 index 000000000..469bd7846 --- /dev/null +++ b/rust/sdk/cocoindex/src/connectors/mod.rs @@ -0,0 +1,32 @@ +//! External-system connectors for CocoIndex pipelines. + +#[cfg(feature = "amazon_s3")] +pub mod amazon_s3; +#[cfg(feature = "doris")] +pub mod doris; +#[cfg(feature = "falkordb")] +pub mod falkordb; +#[cfg(feature = "google_drive")] +pub mod gdrive; +#[cfg(feature = "iggy")] +pub mod iggy; +#[cfg(feature = "kafka")] +pub mod kafka; +#[cfg(feature = "lancedb")] +pub mod lancedb; +#[cfg(feature = "neo4j")] +pub mod neo4j; +#[cfg(feature = "oci_object_storage")] +pub mod oci_object_storage; +#[cfg(feature = "postgres")] +pub mod postgres; +#[cfg(feature = "qdrant")] +pub mod qdrant; +#[cfg(feature = "sqlite")] +pub mod sqlite; +#[cfg(feature = "surrealdb")] +pub mod surrealdb; +#[cfg(feature = "turbopuffer")] +pub mod turbopuffer; +#[cfg(feature = "valkey")] +pub mod valkey; diff --git a/rust/sdk/cocoindex/src/neo4j.rs b/rust/sdk/cocoindex/src/connectors/neo4j.rs similarity index 100% rename from rust/sdk/cocoindex/src/neo4j.rs rename to rust/sdk/cocoindex/src/connectors/neo4j.rs diff --git a/rust/sdk/cocoindex/src/oci_object_storage.rs b/rust/sdk/cocoindex/src/connectors/oci_object_storage.rs similarity index 99% rename from rust/sdk/cocoindex/src/oci_object_storage.rs rename to rust/sdk/cocoindex/src/connectors/oci_object_storage.rs index 3ca92a812..1e1fb7968 100644 --- a/rust/sdk/cocoindex/src/oci_object_storage.rs +++ b/rust/sdk/cocoindex/src/connectors/oci_object_storage.rs @@ -24,7 +24,7 @@ //! Like the Amazon S3 and Google Drive sources, [`OciFile`] serializes only //! stable metadata for memo keys. The clone-cheap client and read cache are //! skipped by serde, while the public type still implements the shared async -//! [`crate::file::FileLike`] trait. +//! [`crate::resources::file::FileLike`] trait. //! //! Not yet supported (parity follow-ups): live bucket watching via OCI //! Events/Streaming, and pass-phrase-encrypted private keys. @@ -46,13 +46,13 @@ use serde::{Deserialize, Serialize}; use sha2::Sha256; use crate::error::{Error, Result}; -use crate::file::{ - FileContentCache, FileLike, FileMetadata, FilePath, FilePathMatcher, FileSourceItem, - MatchAllFilePathMatcher, decode_bytes, -}; use crate::live_component::{ LiveComponentOperator, LiveMapFeed, LiveMapSubscriber, LiveMapView, SingleWatcherGuard, }; +use crate::resources::file::{ + FileContentCache, FileLike, FileMetadata, FilePath, FilePathMatcher, FileSourceItem, + MatchAllFilePathMatcher, decode_bytes, +}; /// Object metadata fields requested from the ListObjects API. const LIST_FIELDS: &str = "name,size,md5,timeModified,etag"; @@ -1234,7 +1234,7 @@ impl LiveMapFeed for OciLiveWalker { #[cfg(test)] mod tests { use super::*; - use crate::file::{FileLike, PatternFilePathMatcher}; + use crate::resources::file::{FileLike, PatternFilePathMatcher}; use rsa::pkcs1v15::{Signature, VerifyingKey}; use rsa::signature::Verifier; use rsa::{RsaPrivateKey, RsaPublicKey}; diff --git a/rust/sdk/cocoindex/src/postgres.rs b/rust/sdk/cocoindex/src/connectors/postgres.rs similarity index 100% rename from rust/sdk/cocoindex/src/postgres.rs rename to rust/sdk/cocoindex/src/connectors/postgres.rs diff --git a/rust/sdk/cocoindex/src/qdrant.rs b/rust/sdk/cocoindex/src/connectors/qdrant.rs similarity index 100% rename from rust/sdk/cocoindex/src/qdrant.rs rename to rust/sdk/cocoindex/src/connectors/qdrant.rs diff --git a/rust/sdk/cocoindex/src/sqlite.rs b/rust/sdk/cocoindex/src/connectors/sqlite.rs similarity index 100% rename from rust/sdk/cocoindex/src/sqlite.rs rename to rust/sdk/cocoindex/src/connectors/sqlite.rs diff --git a/rust/sdk/cocoindex/src/surrealdb.rs b/rust/sdk/cocoindex/src/connectors/surrealdb.rs similarity index 100% rename from rust/sdk/cocoindex/src/surrealdb.rs rename to rust/sdk/cocoindex/src/connectors/surrealdb.rs diff --git a/rust/sdk/cocoindex/src/turbopuffer.rs b/rust/sdk/cocoindex/src/connectors/turbopuffer.rs similarity index 100% rename from rust/sdk/cocoindex/src/turbopuffer.rs rename to rust/sdk/cocoindex/src/connectors/turbopuffer.rs diff --git a/rust/sdk/cocoindex/src/valkey.rs b/rust/sdk/cocoindex/src/connectors/valkey.rs similarity index 100% rename from rust/sdk/cocoindex/src/valkey.rs rename to rust/sdk/cocoindex/src/connectors/valkey.rs diff --git a/rust/sdk/cocoindex/src/lib.rs b/rust/sdk/cocoindex/src/lib.rs index 1bc942edb..049756051 100644 --- a/rust/sdk/cocoindex/src/lib.rs +++ b/rust/sdk/cocoindex/src/lib.rs @@ -1,17 +1,11 @@ -#[cfg(feature = "amazon_s3")] -pub mod amazon_s3; pub mod app; pub mod batched; +pub mod connectors; pub mod ctx; #[cfg(any(feature = "neo4j", feature = "falkordb"))] mod cypher_graph; -#[cfg(feature = "doris")] -pub mod doris; pub mod entity_resolution; pub mod error; -#[cfg(feature = "falkordb")] -pub mod falkordb; -pub mod file; // Rejects non-finite floats before the JSON round-trip in target connectors that // serialize rows through `serde_json` (which maps NaN/±Inf to null). #[cfg(any( @@ -23,49 +17,23 @@ pub mod file; feature = "falkordb" ))] mod finite; -pub mod fs; -#[cfg(feature = "google_drive")] -pub mod gdrive; -pub mod id; -#[cfg(feature = "iggy")] -pub mod iggy; -#[cfg(feature = "kafka")] -pub mod kafka; -#[cfg(feature = "lancedb")] -pub mod lancedb; pub mod live_component; #[doc(hidden)] pub mod logic; pub mod memo; pub mod mount; -#[cfg(feature = "neo4j")] -pub mod neo4j; -#[cfg(feature = "oci_object_storage")] -pub mod oci_object_storage; pub mod ops; -#[cfg(feature = "postgres")] -pub mod postgres; pub mod prelude; pub(crate) mod profile; -#[cfg(feature = "qdrant")] -pub mod qdrant; pub mod resources; pub mod row_schema; #[cfg(any(feature = "doris", feature = "sqlite", feature = "surrealdb"))] pub(crate) mod sql_ident; -#[cfg(feature = "sqlite")] -pub mod sqlite; pub mod statediff; mod stats; -#[cfg(feature = "surrealdb")] -pub mod surrealdb; pub mod target_state; -#[cfg(feature = "turbopuffer")] -pub mod turbopuffer; mod typemap; pub mod user_state; -#[cfg(feature = "valkey")] -pub mod valkey; // Flat re-exports — the public API surface pub use app::{ @@ -80,18 +48,6 @@ pub use entity_resolution::{ resolve_entities_with_events, }; pub use error::{Error, Result}; -pub use file::{ - FileContentCache, FileLike, FileMetadata, FilePath, FilePathMatcher, FileSourceItem, - MatchAllFilePathMatcher, PatternFilePathMatcher, -}; -pub use fs::{ - DirTarget, DirTargetState, DirWalker, FileEntry, declare_dir_target, dir_target, - mount_dir_target, walk, walk_dir, walk_items, -}; -pub use id::{ - IdGenerator, UuidGenerator, generate_id, generate_id_default, generate_uuid, - generate_uuid_default, -}; // Re-exported so `#[cocoindex::function]` output can register each function's // logic fingerprint without the user crate needing a direct `linkme` dependency. #[doc(hidden)] diff --git a/rust/sdk/cocoindex/src/memo.rs b/rust/sdk/cocoindex/src/memo.rs index 77397dcf2..6a61c8cc1 100644 --- a/rust/sdk/cocoindex/src/memo.rs +++ b/rust/sdk/cocoindex/src/memo.rs @@ -13,8 +13,8 @@ use serde::{Deserialize, Serialize}; use crate::ctx::{Ctx, fn_call_guard}; use crate::error::{Error, Result}; -use crate::file::FileLike; use crate::profile::Value; +use crate::resources::file::FileLike; #[derive(Clone)] pub struct MemoStateValue(Value); @@ -262,19 +262,19 @@ pub async fn collect_memo_arg_state( } fn as_file_like(value: &dyn Any) -> Option<&dyn FileLike> { - if let Some(file) = value.downcast_ref::() { + if let Some(file) = value.downcast_ref::() { return Some(file); } #[cfg(feature = "amazon_s3")] - if let Some(file) = value.downcast_ref::() { + if let Some(file) = value.downcast_ref::() { return Some(file); } #[cfg(feature = "google_drive")] - if let Some(file) = value.downcast_ref::() { + if let Some(file) = value.downcast_ref::() { return Some(file); } #[cfg(feature = "oci_object_storage")] - if let Some(file) = value.downcast_ref::() { + if let Some(file) = value.downcast_ref::() { return Some(file); } None diff --git a/rust/sdk/cocoindex/src/ops/api.rs b/rust/sdk/cocoindex/src/ops/api.rs index 4eb3c8811..8417ee50e 100644 --- a/rust/sdk/cocoindex/src/ops/api.rs +++ b/rust/sdk/cocoindex/src/ops/api.rs @@ -21,7 +21,7 @@ use serde_json::json; use tokio::sync::Mutex; use crate::error::{Error, Result}; -use crate::file::FileLike; +use crate::resources::file::FileLike; use crate::resources::schema::{VectorElementType, VectorSchema, VectorSchemaProvider}; const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1"; diff --git a/rust/sdk/cocoindex/src/prelude.rs b/rust/sdk/cocoindex/src/prelude.rs index 9f7a3bce4..029ff9efa 100644 --- a/rust/sdk/cocoindex/src/prelude.rs +++ b/rust/sdk/cocoindex/src/prelude.rs @@ -7,24 +7,24 @@ pub use crate::entity_resolution::{ resolve_entities_with_events, }; pub use crate::error::{Error, Result}; -pub use crate::file::{ +pub use crate::live_component::{ + ExceptionContext, ExceptionHandler, LiveComponent, LiveComponentOperator, LiveMapFeed, + LiveMapSubscriber, LiveMapView, MountKind, SingleWatcherGuard, SingleWatcherToken, +}; +pub use crate::resources::chunk::{Chunk, TextPosition}; +pub use crate::resources::embedder::Embedder; +pub use crate::resources::file::{ FileContentCache, FileLike, FileMetadata, FilePath, FilePathMatcher, FileSourceItem, MatchAllFilePathMatcher, PatternFilePathMatcher, }; -pub use crate::fs::{ +pub use crate::resources::fs::{ DirTarget, DirTargetState, DirWalker, FileEntry, declare_dir_target, dir_target, mount_dir_target, walk, walk_dir, walk_items, }; -pub use crate::id::{ +pub use crate::resources::id::{ IdGenerator, UuidGenerator, generate_id, generate_id_default, generate_uuid, generate_uuid_default, }; -pub use crate::live_component::{ - ExceptionContext, ExceptionHandler, LiveComponent, LiveComponentOperator, LiveMapFeed, - LiveMapSubscriber, LiveMapView, MountKind, SingleWatcherGuard, SingleWatcherToken, -}; -pub use crate::resources::chunk::{Chunk, TextPosition}; -pub use crate::resources::embedder::Embedder; pub use crate::resources::live_map::LiveMap; pub use crate::resources::rate_limit::RateLimiter; pub use crate::resources::schema::{ diff --git a/rust/sdk/cocoindex/src/file.rs b/rust/sdk/cocoindex/src/resources/file.rs similarity index 100% rename from rust/sdk/cocoindex/src/file.rs rename to rust/sdk/cocoindex/src/resources/file.rs diff --git a/rust/sdk/cocoindex/src/fs.rs b/rust/sdk/cocoindex/src/resources/fs.rs similarity index 99% rename from rust/sdk/cocoindex/src/fs.rs rename to rust/sdk/cocoindex/src/resources/fs.rs index f48b968c5..c0613e0b8 100644 --- a/rust/sdk/cocoindex/src/fs.rs +++ b/rust/sdk/cocoindex/src/resources/fs.rs @@ -11,7 +11,7 @@ use walkdir::WalkDir; use crate::ctx::Ctx; use crate::error::{Error, Result}; -pub use crate::file::{ +pub use crate::resources::file::{ FileContentCache, FileLike, FileMetadata, FilePath, FilePathMatcher, FileSourceItem, MatchAllFilePathMatcher, PatternFilePathMatcher, decode_bytes, }; @@ -176,7 +176,7 @@ pub struct FileEntry { root: FilePath, relative: PathBuf, size: u64, - #[serde(with = "crate::file::system_time_serde")] + #[serde(with = "crate::resources::file::system_time_serde")] modified: SystemTime, #[serde(skip)] cache: Arc, @@ -274,7 +274,7 @@ impl FileSourceItem for FileEntry {} /// /// # Examples /// ```ignore -/// let files = cocoindex::fs::walk("./src", &["**/*.rs", "**/*.toml"])?; +/// let files = cocoindex::resources::fs::walk("./src", &["**/*.rs", "**/*.toml"])?; /// ``` pub fn walk(dir: impl AsRef, patterns: &[&str]) -> Result> { let matcher = PatternFilePathMatcher::include(patterns.iter().copied())?; @@ -444,7 +444,7 @@ impl DirTarget { /// # Examples /// /// ```no_run - /// # use cocoindex::{ctx::Ctx, fs::DirTarget}; + /// # use cocoindex::{ctx::Ctx, resources::fs::DirTarget}; /// # async fn doc(ctx: &Ctx) -> cocoindex::error::Result<()> { /// let target = DirTarget::mount(ctx, "./output")?; /// target.declare_file(ctx, "result.txt", b"final output")?; diff --git a/rust/sdk/cocoindex/src/id.rs b/rust/sdk/cocoindex/src/resources/id.rs similarity index 100% rename from rust/sdk/cocoindex/src/id.rs rename to rust/sdk/cocoindex/src/resources/id.rs diff --git a/rust/sdk/cocoindex/src/resources/mod.rs b/rust/sdk/cocoindex/src/resources/mod.rs index dfd4b3950..ac6e39e24 100644 --- a/rust/sdk/cocoindex/src/resources/mod.rs +++ b/rust/sdk/cocoindex/src/resources/mod.rs @@ -5,6 +5,9 @@ pub mod chunk; pub mod embedder; +pub mod file; +pub mod fs; +pub mod id; pub mod live_map; pub mod rate_limit; pub mod schema; diff --git a/rust/sdk/cocoindex/src/row_schema.rs b/rust/sdk/cocoindex/src/row_schema.rs index 7229923eb..d3663ebed 100644 --- a/rust/sdk/cocoindex/src/row_schema.rs +++ b/rust/sdk/cocoindex/src/row_schema.rs @@ -17,7 +17,7 @@ //! } //! // Postgres: id text NOT NULL, title text, views bigint NOT NULL, //! // embedding vector(384) NOT NULL -//! let schema = postgres::TableSchema::from_row::(["id"])?; +//! let schema = cocoindex::connectors::postgres::TableSchema::from_row::(["id"])?; //! ``` //! //! Field attributes: diff --git a/rust/sdk/cocoindex/tests/amazon_s3_source.rs b/rust/sdk/cocoindex/tests/amazon_s3_source.rs index dfb5e065a..838245a92 100644 --- a/rust/sdk/cocoindex/tests/amazon_s3_source.rs +++ b/rust/sdk/cocoindex/tests/amazon_s3_source.rs @@ -11,10 +11,10 @@ use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; -use cocoindex::amazon_s3::aws_sdk_s3::primitives::ByteStream; -use cocoindex::amazon_s3::{self, ListOptions, S3Client, S3File}; -use cocoindex::file::PatternFilePathMatcher; -use cocoindex::{App, Environment, FileLike, Result}; +use cocoindex::connectors::amazon_s3::aws_sdk_s3::primitives::ByteStream; +use cocoindex::connectors::amazon_s3::{self, ListOptions, S3Client, S3File}; +use cocoindex::resources::file::{FileLike, PatternFilePathMatcher}; +use cocoindex::{App, Environment, Result}; /// Build a client against MinIO, or `None` to skip when `AWS_ENDPOINT_URL` is unset. async fn try_client() -> Option { diff --git a/rust/sdk/cocoindex/tests/doris_target.rs b/rust/sdk/cocoindex/tests/doris_target.rs index d35982e33..8297fbcb9 100644 --- a/rust/sdk/cocoindex/tests/doris_target.rs +++ b/rust/sdk/cocoindex/tests/doris_target.rs @@ -13,14 +13,14 @@ //! the no-change skip (which must not duplicate rows), and an inverted index. //! //! Note: `USING ANN` vector indexes require Doris 3.x; the all-in-one 2.1 image -//! rejects that syntax, so the vector-index DDL is covered by a `src/doris.rs` +//! rejects that syntax, so the vector-index DDL is covered by a `src/connectors/doris.rs` //! unit test rather than live here. #![cfg(feature = "doris")] use std::sync::LazyLock; use std::time::{SystemTime, UNIX_EPOCH}; -use cocoindex::doris::{ +use cocoindex::connectors::doris::{ self, ColumnDef, DorisConfig, DorisConnection, DorisTableOptions, InvertedIndexDef, TableSchema, }; use cocoindex::{ContextKey, Environment, Result}; diff --git a/rust/sdk/cocoindex/tests/fs_live.rs b/rust/sdk/cocoindex/tests/fs_live.rs index 5449738c8..c30938564 100644 --- a/rust/sdk/cocoindex/tests/fs_live.rs +++ b/rust/sdk/cocoindex/tests/fs_live.rs @@ -66,14 +66,14 @@ async fn localfs_live_watch_reacts_to_create_and_delete() { let seen = seen.clone(); let src = src.clone(); async move { - let feed = cocoindex::fs::walk_dir(src) + let feed = cocoindex::resources::fs::walk_dir(src) .recursive(true) .live() .poll_interval(Duration::from_millis(200)); ctx.mount_each_live( &"files", feed, - move |_ctx, file: cocoindex::fs::FileEntry| { + move |_ctx, file: cocoindex::resources::fs::FileEntry| { let seen = seen.clone(); async move { seen.lock().unwrap().insert(file.key()); diff --git a/rust/sdk/cocoindex/tests/gdrive_source.rs b/rust/sdk/cocoindex/tests/gdrive_source.rs index 954564b59..e2f9e4b8b 100644 --- a/rust/sdk/cocoindex/tests/gdrive_source.rs +++ b/rust/sdk/cocoindex/tests/gdrive_source.rs @@ -1,7 +1,9 @@ #![cfg(feature = "google_drive")] -use cocoindex::FileSourceItem; -use cocoindex::gdrive::{DriveFile, DriveFileInfo, GoogleDriveClient, GoogleDriveSource}; +use cocoindex::connectors::gdrive::{ + DriveFile, DriveFileInfo, GoogleDriveClient, GoogleDriveSource, +}; +use cocoindex::resources::file::FileSourceItem; use wiremock::matchers::{header, method, path, query_param}; use wiremock::{Mock, MockServer, ResponseTemplate}; diff --git a/rust/sdk/cocoindex/tests/graph_vector_index.rs b/rust/sdk/cocoindex/tests/graph_vector_index.rs index 72fa5c088..44a2519ef 100644 --- a/rust/sdk/cocoindex/tests/graph_vector_index.rs +++ b/rust/sdk/cocoindex/tests/graph_vector_index.rs @@ -24,7 +24,7 @@ fn nonce() -> u128 { #[cfg(feature = "neo4j")] #[tokio::test] async fn neo4j_vector_index_create_then_drop_when_available() { - use cocoindex::neo4j::{self, ColumnDef, TableSchema, VectorMetric}; + use cocoindex::connectors::neo4j::{self, ColumnDef, TableSchema, VectorMetric}; use cocoindex::{App, ContextKey, Environment, Result}; use std::sync::LazyLock; @@ -109,7 +109,7 @@ async fn neo4j_vector_index_create_then_drop_when_available() { #[cfg(feature = "falkordb")] #[tokio::test] async fn falkordb_vector_index_create_then_drop_when_available() { - use cocoindex::falkordb::{self, ColumnDef, TableSchema, VectorMetric}; + use cocoindex::connectors::falkordb::{self, ColumnDef, TableSchema, VectorMetric}; use cocoindex::{App, ContextKey, Environment, Result}; use std::sync::LazyLock; use std::time::{SystemTime, UNIX_EPOCH}; diff --git a/rust/sdk/cocoindex/tests/iggy_source.rs b/rust/sdk/cocoindex/tests/iggy_source.rs index 5e9c8d570..ca898a56a 100644 --- a/rust/sdk/cocoindex/tests/iggy_source.rs +++ b/rust/sdk/cocoindex/tests/iggy_source.rs @@ -15,7 +15,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; -use cocoindex::iggy::{self, IggyConsumer, IggyProducer}; +use cocoindex::connectors::iggy::{self, IggyConsumer, IggyProducer}; use cocoindex::{App, Result, UpdateOptions}; use iggy::prelude::{ CompressionAlgorithm, Identifier, IggyExpiry, MaxTopicSize, StreamClient, TopicClient, diff --git a/rust/sdk/cocoindex/tests/iggy_target.rs b/rust/sdk/cocoindex/tests/iggy_target.rs index 6b96461f4..e50e645b6 100644 --- a/rust/sdk/cocoindex/tests/iggy_target.rs +++ b/rust/sdk/cocoindex/tests/iggy_target.rs @@ -16,7 +16,7 @@ use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; -use cocoindex::iggy::{self, IggyProducer}; +use cocoindex::connectors::iggy::{self, IggyProducer}; use cocoindex::{App, Result}; use iggy::prelude::{ CompressionAlgorithm, Consumer, Identifier, IggyExpiry, MaxTopicSize, MessageClient, diff --git a/rust/sdk/cocoindex/tests/kafka_source.rs b/rust/sdk/cocoindex/tests/kafka_source.rs index 6e34e5f6c..fcbaf53e7 100644 --- a/rust/sdk/cocoindex/tests/kafka_source.rs +++ b/rust/sdk/cocoindex/tests/kafka_source.rs @@ -17,7 +17,8 @@ use std::collections::BTreeMap; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use cocoindex::{App, Result, UpdateOptions, kafka}; +use cocoindex::connectors::kafka; +use cocoindex::{App, Result, UpdateOptions}; use rskafka::client::ClientBuilder; use rskafka::client::partition::{Compression, UnknownTopicHandling}; use rskafka::record::Record; diff --git a/rust/sdk/cocoindex/tests/kafka_target.rs b/rust/sdk/cocoindex/tests/kafka_target.rs index 912394651..a34b9cd77 100644 --- a/rust/sdk/cocoindex/tests/kafka_target.rs +++ b/rust/sdk/cocoindex/tests/kafka_target.rs @@ -18,7 +18,8 @@ use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; -use cocoindex::{App, Result, kafka}; +use cocoindex::connectors::kafka; +use cocoindex::{App, Result}; use rskafka::client::partition::UnknownTopicHandling; use rskafka::client::{Client, ClientBuilder}; diff --git a/rust/sdk/cocoindex/tests/lancedb_target.rs b/rust/sdk/cocoindex/tests/lancedb_target.rs index 4ee70fd38..50d6b2023 100644 --- a/rust/sdk/cocoindex/tests/lancedb_target.rs +++ b/rust/sdk/cocoindex/tests/lancedb_target.rs @@ -8,7 +8,7 @@ //! cargo test -p cocoindex --features lancedb --test lancedb_target #![cfg(feature = "lancedb")] -use cocoindex::lancedb::{self, ColumnDef, ColumnType, LanceDatabase, TableSchema}; +use cocoindex::connectors::lancedb::{self, ColumnDef, ColumnType, LanceDatabase, TableSchema}; use cocoindex::{ContextKey, Environment, ManagedTargetOptions, Result}; use serde::Serialize; use std::sync::LazyLock; diff --git a/rust/sdk/cocoindex/tests/oci_live.rs b/rust/sdk/cocoindex/tests/oci_live.rs index 4b8c7f92e..d7d96dc85 100644 --- a/rust/sdk/cocoindex/tests/oci_live.rs +++ b/rust/sdk/cocoindex/tests/oci_live.rs @@ -16,10 +16,11 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; -use cocoindex::oci_object_storage::{ +use cocoindex::connectors::oci_object_storage::{ ListOptions, OciClient, OciConfig, OciFile, list_objects_live, }; -use cocoindex::{App, LiveMapView, PatternFilePathMatcher, Result, UpdateOptions}; +use cocoindex::resources::file::PatternFilePathMatcher; +use cocoindex::{App, LiveMapView, Result, UpdateOptions}; use rsa::RsaPrivateKey; use rsa::pkcs8::{EncodePrivateKey, LineEnding}; use serde_json::json; diff --git a/rust/sdk/cocoindex/tests/pipeline.rs b/rust/sdk/cocoindex/tests/pipeline.rs index 7c45c1b1a..281070002 100644 --- a/rust/sdk/cocoindex/tests/pipeline.rs +++ b/rust/sdk/cocoindex/tests/pipeline.rs @@ -1,8 +1,7 @@ //! Integration tests for the pipeline: App::update, memo::cached, sync API. -use cocoindex::{ - App, ContextKey, Environment, IdGenerator, UuidGenerator, generate_id, generate_uuid, -}; +use cocoindex::resources::id::{IdGenerator, UuidGenerator, generate_id, generate_uuid}; +use cocoindex::{App, ContextKey, Environment}; use tokio::time::{Duration, sleep}; /// Helper: create an App with a temp LMDB directory (async tests). @@ -1787,8 +1786,9 @@ mod memo_test_file_state { use std::sync::{Arc, OnceLock}; use std::time::Duration; - use cocoindex::fs::{FileEntry, FilePath, walk_dir}; - use cocoindex::{ContextKey, Ctx, FileLike, Result}; + use cocoindex::resources::file::FileLike; + use cocoindex::resources::fs::{FileEntry, FilePath, walk_dir}; + use cocoindex::{ContextKey, Ctx, Result}; fn calls_key() -> &'static ContextKey> { static KEY: OnceLock>> = OnceLock::new(); @@ -2762,7 +2762,7 @@ fn fs_walk_integration() { std::fs::write(dir.path().join("lib.rs"), "pub mod foo;").unwrap(); std::fs::write(dir.path().join("readme.md"), "# Hello").unwrap(); - let files = cocoindex::fs::walk(dir.path(), &["**/*.rs"]).unwrap(); + let files = cocoindex::resources::fs::walk(dir.path(), &["**/*.rs"]).unwrap(); assert_eq!(files.len(), 2); let file = &files[0]; // lib.rs (sorted) @@ -2807,7 +2807,7 @@ async fn ctx_mount_each_rejects_duplicate_keys() { #[tokio::test] async fn dir_target_writes_skips_unchanged_and_reconciles_orphans() { - use cocoindex::DirTarget; + use cocoindex::resources::fs::DirTarget; use std::fs; use std::time::Duration; @@ -2882,7 +2882,7 @@ async fn dir_target_writes_skips_unchanged_and_reconciles_orphans() { #[tokio::test] async fn dir_target_deletes_file_when_source_disappears_via_mount_each() { - use cocoindex::DirTarget; + use cocoindex::resources::fs::DirTarget; let dir = tempfile::tempdir().unwrap(); let db = dir.path().join("lmdb"); diff --git a/rust/sdk/cocoindex/tests/postgres_source.rs b/rust/sdk/cocoindex/tests/postgres_source.rs index 2efe7dc9f..56b70e374 100644 --- a/rust/sdk/cocoindex/tests/postgres_source.rs +++ b/rust/sdk/cocoindex/tests/postgres_source.rs @@ -9,7 +9,8 @@ use std::sync::LazyLock; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; -use cocoindex::{ContextKey, Ctx, Environment, Result, postgres}; +use cocoindex::connectors::postgres; +use cocoindex::{ContextKey, Ctx, Environment, Result}; use serde::{Deserialize, Serialize}; static DB: LazyLock> = LazyLock::new(|| { diff --git a/rust/sdk/cocoindex/tests/postgres_target.rs b/rust/sdk/cocoindex/tests/postgres_target.rs index 12769c2f7..87535cab1 100644 --- a/rust/sdk/cocoindex/tests/postgres_target.rs +++ b/rust/sdk/cocoindex/tests/postgres_target.rs @@ -3,7 +3,8 @@ use std::sync::LazyLock; use std::time::{SystemTime, UNIX_EPOCH}; -use cocoindex::{ContextKey, Ctx, Environment, Result, postgres}; +use cocoindex::connectors::postgres; +use cocoindex::{ContextKey, Ctx, Environment, Result}; use serde::{Deserialize, Serialize}; use sqlx::Row as _; diff --git a/rust/sdk/cocoindex/tests/qdrant_target.rs b/rust/sdk/cocoindex/tests/qdrant_target.rs index 33731e740..33b0092e1 100644 --- a/rust/sdk/cocoindex/tests/qdrant_target.rs +++ b/rust/sdk/cocoindex/tests/qdrant_target.rs @@ -12,7 +12,7 @@ use std::sync::LazyLock; use std::time::{SystemTime, UNIX_EPOCH}; -use cocoindex::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; +use cocoindex::connectors::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; use cocoindex::{ContextKey, Environment, Result}; use serde_json::json; diff --git a/rust/sdk/cocoindex/tests/schema_from_row.rs b/rust/sdk/cocoindex/tests/schema_from_row.rs index 7812fcb26..88e28d9b7 100644 --- a/rust/sdk/cocoindex/tests/schema_from_row.rs +++ b/rust/sdk/cocoindex/tests/schema_from_row.rs @@ -7,7 +7,7 @@ #[test] fn doris_from_row_matches_explicit_schema() { use cocoindex::SchemaFields; - use cocoindex::doris::{ColumnDef, TableSchema}; + use cocoindex::connectors::doris::{ColumnDef, TableSchema}; #[derive(SchemaFields)] #[allow(dead_code)] @@ -47,7 +47,7 @@ fn doris_from_row_matches_explicit_schema() { #[test] fn sqlite_from_row_matches_explicit_schema() { use cocoindex::SchemaFields; - use cocoindex::sqlite::{ColumnDef, TableSchema}; + use cocoindex::connectors::sqlite::{ColumnDef, TableSchema}; #[derive(SchemaFields)] #[allow(dead_code)] @@ -79,7 +79,7 @@ fn sqlite_from_row_matches_explicit_schema() { #[test] fn postgres_from_row_matches_explicit_schema() { use cocoindex::SchemaFields; - use cocoindex::postgres::{ColumnDef, TableSchema}; + use cocoindex::connectors::postgres::{ColumnDef, TableSchema}; #[derive(SchemaFields)] #[allow(dead_code)] @@ -119,7 +119,7 @@ fn postgres_from_row_matches_explicit_schema() { #[cfg(feature = "sqlite")] #[tokio::test] async fn sqlite_from_row_round_trips_a_row() -> cocoindex::Result<()> { - use cocoindex::sqlite::{self, Database, TableSchema}; + use cocoindex::connectors::sqlite::{self, Database, TableSchema}; use cocoindex::{ContextKey, Environment, SchemaFields}; use sqlx::Row as _; @@ -193,7 +193,7 @@ async fn sqlite_from_row_round_trips_a_row() -> cocoindex::Result<()> { #[cfg(feature = "doris")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn doris_from_row_round_trips_a_row() -> cocoindex::Result<()> { - use cocoindex::doris::{self, DorisConfig, DorisConnection, TableSchema}; + use cocoindex::connectors::doris::{self, DorisConfig, DorisConnection, TableSchema}; use cocoindex::{ContextKey, Environment, SchemaFields}; use sqlx::Row as _; diff --git a/rust/sdk/cocoindex/tests/sqlite_target.rs b/rust/sdk/cocoindex/tests/sqlite_target.rs index eacb9fa9b..d47fb68b1 100644 --- a/rust/sdk/cocoindex/tests/sqlite_target.rs +++ b/rust/sdk/cocoindex/tests/sqlite_target.rs @@ -8,7 +8,8 @@ use std::sync::LazyLock; -use cocoindex::{App, ContextKey, Ctx, Environment, Result, sqlite}; +use cocoindex::connectors::sqlite; +use cocoindex::{App, ContextKey, Ctx, Environment, Result}; use serde::Serialize; use serde_json::json; use sqlx::Row as _; diff --git a/rust/sdk/cocoindex/tests/surrealdb_target.rs b/rust/sdk/cocoindex/tests/surrealdb_target.rs index c36a8baed..2625d5e57 100644 --- a/rust/sdk/cocoindex/tests/surrealdb_target.rs +++ b/rust/sdk/cocoindex/tests/surrealdb_target.rs @@ -2,8 +2,8 @@ use std::sync::LazyLock; +use cocoindex::connectors::surrealdb::{self, ColumnDef, Graph, TableSchema, VectorIndexOptions}; use cocoindex::prelude::*; -use cocoindex::surrealdb::{self, ColumnDef, Graph, TableSchema, VectorIndexOptions}; static GRAPH: LazyLock> = LazyLock::new(|| ContextKey::new("surrealdb_smoke_graph")); diff --git a/rust/sdk/cocoindex/tests/turbopuffer_target.rs b/rust/sdk/cocoindex/tests/turbopuffer_target.rs index 95f40d91d..b92cb618f 100644 --- a/rust/sdk/cocoindex/tests/turbopuffer_target.rs +++ b/rust/sdk/cocoindex/tests/turbopuffer_target.rs @@ -11,7 +11,9 @@ use std::sync::LazyLock; use std::time::{SystemTime, UNIX_EPOCH}; -use cocoindex::turbopuffer::{self, DistanceMetric, NamespaceSchema, TurbopufferConnection}; +use cocoindex::connectors::turbopuffer::{ + self, DistanceMetric, NamespaceSchema, TurbopufferConnection, +}; use cocoindex::{ContextKey, Environment, Result}; use serde_json::json; diff --git a/rust/sdk/cocoindex/tests/valkey_target.rs b/rust/sdk/cocoindex/tests/valkey_target.rs index ca1fd7a8f..8459ae135 100644 --- a/rust/sdk/cocoindex/tests/valkey_target.rs +++ b/rust/sdk/cocoindex/tests/valkey_target.rs @@ -7,10 +7,10 @@ use std::collections::BTreeMap; use std::sync::LazyLock; use std::time::{SystemTime, UNIX_EPOCH}; -use cocoindex::resources::schema::VectorSchema; -use cocoindex::valkey::{ +use cocoindex::connectors::valkey::{ self, Distance, Document, FieldDef, FieldType, IndexSchema, VectorAlgorithm, VectorDef, }; +use cocoindex::resources::schema::VectorSchema; use cocoindex::{ContextKey, Environment}; static VK: LazyLock> = From 8c5602f283733f528586becfc82bea505c0e20e3 Mon Sep 17 00:00:00 2001 From: zhhanz Date: Tue, 14 Jul 2026 15:38:56 -0700 Subject: [PATCH 02/15] Add ergonomic Rust function batching --- rust/core/src/engine/context.rs | 16 +- rust/sdk/cocoindex/src/batched.rs | 190 +++++++++- rust/sdk/cocoindex/tests/pipeline.rs | 163 ++++++++- rust/sdk/cocoindex_macros/src/lib.rs | 498 ++++++++++++++++++++++++++- 4 files changed, 849 insertions(+), 18 deletions(-) diff --git a/rust/core/src/engine/context.rs b/rust/core/src/engine/context.rs index a5cfcdd31..bc470998f 100644 --- a/rust/core/src/engine/context.rs +++ b/rust/core/src/engine/context.rs @@ -1000,7 +1000,7 @@ impl ComponentProcessorContext { } } -#[derive(Default)] +#[derive(Clone, Default)] pub struct FnCallContextInner { /// Target states that are declared by the function. pub target_state_paths: Vec, @@ -1044,6 +1044,10 @@ impl FnCallContext { pub fn join_child(&self, child_fn_ctx: &FnCallContext) { // Take the child's inner first to keep lock scope small (and avoid deadlock). let child_inner = child_fn_ctx.update(std::mem::take); + self.merge_child_inner(child_inner); + } + + fn merge_child_inner(&self, child_inner: FnCallContextInner) { self.update(|inner| { inner .target_state_paths @@ -1063,6 +1067,16 @@ impl FnCallContext { }); } + /// Merge a snapshot of `child_fn_ctx` without consuming it. + /// + /// Batched SDK functions execute one shared body for several per-item + /// function-call contexts. Each item must inherit the same dependencies, + /// so the shared context needs to be joined more than once. + pub fn join_child_shared(&self, child_fn_ctx: &FnCallContext) { + let child_inner = child_fn_ctx.update(|inner| inner.clone()); + self.merge_child_inner(child_inner); + } + pub fn add_fn_logic_dep(&self, fp: Fingerprint) { self.update(|inner| { inner.fn_logic_deps.insert(fp); diff --git a/rust/sdk/cocoindex/src/batched.rs b/rust/sdk/cocoindex/src/batched.rs index d8c4a52d8..b11c54b00 100644 --- a/rust/sdk/cocoindex/src/batched.rs +++ b/rust/sdk/cocoindex/src/batched.rs @@ -20,11 +20,13 @@ //! let emb = EMBED.call(&ctx, text).await?; //! ``` +use std::collections::HashMap; use std::future::Future; use std::pin::Pin; -use std::sync::Arc; +use std::sync::{Arc, Mutex, Weak}; use async_trait::async_trait; +use cocoindex_core::engine::context::FnCallContext; use cocoindex_utils::batching::{BatchQueue, Batcher, BatchingOptions, Runner}; use serde::Serialize; use serde::de::DeserializeOwned; @@ -34,7 +36,12 @@ use crate::error::{Error, Result}; type BatchFuture = Pin>> + Send>>; -type BatchFn = Box) -> BatchFuture + Send + Sync>; +type BatchFn = Arc) -> BatchFuture + Send + Sync>; + +struct BatchCall { + ctx: Ctx, + item: In, +} /// Adapts a user closure `Vec -> Result>` to the core batcher's `Runner`. struct FnRunner { @@ -43,25 +50,76 @@ struct FnRunner { #[async_trait] impl Runner for FnRunner { - type Input = In; + type Input = BatchCall; type Output = Out; async fn run( &self, - inputs: Vec, + calls: Vec>, ) -> cocoindex_utils::error::Result> { - let outputs = (self.f)(inputs).await?; + let mut contexts = Vec::with_capacity(calls.len()); + let mut inputs = Vec::with_capacity(calls.len()); + for call in calls { + contexts.push(call.ctx); + inputs.push(call.item); + } + + let batch_fn_ctx = Arc::new(FnCallContext::new(true)); + let batch_ctx = contexts + .first() + .expect("the core batcher never executes an empty batch") + .with_fn_ctx(batch_fn_ctx.clone()); + let result = (self.f)(batch_ctx, inputs).await; + + // The body executes once, but each per-item memo/tracking context must + // inherit the dependencies it observed (context keys, nested function + // logic, target states, and child memo entries). + for ctx in contexts { + if let Some(parent_fn_ctx) = &ctx.fn_ctx { + parent_fn_ctx.join_child_shared(&batch_fn_ctx); + } else if let Some(comp_ctx) = &ctx.comp_ctx { + comp_ctx.join_fn_call(&batch_fn_ctx); + } + } + + let outputs = result?; Ok(outputs.into_iter()) } } -/// A batched, memoized function. See the [module docs](self). +#[derive(Clone, PartialEq, Eq, Hash)] +struct ScheduledBatchKey { + app_id: usize, + code_hash: u64, + extra_args: Vec, +} + +type FunctionBatcher = Batcher>; +type ScheduledBatchers = Mutex>>>; + +enum BatchedMode +where + In: Send + 'static, + Out: Send + 'static, +{ + Fixed(Arc>), + Scheduled { + options: BatchingOptions, + batchers: ScheduledBatchers, + }, +} + +/// A single-item interface over a batch-shaped async function. +/// +/// [`Batched::new`] retains the explicit, memoized adapter API. The +/// `#[cocoindex::function(batching)]` macro uses the hidden scheduled mode so +/// each distinct extra-argument set gets its own short-lived batcher. pub struct Batched where In: Send + 'static, Out: Send + 'static, { - batcher: Arc>>, + mode: BatchedMode, code_hash: u64, } @@ -104,14 +162,15 @@ where F: Fn(Vec) -> Fut + Send + Sync + 'static, Fut: Future>> + Send + 'static, { - let wrapped: BatchFn = Box::new(move |inputs| { + let wrapped: BatchFn = Arc::new(move |_ctx, inputs| { let fut = f(inputs); Box::pin(async move { fut.await.map_err(Error::into_core) }) }); - let runner = FnRunner { f: wrapped }; - let queue = Arc::new(BatchQueue::new()); - let batcher = Arc::new(Batcher::new(runner, queue, options)); - Self { batcher, code_hash } + let batcher = Self::new_batcher(wrapped, options); + Self { + mode: BatchedMode::Fixed(batcher), + code_hash, + } } /// Process one item. On a memo hit the stored result is returned; on a miss @@ -120,14 +179,115 @@ where pub async fn call(&self, ctx: &Ctx, item: In) -> Result { let fp = crate::memo::key_fingerprint_result(&("cocoindex_batched", self.code_hash, &item))?; - let batcher = self.batcher.clone(); + let BatchedMode::Fixed(batcher) = &self.mode else { + return Err(Error::engine( + "scheduled Batched instances must be called by #[cocoindex::function(batching)]", + )); + }; + let batcher = batcher.clone(); // A batch impl is ctx-free, so it makes no tracked child `#[function]` // calls; the `propagate_children_fn_logic` flag is therefore inert here. // Pass `true` (the default) — only the batch impl's own `code_hash` // (folded into the memo key above) tracks its logic. - crate::memo::cached_by_fingerprint(ctx, fp, true, move |_ctx| async move { - batcher.run(item).await.map_err(Error::from) + crate::memo::cached_by_fingerprint(ctx, fp, true, move |scoped_ctx| async move { + batcher + .run(BatchCall { + ctx: scoped_ctx, + item, + }) + .await + .map_err(Error::from) }) .await } } + +impl Batched +where + In: Send + 'static, + Out: Send + 'static, +{ + fn new_batcher(f: BatchFn, options: BatchingOptions) -> Arc> { + Arc::new(Batcher::new( + FnRunner { f }, + Arc::new(BatchQueue::new()), + options, + )) + } + + /// Construct the scheduler used by generated batching wrappers. + #[doc(hidden)] + pub fn __scheduled(code_hash: u64) -> Self { + Self::__scheduled_with_options(code_hash, BatchingOptions::default()) + } + + /// Construct the generated scheduler with a maximum batch size. + #[doc(hidden)] + pub fn __scheduled_with_max_batch(code_hash: u64, max_batch_size: usize) -> Self { + Self::__scheduled_with_options( + code_hash, + BatchingOptions { + max_batch_size: Some(max_batch_size), + }, + ) + } + + fn __scheduled_with_options(code_hash: u64, options: BatchingOptions) -> Self { + Self { + mode: BatchedMode::Scheduled { + options, + batchers: Mutex::new(HashMap::new()), + }, + code_hash, + } + } + + /// Schedule one generated function call. `extra_args_key` keeps calls with + /// different captured arguments in separate batches. + #[doc(hidden)] + pub async fn __call_scheduled( + &self, + ctx: Ctx, + extra_args_key: Vec, + item: In, + f: F, + ) -> Result + where + F: Fn(Ctx, Vec) -> Fut + Send + Sync + 'static, + Fut: Future>> + Send + 'static, + { + let BatchedMode::Scheduled { options, batchers } = &self.mode else { + return Err(Error::engine( + "explicit Batched instances must be called with Batched::call", + )); + }; + + let key = ScheduledBatchKey { + // A static generated scheduler may be used by multiple apps in the + // same process. Never let their context-bound calls share a body. + app_id: Arc::as_ptr(&ctx.state) as usize, + code_hash: self.code_hash, + extra_args: extra_args_key, + }; + let f: BatchFn = Arc::new(move |ctx, inputs| { + let fut = f(ctx, inputs); + Box::pin(async move { fut.await.map_err(Error::into_core) }) + }); + let batcher = { + let mut batchers = batchers.lock().expect("batch scheduler mutex poisoned"); + batchers.retain(|_, batcher| batcher.strong_count() != 0); + if let Some(batcher) = batchers.get(&key).and_then(Weak::upgrade) { + batcher + } else { + let batcher = Self::new_batcher(f, options.clone()); + batchers.insert(key, Arc::downgrade(&batcher)); + batcher + } + }; + + batcher + .run(BatchCall { ctx, item }) + .await + .map_err(Error::from) + } +} diff --git a/rust/sdk/cocoindex/tests/pipeline.rs b/rust/sdk/cocoindex/tests/pipeline.rs index 281070002..08393eb22 100644 --- a/rust/sdk/cocoindex/tests/pipeline.rs +++ b/rust/sdk/cocoindex/tests/pipeline.rs @@ -1602,8 +1602,9 @@ async fn memo_with_function_dep_caches_when_unchanged() { mod batched_test { use super::*; - use std::sync::Arc; + use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, LazyLock}; static ITEMS_PROCESSED: AtomicUsize = AtomicUsize::new(0); @@ -1665,6 +1666,166 @@ mod batched_test { "run 2 should be all cache hits; batch impl must not reprocess items" ); } + + type RecordedBatch = (i64, Vec); + static UNMEMOIZED_BATCHES: LazyLock>> = + LazyLock::new(|| Mutex::new(Vec::new())); + + #[cocoindex::function(batching, max_batch_size = 2)] + async fn macro_unmemoized_batch( + _ctx: &cocoindex::Ctx, + items: Vec, + factor: i64, + ) -> cocoindex::Result> { + UNMEMOIZED_BATCHES + .lock() + .unwrap() + .push((factor, items.clone())); + if items == [0] { + tokio::time::sleep(Duration::from_millis(30)).await; + } + Ok(items.into_iter().map(|item| item * factor).collect()) + } + + #[tokio::test] + async fn function_batching_is_item_shaped_unmemoized_and_capped() { + UNMEMOIZED_BATCHES.lock().unwrap().clear(); + let (app, _dir) = temp_app("function_batching_unmemoized").await; + + for _ in 0..2 { + app.update(|ctx| async move { + let calls = (0..6).map(|item| macro_unmemoized_batch(&ctx, item, 3)); + let outputs = futures::future::join_all(calls).await; + let outputs = outputs.into_iter().collect::>>()?; + assert_eq!(outputs, vec![0, 3, 6, 9, 12, 15]); + Ok(()) + }) + .await + .unwrap(); + } + + let batches = UNMEMOIZED_BATCHES.lock().unwrap(); + assert_eq!( + batches.iter().map(|(_, items)| items.len()).sum::(), + 12 + ); + assert!(batches.iter().any(|(_, items)| items.len() > 1)); + assert!(batches.iter().all(|(_, items)| items.len() <= 2)); + } + + static MEMOIZED_ITEMS_PROCESSED: AtomicUsize = AtomicUsize::new(0); + + #[cocoindex::function(memo, batching, max_batch_size = 8)] + async fn macro_memoized_batch( + _ctx: &cocoindex::Ctx, + items: Vec, + factor: i64, + ) -> cocoindex::Result> { + MEMOIZED_ITEMS_PROCESSED.fetch_add(items.len(), Ordering::SeqCst); + if items == [1] { + tokio::time::sleep(Duration::from_millis(30)).await; + } + Ok(items.into_iter().map(|item| item * factor).collect()) + } + + async fn run_memoized_batch(app: &App, factor: i64) { + app.update(move |ctx| async move { + let calls = [1, 2, 3] + .into_iter() + .map(|item| macro_memoized_batch(&ctx, item, factor)); + let outputs = futures::future::join_all(calls).await; + let outputs = outputs.into_iter().collect::>>()?; + assert_eq!(outputs, vec![factor, factor * 2, factor * 3]); + Ok(()) + }) + .await + .unwrap(); + } + + #[tokio::test] + async fn function_memo_batching_caches_items_and_keys_extra_params() { + MEMOIZED_ITEMS_PROCESSED.store(0, Ordering::SeqCst); + let (app, _dir) = temp_app("function_memo_batching").await; + + run_memoized_batch(&app, 2).await; + assert_eq!(MEMOIZED_ITEMS_PROCESSED.load(Ordering::SeqCst), 3); + + run_memoized_batch(&app, 2).await; + assert_eq!( + MEMOIZED_ITEMS_PROCESSED.load(Ordering::SeqCst), + 3, + "unchanged items should all hit their per-item memo entries" + ); + + run_memoized_batch(&app, 3).await; + assert_eq!( + MEMOIZED_ITEMS_PROCESSED.load(Ordering::SeqCst), + 6, + "changing an extra parameter must invalidate every item's memo key" + ); + } + + static CONTEXT_ITEMS_PROCESSED: AtomicUsize = AtomicUsize::new(0); + + fn batch_factor_key() -> &'static ContextKey { + static KEY: LazyLock> = LazyLock::new(|| { + ContextKey::new_detect_change("pipeline/function_batching_context_factor") + }); + &KEY + } + + #[cocoindex::function(memo, batching)] + async fn macro_context_batch( + ctx: &cocoindex::Ctx, + items: Vec, + ) -> cocoindex::Result> { + let factor = *ctx.get_key(batch_factor_key())?; + CONTEXT_ITEMS_PROCESSED.fetch_add(items.len(), Ordering::SeqCst); + if items == [1] { + tokio::time::sleep(Duration::from_millis(30)).await; + } + Ok(items.into_iter().map(|item| item * factor).collect()) + } + + async fn run_context_batch(db_path: &std::path::Path, factor: i64) { + let app = Environment::builder() + .db_path(db_path) + .provide_key(batch_factor_key(), factor) + .build() + .await + .unwrap() + .app("function_context_batch") + .await + .unwrap(); + app.update(move |ctx| async move { + let calls = [1, 2, 3] + .into_iter() + .map(|item| macro_context_batch(&ctx, item)); + let outputs = futures::future::join_all(calls).await; + let outputs = outputs.into_iter().collect::>>()?; + assert_eq!(outputs, vec![factor, factor * 2, factor * 3]); + Ok(()) + }) + .await + .unwrap(); + } + + #[tokio::test] + async fn function_memo_batching_shares_context_dependencies_with_every_item() { + CONTEXT_ITEMS_PROCESSED.store(0, Ordering::SeqCst); + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("lmdb"); + + run_context_batch(&db_path, 2).await; + assert_eq!(CONTEXT_ITEMS_PROCESSED.load(Ordering::SeqCst), 3); + + run_context_batch(&db_path, 3).await; + assert_eq!( + CONTEXT_ITEMS_PROCESSED.load(Ordering::SeqCst), + 6, + "every cached item must inherit the batch body's context-key dependency" + ); + } } // --------------------------------------------------------------------------- diff --git a/rust/sdk/cocoindex_macros/src/lib.rs b/rust/sdk/cocoindex_macros/src/lib.rs index 308af6eae..6fae65497 100644 --- a/rust/sdk/cocoindex_macros/src/lib.rs +++ b/rust/sdk/cocoindex_macros/src/lib.rs @@ -12,8 +12,10 @@ use syn::{ }; /// Information about a non-ctx parameter. +#[derive(Clone, Debug)] struct ParamInfo { ident: syn::Ident, + ty: Type, is_ref: bool, is_str_ref: bool, } @@ -51,6 +53,7 @@ fn parse_fn_params_opt(func: &ItemFn) -> syn::Result<(Option, Vec, version: Option, memo_key: Vec, logic_tracking: LogicTracking, @@ -186,6 +191,8 @@ impl FunctionArgs { impl Parse for FunctionArgs { fn parse(input: ParseStream) -> syn::Result { let mut memo = false; + let mut batching = false; + let mut max_batch_size = None; let mut version = None; let mut memo_key = Vec::new(); let mut logic_tracking = LogicTracking::Full; @@ -194,6 +201,27 @@ impl Parse for FunctionArgs { let name: Ident = input.parse()?; match name.to_string().as_str() { "memo" => memo = true, + "batching" => batching = true, + "max_batch_size" => { + input.parse::()?; + let value: LitInt = input.parse()?; + if max_batch_size.is_some() { + return Err(SynError::new( + value.span(), + "duplicate `max_batch_size` argument", + )); + } + let value = value + .base10_parse::() + .map_err(|err| SynError::new(value.span(), err.to_string()))?; + if value == 0 { + return Err(SynError::new( + name.span(), + "`max_batch_size` must be greater than zero", + )); + } + max_batch_size = Some(value); + } "logic_tracking" => { input.parse::()?; let mode: LitStr = input.parse()?; @@ -245,7 +273,7 @@ impl Parse for FunctionArgs { _ => { return Err(SynError::new( name.span(), - "unsupported function attribute argument. expected `memo`, `memo_key(...)`, `version = N`, or `logic_tracking = \"full\"|\"self\"|\"none\"`", + "unsupported function attribute argument. expected `memo`, `batching`, `max_batch_size = N`, `memo_key(...)`, `version = N`, or `logic_tracking = \"full\"|\"self\"|\"none\"`", )); } } @@ -258,8 +286,17 @@ impl Parse for FunctionArgs { } } + if max_batch_size.is_some() && !batching { + return Err(SynError::new( + Span::call_site(), + "`max_batch_size` requires `batching`", + )); + } + Ok(Self { memo, + batching, + max_batch_size, version, memo_key, logic_tracking, @@ -369,6 +406,383 @@ fn gen_state_collect_for_param( } } +#[derive(Debug)] +struct BatchingSignature { + ctx_ident: Ident, + item_param: ParamInfo, + extra_params: Vec, + output_ty: Type, + wrapper_sig: syn::Signature, +} + +fn vec_inner_type(ty: &Type) -> Option { + let Type::Path(type_path) = ty else { + return None; + }; + let segment = type_path.path.segments.last()?; + if segment.ident != "Vec" { + return None; + } + let syn::PathArguments::AngleBracketed(args) = &segment.arguments else { + return None; + }; + args.args.iter().find_map(|arg| match arg { + syn::GenericArgument::Type(ty) => Some(ty.clone()), + _ => None, + }) +} + +fn item_shaped_return_type(output: &syn::ReturnType) -> syn::Result<(Type, syn::ReturnType)> { + let syn::ReturnType::Type(_, result_ty) = output else { + return Err(SynError::new_spanned( + output, + "batching function return type must be `Result>`", + )); + }; + let Type::Path(result_path) = result_ty.as_ref() else { + return Err(SynError::new_spanned( + result_ty, + "batching function return type must be `Result>`", + )); + }; + let Some(result_segment) = result_path.path.segments.last() else { + return Err(SynError::new_spanned( + result_ty, + "batching function return type must be `Result>`", + )); + }; + if result_segment.ident != "Result" { + return Err(SynError::new_spanned( + result_ty, + "batching function return type must be `Result>`", + )); + } + let syn::PathArguments::AngleBracketed(result_args) = &result_segment.arguments else { + return Err(SynError::new_spanned( + result_ty, + "batching function return type must be `Result>`", + )); + }; + let Some(syn::GenericArgument::Type(batch_output_ty)) = result_args.args.first() else { + return Err(SynError::new_spanned( + result_ty, + "batching function return type must be `Result>`", + )); + }; + let Some(item_output_ty) = vec_inner_type(batch_output_ty) else { + return Err(SynError::new_spanned( + batch_output_ty, + "batching function return type must be `Result>`", + )); + }; + + let mut wrapper_output = output.clone(); + let syn::ReturnType::Type(_, wrapper_result_ty) = &mut wrapper_output else { + unreachable!(); + }; + let Type::Path(wrapper_result_path) = wrapper_result_ty.as_mut() else { + unreachable!(); + }; + let wrapper_result_segment = wrapper_result_path.path.segments.last_mut().unwrap(); + let syn::PathArguments::AngleBracketed(wrapper_args) = &mut wrapper_result_segment.arguments + else { + unreachable!(); + }; + let Some(syn::GenericArgument::Type(wrapper_item_ty)) = wrapper_args.args.first_mut() else { + unreachable!(); + }; + *wrapper_item_ty = item_output_ty.clone(); + Ok((item_output_ty, wrapper_output)) +} + +fn parse_batching_signature(func: &ItemFn) -> syn::Result { + if func.sig.asyncness.is_none() { + return Err(SynError::new( + func.sig.ident.span(), + "batching requires an async function", + )); + } + if !func.sig.generics.params.is_empty() { + return Err(SynError::new_spanned( + &func.sig.generics, + "batching functions cannot be generic because their scheduler is static", + )); + } + for input in &func.sig.inputs { + let FnArg::Typed(param) = input else { + return Err(SynError::new_spanned( + input, + "batching is supported only on free functions", + )); + }; + if !matches!(param.pat.as_ref(), Pat::Ident(_)) { + return Err(SynError::new_spanned( + ¶m.pat, + "batching function parameters must be identifiers", + )); + } + } + + let (ctx_ident, params) = parse_fn_params(func)?; + let Some(batch_param) = params.first() else { + return Err(SynError::new( + func.sig.ident.span(), + "batching requires a non-context `Vec<_>` parameter", + )); + }; + let Some(item_ty) = vec_inner_type(&batch_param.ty) else { + return Err(SynError::new_spanned( + &batch_param.ty, + "batching requires the first non-context parameter to be `Vec<_>`", + )); + }; + let (output_ty, wrapper_output) = item_shaped_return_type(&func.sig.output)?; + + let mut wrapper_sig = func.sig.clone(); + let mut found_batch_param = false; + for input in &mut wrapper_sig.inputs { + let FnArg::Typed(param) = input else { + return Err(SynError::new_spanned( + input, + "batching is supported only on free functions", + )); + }; + let Pat::Ident(param_ident) = param.pat.as_ref() else { + return Err(SynError::new_spanned( + ¶m.pat, + "batching function parameters must be identifiers", + )); + }; + if param_ident.ident == batch_param.ident { + *param.ty = item_ty.clone(); + found_batch_param = true; + break; + } + } + debug_assert!(found_batch_param); + wrapper_sig.output = wrapper_output; + + Ok(BatchingSignature { + ctx_ident, + item_param: ParamInfo { + ident: batch_param.ident.clone(), + ty: item_ty, + is_ref: false, + is_str_ref: false, + }, + extra_params: params[1..].to_vec(), + output_ty, + wrapper_sig, + }) +} + +fn gen_owned_param_clones(params: &[ParamInfo]) -> Vec { + params + .iter() + .map(|param| { + let ident = ¶m.ident; + quote! { let #ident = ::core::clone::Clone::clone(&#ident); } + }) + .collect() +} + +#[allow(clippy::too_many_arguments)] +fn expand_batching_function( + func: &ItemFn, + args: &FunctionArgs, + hash_const_name: &Ident, + hash_const_value: &TokenStream2, + logic_registration: &TokenStream2, + track_logic: bool, + propagate_children_fn_logic: bool, +) -> syn::Result { + let batching = parse_batching_signature(func)?; + let fn_name = &func.sig.ident; + let vis = &func.vis; + let attrs = &func.attrs; + let body = &func.block; + let wrapper_sig = &batching.wrapper_sig; + let ctx_ident = &batching.ctx_ident; + let item_ident = &batching.item_param.ident; + let item_ty = &batching.item_param.ty; + let output_ty = &batching.output_ty; + let extra_params = &batching.extra_params; + + let batch_impl_name = format_ident!("__coco_batch_impl_{}", fn_name); + let mut batch_impl_sig = func.sig.clone(); + batch_impl_sig.ident = batch_impl_name.clone(); + let batch_static_name = format_ident!("__COCO_BATCHED_{}", fn_name.to_string().to_uppercase()); + let batch_static_init = match args.max_batch_size { + Some(max_batch_size) => quote! { + ::cocoindex::Batched::__scheduled_with_max_batch( + #hash_const_name, + #max_batch_size, + ) + }, + None => quote! { ::cocoindex::Batched::__scheduled(#hash_const_name) }, + }; + + let mut batch_impl_args = Vec::new(); + for input in &func.sig.inputs { + let FnArg::Typed(param) = input else { + return Err(SynError::new_spanned( + input, + "batching is supported only on free functions", + )); + }; + let Pat::Ident(param_ident) = param.pat.as_ref() else { + return Err(SynError::new_spanned( + ¶m.pat, + "batching function parameters must be identifiers", + )); + }; + let ident = ¶m_ident.ident; + if *ident == batching.ctx_ident { + batch_impl_args.push(quote! { &__coco_batch_ctx }); + } else if *ident == batching.item_param.ident { + batch_impl_args.push(quote! { __coco_batch_items }); + } else { + let param = extra_params + .iter() + .find(|param| param.ident == *ident) + .expect("validated non-context parameter"); + if param.is_ref { + batch_impl_args.push(quote! { &#ident }); + } else { + batch_impl_args.push(quote! { #ident }); + } + } + } + + let owned_extra_clones = gen_owned_param_clones(extra_params); + let extra_key_parts = extra_params.iter().map(|param| { + let ident = ¶m.ident; + quote! { &#ident } + }); + let extra_args_key = if extra_params.is_empty() { + quote! { ::std::vec::Vec::new() } + } else { + quote! { + ::cocoindex::memo::key_bytes_result(&(#(#extra_key_parts,)*))? + } + }; + let schedule_call = quote! { + let __coco_extra_args_key = #extra_args_key; + #batch_static_name + .__call_scheduled( + __coco_scoped_ctx, + __coco_extra_args_key, + #item_ident, + move |__coco_batch_ctx, __coco_batch_items| { + #(#owned_extra_clones)* + async move { + #batch_impl_name(#(#batch_impl_args),*).await + } + }, + ) + .await + }; + + let wrapper_body = if args.memo { + let mut key_params = vec![batching.item_param.clone()]; + key_params.extend(extra_params.iter().cloned()); + validate_memo_key_overrides( + &args.memo_key, + key_params.iter().map(|param| param.ident.to_string()), + )?; + let key_writes: Vec = key_params + .iter() + .filter_map(|param| { + gen_key_write_for_param(&format_ident!("__coco_key"), param, &args.memo_key) + }) + .collect(); + let state_collects: Vec = key_params + .iter() + .map(|param| { + gen_state_collect_for_param( + &format_ident!("__coco_states"), + &format_ident!("__coco_state_idx"), + &format_ident!("__coco_prev_states"), + param, + ) + }) + .collect(); + let state_clones = gen_clones(&key_params); + let body_clones = gen_clones(&key_params); + quote! {{ + let __coco_key = { + let mut __coco_key = ::cocoindex::memo::new_key_fingerprinter(); + ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &"cocoindex_fn")?; + ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &::core::module_path!())?; + ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &::core::stringify!(#fn_name))?; + ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &#hash_const_name)?; + #(#key_writes)* + ::cocoindex::memo::finish_key_fingerprinter(__coco_key) + }; + + ::cocoindex::memo::cached_by_fingerprint_with_state( + #ctx_ident, + __coco_key, + #propagate_children_fn_logic, + { + #(#state_clones)* + move |__coco_prev_states| async move { + let mut __coco_states = Vec::new(); + let mut __coco_state_idx = 0usize; + #(#state_collects)* + Ok(__coco_states) + } + }, + { + #(#body_clones)* + move |__coco_scoped_ctx| async move { + #schedule_call + } + }, + ) + .await + }} + } else { + let initial_extra_clones = gen_clones(extra_params); + if track_logic { + quote! {{ + #(#initial_extra_clones)* + #ctx_ident + .__coco_tracked_fn( + ::core::module_path!(), + ::core::stringify!(#fn_name), + #hash_const_name, + #propagate_children_fn_logic, + move |__coco_scoped_ctx| async move { + #schedule_call + }, + ) + .await + }} + } else { + quote! {{ + #(#initial_extra_clones)* + let __coco_scoped_ctx = ::core::clone::Clone::clone(#ctx_ident); + #schedule_call + }} + } + }; + + Ok(quote! { + #[doc(hidden)] + pub const #hash_const_name: u64 = #hash_const_value; + #logic_registration + + static #batch_static_name: ::std::sync::LazyLock<::cocoindex::Batched<#item_ty, #output_ty>> = + ::std::sync::LazyLock::new(|| #batch_static_init); + + #batch_impl_sig #body + + #(#attrs)* + #vis #wrapper_sig #wrapper_body + }) +} + /// `#[cocoindex::function]` — unified macro for cocoindex pipeline functions. /// /// ## Usage @@ -385,6 +799,14 @@ fn gen_state_collect_for_param( /// async fn my_fn(ctx: &Ctx, arg: &String) -> Result { ... } /// ``` /// +/// **With `batching`** — declare a batch-shaped body and call it with one item: +/// ```ignore +/// #[cocoindex::function(memo, batching, max_batch_size = 32)] +/// async fn embed(ctx: &Ctx, texts: Vec) -> Result>> { ... } +/// +/// let embedding: Vec = embed(ctx, text).await?; +/// ``` +/// /// Optional `version` parameter forces cache invalidation: /// ```ignore /// #[cocoindex::function(memo, version = 2)] @@ -457,6 +879,21 @@ pub fn function(attr: TokenStream, item: TokenStream) -> TokenStream { ); } + if args.batching { + return match expand_batching_function( + &func, + &args, + &hash_const_name, + &hash_const_value, + &logic_registration, + track_logic, + propagate_children_fn_logic, + ) { + Ok(expanded) => expanded.into(), + Err(err) => err.to_compile_error().into(), + }; + } + if args.memo { // memo: wrap body in ctx.memo() let (ctx_ident, params) = match parse_fn_params(&func) { @@ -1020,9 +1457,25 @@ mod tests { fn parse_function_args_empty() { let args = FunctionArgs::parse(TokenStream2::new()).unwrap(); assert!(!args.memo); + assert!(!args.batching); + assert!(args.max_batch_size.is_none()); assert!(args.version.is_none()); } + #[test] + fn parse_function_args_accepts_batching_options() { + let args = FunctionArgs::parse(quote!(memo, batching, max_batch_size = 32)).unwrap(); + assert!(args.memo); + assert!(args.batching); + assert_eq!(args.max_batch_size, Some(32)); + } + + #[test] + fn parse_function_args_rejects_max_batch_without_batching() { + let err = FunctionArgs::parse(quote!(max_batch_size = 32)).unwrap_err(); + assert!(err.to_string().contains("requires `batching`")); + } + #[test] fn parse_function_args_accepts_memo_key_overrides() { let args = FunctionArgs::parse(quote!( @@ -1145,6 +1598,49 @@ mod tests { assert!(parse_fn_params(&func).is_err()); } + #[test] + fn batching_signature_rewrites_input_and_output_to_item_shape() { + let func: ItemFn = parse_str( + "async fn embed(ctx: &Ctx, texts: Vec, model: String) -> Result>> { unimplemented!() }", + ) + .unwrap(); + let batching = parse_batching_signature(&func).unwrap(); + assert_eq!(batching.ctx_ident, "ctx"); + assert_eq!(batching.item_param.ident, "texts"); + assert_eq!(batching.extra_params.len(), 1); + let wrapper_sig = &batching.wrapper_sig; + assert_eq!( + quote!(#wrapper_sig).to_string(), + "async fn embed (ctx : & Ctx , texts : String , model : String) -> Result < Vec < f32 > >" + ); + } + + #[test] + fn batching_signature_requires_non_context_parameter() { + let func: ItemFn = + parse_str("async fn bad(ctx: &Ctx) -> Result> { Ok(vec![]) }").unwrap(); + let err = parse_batching_signature(&func).unwrap_err(); + assert!(err.to_string().contains("requires a non-context")); + } + + #[test] + fn batching_signature_requires_vec_input() { + let func: ItemFn = + parse_str("async fn bad(ctx: &Ctx, item: String) -> Result> { Ok(vec![]) }") + .unwrap(); + let err = parse_batching_signature(&func).unwrap_err(); + assert!(err.to_string().contains("first non-context parameter")); + } + + #[test] + fn batching_signature_requires_result_vec_output() { + let func: ItemFn = + parse_str("async fn bad(ctx: &Ctx, items: Vec) -> Result { Ok(1) }") + .unwrap(); + let err = parse_batching_signature(&func).unwrap_err(); + assert!(err.to_string().contains("return type")); + } + #[test] fn explicit_version_replaces_body_hash() { let first: ItemFn = From a2fdfd87187e558294553bd85dccbdc7e863506c Mon Sep 17 00:00:00 2001 From: zhhanz Date: Tue, 14 Jul 2026 15:53:20 -0700 Subject: [PATCH 03/15] Add declarative Rust context keys --- examples/rust/text_embedding/src/main.rs | 19 ++--- rust/sdk/cocoindex/src/connectors/falkordb.rs | 5 +- rust/sdk/cocoindex/src/connectors/neo4j.rs | 5 +- rust/sdk/cocoindex/src/ctx.rs | 74 +++++++++++++++++++ rust/sdk/cocoindex/tests/amazon_s3_source.rs | 12 ++- rust/sdk/cocoindex/tests/doris_target.rs | 6 +- .../sdk/cocoindex/tests/graph_vector_index.rs | 12 +-- rust/sdk/cocoindex/tests/lancedb_target.rs | 12 ++- rust/sdk/cocoindex/tests/postgres_source.rs | 12 ++- rust/sdk/cocoindex/tests/postgres_target.rs | 12 ++- rust/sdk/cocoindex/tests/qdrant_target.rs | 12 ++- rust/sdk/cocoindex/tests/sqlite_target.rs | 13 ++-- rust/sdk/cocoindex/tests/surrealdb_target.rs | 5 +- .../sdk/cocoindex/tests/turbopuffer_target.rs | 12 ++- rust/sdk/cocoindex/tests/valkey_target.rs | 6 +- 15 files changed, 130 insertions(+), 87 deletions(-) diff --git a/examples/rust/text_embedding/src/main.rs b/examples/rust/text_embedding/src/main.rs index 0a5f19611..28e18031b 100644 --- a/examples/rust/text_embedding/src/main.rs +++ b/examples/rust/text_embedding/src/main.rs @@ -13,7 +13,6 @@ //! - target : `postgres::TableTarget` + pgvector index use std::path::PathBuf; -use std::sync::LazyLock; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; @@ -30,16 +29,14 @@ const TOP_K: i64 = 5; const CHUNK_SIZE: usize = 2000; const CHUNK_OVERLAP: usize = 500; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("text_embedding_db", |db: &postgres::Database| { - db.state_id().to_string() - }) -}); -static EMBEDDER: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("embedder", |e: &SentenceTransformerEmbedder| { - e.model_name().to_string() - }) -}); +cocoindex::context_key!( + static DB: postgres::Database = "text_embedding_db", + state = postgres::Database::state_id +); +cocoindex::context_key!( + static EMBEDDER: SentenceTransformerEmbedder = "embedder", + state = SentenceTransformerEmbedder::model_name +); #[derive(Clone, Serialize, Deserialize)] struct DocEmbedding { diff --git a/rust/sdk/cocoindex/src/connectors/falkordb.rs b/rust/sdk/cocoindex/src/connectors/falkordb.rs index 8c07abedd..ab789da4e 100644 --- a/rust/sdk/cocoindex/src/connectors/falkordb.rs +++ b/rust/sdk/cocoindex/src/connectors/falkordb.rs @@ -123,15 +123,14 @@ cypher_graph::graph_target_api!(Graph); #[cfg(test)] mod tests { - use std::sync::LazyLock; use std::time::{SystemTime, UNIX_EPOCH}; use serde::Serialize; use super::*; - use crate::{App, ContextKey, Environment}; + use crate::{App, Environment}; - static GRAPH: LazyLock> = LazyLock::new(|| ContextKey::new("falkordb_graph")); + crate::context_key!(static GRAPH: Graph = "falkordb_graph"); #[derive(Serialize)] struct Person { diff --git a/rust/sdk/cocoindex/src/connectors/neo4j.rs b/rust/sdk/cocoindex/src/connectors/neo4j.rs index b445edb86..9321dd492 100644 --- a/rust/sdk/cocoindex/src/connectors/neo4j.rs +++ b/rust/sdk/cocoindex/src/connectors/neo4j.rs @@ -135,15 +135,14 @@ fn validate_database(database: &str) -> Result<()> { #[cfg(test)] mod tests { - use std::sync::LazyLock; use std::time::{SystemTime, UNIX_EPOCH}; use serde::Serialize; use super::*; - use crate::{App, ContextKey, Environment}; + use crate::{App, Environment}; - static GRAPH: LazyLock> = LazyLock::new(|| ContextKey::new("neo4j_graph")); + crate::context_key!(static GRAPH: Graph = "neo4j_graph"); #[derive(Serialize)] struct Person { diff --git a/rust/sdk/cocoindex/src/ctx.rs b/rust/sdk/cocoindex/src/ctx.rs index bfb1e2954..eee3d1ced 100644 --- a/rust/sdk/cocoindex/src/ctx.rs +++ b/rust/sdk/cocoindex/src/ctx.rs @@ -30,6 +30,49 @@ use crate::user_state::{IntoStateKey, StateHandle}; type ContextFingerprinter = Arc Result + Send + Sync>; +/// Declare a process-wide [`ContextKey`] backed by a [`std::sync::LazyLock`]. +/// +/// The three forms correspond to [`ContextKey::new`], +/// [`ContextKey::new_detect_change`], and [`ContextKey::new_with_state`]: +/// +/// ``` +/// # use cocoindex::ContextKey; +/// # #[derive(serde::Serialize)] +/// # struct AppConfig; +/// # struct Database; +/// # impl Database { fn state_id(&self) -> &str { "database" } } +/// cocoindex::context_key!(static CONFIG: AppConfig = "app_config"); +/// cocoindex::context_key!(static TRACKED_CONFIG: AppConfig = "tracked_config", detect_change); +/// cocoindex::context_key!( +/// static DB: Database = "database", +/// state = Database::state_id +/// ); +/// ``` +/// +/// The key name must be an explicit string literal. It is a stable persistent +/// identity used in target-state and memo keys, so it must not change merely +/// because the Rust item moves to another module. The `state =` expression may +/// be any path or closure implementing `Fn(&T) -> S` where `S: ToString`. +#[macro_export] +macro_rules! context_key { + ($vis:vis static $key:ident: $value:ty = $name:literal, detect_change $(,)?) => { + $vis static $key: ::std::sync::LazyLock<$crate::ContextKey<$value>> = + ::std::sync::LazyLock::new(|| $crate::ContextKey::new_detect_change($name)); + }; + ($vis:vis static $key:ident: $value:ty = $name:literal, state = $state:expr $(,)?) => { + $vis static $key: ::std::sync::LazyLock<$crate::ContextKey<$value>> = + ::std::sync::LazyLock::new(|| { + $crate::ContextKey::new_with_state($name, |value: &$value| { + ($state)(value).to_string() + }) + }); + }; + ($vis:vis static $key:ident: $value:ty = $name:literal $(,)?) => { + $vis static $key: ::std::sync::LazyLock<$crate::ContextKey<$value>> = + ::std::sync::LazyLock::new(|| $crate::ContextKey::new($name)); + }; +} + /// A named context key for app-provided resources. /// /// - [`ContextKey::new`] stores arbitrary `Send + Sync` resources (no change @@ -1200,6 +1243,27 @@ mod tests { static TEST_CLOCK_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + #[derive(Serialize)] + struct TrackedConfig; + + struct StatefulResource; + + impl StatefulResource { + fn state_id(&self) -> u64 { + 42 + } + } + + crate::context_key!(static PLAIN_KEY: String = "ctx_macro_plain"); + crate::context_key!( + static DETECT_CHANGE_KEY: TrackedConfig = "ctx_macro_detect_change", + detect_change + ); + crate::context_key!( + static DERIVED_STATE_KEY: StatefulResource = "ctx_macro_derived_state", + state = StatefulResource::state_id + ); + struct TestClockGuard { _guard: std::sync::MutexGuard<'static, ()>, } @@ -1224,6 +1288,16 @@ mod tests { } } + #[test] + fn context_key_macro_supports_all_three_forms() { + assert_eq!(PLAIN_KEY.name(), "ctx_macro_plain"); + assert!(!PLAIN_KEY.detect_change()); + assert_eq!(DETECT_CHANGE_KEY.name(), "ctx_macro_detect_change"); + assert!(DETECT_CHANGE_KEY.detect_change()); + assert_eq!(DERIVED_STATE_KEY.name(), "ctx_macro_derived_state"); + assert!(DERIVED_STATE_KEY.detect_change()); + } + #[tokio::test(flavor = "current_thread")] async fn next_raw_id_checks_deadline_before_allocating() { let clock = TestClockGuard::new(); diff --git a/rust/sdk/cocoindex/tests/amazon_s3_source.rs b/rust/sdk/cocoindex/tests/amazon_s3_source.rs index 838245a92..ab3872538 100644 --- a/rust/sdk/cocoindex/tests/amazon_s3_source.rs +++ b/rust/sdk/cocoindex/tests/amazon_s3_source.rs @@ -14,7 +14,7 @@ use std::time::{SystemTime, UNIX_EPOCH}; use cocoindex::connectors::amazon_s3::aws_sdk_s3::primitives::ByteStream; use cocoindex::connectors::amazon_s3::{self, ListOptions, S3Client, S3File}; use cocoindex::resources::file::{FileLike, PatternFilePathMatcher}; -use cocoindex::{App, Environment, Result}; +use cocoindex::{Environment, Result}; /// Build a client against MinIO, or `None` to skip when `AWS_ENDPOINT_URL` is unset. async fn try_client() -> Option { @@ -226,12 +226,10 @@ async fn s3_source_mount_each_pipeline_when_available() -> Result<()> { ) .await; - static S3: std::sync::LazyLock> = - std::sync::LazyLock::new(|| { - cocoindex::ContextKey::new_with_state("s3_test_client", |c: &S3Client| { - c.state_id().to_string() - }) - }); + cocoindex::context_key!( + static S3: S3Client = "s3_test_client", + state = S3Client::state_id + ); #[cocoindex::function(memo)] async fn process(ctx: &cocoindex::Ctx, file: &S3File) -> Result { diff --git a/rust/sdk/cocoindex/tests/doris_target.rs b/rust/sdk/cocoindex/tests/doris_target.rs index 8297fbcb9..6fd2be4e4 100644 --- a/rust/sdk/cocoindex/tests/doris_target.rs +++ b/rust/sdk/cocoindex/tests/doris_target.rs @@ -17,18 +17,16 @@ //! unit test rather than live here. #![cfg(feature = "doris")] -use std::sync::LazyLock; use std::time::{SystemTime, UNIX_EPOCH}; use cocoindex::connectors::doris::{ self, ColumnDef, DorisConfig, DorisConnection, DorisTableOptions, InvertedIndexDef, TableSchema, }; -use cocoindex::{ContextKey, Environment, Result}; +use cocoindex::{Environment, Result}; use serde::Serialize; use sqlx::Row; -static DORIS_DB: LazyLock> = - LazyLock::new(|| ContextKey::new("doris_target_test_db")); +cocoindex::context_key!(static DORIS_DB: DorisConnection = "doris_target_test_db"); #[derive(Serialize, Clone)] struct Item { diff --git a/rust/sdk/cocoindex/tests/graph_vector_index.rs b/rust/sdk/cocoindex/tests/graph_vector_index.rs index 44a2519ef..a04d8438f 100644 --- a/rust/sdk/cocoindex/tests/graph_vector_index.rs +++ b/rust/sdk/cocoindex/tests/graph_vector_index.rs @@ -25,11 +25,9 @@ fn nonce() -> u128 { #[tokio::test] async fn neo4j_vector_index_create_then_drop_when_available() { use cocoindex::connectors::neo4j::{self, ColumnDef, TableSchema, VectorMetric}; - use cocoindex::{App, ContextKey, Environment, Result}; - use std::sync::LazyLock; + use cocoindex::{Environment, Result}; - static G: LazyLock> = - LazyLock::new(|| ContextKey::new("neo4j_vidx_graph")); + cocoindex::context_key!(static G: neo4j::Graph = "neo4j_vidx_graph"); let uri = std::env::var("NEO4J_URI").unwrap_or_else(|_| "bolt://localhost:7687".to_string()); let user = std::env::var("NEO4J_USER").unwrap_or_else(|_| "neo4j".to_string()); @@ -110,12 +108,10 @@ async fn neo4j_vector_index_create_then_drop_when_available() { #[tokio::test] async fn falkordb_vector_index_create_then_drop_when_available() { use cocoindex::connectors::falkordb::{self, ColumnDef, TableSchema, VectorMetric}; - use cocoindex::{App, ContextKey, Environment, Result}; - use std::sync::LazyLock; + use cocoindex::{Environment, Result}; use std::time::{SystemTime, UNIX_EPOCH}; - static G: LazyLock> = - LazyLock::new(|| ContextKey::new("falkordb_vidx_graph")); + cocoindex::context_key!(static G: falkordb::Graph = "falkordb_vidx_graph"); let uri = std::env::var("FALKORDB_URI").unwrap_or_else(|_| "falkor://localhost:6379".to_string()); diff --git a/rust/sdk/cocoindex/tests/lancedb_target.rs b/rust/sdk/cocoindex/tests/lancedb_target.rs index 50d6b2023..6590349d0 100644 --- a/rust/sdk/cocoindex/tests/lancedb_target.rs +++ b/rust/sdk/cocoindex/tests/lancedb_target.rs @@ -9,15 +9,13 @@ #![cfg(feature = "lancedb")] use cocoindex::connectors::lancedb::{self, ColumnDef, ColumnType, LanceDatabase, TableSchema}; -use cocoindex::{ContextKey, Environment, ManagedTargetOptions, Result}; +use cocoindex::{Environment, ManagedTargetOptions, Result}; use serde::Serialize; -use std::sync::LazyLock; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("lancedb_test", |db: &LanceDatabase| { - db.state_id().to_string() - }) -}); +cocoindex::context_key!( + static DB: LanceDatabase = "lancedb_test", + state = LanceDatabase::state_id +); const TABLE: &str = "docs"; diff --git a/rust/sdk/cocoindex/tests/postgres_source.rs b/rust/sdk/cocoindex/tests/postgres_source.rs index 56b70e374..0eaac9671 100644 --- a/rust/sdk/cocoindex/tests/postgres_source.rs +++ b/rust/sdk/cocoindex/tests/postgres_source.rs @@ -5,19 +5,17 @@ //! cargo test -p cocoindex --features postgres --test postgres_source #![cfg(feature = "postgres")] -use std::sync::LazyLock; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; use cocoindex::connectors::postgres; -use cocoindex::{ContextKey, Ctx, Environment, Result}; +use cocoindex::{Ctx, Environment, Result}; use serde::{Deserialize, Serialize}; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("postgres_source_test_db", |db: &postgres::Database| { - db.state_id().to_string() - }) -}); +cocoindex::context_key!( + static DB: postgres::Database = "postgres_source_test_db", + state = postgres::Database::state_id +); static CALLS: AtomicUsize = AtomicUsize::new(0); #[derive(Clone, Serialize, Deserialize)] diff --git a/rust/sdk/cocoindex/tests/postgres_target.rs b/rust/sdk/cocoindex/tests/postgres_target.rs index 87535cab1..f6d8fd86f 100644 --- a/rust/sdk/cocoindex/tests/postgres_target.rs +++ b/rust/sdk/cocoindex/tests/postgres_target.rs @@ -1,18 +1,16 @@ #![cfg(feature = "postgres")] -use std::sync::LazyLock; use std::time::{SystemTime, UNIX_EPOCH}; use cocoindex::connectors::postgres; -use cocoindex::{ContextKey, Ctx, Environment, Result}; +use cocoindex::{Ctx, Environment, Result}; use serde::{Deserialize, Serialize}; use sqlx::Row as _; -static PG: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("postgres_target_test_db", |db: &postgres::Database| { - db.state_id().to_string() - }) -}); +cocoindex::context_key!( + static PG: postgres::Database = "postgres_target_test_db", + state = postgres::Database::state_id +); #[derive(Clone, Serialize)] struct TestRow { diff --git a/rust/sdk/cocoindex/tests/qdrant_target.rs b/rust/sdk/cocoindex/tests/qdrant_target.rs index 33b0092e1..6b7a9a32c 100644 --- a/rust/sdk/cocoindex/tests/qdrant_target.rs +++ b/rust/sdk/cocoindex/tests/qdrant_target.rs @@ -9,18 +9,16 @@ //! in-place update, orphan delete, and schema-change collection recreate. #![cfg(feature = "qdrant")] -use std::sync::LazyLock; use std::time::{SystemTime, UNIX_EPOCH}; use cocoindex::connectors::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; -use cocoindex::{ContextKey, Environment, Result}; +use cocoindex::{Environment, Result}; use serde_json::json; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("qdrant_test", |c: &QdrantConnection| { - c.state_id().to_string() - }) -}); +cocoindex::context_key!( + static DB: QdrantConnection = "qdrant_test", + state = QdrantConnection::state_id +); type Point = (u64, Vec, &'static str, &'static str); diff --git a/rust/sdk/cocoindex/tests/sqlite_target.rs b/rust/sdk/cocoindex/tests/sqlite_target.rs index d47fb68b1..c78054450 100644 --- a/rust/sdk/cocoindex/tests/sqlite_target.rs +++ b/rust/sdk/cocoindex/tests/sqlite_target.rs @@ -6,19 +6,16 @@ //! user-managed, multiple tables, and the declare (pending) variant. #![cfg(feature = "sqlite")] -use std::sync::LazyLock; - use cocoindex::connectors::sqlite; -use cocoindex::{App, ContextKey, Ctx, Environment, Result}; +use cocoindex::{App, Ctx, Environment, Result}; use serde::Serialize; use serde_json::json; use sqlx::Row as _; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("sqlite_target_test_db", |db: &sqlite::Database| { - db.state_id().to_string() - }) -}); +cocoindex::context_key!( + static DB: sqlite::Database = "sqlite_target_test_db", + state = sqlite::Database::state_id +); #[derive(Clone, Serialize)] struct Item { diff --git a/rust/sdk/cocoindex/tests/surrealdb_target.rs b/rust/sdk/cocoindex/tests/surrealdb_target.rs index 2625d5e57..d5a67d8bf 100644 --- a/rust/sdk/cocoindex/tests/surrealdb_target.rs +++ b/rust/sdk/cocoindex/tests/surrealdb_target.rs @@ -1,12 +1,9 @@ #![cfg(feature = "surrealdb")] -use std::sync::LazyLock; - use cocoindex::connectors::surrealdb::{self, ColumnDef, Graph, TableSchema, VectorIndexOptions}; use cocoindex::prelude::*; -static GRAPH: LazyLock> = - LazyLock::new(|| ContextKey::new("surrealdb_smoke_graph")); +cocoindex::context_key!(static GRAPH: Graph = "surrealdb_smoke_graph"); async fn try_graph(db_name: &str) -> Option { let url = std::env::var("SURREALDB_URL").unwrap_or_else(|_| "127.0.0.1:8787".to_string()); diff --git a/rust/sdk/cocoindex/tests/turbopuffer_target.rs b/rust/sdk/cocoindex/tests/turbopuffer_target.rs index b92cb618f..d7c48ce64 100644 --- a/rust/sdk/cocoindex/tests/turbopuffer_target.rs +++ b/rust/sdk/cocoindex/tests/turbopuffer_target.rs @@ -8,20 +8,18 @@ //! cargo test -p cocoindex --features turbopuffer --test turbopuffer_target #![cfg(feature = "turbopuffer")] -use std::sync::LazyLock; use std::time::{SystemTime, UNIX_EPOCH}; use cocoindex::connectors::turbopuffer::{ self, DistanceMetric, NamespaceSchema, TurbopufferConnection, }; -use cocoindex::{ContextKey, Environment, Result}; +use cocoindex::{Environment, Result}; use serde_json::json; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("turbopuffer_test", |c: &TurbopufferConnection| { - c.state_id().to_string() - }) -}); +cocoindex::context_key!( + static DB: TurbopufferConnection = "turbopuffer_test", + state = TurbopufferConnection::state_id +); type RowSpec = (&'static str, Vec, &'static str); diff --git a/rust/sdk/cocoindex/tests/valkey_target.rs b/rust/sdk/cocoindex/tests/valkey_target.rs index 8459ae135..76f569fe2 100644 --- a/rust/sdk/cocoindex/tests/valkey_target.rs +++ b/rust/sdk/cocoindex/tests/valkey_target.rs @@ -4,17 +4,15 @@ #![cfg(feature = "valkey")] use std::collections::BTreeMap; -use std::sync::LazyLock; use std::time::{SystemTime, UNIX_EPOCH}; +use cocoindex::Environment; use cocoindex::connectors::valkey::{ self, Distance, Document, FieldDef, FieldType, IndexSchema, VectorAlgorithm, VectorDef, }; use cocoindex::resources::schema::VectorSchema; -use cocoindex::{ContextKey, Environment}; -static VK: LazyLock> = - LazyLock::new(|| ContextKey::new("valkey_target_conn")); +cocoindex::context_key!(static VK: valkey::Valkey = "valkey_target_conn"); fn nonce() -> u128 { SystemTime::now() From bede11f110cc92ac086c3adc34dc9ea602cc1bfe Mon Sep 17 00:00:00 2001 From: zhhanz Date: Tue, 14 Jul 2026 15:57:21 -0700 Subject: [PATCH 04/15] Clarify Rust memoization guidance --- rust/sdk/cocoindex/src/memo.rs | 13 +++++++ rust/sdk/cocoindex/tests/pipeline.rs | 56 ++++++++++++++++++---------- rust/sdk/cocoindex_macros/src/lib.rs | 21 +++++++++++ 3 files changed, 71 insertions(+), 19 deletions(-) diff --git a/rust/sdk/cocoindex/src/memo.rs b/rust/sdk/cocoindex/src/memo.rs index 6a61c8cc1..832fa3be7 100644 --- a/rust/sdk/cocoindex/src/memo.rs +++ b/rust/sdk/cocoindex/src/memo.rs @@ -1,4 +1,17 @@ //! Memoization: skip re-execution when inputs haven't changed. +//! +//! Prefer `#[cocoindex::function(memo)]` for memoizing a whole pipeline +//! function. The attribute builds the argument key, tracks the function body's +//! logic, and passes an owned, memo-scoped [`Ctx`] into the body. Context +//! resources can therefore be read normally with [`Ctx::get_or_err`] or +//! [`Ctx::get_key`]. Use `memo_key(parameter = skip)` for an `Any + Clone` +//! parameter that should not participate in the key; it does not need to +//! implement Serde. +//! +//! [`Ctx::memo`] and [`cached`] are intended for block-level memoization inside +//! a function. Their closure body is not logic-tracked. When a manual memo block +//! must invalidate after its enclosing `#[cocoindex::function]` changes, include +//! that function's generated `__COCO_FN_HASH_` constant in the manual key. use std::any::Any; use std::future::Future; diff --git a/rust/sdk/cocoindex/tests/pipeline.rs b/rust/sdk/cocoindex/tests/pipeline.rs index 08393eb22..ea209b0ab 100644 --- a/rust/sdk/cocoindex/tests/pipeline.rs +++ b/rust/sdk/cocoindex/tests/pipeline.rs @@ -2852,18 +2852,20 @@ mod mock_bare { mod mock_memo { use super::*; - /// Memo function: cached by args. The body is a `'static` closure, - /// so `ctx` is NOT available inside. Clone the API before the body, - /// then use bare `#[cocoindex::function]` + manual `ctx.memo()`. - /// - /// This is the realistic pattern for memoizing an API call. - #[cocoindex::function] - async fn analyze(ctx: &cocoindex::Ctx, input: &str) -> cocoindex::Result { - let api = ctx.get_or_err::().unwrap().clone(); - let key = (__COCO_FN_HASH_ANALYZE, input.to_owned()); - let input = input.to_owned(); - ctx.memo(&key, move |_ctx| async move { api.call(&input).await }) - .await + /// Whole-function memoization should use the attribute. Its body receives + /// the owned memo-scoped `Ctx`, so non-serializable resources remain + /// available through normal context lookup. A skipped argument needs only + /// `Any + Clone`; `MockApi` intentionally does not implement `Serialize`. + #[cocoindex::function(memo, memo_key(client = skip))] + async fn analyze( + ctx: &cocoindex::Ctx, + input: String, + client: &MockApi, + ) -> cocoindex::Result { + // Keep the skipped parameter usable in the body while sourcing the + // configured client from the memo-scoped context. + let _observed_client_calls = client.call_count(); + ctx.get_or_err::()?.call(&input).await } #[tokio::test] @@ -2881,8 +2883,9 @@ mod mock_memo { .unwrap(); // First run — cache miss, API called. - app.update(|ctx| async move { - let result = analyze(&ctx, "hello").await?; + let client = api.clone(); + app.update(move |ctx| async move { + let result = analyze(&ctx, "hello".to_string(), &client).await?; assert_eq!(result, "api:hello"); Ok(()) }) @@ -2891,8 +2894,9 @@ mod mock_memo { assert_eq!(api.call_count(), 1); // Second run — same input, cache hit, API NOT called. - app.update(|ctx| async move { - let result = analyze(&ctx, "hello").await?; + let client = api.clone(); + app.update(move |ctx| async move { + let result = analyze(&ctx, "hello".to_string(), &client).await?; assert_eq!(result, "api:hello"); Ok(()) }) @@ -2900,9 +2904,23 @@ mod mock_memo { .unwrap(); assert_eq!(api.call_count(), 1, "memo should return cached result"); - // Third run — different input, cache miss, API called again. - app.update(|ctx| async move { - let result = analyze(&ctx, "world").await?; + // A distinct non-serializable skipped client does not alter the key. + let skipped_client = MockApi::new(); + let client = skipped_client.clone(); + app.update(move |ctx| async move { + let result = analyze(&ctx, "hello".to_string(), &client).await?; + assert_eq!(result, "api:hello"); + Ok(()) + }) + .await + .unwrap(); + assert_eq!(api.call_count(), 1, "skipped client should not invalidate"); + assert_eq!(skipped_client.call_count(), 0); + + // A keyed argument change still misses the cache. + let client = api.clone(); + app.update(move |ctx| async move { + let result = analyze(&ctx, "world".to_string(), &client).await?; assert_eq!(result, "api:world"); Ok(()) }) diff --git a/rust/sdk/cocoindex_macros/src/lib.rs b/rust/sdk/cocoindex_macros/src/lib.rs index 6fae65497..41c29dfac 100644 --- a/rust/sdk/cocoindex_macros/src/lib.rs +++ b/rust/sdk/cocoindex_macros/src/lib.rs @@ -798,6 +798,27 @@ fn expand_batching_function( /// #[cocoindex::function(memo)] /// async fn my_fn(ctx: &Ctx, arg: &String) -> Result { ... } /// ``` +/// The body receives a memo-scoped `Ctx`, so it can call `ctx.get_or_err()` or +/// `ctx.get_key()` directly. The attribute is the default way to memoize a +/// whole function because it includes the function's logic hash and arguments +/// in the key automatically. +/// +/// Customize individual argument keys with `memo_key(...)`. A transform +/// function replaces the argument's default representation; `skip` (or `None`) +/// excludes it entirely. A skipped parameter needs only `Any + Clone`, not +/// `Serialize`: +/// ```ignore +/// #[cocoindex::function(memo, memo_key(client = skip))] +/// async fn fetch(ctx: &Ctx, url: &str, client: &ApiClient) -> Result { +/// let configured_client = ctx.get_or_err::()?; +/// configured_client.get(url).await +/// } +/// ``` +/// +/// Use `ctx.memo(...)` only for block-level memoization. Its closure body is not +/// logic-tracked; a manual key that should follow edits to the enclosing +/// `#[cocoindex::function]` must include that function's generated +/// `__COCO_FN_HASH_` constant explicitly. /// /// **With `batching`** — declare a batch-shaped body and call it with one item: /// ```ignore From 1854dd5b37bfc8d8df6ebba7c511df65880c3af3 Mon Sep 17 00:00:00 2001 From: zhhanz Date: Tue, 14 Jul 2026 17:09:33 -0700 Subject: [PATCH 05/15] Derive Rust connector schemas from rows --- examples/rust/amazon_s3_embedding/src/main.rs | 17 +- examples/rust/audio_to_text/src/main.rs | 10 +- examples/rust/code_embedding/src/main.rs | 19 +-- .../rust/code_embedding_lancedb/src/main.rs | 18 +- examples/rust/files_to_sqlite/src/main.rs | 11 +- .../rust/gdrive_text_embedding/src/main.rs | 17 +- examples/rust/hn_trending_topics/src/main.rs | 28 +--- .../oci_object_storage_embedding/src/main.rs | 17 +- examples/rust/paper_metadata/src/main.rs | 44 ++--- examples/rust/pdf_embedding/src/main.rs | 19 +-- examples/rust/postgres_source/src/main.rs | 20 +-- examples/rust/text_embedding/src/main.rs | 34 ++-- .../rust/text_embedding_lancedb/src/main.rs | 18 +- rust/sdk/cocoindex/src/connectors/doris.rs | 26 +++ rust/sdk/cocoindex/src/connectors/lancedb.rs | 155 +++++++++++++++++- rust/sdk/cocoindex/src/connectors/postgres.rs | 37 ++++- rust/sdk/cocoindex/src/connectors/qdrant.rs | 80 +++++++++ rust/sdk/cocoindex/src/connectors/sqlite.rs | 21 +++ .../cocoindex/src/connectors/turbopuffer.rs | 69 ++++++++ rust/sdk/cocoindex/src/prelude.rs | 2 +- rust/sdk/cocoindex/src/row_schema.rs | 6 +- rust/sdk/cocoindex/tests/schema_from_row.rs | 120 +++++++++++++- rust/sdk/cocoindex_macros/src/lib.rs | 16 +- 23 files changed, 583 insertions(+), 221 deletions(-) diff --git a/examples/rust/amazon_s3_embedding/src/main.rs b/examples/rust/amazon_s3_embedding/src/main.rs index 6dd1f47d3..0d82a7cd3 100644 --- a/examples/rust/amazon_s3_embedding/src/main.rs +++ b/examples/rust/amazon_s3_embedding/src/main.rs @@ -43,11 +43,12 @@ static EMBEDDER: LazyLock> = LazyLock::n }) }); -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct DocEmbeddingRow { id: i64, filename: String, text: String, + #[coco(vector)] embedding: Vec, } @@ -167,18 +168,8 @@ async fn query_once( } fn doc_embedding_schema() -> Result { - postgres::TableSchema::new( - [ - ("id", postgres::ColumnDef::new("bigint")), - ("filename", postgres::ColumnDef::new("text")), - ("text", postgres::ColumnDef::new("text")), - ( - "embedding", - postgres::ColumnDef::new(format!("vector({EMBED_DIM})")), - ), - ], - ["id"], - ) + postgres::TableSchema::from_row::(["id"])? + .with_vector_dim("embedding", EMBED_DIM) } fn vector_param(vec: &[f32]) -> String { diff --git a/examples/rust/audio_to_text/src/main.rs b/examples/rust/audio_to_text/src/main.rs index b777fd673..2912f7226 100644 --- a/examples/rust/audio_to_text/src/main.rs +++ b/examples/rust/audio_to_text/src/main.rs @@ -43,20 +43,14 @@ static DB: LazyLock> = LazyLock::new(|| { }) }); -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct AudioTranscription { filename: String, text: String, } fn transcription_schema() -> Result { - postgres::TableSchema::new( - [ - ("filename", postgres::ColumnDef::new("text")), - ("text", postgres::ColumnDef::new("text")), - ], - ["filename"], - ) + postgres::TableSchema::from_row::(["filename"]) } /// Transcribe one audio file with OpenAI Whisper via the SDK's diff --git a/examples/rust/code_embedding/src/main.rs b/examples/rust/code_embedding/src/main.rs index a9a0fed76..78f558fbe 100644 --- a/examples/rust/code_embedding/src/main.rs +++ b/examples/rust/code_embedding/src/main.rs @@ -49,11 +49,12 @@ static EMBEDDER: LazyLock> = LazyLock::n }) }); -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct CodeEmbeddingRow { id: i64, filename: String, code: String, + #[coco(vector)] embedding: Vec, start_line: i32, end_line: i32, @@ -188,20 +189,8 @@ async fn query_once( // --------------------------------------------------------------------------- fn code_embedding_schema() -> Result { - postgres::TableSchema::new( - [ - ("id", postgres::ColumnDef::new("bigint")), - ("filename", postgres::ColumnDef::new("text")), - ("code", postgres::ColumnDef::new("text")), - ( - "embedding", - postgres::ColumnDef::new(format!("vector({EMBED_DIM})")), - ), - ("start_line", postgres::ColumnDef::new("integer")), - ("end_line", postgres::ColumnDef::new("integer")), - ], - ["id"], - ) + postgres::TableSchema::from_row::(["id"])? + .with_vector_dim("embedding", EMBED_DIM) } fn vector_param(vec: &[f32]) -> String { diff --git a/examples/rust/code_embedding_lancedb/src/main.rs b/examples/rust/code_embedding_lancedb/src/main.rs index 19f36bb97..ea7331cea 100644 --- a/examples/rust/code_embedding_lancedb/src/main.rs +++ b/examples/rust/code_embedding_lancedb/src/main.rs @@ -13,7 +13,7 @@ use std::path::PathBuf; use std::sync::LazyLock; -use cocoindex::connectors::lancedb::{self, ColumnDef, ColumnType, LanceDatabase, TableSchema}; +use cocoindex::connectors::lancedb::{self, LanceDatabase, TableSchema}; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter, detect_code_language}; use cocoindex::prelude::*; @@ -36,11 +36,12 @@ static EMBEDDER: LazyLock> = LazyLock::n }) }); -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct CodeEmbedding { id: i64, filename: String, code: String, + #[coco(vector)] embedding: Vec, start_line: i64, end_line: i64, @@ -88,17 +89,7 @@ async fn process_file(ctx: &Ctx, file: FileEntry) -> Result> } fn code_embedding_schema() -> Result { - TableSchema::new( - [ - ("id", ColumnDef::new(ColumnType::Int64)), - ("filename", ColumnDef::new(ColumnType::Text)), - ("code", ColumnDef::new(ColumnType::Text)), - ("embedding", ColumnDef::new(ColumnType::Vector(EMBED_DIM))), - ("start_line", ColumnDef::new(ColumnType::Int64)), - ("end_line", ColumnDef::new(ColumnType::Int64)), - ], - ["id"], - ) + TableSchema::from_row::(["id"])?.with_vector_dim("embedding", EMBED_DIM) } fn is_excluded(key: &str) -> bool { @@ -107,7 +98,6 @@ fn is_excluded(key: &str) -> bool { } async fn app_main(ctx: Ctx, sourcedir: PathBuf) -> Result<()> { - let db = ctx.get_key(&DB)?; let table = lancedb::mount_table_target(&ctx, &DB, TABLE, code_embedding_schema()?).await?; let files: Vec<(String, FileEntry)> = walk_items(&sourcedir, INCLUDE_PATTERNS)? diff --git a/examples/rust/files_to_sqlite/src/main.rs b/examples/rust/files_to_sqlite/src/main.rs index add015d15..207dd0ff6 100644 --- a/examples/rust/files_to_sqlite/src/main.rs +++ b/examples/rust/files_to_sqlite/src/main.rs @@ -29,7 +29,7 @@ static DB: LazyLock> = LazyLock::new(|| { const TABLE: &str = "files"; /// One output row: the file path is the primary key. -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct FileRow { path: String, word_count: i64, @@ -37,14 +37,7 @@ struct FileRow { } fn files_schema() -> Result { - sqlite::TableSchema::new( - [ - ("path", sqlite::ColumnDef::new("TEXT")), - ("word_count", sqlite::ColumnDef::new("INTEGER")), - ("first_line", sqlite::ColumnDef::new("TEXT")), - ], - ["path"], - ) + sqlite::TableSchema::from_row::(["path"]) } /// Summarize one file. Logic-tracked, so editing it invalidates cached files. diff --git a/examples/rust/gdrive_text_embedding/src/main.rs b/examples/rust/gdrive_text_embedding/src/main.rs index f0053dce0..747179f7f 100644 --- a/examples/rust/gdrive_text_embedding/src/main.rs +++ b/examples/rust/gdrive_text_embedding/src/main.rs @@ -40,11 +40,12 @@ static EMBEDDER: LazyLock> = LazyLock::n }) }); -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct DocEmbeddingRow { id: i64, filename: String, text: String, + #[coco(vector)] embedding: Vec, } @@ -152,18 +153,8 @@ async fn query_once( } fn doc_embedding_schema() -> Result { - postgres::TableSchema::new( - [ - ("id", postgres::ColumnDef::new("bigint")), - ("filename", postgres::ColumnDef::new("text")), - ("text", postgres::ColumnDef::new("text")), - ( - "embedding", - postgres::ColumnDef::new(format!("vector({EMBED_DIM})")), - ), - ], - ["id"], - ) + postgres::TableSchema::from_row::(["id"])? + .with_vector_dim("embedding", EMBED_DIM) } fn vector_param(vec: &[f32]) -> String { diff --git a/examples/rust/hn_trending_topics/src/main.rs b/examples/rust/hn_trending_topics/src/main.rs index 53274c8ba..5e87c9620 100644 --- a/examples/rust/hn_trending_topics/src/main.rs +++ b/examples/rust/hn_trending_topics/src/main.rs @@ -148,7 +148,7 @@ struct Thread { messages: Vec, // [0] is the thread itself, rest are comments } -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct HnMessageRow { id: String, thread_id: String, @@ -159,7 +159,7 @@ struct HnMessageRow { created_at: Option, } -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct HnTopicRow { topic: String, message_id: String, @@ -351,31 +351,11 @@ async fn app_main(ctx: Ctx, max_threads: usize) -> Result<()> { // --------------------------------------------------------------------------- fn message_schema() -> Result { - postgres::TableSchema::new( - [ - ("id", postgres::ColumnDef::new("text")), - ("thread_id", postgres::ColumnDef::new("text")), - ("content_type", postgres::ColumnDef::new("text")), - ("author", postgres::ColumnDef::new("text").nullable()), - ("text", postgres::ColumnDef::new("text").nullable()), - ("url", postgres::ColumnDef::new("text").nullable()), - ("created_at", postgres::ColumnDef::new("text").nullable()), - ], - ["id"], - ) + postgres::TableSchema::from_row::(["id"]) } fn topic_schema() -> Result { - postgres::TableSchema::new( - [ - ("topic", postgres::ColumnDef::new("text")), - ("message_id", postgres::ColumnDef::new("text")), - ("thread_id", postgres::ColumnDef::new("text")), - ("content_type", postgres::ColumnDef::new("text")), - ("created_at", postgres::ColumnDef::new("text").nullable()), - ], - ["topic", "message_id"], - ) + postgres::TableSchema::from_row::(["topic", "message_id"]) } async fn show_trending(pool: &PgPool, limit: i64) -> Result<()> { diff --git a/examples/rust/oci_object_storage_embedding/src/main.rs b/examples/rust/oci_object_storage_embedding/src/main.rs index 4f8c876c1..ebefb546f 100644 --- a/examples/rust/oci_object_storage_embedding/src/main.rs +++ b/examples/rust/oci_object_storage_embedding/src/main.rs @@ -44,11 +44,12 @@ static EMBEDDER: LazyLock> = LazyLock::n }) }); -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct DocEmbeddingRow { id: i64, filename: String, text: String, + #[coco(vector)] embedding: Vec, } @@ -169,18 +170,8 @@ async fn query_once( } fn doc_embedding_schema() -> Result { - postgres::TableSchema::new( - [ - ("id", postgres::ColumnDef::new("bigint")), - ("filename", postgres::ColumnDef::new("text")), - ("text", postgres::ColumnDef::new("text")), - ( - "embedding", - postgres::ColumnDef::new(format!("vector({EMBED_DIM})")), - ), - ], - ["id"], - ) + postgres::TableSchema::from_row::(["id"])? + .with_vector_dim("embedding", EMBED_DIM) } fn vector_param(vec: &[f32]) -> String { diff --git a/examples/rust/paper_metadata/src/main.rs b/examples/rust/paper_metadata/src/main.rs index 3d3b2c79e..ff283713a 100644 --- a/examples/rust/paper_metadata/src/main.rs +++ b/examples/rust/paper_metadata/src/main.rs @@ -153,28 +153,32 @@ struct PaperMetadataModel { abstract_text: String, } -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct PaperMetadataRow { filename: String, title: String, + #[coco(json)] authors: serde_json::Value, // jsonb: [{name, email, affiliation}] #[serde(rename = "abstract")] + #[coco(rename = "abstract")] abstract_text: String, num_pages: i32, } -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct AuthorPaperRow { author_name: String, filename: String, } -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct MetadataEmbeddingRow { + #[coco(type = "uuid")] id: String, // uuid filename: String, location: String, text: String, + #[coco(vector)] embedding: Vec, } @@ -323,42 +327,16 @@ async fn process_file(ctx: &Ctx, file: &FileEntry) -> Result { // --------------------------------------------------------------------------- fn metadata_schema() -> Result { - postgres::TableSchema::new( - [ - ("filename", postgres::ColumnDef::new("text")), - ("title", postgres::ColumnDef::new("text")), - ("authors", postgres::ColumnDef::new("jsonb")), - ("abstract", postgres::ColumnDef::new("text")), - ("num_pages", postgres::ColumnDef::new("integer")), - ], - ["filename"], - ) + postgres::TableSchema::from_row::(["filename"]) } fn author_schema() -> Result { - postgres::TableSchema::new( - [ - ("author_name", postgres::ColumnDef::new("text")), - ("filename", postgres::ColumnDef::new("text")), - ], - ["author_name", "filename"], - ) + postgres::TableSchema::from_row::(["author_name", "filename"]) } fn embedding_schema() -> Result { - postgres::TableSchema::new( - [ - ("id", postgres::ColumnDef::new("uuid")), - ("filename", postgres::ColumnDef::new("text")), - ("location", postgres::ColumnDef::new("text")), - ("text", postgres::ColumnDef::new("text")), - ( - "embedding", - postgres::ColumnDef::new(format!("vector({EMBED_DIM})")), - ), - ], - ["id"], - ) + postgres::TableSchema::from_row::(["id"])? + .with_vector_dim("embedding", EMBED_DIM) } // --------------------------------------------------------------------------- diff --git a/examples/rust/pdf_embedding/src/main.rs b/examples/rust/pdf_embedding/src/main.rs index 58772d172..f5672dec3 100644 --- a/examples/rust/pdf_embedding/src/main.rs +++ b/examples/rust/pdf_embedding/src/main.rs @@ -42,13 +42,14 @@ static EMBEDDER: LazyLock> = LazyLock::n }) }); -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct PdfEmbedding { id: i64, filename: String, chunk_start: i32, chunk_end: i32, text: String, + #[coco(vector)] embedding: Vec, } @@ -108,20 +109,8 @@ async fn process_file(ctx: &Ctx, file: FileEntry) -> Result> { } fn pdf_embedding_schema() -> Result { - postgres::TableSchema::new( - [ - ("id", postgres::ColumnDef::new("bigint")), - ("filename", postgres::ColumnDef::new("text")), - ("chunk_start", postgres::ColumnDef::new("integer")), - ("chunk_end", postgres::ColumnDef::new("integer")), - ("text", postgres::ColumnDef::new("text")), - ( - "embedding", - postgres::ColumnDef::new(format!("vector({EMBED_DIM})")), - ), - ], - ["id"], - ) + postgres::TableSchema::from_row::(["id"])? + .with_vector_dim("embedding", EMBED_DIM) } async fn app_main(ctx: Ctx, sourcedir: PathBuf) -> Result<()> { diff --git a/examples/rust/postgres_source/src/main.rs b/examples/rust/postgres_source/src/main.rs index d23db9b70..c33709c18 100644 --- a/examples/rust/postgres_source/src/main.rs +++ b/examples/rust/postgres_source/src/main.rs @@ -65,7 +65,7 @@ struct SourceProduct { } /// One row written to the output table. -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct OutputProduct { product_category: String, product_name: String, @@ -73,6 +73,7 @@ struct OutputProduct { price: f64, amount: i64, total_value: f64, + #[coco(vector)] embedding: Vec, } @@ -102,21 +103,8 @@ async fn process_product(ctx: &Ctx, product: SourceProduct) -> Result Result { - postgres::TableSchema::new( - [ - ("product_category", postgres::ColumnDef::new("text")), - ("product_name", postgres::ColumnDef::new("text")), - ("description", postgres::ColumnDef::new("text")), - ("price", postgres::ColumnDef::new("double precision")), - ("amount", postgres::ColumnDef::new("bigint")), - ("total_value", postgres::ColumnDef::new("double precision")), - ( - "embedding", - postgres::ColumnDef::new(format!("vector({EMBED_DIM})")), - ), - ], - ["product_category", "product_name"], - ) + postgres::TableSchema::from_row::(["product_category", "product_name"])? + .with_vector_dim("embedding", EMBED_DIM) } async fn app_main(ctx: Ctx) -> Result<()> { diff --git a/examples/rust/text_embedding/src/main.rs b/examples/rust/text_embedding/src/main.rs index 28e18031b..26e4e6496 100644 --- a/examples/rust/text_embedding/src/main.rs +++ b/examples/rust/text_embedding/src/main.rs @@ -22,7 +22,6 @@ use sqlx::Row; use sqlx::postgres::{PgPool, PgPoolOptions}; const EMBED_MODEL: &str = "sentence-transformers/all-MiniLM-L6-v2"; -const EMBED_DIM: usize = 384; const PG_SCHEMA: &str = "coco_examples"; const TABLE: &str = "doc_embeddings"; const TOP_K: i64 = 5; @@ -38,13 +37,14 @@ cocoindex::context_key!( state = SentenceTransformerEmbedder::model_name ); -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct DocEmbedding { id: i64, filename: String, chunk_start: i32, chunk_end: i32, text: String, + #[coco(vector)] embedding: Vec, } @@ -88,27 +88,21 @@ async fn process_file(ctx: &Ctx, file: FileEntry) -> Result> { Ok(rows) } -fn doc_embedding_schema() -> Result { - postgres::TableSchema::new( - [ - ("id", postgres::ColumnDef::new("bigint")), - ("filename", postgres::ColumnDef::new("text")), - ("chunk_start", postgres::ColumnDef::new("integer")), - ("chunk_end", postgres::ColumnDef::new("integer")), - ("text", postgres::ColumnDef::new("text")), - ( - "embedding", - postgres::ColumnDef::new(format!("vector({EMBED_DIM})")), - ), - ], - ["id"], - ) +fn doc_embedding_schema(vector_dim: usize) -> Result { + postgres::TableSchema::from_row::(["id"])? + .with_vector_dim("embedding", vector_dim) } async fn app_main(ctx: Ctx, sourcedir: PathBuf) -> Result<()> { - let table = - postgres::mount_table_target(&ctx, &DB, TABLE, doc_embedding_schema()?, Some(PG_SCHEMA)) - .await?; + let vector_dim = ctx.get_key(&EMBEDDER)?.dimension(); + let table = postgres::mount_table_target( + &ctx, + &DB, + TABLE, + doc_embedding_schema(vector_dim)?, + Some(PG_SCHEMA), + ) + .await?; table.declare_vector_index( &ctx, "embedding", diff --git a/examples/rust/text_embedding_lancedb/src/main.rs b/examples/rust/text_embedding_lancedb/src/main.rs index 0c784799f..58da20bc3 100644 --- a/examples/rust/text_embedding_lancedb/src/main.rs +++ b/examples/rust/text_embedding_lancedb/src/main.rs @@ -13,7 +13,7 @@ use std::path::PathBuf; use std::sync::LazyLock; -use cocoindex::connectors::lancedb::{self, ColumnDef, ColumnType, LanceDatabase, TableSchema}; +use cocoindex::connectors::lancedb::{self, LanceDatabase, TableSchema}; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; use cocoindex::prelude::*; @@ -36,13 +36,14 @@ static EMBEDDER: LazyLock> = LazyLock::n }) }); -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct DocEmbedding { id: i64, filename: String, chunk_start: i64, chunk_end: i64, text: String, + #[coco(vector)] embedding: Vec, } @@ -87,21 +88,10 @@ async fn process_file(ctx: &Ctx, file: FileEntry) -> Result> { } fn doc_embedding_schema() -> Result { - TableSchema::new( - [ - ("id", ColumnDef::new(ColumnType::Int64)), - ("filename", ColumnDef::new(ColumnType::Text)), - ("chunk_start", ColumnDef::new(ColumnType::Int64)), - ("chunk_end", ColumnDef::new(ColumnType::Int64)), - ("text", ColumnDef::new(ColumnType::Text)), - ("embedding", ColumnDef::new(ColumnType::Vector(EMBED_DIM))), - ], - ["id"], - ) + TableSchema::from_row::(["id"])?.with_vector_dim("embedding", EMBED_DIM) } async fn app_main(ctx: Ctx, sourcedir: PathBuf) -> Result<()> { - let db = ctx.get_key(&DB)?; let table = lancedb::mount_table_target(&ctx, &DB, TABLE, doc_embedding_schema()?).await?; let files = walk_items(&sourcedir, &["**/*.md"])?; diff --git a/rust/sdk/cocoindex/src/connectors/doris.rs b/rust/sdk/cocoindex/src/connectors/doris.rs index 7dd7ede1e..854a72367 100644 --- a/rust/sdk/cocoindex/src/connectors/doris.rs +++ b/rust/sdk/cocoindex/src/connectors/doris.rs @@ -342,6 +342,32 @@ impl TableSchema { .map(|f| (f.name.clone(), doris_column_def(&f))); Self::new(columns, primary_key) } + + /// Resolve or override the dimension of a vector field derived from a row. + pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { + let dim = u32::try_from(dim).map_err(|_| { + Error::engine(format!( + "Doris vector field {field_name:?} dimension does not fit u32" + )) + })?; + if dim == 0 { + return Err(Error::engine(format!( + "Doris vector field {field_name:?} requires a dimension greater than zero" + ))); + } + let def = self.columns.get_mut(field_name).ok_or_else(|| { + Error::engine(format!( + "Doris vector dimension override names unknown field {field_name:?}" + )) + })?; + if !def.is_vector { + return Err(Error::engine(format!( + "Doris field {field_name:?} is not a vector field" + ))); + } + def.vector_dimension = Some(dim); + Ok(self) + } } /// Map a connector-agnostic [`SchemaField`](crate::row_schema::SchemaField) to a diff --git a/rust/sdk/cocoindex/src/connectors/lancedb.rs b/rust/sdk/cocoindex/src/connectors/lancedb.rs index 7b75e0236..4d52fbe47 100644 --- a/rust/sdk/cocoindex/src/connectors/lancedb.rs +++ b/rust/sdk/cocoindex/src/connectors/lancedb.rs @@ -15,7 +15,8 @@ use std::sync::Arc; use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; use arrow_array::{ - Array, ArrayRef, Float64Array, Int64Array, RecordBatch, RecordBatchIterator, StringArray, + Array, ArrayRef, BinaryArray, BooleanArray, Float32Array, Float64Array, Int16Array, Int32Array, + Int64Array, RecordBatch, RecordBatchIterator, StringArray, }; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use cocoindex_utils::fingerprint::Fingerprint; @@ -80,9 +81,16 @@ impl LanceDatabase { /// A LanceDB column type (the subset CocoIndex maps natively). #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] pub enum ColumnType { + Bool, + Int16, + Int32, Int64, + Float32, Float64, Text, + Binary, + /// Complex values encoded as JSON strings, matching Python's fallback. + Json, /// Fixed-size float32 vector of the given dimension. Vector(usize), } @@ -90,9 +98,14 @@ pub enum ColumnType { impl ColumnType { fn arrow_data_type(&self) -> DataType { match self { + ColumnType::Bool => DataType::Boolean, + ColumnType::Int16 => DataType::Int16, + ColumnType::Int32 => DataType::Int32, ColumnType::Int64 => DataType::Int64, + ColumnType::Float32 => DataType::Float32, ColumnType::Float64 => DataType::Float64, - ColumnType::Text => DataType::Utf8, + ColumnType::Text | ColumnType::Json => DataType::Utf8, + ColumnType::Binary => DataType::Binary, ColumnType::Vector(dim) => DataType::FixedSizeList( Arc::new(Field::new("item", DataType::Float32, true)), *dim as i32, @@ -167,6 +180,49 @@ impl TableSchema { &self.primary_key } + pub fn columns(&self) -> &[(String, ColumnDef)] { + &self.columns + } + + /// Derive a schema from a `#[derive(SchemaFields)]` row type. + pub fn from_row( + primary_key: impl IntoIterator>, + ) -> Result { + let columns = T::schema_fields() + .into_iter() + .map(|field| { + let def = lancedb_column_def(&field)?; + Ok((field.name, def)) + }) + .collect::>>()?; + Self::new(columns, primary_key) + } + + /// Resolve or override the dimension of a vector field derived from a row. + pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { + if dim == 0 || i32::try_from(dim).is_err() { + return Err(Error::engine(format!( + "LanceDB vector field {field_name:?} requires a dimension in 1..=i32::MAX" + ))); + } + let (_, def) = self + .columns + .iter_mut() + .find(|(name, _)| name == field_name) + .ok_or_else(|| { + Error::engine(format!( + "LanceDB vector dimension override names unknown field {field_name:?}" + )) + })?; + if !matches!(def.col_type, ColumnType::Vector(_)) { + return Err(Error::engine(format!( + "LanceDB field {field_name:?} is not a vector field" + ))); + } + def.col_type = ColumnType::Vector(dim); + Ok(self) + } + fn column_names(&self) -> impl Iterator { self.columns.iter().map(|(n, _)| n) } @@ -181,6 +237,31 @@ impl TableSchema { } } +fn lancedb_column_def(field: &crate::row_schema::SchemaField) -> Result { + use crate::row_schema::LogicalType as L; + + let col_type = match &field.logical_type { + L::Bool => ColumnType::Bool, + L::Int16 => ColumnType::Int16, + L::Int32 => ColumnType::Int32, + L::Int64 => ColumnType::Int64, + L::Float32 => ColumnType::Float32, + L::Float64 => ColumnType::Float64, + L::Text | L::Uuid | L::Date | L::Time | L::DateTime => ColumnType::Text, + L::Bytes => ColumnType::Binary, + L::Decimal | L::Duration | L::Json => ColumnType::Json, + L::Vector { dim, .. } => ColumnType::Vector(*dim as usize), + L::Custom(custom) => { + return Err(Error::engine(format!( + "LanceDB field {:?} has unsupported custom logical type {custom:?}", + field.name + ))); + } + }; + let def = ColumnDef::new(col_type); + Ok(if field.nullable { def.nullable() } else { def }) +} + // --------------------------------------------------------------------------- // Public target API: constructor / declaration / mount split // --------------------------------------------------------------------------- @@ -1000,11 +1081,37 @@ fn build_record_batch( for (name, def) in &schema.columns { let values = rows.iter().map(|r| r.get(name).unwrap_or(&JsonValue::Null)); let array: ArrayRef = match &def.col_type { + ColumnType::Bool => Arc::new(BooleanArray::from( + values + .map(|v| nullable_value(name, def, v).map(JsonValue::as_bool)) + .collect::>>()?, + )), + ColumnType::Int16 => Arc::new(Int16Array::from( + values + .map(|v| { + nullable_value(name, def, v) + .map(|v| v.as_i64().and_then(|n| i16::try_from(n).ok())) + }) + .collect::>>()?, + )), + ColumnType::Int32 => Arc::new(Int32Array::from( + values + .map(|v| { + nullable_value(name, def, v) + .map(|v| v.as_i64().and_then(|n| i32::try_from(n).ok())) + }) + .collect::>>()?, + )), ColumnType::Int64 => Arc::new(Int64Array::from( values .map(|v| nullable_value(name, def, v).map(|v| v.as_i64())) .collect::>>()?, )), + ColumnType::Float32 => Arc::new(Float32Array::from( + values + .map(|v| nullable_value(name, def, v).map(|v| v.as_f64().map(|n| n as f32))) + .collect::>>()?, + )), ColumnType::Float64 => Arc::new(Float64Array::from( values .map(|v| nullable_value(name, def, v).map(|v| v.as_f64())) @@ -1015,6 +1122,50 @@ fn build_record_batch( .map(|v| nullable_value(name, def, v).map(|v| v.as_str().map(str::to_string))) .collect::>>()?, )), + ColumnType::Binary => { + let bytes = values + .map(|v| { + let value = nullable_value(name, def, v)?; + if value.is_null() { + return Ok(None); + } + let values = value.as_array().ok_or_else(|| { + Error::engine(format!("column {name:?} must be a byte array")) + })?; + values + .iter() + .map(|value| { + value + .as_u64() + .and_then(|n| u8::try_from(n).ok()) + .ok_or_else(|| { + Error::engine(format!( + "column {name:?} has a non-byte element" + )) + }) + }) + .collect::>>() + .map(Some) + }) + .collect::>>()?; + Arc::new(BinaryArray::from_iter( + bytes.iter().map(|value| value.as_deref()), + )) + } + ColumnType::Json => Arc::new(StringArray::from( + values + .map(|v| { + let value = nullable_value(name, def, v)?; + if value.is_null() { + Ok(None) + } else { + serde_json::to_string(value).map(Some).map_err(|e| { + Error::engine(format!("encode LanceDB JSON field {name:?}: {e}")) + }) + } + }) + .collect::>>()?, + )), ColumnType::Vector(dim) => build_vector_array(name, *dim, def.nullable, values)?, }; arrays.push(array); diff --git a/rust/sdk/cocoindex/src/connectors/postgres.rs b/rust/sdk/cocoindex/src/connectors/postgres.rs index f45b6bcd2..a9c5cfc41 100644 --- a/rust/sdk/cocoindex/src/connectors/postgres.rs +++ b/rust/sdk/cocoindex/src/connectors/postgres.rs @@ -142,6 +142,36 @@ impl TableSchema { .map(|f| (f.name.clone(), postgres_column_def(&f))); Self::new(columns, primary_key) } + + /// Resolve or override the dimension of a vector field derived from a row. + pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { + if dim == 0 { + return Err(Error::engine(format!( + "Postgres vector field {field_name:?} requires a dimension greater than zero" + ))); + } + let def = self.columns.get_mut(field_name).ok_or_else(|| { + Error::engine(format!( + "Postgres vector dimension override names unknown field {field_name:?}" + )) + })?; + let base = pgvector_type_base(&def.pg_type).ok_or_else(|| { + Error::engine(format!( + "Postgres field {field_name:?} is not a vector field" + )) + })?; + def.pg_type = format!("{base}({dim})"); + Ok(self) + } +} + +fn pgvector_type_base(pg_type: &str) -> Option<&'static str> { + ["vector", "halfvec"].into_iter().find(|base| { + pg_type == *base + || pg_type + .strip_prefix(*base) + .is_some_and(|suffix| suffix.starts_with('(') && suffix.ends_with(')')) + }) } /// Map a connector-agnostic [`SchemaField`](crate::row_schema::SchemaField) to a @@ -165,10 +195,11 @@ fn postgres_column_def(field: &crate::row_schema::SchemaField) -> ColumnDef { L::Duration => "interval".to_string(), L::Json => "jsonb".to_string(), L::Vector { dim, half } => { - if *half { - format!("halfvec({dim})") + let base = if *half { "halfvec" } else { "vector" }; + if *dim == 0 { + base.to_string() } else { - format!("vector({dim})") + format!("{base}({dim})") } } L::Custom(s) => s.clone(), diff --git a/rust/sdk/cocoindex/src/connectors/qdrant.rs b/rust/sdk/cocoindex/src/connectors/qdrant.rs index aae8d490d..600a774e2 100644 --- a/rust/sdk/cocoindex/src/connectors/qdrant.rs +++ b/rust/sdk/cocoindex/src/connectors/qdrant.rs @@ -317,6 +317,86 @@ impl CollectionSchema { }) } + /// Derive named vector definitions from a `#[derive(SchemaFields)]` row. + /// Non-vector fields remain schemaless Qdrant payload fields, matching the + /// Python connector. Bare `#[coco(vector)]` fields are resolved later with + /// [`CollectionSchema::with_vector_dim`]. + pub fn from_row(distance: Distance) -> Result { + use crate::row_schema::LogicalType; + + let mut vectors = BTreeMap::new(); + for field in T::schema_fields() { + let LogicalType::Vector { dim, half } = field.logical_type else { + continue; + }; + let size = dim as usize; + if size > 0 { + validate_vector_size(size).map_err(|err| { + Error::engine(format!( + "Qdrant vector field {:?} is invalid: {err}", + field.name + )) + })?; + } + let schema = if half { + VectorSchema::f16(size) + } else { + VectorSchema::f32(size) + }; + vectors.insert( + field.name, + QdrantVectorDef { + schema: QdrantVectorSchema::Dense(schema), + distance, + multivector_comparator: MultivectorComparator::MaxSim, + }, + ); + } + if vectors.is_empty() { + return Err(Error::engine( + "Qdrant row schema does not contain a vector field", + )); + } + let first = vectors.values().next().expect("non-empty"); + Ok(Self { + vector_size: first.vector_size(), + distance: first.distance, + multivector: false, + vectors: Some(VectorFields::Named(vectors)), + }) + } + + /// Resolve or override one named vector field's runtime dimension. + pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { + validate_vector_size(dim).map_err(|err| { + Error::engine(format!( + "Qdrant vector field {field_name:?} has invalid dimension: {err}" + )) + })?; + let mut fields = self.vector_fields(); + let def = match &mut fields { + VectorFields::Named(vectors) => vectors.get_mut(field_name), + VectorFields::Single(def) if field_name == "vector" => Some(def), + VectorFields::Single(_) => None, + } + .ok_or_else(|| { + Error::engine(format!( + "Qdrant vector dimension override names unknown vector field {field_name:?}" + )) + })?; + match &mut def.schema { + QdrantVectorSchema::Dense(schema) => schema.size = dim, + QdrantVectorSchema::Multi(schema) => schema.vector_schema.size = dim, + } + self.vectors = Some(fields); + let first = match self.vectors.as_ref().expect("set above") { + VectorFields::Single(def) => def, + VectorFields::Named(vectors) => vectors.values().next().expect("non-empty"), + }; + self.vector_size = first.vector_size(); + Ok(self) + } + fn vectors_config(&self) -> VectorsConfigBuilder { let mut builder = VectorsConfigBuilder::default(); match self.vector_fields() { diff --git a/rust/sdk/cocoindex/src/connectors/sqlite.rs b/rust/sdk/cocoindex/src/connectors/sqlite.rs index db3a8c6a9..2a8afc014 100644 --- a/rust/sdk/cocoindex/src/connectors/sqlite.rs +++ b/rust/sdk/cocoindex/src/connectors/sqlite.rs @@ -165,6 +165,27 @@ impl TableSchema { .map(|f| (f.name.clone(), sqlite_column_def(&f))); Self::new(columns, primary_key) } + + /// Resolve or override the dimension of a vector field derived from a row. + pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { + if dim == 0 { + return Err(Error::engine(format!( + "SQLite vector field {field_name:?} requires a dimension greater than zero" + ))); + } + let def = self.columns.get_mut(field_name).ok_or_else(|| { + Error::engine(format!( + "SQLite vector dimension override names unknown field {field_name:?}" + )) + })?; + if !(def.sqlite_type.starts_with("float[") && def.sqlite_type.ends_with(']')) { + return Err(Error::engine(format!( + "SQLite field {field_name:?} is not a vector field" + ))); + } + def.sqlite_type = format!("float[{dim}]"); + Ok(self) + } } /// Map a connector-agnostic [`SchemaField`](crate::row_schema::SchemaField) to a diff --git a/rust/sdk/cocoindex/src/connectors/turbopuffer.rs b/rust/sdk/cocoindex/src/connectors/turbopuffer.rs index f0386bd21..911bec213 100644 --- a/rust/sdk/cocoindex/src/connectors/turbopuffer.rs +++ b/rust/sdk/cocoindex/src/connectors/turbopuffer.rs @@ -260,6 +260,75 @@ impl NamespaceSchema { }) } + /// Derive named vector definitions from a `#[derive(SchemaFields)]` row. + /// Non-vector fields remain schemaless attributes inferred by Turbopuffer, + /// matching the Python connector. Bare `#[coco(vector)]` fields are + /// resolved later with [`NamespaceSchema::with_vector_dim`]. + pub fn from_row(distance: DistanceMetric) -> Result { + use crate::row_schema::LogicalType; + + let mut vectors = BTreeMap::new(); + for field in T::schema_fields() { + let LogicalType::Vector { dim, half } = field.logical_type else { + continue; + }; + validate_vector_field_name(&field.name)?; + let schema = if half { + VectorSchema::f16(dim as usize) + } else { + VectorSchema::f32(dim as usize) + }; + if schema.size > 0 { + vector_type_str(&schema).map_err(|err| { + Error::engine(format!( + "Turbopuffer vector field {:?} is invalid: {err}", + field.name + )) + })?; + } + vectors.insert(field.name, VectorDef { schema }); + } + if vectors.is_empty() { + return Err(Error::engine( + "Turbopuffer row schema does not contain a vector field", + )); + } + let vector_size = vectors.values().next().expect("non-empty").schema.size; + Ok(Self { + vectors: Some(VectorFields::Named(vectors)), + vector_size, + distance, + }) + } + + /// Resolve or override one named vector field's runtime dimension. + pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { + if dim == 0 { + return Err(Error::engine(format!( + "Turbopuffer vector field {field_name:?} requires a dimension greater than zero" + ))); + } + let mut fields = self.vector_fields(); + let def = match &mut fields { + VectorFields::Named(vectors) => vectors.get_mut(field_name), + VectorFields::Single(def) if field_name == DEFAULT_VECTOR_FIELD => Some(def), + VectorFields::Single(_) => None, + } + .ok_or_else(|| { + Error::engine(format!( + "Turbopuffer vector dimension override names unknown vector field {field_name:?}" + )) + })?; + def.schema.size = dim; + vector_type_str(&def.schema)?; + self.vectors = Some(fields); + self.vector_size = match self.vectors.as_ref().expect("set above") { + VectorFields::Single(def) => def.schema.size, + VectorFields::Named(vectors) => vectors.values().next().expect("non-empty").schema.size, + }; + Ok(self) + } + /// The `schema` payload Turbopuffer's write API expects. fn write_schema(&self) -> Result { let mut fields = Map::new(); diff --git a/rust/sdk/cocoindex/src/prelude.rs b/rust/sdk/cocoindex/src/prelude.rs index 029ff9efa..e832b37c0 100644 --- a/rust/sdk/cocoindex/src/prelude.rs +++ b/rust/sdk/cocoindex/src/prelude.rs @@ -45,7 +45,7 @@ pub use crate::target_state::{ pub use crate::user_state::{IntoStateKey, StateHandle}; pub use crate::{ App, ContextKey, DropHandle, Environment, EnvironmentBuilder, PreviewAction, PreviewValue, - Progress, StatsGroupHandle, StatsGroupOptions, UpdateHandle, UpdateOptions, + Progress, SchemaFields, StatsGroupHandle, StatsGroupOptions, UpdateHandle, UpdateOptions, }; pub use crate::{function, mount_each, use_mount}; diff --git a/rust/sdk/cocoindex/src/row_schema.rs b/rust/sdk/cocoindex/src/row_schema.rs index d3663ebed..20f159ce0 100644 --- a/rust/sdk/cocoindex/src/row_schema.rs +++ b/rust/sdk/cocoindex/src/row_schema.rs @@ -21,6 +21,8 @@ //! ``` //! //! Field attributes: +//! * `#[coco(vector)]` — a dense `f32` vector whose dimension is supplied with +//! the connector schema's `with_vector_dim` builder. //! * `#[coco(vector = N)]` — a dense `f32` vector column of dimension `N`. //! * `#[coco(vector = N, half)]` — a 16-bit (half-precision) vector column. //! * `#[coco(type = "…")]` — a raw connector SQL type, used verbatim (the escape @@ -48,7 +50,9 @@ pub enum LogicalType { Duration, /// A complex value (collection / map / nested struct / `Any`) stored as JSON. Json, - /// A dense float vector of fixed dimension (`half` → 16-bit element type). + /// A dense float vector (`half` → 16-bit element type). A dimension of zero + /// is the unresolved form emitted by `#[coco(vector)]`; connector schemas + /// resolve it through `with_vector_dim` before use. Vector { dim: u32, half: bool, diff --git a/rust/sdk/cocoindex/tests/schema_from_row.rs b/rust/sdk/cocoindex/tests/schema_from_row.rs index 88e28d9b7..5eca3acab 100644 --- a/rust/sdk/cocoindex/tests/schema_from_row.rs +++ b/rust/sdk/cocoindex/tests/schema_from_row.rs @@ -56,11 +56,14 @@ fn sqlite_from_row_matches_explicit_schema() { name: Option, score: f64, blob: Vec, - #[coco(vector = 3)] + #[coco(vector)] embedding: Vec, } - let got = TableSchema::from_row::(["id"]).unwrap(); + let got = TableSchema::from_row::(["id"]) + .unwrap() + .with_vector_dim("embedding", 3) + .unwrap(); let want = TableSchema::new( [ ("id", ColumnDef::new("INTEGER").not_null()), @@ -88,7 +91,7 @@ fn postgres_from_row_matches_explicit_schema() { title: Option, views: i32, ratio: f32, - #[coco(vector = 8)] + #[coco(vector)] embedding: Vec, #[coco(vector = 8, half)] embedding_half: Vec, @@ -96,7 +99,10 @@ fn postgres_from_row_matches_explicit_schema() { meta: Vec, } - let got = TableSchema::from_row::(["id"]).unwrap(); + let got = TableSchema::from_row::(["id"]) + .unwrap() + .with_vector_dim("embedding", 8) + .unwrap(); // Postgres ColumnDef::new is NOT NULL by default; `.nullable()` opts in. let want = TableSchema::new( [ @@ -114,6 +120,112 @@ fn postgres_from_row_matches_explicit_schema() { assert_eq!(got, want); } +#[cfg(feature = "lancedb")] +#[test] +fn lancedb_from_row_matches_explicit_schema() { + use cocoindex::SchemaFields; + use cocoindex::connectors::lancedb::{ColumnDef, ColumnType, TableSchema}; + + #[derive(SchemaFields)] + #[allow(dead_code)] + struct Row { + id: i64, + count: i32, + ratio: f32, + active: bool, + title: Option, + #[coco(json)] + tags: Vec, + #[coco(vector)] + embedding: Vec, + } + + let got = TableSchema::from_row::(["id"]) + .unwrap() + .with_vector_dim("embedding", 4) + .unwrap(); + let want = TableSchema::new( + [ + ("id", ColumnDef::new(ColumnType::Int64)), + ("count", ColumnDef::new(ColumnType::Int32)), + ("ratio", ColumnDef::new(ColumnType::Float32)), + ("active", ColumnDef::new(ColumnType::Bool)), + ("title", ColumnDef::new(ColumnType::Text).nullable()), + ("tags", ColumnDef::new(ColumnType::Json)), + ("embedding", ColumnDef::new(ColumnType::Vector(4))), + ], + ["id"], + ) + .unwrap(); + assert_eq!(got, want); + + #[derive(SchemaFields)] + #[allow(dead_code)] + struct Unsupported { + id: i64, + #[coco(type = "struct")] + details: String, + } + let err = TableSchema::from_row::(["id"]).unwrap_err(); + let message = err.to_string(); + assert!(message.contains("LanceDB"), "{message}"); + assert!(message.contains("details"), "{message}"); +} + +#[cfg(feature = "qdrant")] +#[test] +fn qdrant_from_row_derives_named_vectors_and_runtime_dimension() { + use cocoindex::SchemaFields; + use cocoindex::connectors::qdrant::{CollectionSchema, Distance, QdrantVectorDef}; + + #[derive(SchemaFields)] + #[allow(dead_code)] + struct Row { + id: u64, + text: String, + #[coco(vector)] + embedding: Vec, + } + + let got = CollectionSchema::from_row::(Distance::Cosine) + .unwrap() + .with_vector_dim("embedding", 3) + .unwrap(); + let want = CollectionSchema::named([( + "embedding", + QdrantVectorDef::f32(3, Distance::Cosine).unwrap(), + )]) + .unwrap(); + assert_eq!(got, want); +} + +#[cfg(feature = "turbopuffer")] +#[test] +fn turbopuffer_from_row_derives_named_vectors_and_runtime_dimension() { + use cocoindex::SchemaFields; + use cocoindex::connectors::turbopuffer::{DistanceMetric, NamespaceSchema, VectorDef}; + + #[derive(SchemaFields)] + #[allow(dead_code)] + struct Row { + id: String, + text: String, + #[coco(vector)] + embedding: Vec, + } + + let got = NamespaceSchema::from_row::(DistanceMetric::CosineDistance) + .unwrap() + .with_vector_dim("embedding", 3) + .unwrap(); + let want = NamespaceSchema::named( + [("embedding", VectorDef::f32(3).unwrap())], + DistanceMetric::CosineDistance, + ) + .unwrap(); + assert_eq!(got, want); +} + /// End-to-end: a `from_row`-derived schema actually creates a SQLite table and /// round-trips a row (no server needed). #[cfg(feature = "sqlite")] diff --git a/rust/sdk/cocoindex_macros/src/lib.rs b/rust/sdk/cocoindex_macros/src/lib.rs index 41c29dfac..1916532ae 100644 --- a/rust/sdk/cocoindex_macros/src/lib.rs +++ b/rust/sdk/cocoindex_macros/src/lib.rs @@ -1322,8 +1322,15 @@ fn parse_schema_field_attr(attrs: &[syn::Attribute]) -> syn::Result TokenStream2 { /// type to a connector's `TableSchema::from_row::(primary_key)`. /// /// `Option` fields become nullable columns; everything else is `NOT NULL`. -/// See [`cocoindex::row_schema`] for the `#[coco(...)]` field attributes. +/// Use `#[coco(vector)]` with `TableSchema::with_vector_dim(...)` when the +/// dimension is available only at runtime, or `#[coco(vector = N)]` when it is +/// a compile-time constant. See [`cocoindex::row_schema`] for the other field +/// attributes. #[proc_macro_derive(SchemaFields, attributes(coco))] pub fn derive_schema_fields(input: TokenStream) -> TokenStream { let input = parse_macro_input!(input as syn::DeriveInput); From 898b835b2e537bedb9f3e9ca45a8158cae5cbbaf Mon Sep 17 00:00:00 2001 From: zhhanz Date: Tue, 14 Jul 2026 17:23:21 -0700 Subject: [PATCH 06/15] Batch and memoize Rust text embeddings --- examples/rust/amazon_s3_embedding/src/main.rs | 2 +- examples/rust/code_embedding/src/main.rs | 2 +- .../rust/code_embedding_lancedb/src/main.rs | 2 +- .../rust/gdrive_text_embedding/src/main.rs | 2 +- examples/rust/image_search/src/main.rs | 2 +- .../oci_object_storage_embedding/src/main.rs | 2 +- examples/rust/paper_metadata/src/main.rs | 4 +- examples/rust/pdf_embedding/src/main.rs | 2 +- examples/rust/postgres_source/src/main.rs | 7 +- examples/rust/text_embedding/src/main.rs | 12 +- .../rust/text_embedding_lancedb/src/main.rs | 2 +- .../rust/text_embedding_qdrant/src/main.rs | 2 +- .../text_embedding_turbopuffer/src/main.rs | 2 +- rust/sdk/cocoindex/src/lib.rs | 2 + .../src/ops/sentence_transformers.rs | 127 ++++++++++++++++-- rust/sdk/cocoindex_macros/src/lib.rs | 5 +- 16 files changed, 150 insertions(+), 27 deletions(-) diff --git a/examples/rust/amazon_s3_embedding/src/main.rs b/examples/rust/amazon_s3_embedding/src/main.rs index 0d82a7cd3..192062c58 100644 --- a/examples/rust/amazon_s3_embedding/src/main.rs +++ b/examples/rust/amazon_s3_embedding/src/main.rs @@ -144,7 +144,7 @@ async fn query_once( embedder: &SentenceTransformerEmbedder, query: &str, ) -> Result<()> { - let query_vec = vector_param(&embedder.embed(query).await?); + let query_vec = vector_param(&Embedder::embed(embedder, query).await?); let rows = sqlx::query(&format!( "SELECT filename, text, embedding <=> $1::vector AS distance \ FROM \"{PG_SCHEMA}\".\"{TABLE}\" ORDER BY distance ASC LIMIT $2" diff --git a/examples/rust/code_embedding/src/main.rs b/examples/rust/code_embedding/src/main.rs index 78f558fbe..e54d9d827 100644 --- a/examples/rust/code_embedding/src/main.rs +++ b/examples/rust/code_embedding/src/main.rs @@ -157,7 +157,7 @@ async fn query_once( embedder: &SentenceTransformerEmbedder, query: &str, ) -> Result<()> { - let query_vec = vector_param(&embedder.embed(query).await?); + let query_vec = vector_param(&Embedder::embed(embedder, query).await?); let rows = sqlx::query(&format!( "SELECT filename, code, start_line, end_line, embedding <=> $1::vector AS distance \ diff --git a/examples/rust/code_embedding_lancedb/src/main.rs b/examples/rust/code_embedding_lancedb/src/main.rs index ea7331cea..793a3dd97 100644 --- a/examples/rust/code_embedding_lancedb/src/main.rs +++ b/examples/rust/code_embedding_lancedb/src/main.rs @@ -128,7 +128,7 @@ async fn query_once( embedder: &SentenceTransformerEmbedder, query: &str, ) -> Result<()> { - let query_vec = embedder.embed(query).await?; + let query_vec = Embedder::embed(embedder, query).await?; let results = lancedb::vector_search(db, TABLE, "embedding", query_vec, TOP_K).await?; for r in results { let filename = r.get("filename").and_then(|v| v.as_str()).unwrap_or(""); diff --git a/examples/rust/gdrive_text_embedding/src/main.rs b/examples/rust/gdrive_text_embedding/src/main.rs index 747179f7f..6320b5857 100644 --- a/examples/rust/gdrive_text_embedding/src/main.rs +++ b/examples/rust/gdrive_text_embedding/src/main.rs @@ -127,7 +127,7 @@ async fn query_once( embedder: &SentenceTransformerEmbedder, query: &str, ) -> Result<()> { - let query_vec = vector_param(&embedder.embed(query).await?); + let query_vec = vector_param(&Embedder::embed(embedder, query).await?); let rows = sqlx::query(&format!( "SELECT filename, text, embedding <=> $1::vector AS distance \ diff --git a/examples/rust/image_search/src/main.rs b/examples/rust/image_search/src/main.rs index 7ad304b71..64ed0aa9d 100644 --- a/examples/rust/image_search/src/main.rs +++ b/examples/rust/image_search/src/main.rs @@ -112,7 +112,7 @@ async fn query_once( text_embedder: &SentenceTransformerEmbedder, query: &str, ) -> Result<()> { - let query_vec = text_embedder.embed(query).await?; + let query_vec = Embedder::embed(text_embedder, query).await?; let hits = qdrant::vector_search(conn, COLLECTION, query_vec, TOP_K).await?; for hit in hits { let filename = hit diff --git a/examples/rust/oci_object_storage_embedding/src/main.rs b/examples/rust/oci_object_storage_embedding/src/main.rs index ebefb546f..4da47c9aa 100644 --- a/examples/rust/oci_object_storage_embedding/src/main.rs +++ b/examples/rust/oci_object_storage_embedding/src/main.rs @@ -146,7 +146,7 @@ async fn query_once( embedder: &SentenceTransformerEmbedder, query: &str, ) -> Result<()> { - let query_vec = vector_param(&embedder.embed(query).await?); + let query_vec = vector_param(&Embedder::embed(embedder, query).await?); let rows = sqlx::query(&format!( "SELECT filename, text, embedding <=> $1::vector AS distance \ FROM \"{PG_SCHEMA}\".\"{TABLE}\" ORDER BY distance ASC LIMIT $2" diff --git a/examples/rust/paper_metadata/src/main.rs b/examples/rust/paper_metadata/src/main.rs index ff283713a..afb550a4c 100644 --- a/examples/rust/paper_metadata/src/main.rs +++ b/examples/rust/paper_metadata/src/main.rs @@ -264,7 +264,7 @@ async fn process_file(ctx: &Ctx, file: &FileEntry) -> Result { let mut embeddings = Vec::new(); // Title embedding (one row). - let title_vec = embedder.embed(&metadata.title).await?; + let title_vec = embedder.embed(ctx, &metadata.title).await?; let title_id = uuid_gen .next_uuid(&ctx, &("title", &metadata.title)) .await? @@ -419,7 +419,7 @@ async fn query_once( embedder: &SentenceTransformerEmbedder, query: &str, ) -> Result<()> { - let query_vec = vector_param(&embedder.embed(query).await?); + let query_vec = vector_param(&Embedder::embed(embedder, query).await?); let rows = sqlx::query(&format!( "SELECT filename, location, text, embedding <=> $1::vector AS distance \ FROM \"{PG_SCHEMA}\".\"{TABLE_EMBEDDINGS}\" ORDER BY distance ASC LIMIT $2" diff --git a/examples/rust/pdf_embedding/src/main.rs b/examples/rust/pdf_embedding/src/main.rs index f5672dec3..e74424daf 100644 --- a/examples/rust/pdf_embedding/src/main.rs +++ b/examples/rust/pdf_embedding/src/main.rs @@ -143,7 +143,7 @@ async fn query_once( embedder: &SentenceTransformerEmbedder, query: &str, ) -> Result<()> { - let query_vec = vector_param(&embedder.embed(query).await?); + let query_vec = vector_param(&Embedder::embed(embedder, query).await?); let rows = sqlx::query(&format!( "SELECT filename, text, embedding <=> $1::vector AS distance \ FROM \"{PG_SCHEMA}\".\"{TABLE}\" ORDER BY distance ASC LIMIT $2" diff --git a/examples/rust/postgres_source/src/main.rs b/examples/rust/postgres_source/src/main.rs index c33709c18..e0c84a6b0 100644 --- a/examples/rust/postgres_source/src/main.rs +++ b/examples/rust/postgres_source/src/main.rs @@ -90,7 +90,10 @@ async fn process_product(ctx: &Ctx, product: SourceProduct) -> Result Result<()> { // --------------------------------------------------------------------------- async fn query(pool: &PgPool, embedder: &SentenceTransformerEmbedder, q: &str) -> Result<()> { - let vec = embedder.embed(q).await?; + let vec = Embedder::embed(embedder, q).await?; let vec_lit = format!( "[{}]", vec.iter() diff --git a/examples/rust/text_embedding/src/main.rs b/examples/rust/text_embedding/src/main.rs index 26e4e6496..b1bce8868 100644 --- a/examples/rust/text_embedding/src/main.rs +++ b/examples/rust/text_embedding/src/main.rs @@ -68,7 +68,15 @@ async fn process_file(ctx: &Ctx, file: FileEntry) -> Result> { } let texts: Vec = chunks.iter().map(|c| c.text(&text).to_string()).collect(); - let embeddings = ctx.get_key(&EMBEDDER)?.embed_batch(texts.clone()).await?; + let embedder = ctx.get_key(&EMBEDDER)?.clone(); + let embedding_ctx = ctx.clone(); + let embeddings = ctx + .map(texts.clone(), move |chunk_text| { + let embedder = embedder.clone(); + let ctx = embedding_ctx.clone(); + async move { embedder.embed(&ctx, chunk_text).await } + }) + .await?; let mut id_gen = IdGenerator::new(); let mut rows = Vec::with_capacity(texts.len()); @@ -137,7 +145,7 @@ async fn query_once( embedder: &SentenceTransformerEmbedder, query: &str, ) -> Result<()> { - let query_vec = vector_param(&embedder.embed(query).await?); + let query_vec = vector_param(&Embedder::embed(embedder, query).await?); let rows = sqlx::query(&format!( "SELECT filename, text, embedding <=> $1::vector AS distance \ FROM \"{PG_SCHEMA}\".\"{TABLE}\" ORDER BY distance ASC LIMIT $2" diff --git a/examples/rust/text_embedding_lancedb/src/main.rs b/examples/rust/text_embedding_lancedb/src/main.rs index 58da20bc3..21c748abe 100644 --- a/examples/rust/text_embedding_lancedb/src/main.rs +++ b/examples/rust/text_embedding_lancedb/src/main.rs @@ -119,7 +119,7 @@ async fn query_once( embedder: &SentenceTransformerEmbedder, query: &str, ) -> Result<()> { - let query_vec = embedder.embed(query).await?; + let query_vec = Embedder::embed(embedder, query).await?; let results = lancedb::vector_search(db, TABLE, "embedding", query_vec, TOP_K).await?; for r in results { let filename = r.get("filename").and_then(|v| v.as_str()).unwrap_or(""); diff --git a/examples/rust/text_embedding_qdrant/src/main.rs b/examples/rust/text_embedding_qdrant/src/main.rs index 97c05e935..08bcb229e 100644 --- a/examples/rust/text_embedding_qdrant/src/main.rs +++ b/examples/rust/text_embedding_qdrant/src/main.rs @@ -133,7 +133,7 @@ async fn query_once( embedder: &SentenceTransformerEmbedder, query: &str, ) -> Result<()> { - let query_vec = embedder.embed(query).await?; + let query_vec = Embedder::embed(embedder, query).await?; let hits = qdrant::vector_search(conn, COLLECTION, query_vec, TOP_K).await?; for hit in hits { let filename = hit diff --git a/examples/rust/text_embedding_turbopuffer/src/main.rs b/examples/rust/text_embedding_turbopuffer/src/main.rs index 8751595f2..e9032224b 100644 --- a/examples/rust/text_embedding_turbopuffer/src/main.rs +++ b/examples/rust/text_embedding_turbopuffer/src/main.rs @@ -136,7 +136,7 @@ async fn query_once( namespace: &str, query: &str, ) -> Result<()> { - let query_vec = embedder.embed(query).await?; + let query_vec = Embedder::embed(embedder, query).await?; let hits = turbopuffer::vector_search(conn, namespace, query_vec, TOP_K).await?; for hit in hits { let filename = hit diff --git a/rust/sdk/cocoindex/src/lib.rs b/rust/sdk/cocoindex/src/lib.rs index 049756051..fe37ed441 100644 --- a/rust/sdk/cocoindex/src/lib.rs +++ b/rust/sdk/cocoindex/src/lib.rs @@ -1,3 +1,5 @@ +extern crate self as cocoindex; + pub mod app; pub mod batched; pub mod connectors; diff --git a/rust/sdk/cocoindex/src/ops/sentence_transformers.rs b/rust/sdk/cocoindex/src/ops/sentence_transformers.rs index f9781a03a..005189d72 100644 --- a/rust/sdk/cocoindex/src/ops/sentence_transformers.rs +++ b/rust/sdk/cocoindex/src/ops/sentence_transformers.rs @@ -20,11 +20,43 @@ use crate::resources::schema::{VectorElementType, VectorSchema, VectorSchemaProv /// Cheap to clone: the underlying model is shared behind an [`Arc`]. #[derive(Clone)] pub struct SentenceTransformerEmbedder { - model: Arc, + model: Arc, model_name: String, dimension: usize, } +trait _EmbeddingModel: Send + Sync { + fn embed(&self, texts: Vec) -> Result>>; +} + +impl _EmbeddingModel for TextEmbedding { + fn embed(&self, texts: Vec) -> Result>> { + TextEmbedding::embed(self, texts, None) + .map_err(|e| Error::engine(format!("embedding failed: {e}"))) + } +} + +#[derive(Clone)] +struct _ScheduledEmbedder(SentenceTransformerEmbedder); + +impl serde::Serialize for _ScheduledEmbedder { + fn serialize( + &self, + serializer: S, + ) -> std::result::Result { + serde::Serialize::serialize(&(self.0.model_name(), self.0.dimension()), serializer) + } +} + +#[crate::function(memo, batching, max_batch_size = 64)] +async fn _embed_scheduled( + _ctx: &crate::Ctx, + texts: Vec, + embedder: _ScheduledEmbedder, +) -> Result>> { + embedder.0.embed_batch(texts).await +} + impl SentenceTransformerEmbedder { /// Load a model by name (e.g. `"sentence-transformers/all-MiniLM-L6-v2"`). /// @@ -79,10 +111,14 @@ impl SentenceTransformerEmbedder { self.dimension } - /// Embed a single text into an `f32` vector. - pub async fn embed(&self, text: impl Into) -> Result> { - let mut out = self.embed_batch(vec![text.into()]).await?; - Ok(out.pop().unwrap_or_default()) + /// Embed one text with per-text memoization and automatic batching. + /// + /// Concurrent cache misses in the same update are coalesced into batches of + /// up to 64 texts. Repeated texts are served from the engine memo store. + /// Use [`SentenceTransformerEmbedder::embed_batch`] for an explicit raw + /// batch outside a CocoIndex update. + pub async fn embed(&self, ctx: &crate::Ctx, text: impl Into) -> Result> { + _embed_scheduled(ctx, text.into(), _ScheduledEmbedder(self.clone())).await } /// Embed a batch of texts. Embedding runs on a blocking thread. @@ -91,10 +127,9 @@ impl SentenceTransformerEmbedder { return Ok(Vec::new()); } let model = self.model.clone(); - tokio::task::spawn_blocking(move || model.embed(texts, None)) + tokio::task::spawn_blocking(move || model.embed(texts)) .await .map_err(|e| Error::engine(format!("embedding task panicked: {e}")))? - .map_err(|e| Error::engine(format!("embedding failed: {e}"))) } } @@ -110,12 +145,84 @@ impl VectorSchemaProvider for SentenceTransformerEmbedder { #[async_trait] impl crate::resources::embedder::Embedder for SentenceTransformerEmbedder { - // Delegate to the inherent methods (method-call resolution prefers inherent, - // so these don't recurse). async fn embed(&self, text: &str) -> Result> { - self.embed(text).await + let mut embeddings = self.embed_batch(vec![text.to_string()]).await?; + Ok(embeddings.pop().unwrap_or_default()) } async fn embed_batch(&self, texts: &[String]) -> Result>> { self.embed_batch(texts.to_vec()).await } } + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use futures::future::join_all; + + use super::*; + + #[derive(Default)] + struct _CountingModel { + batches: Mutex>>, + } + + impl _EmbeddingModel for _CountingModel { + fn embed(&self, texts: Vec) -> Result>> { + self.batches.lock().unwrap().push(texts.clone()); + Ok(texts + .into_iter() + .map(|text| vec![text.len() as f32]) + .collect()) + } + } + + #[tokio::test] + async fn single_text_embed_batches_misses_and_memoizes_results() { + let model = Arc::new(_CountingModel::default()); + let embedder = SentenceTransformerEmbedder { + model: model.clone(), + model_name: "counting-model".to_string(), + dimension: 1, + }; + let tempdir = tempfile::tempdir().unwrap(); + let app = crate::Environment::builder() + .db_path(tempdir.path().join("db")) + .build() + .await + .unwrap() + .app("sentence_transformer_batching") + .await + .unwrap(); + let texts: Vec = (0..70).map(|index| format!("text-{index}")).collect(); + + for _ in 0..2 { + let embedder = embedder.clone(); + let texts = texts.clone(); + app.update(move |ctx| async move { + let expected: Vec> = + texts.iter().map(|text| vec![text.len() as f32]).collect(); + let results = join_all(texts.into_iter().map(|text| embedder.embed(&ctx, text))) + .await + .into_iter() + .collect::>>()?; + assert_eq!(results, expected); + Ok(()) + }) + .await + .unwrap(); + } + + let batches = model.batches.lock().unwrap(); + assert_eq!( + batches.iter().map(Vec::len).sum::(), + 70, + "the second update should be served entirely from memo entries" + ); + assert!( + batches.iter().any(|batch| batch.len() > 1), + "concurrent misses should be coalesced into a batch" + ); + assert!(batches.iter().all(|batch| batch.len() <= 64)); + } +} diff --git a/rust/sdk/cocoindex_macros/src/lib.rs b/rust/sdk/cocoindex_macros/src/lib.rs index 1916532ae..ec553687b 100644 --- a/rust/sdk/cocoindex_macros/src/lib.rs +++ b/rust/sdk/cocoindex_macros/src/lib.rs @@ -608,7 +608,10 @@ fn expand_batching_function( let output_ty = &batching.output_ty; let extra_params = &batching.extra_params; - let batch_impl_name = format_ident!("__coco_batch_impl_{}", fn_name); + let batch_impl_name = format_ident!( + "__coco_batch_impl_{}", + fn_name.to_string().trim_start_matches('_') + ); let mut batch_impl_sig = func.sig.clone(); batch_impl_sig.ident = batch_impl_name.clone(); let batch_static_name = format_ident!("__COCO_BATCHED_{}", fn_name.to_string().to_uppercase()); From 4a2738c68a7895f0aaf3c7ef52c6de17898c5b71 Mon Sep 17 00:00:00 2001 From: zhhanz Date: Tue, 14 Jul 2026 17:31:19 -0700 Subject: [PATCH 07/15] Document the ergonomic Rust SDK --- .../docs/getting_started/rust_quickstart.mdx | 252 +++++++++++ docs/src/data/docs-sidebar.ts | 1 + examples/rust/text_embedding/README.md | 23 +- rust/sdk/SHOWCASE.md | 404 +++++++----------- 4 files changed, 411 insertions(+), 269 deletions(-) create mode 100644 docs/src/content/docs/getting_started/rust_quickstart.mdx diff --git a/docs/src/content/docs/getting_started/rust_quickstart.mdx b/docs/src/content/docs/getting_started/rust_quickstart.mdx new file mode 100644 index 000000000..9bb7592c9 --- /dev/null +++ b/docs/src/content/docs/getting_started/rust_quickstart.mdx @@ -0,0 +1,252 @@ +--- +title: Rust SDK *quickstart* +description: > + Build an incremental Rust pipeline that chunks Markdown, batches and caches + local embeddings, and keeps a Postgres/pgvector table in sync. +meta: + time: ~15 minutes + language: Rust 1.89+ + requires: Postgres with pgvector +--- + +This tutorial builds the Rust version of the text-embedding pipeline: read +Markdown files, split them into chunks, embed the chunks, and declare the rows +that should exist in Postgres. + +The important word is **declare**. Your code describes the current target +state; CocoIndex works out which rows to insert, update, retain, or delete on +each run. + +## Create the project + +Create a binary crate with a directory for input files: + +```bash +cargo new cocoindex-rust-quickstart +cd cocoindex-rust-quickstart +mkdir markdown_files +``` + +Until the Rust SDK is published separately, depend on the CocoIndex workspace +from GitHub. Add these dependencies to `Cargo.toml`: + +```toml title="Cargo.toml" +[dependencies] +cocoindex = { git = "https://github.com/cocoindex-io/cocoindex", features = [ + "postgres", + "text", + "fastembed", +] } +dotenvy = "0.15" +serde = { version = "1", features = ["derive"] } +tokio = { version = "1", features = ["full"] } +``` + +You also need a Postgres database where the pgvector extension can be created. +Set its connection URL: + +```bash +export POSTGRES_URL=postgres://cocoindex:cocoindex@localhost/cocoindex +``` + +Add one or more `.md` files under `markdown_files/`. + +## Define resources and the row type + +Start `src/main.rs` with the imports, constants, and two typed context keys: + +```rust title="src/main.rs" +use std::path::PathBuf; + +use cocoindex::connectors::postgres; +use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; +use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; +use cocoindex::prelude::*; + +const EMBED_MODEL: &str = "sentence-transformers/all-MiniLM-L6-v2"; +const PG_SCHEMA: &str = "coco_examples"; +const TABLE: &str = "doc_embeddings"; + +cocoindex::context_key!( + static DB: postgres::Database = "text_embedding_db", + state = postgres::Database::state_id +); +cocoindex::context_key!( + static EMBEDDER: SentenceTransformerEmbedder = "embedder", + state = SentenceTransformerEmbedder::model_name +); + +#[derive(Clone, Serialize, Deserialize, SchemaFields)] +struct DocEmbedding { + id: i64, + filename: String, + chunk_start: i32, + chunk_end: i32, + text: String, + #[coco(vector)] + embedding: Vec, +} +``` + +`context_key!` gives each provided resource a stable name and type. The +`state =` expressions say which durable property should invalidate dependent +memoized work when it changes. Runtime handles such as connection pools do not +need to be serializable. + +`SchemaFields` derives the database columns from `DocEmbedding`. The vector +dimension is intentionally absent because it will come from the loaded model. + +## Process one file + +Add a processing function that splits one file and embeds its chunks: + +```rust title="src/main.rs" +#[cocoindex::function] +async fn process_file(ctx: &Ctx, file: FileEntry) -> Result> { + let filename = file.key(); + let text = file.content_str()?; + let chunks = RecursiveSplitter::new()?.split_with( + &text, + RecursiveChunkConfig { + chunk_size: 2_000, + min_chunk_size: None, + chunk_overlap: Some(500), + language: Some("markdown".to_string()), + }, + ); + + let texts: Vec = chunks + .iter() + .map(|chunk| chunk.text(&text).to_string()) + .collect(); + let embedder = ctx.get_key(&EMBEDDER)?.clone(); + let embedding_ctx = ctx.clone(); + let embeddings = ctx + .map(texts.clone(), move |chunk_text| { + let embedder = embedder.clone(); + let ctx = embedding_ctx.clone(); + async move { embedder.embed(&ctx, chunk_text).await } + }) + .await?; + + let mut id_gen = IdGenerator::new(); + let mut rows = Vec::with_capacity(texts.len()); + for ((chunk, chunk_text), embedding) in chunks.iter().zip(texts).zip(embeddings) { + let id = i64::try_from(id_gen.next_id(ctx, &chunk_text).await?) + .map_err(|_| Error::engine("generated id does not fit in BIGINT"))?; + rows.push(DocEmbedding { + id, + filename: filename.clone(), + chunk_start: chunk.start.char_offset as i32, + chunk_end: chunk.end.char_offset as i32, + text: chunk_text, + embedding, + }); + } + Ok(rows) +} +``` + +`SentenceTransformerEmbedder::embed` is item-shaped at the call site, but its +implementation automatically groups concurrent cache misses into batches of +up to 64. Repeated texts are served from CocoIndex's memo store. + +## Declare the table and rows + +Now add the app's main processing function: + +```rust title="src/main.rs" +async fn app_main(ctx: Ctx, sourcedir: PathBuf) -> Result<()> { + let vector_dim = ctx.get_key(&EMBEDDER)?.dimension(); + let schema = postgres::TableSchema::from_row::(["id"])? + .with_vector_dim("embedding", vector_dim)?; + let table = postgres::mount_table_target( + &ctx, + &DB, + TABLE, + schema, + Some(PG_SCHEMA), + ) + .await?; + table.declare_vector_index( + &ctx, + "embedding", + postgres::VectorIndexOptions { + method: "hnsw", + ..Default::default() + }, + )?; + + let files = walk_items(&sourcedir, &["**/*.md"])?; + let rows_by_file = mount_each!(files, |file| process_file(ctx, file)).await?; + + for rows in rows_by_file { + for row in rows { + table.declare_row(&ctx, &row)?; + } + } + Ok(()) +} +``` + +`mount_each!` gives every relative file path its own stable processing +component. It also fingerprints `process_file` and its arguments, so an +unchanged file can skip the entire component on the next run. If a file or +chunk disappears, the rows its component used to own are removed during table +reconciliation. + +For normal function-shaped components, prefer `use_mount!` and `mount_each!`. +Use `ctx.scope` or `ctx.mount_each` when you deliberately want an explicit-key +closure without the macros' function-and-argument fingerprint. + +## Build the environment and run + +Finish `src/main.rs` by loading the resources and providing them to an +environment: + +```rust title="src/main.rs" +fn database_url() -> String { + std::env::var("POSTGRES_URL") + .unwrap_or_else(|_| "postgres://cocoindex:cocoindex@localhost/cocoindex".to_string()) +} + +#[tokio::main] +async fn main() -> Result<()> { + dotenvy::dotenv().ok(); + let sourcedir = std::env::args() + .nth(1) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("markdown_files")); + + let database = postgres::Database::connect(&database_url()).await?; + let embedder = SentenceTransformerEmbedder::load(EMBED_MODEL).await?; + let app = Environment::builder() + .db_path(".cocoindex_db") + .provide_key(&DB, database) + .provide_key(&EMBEDDER, embedder) + .build() + .await? + .app("RustTextEmbeddingQuickstart") + .await?; + + let stats = app.run(move |ctx| app_main(ctx, sourcedir)).await?; + println!("{stats}"); + Ok(()) +} +``` + +Run the pipeline: + +```bash +cargo run +``` + +Run it again without changing an input: the target rows and embeddings are +skipped. Then edit, add, or delete a Markdown file and rerun; CocoIndex updates +only the affected component and reconciles its rows. + +## Next steps + +- Read the [Rust SDK showcase](https://github.com/cocoindex-io/cocoindex/blob/main/rust/sdk/SHOWCASE.md) for memo keys, batching, context-key forms, and filesystem targets. +- Explore the [complete Rust text-embedding example](https://github.com/cocoindex-io/cocoindex/tree/main/examples/rust/text_embedding), which also implements pgvector similarity queries. +- Review [Core Concepts](../programming_guide/core_concepts) for the target-state and component ownership model shared by the Rust and Python SDKs. diff --git a/docs/src/data/docs-sidebar.ts b/docs/src/data/docs-sidebar.ts index 4cea40299..3850e4add 100644 --- a/docs/src/data/docs-sidebar.ts +++ b/docs/src/data/docs-sidebar.ts @@ -27,6 +27,7 @@ export const sidebar: SidebarItem[] = [ { type: 'doc', slug: 'getting_started/overview', label: 'Overview' }, { type: 'doc', slug: 'getting_started/installation', label: 'Installation' }, { type: 'doc', slug: 'getting_started/quickstart', label: 'Quickstart' }, + { type: 'doc', slug: 'getting_started/rust_quickstart', label: 'Rust SDK quickstart' }, { type: 'doc', slug: 'getting_started/ai_coding_agents', label: 'Use with AI coding agents' }, ], }, diff --git a/examples/rust/text_embedding/README.md b/examples/rust/text_embedding/README.md index ad8ac72bc..aaed531e3 100644 --- a/examples/rust/text_embedding/README.md +++ b/examples/rust/text_embedding/README.md @@ -7,16 +7,19 @@ and stores them in Postgres/pgvector — then serves similarity search. ## Parallel to the Python example -| Concern | Python | Rust (this example) | -| ---------------- | ---------------------------------------- | -------------------------------------------------- | -| Source | `localfs.walk_dir` | `cocoindex::resources::fs::walk` | -| Per-file compute | `@coco.fn(memo=True) process_file` | `#[cocoindex::function(memo)] process_file` | -| Chunking | `RecursiveSplitter` (markdown) | `cocoindex_ops_text` `RecursiveChunker` (markdown) | -| Embeddings | `sentence-transformers/all-MiniLM-L6-v2` | `fastembed` `AllMiniLML6V2` (same model, 384-dim) | -| Target | `postgres.TableTarget` + pgvector index | `postgres::TableTarget` + `declare_vector_index` | - -Incrementality: unchanged files are memo-skipped; chunks of a removed/edited file -are reconciled away (the managed `TableTarget` deletes orphaned rows). +| Concern | Python | Rust (this example) | +| ---------------- | ---------------------------------------- | -------------------------------------------------------- | +| Source | `localfs.walk_dir` | `resources::fs::walk_items` | +| Per-file compute | `coco.mount_each(process_file, ...)` | `mount_each!(files, \|file\| process_file(ctx, file))` | +| Chunking | `RecursiveSplitter` (Markdown) | `ops::text::RecursiveSplitter` (Markdown) | +| Embeddings | `SentenceTransformerEmbedder.embed` | `SentenceTransformerEmbedder::embed(&ctx, text)` | +| Target | `postgres.TableTarget` + pgvector index | `postgres::TableTarget` + `declare_vector_index` | +| Row schema | Dataclass annotations | `#[derive(SchemaFields)]` + runtime model dimension | + +Incrementality: `mount_each!` skips unchanged file components before their +function runs. The embedder memoizes individual texts and automatically groups +concurrent cache misses into batches of up to 64. Chunks from a removed or +edited file are reconciled away by the managed `TableTarget`. ## Run diff --git a/rust/sdk/SHOWCASE.md b/rust/sdk/SHOWCASE.md index a3e47e58e..ef924ce6d 100644 --- a/rust/sdk/SHOWCASE.md +++ b/rust/sdk/SHOWCASE.md @@ -1,324 +1,210 @@ # Ergonomic Rust SDK -**What is the use case?** +CocoIndex's Rust SDK provides the same declarative, incremental processing +model as the Python SDK. Rust pipelines use attribute and function-like macros +for logic tracking, memoization, batching, and mounting, while an explicit +`&Ctx` carries component and resource state. -An idiomatic Rust SDK for CocoIndex with the same incremental pipeline capabilities as the Python SDK — memoized computation, scoped components, file walking with fingerprinting — but using proc macros instead of decorators and explicit `&Ctx` instead of hidden globals. +## Open and run an app -**Describe the solution you'd like** - -### `App` - -Entry point. Open an LMDB-backed app and run a pipeline: - -```rust -let app = cocoindex::App::open("my_app", ".cocoindex_db")?; - -let stats = app.run(|ctx| async move { - // pipeline logic using ctx - Ok(()) -}).await?; - -println!("{stats}"); // "processed 5, wrote 3, skipped 2 in 0.4s" -``` - -| Method | Signature | -|--------|-----------| -| `App::open` | `(name: &str, db_path: impl Into) -> Result` | -| `App::builder` | `(name: &str) -> AppBuilder` — configure `.db_path()`, `.provide()`, `.build()` | -| `app.run` | `(F: FnOnce(Ctx) -> Future>) -> Result` | - -### `Ctx` - -Pipeline context, threaded explicitly through every function: +`App::open` is asynchronous. `App::open_blocking` is available when an async +entry point is not convenient. ```rust -impl Ctx { - /// Create a named sub-component. Scopes track what changed between runs. - pub async fn scope(&self, key: &impl Display, f: impl FnOnce(Ctx) -> Future>) -> Result; - - /// Memoized computation. Skips the closure if the key hasn't changed since last run. - pub async fn memo(&self, key: &impl Serialize, f: impl FnOnce() -> Future>) -> Result; - - /// Batch-process items with per-item memoization. - /// Cache hits return stored values. Cache misses are collected and passed to `f` - /// as a single batch. Results stored back and merged in original order. - pub async fn batch(&self, items: I, key_fn: impl Fn(&Item) -> K, f: impl FnOnce(Vec) -> Future>>) -> Result>; - - /// Run a closure concurrently for each item, creating a child scope per item. - pub async fn mount_each(&self, items: I, key_fn: impl Fn(&Item) -> K, f: impl Fn(Ctx, Item) -> Fut) -> Result>; - - /// Run a closure concurrently for each item within the current scope (no child scopes). - pub async fn map(&self, items: I, f: impl Fn(Item) -> Fut) -> Result>; - - /// Get a shared resource by type, returning a typed error if missing. - pub fn get_or_err(&self) -> Result<&T>; +use cocoindex::prelude::*; - /// Try to get a shared resource by type. - pub fn try_get(&self) -> Option<&T>; +#[tokio::main] +async fn main() -> Result<()> { + let app = App::open("my_app", ".cocoindex_db").await?; + let stats = app + .run(|ctx| async move { + // Declare the desired target state using ctx. + Ok(()) + }) + .await?; - /// Write a file output. CocoIndex tracks it for incremental updates. - pub fn write_file(&self, path: impl AsRef, content: &[u8]) -> Result<()>; + println!("{stats}"); + Ok(()) } ``` -### `FileEntry` - -Returned by `fs::walk()`. Lazy content, eager fingerprint — the fingerprint is used as a memo key so unchanged files skip processing entirely: +Use `App::builder(name).db_path(path).build().await` for single-app +configuration. When an app needs shared resources, build an `Environment`, +provide the resources, and then create the app: ```rust -impl FileEntry { - pub fn path(&self) -> PathBuf; - pub fn relative_path(&self) -> &Path; - pub fn stem(&self) -> &str; - pub fn fingerprint(&self) -> impl Serialize; - pub fn content(&self) -> Result>; - pub fn content_str(&self) -> Result; -} +let app = Environment::builder() + .db_path(".cocoindex_db") + .provide_key(&DB, database) + .build() + .await? + .app("my_app") + .await?; ``` -```rust -let files = cocoindex::resources::fs::walk(&dir, &["**/*.rs", "**/*.py"])?; -``` +## Declare typed context resources -### `RunStats` +`context_key!` gives a resource a stable name and type without requiring users +to write their own `LazyLock>`. The name is persistent identity; +do not derive it from the Rust module path. ```rust -pub struct RunStats { - pub processed: u64, - pub skipped: u64, - pub written: u64, - pub deleted: u64, - pub elapsed: Duration, -} -``` - -### `#[cocoindex::function]` - -Mark a pipeline function. Emits a compile-time code hash constant (`__COCO_FN_HASH_`) for automatic cache invalidation when the function body changes: +use cocoindex::connectors::postgres; -```rust -#[cocoindex::function] -async fn process(ctx: &Ctx, file: &FileEntry) -> Result { - // function body — code hash is computed at compile time - Ok(file.content_str()?) -} +cocoindex::context_key!( + static DB: postgres::Database = "app_database", + state = postgres::Database::state_id +); +cocoindex::context_key!(static CONFIG: AppConfig = "app_config", detect_change); +cocoindex::context_key!(static CLIENT: ApiClient = "api_client"); ``` -The constant can be included in manual `ctx.memo()` / `ctx.batch()` keys so that changing the function body automatically invalidates the cache — even when the function uses non-serializable resources: +- The plain form provides a typed resource without change tracking. +- `detect_change` fingerprints the complete serializable value. +- `state = expression` fingerprints only a stable derived state, which is + useful for connections and model clients whose runtime handles are not + serializable. -```rust -#[cocoindex::function] -async fn analyze(ctx: &Ctx, file: &FileEntry) -> Result { - let client = ctx.get_or_err::()?.clone(); // non-serializable — not in key - let content = file.content_str()?; - ctx.memo(&(__COCO_FN_HASH_ANALYZE, file.fingerprint()), move || async move { - client.call(&content).await // only serializable data in the key - }).await -} -``` +Provide resources with `EnvironmentBuilder::provide_key` and read them with +`ctx.get_key(&KEY)`. Reads inside memoized functions are tracked as +dependencies when the key uses a change-detecting form. -### `#[cocoindex::function(memo)]` +## Define functions -Memoize a function by its arguments. The first `&Ctx` parameter is recognized automatically; remaining parameters become the cache key. The code hash is prepended to the key, so changing the function body automatically invalidates the cache. - -Best for pure computations where all parameters are `Serialize`: +`#[cocoindex::function]` tracks the function's logic. Adding `memo` caches the +result by the function logic, serializable arguments, and context dependencies: ```rust #[cocoindex::function(memo)] -async fn extract_pub_fns(ctx: &Ctx, file: &FileEntry) -> Result> { - Ok(file.content_str()?.lines() - .filter(|l| l.trim_start().starts_with("pub fn ")) - .map(|l| l.trim().to_string()) - .collect()) +async fn parse_file(_ctx: &Ctx, file: FileEntry) -> Result> { + parse(file.content_str()?) } ``` -Conceptually similar to: +Use `memo_key(...)` when the default representation of an argument is either +too broad or not serializable. A transform replaces that argument's memo key; +`skip` (also spelled `None`) excludes it: ```rust -pub const __COCO_FN_HASH_EXTRACT_PUB_FNS: u64 = 0x...; - -async fn extract_pub_fns(ctx: &Ctx, file: &FileEntry) -> Result> { - let __coco_key = build_fingerprint(__COCO_FN_HASH_EXTRACT_PUB_FNS, file); - cached_by_fingerprint(ctx, __coco_key, { - let file = file.clone(); - move || async move { /* original body */ } - }).await +fn entry_identity(entry: &Entry) -> (String, u64) { + (entry.name.clone(), entry.version) } -``` - -Optional `version` parameter for manual cache busting: -```rust -#[cocoindex::function(memo, version = 2)] -async fn analyze(ctx: &Ctx, file: &FileEntry) -> Result { ... } +#[cocoindex::function( + memo, + memo_key(entry = entry_identity, client = skip) +)] +async fn fetch(_ctx: &Ctx, entry: &Entry, client: &ApiClient) -> Result { + client.fetch(&entry.name).await +} ``` -### `#[cocoindex::function(batching)]` +Only skip an argument when changes to it cannot affect the result, or when the +equivalent dependency is tracked through a context key. Use `ctx.memo(...)` +for a memoized block within a function; prefer `#[function(memo)]` for a whole +function because the macro tracks its logic automatically. -Batch processing without caching. The first non-ctx parameter is the items collection. The function body receives **all** items every time — no per-item cache probing. +## Batch item-shaped calls -Use this when you always want to reprocess all items (e.g., a batch API call where caching is handled externally, or when results depend on the full set of items). +A function marked `batching` has a batch-shaped implementation but an +item-shaped call site. Concurrent calls are coalesced automatically, so callers +do not construct or manage a `Batched` value: ```rust -#[cocoindex::function(batching)] -async fn embed_all(ctx: &Ctx, texts: Vec) -> Result>> { - let client = ctx.get::().clone(); - // Always processes all texts — no caching - client.embed_batch(&texts).await +#[cocoindex::function(memo, batching, max_batch_size = 64)] +async fn embed_batch( + ctx: &Ctx, + texts: Vec, + model: String, +) -> Result>> { + ctx.get_key(&EMBEDDING_CLIENT)?.embed(texts, &model).await } + +let embeddings = ctx + .map(texts, { + let ctx = ctx.clone(); + move |text| { + let ctx = ctx.clone(); + let model = model.clone(); + async move { embed_batch(&ctx, text, model).await } + } + }) + .await?; ``` -### `#[cocoindex::function(memo, batching)]` +The body receives only the items in the current batch. With `memo`, cache hits +are returned per item and only misses enter the batch. Without `memo`, every +call is processed. `max_batch_size` caps each physical request. -Batch processing **with** per-item memoization. The first non-ctx parameter is the items collection. The macro wraps the function body in `ctx.batch()` — cache hits return stored values, cache misses are collected and passed to the body as a single batch. +The built-in `SentenceTransformerEmbedder::embed(&ctx, text)` already uses +this pattern: concurrent cache misses are batched up to 64 texts and repeated +texts are memoized. -The code hash is included in each item's cache key, so changing the function body invalidates all cached results. +## Mount processing components -The function body **can** access `ctx` (e.g., `ctx.get::()` for non-serializable resources). +For normal function calls, prefer the macros that include function logic and +arguments in a component-memo fingerprint: ```rust -#[cocoindex::function(memo, batching)] -async fn extract(ctx: &Ctx, files: Vec) -> Result> { - let client = ctx.get::().clone(); - // `files` here is only the cache misses - let mut results = Vec::new(); - for file in &files { - results.push(client.analyze(&file.content_str()?).await?); - } - Ok(results) -} -``` +let summary = use_mount!(summarize(ctx, document)).await?; -Conceptually similar to: - -```rust -pub const __COCO_FN_HASH_EXTRACT: u64 = 0x...; - -async fn extract(ctx: &Ctx, files: Vec) -> Result> { - batch_by_fingerprint( - files, - |__coco_item| build_fingerprint(__COCO_FN_HASH_EXTRACT, __coco_item), - move |files| async move { /* original body */ } - ).await -} +let outputs = mount_each!(files, |file| process_file(ctx, file, target)).await?; ``` -Extra parameters beyond the items collection are cloned into the closure and included in the per-item cache key: - -```rust -#[cocoindex::function(memo, batching)] -async fn extract(ctx: &Ctx, files: Vec, model: String) -> Result> { - // `model` is included in each item's key: (hash, &model, item.clone()) - // ... -} -``` +`mount_each!` accepts `(key, value)` items, creates one child component per +key, and runs them concurrently. Without an explicit prefix, the entry +function's name is used; the prefixed form is +`mount_each!("documents", files, |file| process_file(ctx, file))`. -**Additional context** +Use `ctx.scope(key, body)` or `ctx.mount_each(items, key_fn, body)` only when +the component key and closure are deliberately dynamic and no automatic +function/argument fingerprint is wanted. These methods always execute their +closures; they still provide stable child ownership and reconciliation. -### Example: Multi-Codebase Summarization +## Read files and declare output files -The current checked-in example is `examples/rust/multi_codebase_summarization`. It scans project subdirectories, memoizes per-file extraction, memoizes per-project aggregation, writes markdown summaries, and removes stale outputs when a project disappears. +`walk_items` produces stable `(relative_path, FileEntry)` pairs ready for +`mount_each!`. `FileEntry` is serializable, so a mounted or memoized function +can use it directly as an input. -Key patterns demonstrated: -- **module-level `OnceLock`** — shared `reqwest`-based LLM client usable from `#[cocoindex::function(memo)]` bodies -- **`#[cocoindex::function(memo)]`** — per-file extraction cached by file fingerprint -- **`#[cocoindex::function(memo)]`** — project aggregation cached until the input file summaries change -- **`ctx.mount_each()`** — concurrent per-file extraction within each project -- **`ctx.write_file()` + stale-output cleanup** — incremental markdown writes plus removal of deleted-project outputs +```rust +#[cocoindex::function] +async fn render_file(ctx: &Ctx, file: FileEntry, target: DirTarget) -> Result<()> { + let markdown = render(file.content_str()?); + target.declare_file(ctx, &format!("{}.md", file.stem()), markdown.as_bytes())?; + Ok(()) +} -```toml -# Cargo.toml -[dependencies] -cocoindex = { path = "../../cocoindex" } -dotenvy = "0.15" -reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } -serde = { version = "1", features = ["derive"] } -serde_json = "1" -tokio = { version = "1", features = ["full"] } +let target = DirTarget::mount(&ctx, "./output")?; +let files = walk_items("./input", &["**/*.txt"])?; +mount_each!(files, |file| render_file(ctx, file, target.clone())).await?; ``` -```rust -use cocoindex::prelude::*; -use serde::Deserialize; -use std::collections::HashSet; -use std::path::PathBuf; -use std::sync::OnceLock; +`DirTarget` is declarative: new and changed files are written, unchanged files +are skipped, and files no longer declared by their owning components are +removed during reconciliation. -struct LlmClient { /* fields omitted */ } -static LLM: OnceLock = OnceLock::new(); +## Derive connector schemas from row types -fn init_llm() { /* initialize reqwest client from env */ } -fn llm() -> &'static LlmClient { LLM.get().unwrap() } +`SchemaFields` keeps the Rust row and connector schema in one place. A bare +`#[coco(vector)]` marks a vector whose runtime dimension can be supplied after +the embedding model is loaded: -#[cocoindex::function(memo)] -async fn extract_file_info(_ctx: &Ctx, file: &FileEntry) -> Result { - let content = file.content_str()?; - let file_path = file.key(); - llm().extract(/* prompt using file_path + content */).await -} - -#[cocoindex::function(memo)] -async fn aggregate_project_info( - _ctx: &Ctx, - project_name: String, - file_infos: Vec, -) -> Result { - if file_infos.len() <= 1 { - /* shortcut small projects */ - } - llm().extract(/* aggregation prompt */).await +```rust +#[derive(Clone, Serialize, Deserialize, SchemaFields)] +struct DocEmbedding { + id: i64, + text: String, + #[coco(vector)] + embedding: Vec, } -#[tokio::main] -async fn main() -> cocoindex::Result<()> { - init_llm(); - let app = cocoindex::App::open("multi_codebase_summarization", ".cocoindex_db")?; - - let stats = app - .run(|ctx| async move { - let mut active_projects = HashSet::new(); - for entry in std::fs::read_dir(&root_dir)? { - let project_name = entry?.file_name().to_string_lossy().to_string(); - let project_dir = entry?.path(); - let files = - cocoindex::resources::fs::walk(&project_dir, &["*.py", "**/*.py"])?; - - let file_infos = ctx - .mount_each( - files, - |f| format!("{project_name}/{}", f.key()), - |child_ctx, file| async move { extract_file_info(&child_ctx, &file).await }, - ) - .await?; - - let project_info = - aggregate_project_info(&ctx, project_name.clone(), file_infos.clone()).await?; - let markdown = - generate_markdown(&ctx, project_name.clone(), project_info, file_infos).await?; - ctx.write_file(output_dir.join(format!("{project_name}.md")), markdown.as_bytes())?; - active_projects.insert(project_name); - } - - cleanup_stale_outputs(&output_dir, &active_projects)?; - Ok(()) - }) - .await?; - - println!("{stats}"); - Ok(()) -} +let dim = ctx.get_key(&EMBEDDER)?.dimension(); +let schema = postgres::TableSchema::from_row::(["id"])? + .with_vector_dim("embedding", dim)?; ``` -**What's happening:** - -| Run | Behavior | -|-----|----------| -| First | All Python files are cache misses. Full extraction + aggregation cost. | -| Second (no changes) | Every file fingerprint and project summary hits cache. Zero LLM calls. | -| After editing one `*.py` file | Only that file is re-extracted, and only that project's aggregation reruns. | -| After deleting a project directory | The stale `output/.md` file is removed on the next run. | - ---- -Contributors, please refer to [Contributing Guide](https://cocoindex.io/docs/about/contributing). -Unless the PR can be sent immediately (e.g. just a few lines of code), we recommend you to leave a comment on the issue like **`I'm working on it`** or **`Can I work on this issue?`** to avoid duplicating work. Our [Discord server](https://discord.com/invite/zpA9S2DR7s) is always open and friendly. +The current end-to-end reference is +[`examples/rust/text_embedding`](../../examples/rust/text_embedding). It walks +Markdown files, splits them into chunks, memoizes and batches embeddings, and +declares Postgres/pgvector rows using a schema derived from the Rust row type. From 2c5ecf9758e8f5209db4e3e7e3439719b09c2a32 Mon Sep 17 00:00:00 2001 From: zhhanz Date: Tue, 14 Jul 2026 22:23:04 -0700 Subject: [PATCH 08/15] Address Rust SDK review findings --- Cargo.lock | 1 + rust/core/src/engine/progress_display.rs | 2 +- rust/sdk/cocoindex/Cargo.toml | 9 +- rust/sdk/cocoindex/src/connectors/doris.rs | 11 + rust/sdk/cocoindex/src/connectors/lancedb.rs | 80 ++++-- rust/sdk/cocoindex/src/connectors/postgres.rs | 10 + rust/sdk/cocoindex/src/connectors/qdrant.rs | 25 +- rust/sdk/cocoindex/src/connectors/sqlite.rs | 10 + .../cocoindex/src/connectors/turbopuffer.rs | 21 ++ rust/sdk/cocoindex/src/ctx.rs | 1 - .../src/ops/sentence_transformers.rs | 22 +- rust/sdk/cocoindex/src/row_schema.rs | 32 +++ rust/sdk/cocoindex/tests/lancedb_target.rs | 59 ++++- rust/sdk/cocoindex/tests/pipeline.rs | 21 +- rust/sdk/cocoindex/tests/schema_from_row.rs | 231 ++++++++++++++++-- rust/sdk/cocoindex_macros/src/lib.rs | 167 ++++++++----- 16 files changed, 585 insertions(+), 117 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3f03d8486..6ba9d127e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2006,6 +2006,7 @@ dependencies = [ "futures", "glob", "globset", + "half", "iggy", "lancedb", "linkme", diff --git a/rust/core/src/engine/progress_display.rs b/rust/core/src/engine/progress_display.rs index 9be877a14..09d21d607 100644 --- a/rust/core/src/engine/progress_display.rs +++ b/rust/core/src/engine/progress_display.rs @@ -427,7 +427,7 @@ async fn show_progress_pty( tokio::select! { result = handle.changed() => { let version = result?; - if version == TERMINATED_VERSION { + if version >= TERMINATED_VERSION { break; } } diff --git a/rust/sdk/cocoindex/Cargo.toml b/rust/sdk/cocoindex/Cargo.toml index 0659fdf57..37c90feb1 100644 --- a/rust/sdk/cocoindex/Cargo.toml +++ b/rust/sdk/cocoindex/Cargo.toml @@ -52,7 +52,13 @@ oci_object_storage = [ "dep:chrono", "chrono/clock", ] -lancedb = ["dep:serde_json", "dep:lancedb", "dep:arrow-array", "dep:arrow-schema"] +lancedb = [ + "dep:serde_json", + "dep:lancedb", + "dep:arrow-array", + "dep:arrow-schema", + "dep:half", +] qdrant = ["dep:serde_json", "dep:qdrant-client"] turbopuffer = ["dep:serde_json", "dep:reqwest"] neo4j = ["dep:serde_json", "dep:neo4rs"] @@ -115,6 +121,7 @@ base64 = { version = "0.22", optional = true } lancedb = { version = "0.26", default-features = false, optional = true } arrow-array = { version = "57", optional = true } arrow-schema = { version = "57", optional = true } +half = { version = "2", optional = true } # Qdrant vector-store target (gRPC client; build needs `protoc`). qdrant-client = { version = "1.18", optional = true } surrealdb = { version = "3", default-features = false, optional = true, features = [ diff --git a/rust/sdk/cocoindex/src/connectors/doris.rs b/rust/sdk/cocoindex/src/connectors/doris.rs index 854a72367..b3f186ad4 100644 --- a/rust/sdk/cocoindex/src/connectors/doris.rs +++ b/rust/sdk/cocoindex/src/connectors/doris.rs @@ -368,6 +368,16 @@ impl TableSchema { def.vector_dimension = Some(dim); Ok(self) } + + fn validate_vector_dimensions(&self) -> Result<()> { + for (name, def) in &self.columns { + if def.is_vector { + let dim = def.vector_dimension.unwrap_or(0) as usize; + crate::row_schema::require_resolved_vector_dimension("Doris", name, dim)?; + } + } + Ok(()) + } } /// Map a connector-agnostic [`SchemaField`](crate::row_schema::SchemaField) to a @@ -518,6 +528,7 @@ pub fn table_target_with_options( ) -> Result> { let table_name = table_name.into(); validate_ident(&table_name, "table name")?; + table_schema.validate_vector_dimensions()?; validate_indexes(&table_schema, &options)?; let provider = register_root_target_states_provider( ctx, diff --git a/rust/sdk/cocoindex/src/connectors/lancedb.rs b/rust/sdk/cocoindex/src/connectors/lancedb.rs index 4d52fbe47..f56759259 100644 --- a/rust/sdk/cocoindex/src/connectors/lancedb.rs +++ b/rust/sdk/cocoindex/src/connectors/lancedb.rs @@ -13,13 +13,15 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; -use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; +use arrow_array::builder::{FixedSizeListBuilder, PrimitiveBuilder}; +use arrow_array::types::{ArrowPrimitiveType, Float16Type, Float32Type}; use arrow_array::{ Array, ArrayRef, BinaryArray, BooleanArray, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, RecordBatch, RecordBatchIterator, StringArray, }; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use cocoindex_utils::fingerprint::Fingerprint; +use half::f16; use lancedb::Connection; use lancedb::query::{ExecutableQuery, QueryBase}; use lancedb::table::NewColumnTransform; @@ -93,6 +95,8 @@ pub enum ColumnType { Json, /// Fixed-size float32 vector of the given dimension. Vector(usize), + /// Fixed-size float16 vector of the given dimension. + HalfVector(usize), } impl ColumnType { @@ -110,6 +114,10 @@ impl ColumnType { Arc::new(Field::new("item", DataType::Float32, true)), *dim as i32, ), + ColumnType::HalfVector(dim) => DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float16, true)), + *dim as i32, + ), } } } @@ -214,15 +222,35 @@ impl TableSchema { "LanceDB vector dimension override names unknown field {field_name:?}" )) })?; - if !matches!(def.col_type, ColumnType::Vector(_)) { - return Err(Error::engine(format!( - "LanceDB field {field_name:?} is not a vector field" - ))); + match &mut def.col_type { + ColumnType::Vector(current_dim) | ColumnType::HalfVector(current_dim) => { + *current_dim = dim; + } + _ => { + return Err(Error::engine(format!( + "LanceDB field {field_name:?} is not a vector field" + ))); + } } - def.col_type = ColumnType::Vector(dim); Ok(self) } + fn validate_vector_dimensions(&self) -> Result<()> { + for (name, def) in &self.columns { + let dim = match def.col_type { + ColumnType::Vector(dim) | ColumnType::HalfVector(dim) => dim, + _ => continue, + }; + crate::row_schema::require_resolved_vector_dimension("LanceDB", name, dim)?; + if i32::try_from(dim).is_err() { + return Err(Error::engine(format!( + "LanceDB vector field {name:?} requires a dimension in 1..=i32::MAX" + ))); + } + } + Ok(()) + } + fn column_names(&self) -> impl Iterator { self.columns.iter().map(|(n, _)| n) } @@ -250,7 +278,13 @@ fn lancedb_column_def(field: &crate::row_schema::SchemaField) -> Result ColumnType::Text, L::Bytes => ColumnType::Binary, L::Decimal | L::Duration | L::Json => ColumnType::Json, - L::Vector { dim, .. } => ColumnType::Vector(*dim as usize), + L::Vector { dim, half } => { + if *half { + ColumnType::HalfVector(*dim as usize) + } else { + ColumnType::Vector(*dim as usize) + } + } L::Custom(custom) => { return Err(Error::engine(format!( "LanceDB field {:?} has unsupported custom logical type {custom:?}", @@ -303,6 +337,7 @@ pub fn table_target_with_options( options: ManagedTargetOptions, ) -> Result> { let table_name = table_name.into(); + table_schema.validate_vector_dimensions()?; let provider = register_root_target_states_provider( ctx, format!("cocoindex/lancedb/table/{}/{}", db.name(), table_name), @@ -426,7 +461,7 @@ impl LanceTableTarget { /// Declare a vector index on `column` as an attachment of this table. The /// index is created/recreated/dropped to match the declared options. The - /// column must be a [`ColumnType::Vector`]. + /// column must be a [`ColumnType::Vector`] or [`ColumnType::HalfVector`]. pub fn declare_vector_index( &self, ctx: &Ctx, @@ -435,7 +470,7 @@ impl LanceTableTarget { ) -> Result<()> { validate_identifier(column)?; match self.column_type(column) { - Some(ColumnType::Vector(_)) => {} + Some(ColumnType::Vector(_) | ColumnType::HalfVector(_)) => {} Some(_) => { return Err(Error::engine(format!( "LanceDB vector index column {column:?} is not a vector column" @@ -1166,7 +1201,22 @@ fn build_record_batch( }) .collect::>>()?, )), - ColumnType::Vector(dim) => build_vector_array(name, *dim, def.nullable, values)?, + ColumnType::Vector(dim) => build_vector_array::( + name, + *dim, + def.nullable, + DataType::Float32, + values, + |value| value, + )?, + ColumnType::HalfVector(dim) => build_vector_array::( + name, + *dim, + def.nullable, + DataType::Float16, + values, + f16::from_f32, + )?, }; arrays.push(array); } @@ -1174,14 +1224,16 @@ fn build_record_batch( .map_err(|e| Error::engine(format!("build LanceDB record batch: {e}"))) } -fn build_vector_array<'a>( +fn build_vector_array<'a, T: ArrowPrimitiveType>( column: &str, dim: usize, nullable: bool, + item_type: DataType, values: impl Iterator, + convert: impl Fn(f32) -> T::Native, ) -> Result { - let mut builder = FixedSizeListBuilder::new(Float32Builder::new(), dim as i32) - .with_field(Arc::new(Field::new("item", DataType::Float32, true))); + let mut builder = FixedSizeListBuilder::new(PrimitiveBuilder::::new(), dim as i32) + .with_field(Arc::new(Field::new("item", item_type, true))); let mut count = 0usize; for value in values { if value.is_null() { @@ -1210,7 +1262,7 @@ fn build_vector_array<'a>( let f = v.as_f64().ok_or_else(|| { Error::engine(format!("column {column:?} has non-numeric vector element")) })?; - builder.values().append_value(f as f32); + builder.values().append_value(convert(f as f32)); } builder.append(true); count += 1; diff --git a/rust/sdk/cocoindex/src/connectors/postgres.rs b/rust/sdk/cocoindex/src/connectors/postgres.rs index a9c5cfc41..d67c52833 100644 --- a/rust/sdk/cocoindex/src/connectors/postgres.rs +++ b/rust/sdk/cocoindex/src/connectors/postgres.rs @@ -163,6 +163,15 @@ impl TableSchema { def.pg_type = format!("{base}({dim})"); Ok(self) } + + fn validate_vector_dimensions(&self) -> Result<()> { + for (name, def) in &self.columns { + if matches!(def.pg_type.as_str(), "vector" | "halfvec") { + crate::row_schema::require_resolved_vector_dimension("Postgres", name, 0)?; + } + } + Ok(()) + } } fn pgvector_type_base(pg_type: &str) -> Option<&'static str> { @@ -264,6 +273,7 @@ pub fn table_target_with_options( ) -> Result> { let table_name = table_name.into(); validate_ident(&table_name, "table name")?; + table_schema.validate_vector_dimensions()?; if let Some(schema) = pg_schema_name { validate_ident(schema, "schema name")?; } diff --git a/rust/sdk/cocoindex/src/connectors/qdrant.rs b/rust/sdk/cocoindex/src/connectors/qdrant.rs index 600a774e2..a60d3f5b9 100644 --- a/rust/sdk/cocoindex/src/connectors/qdrant.rs +++ b/rust/sdk/cocoindex/src/connectors/qdrant.rs @@ -366,7 +366,9 @@ impl CollectionSchema { }) } - /// Resolve or override one named vector field's runtime dimension. + /// Resolve or override one vector field's runtime dimension. For schemas + /// built with [`CollectionSchema::new`] or [`CollectionSchema::multivector`], + /// the single unnamed vector uses the field name `"vector"`. pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { validate_vector_size(dim).map_err(|err| { Error::engine(format!( @@ -397,6 +399,26 @@ impl CollectionSchema { Ok(self) } + fn validate_vector_dimensions(&self) -> Result<()> { + match self.vector_fields() { + VectorFields::Single(def) => crate::row_schema::require_resolved_vector_dimension( + "Qdrant", + "vector", + def.vector_size() as usize, + ), + VectorFields::Named(vectors) => { + for (name, def) in vectors { + crate::row_schema::require_resolved_vector_dimension( + "Qdrant", + &name, + def.vector_size() as usize, + )?; + } + Ok(()) + } + } + } + fn vectors_config(&self) -> VectorsConfigBuilder { let mut builder = VectorsConfigBuilder::default(); match self.vector_fields() { @@ -682,6 +704,7 @@ pub fn collection_target_with_options( options: ManagedTargetOptions, ) -> Result> { let collection_name = collection_name.into(); + schema.validate_vector_dimensions()?; let provider = register_root_target_states_provider( ctx, format!( diff --git a/rust/sdk/cocoindex/src/connectors/sqlite.rs b/rust/sdk/cocoindex/src/connectors/sqlite.rs index 2a8afc014..40cf162a0 100644 --- a/rust/sdk/cocoindex/src/connectors/sqlite.rs +++ b/rust/sdk/cocoindex/src/connectors/sqlite.rs @@ -186,6 +186,15 @@ impl TableSchema { def.sqlite_type = format!("float[{dim}]"); Ok(self) } + + fn validate_vector_dimensions(&self) -> Result<()> { + for (name, def) in &self.columns { + if def.sqlite_type == "float[0]" { + crate::row_schema::require_resolved_vector_dimension("SQLite", name, 0)?; + } + } + Ok(()) + } } /// Map a connector-agnostic [`SchemaField`](crate::row_schema::SchemaField) to a @@ -273,6 +282,7 @@ pub fn table_target_with_options( ) -> Result> { let table_name = table_name.into(); validate_ident(&table_name, "table name")?; + table_schema.validate_vector_dimensions()?; if let Some(def) = &options.virtual_table_def { validate_vec0(&table_name, &table_schema, def)?; } diff --git a/rust/sdk/cocoindex/src/connectors/turbopuffer.rs b/rust/sdk/cocoindex/src/connectors/turbopuffer.rs index 911bec213..fc3a34c3e 100644 --- a/rust/sdk/cocoindex/src/connectors/turbopuffer.rs +++ b/rust/sdk/cocoindex/src/connectors/turbopuffer.rs @@ -329,6 +329,26 @@ impl NamespaceSchema { Ok(self) } + fn validate_vector_dimensions(&self) -> Result<()> { + match self.vector_fields() { + VectorFields::Single(def) => crate::row_schema::require_resolved_vector_dimension( + "Turbopuffer", + DEFAULT_VECTOR_FIELD, + def.schema.size, + ), + VectorFields::Named(vectors) => { + for (name, def) in vectors { + crate::row_schema::require_resolved_vector_dimension( + "Turbopuffer", + &name, + def.schema.size, + )?; + } + Ok(()) + } + } + } + /// The `schema` payload Turbopuffer's write API expects. fn write_schema(&self) -> Result { let mut fields = Map::new(); @@ -579,6 +599,7 @@ pub fn namespace_target_with_options( ) -> Result> { let namespace = namespace.into(); validate_namespace(&namespace)?; + schema.validate_vector_dimensions()?; let provider = register_root_target_states_provider( ctx, format!( diff --git a/rust/sdk/cocoindex/src/ctx.rs b/rust/sdk/cocoindex/src/ctx.rs index eee3d1ced..d5ca49674 100644 --- a/rust/sdk/cocoindex/src/ctx.rs +++ b/rust/sdk/cocoindex/src/ctx.rs @@ -36,7 +36,6 @@ type ContextFingerprinter = Arc Result + Sen /// [`ContextKey::new_detect_change`], and [`ContextKey::new_with_state`]: /// /// ``` -/// # use cocoindex::ContextKey; /// # #[derive(serde::Serialize)] /// # struct AppConfig; /// # struct Database; diff --git a/rust/sdk/cocoindex/src/ops/sentence_transformers.rs b/rust/sdk/cocoindex/src/ops/sentence_transformers.rs index 005189d72..9c980ed78 100644 --- a/rust/sdk/cocoindex/src/ops/sentence_transformers.rs +++ b/rust/sdk/cocoindex/src/ops/sentence_transformers.rs @@ -20,16 +20,16 @@ use crate::resources::schema::{VectorElementType, VectorSchema, VectorSchemaProv /// Cheap to clone: the underlying model is shared behind an [`Arc`]. #[derive(Clone)] pub struct SentenceTransformerEmbedder { - model: Arc, + model: Arc, model_name: String, dimension: usize, } -trait _EmbeddingModel: Send + Sync { +trait EmbeddingModel: Send + Sync { fn embed(&self, texts: Vec) -> Result>>; } -impl _EmbeddingModel for TextEmbedding { +impl EmbeddingModel for TextEmbedding { fn embed(&self, texts: Vec) -> Result>> { TextEmbedding::embed(self, texts, None) .map_err(|e| Error::engine(format!("embedding failed: {e}"))) @@ -37,9 +37,9 @@ impl _EmbeddingModel for TextEmbedding { } #[derive(Clone)] -struct _ScheduledEmbedder(SentenceTransformerEmbedder); +struct ScheduledEmbedder(SentenceTransformerEmbedder); -impl serde::Serialize for _ScheduledEmbedder { +impl serde::Serialize for ScheduledEmbedder { fn serialize( &self, serializer: S, @@ -49,10 +49,10 @@ impl serde::Serialize for _ScheduledEmbedder { } #[crate::function(memo, batching, max_batch_size = 64)] -async fn _embed_scheduled( +async fn embed_scheduled( _ctx: &crate::Ctx, texts: Vec, - embedder: _ScheduledEmbedder, + embedder: ScheduledEmbedder, ) -> Result>> { embedder.0.embed_batch(texts).await } @@ -118,7 +118,7 @@ impl SentenceTransformerEmbedder { /// Use [`SentenceTransformerEmbedder::embed_batch`] for an explicit raw /// batch outside a CocoIndex update. pub async fn embed(&self, ctx: &crate::Ctx, text: impl Into) -> Result> { - _embed_scheduled(ctx, text.into(), _ScheduledEmbedder(self.clone())).await + embed_scheduled(ctx, text.into(), ScheduledEmbedder(self.clone())).await } /// Embed a batch of texts. Embedding runs on a blocking thread. @@ -163,11 +163,11 @@ mod tests { use super::*; #[derive(Default)] - struct _CountingModel { + struct CountingModel { batches: Mutex>>, } - impl _EmbeddingModel for _CountingModel { + impl EmbeddingModel for CountingModel { fn embed(&self, texts: Vec) -> Result>> { self.batches.lock().unwrap().push(texts.clone()); Ok(texts @@ -179,7 +179,7 @@ mod tests { #[tokio::test] async fn single_text_embed_batches_misses_and_memoizes_results() { - let model = Arc::new(_CountingModel::default()); + let model = Arc::new(CountingModel::default()); let embedder = SentenceTransformerEmbedder { model: model.clone(), model_name: "counting-model".to_string(), diff --git a/rust/sdk/cocoindex/src/row_schema.rs b/rust/sdk/cocoindex/src/row_schema.rs index 20f159ce0..2397e3ff7 100644 --- a/rust/sdk/cocoindex/src/row_schema.rs +++ b/rust/sdk/cocoindex/src/row_schema.rs @@ -30,6 +30,38 @@ //! * `#[coco(json)]` — force JSON storage for a field. //! * `#[coco(rename = "…")]` — use a different column name. +#[cfg(any( + feature = "postgres", + feature = "sqlite", + feature = "doris", + feature = "lancedb", + feature = "qdrant", + feature = "turbopuffer" +))] +use crate::error::{Error, Result}; + +#[cfg(any( + feature = "postgres", + feature = "sqlite", + feature = "doris", + feature = "lancedb", + feature = "qdrant", + feature = "turbopuffer" +))] +pub(crate) fn require_resolved_vector_dimension( + connector: &str, + field_name: &str, + dimension: usize, +) -> Result<()> { + if dimension == 0 { + return Err(Error::engine(format!( + "{connector} vector field {field_name:?} has an unresolved dimension; call \ + with_vector_dim({field_name:?}, dimension) before declaring the target" + ))); + } + Ok(()) +} + /// A connector-agnostic column type derived from a Rust field type. Each target /// connector maps these to its own SQL type strings. #[derive(Clone, Debug, PartialEq, Eq)] diff --git a/rust/sdk/cocoindex/tests/lancedb_target.rs b/rust/sdk/cocoindex/tests/lancedb_target.rs index 6590349d0..859f2408b 100644 --- a/rust/sdk/cocoindex/tests/lancedb_target.rs +++ b/rust/sdk/cocoindex/tests/lancedb_target.rs @@ -9,7 +9,7 @@ #![cfg(feature = "lancedb")] use cocoindex::connectors::lancedb::{self, ColumnDef, ColumnType, LanceDatabase, TableSchema}; -use cocoindex::{Environment, ManagedTargetOptions, Result}; +use cocoindex::{Environment, ManagedTargetOptions, Result, SchemaFields}; use serde::Serialize; cocoindex::context_key!( @@ -18,6 +18,7 @@ cocoindex::context_key!( ); const TABLE: &str = "docs"; +const HALF_TABLE: &str = "half_docs"; #[derive(Clone, Serialize)] struct Row { @@ -42,6 +43,62 @@ struct RowV2Nullable { embedding: Vec, } +#[derive(Clone, Serialize, SchemaFields)] +struct HalfRow { + id: i64, + #[coco(vector, half)] + embedding: Vec, +} + +#[tokio::test] +async fn lancedb_from_row_writes_float16_vectors_without_downgrading() -> Result<()> { + let tempdir = tempfile::tempdir().unwrap(); + let uri = tempdir.path().join("lancedb_data"); + let db = LanceDatabase::connect(uri.to_str().unwrap()).await?; + let app = Environment::builder() + .db_path(tempdir.path().join(".cocoindex_db")) + .provide_key(&DB, db.clone()) + .build() + .await? + .app("LanceHalfVectorTest") + .await?; + + app.run(move |ctx| async move { + let schema = TableSchema::from_row::(["id"])?.with_vector_dim("embedding", 3)?; + let table = lancedb::mount_table_target(&ctx, &DB, HALF_TABLE, schema).await?; + table.declare_row( + &ctx, + &HalfRow { + id: 1, + embedding: vec![0.25, 0.5, 0.75], + }, + )?; + Ok(()) + }) + .await?; + + let table = db + .connection() + .open_table(HALF_TABLE) + .execute() + .await + .unwrap(); + let schema = table.schema().await.unwrap(); + assert_eq!( + schema.field_with_name("embedding").unwrap().data_type(), + &arrow_schema::DataType::FixedSizeList( + std::sync::Arc::new(arrow_schema::Field::new( + "item", + arrow_schema::DataType::Float16, + true, + )), + 3, + ) + ); + assert_eq!(table.count_rows(None).await.unwrap(), 1); + Ok(()) +} + fn schema() -> TableSchema { TableSchema::new( [ diff --git a/rust/sdk/cocoindex/tests/pipeline.rs b/rust/sdk/cocoindex/tests/pipeline.rs index ea209b0ab..f74222c94 100644 --- a/rust/sdk/cocoindex/tests/pipeline.rs +++ b/rust/sdk/cocoindex/tests/pipeline.rs @@ -1767,19 +1767,17 @@ mod batched_test { static CONTEXT_ITEMS_PROCESSED: AtomicUsize = AtomicUsize::new(0); - fn batch_factor_key() -> &'static ContextKey { - static KEY: LazyLock> = LazyLock::new(|| { - ContextKey::new_detect_change("pipeline/function_batching_context_factor") - }); - &KEY - } + cocoindex::context_key!( + static BATCH_FACTOR: i64 = "pipeline/function_batching_context_factor", + detect_change + ); #[cocoindex::function(memo, batching)] async fn macro_context_batch( ctx: &cocoindex::Ctx, items: Vec, ) -> cocoindex::Result> { - let factor = *ctx.get_key(batch_factor_key())?; + let factor = *ctx.get_key(&BATCH_FACTOR)?; CONTEXT_ITEMS_PROCESSED.fetch_add(items.len(), Ordering::SeqCst); if items == [1] { tokio::time::sleep(Duration::from_millis(30)).await; @@ -1790,7 +1788,7 @@ mod batched_test { async fn run_context_batch(db_path: &std::path::Path, factor: i64) { let app = Environment::builder() .db_path(db_path) - .provide_key(batch_factor_key(), factor) + .provide_key(&BATCH_FACTOR, factor) .build() .await .unwrap() @@ -2856,15 +2854,12 @@ mod mock_memo { /// the owned memo-scoped `Ctx`, so non-serializable resources remain /// available through normal context lookup. A skipped argument needs only /// `Any + Clone`; `MockApi` intentionally does not implement `Serialize`. - #[cocoindex::function(memo, memo_key(client = skip))] + #[cocoindex::function(memo, memo_key(_client = skip))] async fn analyze( ctx: &cocoindex::Ctx, input: String, - client: &MockApi, + _client: &MockApi, ) -> cocoindex::Result { - // Keep the skipped parameter usable in the body while sourcing the - // configured client from the memo-scoped context. - let _observed_client_calls = client.call_count(); ctx.get_or_err::()?.call(&input).await } diff --git a/rust/sdk/cocoindex/tests/schema_from_row.rs b/rust/sdk/cocoindex/tests/schema_from_row.rs index 5eca3acab..ffc168a8e 100644 --- a/rust/sdk/cocoindex/tests/schema_from_row.rs +++ b/rust/sdk/cocoindex/tests/schema_from_row.rs @@ -3,6 +3,47 @@ //! `from_row` must produce exactly the schema a user would write by hand, using //! that connector's leaf-type mapping. +#[cfg(any( + feature = "doris", + feature = "lancedb", + feature = "postgres", + feature = "qdrant", + feature = "sqlite", + feature = "turbopuffer" +))] +async fn capture_target_construction_error( + app_name: &'static str, + check: impl FnOnce(&cocoindex::Ctx) -> cocoindex::Result<()> + Send + 'static, +) -> String { + let tmp = tempfile::tempdir().unwrap(); + let app = cocoindex::App::open(app_name, tmp.path().join("db")) + .await + .unwrap(); + app.update(move |ctx| async move { + let error = match check(&ctx) { + Ok(()) => panic!("unresolved vector schema unexpectedly reached target construction"), + Err(error) => error, + }; + Ok(error.to_string()) + }) + .await + .unwrap() +} + +#[cfg(any( + feature = "doris", + feature = "lancedb", + feature = "postgres", + feature = "qdrant", + feature = "sqlite", + feature = "turbopuffer" +))] +fn assert_unresolved_dimension_error(message: &str, connector: &str) { + assert!(message.contains(connector), "{message}"); + assert!(message.contains("embedding"), "{message}"); + assert!(message.contains("with_vector_dim"), "{message}"); +} + #[cfg(feature = "doris")] #[test] fn doris_from_row_matches_explicit_schema() { @@ -25,7 +66,11 @@ fn doris_from_row_matches_explicit_schema() { code: String, } - let got = TableSchema::from_row::(["id"]).unwrap(); + let unresolved = TableSchema::from_row::(["id"]).unwrap(); + assert!(unresolved.clone().with_vector_dim("embedding", 0).is_err()); + assert!(unresolved.clone().with_vector_dim("missing", 4).is_err()); + assert!(unresolved.clone().with_vector_dim("id", 4).is_err()); + let got = unresolved; let want = TableSchema::new( [ ("id", ColumnDef::new("TEXT").not_null()), @@ -60,10 +105,11 @@ fn sqlite_from_row_matches_explicit_schema() { embedding: Vec, } - let got = TableSchema::from_row::(["id"]) - .unwrap() - .with_vector_dim("embedding", 3) - .unwrap(); + let unresolved = TableSchema::from_row::(["id"]).unwrap(); + assert!(unresolved.clone().with_vector_dim("embedding", 0).is_err()); + assert!(unresolved.clone().with_vector_dim("missing", 3).is_err()); + assert!(unresolved.clone().with_vector_dim("id", 3).is_err()); + let got = unresolved.with_vector_dim("embedding", 3).unwrap(); let want = TableSchema::new( [ ("id", ColumnDef::new("INTEGER").not_null()), @@ -99,10 +145,11 @@ fn postgres_from_row_matches_explicit_schema() { meta: Vec, } - let got = TableSchema::from_row::(["id"]) - .unwrap() - .with_vector_dim("embedding", 8) - .unwrap(); + let unresolved = TableSchema::from_row::(["id"]).unwrap(); + assert!(unresolved.clone().with_vector_dim("embedding", 0).is_err()); + assert!(unresolved.clone().with_vector_dim("missing", 8).is_err()); + assert!(unresolved.clone().with_vector_dim("id", 8).is_err()); + let got = unresolved.with_vector_dim("embedding", 8).unwrap(); // Postgres ColumnDef::new is NOT NULL by default; `.nullable()` opts in. let want = TableSchema::new( [ @@ -138,11 +185,18 @@ fn lancedb_from_row_matches_explicit_schema() { tags: Vec, #[coco(vector)] embedding: Vec, + #[coco(vector, half)] + half_embedding: Vec, } - let got = TableSchema::from_row::(["id"]) - .unwrap() + let unresolved = TableSchema::from_row::(["id"]).unwrap(); + assert!(unresolved.clone().with_vector_dim("embedding", 0).is_err()); + assert!(unresolved.clone().with_vector_dim("missing", 4).is_err()); + assert!(unresolved.clone().with_vector_dim("id", 4).is_err()); + let got = unresolved .with_vector_dim("embedding", 4) + .unwrap() + .with_vector_dim("half_embedding", 4) .unwrap(); let want = TableSchema::new( [ @@ -153,6 +207,7 @@ fn lancedb_from_row_matches_explicit_schema() { ("title", ColumnDef::new(ColumnType::Text).nullable()), ("tags", ColumnDef::new(ColumnType::Json)), ("embedding", ColumnDef::new(ColumnType::Vector(4))), + ("half_embedding", ColumnDef::new(ColumnType::HalfVector(4))), ], ["id"], ) @@ -187,10 +242,10 @@ fn qdrant_from_row_derives_named_vectors_and_runtime_dimension() { embedding: Vec, } - let got = CollectionSchema::from_row::(Distance::Cosine) - .unwrap() - .with_vector_dim("embedding", 3) - .unwrap(); + let unresolved = CollectionSchema::from_row::(Distance::Cosine).unwrap(); + assert!(unresolved.clone().with_vector_dim("embedding", 0).is_err()); + assert!(unresolved.clone().with_vector_dim("missing", 3).is_err()); + let got = unresolved.with_vector_dim("embedding", 3).unwrap(); let want = CollectionSchema::named([( "embedding", QdrantVectorDef::f32(3, Distance::Cosine).unwrap(), @@ -214,10 +269,10 @@ fn turbopuffer_from_row_derives_named_vectors_and_runtime_dimension() { embedding: Vec, } - let got = NamespaceSchema::from_row::(DistanceMetric::CosineDistance) - .unwrap() - .with_vector_dim("embedding", 3) - .unwrap(); + let unresolved = NamespaceSchema::from_row::(DistanceMetric::CosineDistance).unwrap(); + assert!(unresolved.clone().with_vector_dim("embedding", 0).is_err()); + assert!(unresolved.clone().with_vector_dim("missing", 3).is_err()); + let got = unresolved.with_vector_dim("embedding", 3).unwrap(); let want = NamespaceSchema::named( [("embedding", VectorDef::f32(3).unwrap())], DistanceMetric::CosineDistance, @@ -226,6 +281,144 @@ fn turbopuffer_from_row_derives_named_vectors_and_runtime_dimension() { assert_eq!(got, want); } +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn sqlite_target_rejects_unresolved_vector_dimensions() { + use cocoindex::connectors::sqlite::{self, Database, TableSchema}; + use cocoindex::{ContextKey, SchemaFields}; + + #[allow(dead_code)] + #[derive(SchemaFields)] + struct Row { + id: i64, + #[coco(vector)] + embedding: Vec, + } + + let db = ContextKey::::new("sqlite_unresolved_vector_db"); + let message = capture_target_construction_error("sqlite_unresolved_vector", move |ctx| { + let schema = TableSchema::from_row::(["id"])?; + sqlite::table_target(ctx, &db, "docs", schema).map(|_| ()) + }) + .await; + assert_unresolved_dimension_error(&message, "SQLite"); +} + +#[cfg(feature = "postgres")] +#[tokio::test] +async fn postgres_target_rejects_unresolved_vector_dimensions() { + use cocoindex::connectors::postgres::{self, Database, TableSchema}; + use cocoindex::{ContextKey, SchemaFields}; + + #[allow(dead_code)] + #[derive(SchemaFields)] + struct Row { + id: i64, + #[coco(vector)] + embedding: Vec, + } + + let db = ContextKey::::new("postgres_unresolved_vector_db"); + let message = capture_target_construction_error("postgres_unresolved_vector", move |ctx| { + let schema = TableSchema::from_row::(["id"])?; + postgres::table_target(ctx, &db, "docs", schema, None).map(|_| ()) + }) + .await; + assert_unresolved_dimension_error(&message, "Postgres"); +} + +#[cfg(feature = "doris")] +#[tokio::test] +async fn doris_target_rejects_unresolved_vector_dimensions() { + use cocoindex::connectors::doris::{self, DorisConnection, TableSchema}; + use cocoindex::{ContextKey, SchemaFields}; + + #[allow(dead_code)] + #[derive(SchemaFields)] + struct Row { + id: String, + #[coco(vector)] + embedding: Vec, + } + + let db = ContextKey::::new("doris_unresolved_vector_db"); + let message = capture_target_construction_error("doris_unresolved_vector", move |ctx| { + let schema = TableSchema::from_row::(["id"])?; + doris::table_target(ctx, &db, "docs", schema).map(|_| ()) + }) + .await; + assert_unresolved_dimension_error(&message, "Doris"); +} + +#[cfg(feature = "lancedb")] +#[tokio::test] +async fn lancedb_target_rejects_unresolved_vector_dimensions() { + use cocoindex::connectors::lancedb::{self, LanceDatabase, TableSchema}; + use cocoindex::{ContextKey, SchemaFields}; + + #[allow(dead_code)] + #[derive(SchemaFields)] + struct Row { + id: i64, + #[coco(vector)] + embedding: Vec, + } + + let db = ContextKey::::new("lancedb_unresolved_vector_db"); + let message = capture_target_construction_error("lancedb_unresolved_vector", move |ctx| { + let schema = TableSchema::from_row::(["id"])?; + lancedb::table_target(ctx, &db, "docs", schema).map(|_| ()) + }) + .await; + assert_unresolved_dimension_error(&message, "LanceDB"); +} + +#[cfg(feature = "qdrant")] +#[tokio::test] +async fn qdrant_target_rejects_unresolved_vector_dimensions() { + use cocoindex::connectors::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; + use cocoindex::{ContextKey, SchemaFields}; + + #[allow(dead_code)] + #[derive(SchemaFields)] + struct Row { + #[coco(vector)] + embedding: Vec, + } + + let db = ContextKey::::new("qdrant_unresolved_vector_db"); + let message = capture_target_construction_error("qdrant_unresolved_vector", move |ctx| { + let schema = CollectionSchema::from_row::(Distance::Cosine)?; + qdrant::collection_target(ctx, &db, "docs", schema).map(|_| ()) + }) + .await; + assert_unresolved_dimension_error(&message, "Qdrant"); +} + +#[cfg(feature = "turbopuffer")] +#[tokio::test] +async fn turbopuffer_target_rejects_unresolved_vector_dimensions() { + use cocoindex::connectors::turbopuffer::{ + self, DistanceMetric, NamespaceSchema, TurbopufferConnection, + }; + use cocoindex::{ContextKey, SchemaFields}; + + #[allow(dead_code)] + #[derive(SchemaFields)] + struct Row { + #[coco(vector)] + embedding: Vec, + } + + let db = ContextKey::::new("turbopuffer_unresolved_vector_db"); + let message = capture_target_construction_error("turbopuffer_unresolved_vector", move |ctx| { + let schema = NamespaceSchema::from_row::(DistanceMetric::CosineDistance)?; + turbopuffer::namespace_target(ctx, &db, "docs", schema).map(|_| ()) + }) + .await; + assert_unresolved_dimension_error(&message, "Turbopuffer"); +} + /// End-to-end: a `from_row`-derived schema actually creates a SQLite table and /// round-trips a row (no server needed). #[cfg(feature = "sqlite")] diff --git a/rust/sdk/cocoindex_macros/src/lib.rs b/rust/sdk/cocoindex_macros/src/lib.rs index ec553687b..b25747ff8 100644 --- a/rust/sdk/cocoindex_macros/src/lib.rs +++ b/rust/sdk/cocoindex_macros/src/lib.rs @@ -406,11 +406,35 @@ fn gen_state_collect_for_param( } } +fn gen_function_memo_key( + fn_name: &Ident, + hash_const_name: &Ident, + key_writes: &[TokenStream2], +) -> TokenStream2 { + quote! {{ + let mut __coco_key = ::cocoindex::memo::new_key_fingerprinter(); + ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &"cocoindex_fn")?; + ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &::core::module_path!())?; + ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &::core::stringify!(#fn_name))?; + ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &#hash_const_name)?; + #(#key_writes)* + ::cocoindex::memo::finish_key_fingerprinter(__coco_key) + }} +} + +#[derive(Debug)] +enum BatchImplParam { + Context, + Items, + Extra(usize), +} + #[derive(Debug)] struct BatchingSignature { ctx_ident: Ident, item_param: ParamInfo, extra_params: Vec, + batch_impl_params: Vec, output_ty: Type, wrapper_sig: syn::Signature, } @@ -508,6 +532,7 @@ fn parse_batching_signature(func: &ItemFn) -> syn::Result { "batching functions cannot be generic because their scheduler is static", )); } + let mut input_idents = Vec::with_capacity(func.sig.inputs.len()); for input in &func.sig.inputs { let FnArg::Typed(param) = input else { return Err(SynError::new_spanned( @@ -515,12 +540,13 @@ fn parse_batching_signature(func: &ItemFn) -> syn::Result { "batching is supported only on free functions", )); }; - if !matches!(param.pat.as_ref(), Pat::Ident(_)) { + let Pat::Ident(param_ident) = param.pat.as_ref() else { return Err(SynError::new_spanned( ¶m.pat, "batching function parameters must be identifiers", )); - } + }; + input_idents.push(param_ident.ident.clone()); } let (ctx_ident, params) = parse_fn_params(func)?; @@ -562,6 +588,24 @@ fn parse_batching_signature(func: &ItemFn) -> syn::Result { debug_assert!(found_batch_param); wrapper_sig.output = wrapper_output; + let batch_impl_params = input_idents + .into_iter() + .map(|ident| { + if ident == ctx_ident { + BatchImplParam::Context + } else if ident == batch_param.ident { + BatchImplParam::Items + } else { + let param_index = params + .iter() + .position(|param| param.ident == ident) + .expect("all non-context parameters were parsed"); + debug_assert!(param_index > 0); + BatchImplParam::Extra(param_index - 1) + } + }) + .collect(); + Ok(BatchingSignature { ctx_ident, item_param: ParamInfo { @@ -571,6 +615,7 @@ fn parse_batching_signature(func: &ItemFn) -> syn::Result { is_str_ref: false, }, extra_params: params[1..].to_vec(), + batch_impl_params, output_ty, wrapper_sig, }) @@ -608,10 +653,7 @@ fn expand_batching_function( let output_ty = &batching.output_ty; let extra_params = &batching.extra_params; - let batch_impl_name = format_ident!( - "__coco_batch_impl_{}", - fn_name.to_string().trim_start_matches('_') - ); + let batch_impl_name = format_ident!("__coco_batch_impl_{}", fn_name); let mut batch_impl_sig = func.sig.clone(); batch_impl_sig.ident = batch_impl_name.clone(); let batch_static_name = format_ident!("__COCO_BATCHED_{}", fn_name.to_string().to_uppercase()); @@ -625,37 +667,19 @@ fn expand_batching_function( None => quote! { ::cocoindex::Batched::__scheduled(#hash_const_name) }, }; - let mut batch_impl_args = Vec::new(); - for input in &func.sig.inputs { - let FnArg::Typed(param) = input else { - return Err(SynError::new_spanned( - input, - "batching is supported only on free functions", - )); - }; - let Pat::Ident(param_ident) = param.pat.as_ref() else { - return Err(SynError::new_spanned( - ¶m.pat, - "batching function parameters must be identifiers", - )); - }; - let ident = ¶m_ident.ident; - if *ident == batching.ctx_ident { - batch_impl_args.push(quote! { &__coco_batch_ctx }); - } else if *ident == batching.item_param.ident { - batch_impl_args.push(quote! { __coco_batch_items }); - } else { - let param = extra_params - .iter() - .find(|param| param.ident == *ident) - .expect("validated non-context parameter"); + let batch_impl_args = batching.batch_impl_params.iter().map(|param| match param { + BatchImplParam::Context => quote! { &__coco_batch_ctx }, + BatchImplParam::Items => quote! { __coco_batch_items }, + BatchImplParam::Extra(param_index) => { + let param = &extra_params[*param_index]; + let ident = ¶m.ident; if param.is_ref { - batch_impl_args.push(quote! { &#ident }); + quote! { &#ident } } else { - batch_impl_args.push(quote! { #ident }); + quote! { #ident } } } - } + }); let owned_extra_clones = gen_owned_param_clones(extra_params); let extra_key_parts = extra_params.iter().map(|param| { @@ -712,16 +736,9 @@ fn expand_batching_function( .collect(); let state_clones = gen_clones(&key_params); let body_clones = gen_clones(&key_params); + let memo_key = gen_function_memo_key(fn_name, hash_const_name, &key_writes); quote! {{ - let __coco_key = { - let mut __coco_key = ::cocoindex::memo::new_key_fingerprinter(); - ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &"cocoindex_fn")?; - ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &::core::module_path!())?; - ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &::core::stringify!(#fn_name))?; - ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &#hash_const_name)?; - #(#key_writes)* - ::cocoindex::memo::finish_key_fingerprinter(__coco_key) - }; + let __coco_key = #memo_key; ::cocoindex::memo::cached_by_fingerprint_with_state( #ctx_ident, @@ -952,6 +969,7 @@ pub fn function(attr: TokenStream, item: TokenStream) -> TokenStream { ) }) .collect(); + let memo_key = gen_function_memo_key(fn_name, &hash_const_name, &key_writes); let expanded = quote! { #[doc(hidden)] @@ -960,15 +978,7 @@ pub fn function(attr: TokenStream, item: TokenStream) -> TokenStream { #(#attrs)* #vis #sig { - let __coco_key = { - let mut __coco_key = ::cocoindex::memo::new_key_fingerprinter(); - ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &"cocoindex_fn")?; - ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &::core::module_path!())?; - ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &::core::stringify!(#fn_name))?; - ::cocoindex::memo::write_key_fingerprint_part(&mut __coco_key, &#hash_const_name)?; - #(#key_writes)* - ::cocoindex::memo::finish_key_fingerprinter(__coco_key) - }; + let __coco_key = #memo_key; ::cocoindex::memo::cached_by_fingerprint_with_state(#ctx_ident, __coco_key, #propagate_children_fn_logic, { #(#clone_stmts)* @@ -1288,10 +1298,10 @@ pub fn mount_each(input: TokenStream) -> TokenStream { /// Whether `expr` is exactly the identifier `ident` (a bare path with one /// segment), used to spot the `mount_each!` item-value argument. fn expr_is_ident(expr: &Expr, ident: &Ident) -> bool { - if let Expr::Path(path) = expr { - if path.qself.is_none() { - return path.path.is_ident(ident); - } + if let Expr::Path(path) = expr + && path.qself.is_none() + { + return path.path.is_ident(ident); } false } @@ -1701,4 +1711,51 @@ mod tests { compute_code_hash(&func.block, Some(8)) ); } + + #[test] + fn batching_memo_key_includes_body_hash() { + let first: ItemFn = parse_str( + "async fn embed(ctx: &Ctx, items: Vec) -> Result> { Ok(items) }", + ) + .unwrap(); + let second: ItemFn = parse_str( + "async fn embed(ctx: &Ctx, items: Vec) -> Result> { \ + Ok(items.into_iter().map(|item| item * 2).collect()) }", + ) + .unwrap(); + assert_ne!( + compute_code_hash(&first.block, None), + compute_code_hash(&second.block, None), + "editing a batching body must change its generated logic hash" + ); + + let args = FunctionArgs::parse(quote!(memo, batching)).unwrap(); + let hash_const_name = format_ident!("__COCO_FN_HASH_EMBED"); + let first_hash = compute_code_hash(&first.block, None); + let expanded = expand_batching_function( + &first, + &args, + &hash_const_name, + "e!(#first_hash), + "e!(), + true, + true, + ) + .unwrap() + .to_string(); + assert!( + expanded + .contains("write_key_fingerprint_part (& mut __coco_key , & __COCO_FN_HASH_EMBED)"), + "the expanded batching memo key must include the body-derived hash: {expanded}" + ); + } + + #[test] + fn batching_impl_names_preserve_leading_underscores() { + let plain = format_ident!("__coco_batch_impl_{}", format_ident!("embed")); + let underscored = format_ident!("__coco_batch_impl_{}", format_ident!("_embed")); + assert_eq!(plain.to_string(), "__coco_batch_impl_embed"); + assert_eq!(underscored.to_string(), "__coco_batch_impl__embed"); + assert_ne!(plain, underscored); + } } From 4d7d641d094c538c86782b7bfd71102ad47800dd Mon Sep 17 00:00:00 2001 From: zhhanz Date: Tue, 14 Jul 2026 22:38:40 -0700 Subject: [PATCH 09/15] Use context_key macro across Rust examples --- examples/rust/amazon_s3_embedding/src/main.rs | 31 +++++----- examples/rust/audio_to_text/src/main.rs | 12 ++-- examples/rust/code_embedding/README.md | 2 +- examples/rust/code_embedding/src/main.rs | 31 +++++----- .../rust/code_embedding_lancedb/Cargo.lock | 1 + .../rust/code_embedding_lancedb/src/main.rs | 19 +++--- .../rust/conversation_to_knowledge/README.md | 2 +- .../conversation_to_knowledge/src/clients.rs | 58 ++++++++++++------- examples/rust/files_to_sqlite/src/main.rs | 15 +++-- .../rust/gdrive_text_embedding/src/main.rs | 32 +++++----- examples/rust/hn_trending_topics/src/main.rs | 21 ++++--- examples/rust/image_search/src/main.rs | 21 +++---- .../rust/image_search_colpali/src/main.rs | 22 ++++--- .../oci_object_storage_embedding/src/main.rs | 31 +++++----- examples/rust/paper_metadata/src/main.rs | 33 ++++++----- examples/rust/pdf_embedding/src/main.rs | 24 ++++---- examples/rust/postgres_source/src/main.rs | 37 +++++------- .../rust/text_embedding_lancedb/Cargo.lock | 1 + .../rust/text_embedding_lancedb/src/main.rs | 19 +++--- .../rust/text_embedding_qdrant/src/main.rs | 21 +++---- .../text_embedding_turbopuffer/src/main.rs | 24 ++++---- 21 files changed, 225 insertions(+), 232 deletions(-) diff --git a/examples/rust/amazon_s3_embedding/src/main.rs b/examples/rust/amazon_s3_embedding/src/main.rs index 192062c58..c528fb060 100644 --- a/examples/rust/amazon_s3_embedding/src/main.rs +++ b/examples/rust/amazon_s3_embedding/src/main.rs @@ -12,14 +12,14 @@ //! and `AWS_ENDPOINT_URL` for MinIO) plus `S3_BUCKET` (and optional `S3_PREFIX`). use std::path::PathBuf; -use std::sync::{Arc, LazyLock}; +use std::sync::Arc; use cocoindex::connectors::amazon_s3::{self, ListOptions, S3Client, S3File}; -use cocoindex::resources::file::PatternFilePathMatcher; +use cocoindex::connectors::postgres; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; -use cocoindex::connectors::postgres; use cocoindex::prelude::*; +use cocoindex::resources::file::PatternFilePathMatcher; use sqlx::Row; use sqlx::postgres::{PgPool, PgPoolOptions}; @@ -29,19 +29,18 @@ const PG_SCHEMA: &str = "coco_examples"; const TABLE: &str = "amazon_s3_doc_embeddings"; const TOP_K: i64 = 5; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("s3_embedding_db", |db: &postgres::Database| { - db.state_id().to_string() - }) -}); -static S3: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("s3_client", |c: &S3Client| c.state_id().to_string()) -}); -static EMBEDDER: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("embedder", |e: &SentenceTransformerEmbedder| { - e.model_name().to_string() - }) -}); +cocoindex::context_key!( + static DB: postgres::Database = "s3_embedding_db", + state = postgres::Database::state_id +); +cocoindex::context_key!( + static S3: S3Client = "s3_client", + state = S3Client::state_id +); +cocoindex::context_key!( + static EMBEDDER: SentenceTransformerEmbedder = "embedder", + state = SentenceTransformerEmbedder::model_name +); #[derive(Clone, Serialize, Deserialize, SchemaFields)] struct DocEmbeddingRow { diff --git a/examples/rust/audio_to_text/src/main.rs b/examples/rust/audio_to_text/src/main.rs index 2912f7226..12f6c471e 100644 --- a/examples/rust/audio_to_text/src/main.rs +++ b/examples/rust/audio_to_text/src/main.rs @@ -15,10 +15,9 @@ //! - target : `postgres::TableTarget` (cf. `postgres.mount_table_target`) use std::path::PathBuf; -use std::sync::LazyLock; -use cocoindex::ops::api::ApiTranscriber; use cocoindex::connectors::postgres; +use cocoindex::ops::api::ApiTranscriber; use cocoindex::prelude::*; const TABLE: &str = "audio_transcriptions"; @@ -37,11 +36,10 @@ const AUDIO_PATTERNS: &[&str] = &[ "**/*.webm", ]; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("audio_to_text_db", |db: &postgres::Database| { - db.state_id().to_string() - }) -}); +cocoindex::context_key!( + static DB: postgres::Database = "audio_to_text_db", + state = postgres::Database::state_id +); #[derive(Clone, Serialize, Deserialize, SchemaFields)] struct AudioTranscription { diff --git a/examples/rust/code_embedding/README.md b/examples/rust/code_embedding/README.md index 85ffab7a1..a83abe3a6 100644 --- a/examples/rust/code_embedding/README.md +++ b/examples/rust/code_embedding/README.md @@ -13,7 +13,7 @@ Pipeline: **walk → detect language → tree-sitter chunk → embed → store i | Language detection | `detect_code_language` | `cocoindex_ops_text::prog_langs::detect_language` | | Chunking | `RecursiveSplitter` | `cocoindex_ops_text::split::RecursiveChunker` | | Embeddings | `SentenceTransformerEmbedder` (all-MiniLM-L6-v2) | `fastembed` `AllMiniLML6V2` — **the same model**, local ONNX | -| Embedder change-detection | `ContextKey(..., detect_change=True)` | `ContextKey::new_with_state("embedder", \|e\| e.model_name)` | +| Embedder change-detection | `ContextKey(..., detect_change=True)` | `context_key!(..., state = SentenceTransformerEmbedder::model_name)` | | Vector store | `postgres.TableTarget` + `declare_vector_index` | `cocoindex::connectors::postgres` `TableTarget` + `declare_vector_index` | | Stable row ids | `IdGenerator.next_id(chunk.text)` | `IdGenerator::next_id(ctx, chunk_text)` | | Query | pgvector `<=>` | pgvector `<=>` | diff --git a/examples/rust/code_embedding/src/main.rs b/examples/rust/code_embedding/src/main.rs index e54d9d827..2ed1c8de6 100644 --- a/examples/rust/code_embedding/src/main.rs +++ b/examples/rust/code_embedding/src/main.rs @@ -15,11 +15,10 @@ //! - Postgres/pgvector TableTarget sync -> from `cocoindex` use std::path::PathBuf; -use std::sync::LazyLock; +use cocoindex::connectors::postgres; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter, detect_code_language}; -use cocoindex::connectors::postgres; use cocoindex::prelude::*; use serde::{Deserialize, Serialize}; use sqlx::Row; @@ -33,21 +32,19 @@ const TOP_K: i64 = 5; const INCLUDE_PATTERNS: &[&str] = &["**/*.py", "**/*.rs", "**/*.toml", "**/*.md", "**/*.mdx"]; -/// Shared Postgres target database. -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("code_embedding_db", |db: &postgres::Database| { - db.state_id().to_string() - }) -}); - -/// Shared embedder. `new_with_state` tracks the model name, so changing the model -/// invalidates memoized files — the parity for Python's -/// `ContextKey(..., detect_change=True)` + `Annotated[NDArray, EMBEDDER]`. -static EMBEDDER: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("embedder", |e: &SentenceTransformerEmbedder| { - e.model_name().to_string() - }) -}); +// Shared Postgres target database. +cocoindex::context_key!( + static DB: postgres::Database = "code_embedding_db", + state = postgres::Database::state_id +); + +// Shared embedder. The context key tracks the model name, so changing the model +// invalidates memoized files — the parity for Python's +// `ContextKey(..., detect_change=True)` + `Annotated[NDArray, EMBEDDER]`. +cocoindex::context_key!( + static EMBEDDER: SentenceTransformerEmbedder = "embedder", + state = SentenceTransformerEmbedder::model_name +); #[derive(Clone, Serialize, Deserialize, SchemaFields)] struct CodeEmbeddingRow { diff --git a/examples/rust/code_embedding_lancedb/Cargo.lock b/examples/rust/code_embedding_lancedb/Cargo.lock index 1a8738dc1..af749c96f 100644 --- a/examples/rust/code_embedding_lancedb/Cargo.lock +++ b/examples/rust/code_embedding_lancedb/Cargo.lock @@ -899,6 +899,7 @@ dependencies = [ "fastembed", "futures", "globset", + "half", "lancedb", "linkme", "rmp-serde", diff --git a/examples/rust/code_embedding_lancedb/src/main.rs b/examples/rust/code_embedding_lancedb/src/main.rs index 793a3dd97..a7ccf3eb3 100644 --- a/examples/rust/code_embedding_lancedb/src/main.rs +++ b/examples/rust/code_embedding_lancedb/src/main.rs @@ -11,7 +11,6 @@ //! Python example's use of `cocoindex.connectors.lancedb`). use std::path::PathBuf; -use std::sync::LazyLock; use cocoindex::connectors::lancedb::{self, LanceDatabase, TableSchema}; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; @@ -25,16 +24,14 @@ const TOP_K: usize = 5; const INCLUDE_PATTERNS: &[&str] = &["**/*.py", "**/*.rs", "**/*.toml", "**/*.md", "**/*.mdx"]; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("code_embedding_db", |db: &LanceDatabase| { - db.state_id().to_string() - }) -}); -static EMBEDDER: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("embedder", |e: &SentenceTransformerEmbedder| { - e.model_name().to_string() - }) -}); +cocoindex::context_key!( + static DB: LanceDatabase = "code_embedding_db", + state = LanceDatabase::state_id +); +cocoindex::context_key!( + static EMBEDDER: SentenceTransformerEmbedder = "embedder", + state = SentenceTransformerEmbedder::model_name +); #[derive(Clone, Serialize, Deserialize, SchemaFields)] struct CodeEmbedding { diff --git a/examples/rust/conversation_to_knowledge/README.md b/examples/rust/conversation_to_knowledge/README.md index 2ff7a3017..69f695cb5 100644 --- a/examples/rust/conversation_to_knowledge/README.md +++ b/examples/rust/conversation_to_knowledge/README.md @@ -18,7 +18,7 @@ sessions, statements, persons, techs, orgs, and the relationships between them. | Stable ids | `IdGenerator` | `cocoindex::IdGenerator` | | Entity resolution | `ops.entity_resolution` (faiss + LLM) | `cocoindex::entity_resolution` + `fastembed` Snowflake embeddings + LLM pair resolver | | Graph store | `surrealdb` connector (`TableTarget`/`RelationTarget`) | `cocoindex::connectors::surrealdb` targets over the native `surrealdb` crate | -| Embedder change-detection | `ContextKey(..., detect_change=True)` | `ContextKey::new_with_state(...)` | +| Embedder change-detection | `ContextKey(..., detect_change=True)` | `context_key!(..., state = Embedder::model_name)` | ### Design notes / where it differs diff --git a/examples/rust/conversation_to_knowledge/src/clients.rs b/examples/rust/conversation_to_knowledge/src/clients.rs index c2c520776..8e525c1be 100644 --- a/examples/rust/conversation_to_knowledge/src/clients.rs +++ b/examples/rust/conversation_to_knowledge/src/clients.rs @@ -2,10 +2,10 @@ //! ContextKeys used to inject them into the pipeline. use std::path::PathBuf; -use std::sync::{Arc, LazyLock}; +use std::sync::Arc; -use cocoindex::prelude::*; pub use cocoindex::connectors::surrealdb::Graph; +use cocoindex::prelude::*; use fastembed::{ EmbeddingModel, InitOptions, InitOptionsUserDefined, Pooling, TextEmbedding, TokenizerFiles, UserDefinedEmbeddingModel, read_file_to_bytes, @@ -17,26 +17,32 @@ use serde::de::DeserializeOwned; // Context keys // --------------------------------------------------------------------------- -/// LLM used for metadata/statement extraction. State-tracked on the model name, -/// so changing the model invalidates memoized extraction (parity with Python's -/// `LLM_MODEL = ContextKey(..., detect_change=True)`). -pub static LLM: LazyLock> = - LazyLock::new(|| ContextKey::new_with_state("llm_model", |c: &LlmClient| c.model.clone())); - -/// LLM used to confirm entity-resolution pairs. -pub static RESOLVER_LLM: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("resolution_llm_model", |c: &LlmClient| c.model.clone()) -}); - -/// Local embedder for entity-resolution similarity. -pub static EMBEDDER: LazyLock> = - LazyLock::new(|| ContextKey::new_with_state("embedder", |e: &Embedder| e.model_name.clone())); - -/// SurrealDB connection. State-tracked on the target endpoint so changing the -/// external graph database invalidates local target-state reconciliation. -pub static GRAPH: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("surreal_db", |g: &Graph| g.state_id().to_string()) -}); +// LLM used for metadata/statement extraction. State-tracked on the model name, +// so changing the model invalidates memoized extraction (parity with Python's +// `LLM_MODEL = ContextKey(..., detect_change=True)`). +cocoindex::context_key!( + pub static LLM: LlmClient = "llm_model", + state = LlmClient::model_name +); + +// LLM used to confirm entity-resolution pairs. +cocoindex::context_key!( + pub static RESOLVER_LLM: LlmClient = "resolution_llm_model", + state = LlmClient::model_name +); + +// Local embedder for entity-resolution similarity. +cocoindex::context_key!( + pub static EMBEDDER: Embedder = "embedder", + state = Embedder::model_name +); + +// SurrealDB connection. State-tracked on the target endpoint so changing the +// external graph database invalidates local target-state reconciliation. +cocoindex::context_key!( + pub static GRAPH: Graph = "surreal_db", + state = Graph::state_id +); // --------------------------------------------------------------------------- // LLM client (OpenAI-compatible, JSON mode) @@ -51,6 +57,10 @@ pub struct LlmClient { } impl LlmClient { + fn model_name(&self) -> &str { + &self.model + } + pub fn new(model: String) -> Result { let api_key = std::env::var("OPENAI_API_KEY") .or_else(|_| std::env::var("LLM_API_KEY")) @@ -114,6 +124,10 @@ pub struct Embedder { } impl Embedder { + fn model_name(&self) -> &str { + &self.model_name + } + pub fn load(model_name: &str) -> Result { let model = match model_name { "Snowflake/snowflake-arctic-embed-xs" => load_snowflake_arctic_embed_xs()?, diff --git a/examples/rust/files_to_sqlite/src/main.rs b/examples/rust/files_to_sqlite/src/main.rs index 207dd0ff6..f51d1b5e5 100644 --- a/examples/rust/files_to_sqlite/src/main.rs +++ b/examples/rust/files_to_sqlite/src/main.rs @@ -13,18 +13,16 @@ //! Defaults: SOURCE_DIR = this example's `data/`, DB_PATH = `./files.db`. use std::path::PathBuf; -use std::sync::LazyLock; -use cocoindex::prelude::*; use cocoindex::connectors::sqlite; +use cocoindex::prelude::*; use serde::{Deserialize, Serialize}; use sqlx::Row as _; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("files_sqlite_db", |db: &sqlite::Database| { - db.state_id().to_string() - }) -}); +cocoindex::context_key!( + static DB: sqlite::Database = "files_sqlite_db", + state = sqlite::Database::state_id +); const TABLE: &str = "files"; @@ -78,7 +76,8 @@ async fn index(source_dir: PathBuf, db_path: String) -> Result<()> { async move { let table = sqlite::mount_table_target(&ctx, &DB, TABLE, files_schema()?).await?; - let files = cocoindex::resources::fs::walk_items(&source_dir, &["**/*.md", "**/*.txt"])?; + let files = + cocoindex::resources::fs::walk_items(&source_dir, &["**/*.md", "**/*.txt"])?; mount_each!(files, |file| process_file(ctx, file, table)).await?; Ok(()) } diff --git a/examples/rust/gdrive_text_embedding/src/main.rs b/examples/rust/gdrive_text_embedding/src/main.rs index 6320b5857..1f9b3f3ef 100644 --- a/examples/rust/gdrive_text_embedding/src/main.rs +++ b/examples/rust/gdrive_text_embedding/src/main.rs @@ -5,12 +5,11 @@ //! Postgres/pgvector. use std::path::PathBuf; -use std::sync::LazyLock; use cocoindex::connectors::gdrive::{DriveFile, GoogleDriveClient, GoogleDriveSource}; +use cocoindex::connectors::postgres; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; -use cocoindex::connectors::postgres; use cocoindex::prelude::*; use serde::{Deserialize, Serialize}; use sqlx::Row; @@ -22,23 +21,18 @@ const PG_SCHEMA: &str = "coco_examples_v1"; const TABLE: &str = "doc_embeddings"; const TOP_K: i64 = 5; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("gdrive_text_embedding_db", |db: &postgres::Database| { - db.state_id().to_string() - }) -}); - -static GDRIVE: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("gdrive_client", |client: &GoogleDriveClient| { - client.state_id().to_string() - }) -}); - -static EMBEDDER: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("embedder", |e: &SentenceTransformerEmbedder| { - e.model_name().to_string() - }) -}); +cocoindex::context_key!( + static DB: postgres::Database = "gdrive_text_embedding_db", + state = postgres::Database::state_id +); +cocoindex::context_key!( + static GDRIVE: GoogleDriveClient = "gdrive_client", + state = GoogleDriveClient::state_id +); +cocoindex::context_key!( + static EMBEDDER: SentenceTransformerEmbedder = "embedder", + state = SentenceTransformerEmbedder::model_name +); #[derive(Clone, Serialize, Deserialize, SchemaFields)] struct DocEmbeddingRow { diff --git a/examples/rust/hn_trending_topics/src/main.rs b/examples/rust/hn_trending_topics/src/main.rs index 5e87c9620..df0b631d1 100644 --- a/examples/rust/hn_trending_topics/src/main.rs +++ b/examples/rust/hn_trending_topics/src/main.rs @@ -31,13 +31,16 @@ const MAX_TEXT: usize = 4000; static HTTP: LazyLock = LazyLock::new(reqwest::Client::new); -/// Shared Postgres target database. -static PG: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("hn_db", |db: &postgres::Database| db.state_id().to_string()) -}); -/// LLM client; state-tracked on the model so changing it invalidates memos. -static LLM: LazyLock> = - LazyLock::new(|| ContextKey::new_with_state("llm_model", |c: &LlmClient| c.model.clone())); +// Shared Postgres target database. +cocoindex::context_key!( + static PG: postgres::Database = "hn_db", + state = postgres::Database::state_id +); +// LLM client; state-tracked on the model so changing it invalidates memos. +cocoindex::context_key!( + static LLM: LlmClient = "llm_model", + state = LlmClient::model_name +); // --------------------------------------------------------------------------- // LLM client (OpenAI JSON mode) @@ -75,6 +78,10 @@ const TOPICS_PROMPT: &str = "Extract topics from the user's text. Return a JSON Example: \"John Kennedy\", \"JFK\"."; impl LlmClient { + fn model_name(&self) -> &str { + &self.model + } + fn new(model: String) -> Result { let api_key = std::env::var("OPENAI_API_KEY") .or_else(|_| std::env::var("LLM_API_KEY")) diff --git a/examples/rust/image_search/src/main.rs b/examples/rust/image_search/src/main.rs index 64ed0aa9d..9ead5d07d 100644 --- a/examples/rust/image_search/src/main.rs +++ b/examples/rust/image_search/src/main.rs @@ -19,12 +19,11 @@ //! the CLIP ONNX models on first run. use std::path::PathBuf; -use std::sync::LazyLock; +use cocoindex::connectors::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; use cocoindex::ops::image::ImageEmbedder; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::prelude::*; -use cocoindex::connectors::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; use serde_json::json; /// CLIP ViT-B/32 vision tower (images) and text tower (queries). Both output @@ -44,16 +43,14 @@ const IMAGE_GLOBS: &[&str] = &[ "**/*.bmp", ]; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("image_search_db", |c: &QdrantConnection| { - c.state_id().to_string() - }) -}); -static EMBEDDER: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("image_embedder", |e: &ImageEmbedder| { - e.model_name().to_string() - }) -}); +cocoindex::context_key!( + static DB: QdrantConnection = "image_search_db", + state = QdrantConnection::state_id +); +cocoindex::context_key!( + static EMBEDDER: ImageEmbedder = "image_embedder", + state = ImageEmbedder::model_name +); /// A computed point: stable id + image vector + source filename. #[derive(Clone, Serialize, Deserialize)] diff --git a/examples/rust/image_search_colpali/src/main.rs b/examples/rust/image_search_colpali/src/main.rs index d98db4ae7..9a6335c5a 100644 --- a/examples/rust/image_search_colpali/src/main.rs +++ b/examples/rust/image_search_colpali/src/main.rs @@ -24,10 +24,9 @@ //! Build note: `qdrant-client` compiles protobufs, so `protoc` is required. use std::path::PathBuf; -use std::sync::LazyLock; -use cocoindex::prelude::*; use cocoindex::connectors::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; +use cocoindex::prelude::*; use serde::Deserialize; use serde_json::json; @@ -45,13 +44,14 @@ const IMAGE_GLOBS: &[&str] = &[ "**/*.bmp", ]; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("image_search_colpali_db", |c: &QdrantConnection| { - c.state_id().to_string() - }) -}); -static COLPALI: LazyLock> = - LazyLock::new(|| ContextKey::new_with_state("colpali", |c: &ColpaliClient| c.url.clone())); +cocoindex::context_key!( + static DB: QdrantConnection = "image_search_colpali_db", + state = QdrantConnection::state_id +); +cocoindex::context_key!( + static COLPALI: ColpaliClient = "colpali", + state = ColpaliClient::state_id +); /// HTTP client for an external ColPali inference service (see module docs). #[derive(Clone)] @@ -67,6 +67,10 @@ struct EmbeddingResponse { } impl ColpaliClient { + fn state_id(&self) -> &str { + &self.url + } + fn new(url: String) -> Self { Self { http: reqwest::Client::new(), diff --git a/examples/rust/oci_object_storage_embedding/src/main.rs b/examples/rust/oci_object_storage_embedding/src/main.rs index 4da47c9aa..473396694 100644 --- a/examples/rust/oci_object_storage_embedding/src/main.rs +++ b/examples/rust/oci_object_storage_embedding/src/main.rs @@ -13,14 +13,14 @@ //! and profile. use std::path::PathBuf; -use std::sync::{Arc, LazyLock}; +use std::sync::Arc; -use cocoindex::resources::file::PatternFilePathMatcher; use cocoindex::connectors::oci_object_storage::{self, ListOptions, OciClient, OciFile}; +use cocoindex::connectors::postgres; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; -use cocoindex::connectors::postgres; use cocoindex::prelude::*; +use cocoindex::resources::file::PatternFilePathMatcher; use sqlx::Row; use sqlx::postgres::{PgPool, PgPoolOptions}; @@ -30,19 +30,18 @@ const PG_SCHEMA: &str = "coco_examples"; const TABLE: &str = "oci_doc_embeddings"; const TOP_K: i64 = 5; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("oci_embedding_db", |db: &postgres::Database| { - db.state_id().to_string() - }) -}); -static OCI: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("oci_client", |c: &OciClient| c.state_id().to_string()) -}); -static EMBEDDER: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("embedder", |e: &SentenceTransformerEmbedder| { - e.model_name().to_string() - }) -}); +cocoindex::context_key!( + static DB: postgres::Database = "oci_embedding_db", + state = postgres::Database::state_id +); +cocoindex::context_key!( + static OCI: OciClient = "oci_client", + state = OciClient::state_id +); +cocoindex::context_key!( + static EMBEDDER: SentenceTransformerEmbedder = "embedder", + state = SentenceTransformerEmbedder::model_name +); #[derive(Clone, Serialize, Deserialize, SchemaFields)] struct DocEmbeddingRow { diff --git a/examples/rust/paper_metadata/src/main.rs b/examples/rust/paper_metadata/src/main.rs index afb550a4c..9e3a4e420 100644 --- a/examples/rust/paper_metadata/src/main.rs +++ b/examples/rust/paper_metadata/src/main.rs @@ -26,13 +26,12 @@ //! created (the query demo does a sequential cosine scan). use std::path::PathBuf; -use std::sync::LazyLock; -use cocoindex::resources::id::UuidGenerator; +use cocoindex::connectors::postgres; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{CustomLanguageConfig, RecursiveChunkConfig, RecursiveSplitter}; -use cocoindex::connectors::postgres; use cocoindex::prelude::*; +use cocoindex::resources::id::UuidGenerator; use serde::de::DeserializeOwned; use sqlx::Row; use sqlx::postgres::{PgPool, PgPoolOptions}; @@ -51,18 +50,18 @@ const ABSTRACT_MIN_CHUNK_SIZE: usize = 200; const ABSTRACT_CHUNK_OVERLAP: usize = 150; const LLM_INPUT_CHARS: usize = 4000; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("paper_metadata_db", |db: &postgres::Database| { - db.state_id().to_string() - }) -}); -static EMBEDDER: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("embedder", |e: &SentenceTransformerEmbedder| { - e.model_name().to_string() - }) -}); -static LLM: LazyLock> = - LazyLock::new(|| ContextKey::new_with_state("llm_model", |c: &LlmClient| c.model.clone())); +cocoindex::context_key!( + static DB: postgres::Database = "paper_metadata_db", + state = postgres::Database::state_id +); +cocoindex::context_key!( + static EMBEDDER: SentenceTransformerEmbedder = "embedder", + state = SentenceTransformerEmbedder::model_name +); +cocoindex::context_key!( + static LLM: LlmClient = "llm_model", + state = LlmClient::model_name +); // --------------------------------------------------------------------------- // Clients: LLM (OpenAI JSON mode); embedder is `SentenceTransformerEmbedder` @@ -77,6 +76,10 @@ struct LlmClient { } impl LlmClient { + fn model_name(&self) -> &str { + &self.model + } + fn new(model: String) -> Result { let api_key = std::env::var("OPENAI_API_KEY").map_err(|_| Error::engine("set OPENAI_API_KEY"))?; diff --git a/examples/rust/pdf_embedding/src/main.rs b/examples/rust/pdf_embedding/src/main.rs index e74424daf..04735f84a 100644 --- a/examples/rust/pdf_embedding/src/main.rs +++ b/examples/rust/pdf_embedding/src/main.rs @@ -14,11 +14,10 @@ //! mirrors Python. use std::path::PathBuf; -use std::sync::LazyLock; +use cocoindex::connectors::postgres; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; -use cocoindex::connectors::postgres; use cocoindex::prelude::*; use sqlx::Row; use sqlx::postgres::{PgPool, PgPoolOptions}; @@ -31,16 +30,14 @@ const TOP_K: i64 = 5; const CHUNK_SIZE: usize = 2000; const CHUNK_OVERLAP: usize = 500; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("pdf_embedding_db", |db: &postgres::Database| { - db.state_id().to_string() - }) -}); -static EMBEDDER: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("embedder", |e: &SentenceTransformerEmbedder| { - e.model_name().to_string() - }) -}); +cocoindex::context_key!( + static DB: postgres::Database = "pdf_embedding_db", + state = postgres::Database::state_id +); +cocoindex::context_key!( + static EMBEDDER: SentenceTransformerEmbedder = "embedder", + state = SentenceTransformerEmbedder::model_name +); #[derive(Clone, Serialize, Deserialize, SchemaFields)] struct PdfEmbedding { @@ -109,8 +106,7 @@ async fn process_file(ctx: &Ctx, file: FileEntry) -> Result> { } fn pdf_embedding_schema() -> Result { - postgres::TableSchema::from_row::(["id"])? - .with_vector_dim("embedding", EMBED_DIM) + postgres::TableSchema::from_row::(["id"])?.with_vector_dim("embedding", EMBED_DIM) } async fn app_main(ctx: Ctx, sourcedir: PathBuf) -> Result<()> { diff --git a/examples/rust/postgres_source/src/main.rs b/examples/rust/postgres_source/src/main.rs index e0c84a6b0..ec35e41ed 100644 --- a/examples/rust/postgres_source/src/main.rs +++ b/examples/rust/postgres_source/src/main.rs @@ -17,10 +17,8 @@ //! Incrementality: unchanged source rows are memo-skipped; rows deleted from the //! source have their derived output rows reconciled away automatically. -use std::sync::LazyLock; - -use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::connectors::postgres; +use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::prelude::*; use sqlx::Row; use sqlx::postgres::{PgPool, PgPoolOptions}; @@ -31,24 +29,21 @@ const PG_SCHEMA: &str = "coco_examples_v1"; const TABLE: &str = "output"; const TOP_K: i64 = 5; -/// Target database. -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("postgres_source_db", |db: &postgres::Database| { - db.state_id().to_string() - }) -}); -/// Source database. Defaults to the target URL, but can point elsewhere via -/// `SOURCE_DATABASE_URL`, matching the Python example. -static SOURCE_DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("source_pool", |db: &postgres::Database| { - db.state_id().to_string() - }) -}); -static EMBEDDER: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("embedder", |e: &SentenceTransformerEmbedder| { - e.model_name().to_string() - }) -}); +// Target database. +cocoindex::context_key!( + static DB: postgres::Database = "postgres_source_db", + state = postgres::Database::state_id +); +// Source database. Defaults to the target URL, but can point elsewhere via +// `SOURCE_DATABASE_URL`, matching the Python example. +cocoindex::context_key!( + static SOURCE_DB: postgres::Database = "source_pool", + state = postgres::Database::state_id +); +cocoindex::context_key!( + static EMBEDDER: SentenceTransformerEmbedder = "embedder", + state = SentenceTransformerEmbedder::model_name +); // --------------------------------------------------------------------------- // Data models diff --git a/examples/rust/text_embedding_lancedb/Cargo.lock b/examples/rust/text_embedding_lancedb/Cargo.lock index e833b4f1e..b2f3f1935 100644 --- a/examples/rust/text_embedding_lancedb/Cargo.lock +++ b/examples/rust/text_embedding_lancedb/Cargo.lock @@ -899,6 +899,7 @@ dependencies = [ "fastembed", "futures", "globset", + "half", "lancedb", "linkme", "rmp-serde", diff --git a/examples/rust/text_embedding_lancedb/src/main.rs b/examples/rust/text_embedding_lancedb/src/main.rs index 21c748abe..4100b7739 100644 --- a/examples/rust/text_embedding_lancedb/src/main.rs +++ b/examples/rust/text_embedding_lancedb/src/main.rs @@ -11,7 +11,6 @@ //! Python example's use of `cocoindex.connectors.lancedb`. use std::path::PathBuf; -use std::sync::LazyLock; use cocoindex::connectors::lancedb::{self, LanceDatabase, TableSchema}; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; @@ -25,16 +24,14 @@ const TOP_K: usize = 5; const CHUNK_SIZE: usize = 2000; const CHUNK_OVERLAP: usize = 500; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("text_embedding_lancedb_db", |db: &LanceDatabase| { - db.state_id().to_string() - }) -}); -static EMBEDDER: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("embedder", |e: &SentenceTransformerEmbedder| { - e.model_name().to_string() - }) -}); +cocoindex::context_key!( + static DB: LanceDatabase = "text_embedding_lancedb_db", + state = LanceDatabase::state_id +); +cocoindex::context_key!( + static EMBEDDER: SentenceTransformerEmbedder = "embedder", + state = SentenceTransformerEmbedder::model_name +); #[derive(Clone, Serialize, Deserialize, SchemaFields)] struct DocEmbedding { diff --git a/examples/rust/text_embedding_qdrant/src/main.rs b/examples/rust/text_embedding_qdrant/src/main.rs index 08bcb229e..ab063dc91 100644 --- a/examples/rust/text_embedding_qdrant/src/main.rs +++ b/examples/rust/text_embedding_qdrant/src/main.rs @@ -14,12 +14,11 @@ //! binary is required to build this example (set `PROTOC` or put it on `PATH`). use std::path::PathBuf; -use std::sync::LazyLock; +use cocoindex::connectors::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; use cocoindex::prelude::*; -use cocoindex::connectors::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; use serde_json::json; const EMBED_MODEL: &str = "sentence-transformers/all-MiniLM-L6-v2"; @@ -29,16 +28,14 @@ const TOP_K: u64 = 5; const CHUNK_SIZE: usize = 2000; const CHUNK_OVERLAP: usize = 500; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("text_embedding_qdrant_db", |c: &QdrantConnection| { - c.state_id().to_string() - }) -}); -static EMBEDDER: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("embedder", |e: &SentenceTransformerEmbedder| { - e.model_name().to_string() - }) -}); +cocoindex::context_key!( + static DB: QdrantConnection = "text_embedding_qdrant_db", + state = QdrantConnection::state_id +); +cocoindex::context_key!( + static EMBEDDER: SentenceTransformerEmbedder = "embedder", + state = SentenceTransformerEmbedder::model_name +); /// A computed point: id + vector + payload fields. #[derive(Clone, Serialize, Deserialize)] diff --git a/examples/rust/text_embedding_turbopuffer/src/main.rs b/examples/rust/text_embedding_turbopuffer/src/main.rs index e9032224b..bdc755f63 100644 --- a/examples/rust/text_embedding_turbopuffer/src/main.rs +++ b/examples/rust/text_embedding_turbopuffer/src/main.rs @@ -17,12 +17,13 @@ //! TURBOPUFFER_NAMESPACE default `TextEmbedding` use std::path::PathBuf; -use std::sync::LazyLock; +use cocoindex::connectors::turbopuffer::{ + self, DistanceMetric, NamespaceSchema, TurbopufferConnection, +}; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; use cocoindex::prelude::*; -use cocoindex::connectors::turbopuffer::{self, DistanceMetric, NamespaceSchema, TurbopufferConnection}; use serde_json::json; const EMBED_MODEL: &str = "sentence-transformers/all-MiniLM-L6-v2"; @@ -31,17 +32,14 @@ const TOP_K: usize = 5; const CHUNK_SIZE: usize = 2000; const CHUNK_OVERLAP: usize = 500; -static DB: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state( - "text_embedding_turbopuffer_db", - |c: &TurbopufferConnection| c.state_id().to_string(), - ) -}); -static EMBEDDER: LazyLock> = LazyLock::new(|| { - ContextKey::new_with_state("embedder", |e: &SentenceTransformerEmbedder| { - e.model_name().to_string() - }) -}); +cocoindex::context_key!( + static DB: TurbopufferConnection = "text_embedding_turbopuffer_db", + state = TurbopufferConnection::state_id +); +cocoindex::context_key!( + static EMBEDDER: SentenceTransformerEmbedder = "embedder", + state = SentenceTransformerEmbedder::model_name +); #[derive(Clone, Serialize, Deserialize)] struct RowData { From 8a590ab639adbf37e4b2ebd7d615ba3d0ef16468 Mon Sep 17 00:00:00 2001 From: zhhanz Date: Tue, 14 Jul 2026 23:38:47 -0700 Subject: [PATCH 10/15] Address round-two Rust SDK review findings --- Cargo.lock | 1 + rust/sdk/SHOWCASE.md | 6 + rust/sdk/cocoindex/Cargo.toml | 9 +- rust/sdk/cocoindex/src/batched.rs | 44 ++++-- rust/sdk/cocoindex/src/connectors/doris.rs | 95 ++++++++++++- rust/sdk/cocoindex/src/connectors/lancedb.rs | 4 - rust/sdk/cocoindex/src/connectors/postgres.rs | 129 ++++++++++++++++-- rust/sdk/cocoindex/src/connectors/sqlite.rs | 66 ++++++++- rust/sdk/cocoindex/src/lib.rs | 12 ++ rust/sdk/cocoindex/src/row_schema.rs | 11 ++ rust/sdk/cocoindex/tests/pipeline.rs | 61 +++++++++ rust/sdk/cocoindex/tests/postgres_target.rs | 85 +++++++++++- rust/sdk/cocoindex/tests/schema_from_row.rs | 43 +++++- rust/sdk/cocoindex_macros/src/lib.rs | 29 +++- 14 files changed, 555 insertions(+), 40 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6ba9d127e..ce23528a1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2019,6 +2019,7 @@ dependencies = [ "rmp-serde", "rsa", "rskafka", + "rust_decimal", "rustc-hash", "serde", "serde_json", diff --git a/rust/sdk/SHOWCASE.md b/rust/sdk/SHOWCASE.md index ef924ce6d..79a319afb 100644 --- a/rust/sdk/SHOWCASE.md +++ b/rust/sdk/SHOWCASE.md @@ -136,6 +136,12 @@ The body receives only the items in the current batch. With `memo`, cache hits are returned per item and only misses enter the batch. Without `memo`, every call is processed. `max_batch_size` caps each physical request. +A physical batch does not inherit any individual caller's deadline. A batch +error currently fails every item in that physical batch; the Rust SDK does not +yet retry smaller sub-batches automatically. A batching function must also not +call itself recursively from its own body, because that waits on the batcher +which is already executing the body. + The built-in `SentenceTransformerEmbedder::embed(&ctx, text)` already uses this pattern: concurrent cache misses are batched up to 64 texts and repeated texts are memoized. diff --git a/rust/sdk/cocoindex/Cargo.toml b/rust/sdk/cocoindex/Cargo.toml index 37c90feb1..6bcc91d4e 100644 --- a/rust/sdk/cocoindex/Cargo.toml +++ b/rust/sdk/cocoindex/Cargo.toml @@ -27,7 +27,13 @@ postgres = ["dep:serde_json", "dep:sqlx", "dep:chrono"] sqlite = ["dep:serde_json", "dep:sqlx", "sqlx/sqlite"] # Apache Doris target: MySQL protocol (DDL/DELETE) via sqlx + Stream Load (row # ingestion) over HTTP via reqwest. -doris = ["dep:serde_json", "dep:sqlx", "sqlx/mysql", "dep:reqwest"] +doris = [ + "dep:serde_json", + "dep:sqlx", + "sqlx/mysql", + "dep:reqwest", + "dep:base64", +] surrealdb = ["dep:serde_json", "dep:surrealdb"] kafka = ["dep:rskafka", "dep:chrono"] iggy = ["dep:iggy"] @@ -146,6 +152,7 @@ rskafka = { version = "0.6", default-features = false } wiremock = "0.6" # Generate an ephemeral RSA key for the OCI signing unit test. rand = "0.8" +rust_decimal = "1" [[bench]] name = "sdk_microbench" diff --git a/rust/sdk/cocoindex/src/batched.rs b/rust/sdk/cocoindex/src/batched.rs index b11c54b00..6285a500b 100644 --- a/rust/sdk/cocoindex/src/batched.rs +++ b/rust/sdk/cocoindex/src/batched.rs @@ -2,23 +2,30 @@ //! memoization and automatic coalescing of concurrent cache-misses into one //! batch call (via the core batcher, `cocoindex_utils::batching`). //! -//! This is the single batching mechanism, composable with memoization: each -//! `call` memo-probes its own item; only misses execute, and concurrent misses -//! (even across different components) are combined into one invocation of the -//! batch implementation. You never assemble a list yourself. +//! The `#[cocoindex::function(batching)]` wrapper composes batching with +//! per-item memoization: each call memo-probes its own item; only misses +//! execute, and concurrent misses (even across different components) are +//! combined into one invocation of the batch implementation. //! //! ```ignore -//! #[cocoindex::function] // ctx-free batch impl; emits a logic hash -//! async fn embed_batch(texts: Vec) -> coco::Result>> { +//! #[cocoindex::function(memo, batching, max_batch_size = 64)] +//! async fn embed_batch( +//! _ctx: &cocoindex::Ctx, +//! texts: Vec, +//! ) -> cocoindex::Result>> { //! model.encode(&texts) //! } //! -//! static EMBED: std::sync::LazyLock>> = -//! std::sync::LazyLock::new(|| coco::Batched::new(embed_batch, __COCO_FN_HASH_EMBED_BATCH)); -//! //! // Call per single item (e.g. inside ctx.map over chunks): -//! let emb = EMBED.call(&ctx, text).await?; +//! let emb = embed_batch(&ctx, text).await?; //! ``` +//! +//! A physical-batch error currently fails every item in that batch; the Rust +//! SDK does not automatically retry smaller sub-batches. A batching function +//! must not call itself recursively from its body because it would wait on the +//! batcher that is already executing that body. [`Batched`] remains available +//! as the lower-level explicit adapter when a generated function wrapper is not +//! suitable. use std::collections::HashMap; use std::future::Future; @@ -27,6 +34,7 @@ use std::sync::{Arc, Mutex, Weak}; use async_trait::async_trait; use cocoindex_core::engine::context::FnCallContext; +use cocoindex_core::engine::deadline::DeadlineContext; use cocoindex_utils::batching::{BatchQueue, Batcher, BatchingOptions, Runner}; use serde::Serialize; use serde::de::DeserializeOwned; @@ -65,10 +73,13 @@ impl Runner for FnRunner { } let batch_fn_ctx = Arc::new(FnCallContext::new(true)); - let batch_ctx = contexts + let mut batch_ctx = contexts .first() .expect("the core batcher never executes an empty batch") .with_fn_ctx(batch_fn_ctx.clone()); + // A physical batch belongs to all of its callers, so it must not + // inherit any one caller's deadline. + batch_ctx.deadline = DeadlineContext::NONE; let result = (self.f)(batch_ctx, inputs).await; // The body executes once, but each per-item memo/tracking context must @@ -208,6 +219,7 @@ where Out: Send + 'static, { fn new_batcher(f: BatchFn, options: BatchingOptions) -> Arc> { + Self::validate_options(&options); Arc::new(Batcher::new( FnRunner { f }, Arc::new(BatchQueue::new()), @@ -233,6 +245,7 @@ where } fn __scheduled_with_options(code_hash: u64, options: BatchingOptions) -> Self { + Self::validate_options(&options); Self { mode: BatchedMode::Scheduled { options, @@ -242,6 +255,15 @@ where } } + fn validate_options(options: &BatchingOptions) { + if let Some(max_batch_size) = options.max_batch_size { + assert!( + max_batch_size > 0, + "max_batch_size must be greater than zero" + ); + } + } + /// Schedule one generated function call. `extra_args_key` keeps calls with /// different captured arguments in separate batches. #[doc(hidden)] diff --git a/rust/sdk/cocoindex/src/connectors/doris.rs b/rust/sdk/cocoindex/src/connectors/doris.rs index b3f186ad4..be754a597 100644 --- a/rust/sdk/cocoindex/src/connectors/doris.rs +++ b/rust/sdk/cocoindex/src/connectors/doris.rs @@ -25,6 +25,7 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use base64::Engine as _; use cocoindex_utils::fingerprint::Fingerprint; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value as JsonValue}; @@ -303,7 +304,11 @@ impl TableSchema { let name = name.into(); validate_ident(&name, "column name")?; validate_doris_type(&def.doris_type)?; - out.insert(name, def); + if out.insert(name.clone(), def).is_some() { + return Err(Error::engine(format!( + "Doris table schema contains duplicate column {name:?}" + ))); + } } let primary_key: Vec = primary_key.into_iter().map(Into::into).collect(); if primary_key.is_empty() { @@ -1493,12 +1498,52 @@ fn row_state(row: &R, schema: &TableSchema) -> Result Option> { + let JsonValue::Array(items) = value else { + return None; + }; + let mut bytes = Vec::with_capacity(items.len()); + for item in items { + let byte = item.as_u64().filter(|byte| *byte <= 255)?; + bytes.push(byte as u8); + } + Some(bytes) +} + +fn is_string_type(doris_type: &str) -> bool { + doris_type == "STRING" + || doris_type == "TEXT" + || doris_type.starts_with("VARCHAR(") + || doris_type.starts_with("CHAR(") +} + +fn is_integer_doris_type(doris_type: &str) -> bool { + matches!( + doris_type, + "TINYINT" | "SMALLINT" | "INT" | "BIGINT" | "LARGEINT" + ) +} + fn pk_values(fields: &Map, pk_cols: &[String]) -> Result> { pk_cols .iter() @@ -1666,6 +1711,50 @@ mod tests { assert!(sql.contains("\"replication_num\" = \"1\""), "{sql}"); } + #[test] + fn row_state_encodes_bytes_and_duration_for_stream_load() { + #[derive(Serialize)] + struct Row { + id: i64, + data: Vec, + elapsed: std::time::Duration, + } + + let schema = TableSchema::new( + [ + ("id", ColumnDef::new("BIGINT")), + ("data", ColumnDef::new("STRING")), + ("elapsed", ColumnDef::new("BIGINT")), + ], + ["id"], + ) + .unwrap(); + let fields = row_state( + &Row { + id: 1, + data: b"hi".to_vec(), + elapsed: std::time::Duration::new(2, 345_678_000), + }, + &schema, + ) + .unwrap(); + assert_eq!(fields["data"], JsonValue::String("aGk=".to_string())); + assert_eq!(fields["elapsed"], JsonValue::from(2_345_678)); + } + + #[test] + fn table_schema_rejects_duplicate_columns() { + let error = TableSchema::new( + [ + ("id", ColumnDef::new("BIGINT")), + ("id", ColumnDef::new("TEXT")), + ], + ["id"], + ) + .unwrap_err(); + assert!(error.to_string().contains("duplicate column \"id\"")); + } + #[test] fn create_table_sql_emits_vector_and_inverted_index_clauses() { let schema = TableSchema::new( diff --git a/rust/sdk/cocoindex/src/connectors/lancedb.rs b/rust/sdk/cocoindex/src/connectors/lancedb.rs index f56759259..46ac3c996 100644 --- a/rust/sdk/cocoindex/src/connectors/lancedb.rs +++ b/rust/sdk/cocoindex/src/connectors/lancedb.rs @@ -188,10 +188,6 @@ impl TableSchema { &self.primary_key } - pub fn columns(&self) -> &[(String, ColumnDef)] { - &self.columns - } - /// Derive a schema from a `#[derive(SchemaFields)]` row type. pub fn from_row( primary_key: impl IntoIterator>, diff --git a/rust/sdk/cocoindex/src/connectors/postgres.rs b/rust/sdk/cocoindex/src/connectors/postgres.rs index d67c52833..4943fcaa5 100644 --- a/rust/sdk/cocoindex/src/connectors/postgres.rs +++ b/rust/sdk/cocoindex/src/connectors/postgres.rs @@ -10,7 +10,7 @@ //! declared immediately. [`read_table`] and [`read_table_items`] read source rows //! for use with `Ctx::mount_each`. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; use cocoindex_utils::fingerprint::Fingerprint; @@ -89,6 +89,8 @@ impl ColumnDef { pub struct TableSchema { columns: BTreeMap, primary_key: Vec, + #[serde(default, skip_serializing_if = "BTreeSet::is_empty")] + unresolved_vector_columns: BTreeSet, } impl TableSchema { @@ -101,7 +103,11 @@ impl TableSchema { let name = name.into(); validate_ident(&name, "column name")?; validate_pg_type(&def.pg_type)?; - out.insert(name, def); + if out.insert(name.clone(), def).is_some() { + return Err(Error::engine(format!( + "Postgres table schema contains duplicate column {name:?}" + ))); + } } let primary_key: Vec = primary_key.into_iter().map(Into::into).collect::>(); if primary_key.is_empty() { @@ -118,6 +124,7 @@ impl TableSchema { Ok(Self { columns: out, primary_key, + unresolved_vector_columns: BTreeSet::new(), }) } @@ -137,10 +144,20 @@ impl TableSchema { pub fn from_row( primary_key: impl IntoIterator>, ) -> Result { - let columns = T::schema_fields() + let fields = T::schema_fields(); + let unresolved_vector_columns = fields + .iter() + .filter_map(|field| match field.logical_type { + crate::row_schema::LogicalType::Vector { dim: 0, .. } => Some(field.name.clone()), + _ => None, + }) + .collect(); + let columns = fields .into_iter() .map(|f| (f.name.clone(), postgres_column_def(&f))); - Self::new(columns, primary_key) + let mut schema = Self::new(columns, primary_key)?; + schema.unresolved_vector_columns = unresolved_vector_columns; + Ok(schema) } /// Resolve or override the dimension of a vector field derived from a row. @@ -161,14 +178,13 @@ impl TableSchema { )) })?; def.pg_type = format!("{base}({dim})"); + self.unresolved_vector_columns.remove(field_name); Ok(self) } fn validate_vector_dimensions(&self) -> Result<()> { - for (name, def) in &self.columns { - if matches!(def.pg_type.as_str(), "vector" | "halfvec") { - crate::row_schema::require_resolved_vector_dimension("Postgres", name, 0)?; - } + for name in &self.unresolved_vector_columns { + crate::row_schema::require_resolved_vector_dimension("Postgres", name, 0)?; } Ok(()) } @@ -1413,9 +1429,15 @@ fn sql_literal(value: &JsonValue, col: &ColumnDef) -> Result { .map(|b| if b { "TRUE" } else { "FALSE" }.to_string()) .ok_or_else(|| Error::engine("boolean column requires bool JSON value")); } + if lower.starts_with("interval") { + return interval_literal(value, &col.pg_type); + } if is_numeric_type(&lower) { return match value { JsonValue::Number(n) => Ok(n.to_string()), + JsonValue::String(s) if is_decimal_type(&lower) && is_decimal_literal(s) => { + Ok(s.clone()) + } _ => Err(Error::engine(format!( "numeric column {} requires numeric JSON value", col.pg_type @@ -1458,6 +1480,21 @@ fn sql_literal(value: &JsonValue, col: &ColumnDef) -> Result { } } +fn interval_literal(value: &JsonValue, pg_type: &str) -> Result { + if let Some((secs, nanos)) = crate::row_schema::serialized_duration_parts(value) { + return Ok(format!( + "{}::{pg_type}", + quote_string(format!("{secs}.{nanos:09} seconds")) + )); + } + match value { + JsonValue::String(s) => Ok(format!("{}::{pg_type}", quote_string(s))), + _ => Err(Error::engine( + "interval column requires a duration or string JSON value", + )), + } +} + fn value_to_string(value: &JsonValue) -> Result<&str> { value .as_str() @@ -1668,7 +1705,27 @@ fn is_numeric_type(lower: &str) -> bool { | "double precision" | "float8" | "numeric" - ) + | "decimal" + ) || lower.starts_with("numeric(") + || lower.starts_with("decimal(") +} + +fn is_decimal_type(lower: &str) -> bool { + lower == "numeric" + || lower == "decimal" + || lower.starts_with("numeric(") + || lower.starts_with("decimal(") +} + +fn is_decimal_literal(value: &str) -> bool { + let unsigned = value.strip_prefix('-').unwrap_or(value); + let mut parts = unsigned.split('.'); + let whole = parts.next().unwrap_or_default(); + let fraction = parts.next(); + !whole.is_empty() + && whole.bytes().all(|b| b.is_ascii_digit()) + && fraction.is_none_or(|part| !part.is_empty() && part.bytes().all(|b| b.is_ascii_digit())) + && parts.next().is_none() } fn pg_err(e: sqlx::Error) -> Error { @@ -1941,6 +1998,8 @@ fn pg_col_to_json( #[cfg(test)] mod review_fix_tests { use super::*; + use rust_decimal::Decimal; + use std::str::FromStr as _; // --- vector index WITH-clause gating (pgvector) --- @@ -1996,4 +2055,56 @@ mod review_fix_tests { let value = serde_json::json!({"a": [1, 2], "b": "ok"}); assert_eq!(sanitize_json_nul(&value), value); } + + #[test] + fn numeric_literal_accepts_rust_decimal_string_encoding() { + let value = serde_json::to_value(Decimal::from_str("1234567890.012300").unwrap()).unwrap(); + assert_eq!(value, JsonValue::String("1234567890.012300".to_string())); + assert_eq!( + sql_literal(&value, &ColumnDef::new("numeric(28, 6)")).unwrap(), + "1234567890.012300" + ); + assert!( + sql_literal( + &JsonValue::String("1); DROP TABLE rows; --".to_string()), + &ColumnDef::new("numeric") + ) + .is_err() + ); + } + + #[test] + fn interval_literal_accepts_std_duration_encoding() { + let value = serde_json::to_value(std::time::Duration::new(12, 345_000_000)).unwrap(); + assert_eq!( + sql_literal(&value, &ColumnDef::new("interval")).unwrap(), + "'12.345000000 seconds'::interval" + ); + } + + #[test] + fn explicit_dimensionless_vector_type_is_not_an_unresolved_derive_sentinel() { + let schema = TableSchema::new( + [ + ("id", ColumnDef::new("bigint")), + ("embedding", ColumnDef::new("vector")), + ], + ["id"], + ) + .unwrap(); + schema.validate_vector_dimensions().unwrap(); + } + + #[test] + fn table_schema_rejects_duplicate_columns() { + let error = TableSchema::new( + [ + ("id", ColumnDef::new("bigint")), + ("id", ColumnDef::new("text")), + ], + ["id"], + ) + .unwrap_err(); + assert!(error.to_string().contains("duplicate column \"id\"")); + } } diff --git a/rust/sdk/cocoindex/src/connectors/sqlite.rs b/rust/sdk/cocoindex/src/connectors/sqlite.rs index 40cf162a0..1ec3d1301 100644 --- a/rust/sdk/cocoindex/src/connectors/sqlite.rs +++ b/rust/sdk/cocoindex/src/connectors/sqlite.rs @@ -124,7 +124,11 @@ impl TableSchema { let name = name.into(); validate_ident(&name, "column name")?; validate_sqlite_type(&def.sqlite_type)?; - out.insert(name, def); + if out.insert(name.clone(), def).is_some() { + return Err(Error::engine(format!( + "SQLite table schema contains duplicate column {name:?}" + ))); + } } let primary_key: Vec = primary_key.into_iter().map(Into::into).collect(); if primary_key.is_empty() { @@ -1175,6 +1179,9 @@ fn sql_literal(value: &JsonValue, col: &ColumnDef) -> Result { }; } if is_real_type(&t) { + if let Some((secs, nanos)) = crate::row_schema::serialized_duration_parts(value) { + return Ok(format!("{secs}.{nanos:09}")); + } return match value { JsonValue::Number(n) => Ok(n.to_string()), _ => Err(Error::engine(format!( @@ -1186,6 +1193,9 @@ fn sql_literal(value: &JsonValue, col: &ColumnDef) -> Result { if is_text_type(&t) { return Ok(quote_string(value_to_string(value))); } + if is_blob_type(&t) { + return blob_literal(value); + } // Fallback (BLOB / unknown): store scalars directly, complex as JSON text. match value { JsonValue::String(s) => Ok(quote_string(s)), @@ -1195,6 +1205,33 @@ fn sql_literal(value: &JsonValue, col: &ColumnDef) -> Result { } } +fn blob_literal(value: &JsonValue) -> Result { + let bytes: Vec = match value { + JsonValue::Array(items) => items + .iter() + .map(|item| { + item.as_u64() + .filter(|value| *value <= 255) + .map(|value| value as u8) + .ok_or_else(|| Error::engine("BLOB array elements must be integers 0..=255")) + }) + .collect::>()?, + JsonValue::String(value) => value.as_bytes().to_vec(), + _ => { + return Err(Error::engine( + "BLOB column requires a byte array or string JSON value", + )); + } + }; + let mut literal = String::with_capacity(3 + bytes.len() * 2); + literal.push_str("X'"); + for byte in bytes { + literal.push_str(&format!("{byte:02x}")); + } + literal.push('\''); + Ok(literal) +} + fn value_to_string(value: &JsonValue) -> String { match value { JsonValue::String(s) => s.clone(), @@ -1243,6 +1280,10 @@ fn is_text_type(t: &str) -> bool { ["char", "text", "clob"].iter().any(|k| t.contains(k)) } +fn is_blob_type(t: &str) -> bool { + t.contains("blob") +} + fn validate_sqlite_type(value: &str) -> Result<()> { if value.is_empty() || !value @@ -1383,6 +1424,29 @@ mod tests { assert!(sql_literal(&JsonValue::from(1.5), &col).is_err()); } + #[test] + fn blob_literal_encodes_byte_arrays_as_blob_storage() { + let col = ColumnDef::new("BLOB"); + assert_eq!( + sql_literal(&serde_json::json!([0, 104, 105, 255]), &col).unwrap(), + "X'006869ff'" + ); + assert!(sql_literal(&serde_json::json!([256]), &col).is_err()); + } + + #[test] + fn table_schema_rejects_duplicate_columns() { + let error = TableSchema::new( + [ + ("id", ColumnDef::new("INTEGER")), + ("id", ColumnDef::new("TEXT")), + ], + ["id"], + ) + .unwrap_err(); + assert!(error.to_string().contains("duplicate column \"id\"")); + } + #[test] fn single_column_table_upsert_does_nothing_on_conflict() { let schema = TableSchema::new([("id", ColumnDef::new("INTEGER"))], ["id"]).unwrap(); diff --git a/rust/sdk/cocoindex/src/lib.rs b/rust/sdk/cocoindex/src/lib.rs index fe37ed441..7b8b437b3 100644 --- a/rust/sdk/cocoindex/src/lib.rs +++ b/rust/sdk/cocoindex/src/lib.rs @@ -50,6 +50,18 @@ pub use entity_resolution::{ resolve_entities_with_events, }; pub use error::{Error, Result}; +pub use resources::file::{ + FileContentCache, FileLike, FileMetadata, FilePath, FilePathMatcher, FileSourceItem, + MatchAllFilePathMatcher, PatternFilePathMatcher, +}; +pub use resources::fs::{ + DirTarget, DirTargetState, DirWalker, FileEntry, declare_dir_target, dir_target, + mount_dir_target, walk, walk_dir, walk_items, +}; +pub use resources::id::{ + IdGenerator, UuidGenerator, generate_id, generate_id_default, generate_uuid, + generate_uuid_default, +}; // Re-exported so `#[cocoindex::function]` output can register each function's // logic fingerprint without the user crate needing a direct `linkme` dependency. #[doc(hidden)] diff --git a/rust/sdk/cocoindex/src/row_schema.rs b/rust/sdk/cocoindex/src/row_schema.rs index 2397e3ff7..d839dca34 100644 --- a/rust/sdk/cocoindex/src/row_schema.rs +++ b/rust/sdk/cocoindex/src/row_schema.rs @@ -62,6 +62,17 @@ pub(crate) fn require_resolved_vector_dimension( Ok(()) } +#[cfg(any(feature = "postgres", feature = "sqlite", feature = "doris"))] +pub(crate) fn serialized_duration_parts(value: &serde_json::Value) -> Option<(u64, u32)> { + let object = value.as_object()?; + if object.len() != 2 { + return None; + } + let secs = object.get("secs")?.as_u64()?; + let nanos = u32::try_from(object.get("nanos")?.as_u64()?).ok()?; + (nanos < 1_000_000_000).then_some((secs, nanos)) +} + /// A connector-agnostic column type derived from a Rust field type. Each target /// connector maps these to its own SQL type strings. #[derive(Clone, Debug, PartialEq, Eq)] diff --git a/rust/sdk/cocoindex/tests/pipeline.rs b/rust/sdk/cocoindex/tests/pipeline.rs index f74222c94..6c477f572 100644 --- a/rust/sdk/cocoindex/tests/pipeline.rs +++ b/rust/sdk/cocoindex/tests/pipeline.rs @@ -1605,6 +1605,7 @@ mod batched_test { use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, LazyLock}; + use tokio::sync::Notify; static ITEMS_PROCESSED: AtomicUsize = AtomicUsize::new(0); @@ -1713,6 +1714,66 @@ mod batched_test { assert!(batches.iter().all(|(_, items)| items.len() <= 2)); } + static DEADLINE_BLOCKER_STARTED: Notify = Notify::const_new(); + static DEADLINE_BLOCKER_RELEASE: Notify = Notify::const_new(); + + #[cocoindex::function(batching, max_batch_size = 2)] + async fn deadline_isolated_batch( + ctx: &cocoindex::Ctx, + items: Vec, + ) -> cocoindex::Result> { + if items == [0] { + DEADLINE_BLOCKER_STARTED.notify_one(); + DEADLINE_BLOCKER_RELEASE.notified().await; + } else { + ctx.check_cancellation()?; + assert!( + !ctx.has_deadline(), + "a shared batch body must not inherit one caller's deadline" + ); + } + Ok(items) + } + + #[tokio::test] + async fn scheduled_batch_body_does_not_inherit_first_callers_deadline() { + let (app, _dir) = temp_app("batching_deadline_isolation").await; + + app.update(|ctx| async move { + let blocker_ctx = ctx.clone(); + let blocker = + tokio::spawn(async move { deadline_isolated_batch(&blocker_ctx, 0).await }); + DEADLINE_BLOCKER_STARTED.notified().await; + + // Queue a call whose deadline expires while the first physical + // batch is blocked, then fill the pending batch with an ordinary + // caller. The shared body must use neither caller's deadline. + let expiring_ctx = ctx.with_timeout(Duration::from_millis(10)); + let expiring = + tokio::spawn(async move { deadline_isolated_batch(&expiring_ctx, 1).await }); + sleep(Duration::from_millis(30)).await; + let ordinary = deadline_isolated_batch(&ctx, 2).await; + let expired = expiring.await.unwrap(); + + DEADLINE_BLOCKER_RELEASE.notify_one(); + let blocked = blocker.await.unwrap(); + + assert_eq!(expired?, 1); + assert_eq!(ordinary?, 2); + assert_eq!(blocked?, 0); + Ok(()) + }) + .await + .unwrap(); + } + + #[test] + #[should_panic(expected = "max_batch_size must be greater than zero")] + fn explicit_batched_rejects_zero_max_batch_size() { + let _ = + cocoindex::Batched::with_max_batch(|items: Vec| async move { Ok(items) }, 1, 0); + } + static MEMOIZED_ITEMS_PROCESSED: AtomicUsize = AtomicUsize::new(0); #[cocoindex::function(memo, batching, max_batch_size = 8)] diff --git a/rust/sdk/cocoindex/tests/postgres_target.rs b/rust/sdk/cocoindex/tests/postgres_target.rs index f6d8fd86f..a72d41666 100644 --- a/rust/sdk/cocoindex/tests/postgres_target.rs +++ b/rust/sdk/cocoindex/tests/postgres_target.rs @@ -1,9 +1,11 @@ #![cfg(feature = "postgres")] -use std::time::{SystemTime, UNIX_EPOCH}; +use std::str::FromStr as _; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use cocoindex::connectors::postgres; -use cocoindex::{Ctx, Environment, Result}; +use cocoindex::{Ctx, Environment, Result, SchemaFields}; +use rust_decimal::Decimal; use serde::{Deserialize, Serialize}; use sqlx::Row as _; @@ -811,6 +813,85 @@ async fn postgres_bytea_round_trips_byte_arrays_when_available() -> Result<()> { Ok(()) } +#[derive(Clone, Serialize, SchemaFields)] +struct DecimalRow { + id: i64, + amount: Decimal, + elapsed: Duration, +} + +async fn declare_decimals(ctx: Ctx, schema: String, rows: Vec) -> Result<()> { + let table = postgres::mount_table_target( + &ctx, + &PG, + "decimals", + postgres::TableSchema::from_row::(["id"])?, + Some(&schema), + ) + .await?; + for row in &rows { + table.declare_row(&ctx, row)?; + } + Ok(()) +} + +#[tokio::test] +async fn postgres_decimal_and_duration_round_trip_when_available() -> Result<()> { + let Ok(url) = std::env::var("POSTGRES_URL") else { + eprintln!("skipping live Postgres decimal test; POSTGRES_URL is not set"); + return Ok(()); + }; + let nonce = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + let schema = format!("cocoindex_rust_decimal_{nonce}"); + let db = postgres::Database::connect(&url).await?; + let tempdir = tempfile::tempdir().unwrap(); + let app = Environment::builder() + .db_path(tempdir.path().join(".cocoindex_db")) + .provide_key(&PG, db.clone()) + .build() + .await? + .app("PostgresDecimalTest") + .await?; + + app.run({ + let schema = schema.clone(); + move |ctx| { + declare_decimals( + ctx, + schema, + vec![DecimalRow { + id: 1, + amount: Decimal::from_str("1234567890.012300").unwrap(), + elapsed: Duration::new(2, 345_678_000), + }], + ) + } + }) + .await?; + + let row = sqlx::query(&format!( + "SELECT amount::text AS amount, EXTRACT(EPOCH FROM elapsed)::double precision AS elapsed \ + FROM \"{schema}\".\"decimals\" WHERE id = 1" + )) + .fetch_one(db.pool()) + .await + .map_err(|e| cocoindex::Error::engine(format!("postgres query: {e}")))?; + assert_eq!( + row.try_get::("amount").unwrap(), + "1234567890.012300" + ); + assert_eq!(row.try_get::("elapsed").unwrap(), 2.345678); + + sqlx::query(&format!("DROP SCHEMA IF EXISTS \"{schema}\" CASCADE")) + .execute(db.pool()) + .await + .map_err(|e| cocoindex::Error::engine(format!("postgres cleanup: {e}")))?; + Ok(()) +} + async fn declare_evo_table( ctx: Ctx, schema: String, diff --git a/rust/sdk/cocoindex/tests/schema_from_row.rs b/rust/sdk/cocoindex/tests/schema_from_row.rs index ffc168a8e..852abd46f 100644 --- a/rust/sdk/cocoindex/tests/schema_from_row.rs +++ b/rust/sdk/cocoindex/tests/schema_from_row.rs @@ -58,6 +58,8 @@ fn doris_from_row_matches_explicit_schema() { count: i64, score: f64, active: bool, + blob: Vec, + elapsed: std::time::Duration, #[coco(vector = 4)] embedding: Vec, #[coco(json)] @@ -78,6 +80,8 @@ fn doris_from_row_matches_explicit_schema() { ("count", ColumnDef::new("BIGINT").not_null()), ("score", ColumnDef::new("DOUBLE").not_null()), ("active", ColumnDef::new("BOOLEAN").not_null()), + ("blob", ColumnDef::new("STRING").not_null()), + ("elapsed", ColumnDef::new("BIGINT").not_null()), ("embedding", ColumnDef::vector(4)), ("tags", ColumnDef::new("JSON").not_null()), ("code", ColumnDef::new("VARCHAR(10)").not_null()), @@ -101,6 +105,7 @@ fn sqlite_from_row_matches_explicit_schema() { name: Option, score: f64, blob: Vec, + elapsed: std::time::Duration, #[coco(vector)] embedding: Vec, } @@ -116,6 +121,7 @@ fn sqlite_from_row_matches_explicit_schema() { ("name", ColumnDef::new("TEXT")), ("score", ColumnDef::new("REAL").not_null()), ("blob", ColumnDef::new("BLOB").not_null()), + ("elapsed", ColumnDef::new("REAL").not_null()), ("embedding", ColumnDef::new("float[3]").not_null()), ], ["id"], @@ -124,6 +130,24 @@ fn sqlite_from_row_matches_explicit_schema() { assert_eq!(got, want); } +#[cfg(feature = "sqlite")] +#[test] +fn schema_from_row_rejects_duplicate_renamed_columns() { + use cocoindex::SchemaFields; + use cocoindex::connectors::sqlite::TableSchema; + + #[derive(SchemaFields)] + #[allow(dead_code)] + struct Row { + id: i64, + #[coco(rename = "id")] + other_id: i64, + } + + let error = TableSchema::from_row::(["id"]).unwrap_err(); + assert!(error.to_string().contains("duplicate column \"id\"")); +} + #[cfg(feature = "postgres")] #[test] fn postgres_from_row_matches_explicit_schema() { @@ -137,6 +161,7 @@ fn postgres_from_row_matches_explicit_schema() { title: Option, views: i32, ratio: f32, + elapsed: std::time::Duration, #[coco(vector)] embedding: Vec, #[coco(vector = 8, half)] @@ -157,6 +182,7 @@ fn postgres_from_row_matches_explicit_schema() { ("title", ColumnDef::new("text").nullable()), ("views", ColumnDef::new("integer")), ("ratio", ColumnDef::new("real")), + ("elapsed", ColumnDef::new("interval")), ("embedding", ColumnDef::new("vector(8)")), ("embedding_half", ColumnDef::new("halfvec(8)")), ("meta", ColumnDef::new("jsonb")), @@ -433,6 +459,8 @@ async fn sqlite_from_row_round_trips_a_row() -> cocoindex::Result<()> { id: i64, name: String, score: f64, + data: Vec, + elapsed: std::time::Duration, } let tmp = tempfile::tempdir().unwrap(); @@ -444,11 +472,15 @@ async fn sqlite_from_row_round_trips_a_row() -> cocoindex::Result<()> { id: 1, name: "a".into(), score: 1.5, + data: vec![0, 104, 105, 255], + elapsed: std::time::Duration::new(1, 250_000_000), }, Item { id: 2, name: "b".into(), score: 2.5, + data: Vec::new(), + elapsed: std::time::Duration::new(2, 500_000_000), }, ]; @@ -472,23 +504,28 @@ async fn sqlite_from_row_round_trips_a_row() -> cocoindex::Result<()> { }) .await?; - let fetched = sqlx::query("SELECT id, name, score FROM items ORDER BY id") + let fetched = sqlx::query("SELECT id, name, score, data, elapsed FROM items ORDER BY id") .fetch_all(db.pool()) .await .unwrap(); - let got: Vec<(i64, String, f64)> = fetched + let got: Vec<(i64, String, f64, Vec, f64)> = fetched .iter() .map(|r| { ( r.get::("id"), r.get::("name"), r.get::("score"), + r.get::, _>("data"), + r.get::("elapsed"), ) }) .collect(); assert_eq!( got, - vec![(1, "a".to_string(), 1.5), (2, "b".to_string(), 2.5)] + vec![ + (1, "a".to_string(), 1.5, vec![0, 104, 105, 255], 1.25), + (2, "b".to_string(), 2.5, Vec::new(), 2.5), + ] ); Ok(()) } diff --git a/rust/sdk/cocoindex_macros/src/lib.rs b/rust/sdk/cocoindex_macros/src/lib.rs index b25747ff8..c29b18192 100644 --- a/rust/sdk/cocoindex_macros/src/lib.rs +++ b/rust/sdk/cocoindex_macros/src/lib.rs @@ -1421,6 +1421,7 @@ fn logical_type_tokens(ty: &Type, attr: &SchemaFieldAttr) -> TokenStream2 { Some("NaiveTime") => quote! { Time }, Some("NaiveDateTime" | "DateTime") => quote! { DateTime }, Some("Decimal") => quote! { Decimal }, + Some("Duration") => quote! { Duration }, // `Vec` is bytes; any other `Vec<_>` falls through to JSON. Some("Vec") if first_generic_ident(ty).as_deref() == Some("u8") => quote! { Bytes }, // Everything else (collections, maps, nested structs, enums) → JSON. @@ -1713,7 +1714,7 @@ mod tests { } #[test] - fn batching_memo_key_includes_body_hash() { + fn batching_memo_key_changes_with_body() { let first: ItemFn = parse_str( "async fn embed(ctx: &Ctx, items: Vec) -> Result> { Ok(items) }", ) @@ -1732,7 +1733,8 @@ mod tests { let args = FunctionArgs::parse(quote!(memo, batching)).unwrap(); let hash_const_name = format_ident!("__COCO_FN_HASH_EMBED"); let first_hash = compute_code_hash(&first.block, None); - let expanded = expand_batching_function( + let second_hash = compute_code_hash(&second.block, None); + let first_expanded = expand_batching_function( &first, &args, &hash_const_name, @@ -1743,10 +1745,25 @@ mod tests { ) .unwrap() .to_string(); - assert!( - expanded - .contains("write_key_fingerprint_part (& mut __coco_key , & __COCO_FN_HASH_EMBED)"), - "the expanded batching memo key must include the body-derived hash: {expanded}" + let second_expanded = expand_batching_function( + &second, + &args, + &hash_const_name, + "e!(#second_hash), + "e!(), + true, + true, + ) + .unwrap() + .to_string(); + let key_part = "write_key_fingerprint_part (& mut __coco_key , & __COCO_FN_HASH_EMBED)"; + assert!(first_expanded.contains(key_part)); + assert!(second_expanded.contains(key_part)); + assert!(first_expanded.contains(&first_hash.to_string())); + assert!(second_expanded.contains(&second_hash.to_string())); + assert_ne!( + first_hash, second_hash, + "two batching bodies must produce different memo-key constants" ); } From 587c6f15be1e6f7b4c4391aa31a8eb52380b0ffb Mon Sep 17 00:00:00 2001 From: zhhanz Date: Wed, 15 Jul 2026 00:29:41 -0700 Subject: [PATCH 11/15] Simplify Rust SDK batching and schema plumbing --- .../docs/getting_started/rust_quickstart.mdx | 6 +- .../rust/text_embedding_qdrant/src/main.rs | 22 +- .../text_embedding_turbopuffer/src/main.rs | 17 +- rust/sdk/cocoindex/src/batched.rs | 115 +++---- rust/sdk/cocoindex/src/connectors/doris.rs | 35 +- rust/sdk/cocoindex/src/connectors/lancedb.rs | 32 +- rust/sdk/cocoindex/src/connectors/postgres.rs | 32 +- rust/sdk/cocoindex/src/connectors/qdrant.rs | 37 +- rust/sdk/cocoindex/src/connectors/sqlite.rs | 44 +-- .../cocoindex/src/connectors/turbopuffer.rs | 34 +- rust/sdk/cocoindex/src/row_schema.rs | 59 +++- rust/sdk/cocoindex/tests/pipeline.rs | 7 - rust/sdk/cocoindex/tests/schema_from_row.rs | 238 +++++-------- rust/sdk/cocoindex_macros/src/lib.rs | 323 +++++++----------- 14 files changed, 415 insertions(+), 586 deletions(-) diff --git a/docs/src/content/docs/getting_started/rust_quickstart.mdx b/docs/src/content/docs/getting_started/rust_quickstart.mdx index 9bb7592c9..46ace0062 100644 --- a/docs/src/content/docs/getting_started/rust_quickstart.mdx +++ b/docs/src/content/docs/getting_started/rust_quickstart.mdx @@ -1,5 +1,5 @@ --- -title: Rust SDK *quickstart* +title: Rust SDK quickstart description: > Build an incremental Rust pipeline that chunks Markdown, batches and caches local embeddings, and keeps a Postgres/pgvector table in sync. @@ -195,10 +195,6 @@ unchanged file can skip the entire component on the next run. If a file or chunk disappears, the rows its component used to own are removed during table reconciliation. -For normal function-shaped components, prefer `use_mount!` and `mount_each!`. -Use `ctx.scope` or `ctx.mount_each` when you deliberately want an explicit-key -closure without the macros' function-and-argument fingerprint. - ## Build the environment and run Finish `src/main.rs` by loading the resources and providing them to an diff --git a/examples/rust/text_embedding_qdrant/src/main.rs b/examples/rust/text_embedding_qdrant/src/main.rs index ab063dc91..9d4635eb6 100644 --- a/examples/rust/text_embedding_qdrant/src/main.rs +++ b/examples/rust/text_embedding_qdrant/src/main.rs @@ -15,14 +15,16 @@ use std::path::PathBuf; -use cocoindex::connectors::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; +use cocoindex::connectors::qdrant::{ + self, CollectionSchema, Distance, NamedPointVector, QdrantConnection, +}; use cocoindex::ops::sentence_transformers::SentenceTransformerEmbedder; use cocoindex::ops::text::{RecursiveChunkConfig, RecursiveSplitter}; use cocoindex::prelude::*; use serde_json::json; const EMBED_MODEL: &str = "sentence-transformers/all-MiniLM-L6-v2"; -const EMBED_DIM: u64 = 384; +const EMBED_DIM: usize = 384; const COLLECTION: &str = "TextEmbedding"; const TOP_K: u64 = 5; const CHUNK_SIZE: usize = 2000; @@ -38,9 +40,10 @@ cocoindex::context_key!( ); /// A computed point: id + vector + payload fields. -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct PointData { id: u64, + #[coco(vector)] vector: Vec, filename: String, chunk_start: i64, @@ -87,12 +90,12 @@ async fn process_file(ctx: &Ctx, file: FileEntry) -> Result> { } async fn app_main(ctx: Ctx, sourcedir: PathBuf) -> Result<()> { - let conn = ctx.get_key(&DB)?; let target = qdrant::mount_collection_target( &ctx, &DB, COLLECTION, - CollectionSchema::new(EMBED_DIM, Distance::Cosine), + CollectionSchema::from_row::(Distance::Cosine)? + .with_vector_dim("vector", EMBED_DIM)?, ) .await?; @@ -118,7 +121,12 @@ async fn app_main(ctx: Ctx, sourcedir: PathBuf) -> Result<()> { .as_object() .unwrap() .clone(); - target.declare_point(&ctx, p.id, p.vector.clone(), payload)?; + target.declare_named_vectors_point( + &ctx, + p.id, + [("vector", NamedPointVector::Single(p.vector.clone()))], + payload, + )?; } } println!("indexed {count} chunk(s) total"); @@ -131,7 +139,7 @@ async fn query_once( query: &str, ) -> Result<()> { let query_vec = Embedder::embed(embedder, query).await?; - let hits = qdrant::vector_search(conn, COLLECTION, query_vec, TOP_K).await?; + let hits = qdrant::named_vector_search(conn, COLLECTION, "vector", query_vec, TOP_K).await?; for hit in hits { let filename = hit .payload diff --git a/examples/rust/text_embedding_turbopuffer/src/main.rs b/examples/rust/text_embedding_turbopuffer/src/main.rs index bdc755f63..9f6a9a8f6 100644 --- a/examples/rust/text_embedding_turbopuffer/src/main.rs +++ b/examples/rust/text_embedding_turbopuffer/src/main.rs @@ -41,9 +41,10 @@ cocoindex::context_key!( state = SentenceTransformerEmbedder::model_name ); -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, SchemaFields)] struct RowData { id: String, + #[coco(vector)] vector: Vec, filename: String, chunk_start: i64, @@ -90,12 +91,12 @@ async fn process_file(ctx: &Ctx, file: FileEntry) -> Result> { } async fn app_main(ctx: Ctx, sourcedir: PathBuf, namespace: String) -> Result<()> { - let conn = ctx.get_key(&DB)?; let target = turbopuffer::mount_namespace_target( &ctx, &DB, &namespace, - NamespaceSchema::new(EMBED_DIM, DistanceMetric::CosineDistance), + NamespaceSchema::from_row::(DistanceMetric::CosineDistance)? + .with_vector_dim("vector", EMBED_DIM)?, ) .await?; @@ -121,7 +122,12 @@ async fn app_main(ctx: Ctx, sourcedir: PathBuf, namespace: String) -> Result<()> .as_object() .unwrap() .clone(); - target.declare_row(&ctx, r.id.clone(), r.vector.clone(), attributes)?; + target.declare_named_row( + &ctx, + r.id.clone(), + [("vector", r.vector.clone())], + attributes, + )?; } } println!("indexed {count} chunk(s) total"); @@ -135,7 +141,8 @@ async fn query_once( query: &str, ) -> Result<()> { let query_vec = Embedder::embed(embedder, query).await?; - let hits = turbopuffer::vector_search(conn, namespace, query_vec, TOP_K).await?; + let hits = turbopuffer::vector_search_by_field(conn, namespace, "vector", query_vec, TOP_K) + .await?; for hit in hits { let filename = hit .attributes diff --git a/rust/sdk/cocoindex/src/batched.rs b/rust/sdk/cocoindex/src/batched.rs index 6285a500b..0e5141c55 100644 --- a/rust/sdk/cocoindex/src/batched.rs +++ b/rust/sdk/cocoindex/src/batched.rs @@ -108,29 +108,31 @@ struct ScheduledBatchKey { type FunctionBatcher = Batcher>; type ScheduledBatchers = Mutex>>>; -enum BatchedMode +fn new_batcher( + f: BatchFn, + options: BatchingOptions, +) -> Arc> where In: Send + 'static, Out: Send + 'static, { - Fixed(Arc>), - Scheduled { - options: BatchingOptions, - batchers: ScheduledBatchers, - }, + Arc::new(Batcher::new( + FnRunner { f }, + Arc::new(BatchQueue::new()), + options, + )) } /// A single-item interface over a batch-shaped async function. /// -/// [`Batched::new`] retains the explicit, memoized adapter API. The -/// `#[cocoindex::function(batching)]` macro uses the hidden scheduled mode so -/// each distinct extra-argument set gets its own short-lived batcher. +/// `#[cocoindex::function(batching)]` is the usual entry point. [`Batched::new`] +/// is the explicit, memoized adapter for code that cannot use the macro. pub struct Batched where In: Send + 'static, Out: Send + 'static, { - mode: BatchedMode, + batcher: Arc>, code_hash: u64, } @@ -146,29 +148,6 @@ where /// folded into each item's memo key, so editing the batch logic invalidates /// cached results. pub fn new(f: F, code_hash: u64) -> Self - where - F: Fn(Vec) -> Fut + Send + Sync + 'static, - Fut: Future>> + Send + 'static, - { - Self::with_options(f, code_hash, BatchingOptions::default()) - } - - /// Like [`Batched::new`], but caps how many items are processed per batch. - pub fn with_max_batch(f: F, code_hash: u64, max_batch_size: usize) -> Self - where - F: Fn(Vec) -> Fut + Send + Sync + 'static, - Fut: Future>> + Send + 'static, - { - Self::with_options( - f, - code_hash, - BatchingOptions { - max_batch_size: Some(max_batch_size), - }, - ) - } - - fn with_options(f: F, code_hash: u64, options: BatchingOptions) -> Self where F: Fn(Vec) -> Fut + Send + Sync + 'static, Fut: Future>> + Send + 'static, @@ -177,9 +156,8 @@ where let fut = f(inputs); Box::pin(async move { fut.await.map_err(Error::into_core) }) }); - let batcher = Self::new_batcher(wrapped, options); Self { - mode: BatchedMode::Fixed(batcher), + batcher: new_batcher(wrapped, BatchingOptions::default()), code_hash, } } @@ -190,12 +168,7 @@ where pub async fn call(&self, ctx: &Ctx, item: In) -> Result { let fp = crate::memo::key_fingerprint_result(&("cocoindex_batched", self.code_hash, &item))?; - let BatchedMode::Fixed(batcher) = &self.mode else { - return Err(Error::engine( - "scheduled Batched instances must be called by #[cocoindex::function(batching)]", - )); - }; - let batcher = batcher.clone(); + let batcher = self.batcher.clone(); // A batch impl is ctx-free, so it makes no tracked child `#[function]` // calls; the `propagate_children_fn_logic` flag is therefore inert here. // Pass `true` (the default) — only the batch impl's own `code_hash` @@ -213,30 +186,33 @@ where } } -impl Batched +/// Scheduler backing generated `#[cocoindex::function(batching)]` wrappers. +#[doc(hidden)] +pub struct __ScheduledBatched where In: Send + 'static, Out: Send + 'static, { - fn new_batcher(f: BatchFn, options: BatchingOptions) -> Arc> { - Self::validate_options(&options); - Arc::new(Batcher::new( - FnRunner { f }, - Arc::new(BatchQueue::new()), - options, - )) - } + options: BatchingOptions, + batchers: ScheduledBatchers, + code_hash: u64, +} +impl __ScheduledBatched +where + In: Send + 'static, + Out: Send + 'static, +{ /// Construct the scheduler used by generated batching wrappers. #[doc(hidden)] - pub fn __scheduled(code_hash: u64) -> Self { - Self::__scheduled_with_options(code_hash, BatchingOptions::default()) + pub fn __new(code_hash: u64) -> Self { + Self::with_options(code_hash, BatchingOptions::default()) } /// Construct the generated scheduler with a maximum batch size. #[doc(hidden)] - pub fn __scheduled_with_max_batch(code_hash: u64, max_batch_size: usize) -> Self { - Self::__scheduled_with_options( + pub fn __with_max_batch(code_hash: u64, max_batch_size: usize) -> Self { + Self::with_options( code_hash, BatchingOptions { max_batch_size: Some(max_batch_size), @@ -244,26 +220,14 @@ where ) } - fn __scheduled_with_options(code_hash: u64, options: BatchingOptions) -> Self { - Self::validate_options(&options); + fn with_options(code_hash: u64, options: BatchingOptions) -> Self { Self { - mode: BatchedMode::Scheduled { - options, - batchers: Mutex::new(HashMap::new()), - }, + options, + batchers: Mutex::new(HashMap::new()), code_hash, } } - fn validate_options(options: &BatchingOptions) { - if let Some(max_batch_size) = options.max_batch_size { - assert!( - max_batch_size > 0, - "max_batch_size must be greater than zero" - ); - } - } - /// Schedule one generated function call. `extra_args_key` keeps calls with /// different captured arguments in separate batches. #[doc(hidden)] @@ -278,12 +242,6 @@ where F: Fn(Ctx, Vec) -> Fut + Send + Sync + 'static, Fut: Future>> + Send + 'static, { - let BatchedMode::Scheduled { options, batchers } = &self.mode else { - return Err(Error::engine( - "explicit Batched instances must be called with Batched::call", - )); - }; - let key = ScheduledBatchKey { // A static generated scheduler may be used by multiple apps in the // same process. Never let their context-bound calls share a body. @@ -296,12 +254,15 @@ where Box::pin(async move { fut.await.map_err(Error::into_core) }) }); let batcher = { - let mut batchers = batchers.lock().expect("batch scheduler mutex poisoned"); + let mut batchers = self + .batchers + .lock() + .expect("batch scheduler mutex poisoned"); batchers.retain(|_, batcher| batcher.strong_count() != 0); if let Some(batcher) = batchers.get(&key).and_then(Weak::upgrade) { batcher } else { - let batcher = Self::new_batcher(f, options.clone()); + let batcher = new_batcher(f, self.options.clone()); batchers.insert(key, Arc::downgrade(&batcher)); batcher } diff --git a/rust/sdk/cocoindex/src/connectors/doris.rs b/rust/sdk/cocoindex/src/connectors/doris.rs index be754a597..406748d19 100644 --- a/rust/sdk/cocoindex/src/connectors/doris.rs +++ b/rust/sdk/cocoindex/src/connectors/doris.rs @@ -351,24 +351,27 @@ impl TableSchema { /// Resolve or override the dimension of a vector field derived from a row. pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { let dim = u32::try_from(dim).map_err(|_| { - Error::engine(format!( - "Doris vector field {field_name:?} dimension does not fit u32" - )) + crate::row_schema::vector_dimension_error( + "Doris", + field_name, + "dimension does not fit u32", + ) })?; if dim == 0 { - return Err(Error::engine(format!( - "Doris vector field {field_name:?} requires a dimension greater than zero" - ))); + return Err(crate::row_schema::vector_dimension_error( + "Doris", + field_name, + "requires a dimension greater than zero", + )); } - let def = self.columns.get_mut(field_name).ok_or_else(|| { - Error::engine(format!( - "Doris vector dimension override names unknown field {field_name:?}" - )) - })?; + let def = self + .columns + .get_mut(field_name) + .ok_or_else(|| crate::row_schema::unknown_vector_field_error("Doris", field_name))?; if !def.is_vector { - return Err(Error::engine(format!( - "Doris field {field_name:?} is not a vector field" - ))); + return Err(crate::row_schema::not_vector_field_error( + "Doris", field_name, + )); } def.vector_dimension = Some(dim); Ok(self) @@ -378,7 +381,9 @@ impl TableSchema { for (name, def) in &self.columns { if def.is_vector { let dim = def.vector_dimension.unwrap_or(0) as usize; - crate::row_schema::require_resolved_vector_dimension("Doris", name, dim)?; + if dim == 0 { + crate::row_schema::require_resolved_vector_dimension("Doris", name)?; + } } } Ok(()) diff --git a/rust/sdk/cocoindex/src/connectors/lancedb.rs b/rust/sdk/cocoindex/src/connectors/lancedb.rs index 46ac3c996..3ae473645 100644 --- a/rust/sdk/cocoindex/src/connectors/lancedb.rs +++ b/rust/sdk/cocoindex/src/connectors/lancedb.rs @@ -205,27 +205,25 @@ impl TableSchema { /// Resolve or override the dimension of a vector field derived from a row. pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { if dim == 0 || i32::try_from(dim).is_err() { - return Err(Error::engine(format!( - "LanceDB vector field {field_name:?} requires a dimension in 1..=i32::MAX" - ))); + return Err(crate::row_schema::vector_dimension_error( + "LanceDB", + field_name, + "requires a dimension in 1..=i32::MAX", + )); } let (_, def) = self .columns .iter_mut() .find(|(name, _)| name == field_name) - .ok_or_else(|| { - Error::engine(format!( - "LanceDB vector dimension override names unknown field {field_name:?}" - )) - })?; + .ok_or_else(|| crate::row_schema::unknown_vector_field_error("LanceDB", field_name))?; match &mut def.col_type { ColumnType::Vector(current_dim) | ColumnType::HalfVector(current_dim) => { *current_dim = dim; } _ => { - return Err(Error::engine(format!( - "LanceDB field {field_name:?} is not a vector field" - ))); + return Err(crate::row_schema::not_vector_field_error( + "LanceDB", field_name, + )); } } Ok(self) @@ -237,11 +235,15 @@ impl TableSchema { ColumnType::Vector(dim) | ColumnType::HalfVector(dim) => dim, _ => continue, }; - crate::row_schema::require_resolved_vector_dimension("LanceDB", name, dim)?; + if dim == 0 { + crate::row_schema::require_resolved_vector_dimension("LanceDB", name)?; + } if i32::try_from(dim).is_err() { - return Err(Error::engine(format!( - "LanceDB vector field {name:?} requires a dimension in 1..=i32::MAX" - ))); + return Err(crate::row_schema::vector_dimension_error( + "LanceDB", + name, + "requires a dimension in 1..=i32::MAX", + )); } } Ok(()) diff --git a/rust/sdk/cocoindex/src/connectors/postgres.rs b/rust/sdk/cocoindex/src/connectors/postgres.rs index 4943fcaa5..14e055e3f 100644 --- a/rust/sdk/cocoindex/src/connectors/postgres.rs +++ b/rust/sdk/cocoindex/src/connectors/postgres.rs @@ -163,20 +163,18 @@ impl TableSchema { /// Resolve or override the dimension of a vector field derived from a row. pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { if dim == 0 { - return Err(Error::engine(format!( - "Postgres vector field {field_name:?} requires a dimension greater than zero" - ))); + return Err(crate::row_schema::vector_dimension_error( + "Postgres", + field_name, + "requires a dimension greater than zero", + )); } - let def = self.columns.get_mut(field_name).ok_or_else(|| { - Error::engine(format!( - "Postgres vector dimension override names unknown field {field_name:?}" - )) - })?; - let base = pgvector_type_base(&def.pg_type).ok_or_else(|| { - Error::engine(format!( - "Postgres field {field_name:?} is not a vector field" - )) - })?; + let def = self + .columns + .get_mut(field_name) + .ok_or_else(|| crate::row_schema::unknown_vector_field_error("Postgres", field_name))?; + let base = pgvector_type_base(&def.pg_type) + .ok_or_else(|| crate::row_schema::not_vector_field_error("Postgres", field_name))?; def.pg_type = format!("{base}({dim})"); self.unresolved_vector_columns.remove(field_name); Ok(self) @@ -184,7 +182,7 @@ impl TableSchema { fn validate_vector_dimensions(&self) -> Result<()> { for name in &self.unresolved_vector_columns { - crate::row_schema::require_resolved_vector_dimension("Postgres", name, 0)?; + crate::row_schema::require_resolved_vector_dimension("Postgres", name)?; } Ok(()) } @@ -1704,10 +1702,7 @@ fn is_numeric_type(lower: &str) -> bool { | "float4" | "double precision" | "float8" - | "numeric" - | "decimal" - ) || lower.starts_with("numeric(") - || lower.starts_with("decimal(") + ) || is_decimal_type(lower) } fn is_decimal_type(lower: &str) -> bool { @@ -2059,7 +2054,6 @@ mod review_fix_tests { #[test] fn numeric_literal_accepts_rust_decimal_string_encoding() { let value = serde_json::to_value(Decimal::from_str("1234567890.012300").unwrap()).unwrap(); - assert_eq!(value, JsonValue::String("1234567890.012300".to_string())); assert_eq!( sql_literal(&value, &ColumnDef::new("numeric(28, 6)")).unwrap(), "1234567890.012300" diff --git a/rust/sdk/cocoindex/src/connectors/qdrant.rs b/rust/sdk/cocoindex/src/connectors/qdrant.rs index a60d3f5b9..352505dbb 100644 --- a/rust/sdk/cocoindex/src/connectors/qdrant.rs +++ b/rust/sdk/cocoindex/src/connectors/qdrant.rs @@ -370,22 +370,20 @@ impl CollectionSchema { /// built with [`CollectionSchema::new`] or [`CollectionSchema::multivector`], /// the single unnamed vector uses the field name `"vector"`. pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { - validate_vector_size(dim).map_err(|err| { - Error::engine(format!( - "Qdrant vector field {field_name:?} has invalid dimension: {err}" - )) - })?; + if dim == 0 { + return Err(crate::row_schema::vector_dimension_error( + "Qdrant", + field_name, + "has invalid dimension: qdrant vector size must be greater than zero", + )); + } let mut fields = self.vector_fields(); let def = match &mut fields { VectorFields::Named(vectors) => vectors.get_mut(field_name), VectorFields::Single(def) if field_name == "vector" => Some(def), VectorFields::Single(_) => None, } - .ok_or_else(|| { - Error::engine(format!( - "Qdrant vector dimension override names unknown vector field {field_name:?}" - )) - })?; + .ok_or_else(|| crate::row_schema::unknown_vector_field_error("Qdrant", field_name))?; match &mut def.schema { QdrantVectorSchema::Dense(schema) => schema.size = dim, QdrantVectorSchema::Multi(schema) => schema.vector_schema.size = dim, @@ -401,18 +399,17 @@ impl CollectionSchema { fn validate_vector_dimensions(&self) -> Result<()> { match self.vector_fields() { - VectorFields::Single(def) => crate::row_schema::require_resolved_vector_dimension( - "Qdrant", - "vector", - def.vector_size() as usize, - ), + VectorFields::Single(def) => { + if def.vector_size() == 0 { + crate::row_schema::require_resolved_vector_dimension("Qdrant", "vector")?; + } + Ok(()) + } VectorFields::Named(vectors) => { for (name, def) in vectors { - crate::row_schema::require_resolved_vector_dimension( - "Qdrant", - &name, - def.vector_size() as usize, - )?; + if def.vector_size() == 0 { + crate::row_schema::require_resolved_vector_dimension("Qdrant", &name)?; + } } Ok(()) } diff --git a/rust/sdk/cocoindex/src/connectors/sqlite.rs b/rust/sdk/cocoindex/src/connectors/sqlite.rs index 1ec3d1301..cb5c2956f 100644 --- a/rust/sdk/cocoindex/src/connectors/sqlite.rs +++ b/rust/sdk/cocoindex/src/connectors/sqlite.rs @@ -173,19 +173,20 @@ impl TableSchema { /// Resolve or override the dimension of a vector field derived from a row. pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { if dim == 0 { - return Err(Error::engine(format!( - "SQLite vector field {field_name:?} requires a dimension greater than zero" - ))); + return Err(crate::row_schema::vector_dimension_error( + "SQLite", + field_name, + "requires a dimension greater than zero", + )); } - let def = self.columns.get_mut(field_name).ok_or_else(|| { - Error::engine(format!( - "SQLite vector dimension override names unknown field {field_name:?}" - )) - })?; + let def = self + .columns + .get_mut(field_name) + .ok_or_else(|| crate::row_schema::unknown_vector_field_error("SQLite", field_name))?; if !(def.sqlite_type.starts_with("float[") && def.sqlite_type.ends_with(']')) { - return Err(Error::engine(format!( - "SQLite field {field_name:?} is not a vector field" - ))); + return Err(crate::row_schema::not_vector_field_error( + "SQLite", field_name, + )); } def.sqlite_type = format!("float[{dim}]"); Ok(self) @@ -194,7 +195,7 @@ impl TableSchema { fn validate_vector_dimensions(&self) -> Result<()> { for (name, def) in &self.columns { if def.sqlite_type == "float[0]" { - crate::row_schema::require_resolved_vector_dimension("SQLite", name, 0)?; + crate::row_schema::require_resolved_vector_dimension("SQLite", name)?; } } Ok(()) @@ -1425,28 +1426,11 @@ mod tests { } #[test] - fn blob_literal_encodes_byte_arrays_as_blob_storage() { + fn blob_literal_rejects_out_of_range_bytes() { let col = ColumnDef::new("BLOB"); - assert_eq!( - sql_literal(&serde_json::json!([0, 104, 105, 255]), &col).unwrap(), - "X'006869ff'" - ); assert!(sql_literal(&serde_json::json!([256]), &col).is_err()); } - #[test] - fn table_schema_rejects_duplicate_columns() { - let error = TableSchema::new( - [ - ("id", ColumnDef::new("INTEGER")), - ("id", ColumnDef::new("TEXT")), - ], - ["id"], - ) - .unwrap_err(); - assert!(error.to_string().contains("duplicate column \"id\"")); - } - #[test] fn single_column_table_upsert_does_nothing_on_conflict() { let schema = TableSchema::new([("id", ColumnDef::new("INTEGER"))], ["id"]).unwrap(); diff --git a/rust/sdk/cocoindex/src/connectors/turbopuffer.rs b/rust/sdk/cocoindex/src/connectors/turbopuffer.rs index fc3a34c3e..703b1a135 100644 --- a/rust/sdk/cocoindex/src/connectors/turbopuffer.rs +++ b/rust/sdk/cocoindex/src/connectors/turbopuffer.rs @@ -304,9 +304,11 @@ impl NamespaceSchema { /// Resolve or override one named vector field's runtime dimension. pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { if dim == 0 { - return Err(Error::engine(format!( - "Turbopuffer vector field {field_name:?} requires a dimension greater than zero" - ))); + return Err(crate::row_schema::vector_dimension_error( + "Turbopuffer", + field_name, + "requires a dimension greater than zero", + )); } let mut fields = self.vector_fields(); let def = match &mut fields { @@ -314,11 +316,7 @@ impl NamespaceSchema { VectorFields::Single(def) if field_name == DEFAULT_VECTOR_FIELD => Some(def), VectorFields::Single(_) => None, } - .ok_or_else(|| { - Error::engine(format!( - "Turbopuffer vector dimension override names unknown vector field {field_name:?}" - )) - })?; + .ok_or_else(|| crate::row_schema::unknown_vector_field_error("Turbopuffer", field_name))?; def.schema.size = dim; vector_type_str(&def.schema)?; self.vectors = Some(fields); @@ -331,21 +329,23 @@ impl NamespaceSchema { fn validate_vector_dimensions(&self) -> Result<()> { match self.vector_fields() { - VectorFields::Single(def) => crate::row_schema::require_resolved_vector_dimension( - "Turbopuffer", - DEFAULT_VECTOR_FIELD, - def.schema.size, - ), - VectorFields::Named(vectors) => { - for (name, def) in vectors { + VectorFields::Single(def) => { + if def.schema.size == 0 { crate::row_schema::require_resolved_vector_dimension( "Turbopuffer", - &name, - def.schema.size, + DEFAULT_VECTOR_FIELD, )?; } Ok(()) } + VectorFields::Named(vectors) => { + for (name, def) in vectors { + if def.schema.size == 0 { + crate::row_schema::require_resolved_vector_dimension("Turbopuffer", &name)?; + } + } + Ok(()) + } } } diff --git a/rust/sdk/cocoindex/src/row_schema.rs b/rust/sdk/cocoindex/src/row_schema.rs index d839dca34..2f548580c 100644 --- a/rust/sdk/cocoindex/src/row_schema.rs +++ b/rust/sdk/cocoindex/src/row_schema.rs @@ -48,18 +48,57 @@ use crate::error::{Error, Result}; feature = "qdrant", feature = "turbopuffer" ))] -pub(crate) fn require_resolved_vector_dimension( +pub(crate) fn require_resolved_vector_dimension(connector: &str, field_name: &str) -> Result<()> { + Err(vector_dimension_error( + connector, + field_name, + format_args!( + "has an unresolved dimension; call with_vector_dim({field_name:?}, dimension) before \ + declaring the target" + ), + )) +} + +#[cfg(any( + feature = "postgres", + feature = "sqlite", + feature = "doris", + feature = "lancedb", + feature = "qdrant", + feature = "turbopuffer" +))] +pub(crate) fn vector_dimension_error( connector: &str, field_name: &str, - dimension: usize, -) -> Result<()> { - if dimension == 0 { - return Err(Error::engine(format!( - "{connector} vector field {field_name:?} has an unresolved dimension; call \ - with_vector_dim({field_name:?}, dimension) before declaring the target" - ))); - } - Ok(()) + detail: impl std::fmt::Display, +) -> Error { + Error::engine(format!("{connector} vector field {field_name:?} {detail}")) +} + +#[cfg(any( + feature = "postgres", + feature = "sqlite", + feature = "doris", + feature = "lancedb", + feature = "qdrant", + feature = "turbopuffer" +))] +pub(crate) fn unknown_vector_field_error(connector: &str, field_name: &str) -> Error { + Error::engine(format!( + "{connector} vector dimension override names unknown field {field_name:?}" + )) +} + +#[cfg(any( + feature = "postgres", + feature = "sqlite", + feature = "doris", + feature = "lancedb" +))] +pub(crate) fn not_vector_field_error(connector: &str, field_name: &str) -> Error { + Error::engine(format!( + "{connector} field {field_name:?} is not a vector field" + )) } #[cfg(any(feature = "postgres", feature = "sqlite", feature = "doris"))] diff --git a/rust/sdk/cocoindex/tests/pipeline.rs b/rust/sdk/cocoindex/tests/pipeline.rs index 6c477f572..3e0c3fde8 100644 --- a/rust/sdk/cocoindex/tests/pipeline.rs +++ b/rust/sdk/cocoindex/tests/pipeline.rs @@ -1767,13 +1767,6 @@ mod batched_test { .unwrap(); } - #[test] - #[should_panic(expected = "max_batch_size must be greater than zero")] - fn explicit_batched_rejects_zero_max_batch_size() { - let _ = - cocoindex::Batched::with_max_batch(|items: Vec| async move { Ok(items) }, 1, 0); - } - static MEMOIZED_ITEMS_PROCESSED: AtomicUsize = AtomicUsize::new(0); #[cocoindex::function(memo, batching, max_batch_size = 8)] diff --git a/rust/sdk/cocoindex/tests/schema_from_row.rs b/rust/sdk/cocoindex/tests/schema_from_row.rs index 852abd46f..4a24c323a 100644 --- a/rust/sdk/cocoindex/tests/schema_from_row.rs +++ b/rust/sdk/cocoindex/tests/schema_from_row.rs @@ -11,23 +11,13 @@ feature = "sqlite", feature = "turbopuffer" ))] -async fn capture_target_construction_error( - app_name: &'static str, - check: impl FnOnce(&cocoindex::Ctx) -> cocoindex::Result<()> + Send + 'static, -) -> String { - let tmp = tempfile::tempdir().unwrap(); - let app = cocoindex::App::open(app_name, tmp.path().join("db")) - .await - .unwrap(); - app.update(move |ctx| async move { - let error = match check(&ctx) { - Ok(()) => panic!("unresolved vector schema unexpectedly reached target construction"), - Err(error) => error, - }; - Ok(error.to_string()) - }) - .await - .unwrap() +fn assert_unresolved_dimension_error(result: cocoindex::Result<()>, connector: &str) { + let message = result + .expect_err("unresolved vector schema unexpectedly reached target construction") + .to_string(); + assert!(message.contains(connector), "{message}"); + assert!(message.contains("embedding"), "{message}"); + assert!(message.contains("with_vector_dim"), "{message}"); } #[cfg(any( @@ -38,10 +28,12 @@ async fn capture_target_construction_error( feature = "sqlite", feature = "turbopuffer" ))] -fn assert_unresolved_dimension_error(message: &str, connector: &str) { - assert!(message.contains(connector), "{message}"); - assert!(message.contains("embedding"), "{message}"); - assert!(message.contains("with_vector_dim"), "{message}"); +#[derive(cocoindex::SchemaFields)] +#[allow(dead_code)] +struct UnresolvedVectorRow { + id: i64, + #[coco(vector)] + embedding: Vec, } #[cfg(feature = "doris")] @@ -307,142 +299,82 @@ fn turbopuffer_from_row_derives_named_vectors_and_runtime_dimension() { assert_eq!(got, want); } -#[cfg(feature = "sqlite")] -#[tokio::test] -async fn sqlite_target_rejects_unresolved_vector_dimensions() { - use cocoindex::connectors::sqlite::{self, Database, TableSchema}; - use cocoindex::{ContextKey, SchemaFields}; - - #[allow(dead_code)] - #[derive(SchemaFields)] - struct Row { - id: i64, - #[coco(vector)] - embedding: Vec, - } - - let db = ContextKey::::new("sqlite_unresolved_vector_db"); - let message = capture_target_construction_error("sqlite_unresolved_vector", move |ctx| { - let schema = TableSchema::from_row::(["id"])?; - sqlite::table_target(ctx, &db, "docs", schema).map(|_| ()) - }) - .await; - assert_unresolved_dimension_error(&message, "SQLite"); -} - -#[cfg(feature = "postgres")] -#[tokio::test] -async fn postgres_target_rejects_unresolved_vector_dimensions() { - use cocoindex::connectors::postgres::{self, Database, TableSchema}; - use cocoindex::{ContextKey, SchemaFields}; - - #[allow(dead_code)] - #[derive(SchemaFields)] - struct Row { - id: i64, - #[coco(vector)] - embedding: Vec, - } - - let db = ContextKey::::new("postgres_unresolved_vector_db"); - let message = capture_target_construction_error("postgres_unresolved_vector", move |ctx| { - let schema = TableSchema::from_row::(["id"])?; - postgres::table_target(ctx, &db, "docs", schema, None).map(|_| ()) - }) - .await; - assert_unresolved_dimension_error(&message, "Postgres"); -} - -#[cfg(feature = "doris")] -#[tokio::test] -async fn doris_target_rejects_unresolved_vector_dimensions() { - use cocoindex::connectors::doris::{self, DorisConnection, TableSchema}; - use cocoindex::{ContextKey, SchemaFields}; - - #[allow(dead_code)] - #[derive(SchemaFields)] - struct Row { - id: String, - #[coco(vector)] - embedding: Vec, - } - - let db = ContextKey::::new("doris_unresolved_vector_db"); - let message = capture_target_construction_error("doris_unresolved_vector", move |ctx| { - let schema = TableSchema::from_row::(["id"])?; - doris::table_target(ctx, &db, "docs", schema).map(|_| ()) - }) - .await; - assert_unresolved_dimension_error(&message, "Doris"); -} - -#[cfg(feature = "lancedb")] -#[tokio::test] -async fn lancedb_target_rejects_unresolved_vector_dimensions() { - use cocoindex::connectors::lancedb::{self, LanceDatabase, TableSchema}; - use cocoindex::{ContextKey, SchemaFields}; - - #[allow(dead_code)] - #[derive(SchemaFields)] - struct Row { - id: i64, - #[coco(vector)] - embedding: Vec, - } - - let db = ContextKey::::new("lancedb_unresolved_vector_db"); - let message = capture_target_construction_error("lancedb_unresolved_vector", move |ctx| { - let schema = TableSchema::from_row::(["id"])?; - lancedb::table_target(ctx, &db, "docs", schema).map(|_| ()) - }) - .await; - assert_unresolved_dimension_error(&message, "LanceDB"); -} - -#[cfg(feature = "qdrant")] -#[tokio::test] -async fn qdrant_target_rejects_unresolved_vector_dimensions() { - use cocoindex::connectors::qdrant::{self, CollectionSchema, Distance, QdrantConnection}; - use cocoindex::{ContextKey, SchemaFields}; - - #[allow(dead_code)] - #[derive(SchemaFields)] - struct Row { - #[coco(vector)] - embedding: Vec, - } - - let db = ContextKey::::new("qdrant_unresolved_vector_db"); - let message = capture_target_construction_error("qdrant_unresolved_vector", move |ctx| { - let schema = CollectionSchema::from_row::(Distance::Cosine)?; - qdrant::collection_target(ctx, &db, "docs", schema).map(|_| ()) - }) - .await; - assert_unresolved_dimension_error(&message, "Qdrant"); -} - -#[cfg(feature = "turbopuffer")] +#[cfg(any( + feature = "doris", + feature = "lancedb", + feature = "postgres", + feature = "qdrant", + feature = "sqlite", + feature = "turbopuffer" +))] #[tokio::test] -async fn turbopuffer_target_rejects_unresolved_vector_dimensions() { - use cocoindex::connectors::turbopuffer::{ - self, DistanceMetric, NamespaceSchema, TurbopufferConnection, - }; - use cocoindex::{ContextKey, SchemaFields}; - - #[allow(dead_code)] - #[derive(SchemaFields)] - struct Row { - #[coco(vector)] - embedding: Vec, +async fn targets_reject_unresolved_vector_dimensions() { + macro_rules! check_target { + ($ctx:expr, $connector:literal, $schema:expr, $target:path $(, $extra:expr)*) => {{ + let key = cocoindex::ContextKey::new($connector); + assert_unresolved_dimension_error( + $target($ctx, &key, "docs", $schema $(, $extra)*).map(|_| ()), + $connector, + ); + }}; } - let db = ContextKey::::new("turbopuffer_unresolved_vector_db"); - let message = capture_target_construction_error("turbopuffer_unresolved_vector", move |ctx| { - let schema = NamespaceSchema::from_row::(DistanceMetric::CosineDistance)?; - turbopuffer::namespace_target(ctx, &db, "docs", schema).map(|_| ()) + let tmp = tempfile::tempdir().unwrap(); + let app = cocoindex::App::open("unresolved_vector_dimensions", tmp.path().join("db")) + .await + .unwrap(); + app.update(|ctx| async move { + #[cfg(feature = "sqlite")] + check_target!( + &ctx, + "SQLite", + cocoindex::connectors::sqlite::TableSchema::from_row::(["id"])?, + cocoindex::connectors::sqlite::table_target + ); + #[cfg(feature = "postgres")] + check_target!( + &ctx, + "Postgres", + cocoindex::connectors::postgres::TableSchema::from_row::(["id"])?, + cocoindex::connectors::postgres::table_target, + None + ); + #[cfg(feature = "doris")] + check_target!( + &ctx, + "Doris", + cocoindex::connectors::doris::TableSchema::from_row::(["id"])?, + cocoindex::connectors::doris::table_target + ); + #[cfg(feature = "lancedb")] + check_target!( + &ctx, + "LanceDB", + cocoindex::connectors::lancedb::TableSchema::from_row::(["id"])?, + cocoindex::connectors::lancedb::table_target + ); + #[cfg(feature = "qdrant")] + check_target!( + &ctx, + "Qdrant", + cocoindex::connectors::qdrant::CollectionSchema::from_row::( + cocoindex::connectors::qdrant::Distance::Cosine + )?, + cocoindex::connectors::qdrant::collection_target + ); + #[cfg(feature = "turbopuffer")] + check_target!( + &ctx, + "Turbopuffer", + cocoindex::connectors::turbopuffer::NamespaceSchema::from_row::( + cocoindex::connectors::turbopuffer::DistanceMetric::CosineDistance + )?, + cocoindex::connectors::turbopuffer::namespace_target + ); + Ok(()) }) - .await; - assert_unresolved_dimension_error(&message, "Turbopuffer"); + .await + .unwrap(); } /// End-to-end: a `from_row`-derived schema actually creates a SQLite table and diff --git a/rust/sdk/cocoindex_macros/src/lib.rs b/rust/sdk/cocoindex_macros/src/lib.rs index c29b18192..a4cca17f2 100644 --- a/rust/sdk/cocoindex_macros/src/lib.rs +++ b/rust/sdk/cocoindex_macros/src/lib.rs @@ -2,7 +2,7 @@ use proc_macro::TokenStream; use proc_macro2::{Span, TokenStream as TokenStream2}; -use quote::{format_ident, quote}; +use quote::{ToTokens, format_ident, quote}; use syn::{ Error as SynError, Expr, FnArg, Ident, ItemFn, LitInt, LitStr, Pat, PatType, Path, Stmt, Token, Type, TypeReference, parenthesized, @@ -38,7 +38,7 @@ fn parse_fn_params_opt(func: &ItemFn) -> syn::Result<(Option, Vec syn::Result<(Option, Vec syn::Result<(syn::Ident, Vec)> { Ok((ctx_ident, params)) } -fn is_str_ref_type(ty: &Type) -> bool { +fn is_ref_to(ty: &Type, name: &str) -> bool { if let Type::Reference(TypeReference { elem, .. }) = ty && let Type::Path(type_path) = elem.as_ref() { @@ -83,21 +83,7 @@ fn is_str_ref_type(ty: &Type) -> bool { .path .segments .last() - .is_some_and(|seg| seg.ident == "str"); - } - false -} - -/// Check if a type is a `&Ctx` reference. -fn is_ctx_type(ty: &Type) -> bool { - if let Type::Reference(TypeReference { elem, .. }) = ty - && let Type::Path(type_path) = elem.as_ref() - { - return type_path - .path - .segments - .last() - .is_some_and(|seg| seg.ident == "Ctx"); + .is_some_and(|segment| segment.ident == name); } false } @@ -138,18 +124,6 @@ fn compute_code_hash(block: &syn::Block, version: Option) -> u64 { hash } -use proc_macro2::TokenStream as Pm2TokenStream; - -trait ToTokenStream { - fn to_token_stream(&self) -> Pm2TokenStream; -} - -impl ToTokenStream for syn::Block { - fn to_token_stream(&self) -> Pm2TokenStream { - quote! { #self } - } -} - /// Parsed arguments for `#[function(...)]`. #[derive(Debug)] struct FunctionArgs { @@ -384,7 +358,6 @@ fn gen_key_write_for_param( fn gen_state_collect_for_param( states_ident: &Ident, - state_idx_ident: &Ident, prev_states_ident: &Ident, param: &ParamInfo, ) -> TokenStream2 { @@ -396,12 +369,11 @@ fn gen_state_collect_for_param( quote! { let __coco_prev_state = #prev_states_ident .as_ref() - .and_then(|__coco_states| __coco_states.get(#state_idx_ident)); + .and_then(|__coco_previous_states| __coco_previous_states.get(#states_ident.len())); if let Some(__coco_state) = ::cocoindex::memo::collect_memo_arg_state(#default_arg, __coco_prev_state).await? { #states_ident.push(__coco_state); - #state_idx_ident += 1; } } } @@ -422,6 +394,57 @@ fn gen_function_memo_key( }} } +fn gen_memo_wrapper_block( + fn_name: &Ident, + hash_const_name: &Ident, + ctx_ident: &Ident, + params: &[ParamInfo], + overrides: &[MemoKeyOverride], + propagate_children_fn_logic: bool, + body_closure: TokenStream2, +) -> syn::Result { + validate_memo_key_overrides( + overrides, + params.iter().map(|param| param.ident.to_string()), + )?; + let key_ident = format_ident!("__coco_key"); + let states_ident = format_ident!("__coco_states"); + let prev_states_ident = format_ident!("__coco_prev_states"); + let key_writes = params + .iter() + .filter_map(|param| gen_key_write_for_param(&key_ident, param, overrides)) + .collect::>(); + let state_collects = params + .iter() + .map(|param| gen_state_collect_for_param(&states_ident, &prev_states_ident, param)) + .collect::>(); + let clone_stmts = gen_clones(params); + let memo_key = gen_function_memo_key(fn_name, hash_const_name, &key_writes); + + Ok(quote! {{ + let __coco_key = #memo_key; + + ::cocoindex::memo::cached_by_fingerprint_with_state( + #ctx_ident, + __coco_key, + #propagate_children_fn_logic, + { + #(#clone_stmts)* + move |__coco_prev_states| async move { + let mut __coco_states = Vec::new(); + #(#state_collects)* + Ok(__coco_states) + } + }, + { + #(#clone_stmts)* + #body_closure + }, + ) + .await + }}) +} + #[derive(Debug)] enum BatchImplParam { Context, @@ -457,65 +480,36 @@ fn vec_inner_type(ty: &Type) -> Option { } fn item_shaped_return_type(output: &syn::ReturnType) -> syn::Result<(Type, syn::ReturnType)> { - let syn::ReturnType::Type(_, result_ty) = output else { - return Err(SynError::new_spanned( + let mut wrapper_output = output.clone(); + let invalid = || { + SynError::new_spanned( output, "batching function return type must be `Result>`", - )); - }; - let Type::Path(result_path) = result_ty.as_ref() else { - return Err(SynError::new_spanned( - result_ty, - "batching function return type must be `Result>`", - )); - }; - let Some(result_segment) = result_path.path.segments.last() else { - return Err(SynError::new_spanned( - result_ty, - "batching function return type must be `Result>`", - )); - }; - if result_segment.ident != "Result" { - return Err(SynError::new_spanned( - result_ty, - "batching function return type must be `Result>`", - )); - } - let syn::PathArguments::AngleBracketed(result_args) = &result_segment.arguments else { - return Err(SynError::new_spanned( - result_ty, - "batching function return type must be `Result>`", - )); - }; - let Some(syn::GenericArgument::Type(batch_output_ty)) = result_args.args.first() else { - return Err(SynError::new_spanned( - result_ty, - "batching function return type must be `Result>`", - )); - }; - let Some(item_output_ty) = vec_inner_type(batch_output_ty) else { - return Err(SynError::new_spanned( - batch_output_ty, - "batching function return type must be `Result>`", - )); + ) }; - - let mut wrapper_output = output.clone(); let syn::ReturnType::Type(_, wrapper_result_ty) = &mut wrapper_output else { - unreachable!(); + return Err(invalid()); }; let Type::Path(wrapper_result_path) = wrapper_result_ty.as_mut() else { - unreachable!(); + return Err(invalid()); }; - let wrapper_result_segment = wrapper_result_path.path.segments.last_mut().unwrap(); + let Some(wrapper_result_segment) = wrapper_result_path.path.segments.last_mut() else { + return Err(invalid()); + }; + if wrapper_result_segment.ident != "Result" { + return Err(invalid()); + } let syn::PathArguments::AngleBracketed(wrapper_args) = &mut wrapper_result_segment.arguments else { - unreachable!(); + return Err(invalid()); }; - let Some(syn::GenericArgument::Type(wrapper_item_ty)) = wrapper_args.args.first_mut() else { - unreachable!(); + let Some(syn::GenericArgument::Type(batch_output_ty)) = wrapper_args.args.first_mut() else { + return Err(invalid()); + }; + let Some(item_output_ty) = vec_inner_type(batch_output_ty) else { + return Err(invalid()); }; - *wrapper_item_ty = item_output_ty.clone(); + *batch_output_ty = item_output_ty.clone(); Ok((item_output_ty, wrapper_output)) } @@ -565,27 +559,19 @@ fn parse_batching_signature(func: &ItemFn) -> syn::Result { let (output_ty, wrapper_output) = item_shaped_return_type(&func.sig.output)?; let mut wrapper_sig = func.sig.clone(); - let mut found_batch_param = false; - for input in &mut wrapper_sig.inputs { - let FnArg::Typed(param) = input else { - return Err(SynError::new_spanned( - input, - "batching is supported only on free functions", - )); - }; - let Pat::Ident(param_ident) = param.pat.as_ref() else { - return Err(SynError::new_spanned( - ¶m.pat, - "batching function parameters must be identifiers", - )); - }; - if param_ident.ident == batch_param.ident { - *param.ty = item_ty.clone(); - found_batch_param = true; - break; - } - } - debug_assert!(found_batch_param); + let batch_input = wrapper_sig + .inputs + .iter_mut() + .find_map(|input| match input { + FnArg::Typed(param) + if matches!(param.pat.as_ref(), Pat::Ident(ident) if ident.ident == batch_param.ident) => + { + Some(param) + } + _ => None, + }) + .expect("batch parameter was validated above"); + *batch_input.ty = item_ty.clone(); wrapper_sig.output = wrapper_output; let batch_impl_params = input_idents @@ -659,12 +645,12 @@ fn expand_batching_function( let batch_static_name = format_ident!("__COCO_BATCHED_{}", fn_name.to_string().to_uppercase()); let batch_static_init = match args.max_batch_size { Some(max_batch_size) => quote! { - ::cocoindex::Batched::__scheduled_with_max_batch( + ::cocoindex::batched::__ScheduledBatched::__with_max_batch( #hash_const_name, #max_batch_size, ) }, - None => quote! { ::cocoindex::Batched::__scheduled(#hash_const_name) }, + None => quote! { ::cocoindex::batched::__ScheduledBatched::__new(#hash_const_name) }, }; let batch_impl_args = batching.batch_impl_params.iter().map(|param| match param { @@ -713,55 +699,19 @@ fn expand_batching_function( let wrapper_body = if args.memo { let mut key_params = vec![batching.item_param.clone()]; key_params.extend(extra_params.iter().cloned()); - validate_memo_key_overrides( + gen_memo_wrapper_block( + fn_name, + hash_const_name, + ctx_ident, + &key_params, &args.memo_key, - key_params.iter().map(|param| param.ident.to_string()), - )?; - let key_writes: Vec = key_params - .iter() - .filter_map(|param| { - gen_key_write_for_param(&format_ident!("__coco_key"), param, &args.memo_key) - }) - .collect(); - let state_collects: Vec = key_params - .iter() - .map(|param| { - gen_state_collect_for_param( - &format_ident!("__coco_states"), - &format_ident!("__coco_state_idx"), - &format_ident!("__coco_prev_states"), - param, - ) - }) - .collect(); - let state_clones = gen_clones(&key_params); - let body_clones = gen_clones(&key_params); - let memo_key = gen_function_memo_key(fn_name, hash_const_name, &key_writes); - quote! {{ - let __coco_key = #memo_key; - - ::cocoindex::memo::cached_by_fingerprint_with_state( - #ctx_ident, - __coco_key, - #propagate_children_fn_logic, - { - #(#state_clones)* - move |__coco_prev_states| async move { - let mut __coco_states = Vec::new(); - let mut __coco_state_idx = 0usize; - #(#state_collects)* - Ok(__coco_states) - } - }, - { - #(#body_clones)* - move |__coco_scoped_ctx| async move { - #schedule_call - } - }, - ) - .await - }} + propagate_children_fn_logic, + quote! { + move |__coco_scoped_ctx| async move { + #schedule_call + } + }, + )? } else { let initial_extra_clones = gen_clones(extra_params); if track_logic { @@ -793,7 +743,7 @@ fn expand_batching_function( pub const #hash_const_name: u64 = #hash_const_value; #logic_registration - static #batch_static_name: ::std::sync::LazyLock<::cocoindex::Batched<#item_ty, #output_ty>> = + static #batch_static_name: ::std::sync::LazyLock<::cocoindex::batched::__ScheduledBatched<#item_ty, #output_ty>> = ::std::sync::LazyLock::new(|| #batch_static_init); #batch_impl_sig #body @@ -945,31 +895,18 @@ pub fn function(attr: TokenStream, item: TokenStream) -> TokenStream { let sig = &func.sig; let attrs = &func.attrs; let body = &func.block; - let clone_stmts = gen_clones(¶ms); - if let Err(err) = - validate_memo_key_overrides(&args.memo_key, params.iter().map(|p| p.ident.to_string())) - { - return TokenStream::from(err.to_compile_error()); - } - - let key_writes: Vec = params - .iter() - .filter_map(|p| { - gen_key_write_for_param(&format_ident!("__coco_key"), p, &args.memo_key) - }) - .collect(); - let state_collects: Vec = params - .iter() - .map(|p| { - gen_state_collect_for_param( - &format_ident!("__coco_states"), - &format_ident!("__coco_state_idx"), - &format_ident!("__coco_prev_states"), - p, - ) - }) - .collect(); - let memo_key = gen_function_memo_key(fn_name, &hash_const_name, &key_writes); + let wrapper_body = match gen_memo_wrapper_block( + fn_name, + &hash_const_name, + &ctx_ident, + ¶ms, + &args.memo_key, + propagate_children_fn_logic, + quote! { move |#ctx_ident| async move #body }, + ) { + Ok(wrapper) => wrapper, + Err(err) => return TokenStream::from(err.to_compile_error()), + }; let expanded = quote! { #[doc(hidden)] @@ -977,22 +914,7 @@ pub fn function(attr: TokenStream, item: TokenStream) -> TokenStream { #logic_registration #(#attrs)* - #vis #sig { - let __coco_key = #memo_key; - - ::cocoindex::memo::cached_by_fingerprint_with_state(#ctx_ident, __coco_key, #propagate_children_fn_logic, { - #(#clone_stmts)* - move |__coco_prev_states| async move { - let mut __coco_states = Vec::new(); - let mut __coco_state_idx = 0usize; - #(#state_collects)* - Ok(__coco_states) - } - }, { - #(#clone_stmts)* - move |#ctx_ident| async move #body - }).await - } + #vis #sig #wrapper_body }; expanded.into() @@ -1045,8 +967,6 @@ pub fn function(attr: TokenStream, item: TokenStream) -> TokenStream { } } } else { - // No `&Ctx`: a pure logic-tracked helper, typically a batch impl for - // `coco::Batched`. quote! { #[doc(hidden)] pub const #hash_const_name: u64 = #hash_const_value; @@ -1766,13 +1686,4 @@ mod tests { "two batching bodies must produce different memo-key constants" ); } - - #[test] - fn batching_impl_names_preserve_leading_underscores() { - let plain = format_ident!("__coco_batch_impl_{}", format_ident!("embed")); - let underscored = format_ident!("__coco_batch_impl_{}", format_ident!("_embed")); - assert_eq!(plain.to_string(), "__coco_batch_impl_embed"); - assert_eq!(underscored.to_string(), "__coco_batch_impl__embed"); - assert_ne!(plain, underscored); - } } From 46ab58bd0cf120c4b84d1c17249a72a2dc22ad33 Mon Sep 17 00:00:00 2001 From: zhhanz Date: Wed, 15 Jul 2026 00:48:04 -0700 Subject: [PATCH 12/15] Document Doris byte encoding --- docs/src/content/docs/connectors/doris.mdx | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/src/content/docs/connectors/doris.mdx b/docs/src/content/docs/connectors/doris.mdx index db60ae080..966584816 100644 --- a/docs/src/content/docs/connectors/doris.mdx +++ b/docs/src/content/docs/connectors/doris.mdx @@ -199,6 +199,8 @@ Python types are automatically mapped to Doris types: | `list`, `dict`, nested structs | `JSON` | | `NDArray` (with vector schema) | `ARRAY` | +`bytes` values are intentionally base64-encoded before writing to `STRING` so arbitrary binary data survives Doris's JSON Stream Load path. + #### DorisType Use `DorisType` to specify a custom Doris type: From 9ad010f767bb0265c60be88f3f94fc47efe110cd Mon Sep 17 00:00:00 2001 From: zhhanz Date: Wed, 15 Jul 2026 13:20:16 -0700 Subject: [PATCH 13/15] Finish Rust SDK cleanup pass --- examples/rust/text_embedding/src/main.rs | 10 +-- examples/rust/text_embedding_qdrant/README.md | 2 +- .../rust/text_embedding_qdrant/src/main.rs | 14 ++- .../rust/text_embedding_turbopuffer/README.md | 2 +- .../text_embedding_turbopuffer/src/main.rs | 14 ++- rust/sdk/cocoindex/src/batched.rs | 19 +--- rust/sdk/cocoindex/src/connectors/doris.rs | 6 +- rust/sdk/cocoindex/src/connectors/postgres.rs | 8 +- rust/sdk/cocoindex/src/connectors/qdrant.rs | 6 +- rust/sdk/cocoindex/src/connectors/sqlite.rs | 6 +- .../cocoindex/src/connectors/turbopuffer.rs | 3 +- rust/sdk/cocoindex/src/memo.rs | 90 +------------------ rust/sdk/cocoindex/src/row_schema.rs | 15 ++++ rust/sdk/cocoindex_macros/src/lib.rs | 19 ++-- 14 files changed, 65 insertions(+), 149 deletions(-) diff --git a/examples/rust/text_embedding/src/main.rs b/examples/rust/text_embedding/src/main.rs index b1bce8868..708873ae1 100644 --- a/examples/rust/text_embedding/src/main.rs +++ b/examples/rust/text_embedding/src/main.rs @@ -7,7 +7,7 @@ //! //! Parallels the Python example: //! - source : `cocoindex::resources::fs::walk` (cf. `localfs.walk_dir`) -//! - per-file compute : `#[cocoindex::function(memo)]` (cf. `@coco.fn(memo=True)`) +//! - per-file compute : `#[cocoindex::function]` (cf. `@coco.fn`) //! - chunking : `cocoindex::ops::text::RecursiveSplitter` (cf. `RecursiveSplitter`) //! - embeddings : `cocoindex::ops::sentence_transformers` all-MiniLM-L6-v2 //! - target : `postgres::TableTarget` + pgvector index @@ -96,18 +96,14 @@ async fn process_file(ctx: &Ctx, file: FileEntry) -> Result> { Ok(rows) } -fn doc_embedding_schema(vector_dim: usize) -> Result { - postgres::TableSchema::from_row::(["id"])? - .with_vector_dim("embedding", vector_dim) -} - async fn app_main(ctx: Ctx, sourcedir: PathBuf) -> Result<()> { let vector_dim = ctx.get_key(&EMBEDDER)?.dimension(); let table = postgres::mount_table_target( &ctx, &DB, TABLE, - doc_embedding_schema(vector_dim)?, + postgres::TableSchema::from_row::(["id"])? + .with_vector_dim("embedding", vector_dim)?, Some(PG_SCHEMA), ) .await?; diff --git a/examples/rust/text_embedding_qdrant/README.md b/examples/rust/text_embedding_qdrant/README.md index badd48d9e..43af0990e 100644 --- a/examples/rust/text_embedding_qdrant/README.md +++ b/examples/rust/text_embedding_qdrant/README.md @@ -13,7 +13,7 @@ Same pipeline as [`text_embedding`](../text_embedding), but the vector store is | Chunking | `RecursiveSplitter` (markdown) | `cocoindex_ops_text` `RecursiveChunker` (markdown) | | Embeddings | `sentence-transformers/all-MiniLM-L6-v2` | `fastembed` `AllMiniLML6V2` (same model, 384-dim) | | Target | `qdrant.CollectionTarget` | `cocoindex::connectors::qdrant::CollectionTarget` | -| Search | `client.query_points(...)` | `cocoindex::connectors::qdrant::vector_search` (cosine score) | +| Search | `client.query_points(...)` | `cocoindex::connectors::qdrant::named_vector_search` (cosine score) | The `cocoindex::connectors::qdrant` connector is a declarative two-level **managed target** (collection → points) built on CocoIndex's public target-state facade: it diff --git a/examples/rust/text_embedding_qdrant/src/main.rs b/examples/rust/text_embedding_qdrant/src/main.rs index 9d4635eb6..03fe2f4cc 100644 --- a/examples/rust/text_embedding_qdrant/src/main.rs +++ b/examples/rust/text_embedding_qdrant/src/main.rs @@ -24,7 +24,6 @@ use cocoindex::prelude::*; use serde_json::json; const EMBED_MODEL: &str = "sentence-transformers/all-MiniLM-L6-v2"; -const EMBED_DIM: usize = 384; const COLLECTION: &str = "TextEmbedding"; const TOP_K: u64 = 5; const CHUNK_SIZE: usize = 2000; @@ -71,7 +70,15 @@ async fn process_file(ctx: &Ctx, file: FileEntry) -> Result> { } let texts: Vec = chunks.iter().map(|c| c.text(&text).to_string()).collect(); - let embeddings = ctx.get_key(&EMBEDDER)?.embed_batch(texts.clone()).await?; + let embedder = ctx.get_key(&EMBEDDER)?.clone(); + let embedding_ctx = ctx.clone(); + let embeddings = ctx + .map(texts.clone(), move |chunk_text| { + let embedder = embedder.clone(); + let ctx = embedding_ctx.clone(); + async move { embedder.embed(&ctx, chunk_text).await } + }) + .await?; let mut id_gen = IdGenerator::new(); let mut points = Vec::with_capacity(texts.len()); @@ -90,12 +97,13 @@ async fn process_file(ctx: &Ctx, file: FileEntry) -> Result> { } async fn app_main(ctx: Ctx, sourcedir: PathBuf) -> Result<()> { + let vector_dim = ctx.get_key(&EMBEDDER)?.dimension(); let target = qdrant::mount_collection_target( &ctx, &DB, COLLECTION, CollectionSchema::from_row::(Distance::Cosine)? - .with_vector_dim("vector", EMBED_DIM)?, + .with_vector_dim("vector", vector_dim)?, ) .await?; diff --git a/examples/rust/text_embedding_turbopuffer/README.md b/examples/rust/text_embedding_turbopuffer/README.md index 91d568680..1508401ee 100644 --- a/examples/rust/text_embedding_turbopuffer/README.md +++ b/examples/rust/text_embedding_turbopuffer/README.md @@ -13,7 +13,7 @@ store is **Turbopuffer** (a hosted vector database) via the native | Chunking | `RecursiveSplitter` (markdown) | `cocoindex_ops_text` `RecursiveChunker` (markdown) | | Embeddings | `sentence-transformers/all-MiniLM-L6-v2` | `fastembed` `AllMiniLML6V2` (same model, 384-dim) | | Target | `turbopuffer.NamespaceTarget` | `cocoindex::connectors::turbopuffer::NamespaceTarget` | -| Search | `ns.query(rank_by=("vector","ANN",...))` | `cocoindex::connectors::turbopuffer::vector_search` | +| Search | `ns.query(rank_by=("vector","ANN",...))` | `cocoindex::connectors::turbopuffer::vector_search_by_field` | The `cocoindex::connectors::turbopuffer` connector is a declarative two-level **managed target** (namespace → rows) built on CocoIndex's public target-state facade: it diff --git a/examples/rust/text_embedding_turbopuffer/src/main.rs b/examples/rust/text_embedding_turbopuffer/src/main.rs index 9f6a9a8f6..a2de6a97e 100644 --- a/examples/rust/text_embedding_turbopuffer/src/main.rs +++ b/examples/rust/text_embedding_turbopuffer/src/main.rs @@ -27,7 +27,6 @@ use cocoindex::prelude::*; use serde_json::json; const EMBED_MODEL: &str = "sentence-transformers/all-MiniLM-L6-v2"; -const EMBED_DIM: usize = 384; const TOP_K: usize = 5; const CHUNK_SIZE: usize = 2000; const CHUNK_OVERLAP: usize = 500; @@ -72,7 +71,15 @@ async fn process_file(ctx: &Ctx, file: FileEntry) -> Result> { } let texts: Vec = chunks.iter().map(|c| c.text(&text).to_string()).collect(); - let embeddings = ctx.get_key(&EMBEDDER)?.embed_batch(texts.clone()).await?; + let embedder = ctx.get_key(&EMBEDDER)?.clone(); + let embedding_ctx = ctx.clone(); + let embeddings = ctx + .map(texts.clone(), move |chunk_text| { + let embedder = embedder.clone(); + let ctx = embedding_ctx.clone(); + async move { embedder.embed(&ctx, chunk_text).await } + }) + .await?; let mut id_gen = IdGenerator::new(); let mut rows = Vec::with_capacity(texts.len()); @@ -91,12 +98,13 @@ async fn process_file(ctx: &Ctx, file: FileEntry) -> Result> { } async fn app_main(ctx: Ctx, sourcedir: PathBuf, namespace: String) -> Result<()> { + let vector_dim = ctx.get_key(&EMBEDDER)?.dimension(); let target = turbopuffer::mount_namespace_target( &ctx, &DB, &namespace, NamespaceSchema::from_row::(DistanceMetric::CosineDistance)? - .with_vector_dim("vector", EMBED_DIM)?, + .with_vector_dim("vector", vector_dim)?, ) .await?; diff --git a/rust/sdk/cocoindex/src/batched.rs b/rust/sdk/cocoindex/src/batched.rs index 0e5141c55..15ffcb30a 100644 --- a/rust/sdk/cocoindex/src/batched.rs +++ b/rust/sdk/cocoindex/src/batched.rs @@ -205,24 +205,9 @@ where { /// Construct the scheduler used by generated batching wrappers. #[doc(hidden)] - pub fn __new(code_hash: u64) -> Self { - Self::with_options(code_hash, BatchingOptions::default()) - } - - /// Construct the generated scheduler with a maximum batch size. - #[doc(hidden)] - pub fn __with_max_batch(code_hash: u64, max_batch_size: usize) -> Self { - Self::with_options( - code_hash, - BatchingOptions { - max_batch_size: Some(max_batch_size), - }, - ) - } - - fn with_options(code_hash: u64, options: BatchingOptions) -> Self { + pub fn __new(code_hash: u64, max_batch_size: Option) -> Self { Self { - options, + options: BatchingOptions { max_batch_size }, batchers: Mutex::new(HashMap::new()), code_hash, } diff --git a/rust/sdk/cocoindex/src/connectors/doris.rs b/rust/sdk/cocoindex/src/connectors/doris.rs index 406748d19..d34fce858 100644 --- a/rust/sdk/cocoindex/src/connectors/doris.rs +++ b/rust/sdk/cocoindex/src/connectors/doris.rs @@ -358,10 +358,8 @@ impl TableSchema { ) })?; if dim == 0 { - return Err(crate::row_schema::vector_dimension_error( - "Doris", - field_name, - "requires a dimension greater than zero", + return Err(crate::row_schema::zero_vector_dimension_error( + "Doris", field_name, )); } let def = self diff --git a/rust/sdk/cocoindex/src/connectors/postgres.rs b/rust/sdk/cocoindex/src/connectors/postgres.rs index 14e055e3f..d995bd96b 100644 --- a/rust/sdk/cocoindex/src/connectors/postgres.rs +++ b/rust/sdk/cocoindex/src/connectors/postgres.rs @@ -163,10 +163,8 @@ impl TableSchema { /// Resolve or override the dimension of a vector field derived from a row. pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { if dim == 0 { - return Err(crate::row_schema::vector_dimension_error( - "Postgres", - field_name, - "requires a dimension greater than zero", + return Err(crate::row_schema::zero_vector_dimension_error( + "Postgres", field_name, )); } let def = self @@ -1453,7 +1451,7 @@ fn sql_literal(value: &JsonValue, col: &ColumnDef) -> Result { } if lower == "json" || lower == "jsonb" { // Postgres rejects U+0000 in json/jsonb on parse, and serde serializes a - // NUL inside a nested string/key as the `` escape (which + // NUL inside a nested string/key as the `\x00` escape (which // `quote_string`'s byte-level strip can't catch). Remove NULs from every // string and object key before serializing. return Ok(format!( diff --git a/rust/sdk/cocoindex/src/connectors/qdrant.rs b/rust/sdk/cocoindex/src/connectors/qdrant.rs index 352505dbb..7a08228cf 100644 --- a/rust/sdk/cocoindex/src/connectors/qdrant.rs +++ b/rust/sdk/cocoindex/src/connectors/qdrant.rs @@ -371,10 +371,8 @@ impl CollectionSchema { /// the single unnamed vector uses the field name `"vector"`. pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { if dim == 0 { - return Err(crate::row_schema::vector_dimension_error( - "Qdrant", - field_name, - "has invalid dimension: qdrant vector size must be greater than zero", + return Err(crate::row_schema::zero_vector_dimension_error( + "Qdrant", field_name, )); } let mut fields = self.vector_fields(); diff --git a/rust/sdk/cocoindex/src/connectors/sqlite.rs b/rust/sdk/cocoindex/src/connectors/sqlite.rs index cb5c2956f..6b46e5b9b 100644 --- a/rust/sdk/cocoindex/src/connectors/sqlite.rs +++ b/rust/sdk/cocoindex/src/connectors/sqlite.rs @@ -173,10 +173,8 @@ impl TableSchema { /// Resolve or override the dimension of a vector field derived from a row. pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { if dim == 0 { - return Err(crate::row_schema::vector_dimension_error( - "SQLite", - field_name, - "requires a dimension greater than zero", + return Err(crate::row_schema::zero_vector_dimension_error( + "SQLite", field_name, )); } let def = self diff --git a/rust/sdk/cocoindex/src/connectors/turbopuffer.rs b/rust/sdk/cocoindex/src/connectors/turbopuffer.rs index 703b1a135..1d7f82644 100644 --- a/rust/sdk/cocoindex/src/connectors/turbopuffer.rs +++ b/rust/sdk/cocoindex/src/connectors/turbopuffer.rs @@ -304,10 +304,9 @@ impl NamespaceSchema { /// Resolve or override one named vector field's runtime dimension. pub fn with_vector_dim(mut self, field_name: &str, dim: usize) -> Result { if dim == 0 { - return Err(crate::row_schema::vector_dimension_error( + return Err(crate::row_schema::zero_vector_dimension_error( "Turbopuffer", field_name, - "requires a dimension greater than zero", )); } let mut fields = self.vector_fields(); diff --git a/rust/sdk/cocoindex/src/memo.rs b/rust/sdk/cocoindex/src/memo.rs index 832fa3be7..08841e865 100644 --- a/rust/sdk/cocoindex/src/memo.rs +++ b/rust/sdk/cocoindex/src/memo.rs @@ -16,7 +16,6 @@ use std::any::Any; use std::future::Future; use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; use cocoindex_core::engine::context::{FnCallContext, MemoStatesPayload}; @@ -203,24 +202,8 @@ where /// Internal helper for stable memo keys in generated macros. /// -/// Serializes key parts immediately into owned bytes to avoid leaking -/// temporary references from generated closures. -/// -/// This infallible helper never panics. Unsupported key types produce a unique, -/// stable fallback payload that includes the key type and an incrementing -/// sequence identifier to avoid collisions. -#[doc(hidden)] -pub fn key_bytes(value: &T) -> Vec { - match key_bytes_result(value) { - Ok(bytes) => bytes, - Err(err) => fallback_key_bytes::(err), - } -} - -/// Internal helper for stable memo keys in generated macros. -/// -/// This variant returns an error instead of panicking so callers can fail -/// fast with a typed error when serialization is not supported. +/// Returns an error instead of panicking so callers can fail fast with a typed +/// error when serialization is not supported. #[doc(hidden)] pub fn key_bytes_result(value: &T) -> Result> { rmp_serde::to_vec_named(value).map_err(Error::from) @@ -329,72 +312,3 @@ fn system_time_nanos(time: SystemTime) -> u128 { .unwrap_or_default() .as_nanos() } - -fn fallback_key_bytes(error: Error) -> Vec { - static KEY_BYTES_FALLBACK_SEQ: AtomicU64 = AtomicU64::new(0); - - let seq = KEY_BYTES_FALLBACK_SEQ.fetch_add(1, Ordering::AcqRel); - match rmp_serde::to_vec_named(&( - "__cocoindex_key_bytes_unsupported__", - std::any::type_name::(), - seq, - error.to_string(), - )) { - Ok(bytes) => bytes, - Err(_) => { - let mut bytes = b"__cocoindex_key_bytes_unsupported__".to_vec(); - bytes.extend_from_slice(std::any::type_name::().as_bytes()); - bytes.extend_from_slice(&seq.to_le_bytes()); - bytes.extend_from_slice(error.to_string().as_bytes()); - bytes - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use serde::ser::Error as SerError; - use serde::{Serialize, Serializer}; - - #[derive(Debug)] - struct FailingSerialize; - - impl Serialize for FailingSerialize { - fn serialize(&self, _serializer: S) -> std::result::Result { - Err(S::Error::custom("forced serialization failure")) - } - } - - #[test] - fn key_bytes_result_serializes_supported_types() { - let value = ("hello", 123); - let bytes = key_bytes_result(&value).unwrap(); - let decoded: (String, i32) = rmp_serde::from_slice(&bytes).unwrap(); - assert_eq!(decoded, (value.0.to_string(), value.1)); - } - - #[test] - fn key_bytes_uses_fallback_on_serialization_error() { - let bytes = key_bytes(&FailingSerialize); - let (marker, type_name, _, message) = - rmp_serde::from_slice::<(String, String, u64, String)>(&bytes).unwrap(); - - assert_eq!(marker, "__cocoindex_key_bytes_unsupported__"); - assert_eq!(type_name, std::any::type_name::()); - assert!(message.contains("forced serialization failure")); - } - - #[test] - fn key_bytes_fallback_is_unique_per_failure() { - let first = key_bytes(&FailingSerialize); - let second = key_bytes(&FailingSerialize); - - let (_, _, first_seq, _) = - rmp_serde::from_slice::<(String, String, u64, String)>(&first).unwrap(); - let (_, _, second_seq, _) = - rmp_serde::from_slice::<(String, String, u64, String)>(&second).unwrap(); - - assert_ne!(first_seq, second_seq); - } -} diff --git a/rust/sdk/cocoindex/src/row_schema.rs b/rust/sdk/cocoindex/src/row_schema.rs index 2f548580c..abe5a8986 100644 --- a/rust/sdk/cocoindex/src/row_schema.rs +++ b/rust/sdk/cocoindex/src/row_schema.rs @@ -75,6 +75,21 @@ pub(crate) fn vector_dimension_error( Error::engine(format!("{connector} vector field {field_name:?} {detail}")) } +#[cfg(any( + feature = "postgres", + feature = "sqlite", + feature = "doris", + feature = "qdrant", + feature = "turbopuffer" +))] +pub(crate) fn zero_vector_dimension_error(connector: &str, field_name: &str) -> Error { + vector_dimension_error( + connector, + field_name, + "requires a dimension greater than zero", + ) +} + #[cfg(any( feature = "postgres", feature = "sqlite", diff --git a/rust/sdk/cocoindex_macros/src/lib.rs b/rust/sdk/cocoindex_macros/src/lib.rs index a4cca17f2..266998b88 100644 --- a/rust/sdk/cocoindex_macros/src/lib.rs +++ b/rust/sdk/cocoindex_macros/src/lib.rs @@ -617,7 +617,6 @@ fn gen_owned_param_clones(params: &[ParamInfo]) -> Vec { .collect() } -#[allow(clippy::too_many_arguments)] fn expand_batching_function( func: &ItemFn, args: &FunctionArgs, @@ -643,14 +642,9 @@ fn expand_batching_function( let mut batch_impl_sig = func.sig.clone(); batch_impl_sig.ident = batch_impl_name.clone(); let batch_static_name = format_ident!("__COCO_BATCHED_{}", fn_name.to_string().to_uppercase()); - let batch_static_init = match args.max_batch_size { - Some(max_batch_size) => quote! { - ::cocoindex::batched::__ScheduledBatched::__with_max_batch( - #hash_const_name, - #max_batch_size, - ) - }, - None => quote! { ::cocoindex::batched::__ScheduledBatched::__new(#hash_const_name) }, + let max_batch_size = match args.max_batch_size { + Some(max_batch_size) => quote! { ::core::option::Option::Some(#max_batch_size) }, + None => quote! { ::core::option::Option::None }, }; let batch_impl_args = batching.batch_impl_params.iter().map(|param| match param { @@ -744,7 +738,12 @@ fn expand_batching_function( #logic_registration static #batch_static_name: ::std::sync::LazyLock<::cocoindex::batched::__ScheduledBatched<#item_ty, #output_ty>> = - ::std::sync::LazyLock::new(|| #batch_static_init); + ::std::sync::LazyLock::new(|| { + ::cocoindex::batched::__ScheduledBatched::__new( + #hash_const_name, + #max_batch_size, + ) + }); #batch_impl_sig #body From bb91e2fc665b16b334d88f21f87575ff0c2bb54a Mon Sep 17 00:00:00 2001 From: zhhanz Date: Wed, 15 Jul 2026 13:35:52 -0700 Subject: [PATCH 14/15] Fix private-key hook after Rust SDK regroup --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7347f3972..e6949266c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -16,7 +16,7 @@ repos: # oci_object_storage.rs only references the PEM key header in doc comments. # (This comment must avoid the literal key-header phrase, or the hook would # flag its own config file.) - exclude: "^rust/sdk/cocoindex/src/(gdrive|oci_object_storage)\\.rs$" + exclude: "^rust/sdk/cocoindex/src/connectors/(gdrive|oci_object_storage)\\.rs$" - id: end-of-file-fixer # Makes sure files end in a newline and only a newline. exclude: ".*(data.*|licenses.*|_static.*|\\.ya?ml|\\.jpe?g|\\.png|\\.svg|\\.webp)$" From 1e28f88aaddb188d94775c03d3e729ad4c463c76 Mon Sep 17 00:00:00 2001 From: zhhanz Date: Mon, 27 Jul 2026 11:53:59 -0700 Subject: [PATCH 15/15] Add recursive MemoInput and redesign Rust context keys Introduce MemoInput and #[derive(MemoInput)] so memo identity and optional async freshness validation are owned by types. Compose both through nested containers and supported ecosystem leaves, while preserving existing single-file resource memo keys and treating legacy or malformed nested state as cache misses. Redesign context_key! around identifier-default names, an optional key override, and opt-in detect_change matching Python defaults. Connection handles remain untracked unless explicitly enabled, the advanced context-site projection stays out of the initial API, and function-level memo_key overrides remain available. Migrate the Rust examples and documentation to the new contracts and refresh the standalone example lockfiles. --- benchmarks/file_summarization/rust/Cargo.lock | 3 + .../file_summarization/rust/src/main.rs | 4 +- dev/plans/rust-sdk-ergonomics.md | 218 ++++ .../docs/advanced_topics/memoization_keys.mdx | 121 ++ .../docs/getting_started/rust_quickstart.mdx | 18 +- examples/rust/amazon_s3_embedding/Cargo.lock | 17 + examples/rust/amazon_s3_embedding/src/main.rs | 16 +- examples/rust/audio_to_text/Cargo.lock | 23 + examples/rust/audio_to_text/src/main.rs | 7 +- examples/rust/code_embedding/Cargo.lock | 17 + examples/rust/code_embedding/README.md | 2 +- examples/rust/code_embedding/src/main.rs | 9 +- .../rust/code_embedding_lancedb/Cargo.lock | 18 + .../rust/code_embedding_lancedb/src/main.rs | 9 +- .../rust/conversation_to_knowledge/Cargo.lock | 2 + .../rust/conversation_to_knowledge/README.md | 2 +- .../conversation_to_knowledge/src/clients.rs | 42 +- .../conversation_to_knowledge/src/models.rs | 6 +- .../conversation_to_knowledge/src/pipeline.rs | 2 +- examples/rust/csv_to_iggy/Cargo.lock | 3 + examples/rust/csv_to_kafka/Cargo.lock | 3 + examples/rust/csv_to_kafka/src/main.rs | 2 +- examples/rust/files_to_sqlite/Cargo.lock | 24 + examples/rust/files_to_sqlite/src/main.rs | 5 +- examples/rust/files_transform/Cargo.lock | 3 + .../rust/gdrive_text_embedding/Cargo.lock | 17 + .../rust/gdrive_text_embedding/src/main.rs | 14 +- examples/rust/hn_trending_topics/Cargo.lock | 23 + examples/rust/hn_trending_topics/src/main.rs | 21 +- examples/rust/image_search/Cargo.lock | 18 + examples/rust/image_search/src/main.rs | 10 +- examples/rust/image_search_colpali/Cargo.lock | 24 + .../rust/image_search_colpali/src/main.rs | 20 +- examples/rust/kafka_consume/Cargo.lock | 3 + .../meeting_notes_graph_falkordb/Cargo.lock | 4 + .../rust/meeting_notes_graph_neo4j/Cargo.lock | 1 + .../multi_codebase_summarization/Cargo.lock | 3 + .../src/models.rs | 6 +- .../oci_object_storage_embedding/Cargo.lock | 17 + .../oci_object_storage_embedding/src/main.rs | 14 +- examples/rust/paper_metadata/Cargo.lock | 17 + examples/rust/paper_metadata/src/main.rs | 27 +- examples/rust/pdf_embedding/Cargo.lock | 17 + examples/rust/pdf_embedding/src/main.rs | 9 +- examples/rust/pdf_to_markdown/Cargo.lock | 3 + examples/rust/postgres_source/Cargo.lock | 17 + examples/rust/postgres_source/src/main.rs | 20 +- examples/rust/text_embedding/Cargo.lock | 17 + examples/rust/text_embedding/src/main.rs | 11 +- .../rust/text_embedding_lancedb/Cargo.lock | 18 + .../rust/text_embedding_lancedb/src/main.rs | 9 +- .../rust/text_embedding_qdrant/Cargo.lock | 18 + .../rust/text_embedding_qdrant/src/main.rs | 9 +- .../text_embedding_turbopuffer/Cargo.lock | 18 + .../text_embedding_turbopuffer/src/main.rs | 13 +- rust/sdk/SHOWCASE.md | 61 +- rust/sdk/cocoindex/Cargo.toml | 68 +- rust/sdk/cocoindex/src/app.rs | 3 +- .../sdk/cocoindex/src/connectors/amazon_s3.rs | 2 + rust/sdk/cocoindex/src/connectors/falkordb.rs | 2 +- rust/sdk/cocoindex/src/connectors/gdrive.rs | 2 + rust/sdk/cocoindex/src/connectors/neo4j.rs | 2 +- .../src/connectors/oci_object_storage.rs | 2 + rust/sdk/cocoindex/src/ctx.rs | 246 ++-- rust/sdk/cocoindex/src/lib.rs | 3 +- rust/sdk/cocoindex/src/memo.rs | 1051 ++++++++++++++++- rust/sdk/cocoindex/src/ops/image.rs | 7 + .../src/ops/sentence_transformers.rs | 9 +- rust/sdk/cocoindex/src/prelude.rs | 5 +- rust/sdk/cocoindex/src/profile.rs | 57 +- rust/sdk/cocoindex/src/resources/fs.rs | 2 + rust/sdk/cocoindex/tests/amazon_s3_source.rs | 3 +- rust/sdk/cocoindex/tests/doris_target.rs | 2 +- .../sdk/cocoindex/tests/graph_vector_index.rs | 4 +- rust/sdk/cocoindex/tests/lancedb_target.rs | 3 +- rust/sdk/cocoindex/tests/pipeline.rs | 362 +++++- rust/sdk/cocoindex/tests/postgres_source.rs | 5 +- rust/sdk/cocoindex/tests/postgres_target.rs | 3 +- rust/sdk/cocoindex/tests/qdrant_target.rs | 3 +- rust/sdk/cocoindex/tests/sqlite_target.rs | 3 +- rust/sdk/cocoindex/tests/surrealdb_target.rs | 2 +- .../sdk/cocoindex/tests/turbopuffer_target.rs | 3 +- rust/sdk/cocoindex/tests/valkey_target.rs | 2 +- rust/sdk/cocoindex_macros/src/lib.rs | 274 ++++- rust/utils/src/fingerprint.rs | 12 + 85 files changed, 2748 insertions(+), 449 deletions(-) create mode 100644 dev/plans/rust-sdk-ergonomics.md diff --git a/benchmarks/file_summarization/rust/Cargo.lock b/benchmarks/file_summarization/rust/Cargo.lock index 395791a16..146c0fd27 100644 --- a/benchmarks/file_summarization/rust/Cargo.lock +++ b/benchmarks/file_summarization/rust/Cargo.lock @@ -235,6 +235,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "cc" diff --git a/benchmarks/file_summarization/rust/src/main.rs b/benchmarks/file_summarization/rust/src/main.rs index 62fcdf69b..c448f0897 100644 --- a/benchmarks/file_summarization/rust/src/main.rs +++ b/benchmarks/file_summarization/rust/src/main.rs @@ -146,7 +146,7 @@ struct OutputSyncStats { output_hash: String, } -#[derive(Clone, Debug, Serialize, Deserialize)] +#[derive(Clone, Debug, Serialize, Deserialize, cocoindex::MemoInput)] struct SectionInput { stable_id: String, file_path: String, @@ -155,7 +155,7 @@ struct SectionInput { text: String, } -#[derive(Clone, Debug, Serialize, Deserialize)] +#[derive(Clone, Debug, Serialize, Deserialize, cocoindex::MemoInput)] struct SectionAnalysis { stable_id: String, file_path: String, diff --git a/dev/plans/rust-sdk-ergonomics.md b/dev/plans/rust-sdk-ergonomics.md new file mode 100644 index 000000000..7e8badae3 --- /dev/null +++ b/dev/plans/rust-sdk-ergonomics.md @@ -0,0 +1,218 @@ +# Plan: ergonomic Rust SDK — full workstream spec + +Tracking issue: [#2273](https://github.com/cocoindex-io/cocoindex/issues/2273) (plus maintainer +follow-ups in the thread). Goal: a Rust user should write pipelines that read line-for-line +like the Python SDK, without hand-wiring engine machinery (hash constants, `LazyLock` statics, +stringly-typed schemas, custom `main` plumbing). + +Each workstream below is independently implementable by an agent. Dependencies and suggested +PR order are at the end. The CLI workstream has its own spec: [rust-cli.md](rust-cli.md). + +Reference points used throughout: +- Macro crate: `rust/sdk/cocoindex_macros/src/lib.rs` (existing `#[function]`, `use_mount!`, + `mount_each!`, `#[derive(SchemaFields)]`). +- SDK crate: `rust/sdk/cocoindex/`. +- Python parity targets: `python/cocoindex/` and `examples/text_embedding/main.py`. + +--- + +## WS1 — `#[cocoindex::function(batching)]` and `(memo, batching)` + +**Problem.** Python: `@coco.fn(batching=True)` and done. Rust today: write a ctx-free batch fn, +then hand-wire `static EMBED: LazyLock>> = +LazyLock::new(|| Batched::new(embed_batch, __COCO_FN_HASH_EMBED_BATCH))` and call +`EMBED.call(&ctx, item)` (`rust/sdk/cocoindex/src/batched.rs`). `SHOWCASE.md` documents the +macro form but it was never implemented (the parser accepts only `memo`, `memo_key`, `version`, +`logic_tracking`). + +**Design.** Match Python's contract exactly: the function is *declared* batch-shaped, *called* +item-shaped. + +```rust +#[cocoindex::function(memo, batching, max_batch_size = 32)] +async fn embed(ctx: &Ctx, texts: Vec) -> Result>> { + ctx.get_key(&EMBEDDER)?.embed_batch(texts).await // body sees only the cache misses +} +// call site: +let v: Vec = embed(&ctx, text).await?; +``` + +The macro emits: the hash const + logic registration (as today), a hidden module-level +`Batched` static wired with the hash const, a private batch-impl fn holding the user's body, +and a public wrapper with the rewritten signature — first non-ctx param `Vec` → `T`, return +`Result>` → `Result`. Extra params beyond the items collection are cloned into the +closure and folded into each item's memo key (parity with `Batched` + Python). +`(batching)` without `memo` uses a pass-through (no per-item cache probe; body gets all items) +— mirror Python's `batching=True` without memo. `max_batch_size = N` maps to +`Batched::with_max_batch`. + +Compile errors (via `syn::Error`): `batching` with no non-ctx param; first non-ctx param not +`Vec<_>`; return type not `Result>`; `max_batch_size` without `batching`. + +**Acceptance.** The `SHOWCASE.md` batching examples compile as written (after WS7 fixes the +rest of that file). Test in `rust/sdk/cocoindex/tests/pipeline.rs`: per-item cache hits skip +the body, misses batch, body-edit (hash change) invalidates, extra-param change invalidates. + +--- + +## WS2 — `context_key!` declarative macro + +**Problem.** Python: `PG_DB = coco.ContextKey[asyncpg.Pool]("db")`. Rust: a 5-line +`static LazyLock>` ritual (see `examples/rust/text_embedding/src/main.rs:33-42`). + +**Design.** Plain `macro_rules!` in the SDK crate (no proc macro). The Rust identifier is the +context key's default stable name; `key = "..."` is available when the persisted identity must +differ from the identifier or survive a rename. Change detection is disabled by default and uses +the value type's `MemoInput` behavior when enabled: + +```rust +cocoindex::context_key!(static CONFIG: AppConfig); +cocoindex::context_key!(static EMBEDDER: SentenceTransformerEmbedder, detect_change); +cocoindex::context_key!(static DB: postgres::Database, key = "text_embedding_db"); +``` + +`MemoInput` is structural: the derive visits fields, standard containers recurse into their +children, and nested resource state validation is preserved. UUID is always available; standalone +`chrono`, `serde_json`, and `rust_decimal` features cover those ecosystem leaves, and connectors +enable the corresponding leaf features automatically. Unknown foreign leaves require a local +newtype. Ordered containers write framed children into one fingerprinter; maps and sets sort by +child fingerprints. Nested states remain one opaque state per top-level argument/context slot, so +the engine storage schema does not change. `Vec` follows the generic sequence path; large +buffers should use `bytes::Bytes` for bulk hashing or a precomputed fingerprint. + +When both optional arguments are present, `key =` appears before `detect_change`. + +**Acceptance.** Expansion produces the typed `LazyLock` pattern; migrate the statics in +`examples/rust/text_embedding` and connector tests (`tests/sqlite_target.rs:16` etc.); the macro +documentation explains identifier-default naming, stable-name overrides, default-off change +detection, `MemoInput`, the foreign-type newtype pattern, and the argument order. + +--- + +## WS3 — memo guidance fix (documentation, not a new macro) + +**Problem.** `rust/sdk/cocoindex/tests/pipeline.rs` (~2691) calls hand-rolled +`ctx.memo(&(__COCO_FN_HASH_ANALYZE, input), ...)` "the realistic pattern" for memoizing calls +that use non-serializable resources. This is wrong on both counts: the memo body closure +receives an owned `Ctx` (`cached_by_fingerprint_with_state`, `memo.rs:117-126`), so +`ctx.get_or_err::()`/`get_key` work inside `(memo)` bodies; and `memo_key(param = skip)` +params need only `Clone`, not `MemoInput` or `Serialize`, +`memo.rs:254`). Manual `ctx.memo` is also a foot-gun: its closure is not logic-tracked, so +forgetting the hash const silently serves stale results after a code edit. + +**Work.** +1. Rewrite the misleading test comment; add a test demonstrating `(memo)` + `ctx.get_or_err` + inside the body and `(memo, memo_key(client = skip))` with a non-`Serialize` `Clone` client. +2. Rustdoc on the `function` macro and `memo` module: the default is the attribute; manual + `ctx.memo` is for block-level memoization only, and must fold `__COCO_FN_HASH_*` into the + key by hand. +3. Do **not** implement a `memo!` block macro now; revisit only if block-level memoization + demand shows up. + +--- + +## WS4 — `SchemaFields` wiring into remaining connectors + +**Problem.** `#[derive(SchemaFields)]` exists (`cocoindex_macros/src/lib.rs:960`), but +`TableSchema::from_row` is implemented only for Doris (`doris.rs:337`) and SQLite +(`sqlite.rs:160`). Postgres, LanceDB, Qdrant, and Turbopuffer still require hand-written +column strings, and zero `examples/rust/` projects use the derive. + +**Work.** +1. Implement `TableSchema::from_row::(primary_key)` for `postgres`, + `lancedb`, `qdrant`, `turbopuffer`, mapping `LogicalType` (`row_schema.rs:34`) to each + connector's native types — mirror the per-connector `_LEAF_TYPE_MAPPINGS` in the Python + connectors (`python/cocoindex/connectors/*/`). Follow the sqlite impl as the template. +2. Add a runtime vector-dim override so the dim need not be hardcoded in the attribute + (Python parity: `Annotated[NDArray, EMBEDDER]` infers dim from the provided embedder): + `schema.with_vector_dim("embedding", embedder.dim())` on each connector's `TableSchema`. +3. Convert `examples/rust/*` that declare table schemas to the derive. + +**Acceptance.** `examples/rust/text_embedding` has no `ColumnDef::new("...")` strings left; +`tests/schema_from_row.rs` gains per-connector cases; unknown-`LogicalType` errors name the +field and connector. + +--- + +## WS5 — module regroup (mirror Python package layout) + +**Problem.** Connectors (`postgres`, `qdrant`, `lancedb`, `kafka`, …), resources (`file`, +`fs`, `id`), and machinery (`batched`, `memo`, `mount`, `statediff`) all sit at the crate root +(`rust/sdk/cocoindex/src/lib.rs:1-51`). Python groups them: `connectors/`, `resources/`, +`ops/`. + +**Work.** Move to `connectors::{postgres, sqlite, qdrant, lancedb, turbopuffer, doris, +surrealdb, kafka, iggy, valkey, neo4j, falkordb, amazon_s3, gdrive, oci_object_storage}`, +`resources::{file, fs, id}`; keep `ops::*` as-is; keep engine machinery (`app`, `ctx`, +`memo`, `mount`, `batched`, `logic`, `statediff`, `target_state`, `live_component`) at root. +Feature gates unchanged. Update `prelude`, all tests, all `examples/rust/*`, and +`dev/agent-skills/target-connector/SKILL.md`. No deprecated re-exports — the crate is +unreleased (workspace version is a placeholder), this is the moment for the breaking change. + +**Do this workstream first**: every other WS touches paths it moves. + +--- + +## WS6 — `ops::sentence_transformers` adopts engine batching + memoization + +**Problem.** Python's `SentenceTransformerEmbedder` uses +`@coco.fn.as_async(batching=True, runner=coco.GPU, max_batch_size=64)` and +`@coco.fn(memo=True, ...)` internally (`python/cocoindex/ops/sentence_transformers.py:122,167,205`) +— callers get batching and caching for free. The Rust op (`ops/sentence_transformers.rs`) is a +plain method. + +**Work (depends on WS1).** Restructure the op so single-text embedding goes through a +`#[cocoindex::function(memo, batching, max_batch_size = 64)]` entry; keep the raw +`embed_batch` public for direct use. GPU-pool runners are Python-only for now — out of scope. + +**Acceptance.** In an update run, repeated/unchanged texts hit cache; distinct texts group +into batches (assert via a counting fake embedder in tests). + +--- + +## WS7 — docs truth pass + +**Problem.** `rust/sdk/SHOWCASE.md` documents APIs that do not exist (`ctx.write_file`, +`ctx.batch`, sync `App::open`); the docs site has zero Rust SDK pages; the accurate reference +today is `tests/pipeline.rs`. + +**Work.** +1. Fix `SHOWCASE.md` to match the implemented API (after WS1 the batching sections become + true; fix `App::open` async-ness, replace `ctx.write_file` with `DirTarget`). +2. Add one docs-site page: Rust quickstart mirroring the Python `text_embedding` walkthrough + (docs live under `docs/src/content/docs/`; follow existing page conventions). +3. Document `memo_key(...)` (WS3), `context_key!` (WS2), and the mount spellings guidance: + `use_mount!`/`mount_each!` are the default; `ctx.scope`/`ctx.mount_each` are the + explicit-key/no-fingerprint variants. + +**Run last** — docs describe what has landed. + +--- + +## WS8 — unified Rust CLI + +Own spec: [rust-cli.md](rust-cli.md). Phase 1 = Rust CLI for Rust projects, `#[app]` / +`#[lifespan]` / `#[main]` macros, stdio protocol, `Environment::from_registered()`. Phase 2 = +replace the Python click CLI (binary shipped in wheels, `python -m cocoindex` protocol child). + +--- + +## Dependencies and suggested PR order + +``` +WS5 (module regroup) ──► everything else (paths) +WS1 (batching macro) ──► WS6 (ops adoption) +WS2, WS3, WS4 ──► independent of each other; after WS5 +WS8 (CLI) ──► independent (new crate + macros); after WS5 to avoid path churn +WS7 (docs) ──► last +``` + +1. PR 1: WS5 (mechanical, review-once) +2. PR 2: WS1 → PR 3: WS6 +3. PR 4: WS2 + WS3 (small, can share a PR) +4. PR 5: WS4 +5. PR 6+: WS8 phase 1 (SDK macros + dispatch, then CLI crate — see rust-cli.md milestones) +6. PR final: WS7 + +Every PR: `cargo test` + `cargo clippy` clean; when Python files are touched (WS7 docs only), +`uv run mypy && uv run pytest python/` per AGENTS.md. diff --git a/docs/src/content/docs/advanced_topics/memoization_keys.mdx b/docs/src/content/docs/advanced_topics/memoization_keys.mdx index 0b59b8223..4918e4aa9 100644 --- a/docs/src/content/docs/advanced_topics/memoization_keys.mdx +++ b/docs/src/content/docs/advanced_topics/memoization_keys.mdx @@ -234,6 +234,127 @@ class S3Object: ... # boto3 / aioboto3 HEAD call ``` +## Rust SDK + +The Rust SDK follows the same split between **identity** and **freshness**, but +expresses type-level behavior with the `MemoInput` trait. Arguments of a +`#[cocoindex::function(memo)]` function use `MemoInput` automatically. A +`context_key!` value participates only when the declaration includes +`detect_change`; plain context keys do not invalidate memoized work, matching +Python's default `detect_change=False`. + +### Derive structural keys + +Derive `MemoInput` for ordinary structs and enums. The derive visits fields +recursively and does not require the containing type to implement `Serialize`: + +```rust +use std::sync::Arc; + +#[derive(cocoindex::MemoInput)] +struct Batch { + model: String, + documents: Vec>, +} + +#[derive(cocoindex::MemoInput)] +struct Document { + id: uuid::Uuid, + tags: std::collections::BTreeSet, +} +``` + +`Option`, arrays, slices, `Vec`, tuples, maps, sets, `Box`, and `Arc` all +compose their children's `MemoInput` behavior. Ordered containers preserve +order; maps and sets use a deterministic order. If a nested resource defines +freshness validation, that validation is retained through every enclosing +container and derived type. + +The SDK always provides a UUID leaf implementation. Enable the standalone +`serde_json`, `chrono`, or `rust_decimal` feature for JSON values, Chrono +date/time values, or `rust_decimal::Decimal`; connector features enable the +corresponding leaf features automatically. Rust's orphan rules prevent an +application from implementing `MemoInput` directly for an unrelated +third-party type, so wrap any other foreign leaf in a local newtype: + +```rust +struct ForeignId(third_party::Id); + +#[cocoindex::async_trait] +impl cocoindex::MemoInput for ForeignId { + fn write_memo_key( + &self, + writer: &mut cocoindex::memo::MemoKeyWriter<'_>, + ) -> cocoindex::Result<()> { + writer.write(&self.0.to_string()) + } +} +``` + +Moving or renaming a derived type, or renaming one of its named fields, changes +its structural identity and causes a one-time cache miss. It does not rename +context keys: those use the identifier passed to `context_key!`, or an explicit +`key = "..."` override. + +### Add freshness validation to resource types + +A resource implementation can keep stable identity in `write_memo_key` and +perform async validation in `memo_state`. The previous opaque state is supplied +after an identity match: + +```rust +#[cocoindex::async_trait] +impl cocoindex::MemoInput for RemoteObject { + fn write_memo_key( + &self, + writer: &mut cocoindex::memo::MemoKeyWriter<'_>, + ) -> cocoindex::Result<()> { + writer.write(&(self.bucket(), self.key())) + } + + async fn memo_state( + &self, + previous: Option<&cocoindex::memo::MemoStateValue>, + ) -> cocoindex::Result> { + let etag = self.head_etag().await?; + let unchanged = previous + .and_then(|state| state.deserialize::().ok()) + .is_some_and(|old| old == etag); + Ok(Some(cocoindex::memo::MemoStateDecision::new( + &etag, + unchanged, + )?)) + } +} +``` + +Nested state checks run concurrently. CocoIndex stores their results as one +opaque state per top-level function argument or context dependency, so this +composition does not change the engine's memo-state storage schema. A malformed +or older nested state is treated as a cache miss rather than an application +error. + +### Override a function argument at the call site + +Use a projection when one function needs different identity semantics, or when +an unwrapped foreign value appears directly in its arguments: + +```rust +#[cocoindex::function(memo, memo_key(client = |c: &Client| c.endpoint()))] +async fn fetch(ctx: &cocoindex::Ctx, client: &Client) -> cocoindex::Result { + // ... +} +``` + +As in Python, a projection **replaces** the original value before +fingerprinting. In Rust it also suppresses that value type's `memo_state` +validation. Use a newtype instead when freshness validation must be preserved +or the behavior should be reusable. + +`Vec` intentionally follows the normal recursive `Vec` representation. +For large byte buffers, use `bytes::Bytes` for bulk hashing or pass a precomputed +`Fingerprint` when the content has already been hashed. + ## Preventing memoization Some types maintain internal state that makes memoization semantically incorrect. For example, a generator that tracks call counts would produce wrong results if memoized. diff --git a/docs/src/content/docs/getting_started/rust_quickstart.mdx b/docs/src/content/docs/getting_started/rust_quickstart.mdx index 46ace0062..33f3a6301 100644 --- a/docs/src/content/docs/getting_started/rust_quickstart.mdx +++ b/docs/src/content/docs/getting_started/rust_quickstart.mdx @@ -67,13 +67,10 @@ const EMBED_MODEL: &str = "sentence-transformers/all-MiniLM-L6-v2"; const PG_SCHEMA: &str = "coco_examples"; const TABLE: &str = "doc_embeddings"; +cocoindex::context_key!(static DB: postgres::Database); cocoindex::context_key!( - static DB: postgres::Database = "text_embedding_db", - state = postgres::Database::state_id -); -cocoindex::context_key!( - static EMBEDDER: SentenceTransformerEmbedder = "embedder", - state = SentenceTransformerEmbedder::model_name + static EMBEDDER: SentenceTransformerEmbedder, + detect_change ); #[derive(Clone, Serialize, Deserialize, SchemaFields)] @@ -88,10 +85,11 @@ struct DocEmbedding { } ``` -`context_key!` gives each provided resource a stable name and type. The -`state =` expressions say which durable property should invalidate dependent -memoized work when it changes. Runtime handles such as connection pools do not -need to be serializable. +`context_key!` gives each provided resource a stable name and type, using the +Rust identifier as the key name. Change detection is disabled for the database +connection. It is enabled for the embedder, whose type defines the model +identity that invalidates dependent memoized work. Runtime handles such as +connection pools do not need to be serializable. `SchemaFields` derives the database columns from `DocEmbedding`. The vector dimension is intentionally absent because it will come from the loaded model. diff --git a/examples/rust/amazon_s3_embedding/Cargo.lock b/examples/rust/amazon_s3_embedding/Cargo.lock index 66e5b5ee5..e53ad500f 100644 --- a/examples/rust/amazon_s3_embedding/Cargo.lock +++ b/examples/rust/amazon_s3_embedding/Cargo.lock @@ -929,6 +929,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "bytes-utils" @@ -1020,6 +1023,7 @@ dependencies = [ "globset", "linkme", "rmp-serde", + "rust_decimal", "rustc-hash", "serde", "serde_json", @@ -4111,6 +4115,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rust_decimal" +version = "1.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" +dependencies = [ + "arrayvec", + "num-traits", + "serde", + "wasm-bindgen", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -5895,6 +5911,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] diff --git a/examples/rust/amazon_s3_embedding/src/main.rs b/examples/rust/amazon_s3_embedding/src/main.rs index c528fb060..baaddc8c2 100644 --- a/examples/rust/amazon_s3_embedding/src/main.rs +++ b/examples/rust/amazon_s3_embedding/src/main.rs @@ -29,17 +29,11 @@ const PG_SCHEMA: &str = "coco_examples"; const TABLE: &str = "amazon_s3_doc_embeddings"; const TOP_K: i64 = 5; +cocoindex::context_key!(static DB: postgres::Database); +cocoindex::context_key!(static S3: S3Client); cocoindex::context_key!( - static DB: postgres::Database = "s3_embedding_db", - state = postgres::Database::state_id -); -cocoindex::context_key!( - static S3: S3Client = "s3_client", - state = S3Client::state_id -); -cocoindex::context_key!( - static EMBEDDER: SentenceTransformerEmbedder = "embedder", - state = SentenceTransformerEmbedder::model_name + static EMBEDDER: SentenceTransformerEmbedder, + detect_change ); #[derive(Clone, Serialize, Deserialize, SchemaFields)] @@ -51,7 +45,7 @@ struct DocEmbeddingRow { embedding: Vec, } -#[cocoindex::function] +#[cocoindex::function(memo)] async fn process_file(ctx: &Ctx, file: S3File) -> Result> { let text = ctx.get_key(&S3)?.read_text(&file).await?; let splitter = RecursiveSplitter::new()?; diff --git a/examples/rust/audio_to_text/Cargo.lock b/examples/rust/audio_to_text/Cargo.lock index 21dee8cf7..3bd5905ff 100644 --- a/examples/rust/audio_to_text/Cargo.lock +++ b/examples/rust/audio_to_text/Cargo.lock @@ -91,6 +91,12 @@ dependencies = [ "triomphe", ] +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + [[package]] name = "async-trait" version = "0.1.89" @@ -269,6 +275,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "cc" @@ -320,6 +329,7 @@ dependencies = [ "linkme", "reqwest", "rmp-serde", + "rust_decimal", "rustc-hash", "serde", "serde_json", @@ -2250,6 +2260,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rust_decimal" +version = "1.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" +dependencies = [ + "arrayvec", + "num-traits", + "serde", + "wasm-bindgen", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -3377,6 +3399,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] diff --git a/examples/rust/audio_to_text/src/main.rs b/examples/rust/audio_to_text/src/main.rs index 12f6c471e..b9dc78019 100644 --- a/examples/rust/audio_to_text/src/main.rs +++ b/examples/rust/audio_to_text/src/main.rs @@ -36,10 +36,7 @@ const AUDIO_PATTERNS: &[&str] = &[ "**/*.webm", ]; -cocoindex::context_key!( - static DB: postgres::Database = "audio_to_text_db", - state = postgres::Database::state_id -); +cocoindex::context_key!(static DB: postgres::Database); #[derive(Clone, Serialize, Deserialize, SchemaFields)] struct AudioTranscription { @@ -54,7 +51,7 @@ fn transcription_schema() -> Result { /// Transcribe one audio file with OpenAI Whisper via the SDK's /// `ApiTranscriber`. Memoized so the expensive API call only runs when the /// file's content changes (or is first seen). -#[cocoindex::function] +#[cocoindex::function(memo)] async fn transcribe(_ctx: &Ctx, file: &FileEntry) -> Result { let bytes = file.content()?; // Whisper sniffs the container from the upload filename's extension, so diff --git a/examples/rust/code_embedding/Cargo.lock b/examples/rust/code_embedding/Cargo.lock index 8ab404736..192155183 100644 --- a/examples/rust/code_embedding/Cargo.lock +++ b/examples/rust/code_embedding/Cargo.lock @@ -411,6 +411,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "castaway" @@ -475,6 +478,7 @@ dependencies = [ "globset", "linkme", "rmp-serde", + "rust_decimal", "rustc-hash", "serde", "serde_json", @@ -3275,6 +3279,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rust_decimal" +version = "1.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" +dependencies = [ + "arrayvec", + "num-traits", + "serde", + "wasm-bindgen", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -4930,6 +4946,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] diff --git a/examples/rust/code_embedding/README.md b/examples/rust/code_embedding/README.md index a83abe3a6..8bf9e26e6 100644 --- a/examples/rust/code_embedding/README.md +++ b/examples/rust/code_embedding/README.md @@ -13,7 +13,7 @@ Pipeline: **walk → detect language → tree-sitter chunk → embed → store i | Language detection | `detect_code_language` | `cocoindex_ops_text::prog_langs::detect_language` | | Chunking | `RecursiveSplitter` | `cocoindex_ops_text::split::RecursiveChunker` | | Embeddings | `SentenceTransformerEmbedder` (all-MiniLM-L6-v2) | `fastembed` `AllMiniLML6V2` — **the same model**, local ONNX | -| Embedder change-detection | `ContextKey(..., detect_change=True)` | `context_key!(..., state = SentenceTransformerEmbedder::model_name)` | +| Embedder change-detection | `ContextKey(..., detect_change=True)` | `context_key!(..., detect_change)` with type-owned `MemoInput` behavior | | Vector store | `postgres.TableTarget` + `declare_vector_index` | `cocoindex::connectors::postgres` `TableTarget` + `declare_vector_index` | | Stable row ids | `IdGenerator.next_id(chunk.text)` | `IdGenerator::next_id(ctx, chunk_text)` | | Query | pgvector `<=>` | pgvector `<=>` | diff --git a/examples/rust/code_embedding/src/main.rs b/examples/rust/code_embedding/src/main.rs index 2ed1c8de6..32417b5db 100644 --- a/examples/rust/code_embedding/src/main.rs +++ b/examples/rust/code_embedding/src/main.rs @@ -33,17 +33,14 @@ const TOP_K: i64 = 5; const INCLUDE_PATTERNS: &[&str] = &["**/*.py", "**/*.rs", "**/*.toml", "**/*.md", "**/*.mdx"]; // Shared Postgres target database. -cocoindex::context_key!( - static DB: postgres::Database = "code_embedding_db", - state = postgres::Database::state_id -); +cocoindex::context_key!(static DB: postgres::Database); // Shared embedder. The context key tracks the model name, so changing the model // invalidates memoized files — the parity for Python's // `ContextKey(..., detect_change=True)` + `Annotated[NDArray, EMBEDDER]`. cocoindex::context_key!( - static EMBEDDER: SentenceTransformerEmbedder = "embedder", - state = SentenceTransformerEmbedder::model_name + static EMBEDDER: SentenceTransformerEmbedder, + detect_change ); #[derive(Clone, Serialize, Deserialize, SchemaFields)] diff --git a/examples/rust/code_embedding_lancedb/Cargo.lock b/examples/rust/code_embedding_lancedb/Cargo.lock index af749c96f..43e5922fd 100644 --- a/examples/rust/code_embedding_lancedb/Cargo.lock +++ b/examples/rust/code_embedding_lancedb/Cargo.lock @@ -818,6 +818,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "castaway" @@ -890,6 +893,7 @@ dependencies = [ "arrow-schema", "async-trait", "bytes", + "chrono", "cocoindex_code_ast", "cocoindex_core", "cocoindex_macros", @@ -903,6 +907,7 @@ dependencies = [ "lancedb", "linkme", "rmp-serde", + "rust_decimal", "rustc-hash", "serde", "serde_json", @@ -5712,6 +5717,18 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "rust_decimal" +version = "1.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" +dependencies = [ + "arrayvec", + "num-traits", + "serde", + "wasm-bindgen", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -7518,6 +7535,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] diff --git a/examples/rust/code_embedding_lancedb/src/main.rs b/examples/rust/code_embedding_lancedb/src/main.rs index a7ccf3eb3..659203b24 100644 --- a/examples/rust/code_embedding_lancedb/src/main.rs +++ b/examples/rust/code_embedding_lancedb/src/main.rs @@ -24,13 +24,10 @@ const TOP_K: usize = 5; const INCLUDE_PATTERNS: &[&str] = &["**/*.py", "**/*.rs", "**/*.toml", "**/*.md", "**/*.mdx"]; +cocoindex::context_key!(static DB: LanceDatabase); cocoindex::context_key!( - static DB: LanceDatabase = "code_embedding_db", - state = LanceDatabase::state_id -); -cocoindex::context_key!( - static EMBEDDER: SentenceTransformerEmbedder = "embedder", - state = SentenceTransformerEmbedder::model_name + static EMBEDDER: SentenceTransformerEmbedder, + detect_change ); #[derive(Clone, Serialize, Deserialize, SchemaFields)] diff --git a/examples/rust/conversation_to_knowledge/Cargo.lock b/examples/rust/conversation_to_knowledge/Cargo.lock index f7b026477..ab843daa8 100644 --- a/examples/rust/conversation_to_knowledge/Cargo.lock +++ b/examples/rust/conversation_to_knowledge/Cargo.lock @@ -770,6 +770,7 @@ version = "0.1.0" dependencies = [ "async-trait", "bytes", + "chrono", "cocoindex_core", "cocoindex_macros", "cocoindex_utils", @@ -778,6 +779,7 @@ dependencies = [ "globset", "linkme", "rmp-serde", + "rust_decimal", "rustc-hash", "serde", "serde_json", diff --git a/examples/rust/conversation_to_knowledge/README.md b/examples/rust/conversation_to_knowledge/README.md index 69f695cb5..4caf8754d 100644 --- a/examples/rust/conversation_to_knowledge/README.md +++ b/examples/rust/conversation_to_knowledge/README.md @@ -18,7 +18,7 @@ sessions, statements, persons, techs, orgs, and the relationships between them. | Stable ids | `IdGenerator` | `cocoindex::IdGenerator` | | Entity resolution | `ops.entity_resolution` (faiss + LLM) | `cocoindex::entity_resolution` + `fastembed` Snowflake embeddings + LLM pair resolver | | Graph store | `surrealdb` connector (`TableTarget`/`RelationTarget`) | `cocoindex::connectors::surrealdb` targets over the native `surrealdb` crate | -| Embedder change-detection | `ContextKey(..., detect_change=True)` | `context_key!(..., state = Embedder::model_name)` | +| Embedder change-detection | `ContextKey(..., detect_change=True)` | `context_key!(..., detect_change)` with type-owned `MemoInput` behavior | ### Design notes / where it differs diff --git a/examples/rust/conversation_to_knowledge/src/clients.rs b/examples/rust/conversation_to_knowledge/src/clients.rs index 8e525c1be..4bd058501 100644 --- a/examples/rust/conversation_to_knowledge/src/clients.rs +++ b/examples/rust/conversation_to_knowledge/src/clients.rs @@ -17,32 +17,28 @@ use serde::de::DeserializeOwned; // Context keys // --------------------------------------------------------------------------- -// LLM used for metadata/statement extraction. State-tracked on the model name, +// LLM used for metadata/statement extraction. Its MemoInput tracks the model, // so changing the model invalidates memoized extraction (parity with Python's // `LLM_MODEL = ContextKey(..., detect_change=True)`). cocoindex::context_key!( - pub static LLM: LlmClient = "llm_model", - state = LlmClient::model_name + pub static LLM: LlmClient, + detect_change ); // LLM used to confirm entity-resolution pairs. cocoindex::context_key!( - pub static RESOLVER_LLM: LlmClient = "resolution_llm_model", - state = LlmClient::model_name + pub static RESOLVER_LLM: LlmClient, + detect_change ); // Local embedder for entity-resolution similarity. cocoindex::context_key!( - pub static EMBEDDER: Embedder = "embedder", - state = Embedder::model_name + pub static EMBEDDER: Embedder, + detect_change ); -// SurrealDB connection. State-tracked on the target endpoint so changing the -// external graph database invalidates local target-state reconciliation. -cocoindex::context_key!( - pub static GRAPH: Graph = "surreal_db", - state = Graph::state_id -); +// SurrealDB connection. Connection resources are not memo dependencies. +cocoindex::context_key!(pub static GRAPH: Graph); // --------------------------------------------------------------------------- // LLM client (OpenAI-compatible, JSON mode) @@ -57,10 +53,6 @@ pub struct LlmClient { } impl LlmClient { - fn model_name(&self) -> &str { - &self.model - } - pub fn new(model: String) -> Result { let api_key = std::env::var("OPENAI_API_KEY") .or_else(|_| std::env::var("LLM_API_KEY")) @@ -113,6 +105,12 @@ impl LlmClient { } } +impl MemoInput for LlmClient { + fn write_memo_key(&self, writer: &mut cocoindex::memo::MemoKeyWriter<'_>) -> Result<()> { + writer.write(&self.model) + } +} + // --------------------------------------------------------------------------- // Embedder (local fastembed model) // --------------------------------------------------------------------------- @@ -124,10 +122,6 @@ pub struct Embedder { } impl Embedder { - fn model_name(&self) -> &str { - &self.model_name - } - pub fn load(model_name: &str) -> Result { let model = match model_name { "Snowflake/snowflake-arctic-embed-xs" => load_snowflake_arctic_embed_xs()?, @@ -160,6 +154,12 @@ impl Embedder { } } +impl MemoInput for Embedder { + fn write_memo_key(&self, writer: &mut cocoindex::memo::MemoKeyWriter<'_>) -> Result<()> { + writer.write(&self.model_name) + } +} + fn load_snowflake_arctic_embed_xs() -> Result { let cache_dir = std::env::var("HF_HOME") .map(PathBuf::from) diff --git a/examples/rust/conversation_to_knowledge/src/models.rs b/examples/rust/conversation_to_knowledge/src/models.rs index f9d5df041..fcf6f6b13 100644 --- a/examples/rust/conversation_to_knowledge/src/models.rs +++ b/examples/rust/conversation_to_knowledge/src/models.rs @@ -36,14 +36,14 @@ pub fn entity_guidance(kind: &str) -> (&'static str, &'static str) { // Transcript / source // --------------------------------------------------------------------------- -#[derive(Clone, Debug, Serialize, Deserialize)] +#[derive(Clone, Debug, Serialize, Deserialize, cocoindex::MemoInput)] pub struct Utterance { /// Diarization label (e.g. "A", "B") or a resolved speaker name. pub speaker: String, pub text: String, } -#[derive(Clone, Debug, Serialize, Deserialize)] +#[derive(Clone, Debug, Serialize, Deserialize, cocoindex::MemoInput)] pub struct SessionTranscript { pub utterances: Vec, pub yt_channel: String, @@ -54,7 +54,7 @@ pub struct SessionTranscript { /// A unit of input: either a YouTube video to fetch+transcribe, or a /// pre-transcribed local session (cheap, audio-free — for testing/demo). -#[derive(Clone, Debug, Serialize, Deserialize)] +#[derive(Clone, Debug, Serialize, Deserialize, cocoindex::MemoInput)] pub enum SessionSource { YouTube { youtube_id: String, diff --git a/examples/rust/conversation_to_knowledge/src/pipeline.rs b/examples/rust/conversation_to_knowledge/src/pipeline.rs index 59f4b3caa..6898436e1 100644 --- a/examples/rust/conversation_to_knowledge/src/pipeline.rs +++ b/examples/rust/conversation_to_knowledge/src/pipeline.rs @@ -3,9 +3,9 @@ use std::collections::{HashMap, HashSet}; use async_trait::async_trait; +use cocoindex::connectors::surrealdb; use cocoindex::entity_resolution; use cocoindex::prelude::*; -use cocoindex::connectors::surrealdb; use serde::Deserialize; use crate::clients::{EMBEDDER, Embedder, GRAPH, LLM, LlmClient, RESOLVER_LLM}; diff --git a/examples/rust/csv_to_iggy/Cargo.lock b/examples/rust/csv_to_iggy/Cargo.lock index e3061cf3d..75267949e 100644 --- a/examples/rust/csv_to_iggy/Cargo.lock +++ b/examples/rust/csv_to_iggy/Cargo.lock @@ -589,6 +589,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "cc" diff --git a/examples/rust/csv_to_kafka/Cargo.lock b/examples/rust/csv_to_kafka/Cargo.lock index 585fb5717..c9db1af19 100644 --- a/examples/rust/csv_to_kafka/Cargo.lock +++ b/examples/rust/csv_to_kafka/Cargo.lock @@ -244,6 +244,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "cc" diff --git a/examples/rust/csv_to_kafka/src/main.rs b/examples/rust/csv_to_kafka/src/main.rs index b01d0d97d..c4396d0b4 100644 --- a/examples/rust/csv_to_kafka/src/main.rs +++ b/examples/rust/csv_to_kafka/src/main.rs @@ -38,7 +38,7 @@ fn topic_name() -> String { /// Parse one CSV file into `(message_key, json_value)` pairs. Memoized: a file /// whose content is unchanged since the last run is not re-parsed. -#[cocoindex::function] +#[cocoindex::function(memo)] async fn process_csv(_ctx: &Ctx, file: FileEntry) -> Result> { let text = file.content_str()?; let mut reader = csv::Reader::from_reader(text.as_bytes()); diff --git a/examples/rust/files_to_sqlite/Cargo.lock b/examples/rust/files_to_sqlite/Cargo.lock index 597afddd4..4da06225b 100644 --- a/examples/rust/files_to_sqlite/Cargo.lock +++ b/examples/rust/files_to_sqlite/Cargo.lock @@ -91,6 +91,12 @@ dependencies = [ "triomphe", ] +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + [[package]] name = "async-trait" version = "0.1.89" @@ -259,6 +265,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "cc" @@ -300,6 +309,7 @@ version = "0.1.0" dependencies = [ "async-trait", "bytes", + "chrono", "cocoindex_core", "cocoindex_macros", "cocoindex_utils", @@ -308,6 +318,7 @@ dependencies = [ "globset", "linkme", "rmp-serde", + "rust_decimal", "rustc-hash", "serde", "serde_json", @@ -2080,6 +2091,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rust_decimal" +version = "1.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" +dependencies = [ + "arrayvec", + "num-traits", + "serde", + "wasm-bindgen", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -3112,6 +3135,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] diff --git a/examples/rust/files_to_sqlite/src/main.rs b/examples/rust/files_to_sqlite/src/main.rs index f51d1b5e5..85a9e2b75 100644 --- a/examples/rust/files_to_sqlite/src/main.rs +++ b/examples/rust/files_to_sqlite/src/main.rs @@ -19,10 +19,7 @@ use cocoindex::prelude::*; use serde::{Deserialize, Serialize}; use sqlx::Row as _; -cocoindex::context_key!( - static DB: sqlite::Database = "files_sqlite_db", - state = sqlite::Database::state_id -); +cocoindex::context_key!(static DB: sqlite::Database); const TABLE: &str = "files"; diff --git a/examples/rust/files_transform/Cargo.lock b/examples/rust/files_transform/Cargo.lock index 2412510b4..9809b5526 100644 --- a/examples/rust/files_transform/Cargo.lock +++ b/examples/rust/files_transform/Cargo.lock @@ -235,6 +235,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "cc" diff --git a/examples/rust/gdrive_text_embedding/Cargo.lock b/examples/rust/gdrive_text_embedding/Cargo.lock index 4dfd1d5bf..9006eec7c 100644 --- a/examples/rust/gdrive_text_embedding/Cargo.lock +++ b/examples/rust/gdrive_text_embedding/Cargo.lock @@ -411,6 +411,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "castaway" @@ -478,6 +481,7 @@ dependencies = [ "reqwest", "rmp-serde", "rsa", + "rust_decimal", "rustc-hash", "serde", "serde_json", @@ -3280,6 +3284,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rust_decimal" +version = "1.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" +dependencies = [ + "arrayvec", + "num-traits", + "serde", + "wasm-bindgen", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -4935,6 +4951,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] diff --git a/examples/rust/gdrive_text_embedding/src/main.rs b/examples/rust/gdrive_text_embedding/src/main.rs index 1f9b3f3ef..7085fa199 100644 --- a/examples/rust/gdrive_text_embedding/src/main.rs +++ b/examples/rust/gdrive_text_embedding/src/main.rs @@ -21,17 +21,11 @@ const PG_SCHEMA: &str = "coco_examples_v1"; const TABLE: &str = "doc_embeddings"; const TOP_K: i64 = 5; +cocoindex::context_key!(static DB: postgres::Database); +cocoindex::context_key!(static GDRIVE: GoogleDriveClient); cocoindex::context_key!( - static DB: postgres::Database = "gdrive_text_embedding_db", - state = postgres::Database::state_id -); -cocoindex::context_key!( - static GDRIVE: GoogleDriveClient = "gdrive_client", - state = GoogleDriveClient::state_id -); -cocoindex::context_key!( - static EMBEDDER: SentenceTransformerEmbedder = "embedder", - state = SentenceTransformerEmbedder::model_name + static EMBEDDER: SentenceTransformerEmbedder, + detect_change ); #[derive(Clone, Serialize, Deserialize, SchemaFields)] diff --git a/examples/rust/hn_trending_topics/Cargo.lock b/examples/rust/hn_trending_topics/Cargo.lock index b4585eb61..7536fdc59 100644 --- a/examples/rust/hn_trending_topics/Cargo.lock +++ b/examples/rust/hn_trending_topics/Cargo.lock @@ -91,6 +91,12 @@ dependencies = [ "triomphe", ] +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + [[package]] name = "async-trait" version = "0.1.89" @@ -259,6 +265,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "cc" @@ -309,6 +318,7 @@ dependencies = [ "globset", "linkme", "rmp-serde", + "rust_decimal", "rustc-hash", "serde", "serde_json", @@ -2240,6 +2250,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rust_decimal" +version = "1.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" +dependencies = [ + "arrayvec", + "num-traits", + "serde", + "wasm-bindgen", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -3361,6 +3383,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] diff --git a/examples/rust/hn_trending_topics/src/main.rs b/examples/rust/hn_trending_topics/src/main.rs index df0b631d1..0034eb0e3 100644 --- a/examples/rust/hn_trending_topics/src/main.rs +++ b/examples/rust/hn_trending_topics/src/main.rs @@ -32,14 +32,11 @@ const MAX_TEXT: usize = 4000; static HTTP: LazyLock = LazyLock::new(reqwest::Client::new); // Shared Postgres target database. +cocoindex::context_key!(static PG: postgres::Database); +// LLM client; its MemoInput implementation tracks the model name. cocoindex::context_key!( - static PG: postgres::Database = "hn_db", - state = postgres::Database::state_id -); -// LLM client; state-tracked on the model so changing it invalidates memos. -cocoindex::context_key!( - static LLM: LlmClient = "llm_model", - state = LlmClient::model_name + static LLM: LlmClient, + detect_change ); // --------------------------------------------------------------------------- @@ -78,10 +75,6 @@ const TOPICS_PROMPT: &str = "Extract topics from the user's text. Return a JSON Example: \"John Kennedy\", \"JFK\"."; impl LlmClient { - fn model_name(&self) -> &str { - &self.model - } - fn new(model: String) -> Result { let api_key = std::env::var("OPENAI_API_KEY") .or_else(|_| std::env::var("LLM_API_KEY")) @@ -137,6 +130,12 @@ impl LlmClient { } } +impl MemoInput for LlmClient { + fn write_memo_key(&self, writer: &mut cocoindex::memo::MemoKeyWriter<'_>) -> Result<()> { + writer.write(&self.model) + } +} + // --------------------------------------------------------------------------- // HackerNews (Algolia) API // --------------------------------------------------------------------------- diff --git a/examples/rust/image_search/Cargo.lock b/examples/rust/image_search/Cargo.lock index db531f272..2b464ef21 100644 --- a/examples/rust/image_search/Cargo.lock +++ b/examples/rust/image_search/Cargo.lock @@ -465,6 +465,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "castaway" @@ -517,6 +520,7 @@ version = "0.1.0" dependencies = [ "async-trait", "bytes", + "chrono", "cocoindex_core", "cocoindex_macros", "cocoindex_utils", @@ -527,6 +531,7 @@ dependencies = [ "linkme", "qdrant-client", "rmp-serde", + "rust_decimal", "rustc-hash", "serde", "serde_json", @@ -3104,6 +3109,18 @@ dependencies = [ "smallvec", ] +[[package]] +name = "rust_decimal" +version = "1.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" +dependencies = [ + "arrayvec", + "num-traits", + "serde", + "wasm-bindgen", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -4216,6 +4233,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] diff --git a/examples/rust/image_search/src/main.rs b/examples/rust/image_search/src/main.rs index 9ead5d07d..69c40fbfe 100644 --- a/examples/rust/image_search/src/main.rs +++ b/examples/rust/image_search/src/main.rs @@ -43,13 +43,10 @@ const IMAGE_GLOBS: &[&str] = &[ "**/*.bmp", ]; +cocoindex::context_key!(static DB: QdrantConnection); cocoindex::context_key!( - static DB: QdrantConnection = "image_search_db", - state = QdrantConnection::state_id -); -cocoindex::context_key!( - static EMBEDDER: ImageEmbedder = "image_embedder", - state = ImageEmbedder::model_name + static EMBEDDER: ImageEmbedder, + detect_change ); /// A computed point: stable id + image vector + source filename. @@ -75,7 +72,6 @@ async fn process_image(ctx: &Ctx, file: FileEntry) -> Result { } async fn app_main(ctx: Ctx, sourcedir: PathBuf) -> Result<()> { - let conn = ctx.get_key(&DB)?; let target = qdrant::mount_collection_target( &ctx, &DB, diff --git a/examples/rust/image_search_colpali/Cargo.lock b/examples/rust/image_search_colpali/Cargo.lock index c81b3d3d0..463d32b91 100644 --- a/examples/rust/image_search_colpali/Cargo.lock +++ b/examples/rust/image_search_colpali/Cargo.lock @@ -97,6 +97,12 @@ dependencies = [ "triomphe", ] +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + [[package]] name = "async-stream" version = "0.3.6" @@ -319,6 +325,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "cc" @@ -360,6 +369,7 @@ version = "0.1.0" dependencies = [ "async-trait", "bytes", + "chrono", "cocoindex_core", "cocoindex_macros", "cocoindex_utils", @@ -369,6 +379,7 @@ dependencies = [ "linkme", "qdrant-client", "rmp-serde", + "rust_decimal", "rustc-hash", "serde", "serde_json", @@ -2204,6 +2215,18 @@ dependencies = [ "smallvec", ] +[[package]] +name = "rust_decimal" +version = "1.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" +dependencies = [ + "arrayvec", + "num-traits", + "serde", + "wasm-bindgen", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -3151,6 +3174,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] diff --git a/examples/rust/image_search_colpali/src/main.rs b/examples/rust/image_search_colpali/src/main.rs index 9a6335c5a..d894e5adb 100644 --- a/examples/rust/image_search_colpali/src/main.rs +++ b/examples/rust/image_search_colpali/src/main.rs @@ -44,13 +44,10 @@ const IMAGE_GLOBS: &[&str] = &[ "**/*.bmp", ]; +cocoindex::context_key!(static DB: QdrantConnection); cocoindex::context_key!( - static DB: QdrantConnection = "image_search_colpali_db", - state = QdrantConnection::state_id -); -cocoindex::context_key!( - static COLPALI: ColpaliClient = "colpali", - state = ColpaliClient::state_id + static COLPALI: ColpaliClient, + detect_change ); /// HTTP client for an external ColPali inference service (see module docs). @@ -67,10 +64,6 @@ struct EmbeddingResponse { } impl ColpaliClient { - fn state_id(&self) -> &str { - &self.url - } - fn new(url: String) -> Self { Self { http: reqwest::Client::new(), @@ -112,6 +105,12 @@ impl ColpaliClient { } } +impl MemoInput for ColpaliClient { + fn write_memo_key(&self, writer: &mut cocoindex::memo::MemoKeyWriter<'_>) -> Result<()> { + writer.write(&self.url) + } +} + /// A computed point: stable id + multi-vector embedding + source filename. #[derive(Clone, Serialize, Deserialize)] struct PointData { @@ -135,7 +134,6 @@ async fn process_image(ctx: &Ctx, file: FileEntry) -> Result { } async fn app_main(ctx: Ctx, sourcedir: PathBuf) -> Result<()> { - let conn = ctx.get_key(&DB)?; let target = qdrant::mount_collection_target( &ctx, &DB, diff --git a/examples/rust/kafka_consume/Cargo.lock b/examples/rust/kafka_consume/Cargo.lock index f134b9e4f..a461a0538 100644 --- a/examples/rust/kafka_consume/Cargo.lock +++ b/examples/rust/kafka_consume/Cargo.lock @@ -244,6 +244,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "cc" diff --git a/examples/rust/meeting_notes_graph_falkordb/Cargo.lock b/examples/rust/meeting_notes_graph_falkordb/Cargo.lock index f55393ed4..08966b06d 100644 --- a/examples/rust/meeting_notes_graph_falkordb/Cargo.lock +++ b/examples/rust/meeting_notes_graph_falkordb/Cargo.lock @@ -244,6 +244,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "cc" @@ -285,6 +288,7 @@ version = "0.1.0" dependencies = [ "async-trait", "bytes", + "chrono", "cocoindex_core", "cocoindex_macros", "cocoindex_utils", diff --git a/examples/rust/meeting_notes_graph_neo4j/Cargo.lock b/examples/rust/meeting_notes_graph_neo4j/Cargo.lock index 0deb8ebc1..934f0e053 100644 --- a/examples/rust/meeting_notes_graph_neo4j/Cargo.lock +++ b/examples/rust/meeting_notes_graph_neo4j/Cargo.lock @@ -324,6 +324,7 @@ version = "0.1.0" dependencies = [ "async-trait", "bytes", + "chrono", "cocoindex_core", "cocoindex_macros", "cocoindex_utils", diff --git a/examples/rust/multi_codebase_summarization/Cargo.lock b/examples/rust/multi_codebase_summarization/Cargo.lock index 916a76e95..df12c1f10 100644 --- a/examples/rust/multi_codebase_summarization/Cargo.lock +++ b/examples/rust/multi_codebase_summarization/Cargo.lock @@ -235,6 +235,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "cc" diff --git a/examples/rust/multi_codebase_summarization/src/models.rs b/examples/rust/multi_codebase_summarization/src/models.rs index 0e3811b55..6b30b2c74 100644 --- a/examples/rust/multi_codebase_summarization/src/models.rs +++ b/examples/rust/multi_codebase_summarization/src/models.rs @@ -3,7 +3,7 @@ use serde::{Deserialize, Serialize}; /// Information about a public function. -#[derive(Clone, Debug, Serialize, Deserialize)] +#[derive(Clone, Debug, Serialize, Deserialize, cocoindex::MemoInput)] pub struct FunctionInfo { /// Function name. pub name: String, @@ -16,7 +16,7 @@ pub struct FunctionInfo { } /// Information about a public class. -#[derive(Clone, Debug, Serialize, Deserialize)] +#[derive(Clone, Debug, Serialize, Deserialize, cocoindex::MemoInput)] pub struct ClassInfo { /// Class name. pub name: String, @@ -25,7 +25,7 @@ pub struct ClassInfo { } /// Extracted information from Python code (file or project level). -#[derive(Clone, Debug, Default, Serialize, Deserialize)] +#[derive(Clone, Debug, Default, Serialize, Deserialize, cocoindex::MemoInput)] pub struct CodebaseInfo { /// File path (for files) or project name (for projects). pub name: String, diff --git a/examples/rust/oci_object_storage_embedding/Cargo.lock b/examples/rust/oci_object_storage_embedding/Cargo.lock index 8dd7ca63e..3c3bd14f0 100644 --- a/examples/rust/oci_object_storage_embedding/Cargo.lock +++ b/examples/rust/oci_object_storage_embedding/Cargo.lock @@ -411,6 +411,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "castaway" @@ -478,6 +481,7 @@ dependencies = [ "reqwest", "rmp-serde", "rsa", + "rust_decimal", "rustc-hash", "serde", "serde_json", @@ -3280,6 +3284,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rust_decimal" +version = "1.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" +dependencies = [ + "arrayvec", + "num-traits", + "serde", + "wasm-bindgen", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -4935,6 +4951,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] diff --git a/examples/rust/oci_object_storage_embedding/src/main.rs b/examples/rust/oci_object_storage_embedding/src/main.rs index 473396694..8b3e026cb 100644 --- a/examples/rust/oci_object_storage_embedding/src/main.rs +++ b/examples/rust/oci_object_storage_embedding/src/main.rs @@ -30,17 +30,11 @@ const PG_SCHEMA: &str = "coco_examples"; const TABLE: &str = "oci_doc_embeddings"; const TOP_K: i64 = 5; +cocoindex::context_key!(static DB: postgres::Database); +cocoindex::context_key!(static OCI: OciClient); cocoindex::context_key!( - static DB: postgres::Database = "oci_embedding_db", - state = postgres::Database::state_id -); -cocoindex::context_key!( - static OCI: OciClient = "oci_client", - state = OciClient::state_id -); -cocoindex::context_key!( - static EMBEDDER: SentenceTransformerEmbedder = "embedder", - state = SentenceTransformerEmbedder::model_name + static EMBEDDER: SentenceTransformerEmbedder, + detect_change ); #[derive(Clone, Serialize, Deserialize, SchemaFields)] diff --git a/examples/rust/paper_metadata/Cargo.lock b/examples/rust/paper_metadata/Cargo.lock index f870c1bef..b93eaf587 100644 --- a/examples/rust/paper_metadata/Cargo.lock +++ b/examples/rust/paper_metadata/Cargo.lock @@ -437,6 +437,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "castaway" @@ -531,6 +534,7 @@ dependencies = [ "globset", "linkme", "rmp-serde", + "rust_decimal", "rustc-hash", "serde", "serde_json", @@ -3446,6 +3450,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rust_decimal" +version = "1.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" +dependencies = [ + "arrayvec", + "num-traits", + "serde", + "wasm-bindgen", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -5107,6 +5123,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] diff --git a/examples/rust/paper_metadata/src/main.rs b/examples/rust/paper_metadata/src/main.rs index 9e3a4e420..c58eb5c91 100644 --- a/examples/rust/paper_metadata/src/main.rs +++ b/examples/rust/paper_metadata/src/main.rs @@ -50,17 +50,14 @@ const ABSTRACT_MIN_CHUNK_SIZE: usize = 200; const ABSTRACT_CHUNK_OVERLAP: usize = 150; const LLM_INPUT_CHARS: usize = 4000; +cocoindex::context_key!(static DB: postgres::Database); cocoindex::context_key!( - static DB: postgres::Database = "paper_metadata_db", - state = postgres::Database::state_id + static EMBEDDER: SentenceTransformerEmbedder, + detect_change ); cocoindex::context_key!( - static EMBEDDER: SentenceTransformerEmbedder = "embedder", - state = SentenceTransformerEmbedder::model_name -); -cocoindex::context_key!( - static LLM: LlmClient = "llm_model", - state = LlmClient::model_name + static LLM: LlmClient, + detect_change ); // --------------------------------------------------------------------------- @@ -76,10 +73,6 @@ struct LlmClient { } impl LlmClient { - fn model_name(&self) -> &str { - &self.model - } - fn new(model: String) -> Result { let api_key = std::env::var("OPENAI_API_KEY").map_err(|_| Error::engine("set OPENAI_API_KEY"))?; @@ -132,6 +125,12 @@ impl LlmClient { } } +impl MemoInput for LlmClient { + fn write_memo_key(&self, writer: &mut cocoindex::memo::MemoKeyWriter<'_>) -> Result<()> { + writer.write(&self.model) + } +} + // --------------------------------------------------------------------------- // Data models // --------------------------------------------------------------------------- @@ -236,7 +235,7 @@ async fn extract_metadata(llm: &LlmClient, first_page_text: &str) -> Result Result { let filename = file.key(); let content = file.content()?; @@ -267,7 +266,7 @@ async fn process_file(ctx: &Ctx, file: &FileEntry) -> Result { let mut embeddings = Vec::new(); // Title embedding (one row). - let title_vec = embedder.embed(ctx, &metadata.title).await?; + let title_vec = embedder.embed(&ctx, &metadata.title).await?; let title_id = uuid_gen .next_uuid(&ctx, &("title", &metadata.title)) .await? diff --git a/examples/rust/pdf_embedding/Cargo.lock b/examples/rust/pdf_embedding/Cargo.lock index 974e62906..605fa14b2 100644 --- a/examples/rust/pdf_embedding/Cargo.lock +++ b/examples/rust/pdf_embedding/Cargo.lock @@ -437,6 +437,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "castaway" @@ -531,6 +534,7 @@ dependencies = [ "globset", "linkme", "rmp-serde", + "rust_decimal", "rustc-hash", "serde", "serde_json", @@ -3444,6 +3448,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rust_decimal" +version = "1.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" +dependencies = [ + "arrayvec", + "num-traits", + "serde", + "wasm-bindgen", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -5105,6 +5121,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] diff --git a/examples/rust/pdf_embedding/src/main.rs b/examples/rust/pdf_embedding/src/main.rs index 04735f84a..fc888ddfb 100644 --- a/examples/rust/pdf_embedding/src/main.rs +++ b/examples/rust/pdf_embedding/src/main.rs @@ -30,13 +30,10 @@ const TOP_K: i64 = 5; const CHUNK_SIZE: usize = 2000; const CHUNK_OVERLAP: usize = 500; +cocoindex::context_key!(static DB: postgres::Database); cocoindex::context_key!( - static DB: postgres::Database = "pdf_embedding_db", - state = postgres::Database::state_id -); -cocoindex::context_key!( - static EMBEDDER: SentenceTransformerEmbedder = "embedder", - state = SentenceTransformerEmbedder::model_name + static EMBEDDER: SentenceTransformerEmbedder, + detect_change ); #[derive(Clone, Serialize, Deserialize, SchemaFields)] diff --git a/examples/rust/pdf_to_markdown/Cargo.lock b/examples/rust/pdf_to_markdown/Cargo.lock index 88d8ae7df..80078857f 100644 --- a/examples/rust/pdf_to_markdown/Cargo.lock +++ b/examples/rust/pdf_to_markdown/Cargo.lock @@ -276,6 +276,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "cbc" diff --git a/examples/rust/postgres_source/Cargo.lock b/examples/rust/postgres_source/Cargo.lock index debddad4a..1e5ce57d4 100644 --- a/examples/rust/postgres_source/Cargo.lock +++ b/examples/rust/postgres_source/Cargo.lock @@ -411,6 +411,9 @@ name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +dependencies = [ + "serde", +] [[package]] name = "castaway" @@ -473,6 +476,7 @@ dependencies = [ "globset", "linkme", "rmp-serde", + "rust_decimal", "rustc-hash", "serde", "serde_json", @@ -3216,6 +3220,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rust_decimal" +version = "1.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" +dependencies = [ + "arrayvec", + "num-traits", + "serde", + "wasm-bindgen", +] + [[package]] name = "rustc-hash" version = "2.1.2" @@ -4498,6 +4514,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] diff --git a/examples/rust/postgres_source/src/main.rs b/examples/rust/postgres_source/src/main.rs index ec35e41ed..ee9f6c6db 100644 --- a/examples/rust/postgres_source/src/main.rs +++ b/examples/rust/postgres_source/src/main.rs @@ -30,19 +30,13 @@ const TABLE: &str = "output"; const TOP_K: i64 = 5; // Target database. -cocoindex::context_key!( - static DB: postgres::Database = "postgres_source_db", - state = postgres::Database::state_id -); +cocoindex::context_key!(static DB: postgres::Database); // Source database. Defaults to the target URL, but can point elsewhere via // `SOURCE_DATABASE_URL`, matching the Python example. +cocoindex::context_key!(static SOURCE_DB: postgres::Database); cocoindex::context_key!( - static SOURCE_DB: postgres::Database = "source_pool", - state = postgres::Database::state_id -); -cocoindex::context_key!( - static EMBEDDER: SentenceTransformerEmbedder = "embedder", - state = SentenceTransformerEmbedder::model_name + static EMBEDDER: SentenceTransformerEmbedder, + detect_change ); // --------------------------------------------------------------------------- @@ -50,7 +44,7 @@ cocoindex::context_key!( // --------------------------------------------------------------------------- /// One row of the `source_products` source table (extra columns are ignored). -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, MemoInput)] struct SourceProduct { product_category: String, product_name: String, @@ -78,7 +72,7 @@ struct OutputProduct { /// Compute the derived fields + embedding for one source row. Memoized by the /// row content, so unchanged rows skip the embedding work on re-runs. -#[cocoindex::function] +#[cocoindex::function(memo)] async fn process_product(ctx: &Ctx, product: SourceProduct) -> Result { let full_description = format!( "Category: {}\nName: {}\n\n{}", @@ -87,7 +81,7 @@ async fn process_product(ctx: &Ctx, product: SourceProduct) -> Result Result<()> { let query_vec = Embedder::embed(embedder, query).await?; - let hits = turbopuffer::vector_search_by_field(conn, namespace, "vector", query_vec, TOP_K) - .await?; + let hits = + turbopuffer::vector_search_by_field(conn, namespace, "vector", query_vec, TOP_K).await?; for hit in hits { let filename = hit .attributes diff --git a/rust/sdk/SHOWCASE.md b/rust/sdk/SHOWCASE.md index 79a319afb..1b443d720 100644 --- a/rust/sdk/SHOWCASE.md +++ b/rust/sdk/SHOWCASE.md @@ -45,25 +45,48 @@ let app = Environment::builder() ## Declare typed context resources `context_key!` gives a resource a stable name and type without requiring users -to write their own `LazyLock>`. The name is persistent identity; -do not derive it from the Rust module path. +to write their own `LazyLock>`. The Rust identifier is the key +name by default. Use `key = "..."` when the persistent identity should differ +from the identifier or survive an identifier rename. ```rust use cocoindex::connectors::postgres; +#[derive(cocoindex::MemoInput)] +struct AppConfig { + model: String, + debug: bool, +} + +cocoindex::context_key!(static DB: postgres::Database); +cocoindex::context_key!(static CONFIG: AppConfig, detect_change); +cocoindex::context_key!(static CLIENT: ApiClient); + +// Use an explicit key only when its persisted identity should differ from +// the Rust identifier or remain stable across an identifier rename. cocoindex::context_key!( - static DB: postgres::Database = "app_database", - state = postgres::Database::state_id + static LEGACY_DB: postgres::Database, + key = "app_database" ); -cocoindex::context_key!(static CONFIG: AppConfig = "app_config", detect_change); -cocoindex::context_key!(static CLIENT: ApiClient = "api_client"); ``` - The plain form provides a typed resource without change tracking. -- `detect_change` fingerprints the complete serializable value. -- `state = expression` fingerprints only a stable derived state, which is - useful for connections and model clients whose runtime handles are not - serializable. +- `detect_change` uses the value type's `MemoInput` implementation. Derive + `MemoInput` for ordinary structs and enums; each field is handled + recursively, so the containing type does not need `Serialize`. Resource + types can define a stable identity and external-state validation themselves. +- `MemoInput` composes through `Option`, sequences, tuples, maps, sets, `Box`, + and `Arc`. UUID support is always available; enable the standalone + `serde_json`, `chrono`, or `rust_decimal` feature for those leaf types. + Connector features enable the corresponding leaf feature automatically. + Wrap other third-party leaf types in a newtype when they appear inside + derived data. +- When both optional arguments are present, write `key = "..."` before + `detect_change`. + +`Vec` follows normal element-by-element `Vec` semantics. Use +`bytes::Bytes` when a large byte buffer should be hashed in bulk, or pass a +precomputed `Fingerprint` when one is already available. Provide resources with `EnvironmentBuilder::provide_key` and read them with `ctx.get_key(&KEY)`. Reads inside memoized functions are tracked as @@ -72,7 +95,7 @@ dependencies when the key uses a change-detecting form. ## Define functions `#[cocoindex::function]` tracks the function's logic. Adding `memo` caches the -result by the function logic, serializable arguments, and context dependencies: +result by the function logic, `MemoInput` arguments, and context dependencies: ```rust #[cocoindex::function(memo)] @@ -81,9 +104,10 @@ async fn parse_file(_ctx: &Ctx, file: FileEntry) -> Result> { } ``` -Use `memo_key(...)` when the default representation of an argument is either -too broad or not serializable. A transform replaces that argument's memo key; -`skip` (also spelled `None`) excludes it: +Use `memo_key(...)` when an argument's default `MemoInput` identity is too +broad or the type does not implement `MemoInput`. A transform replaces that +argument's memo key and external-state validation; `skip` (also spelled +`None`) excludes it: ```rust fn entry_identity(entry: &Entry) -> (String, u64) { @@ -134,7 +158,10 @@ let embeddings = ctx The body receives only the items in the current batch. With `memo`, cache hits are returned per item and only misses enter the batch. Without `memo`, every -call is processed. `max_batch_size` caps each physical request. +call is processed. `max_batch_size` caps each physical request. Additional +parameters must implement `Serialize` because they identify compatible calls +for batching; with `memo`, they must also implement `MemoInput` unless their +memo identity is transformed or skipped. A physical batch does not inherit any individual caller's deadline. A batch error currently fails every item in that physical batch; the Rust SDK does not @@ -170,8 +197,8 @@ closures; they still provide stable child ownership and reconciliation. ## Read files and declare output files `walk_items` produces stable `(relative_path, FileEntry)` pairs ready for -`mount_each!`. `FileEntry` is serializable, so a mounted or memoized function -can use it directly as an input. +`mount_each!`. `FileEntry` implements `MemoInput`, including content freshness +validation, so a memoized function can use it directly or inside a container. ```rust #[cocoindex::function] diff --git a/rust/sdk/cocoindex/Cargo.toml b/rust/sdk/cocoindex/Cargo.toml index 6bcc91d4e..2c1666721 100644 --- a/rust/sdk/cocoindex/Cargo.toml +++ b/rust/sdk/cocoindex/Cargo.toml @@ -21,56 +21,86 @@ code_match = [ ] # Local sentence-transformer embeddings via fastembed (`ops::sentence_transformers`). fastembed = ["dep:fastembed"] +# MemoInput leaves for ecosystem data types. Connector features below depend on +# these public features rather than duplicating implementation cfg lists. +serde_json = ["dep:serde_json"] +chrono = ["dep:chrono"] +rust_decimal = ["dep:rust_decimal"] # Remote embeddings / transcription over an OpenAI-compatible HTTP API (`ops::api`). -embed_api = ["dep:serde_json", "dep:reqwest", "reqwest/multipart"] -postgres = ["dep:serde_json", "dep:sqlx", "dep:chrono"] -sqlite = ["dep:serde_json", "dep:sqlx", "sqlx/sqlite"] +embed_api = ["serde_json", "dep:reqwest", "reqwest/multipart"] +postgres = ["serde_json", "dep:sqlx", "chrono", "rust_decimal"] +sqlite = [ + "serde_json", + "dep:sqlx", + "sqlx/sqlite", + "chrono", + "rust_decimal", +] # Apache Doris target: MySQL protocol (DDL/DELETE) via sqlx + Stream Load (row # ingestion) over HTTP via reqwest. doris = [ - "dep:serde_json", + "serde_json", "dep:sqlx", "sqlx/mysql", "dep:reqwest", "dep:base64", + "chrono", + "rust_decimal", ] -surrealdb = ["dep:serde_json", "dep:surrealdb"] -kafka = ["dep:rskafka", "dep:chrono"] +surrealdb = [ + "serde_json", + "dep:surrealdb", + "chrono", + "rust_decimal", +] +kafka = ["dep:rskafka", "chrono"] iggy = ["dep:iggy"] amazon_s3 = ["dep:aws-config", "dep:aws-sdk-s3"] google_drive = [ - "dep:serde_json", + "serde_json", "dep:reqwest", "dep:rsa", "dep:sha2", "dep:base64", - "dep:chrono", + "chrono", ] # OCI Object Storage source: REST API + RSA-SHA256 HTTP Signature request # signing (no official Oracle Rust SDK). Reuses the google_drive crypto/HTTP # crates; `chrono/clock` provides the request `date` header + timestamp parsing. oci_object_storage = [ - "dep:serde_json", + "serde_json", "dep:reqwest", "dep:rsa", "dep:sha2", "dep:base64", - "dep:chrono", + "chrono", "chrono/clock", ] lancedb = [ - "dep:serde_json", + "serde_json", "dep:lancedb", "dep:arrow-array", "dep:arrow-schema", "dep:half", + "chrono", + "rust_decimal", +] +qdrant = [ + "serde_json", + "dep:qdrant-client", + "chrono", + "rust_decimal", +] +turbopuffer = [ + "serde_json", + "dep:reqwest", + "chrono", + "rust_decimal", ] -qdrant = ["dep:serde_json", "dep:qdrant-client"] -turbopuffer = ["dep:serde_json", "dep:reqwest"] -neo4j = ["dep:serde_json", "dep:neo4rs"] -falkordb = ["dep:serde_json", "dep:redis"] +neo4j = ["serde_json", "dep:neo4rs", "chrono"] +falkordb = ["serde_json", "dep:redis", "chrono"] # Valkey/Redis (RediSearch) vector-search target. Reuses the `redis` crate. -valkey = ["dep:redis", "dep:serde_json"] +valkey = ["dep:redis", "serde_json"] [dependencies] cocoindex_macros = { path = "../cocoindex_macros" } @@ -86,7 +116,7 @@ cocoindex_code_match = { path = "../../code_match", optional = true } # the ONNX runtime, so it is gated behind the `fastembed` feature. fastembed = { version = "4", optional = true } thiserror = "2" -bytes = "1" +bytes = { version = "1", features = ["serde"] } serde = { version = "1", features = ["derive"] } rmp-serde = "1" tokio = { version = "1", features = ["fs", "sync"] } @@ -110,7 +140,8 @@ sqlx = { version = "0.8", default-features = false, optional = true, features = "postgres", "chrono", ] } -chrono = { version = "0.4", default-features = false, optional = true } +chrono = { version = "0.4", default-features = false, optional = true, features = ["serde"] } +rust_decimal = { version = "1", default-features = false, optional = true, features = ["serde"] } rskafka = { version = "0.6", default-features = false, optional = true } # Apache Iggy message-streaming target (pure-Rust client; TCP/QUIC/HTTP). iggy = { version = "0.10", optional = true } @@ -152,7 +183,6 @@ rskafka = { version = "0.6", default-features = false } wiremock = "0.6" # Generate an ephemeral RSA key for the OCI signing unit test. rand = "0.8" -rust_decimal = "1" [[bench]] name = "sdk_microbench" diff --git a/rust/sdk/cocoindex/src/app.rs b/rust/sdk/cocoindex/src/app.rs index c9606779f..b64082203 100644 --- a/rust/sdk/cocoindex/src/app.rs +++ b/rust/sdk/cocoindex/src/app.rs @@ -180,7 +180,7 @@ impl EnvironmentBuilder { /// /// Returns an error if the LMDB database environment fails to initialize /// (e.g., due to permissions, disk space, or a corrupted state directory). - pub async fn build(self) -> Result { + pub async fn build(mut self) -> Result { // Register every `#[coco::function]`'s logic fingerprint into the engine's // logic set, so memo entries that depend on them validate correctly // (see `crate::logic`). Idempotent across builds. @@ -202,6 +202,7 @@ impl EnvironmentBuilder { .await .map_err(|e| Error::engine(format!("failed to open LMDB: {e}")))?; self.context.register_logic(&core_env); + self.context.register_initial_states(&core_env).await?; Ok(Environment { inner: Arc::new(EnvironmentInner { diff --git a/rust/sdk/cocoindex/src/connectors/amazon_s3.rs b/rust/sdk/cocoindex/src/connectors/amazon_s3.rs index 045cca161..414507c67 100644 --- a/rust/sdk/cocoindex/src/connectors/amazon_s3.rs +++ b/rust/sdk/cocoindex/src/connectors/amazon_s3.rs @@ -213,6 +213,8 @@ impl FileLike for S3File { impl FileSourceItem for S3File {} +crate::memo::impl_file_memo_input!(S3File); + fn default_file_cache() -> Arc { Arc::new(FileContentCache::new()) } diff --git a/rust/sdk/cocoindex/src/connectors/falkordb.rs b/rust/sdk/cocoindex/src/connectors/falkordb.rs index ab789da4e..360768212 100644 --- a/rust/sdk/cocoindex/src/connectors/falkordb.rs +++ b/rust/sdk/cocoindex/src/connectors/falkordb.rs @@ -130,7 +130,7 @@ mod tests { use super::*; use crate::{App, Environment}; - crate::context_key!(static GRAPH: Graph = "falkordb_graph"); + crate::context_key!(static GRAPH: Graph, key = "falkordb_graph"); #[derive(Serialize)] struct Person { diff --git a/rust/sdk/cocoindex/src/connectors/gdrive.rs b/rust/sdk/cocoindex/src/connectors/gdrive.rs index 6b4832c9d..2a73fbdee 100644 --- a/rust/sdk/cocoindex/src/connectors/gdrive.rs +++ b/rust/sdk/cocoindex/src/connectors/gdrive.rs @@ -220,6 +220,8 @@ impl FileSourceItem for DriveFile { } } +crate::memo::impl_file_memo_input!(DriveFile); + fn default_file_cache() -> Arc { Arc::new(FileContentCache::new()) } diff --git a/rust/sdk/cocoindex/src/connectors/neo4j.rs b/rust/sdk/cocoindex/src/connectors/neo4j.rs index 9321dd492..c19e9b55f 100644 --- a/rust/sdk/cocoindex/src/connectors/neo4j.rs +++ b/rust/sdk/cocoindex/src/connectors/neo4j.rs @@ -142,7 +142,7 @@ mod tests { use super::*; use crate::{App, Environment}; - crate::context_key!(static GRAPH: Graph = "neo4j_graph"); + crate::context_key!(static GRAPH: Graph, key = "neo4j_graph"); #[derive(Serialize)] struct Person { diff --git a/rust/sdk/cocoindex/src/connectors/oci_object_storage.rs b/rust/sdk/cocoindex/src/connectors/oci_object_storage.rs index 1e1fb7968..0d8207632 100644 --- a/rust/sdk/cocoindex/src/connectors/oci_object_storage.rs +++ b/rust/sdk/cocoindex/src/connectors/oci_object_storage.rs @@ -351,6 +351,8 @@ impl FileLike for OciFile { impl FileSourceItem for OciFile {} +crate::memo::impl_file_memo_input!(OciFile); + fn default_file_cache() -> Arc { Arc::new(FileContentCache::new()) } diff --git a/rust/sdk/cocoindex/src/ctx.rs b/rust/sdk/cocoindex/src/ctx.rs index d5ca49674..81cfa0ab6 100644 --- a/rust/sdk/cocoindex/src/ctx.rs +++ b/rust/sdk/cocoindex/src/ctx.rs @@ -5,6 +5,7 @@ use std::collections::{HashMap, HashSet}; use std::fmt::Display; use std::future::Future; use std::marker::PhantomData; +use std::pin::Pin; use std::sync::{Arc, Mutex, OnceLock}; use std::time::Duration; @@ -25,65 +26,83 @@ use crate::live_component::{ ExceptionContext, ExceptionHandler, LiveComponent, LiveMapView, MountEachLiveComponent, MountKind, build_chained_on_error, new_operator, start_process_live, }; +use crate::memo::{MemoInput, MemoStateDecision, MemoStateValue}; use crate::profile::{BoxedHandler, BoxedProcessor, RustProfile, Value}; use crate::user_state::{IntoStateKey, StateHandle}; type ContextFingerprinter = Arc Result + Send + Sync>; +type ContextStateFuture = + Pin>> + Send + 'static>>; +type ContextStateFn = + Arc, Option) -> ContextStateFuture + Send + Sync>; +type StoredContextStateFn = Arc) -> ContextStateFuture + Send + Sync>; /// Declare a process-wide [`ContextKey`] backed by a [`std::sync::LazyLock`]. /// -/// The three forms correspond to [`ContextKey::new`], -/// [`ContextKey::new_detect_change`], and [`ContextKey::new_with_state`]: +/// By default, the Rust identifier is also the context key's stable name. Use +/// `key = "..."` when the persisted name must differ from the identifier or +/// remain stable across an identifier rename. Change detection is disabled by +/// default. With `detect_change`, the value's [`MemoInput`] implementation +/// defines its memo identity and optional external-state validation. For a +/// foreign type that cannot implement [`MemoInput`] because of Rust's orphan +/// rules, use a local newtype with a manual [`MemoInput`] implementation. /// /// ``` -/// # #[derive(serde::Serialize)] +/// # #[derive(cocoindex::MemoInput)] /// # struct AppConfig; /// # struct Database; -/// # impl Database { fn state_id(&self) -> &str { "database" } } -/// cocoindex::context_key!(static CONFIG: AppConfig = "app_config"); -/// cocoindex::context_key!(static TRACKED_CONFIG: AppConfig = "tracked_config", detect_change); +/// cocoindex::context_key!(static CONFIG: AppConfig); +/// cocoindex::context_key!(static TRACKED_CONFIG: AppConfig, detect_change); /// cocoindex::context_key!( -/// static DB: Database = "database", -/// state = Database::state_id +/// static DB: Database, +/// key = "my_package/database" /// ); /// ``` /// -/// The key name must be an explicit string literal. It is a stable persistent -/// identity used in target-state and memo keys, so it must not change merely -/// because the Rust item moves to another module. The `state =` expression may -/// be any path or closure implementing `Fn(&T) -> S` where `S: ToString`. +/// The key name is a stable persistent identity used in target-state and memo +/// keys. Renaming an identifier that uses the default name therefore changes +/// its identity; use `key = "..."` when that is not intended. `key =` changes +/// only this logical name. When both optional arguments are present, they must +/// appear in this order: `key = "..."`, then `detect_change`. #[macro_export] macro_rules! context_key { - ($vis:vis static $key:ident: $value:ty = $name:literal, detect_change $(,)?) => { + ($vis:vis static $key:ident: $value:ty, key = $name:literal, detect_change $(,)?) => { $vis static $key: ::std::sync::LazyLock<$crate::ContextKey<$value>> = ::std::sync::LazyLock::new(|| $crate::ContextKey::new_detect_change($name)); }; - ($vis:vis static $key:ident: $value:ty = $name:literal, state = $state:expr $(,)?) => { + ($vis:vis static $key:ident: $value:ty, detect_change $(,)?) => { $vis static $key: ::std::sync::LazyLock<$crate::ContextKey<$value>> = ::std::sync::LazyLock::new(|| { - $crate::ContextKey::new_with_state($name, |value: &$value| { - ($state)(value).to_string() - }) + $crate::ContextKey::new_detect_change(::std::stringify!($key)) }); }; - ($vis:vis static $key:ident: $value:ty = $name:literal $(,)?) => { + ($vis:vis static $key:ident: $value:ty, key = $name:literal $(,)?) => { $vis static $key: ::std::sync::LazyLock<$crate::ContextKey<$value>> = ::std::sync::LazyLock::new(|| $crate::ContextKey::new($name)); }; + ($vis:vis static $key:ident: $value:ty $(,)?) => { + $vis static $key: ::std::sync::LazyLock<$crate::ContextKey<$value>> = + ::std::sync::LazyLock::new(|| $crate::ContextKey::new(::std::stringify!($key))); + }; + ($($invalid:tt)*) => { + ::core::compile_error!( + "invalid context_key! syntax; expected `static NAME: Type` followed by optional \ + `key = \"...\"` and `detect_change` in that order" + ); + }; } /// A named context key for app-provided resources. /// /// - [`ContextKey::new`] stores arbitrary `Send + Sync` resources (no change /// tracking). -/// - [`ContextKey::new_detect_change`] tracks a serializable value: memoized -/// work is invalidated when the whole value's fingerprint changes. -/// - [`ContextKey::new_with_state`] tracks a derived state of an arbitrary -/// value. Only changes to the extracted state invalidate memoized work. +/// - [`ContextKey::new_detect_change`] uses the value type's [`MemoInput`] +/// behavior to invalidate memoized work. pub struct ContextKey { name: Arc, detect_change: bool, fingerprint_fn: Option>, + state_fn: Option>, _marker: PhantomData T>, } @@ -93,6 +112,7 @@ impl Clone for ContextKey { name: self.name.clone(), detect_change: self.detect_change, fingerprint_fn: self.fingerprint_fn.clone(), + state_fn: self.state_fn.clone(), _marker: PhantomData, } } @@ -105,7 +125,7 @@ impl ContextKey { /// Panics if the same key name has already been constructed in this /// process. pub fn new(name: impl Into) -> Self { - Self::with_parts(name.into(), false, None) + Self::with_parts(name.into(), false, None, None) } /// The stable key name. @@ -122,56 +142,39 @@ impl ContextKey { name: String, detect_change: bool, fingerprint_fn: Option>, + state_fn: Option>, ) -> Self { let used = USED_CONTEXT_KEYS.get_or_init(|| Mutex::new(HashSet::new())); let duplicate = { let mut used = used.lock().expect("context key registry poisoned"); !used.insert(name.clone()) }; - assert!(!duplicate, "Context key {name} already used"); + assert!( + !duplicate, + "Context key {name:?} is already used; choose a unique identifier or set \ + `key = \"stable_unique_name\"`" + ); Self { name: Arc::from(name), detect_change, fingerprint_fn, + state_fn, _marker: PhantomData, } } } -impl ContextKey { +impl ContextKey { /// Create a named context key whose provided values invalidate memoized - /// work when their serialized fingerprint changes. + /// work when their type-defined memo identity or external state changes. pub fn new_detect_change(name: impl Into) -> Self { let fingerprint_fn: ContextFingerprinter = Arc::new(|name: &str, value: &T| { - Fingerprint::from(&("context_key", name, value)) - .map_err(|e| Error::engine(format!("context key fingerprint error: {e}"))) + crate::memo::memo_input_fingerprint("context_key", name, value) }); - Self::with_parts(name.into(), true, Some(fingerprint_fn)) - } -} - -impl ContextKey { - /// Create a named context key whose memo invalidation is driven by a - /// *derived state* rather than the whole value. `state_fn` extracts a - /// serializable state from the provided value; memoized work that reads - /// this key (via [`Ctx::get_key`]) is invalidated only when that state's - /// fingerprint changes. - /// - /// Use this for resources that are not serializable, such as DB pools or - /// clients, or when only a narrow identity like a connection string or - /// schema version should affect memoization. The value type `T` need not be - /// `Serialize`; only the extracted state must be. - pub fn new_with_state(name: impl Into, state_fn: SF) -> Self - where - S: Serialize, - SF: Fn(&T) -> S + Send + Sync + 'static, - { - let fingerprint_fn: ContextFingerprinter = Arc::new(move |name: &str, value: &T| { - let state = state_fn(value); - Fingerprint::from(&("context_key", name, &state)) - .map_err(|e| Error::engine(format!("context key state fingerprint error: {e}"))) + let state_fn: ContextStateFn = Arc::new(|value, previous| { + Box::pin(async move { value.memo_state(previous.as_ref()).await }) }); - Self::with_parts(name.into(), true, Some(fingerprint_fn)) + Self::with_parts(name.into(), true, Some(fingerprint_fn), Some(state_fn)) } } @@ -181,6 +184,7 @@ static USED_CONTEXT_KEYS: OnceLock>> = OnceLock::new(); pub(crate) struct ContextStore { values: HashMap, Arc>, fingerprints: HashMap, Fingerprint>, + state_fns: HashMap, } impl ContextStore { @@ -189,11 +193,22 @@ impl ContextStore { key: &ContextKey, value: T, ) -> Result<()> { + let value = Arc::new(value); if let Some(fingerprint_fn) = &key.fingerprint_fn { - let fp = fingerprint_fn(&key.name, &value)?; - self.fingerprints.insert(key.name.clone(), fp); + let fp = fingerprint_fn(&key.name, value.as_ref())?; + if let Some(previous_fp) = self.fingerprints.insert(key.name.clone(), fp) { + self.state_fns.remove(&previous_fp); + } + if let Some(state_fn) = &key.state_fn { + let state_fn = Arc::clone(state_fn); + let value = Arc::clone(&value); + self.state_fns.insert( + fp, + Arc::new(move |previous| state_fn(Arc::clone(&value), previous)), + ); + } } - self.values.insert(key.name.clone(), Arc::new(value)); + self.values.insert(key.name.clone(), value); Ok(()) } @@ -203,6 +218,78 @@ impl ContextStore { } } + pub(crate) async fn register_initial_states( + &mut self, + env: &Environment, + ) -> Result<()> { + let fingerprints = self.state_fns.keys().copied().collect::>(); + for fingerprint in fingerprints { + let state_fn = self + .state_fns + .get(&fingerprint) + .expect("context state function disappeared during initialization"); + if let Some(decision) = state_fn(None).await? { + let (state, _) = decision.into_parts(); + env.register_context_initial_states(fingerprint, vec![state]); + } else { + self.state_fns.remove(&fingerprint); + } + } + Ok(()) + } + + pub(crate) fn has_memo_state_fns(&self) -> bool { + !self.state_fns.is_empty() + } + + pub(crate) async fn validate_memo_states( + &self, + stored: &[(Fingerprint, Vec)], + ) -> Result { + let mut updated = Vec::with_capacity(stored.len()); + let mut memo_valid = true; + + for (fingerprint, states) in stored { + let Some(state_fn) = self.state_fns.get(fingerprint) else { + memo_valid = false; + updated.push((*fingerprint, states.clone())); + continue; + }; + if states.len() != 1 { + memo_valid = false; + updated.push((*fingerprint, states.clone())); + continue; + } + let previous = MemoStateValue::from_profile_value(states[0].clone()); + let Some(decision) = state_fn(Some(previous)).await? else { + memo_valid = false; + updated.push((*fingerprint, states.clone())); + continue; + }; + let (state, valid) = decision.into_parts(); + memo_valid &= valid; + updated.push((*fingerprint, vec![state])); + } + + let states_changed = + updated + .iter() + .zip(stored) + .any(|((new_fp, new_values), (old_fp, old_values))| { + new_fp != old_fp + || new_values.len() != old_values.len() + || new_values + .iter() + .zip(old_values) + .any(|(new, old)| new.0 != old.0) + }); + Ok(ContextMemoValidation { + states: updated, + memo_valid, + states_changed, + }) + } + fn get(&self, key: &ContextKey) -> Option<&T> { self.values .get(&key.name) @@ -227,6 +314,12 @@ impl ContextStore { } } +pub(crate) struct ContextMemoValidation { + pub(crate) states: Vec<(Fingerprint, Vec)>, + pub(crate) memo_valid: bool, + pub(crate) states_changed: bool, +} + /// Pipeline context passed to closures inside `App::update()` / `App::run()`. #[derive(Clone)] pub struct Ctx { @@ -474,10 +567,11 @@ impl Ctx { { if let Some(fn_ctx) = &self.fn_ctx { fn_ctx.add_context_change_dep(fp); - } else if let Some(comp_ctx) = &self.comp_ctx { - let fn_ctx = FnCallContext::default(); - fn_ctx.add_context_change_dep(fp); - comp_ctx.join_fn_call(&fn_ctx); + } + if let Some(comp_ctx) = &self.comp_ctx { + let component_dep = FnCallContext::default(); + component_dep.add_context_change_dep(fp); + comp_ctx.join_fn_call(&component_dep); } } Ok(value) @@ -948,6 +1042,7 @@ impl Ctx { let _guard = fn_call_guard(comp_ctx, fn_ctx.clone()); let state = self.state.clone(); + let has_context_memo_state = state.context.has_memo_state_fns(); let scope_fn_ctx = fn_ctx.clone(); let scope_handler_chain = self.handler_chain.clone(); // Foreground child inherits the caller's current scoped deadline — @@ -970,7 +1065,8 @@ impl Ctx { }, memo_fp, format!("mount:{key}"), - ); + ) + .with_memo_state_handler(has_context_memo_state); let handle = match child_component .use_mount(comp_ctx, processor, self.deadline) @@ -1242,26 +1338,16 @@ mod tests { static TEST_CLOCK_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); - #[derive(Serialize)] + #[derive(Serialize, crate::MemoInput)] struct TrackedConfig; - struct StatefulResource; - - impl StatefulResource { - fn state_id(&self) -> u64 { - 42 - } - } - - crate::context_key!(static PLAIN_KEY: String = "ctx_macro_plain"); + crate::context_key!(static PLAIN_KEY: String); + crate::context_key!(static NAMED_KEY: String, key = "ctx_macro_named"); crate::context_key!( - static DETECT_CHANGE_KEY: TrackedConfig = "ctx_macro_detect_change", + static DETECT_CHANGE_KEY: TrackedConfig, + key = "ctx_macro_detect_change", detect_change ); - crate::context_key!( - static DERIVED_STATE_KEY: StatefulResource = "ctx_macro_derived_state", - state = StatefulResource::state_id - ); struct TestClockGuard { _guard: std::sync::MutexGuard<'static, ()>, @@ -1288,13 +1374,13 @@ mod tests { } #[test] - fn context_key_macro_supports_all_three_forms() { - assert_eq!(PLAIN_KEY.name(), "ctx_macro_plain"); + fn context_key_macro_uses_identifier_by_default_and_supports_override() { + assert_eq!(PLAIN_KEY.name(), "PLAIN_KEY"); assert!(!PLAIN_KEY.detect_change()); + assert_eq!(NAMED_KEY.name(), "ctx_macro_named"); + assert!(!NAMED_KEY.detect_change()); assert_eq!(DETECT_CHANGE_KEY.name(), "ctx_macro_detect_change"); assert!(DETECT_CHANGE_KEY.detect_change()); - assert_eq!(DERIVED_STATE_KEY.name(), "ctx_macro_derived_state"); - assert!(DERIVED_STATE_KEY.detect_change()); } #[tokio::test(flavor = "current_thread")] diff --git a/rust/sdk/cocoindex/src/lib.rs b/rust/sdk/cocoindex/src/lib.rs index 7b8b437b3..d7226091e 100644 --- a/rust/sdk/cocoindex/src/lib.rs +++ b/rust/sdk/cocoindex/src/lib.rs @@ -50,6 +50,7 @@ pub use entity_resolution::{ resolve_entities_with_events, }; pub use error::{Error, Result}; +pub use memo::MemoInput; pub use resources::file::{ FileContentCache, FileLike, FileMetadata, FilePath, FilePathMatcher, FileSourceItem, MatchAllFilePathMatcher, PatternFilePathMatcher, @@ -94,7 +95,7 @@ pub use target_state::{ pub use user_state::{IntoStateKey, StateHandle}; // Re-export proc macros -pub use cocoindex_macros::{SchemaFields, function, mount_each, use_mount}; +pub use cocoindex_macros::{MemoInput, SchemaFields, function, mount_each, use_mount}; pub use row_schema::{LogicalType, SchemaField, SchemaFields}; // Re-exported so users can implement the async `LiveComponent` / `LiveMapFeed` diff --git a/rust/sdk/cocoindex/src/memo.rs b/rust/sdk/cocoindex/src/memo.rs index 08841e865..a94ba41ba 100644 --- a/rust/sdk/cocoindex/src/memo.rs +++ b/rust/sdk/cocoindex/src/memo.rs @@ -4,20 +4,38 @@ //! function. The attribute builds the argument key, tracks the function body's //! logic, and passes an owned, memo-scoped [`Ctx`] into the body. Context //! resources can therefore be read normally with [`Ctx::get_or_err`] or -//! [`Ctx::get_key`]. Use `memo_key(parameter = skip)` for an `Any + Clone` -//! parameter that should not participate in the key; it does not need to -//! implement Serde. +//! [`Ctx::get_key`]. Use `memo_key(parameter = skip)` for a `Clone` parameter +//! that should not participate in the key; it does not need to implement +//! [`MemoInput`] or Serde. A transform override similarly replaces the +//! parameter's normal [`MemoInput`] identity and suppresses its type-owned +//! external-state validation. +//! +//! `MemoInput` composes recursively through `Option`, sequences, tuples, +//! maps, sets, `Box`, and `Arc`. The SDK also implements it for its standard +//! scalar schema types, including UUIDs and, when their corresponding SDK +//! features are enabled, `serde_json::Value`, Chrono values, and +//! `rust_decimal::Decimal`. +//! Other third-party types need a local newtype or a declaration/call-site +//! `memo_key` projection because Rust's orphan rules prevent downstream crates +//! from implementing this trait for a foreign type. +//! +//! `Vec` uses the same recursive sequence representation as every other +//! `Vec`. For large byte payloads, prefer [`bytes::Bytes`] for bulk hashing, +//! or pass a precomputed [`Fingerprint`] when the content is already hashed. //! //! [`Ctx::memo`] and [`cached`] are intended for block-level memoization inside //! a function. Their closure body is not logic-tracked. When a manual memo block //! must invalidate after its enclosing `#[cocoindex::function]` changes, include //! that function's generated `__COCO_FN_HASH_` constant in the manual key. -use std::any::Any; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::future::Future; +use std::hash::{BuildHasher, Hash}; +use std::path::{Path, PathBuf}; use std::sync::Arc; -use std::time::{SystemTime, UNIX_EPOCH}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use async_trait::async_trait; use cocoindex_core::engine::context::{FnCallContext, MemoStatesPayload}; use cocoindex_core::engine::function::reserve_memoization; use cocoindex_utils::fingerprint::{Fingerprint, Fingerprinter}; @@ -28,6 +46,426 @@ use crate::error::{Error, Result}; use crate::profile::Value; use crate::resources::file::FileLike; +/// Type-owned memoization behavior for function arguments and tracked context +/// values. +/// +/// Use the [`crate::MemoInput`] derive for normal structs and enums. It adds a +/// `MemoInput` bound to each field type and recursively visits every field; the +/// containing type itself does not need to implement [`Serialize`]. Nested +/// resources therefore preserve their type-owned identity and freshness +/// validation. Resource types can implement this trait directly to provide a +/// small stable identity and asynchronous validation. +#[async_trait] +#[diagnostic::on_unimplemented( + message = "`{Self}` cannot be used as a memoized input", + note = "add `#[derive(cocoindex::MemoInput)]`, implement `MemoInput` manually, or provide a `memo_key` override" +)] +pub trait MemoInput: Send + Sync { + /// Add this value's semantic identity to a function memo key. + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()>; + + /// Validate external state before a matching memo is reused. + /// + /// `None` means that this value needs no state validation. + async fn memo_state( + &self, + _previous: Option<&MemoStateValue>, + ) -> Result> { + Ok(None) + } +} + +macro_rules! impl_serde_memo_input { + ($($ty:ty),+ $(,)?) => { + $( + #[async_trait] + impl MemoInput for $ty { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + writer.write(self) + } + } + )+ + }; +} + +impl_serde_memo_input!( + (), + bool, + char, + i8, + i16, + i32, + i64, + i128, + isize, + u8, + u16, + u32, + u64, + u128, + usize, + f32, + f64, + String, + str, + PathBuf, + Path, + Duration, + SystemTime, + uuid::Uuid, + Fingerprint, +); + +#[async_trait] +impl MemoInput for bytes::Bytes { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + writer.write(self) + } +} + +#[cfg(feature = "serde_json")] +impl_serde_memo_input!(serde_json::Value); + +#[cfg(feature = "chrono")] +impl_serde_memo_input!(chrono::NaiveDate, chrono::NaiveTime, chrono::NaiveDateTime); + +#[cfg(feature = "chrono")] +#[async_trait] +impl MemoInput for chrono::DateTime +where + Tz: chrono::TimeZone + Send + Sync, + Tz::Offset: Send + Sync, + chrono::DateTime: Serialize, +{ + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + writer.write(self) + } +} + +#[cfg(feature = "rust_decimal")] +impl_serde_memo_input!(rust_decimal::Decimal); + +#[async_trait] +impl MemoInput for &T { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + (*self).write_memo_key(writer) + } + + async fn memo_state( + &self, + previous: Option<&MemoStateValue>, + ) -> Result> { + (*self).memo_state(previous).await + } +} + +#[async_trait] +impl MemoInput for Box { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + self.as_ref().write_memo_key(writer) + } + + async fn memo_state( + &self, + previous: Option<&MemoStateValue>, + ) -> Result> { + self.as_ref().memo_state(previous).await + } +} + +#[async_trait] +impl MemoInput for Arc { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + self.as_ref().write_memo_key(writer) + } + + async fn memo_state( + &self, + previous: Option<&MemoStateValue>, + ) -> Result> { + self.as_ref().memo_state(previous).await + } +} + +#[async_trait] +impl MemoInput for Option { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + match self { + Some(value) => { + writer.write(&("option", true))?; + writer.write_input(value) + } + None => writer.write(&("option", false)), + } + } + + async fn memo_state( + &self, + previous: Option<&MemoStateValue>, + ) -> Result> { + match self { + Some(value) => value.memo_state(previous).await, + None => Ok(None), + } + } +} + +#[async_trait] +impl MemoInput for [T] { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + writer.write(&("sequence", self.len()))?; + for item in self { + writer.write_input(item)?; + } + Ok(()) + } + + async fn memo_state( + &self, + previous: Option<&MemoStateValue>, + ) -> Result> { + let children = self + .iter() + .map(|item| item as &dyn MemoInput) + .collect::>(); + memo_state_group(&children, previous, false).await + } +} + +#[async_trait] +impl MemoInput for Vec { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + self.as_slice().write_memo_key(writer) + } + + async fn memo_state( + &self, + previous: Option<&MemoStateValue>, + ) -> Result> { + self.as_slice().memo_state(previous).await + } +} + +#[async_trait] +impl MemoInput for [T; N] { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + self.as_slice().write_memo_key(writer) + } + + async fn memo_state( + &self, + previous: Option<&MemoStateValue>, + ) -> Result> { + self.as_slice().memo_state(previous).await + } +} + +fn memo_value_fingerprint(value: &T) -> Result { + let mut fingerprinter = Fingerprinter::default(); + value.write_memo_key(&mut MemoKeyWriter { + fingerprinter: &mut fingerprinter, + })?; + Ok(fingerprinter.into_fingerprint()) +} + +type SortedMapEntries<'a, K, V> = (Vec<(Fingerprint, Fingerprint, &'a K, &'a V)>, bool); + +fn sorted_map_entries<'a, K, V>( + entries: impl Iterator, +) -> Result> +where + K: MemoInput + 'a, + V: MemoInput + 'a, +{ + let mut entries = entries + .map(|(key, value)| { + Ok(( + memo_value_fingerprint(key)?, + memo_value_fingerprint(value)?, + key, + value, + )) + }) + .collect::>>()?; + entries.sort_by_key(|(key_fp, value_fp, _, _)| (*key_fp, *value_fp)); + let collision = entries.windows(2).any(|pair| pair[0].0 == pair[1].0); + Ok((entries, collision)) +} + +macro_rules! impl_map_memo_input { + ($ty:ty) => { + #[async_trait] + impl MemoInput for $ty + where + K: MemoInput + Eq + Hash, + V: MemoInput, + S: BuildHasher + Send + Sync, + { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + let (entries, collision) = sorted_map_entries(self.iter())?; + writer.write(&("map", entries.len(), collision))?; + for (_, _, key, value) in entries { + writer.write_input(key)?; + writer.write_input(value)?; + } + Ok(()) + } + + async fn memo_state( + &self, + previous: Option<&MemoStateValue>, + ) -> Result> { + let (entries, collision) = sorted_map_entries(self.iter())?; + let children = entries + .iter() + .flat_map(|(_, _, key, value)| { + [*key as &dyn MemoInput, *value as &dyn MemoInput] + }) + .collect::>(); + memo_state_group(&children, previous, collision).await + } + } + }; +} + +impl_map_memo_input!(HashMap); + +#[async_trait] +impl MemoInput for BTreeMap { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + let (entries, collision) = sorted_map_entries(self.iter())?; + writer.write(&("map", entries.len(), collision))?; + for (_, _, key, value) in entries { + writer.write_input(key)?; + writer.write_input(value)?; + } + Ok(()) + } + + async fn memo_state( + &self, + previous: Option<&MemoStateValue>, + ) -> Result> { + let (entries, collision) = sorted_map_entries(self.iter())?; + let children = entries + .iter() + .flat_map(|(_, _, key, value)| [*key as &dyn MemoInput, *value as &dyn MemoInput]) + .collect::>(); + memo_state_group(&children, previous, collision).await + } +} + +fn sorted_set_entries<'a, T: MemoInput + 'a>( + entries: impl Iterator, +) -> Result<(Vec<(Fingerprint, &'a T)>, bool)> { + let mut entries = entries + .map(|value| Ok((memo_value_fingerprint(value)?, value))) + .collect::>>()?; + entries.sort_by_key(|(fingerprint, _)| *fingerprint); + let collision = entries.windows(2).any(|pair| pair[0].0 == pair[1].0); + Ok((entries, collision)) +} + +#[async_trait] +impl MemoInput for HashSet +where + T: MemoInput + Eq + Hash, + S: BuildHasher + Send + Sync, +{ + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + let (entries, collision) = sorted_set_entries(self.iter())?; + writer.write(&("set", entries.len(), collision))?; + for (_, value) in entries { + writer.write_input(value)?; + } + Ok(()) + } + + async fn memo_state( + &self, + previous: Option<&MemoStateValue>, + ) -> Result> { + let (entries, collision) = sorted_set_entries(self.iter())?; + let children = entries + .iter() + .map(|(_, value)| *value as &dyn MemoInput) + .collect::>(); + memo_state_group(&children, previous, collision).await + } +} + +#[async_trait] +impl MemoInput for BTreeSet { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + let (entries, collision) = sorted_set_entries(self.iter())?; + writer.write(&("set", entries.len(), collision))?; + for (_, value) in entries { + writer.write_input(value)?; + } + Ok(()) + } + + async fn memo_state( + &self, + previous: Option<&MemoStateValue>, + ) -> Result> { + let (entries, collision) = sorted_set_entries(self.iter())?; + let children = entries + .iter() + .map(|(_, value)| *value as &dyn MemoInput) + .collect::>(); + memo_state_group(&children, previous, collision).await + } +} + +macro_rules! impl_tuple_memo_input { + ($len:expr; $(($index:tt, $name:ident)),+ $(,)?) => { + #[async_trait] + impl<$($name: MemoInput),+> MemoInput for ($($name,)+) { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + writer.write(&("tuple", $len as usize))?; + $(writer.write_input(&self.$index)?;)+ + Ok(()) + } + + async fn memo_state( + &self, + previous: Option<&MemoStateValue>, + ) -> Result> { + let children: &[&dyn MemoInput] = &[$(&self.$index),+]; + memo_state_group(children, previous, false).await + } + } + }; +} + +impl_tuple_memo_input!(1; (0, A)); +impl_tuple_memo_input!(2; (0, A), (1, B)); +impl_tuple_memo_input!(3; (0, A), (1, B), (2, C)); +impl_tuple_memo_input!(4; (0, A), (1, B), (2, C), (3, D)); +impl_tuple_memo_input!(5; (0, A), (1, B), (2, C), (3, D), (4, E)); +impl_tuple_memo_input!(6; (0, A), (1, B), (2, C), (3, D), (4, E), (5, F)); +impl_tuple_memo_input!(7; (0, A), (1, B), (2, C), (3, D), (4, E), (5, F), (6, G)); +impl_tuple_memo_input!(8; (0, A), (1, B), (2, C), (3, D), (4, E), (5, F), (6, G), (7, H)); + +/// Opaque writer passed to [`MemoInput::write_memo_key`]. +pub struct MemoKeyWriter<'a> { + fingerprinter: &'a mut Fingerprinter, +} + +impl MemoKeyWriter<'_> { + /// Write a serializable value as one framed part of the memo identity. + pub fn write(&mut self, value: &T) -> Result<()> { + write_key_fingerprint_part(self.fingerprinter, value) + } + + /// Recursively write another memo input into the same framed fingerprint + /// stream. + pub fn write_input(&mut self, value: &T) -> Result<()> { + value.write_memo_key(self) + } +} + #[derive(Clone)] pub struct MemoStateValue(Value); @@ -36,6 +474,110 @@ pub struct MemoStateDecision { memo_valid: bool, } +#[derive(Serialize, Deserialize)] +struct NestedMemoStatesV1 { + version: u8, + children: Vec>>, +} + +impl MemoStateValue { + /// Serialize a state value for storage with a memo entry. + pub fn from_serializable(value: &T) -> Result { + Ok(Self(Value::from_serializable(value)?)) + } + + /// Deserialize a state stored by a previous memo run. + pub fn deserialize Deserialize<'de>>(&self) -> Result { + self.0.deserialize() + } + + pub(crate) fn from_profile_value(value: Value) -> Self { + Self(value) + } +} + +impl MemoStateDecision { + /// Build a state-validation decision. + pub fn new(state: &T, memo_valid: bool) -> Result { + Ok(Self { + state: MemoStateValue::from_serializable(state)?, + memo_valid, + }) + } + + pub(crate) fn into_parts(self) -> (Value, bool) { + (self.state.0, self.memo_valid) + } +} + +/// Validate the memo states of recursively nested inputs and bundle them into +/// one opaque state value for the containing top-level argument or context +/// value. +#[doc(hidden)] +pub async fn memo_state_group( + children: &[&dyn MemoInput], + previous: Option<&MemoStateValue>, + force_miss: bool, +) -> Result> { + let previous_group = previous + .and_then(|value| value.deserialize::().ok()) + .filter(|group| group.version == 1); + let shape_matches = previous_group + .as_ref() + .is_some_and(|group| group.children.len() == children.len()); + + let previous_values = (0..children.len()) + .map(|index| { + previous_group + .as_ref() + .and_then(|group| group.children.get(index)) + .and_then(|state| state.as_ref()) + .map(|bytes| MemoStateValue(Value(bytes::Bytes::copy_from_slice(bytes.as_slice())))) + }) + .collect::>(); + + let decisions = futures::future::try_join_all( + children + .iter() + .zip(&previous_values) + .map(|(child, previous)| child.memo_state(previous.as_ref())), + ) + .await?; + + let mut memo_valid = previous.is_some() && shape_matches && !force_miss; + // Preserve a marker after malformed/legacy state so this slot forces one + // miss without shifting the following top-level argument states. + let mut has_state = force_miss || previous.is_some(); + let mut new_children = Vec::with_capacity(decisions.len()); + for (index, decision) in decisions.into_iter().enumerate() { + let previous_present = previous_values[index].is_some(); + match decision { + Some(decision) => { + has_state = true; + memo_valid &= previous_present && decision.memo_valid; + new_children.push(Some(decision.state.0.0.to_vec())); + } + None => { + memo_valid &= !previous_present; + new_children.push(None); + } + } + } + + if !has_state { + return Ok(None); + } + + MemoStateDecision::new( + &NestedMemoStatesV1 { + version: 1, + children: new_children, + }, + memo_valid, + ) + .map(Some) +} + #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] struct FileMemoState { modified_nanos: u128, @@ -97,15 +639,34 @@ where return Ok(result); }; - let guard = reserve_memoization(comp_ctx, fp) + let mut guard = reserve_memoization(comp_ctx, fp) .await .map_err(|e| Error::engine(format!("reserve_memoization: {e}")))?; - // Cache hit — deserialize and return the stored value. - if let Some(cached) = guard.cached() { - let value: T = cached.ret.deserialize()?; - ctx.check_cancellation()?; - return Ok(value); + let cached_context_states = guard + .cached() + .map(|cached| cached.context_memo_states.to_vec()); + if let Some(cached_context_states) = cached_context_states { + let context_validation = ctx + .state + .context + .validate_memo_states(&cached_context_states) + .await?; + if context_validation.memo_valid { + let cached = guard + .cached() + .expect("cached memo disappeared while reserved"); + let value: T = cached.ret.deserialize()?; + let positional_states = cached.memo_states.to_vec(); + ctx.check_cancellation()?; + if context_validation.states_changed { + guard.update_memo_states(MemoStatesPayload { + positional: positional_states, + by_context_fp: context_validation.states, + }); + } + return Ok(value); + } } // Cache miss (or memo disabled) — we are the resolver. Execute and commit. @@ -164,21 +725,30 @@ where .iter() .map(|decision| decision.state.0.clone()) .collect::>(); + let cached_context_states = guard + .cached() + .map(|cached| cached.context_memo_states.to_vec()) + .unwrap_or_default(); + let context_validation = ctx + .state + .context + .validate_memo_states(&cached_context_states) + .await?; if let Some(cached) = guard.cached() && state_decisions.iter().all(|decision| decision.memo_valid) + && context_validation.memo_valid { let states_changed = state_decisions .iter() .map(|decision| &decision.state.0.0) .ne(cached.memo_states.iter().map(|value| &value.0)); - let cached_context_states = cached.context_memo_states.to_vec(); let value: T = cached.ret.deserialize()?; ctx.check_cancellation()?; - if states_changed { + if states_changed || context_validation.states_changed { guard.update_memo_states(MemoStatesPayload { positional: memo_states_for_resolve, - by_context_fp: cached_context_states, + by_context_fp: context_validation.states, }); } return Ok(value); @@ -230,15 +800,11 @@ pub fn write_key_fingerprint_part( } #[doc(hidden)] -pub fn write_key_fingerprint_part_for_arg( +pub fn write_key_fingerprint_part_for_arg( fingerprinter: &mut Fingerprinter, value: &T, ) -> Result<()> { - if let Some(file) = as_file_like(value) { - let file_path = file.file_path(); - return write_key_fingerprint_part(fingerprinter, &file_path.memo_key()); - } - write_key_fingerprint_part(fingerprinter, value) + value.write_memo_key(&mut MemoKeyWriter { fingerprinter }) } #[doc(hidden)] @@ -246,51 +812,45 @@ pub fn finish_key_fingerprinter(fingerprinter: Fingerprinter) -> Fingerprint { fingerprinter.into_fingerprint() } +pub(crate) fn memo_input_fingerprint( + namespace: &str, + name: &str, + value: &T, +) -> Result { + let mut fingerprinter = new_key_fingerprinter(); + write_key_fingerprint_part(&mut fingerprinter, namespace)?; + write_key_fingerprint_part(&mut fingerprinter, name)?; + write_key_fingerprint_part_for_arg(&mut fingerprinter, value)?; + Ok(finish_key_fingerprinter(fingerprinter)) +} + +pub(crate) fn write_file_memo_key( + file: &dyn FileLike, + writer: &mut MemoKeyWriter<'_>, +) -> Result<()> { + writer.write(&file.file_path().memo_key()) +} + #[doc(hidden)] -pub async fn collect_memo_arg_state( +pub async fn collect_memo_arg_state( value: &T, prev: Option<&MemoStateValue>, ) -> Result> { - if let Some(file) = as_file_like(value) { - return file_memo_state(file, prev).await.map(Some); - } - Ok(None) + value.memo_state(prev).await } -fn as_file_like(value: &dyn Any) -> Option<&dyn FileLike> { - if let Some(file) = value.downcast_ref::() { - return Some(file); - } - #[cfg(feature = "amazon_s3")] - if let Some(file) = value.downcast_ref::() { - return Some(file); - } - #[cfg(feature = "google_drive")] - if let Some(file) = value.downcast_ref::() { - return Some(file); - } - #[cfg(feature = "oci_object_storage")] - if let Some(file) = value.downcast_ref::() { - return Some(file); - } - None -} - -async fn file_memo_state( +pub(crate) async fn file_memo_state( file: &dyn FileLike, prev: Option<&MemoStateValue>, ) -> Result { let metadata = file.metadata().await?; let modified_nanos = system_time_nanos(metadata.modified); - let prev_state = prev.and_then(|value| value.0.deserialize::().ok()); + let prev_state = prev.and_then(|value| value.deserialize::().ok()); if let Some(ref prev_state) = prev_state && prev_state.modified_nanos == modified_nanos { - return Ok(MemoStateDecision { - state: MemoStateValue(Value::from_serializable(&prev_state)?), - memo_valid: true, - }); + return MemoStateDecision::new(&prev_state, true); } let content_fingerprint = file.content_fingerprint().await?; @@ -301,14 +861,397 @@ async fn file_memo_state( let memo_valid = prev_state .as_ref() .is_some_and(|prev_state| prev_state.content_fingerprint == content_fingerprint); - Ok(MemoStateDecision { - state: MemoStateValue(Value::from_serializable(&state)?), - memo_valid, - }) + MemoStateDecision::new(&state, memo_valid) +} + +macro_rules! impl_file_memo_input { + ($ty:ty) => { + #[async_trait::async_trait] + impl $crate::memo::MemoInput for $ty { + fn write_memo_key( + &self, + writer: &mut $crate::memo::MemoKeyWriter<'_>, + ) -> $crate::Result<()> { + $crate::memo::write_file_memo_key(self, writer) + } + + async fn memo_state( + &self, + previous: Option<&$crate::memo::MemoStateValue>, + ) -> $crate::Result> { + $crate::memo::file_memo_state(self, previous) + .await + .map(Some) + } + } + }; } +pub(crate) use impl_file_memo_input; + fn system_time_nanos(time: SystemTime) -> u128 { time.duration_since(UNIX_EPOCH) .unwrap_or_default() .as_nanos() } + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Clone)] + struct Tracked { + id: &'static str, + version: u64, + } + + #[derive(Eq, Hash, PartialEq)] + struct CollidingKey(&'static str); + + #[async_trait] + impl MemoInput for Tracked { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + writer.write(&self.id) + } + + async fn memo_state( + &self, + previous: Option<&MemoStateValue>, + ) -> Result> { + let previous_version = previous.and_then(|value| value.deserialize::().ok()); + MemoStateDecision::new(&self.version, previous_version == Some(self.version)).map(Some) + } + } + + #[async_trait] + impl MemoInput for CollidingKey { + fn write_memo_key(&self, writer: &mut MemoKeyWriter<'_>) -> Result<()> { + writer.write(&"deliberate collision") + } + } + + #[derive(crate::MemoInput)] + struct TrackedBatch { + name: String, + items: Vec>, + } + + #[derive(crate::MemoInput)] + enum TrackedSource { + Empty, + Batch(TrackedBatch), + Named { item: Tracked }, + } + + #[derive(crate::MemoInput)] + enum ShadowedNames { + Named { writer: String, previous: u64 }, + } + + fn fingerprint(value: &impl MemoInput) -> Fingerprint { + memo_value_fingerprint(value).unwrap() + } + + fn previous_state(decision: &MemoStateDecision) -> MemoStateValue { + decision.state.clone() + } + + #[test] + fn recursive_sequence_fingerprints_are_framed_and_ordered() { + assert_ne!( + fingerprint(&vec!["ab".to_string(), "c".to_string()]), + fingerprint(&vec!["a".to_string(), "bc".to_string()]), + ); + assert_ne!( + fingerprint(&vec!["a".to_string(), "b".to_string()]), + fingerprint(&vec!["b".to_string(), "a".to_string()]), + ); + } + + #[test] + fn bytes_use_bulk_encoding_while_vec_u8_remains_a_sequence() { + let value = bytes::Bytes::from_static(b"bulk bytes"); + let mut expected = Fingerprinter::default(); + serde::Serializer::serialize_bytes(&mut expected, value.as_ref()).unwrap(); + + assert_eq!(fingerprint(&value), expected.into_fingerprint()); + assert_ne!(fingerprint(&value), fingerprint(&value.to_vec())); + } + + #[test] + fn all_advertised_integer_leaves_are_fingerprintable() { + assert_ne!(fingerprint(&i128::MIN), fingerprint(&i128::MAX)); + assert_ne!(fingerprint(&u128::MIN), fingerprint(&u128::MAX)); + assert_ne!(fingerprint(&i128::MAX), fingerprint(&(i128::MAX as u128))); + } + + #[test] + fn tuple_identity_is_tagged_ordered_and_distinct_from_vec() { + let tuple = ("a".to_string(), "b".to_string()); + assert_ne!( + fingerprint(&tuple), + fingerprint(&(tuple.1.clone(), tuple.0.clone())) + ); + assert_ne!(fingerprint(&tuple), fingerprint(&vec![tuple.0, tuple.1])); + } + + #[test] + fn option_pointer_and_set_framing_are_structural() { + assert_ne!( + fingerprint(&None::), + fingerprint(&Vec::::new()) + ); + assert_ne!( + fingerprint(&None::), + fingerprint(&Some(String::new())) + ); + + let value = "value".to_string(); + assert_eq!(fingerprint(&value), fingerprint(&Box::new(value.clone()))); + assert_eq!(fingerprint(&value), fingerprint(&Arc::new(value))); + + let first = HashSet::from(["a".to_string(), "b".to_string()]); + let mut second = HashSet::new(); + second.insert("b".to_string()); + second.insert("a".to_string()); + assert_eq!(fingerprint(&first), fingerprint(&second)); + let ordered = BTreeSet::from(["a".to_string(), "b".to_string()]); + assert_eq!(fingerprint(&first), fingerprint(&ordered)); + } + + #[test] + fn derive_tags_types_and_handles_reserved_named_enum_fields() { + #[derive(crate::MemoInput)] + struct First { + value: String, + } + + #[derive(crate::MemoInput)] + struct Second { + value: String, + } + + assert_ne!( + fingerprint(&First { + value: "same".to_string(), + }), + fingerprint(&Second { + value: "same".to_string(), + }), + ); + let _ = fingerprint(&ShadowedNames::Named { + writer: "value".to_string(), + previous: 1, + }); + } + + #[test] + fn hash_map_fingerprint_is_insertion_order_independent() { + let first = HashMap::from([("a".to_string(), 1_u64), ("b".to_string(), 2)]); + let mut second = HashMap::new(); + second.insert("b".to_string(), 2_u64); + second.insert("a".to_string(), 1_u64); + assert_eq!(fingerprint(&first), fingerprint(&second)); + } + + #[tokio::test] + async fn unordered_identity_collision_forces_recomputation_without_error() { + let values = HashMap::from([ + (CollidingKey("a"), "first".to_string()), + (CollidingKey("b"), "second".to_string()), + ]); + let first = values.memo_state(None).await.unwrap().unwrap(); + assert!(!first.memo_valid); + let second = values + .memo_state(Some(&previous_state(&first))) + .await + .unwrap() + .unwrap(); + assert!(!second.memo_valid); + } + + #[tokio::test] + async fn nested_sequence_preserves_child_state_validation() { + let initial = vec![ + Tracked { + id: "a", + version: 1, + }, + Tracked { + id: "b", + version: 1, + }, + ]; + let first = initial.memo_state(None).await.unwrap().unwrap(); + assert!(!first.memo_valid); + + let unchanged = initial + .memo_state(Some(&previous_state(&first))) + .await + .unwrap() + .unwrap(); + assert!(unchanged.memo_valid); + + let changed = vec![ + Tracked { + id: "a", + version: 1, + }, + Tracked { + id: "b", + version: 2, + }, + ] + .memo_state(Some(&previous_state(&unchanged))) + .await + .unwrap() + .unwrap(); + assert!(!changed.memo_valid); + } + + #[tokio::test] + async fn unordered_map_reordering_preserves_child_state_alignment() { + let mut initial = HashMap::new(); + initial.insert( + "a".to_string(), + Tracked { + id: "a", + version: 1, + }, + ); + initial.insert( + "b".to_string(), + Tracked { + id: "b", + version: 1, + }, + ); + let first = initial.memo_state(None).await.unwrap().unwrap(); + + let mut reordered = HashMap::new(); + reordered.insert( + "b".to_string(), + Tracked { + id: "b", + version: 1, + }, + ); + reordered.insert( + "a".to_string(), + Tracked { + id: "a", + version: 1, + }, + ); + let unchanged = reordered + .memo_state(Some(&previous_state(&first))) + .await + .unwrap() + .unwrap(); + assert!(unchanged.memo_valid); + + reordered.get_mut("b").unwrap().version = 2; + let changed = reordered + .memo_state(Some(&previous_state(&unchanged))) + .await + .unwrap() + .unwrap(); + assert!(!changed.memo_valid); + } + + #[tokio::test] + async fn derive_recurses_through_struct_enum_vec_and_arc() { + let source = TrackedSource::Batch(TrackedBatch { + name: "batch".to_string(), + items: vec![Arc::new(Tracked { + id: "item", + version: 3, + })], + }); + let first = source.memo_state(None).await.unwrap().unwrap(); + let unchanged = source + .memo_state(Some(&previous_state(&first))) + .await + .unwrap() + .unwrap(); + assert!(unchanged.memo_valid); + + let changed = TrackedSource::Batch(TrackedBatch { + name: "batch".to_string(), + items: vec![Arc::new(Tracked { + id: "item", + version: 4, + })], + }) + .memo_state(Some(&previous_state(&unchanged))) + .await + .unwrap() + .unwrap(); + assert!(!changed.memo_valid); + + assert_ne!(fingerprint(&TrackedSource::Empty), fingerprint(&source)); + assert_ne!( + fingerprint(&TrackedSource::Named { + item: Tracked { + id: "item", + version: 3, + }, + }), + fingerprint(&source), + ); + } + + #[tokio::test] + async fn malformed_nested_state_causes_a_miss_not_an_error() { + let values = vec![Tracked { + id: "item", + version: 1, + }]; + let malformed = MemoStateValue::from_serializable(&"not a nested state").unwrap(); + let decision = values.memo_state(Some(&malformed)).await.unwrap().unwrap(); + assert!(!decision.memo_valid); + } + + #[tokio::test] + async fn malformed_state_for_stateless_container_still_causes_one_miss() { + let values = vec!["item".to_string()]; + let malformed = MemoStateValue::from_serializable(&"not a nested state").unwrap(); + let replacement = values.memo_state(Some(&malformed)).await.unwrap().unwrap(); + assert!(!replacement.memo_valid); + + let unchanged = values + .memo_state(Some(&previous_state(&replacement))) + .await + .unwrap() + .unwrap(); + assert!(unchanged.memo_valid); + } + + #[tokio::test] + async fn stateless_nested_values_do_not_store_state() { + let values = vec![Some("a".to_string()), None]; + assert!(values.memo_state(None).await.unwrap().is_none()); + } + + #[cfg(all(feature = "serde_json", feature = "chrono", feature = "rust_decimal"))] + #[test] + fn ecosystem_schema_leaf_types_compose_in_derived_inputs() { + #[derive(crate::MemoInput)] + struct EcosystemLeaves { + timestamp: chrono::DateTime, + date: chrono::NaiveDate, + id: uuid::Uuid, + json: serde_json::Value, + decimal: rust_decimal::Decimal, + } + + let value = EcosystemLeaves { + timestamp: "2024-01-02T03:04:05Z".parse().unwrap(), + date: chrono::NaiveDate::from_ymd_opt(2024, 1, 2).unwrap(), + id: uuid::Uuid::nil(), + json: serde_json::json!({"nested": [1, 2, 3]}), + decimal: rust_decimal::Decimal::new(1234, 2), + }; + let _ = fingerprint(&value); + } +} diff --git a/rust/sdk/cocoindex/src/ops/image.rs b/rust/sdk/cocoindex/src/ops/image.rs index 93426c68e..ebcf48c1d 100644 --- a/rust/sdk/cocoindex/src/ops/image.rs +++ b/rust/sdk/cocoindex/src/ops/image.rs @@ -104,6 +104,13 @@ impl ImageEmbedder { } } +#[async_trait] +impl crate::memo::MemoInput for ImageEmbedder { + fn write_memo_key(&self, writer: &mut crate::memo::MemoKeyWriter<'_>) -> Result<()> { + writer.write(&(self.model_name(), self.dimension())) + } +} + #[async_trait] impl VectorSchemaProvider for ImageEmbedder { async fn vector_schema(&self) -> Result { diff --git a/rust/sdk/cocoindex/src/ops/sentence_transformers.rs b/rust/sdk/cocoindex/src/ops/sentence_transformers.rs index 9c980ed78..3667f0522 100644 --- a/rust/sdk/cocoindex/src/ops/sentence_transformers.rs +++ b/rust/sdk/cocoindex/src/ops/sentence_transformers.rs @@ -36,7 +36,7 @@ impl EmbeddingModel for TextEmbedding { } } -#[derive(Clone)] +#[derive(Clone, crate::MemoInput)] struct ScheduledEmbedder(SentenceTransformerEmbedder); impl serde::Serialize for ScheduledEmbedder { @@ -133,6 +133,13 @@ impl SentenceTransformerEmbedder { } } +#[async_trait] +impl crate::memo::MemoInput for SentenceTransformerEmbedder { + fn write_memo_key(&self, writer: &mut crate::memo::MemoKeyWriter<'_>) -> Result<()> { + writer.write(&(self.model_name(), self.dimension())) + } +} + #[async_trait] impl VectorSchemaProvider for SentenceTransformerEmbedder { async fn vector_schema(&self) -> Result { diff --git a/rust/sdk/cocoindex/src/prelude.rs b/rust/sdk/cocoindex/src/prelude.rs index e832b37c0..ca151f673 100644 --- a/rust/sdk/cocoindex/src/prelude.rs +++ b/rust/sdk/cocoindex/src/prelude.rs @@ -44,8 +44,9 @@ pub use crate::target_state::{ }; pub use crate::user_state::{IntoStateKey, StateHandle}; pub use crate::{ - App, ContextKey, DropHandle, Environment, EnvironmentBuilder, PreviewAction, PreviewValue, - Progress, SchemaFields, StatsGroupHandle, StatsGroupOptions, UpdateHandle, UpdateOptions, + App, ContextKey, DropHandle, Environment, EnvironmentBuilder, MemoInput, PreviewAction, + PreviewValue, Progress, SchemaFields, StatsGroupHandle, StatsGroupOptions, UpdateHandle, + UpdateOptions, }; pub use crate::{function, mount_each, use_mount}; diff --git a/rust/sdk/cocoindex/src/profile.rs b/rust/sdk/cocoindex/src/profile.rs index a208660f0..cb471cc96 100644 --- a/rust/sdk/cocoindex/src/profile.rs +++ b/rust/sdk/cocoindex/src/profile.rs @@ -7,7 +7,7 @@ use std::sync::Arc; use async_trait::async_trait; use cocoindex_core::engine::component::{ComponentProcessor, ComponentProcessorInfo}; -use cocoindex_core::engine::context::ComponentProcessorContext; +use cocoindex_core::engine::context::{ComponentProcessorContext, MemoStatesPayload}; use cocoindex_core::engine::profile::{EngineProfile, Persist}; use cocoindex_core::engine::target_state::{ ChildTargetDef, TargetActionSink, TargetHandler, TargetReconcileOutput, @@ -17,7 +17,7 @@ use cocoindex_utils::fingerprint::Fingerprint; use serde::{Deserialize, Serialize}; use crate::ctx::ContextStore; -use crate::error::Result; +use crate::error::{Error, Result}; // --------------------------------------------------------------------------- // RustProfile — the sealed EngineProfile implementation @@ -98,6 +98,7 @@ type ProcessFn = Box< pub(crate) struct BoxedProcessor { process_fn: std::sync::Mutex>, memo_fp: Option, + has_memo_state_handler: bool, info: Arc, } @@ -115,9 +116,15 @@ impl BoxedProcessor { Self { process_fn: std::sync::Mutex::new(Some(Box::new(process_fn))), memo_fp, + has_memo_state_handler: false, info: Arc::new(ComponentProcessorInfo::new(name)), } } + + pub(crate) fn with_memo_state_handler(mut self, enabled: bool) -> Self { + self.has_memo_state_handler = enabled; + self + } } impl ComponentProcessor for BoxedProcessor { @@ -144,6 +151,52 @@ impl ComponentProcessor for BoxedProcessor { self.memo_fp } + fn has_memo_state_handler(&self) -> bool { + self.has_memo_state_handler + } + + fn handle_memo_states( + &self, + _host_runtime_ctx: &(), + comp_ctx: &ComponentProcessorContext, + stored_states: Option>, + ) -> cocoindex_utils::error::Result< + impl Future< + Output = cocoindex_utils::error::Result<(MemoStatesPayload, bool, bool)>, + > + Send + + 'static, + > { + let context = Arc::clone(comp_ctx.host_ctx()); + let initial_context_states = stored_states + .is_none() + .then(|| comp_ctx.collect_context_initial_states()); + let stored_states = stored_states.unwrap_or_default(); + Ok(async move { + let Some(initial_context_states) = initial_context_states else { + let validation = context + .validate_memo_states(&stored_states.by_context_fp) + .await + .map_err(Error::into_core)?; + return Ok(( + MemoStatesPayload { + positional: stored_states.positional, + by_context_fp: validation.states, + }, + validation.memo_valid, + validation.states_changed, + )); + }; + Ok(( + MemoStatesPayload { + positional: stored_states.positional, + by_context_fp: initial_context_states, + }, + true, + false, + )) + }) + } + fn processor_info(&self) -> &ComponentProcessorInfo { &self.info } diff --git a/rust/sdk/cocoindex/src/resources/fs.rs b/rust/sdk/cocoindex/src/resources/fs.rs index c0613e0b8..829c22ad8 100644 --- a/rust/sdk/cocoindex/src/resources/fs.rs +++ b/rust/sdk/cocoindex/src/resources/fs.rs @@ -269,6 +269,8 @@ impl FileLike for FileEntry { impl FileSourceItem for FileEntry {} +crate::memo::impl_file_memo_input!(FileEntry); + /// Walk a directory matching multiple glob patterns. Returns all matching files /// sorted by relative path. /// diff --git a/rust/sdk/cocoindex/tests/amazon_s3_source.rs b/rust/sdk/cocoindex/tests/amazon_s3_source.rs index ab3872538..dde9a25d1 100644 --- a/rust/sdk/cocoindex/tests/amazon_s3_source.rs +++ b/rust/sdk/cocoindex/tests/amazon_s3_source.rs @@ -227,8 +227,7 @@ async fn s3_source_mount_each_pipeline_when_available() -> Result<()> { .await; cocoindex::context_key!( - static S3: S3Client = "s3_test_client", - state = S3Client::state_id + static S3: S3Client, key = "s3_test_client" ); #[cocoindex::function(memo)] diff --git a/rust/sdk/cocoindex/tests/doris_target.rs b/rust/sdk/cocoindex/tests/doris_target.rs index 6fd2be4e4..980c195ef 100644 --- a/rust/sdk/cocoindex/tests/doris_target.rs +++ b/rust/sdk/cocoindex/tests/doris_target.rs @@ -26,7 +26,7 @@ use cocoindex::{Environment, Result}; use serde::Serialize; use sqlx::Row; -cocoindex::context_key!(static DORIS_DB: DorisConnection = "doris_target_test_db"); +cocoindex::context_key!(static DORIS_DB: DorisConnection, key = "doris_target_test_db"); #[derive(Serialize, Clone)] struct Item { diff --git a/rust/sdk/cocoindex/tests/graph_vector_index.rs b/rust/sdk/cocoindex/tests/graph_vector_index.rs index a04d8438f..6f0af134b 100644 --- a/rust/sdk/cocoindex/tests/graph_vector_index.rs +++ b/rust/sdk/cocoindex/tests/graph_vector_index.rs @@ -27,7 +27,7 @@ async fn neo4j_vector_index_create_then_drop_when_available() { use cocoindex::connectors::neo4j::{self, ColumnDef, TableSchema, VectorMetric}; use cocoindex::{Environment, Result}; - cocoindex::context_key!(static G: neo4j::Graph = "neo4j_vidx_graph"); + cocoindex::context_key!(static G: neo4j::Graph, key = "neo4j_vidx_graph"); let uri = std::env::var("NEO4J_URI").unwrap_or_else(|_| "bolt://localhost:7687".to_string()); let user = std::env::var("NEO4J_USER").unwrap_or_else(|_| "neo4j".to_string()); @@ -111,7 +111,7 @@ async fn falkordb_vector_index_create_then_drop_when_available() { use cocoindex::{Environment, Result}; use std::time::{SystemTime, UNIX_EPOCH}; - cocoindex::context_key!(static G: falkordb::Graph = "falkordb_vidx_graph"); + cocoindex::context_key!(static G: falkordb::Graph, key = "falkordb_vidx_graph"); let uri = std::env::var("FALKORDB_URI").unwrap_or_else(|_| "falkor://localhost:6379".to_string()); diff --git a/rust/sdk/cocoindex/tests/lancedb_target.rs b/rust/sdk/cocoindex/tests/lancedb_target.rs index 859f2408b..c06a70732 100644 --- a/rust/sdk/cocoindex/tests/lancedb_target.rs +++ b/rust/sdk/cocoindex/tests/lancedb_target.rs @@ -13,8 +13,7 @@ use cocoindex::{Environment, ManagedTargetOptions, Result, SchemaFields}; use serde::Serialize; cocoindex::context_key!( - static DB: LanceDatabase = "lancedb_test", - state = LanceDatabase::state_id + static DB: LanceDatabase, key = "lancedb_test" ); const TABLE: &str = "docs"; diff --git a/rust/sdk/cocoindex/tests/pipeline.rs b/rust/sdk/cocoindex/tests/pipeline.rs index 3e0c3fde8..a11f5a61d 100644 --- a/rust/sdk/cocoindex/tests/pipeline.rs +++ b/rust/sdk/cocoindex/tests/pipeline.rs @@ -675,90 +675,126 @@ async fn no_detect_change_context_key_does_not_invalidate_memo() { assert_eq!(call_count.load(Ordering::SeqCst), 1); } -/// `ContextKey::new_with_state` drives memo invalidation from a derived state, -/// not the whole value. Here the resource is non-serializable and only its -/// `version` is tracked: -/// - changing a non-state field (`_tag`) must NOT invalidate the memo, -/// - changing the state field (`version`) MUST invalidate it. +/// A tracked context uses the value type's `MemoInput` implementation for both +/// stable identity and external-state validation. #[tokio::test] -async fn state_fn_context_key_invalidates_on_state_change_only() { +async fn type_owned_context_memo_state_invalidates_cached_function() { use std::sync::OnceLock; use std::sync::atomic::{AtomicUsize, Ordering}; - // Deliberately NOT `Serialize` — proves new_with_state works for resources - // that cannot be fingerprinted directly. struct Resource { + id: &'static str, version: u64, - _tag: &'static str, } - static KEY: OnceLock> = OnceLock::new(); - let key = KEY - .get_or_init(|| ContextKey::new_with_state("pipeline/state_fn", |r: &Resource| r.version)); + #[cocoindex::async_trait] + impl cocoindex::MemoInput for Resource { + fn write_memo_key( + &self, + writer: &mut cocoindex::memo::MemoKeyWriter<'_>, + ) -> cocoindex::Result<()> { + writer.write(&self.id) + } + + async fn memo_state( + &self, + previous: Option<&cocoindex::memo::MemoStateValue>, + ) -> cocoindex::Result> { + let previous_version = previous.and_then(|state| state.deserialize::().ok()); + Ok(Some(cocoindex::memo::MemoStateDecision::new( + &self.version, + previous_version == Some(self.version), + )?)) + } + } + static KEY: OnceLock> = OnceLock::new(); + let key = KEY.get_or_init(|| ContextKey::new_detect_change("pipeline/type_owned_state")); let dir = tempfile::tempdir().unwrap(); - let call_count = std::sync::Arc::new(AtomicUsize::new(0)); + let function_calls = std::sync::Arc::new(AtomicUsize::new(0)); + let component_calls = std::sync::Arc::new(AtomicUsize::new(0)); async fn run( path: std::path::PathBuf, key: ContextKey, version: u64, - tag: &'static str, - call_count: std::sync::Arc, - ) -> u64 { + function_calls: std::sync::Arc, + component_calls: std::sync::Arc, + ) { let app = Environment::builder() .db_path(path) - .provide_key(&key, Resource { version, _tag: tag }) + .provide_key( + &key, + Resource { + id: "resource", + version, + }, + ) .build() .await .unwrap() - .app("state_fn_context_key") + .app("type_owned_context_state") .await .unwrap(); - let count = call_count.clone(); app.update(move |ctx| async move { - let v: u64 = ctx - .memo(&"stable", move |ctx| { - let count = count.clone(); - async move { - count.fetch_add(1, Ordering::SeqCst); - Ok(ctx.get_key(&key)?.version) - } - }) - .await?; - Ok::<_, cocoindex::Error>(v) + let function_key = key.clone(); + ctx.memo(&"stable", move |ctx| { + let function_calls = function_calls.clone(); + async move { + function_calls.fetch_add(1, Ordering::SeqCst); + Ok(ctx.get_key(&function_key)?.version) + } + }) + .await?; + + let memo_fp = cocoindex_utils::fingerprint::Fingerprint::from( + &"type_owned_context_state_component", + ) + .unwrap(); + ctx.__use_mount_fp( + "stateful_child".to_string(), + Some(memo_fp), + move |child_ctx| async move { + component_calls.fetch_add(1, Ordering::SeqCst); + Ok(child_ctx.get_key(&key)?.version) + }, + ) + .await?; + Ok(()) }) .await - .unwrap() + .unwrap(); } - let path = dir.path().join("lmdb"); - - // Run 1: version=1, tag="a" -> miss (executes). - assert_eq!( - run(path.clone(), key.clone(), 1, "a", call_count.clone()).await, - 1 - ); - assert_eq!(call_count.load(Ordering::SeqCst), 1); - - // Run 2: version=1, tag="b" -> state (version) unchanged -> cache HIT. - assert_eq!( - run(path.clone(), key.clone(), 1, "b", call_count.clone()).await, - 1 - ); - assert_eq!( - call_count.load(Ordering::SeqCst), + run( + dir.path().join("lmdb"), + key.clone(), 1, - "changing a non-state field must not invalidate the memo" - ); + function_calls.clone(), + component_calls.clone(), + ) + .await; + run( + dir.path().join("lmdb"), + key.clone(), + 1, + function_calls.clone(), + component_calls.clone(), + ) + .await; + assert_eq!(function_calls.load(Ordering::SeqCst), 1); + assert_eq!(component_calls.load(Ordering::SeqCst), 1); - // Run 3: version=2 -> state changed -> cache MISS (re-executes). - assert_eq!(run(path, key.clone(), 2, "b", call_count.clone()).await, 2); - assert_eq!( - call_count.load(Ordering::SeqCst), + run( + dir.path().join("lmdb"), + key.clone(), 2, - "changing the tracked state must invalidate the memo" - ); + function_calls.clone(), + component_calls.clone(), + ) + .await; + assert_eq!(function_calls.load(Ordering::SeqCst), 2); + assert_eq!(component_calls.load(Ordering::SeqCst), 2); } /// Regression test: two memo bodies running concurrently, each reading a @@ -1822,7 +1858,7 @@ mod batched_test { static CONTEXT_ITEMS_PROCESSED: AtomicUsize = AtomicUsize::new(0); cocoindex::context_key!( - static BATCH_FACTOR: i64 = "pipeline/function_batching_context_factor", + static BATCH_FACTOR: i64, key = "pipeline/function_batching_context_factor", detect_change ); @@ -1992,6 +2028,128 @@ mod memo_test_borrowed_str { } } +mod memo_test_nested_inputs { + use super::*; + use std::path::PathBuf; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, OnceLock}; + + #[derive(Clone)] + struct TrackedResource { + id: &'static str, + version: u64, + } + + #[cocoindex::async_trait] + impl cocoindex::MemoInput for TrackedResource { + fn write_memo_key( + &self, + writer: &mut cocoindex::memo::MemoKeyWriter<'_>, + ) -> cocoindex::Result<()> { + writer.write(&self.id) + } + + async fn memo_state( + &self, + previous: Option<&cocoindex::memo::MemoStateValue>, + ) -> cocoindex::Result> { + let previous_version = previous.and_then(|state| state.deserialize::().ok()); + Ok(Some(cocoindex::memo::MemoStateDecision::new( + &self.version, + previous_version == Some(self.version), + )?)) + } + } + + #[derive(Clone, cocoindex::MemoInput)] + struct NestedInput { + resources: Vec>, + } + + static ARG_CALLS: AtomicUsize = AtomicUsize::new(0); + static CONTEXT_CALLS: AtomicUsize = AtomicUsize::new(0); + + #[cocoindex::function(memo)] + async fn read_argument(_ctx: &cocoindex::Ctx, input: &NestedInput) -> cocoindex::Result { + ARG_CALLS.fetch_add(1, Ordering::SeqCst); + Ok(input.resources.iter().map(|item| item.version).sum()) + } + + fn resources_key() -> &'static ContextKey>> { + static KEY: OnceLock>>> = OnceLock::new(); + KEY.get_or_init(|| ContextKey::new_detect_change("pipeline/nested_memo_input_context")) + } + + #[cocoindex::function(memo)] + async fn read_context(ctx: &cocoindex::Ctx) -> cocoindex::Result { + CONTEXT_CALLS.fetch_add(1, Ordering::SeqCst); + read_context_helper(&ctx).await + } + + #[cocoindex::function] + async fn read_context_helper(ctx: &cocoindex::Ctx) -> cocoindex::Result { + Ok(ctx + .get_key(resources_key())? + .iter() + .map(|item| item.version) + .sum()) + } + + fn nested(version: u64) -> NestedInput { + NestedInput { + resources: vec![Arc::new(TrackedResource { + id: "resource", + version, + })], + } + } + + #[tokio::test] + async fn nested_argument_state_controls_real_function_memo_reuse() { + ARG_CALLS.store(0, Ordering::SeqCst); + let (app, _dir) = temp_app("nested_argument_memo_state").await; + + for (version, expected_calls) in [(1, 1), (1, 1), (2, 2)] { + app.update(move |ctx| async move { + assert_eq!(read_argument(&ctx, &nested(version)).await?, version); + Ok(()) + }) + .await + .unwrap(); + assert_eq!(ARG_CALLS.load(Ordering::SeqCst), expected_calls); + } + } + + #[tokio::test] + async fn nested_context_state_controls_real_function_memo_reuse() { + CONTEXT_CALLS.store(0, Ordering::SeqCst); + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("lmdb"); + + async fn run(path: PathBuf, version: u64) -> u64 { + let app = Environment::builder() + .db_path(path) + .provide_key(resources_key(), nested(version).resources) + .build() + .await + .unwrap() + .app("nested_context_memo_state") + .await + .unwrap(); + app.update(|ctx| async move { read_context(&ctx).await }) + .await + .unwrap() + } + + assert_eq!(run(db_path.clone(), 1).await, 1); + assert_eq!(CONTEXT_CALLS.load(Ordering::SeqCst), 1); + assert_eq!(run(db_path.clone(), 1).await, 1); + assert_eq!(CONTEXT_CALLS.load(Ordering::SeqCst), 1); + assert_eq!(run(db_path, 2).await, 2); + assert_eq!(CONTEXT_CALLS.load(Ordering::SeqCst), 2); + } +} + mod memo_test_file_state { use super::*; use std::path::{Path, PathBuf}; @@ -2022,12 +2180,50 @@ mod memo_test_file_state { .unwrap() } + fn read_source_files(root: &Path) -> Vec { + let mut files = walk_dir(FilePath::with_base_dir( + "docs", + root.to_path_buf(), + PathBuf::new(), + )) + .recursive(true) + .walk() + .unwrap(); + files.sort_by_key(FileEntry::key); + files + } + #[cocoindex::function(memo)] async fn read_file(ctx: &Ctx, file: &FileEntry) -> Result { ctx.get_key(calls_key())?.fetch_add(1, Ordering::SeqCst); file.read_text().await } + #[cocoindex::function(memo)] + async fn read_files(ctx: &Ctx, files: &Vec) -> Result> { + ctx.get_key(calls_key())?.fetch_add(1, Ordering::SeqCst); + let mut result = Vec::with_capacity(files.len()); + for file in files { + result.push(format!("{}={}", file.key(), file.read_text().await?)); + } + Ok(result) + } + + async fn assert_file_list(app: &cocoindex::App, root: &Path, reverse: bool, expected: &[&str]) { + let root = root.to_path_buf(); + let actual = app + .update(move |ctx| async move { + let mut files = read_source_files(&root); + if reverse { + files.reverse(); + } + read_files(&ctx, &files).await + }) + .await + .unwrap(); + assert_eq!(actual, expected); + } + #[tokio::test] async fn file_memo_state_reuses_cache_when_mtime_changes_but_content_is_same() { let source = tempfile::tempdir().unwrap(); @@ -2095,6 +2291,60 @@ mod memo_test_file_state { "changed content must invalidate the memoized result" ); } + + #[tokio::test] + async fn file_container_tracks_unchanged_content_add_remove_and_reorder() { + let source = tempfile::tempdir().unwrap(); + let root = source.path(); + std::fs::write(root.join("a.txt"), "A").unwrap(); + std::fs::write(root.join("b.txt"), "B").unwrap(); + + let calls = Arc::new(AtomicUsize::new(0)); + let dir = tempfile::tempdir().unwrap(); + let app = Environment::builder() + .db_path(dir.path().join("lmdb")) + .provide_key(calls_key(), calls.clone()) + .build() + .await + .unwrap() + .app("file_container_memo_state") + .await + .unwrap(); + + assert_file_list(&app, root, false, &["a.txt=A", "b.txt=B"]).await; + assert_eq!(calls.load(Ordering::SeqCst), 1); + + assert_file_list(&app, root, false, &["a.txt=A", "b.txt=B"]).await; + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "unchanged input should hit" + ); + + tokio::time::sleep(Duration::from_millis(20)).await; + std::fs::write(root.join("b.txt"), "B2").unwrap(); + assert_file_list(&app, root, false, &["a.txt=A", "b.txt=B2"]).await; + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "changed content should miss" + ); + + std::fs::write(root.join("c.txt"), "C").unwrap(); + assert_file_list(&app, root, false, &["a.txt=A", "b.txt=B2", "c.txt=C"]).await; + assert_eq!(calls.load(Ordering::SeqCst), 3, "added file should miss"); + + assert_file_list(&app, root, true, &["c.txt=C", "b.txt=B2", "a.txt=A"]).await; + assert_eq!( + calls.load(Ordering::SeqCst), + 4, + "reordered files should miss" + ); + + std::fs::remove_file(root.join("c.txt")).unwrap(); + assert_file_list(&app, root, false, &["a.txt=A", "b.txt=B2"]).await; + assert_eq!(calls.load(Ordering::SeqCst), 5, "removed file should miss"); + } } // Two functions with same signature but different bodies produce different code hashes. @@ -2907,7 +3157,7 @@ mod mock_memo { /// Whole-function memoization should use the attribute. Its body receives /// the owned memo-scoped `Ctx`, so non-serializable resources remain /// available through normal context lookup. A skipped argument needs only - /// `Any + Clone`; `MockApi` intentionally does not implement `Serialize`. + /// `Clone`; `MockApi` intentionally does not implement `Serialize`. #[cocoindex::function(memo, memo_key(_client = skip))] async fn analyze( ctx: &cocoindex::Ctx, diff --git a/rust/sdk/cocoindex/tests/postgres_source.rs b/rust/sdk/cocoindex/tests/postgres_source.rs index 0eaac9671..736c7019e 100644 --- a/rust/sdk/cocoindex/tests/postgres_source.rs +++ b/rust/sdk/cocoindex/tests/postgres_source.rs @@ -13,12 +13,11 @@ use cocoindex::{Ctx, Environment, Result}; use serde::{Deserialize, Serialize}; cocoindex::context_key!( - static DB: postgres::Database = "postgres_source_test_db", - state = postgres::Database::state_id + static DB: postgres::Database, key = "postgres_source_test_db" ); static CALLS: AtomicUsize = AtomicUsize::new(0); -#[derive(Clone, Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize, cocoindex::MemoInput)] struct SourceRow { category: String, name: String, diff --git a/rust/sdk/cocoindex/tests/postgres_target.rs b/rust/sdk/cocoindex/tests/postgres_target.rs index a72d41666..76fdcb00d 100644 --- a/rust/sdk/cocoindex/tests/postgres_target.rs +++ b/rust/sdk/cocoindex/tests/postgres_target.rs @@ -10,8 +10,7 @@ use serde::{Deserialize, Serialize}; use sqlx::Row as _; cocoindex::context_key!( - static PG: postgres::Database = "postgres_target_test_db", - state = postgres::Database::state_id + static PG: postgres::Database, key = "postgres_target_test_db" ); #[derive(Clone, Serialize)] diff --git a/rust/sdk/cocoindex/tests/qdrant_target.rs b/rust/sdk/cocoindex/tests/qdrant_target.rs index 6b7a9a32c..3b43e3dac 100644 --- a/rust/sdk/cocoindex/tests/qdrant_target.rs +++ b/rust/sdk/cocoindex/tests/qdrant_target.rs @@ -16,8 +16,7 @@ use cocoindex::{Environment, Result}; use serde_json::json; cocoindex::context_key!( - static DB: QdrantConnection = "qdrant_test", - state = QdrantConnection::state_id + static DB: QdrantConnection, key = "qdrant_test" ); type Point = (u64, Vec, &'static str, &'static str); diff --git a/rust/sdk/cocoindex/tests/sqlite_target.rs b/rust/sdk/cocoindex/tests/sqlite_target.rs index c78054450..c62c27bd7 100644 --- a/rust/sdk/cocoindex/tests/sqlite_target.rs +++ b/rust/sdk/cocoindex/tests/sqlite_target.rs @@ -13,8 +13,7 @@ use serde_json::json; use sqlx::Row as _; cocoindex::context_key!( - static DB: sqlite::Database = "sqlite_target_test_db", - state = sqlite::Database::state_id + static DB: sqlite::Database, key = "sqlite_target_test_db" ); #[derive(Clone, Serialize)] diff --git a/rust/sdk/cocoindex/tests/surrealdb_target.rs b/rust/sdk/cocoindex/tests/surrealdb_target.rs index d5a67d8bf..afc860d66 100644 --- a/rust/sdk/cocoindex/tests/surrealdb_target.rs +++ b/rust/sdk/cocoindex/tests/surrealdb_target.rs @@ -3,7 +3,7 @@ use cocoindex::connectors::surrealdb::{self, ColumnDef, Graph, TableSchema, VectorIndexOptions}; use cocoindex::prelude::*; -cocoindex::context_key!(static GRAPH: Graph = "surrealdb_smoke_graph"); +cocoindex::context_key!(static GRAPH: Graph, key = "surrealdb_smoke_graph"); async fn try_graph(db_name: &str) -> Option { let url = std::env::var("SURREALDB_URL").unwrap_or_else(|_| "127.0.0.1:8787".to_string()); diff --git a/rust/sdk/cocoindex/tests/turbopuffer_target.rs b/rust/sdk/cocoindex/tests/turbopuffer_target.rs index d7c48ce64..a9e3f4630 100644 --- a/rust/sdk/cocoindex/tests/turbopuffer_target.rs +++ b/rust/sdk/cocoindex/tests/turbopuffer_target.rs @@ -17,8 +17,7 @@ use cocoindex::{Environment, Result}; use serde_json::json; cocoindex::context_key!( - static DB: TurbopufferConnection = "turbopuffer_test", - state = TurbopufferConnection::state_id + static DB: TurbopufferConnection, key = "turbopuffer_test" ); type RowSpec = (&'static str, Vec, &'static str); diff --git a/rust/sdk/cocoindex/tests/valkey_target.rs b/rust/sdk/cocoindex/tests/valkey_target.rs index 76f569fe2..1734dc3bb 100644 --- a/rust/sdk/cocoindex/tests/valkey_target.rs +++ b/rust/sdk/cocoindex/tests/valkey_target.rs @@ -12,7 +12,7 @@ use cocoindex::connectors::valkey::{ }; use cocoindex::resources::schema::VectorSchema; -cocoindex::context_key!(static VK: valkey::Valkey = "valkey_target_conn"); +cocoindex::context_key!(static VK: valkey::Valkey, key = "valkey_target_conn"); fn nonce() -> u128 { SystemTime::now() diff --git a/rust/sdk/cocoindex_macros/src/lib.rs b/rust/sdk/cocoindex_macros/src/lib.rs index 266998b88..b485a0da5 100644 --- a/rust/sdk/cocoindex_macros/src/lib.rs +++ b/rust/sdk/cocoindex_macros/src/lib.rs @@ -4,13 +4,253 @@ use proc_macro::TokenStream; use proc_macro2::{Span, TokenStream as TokenStream2}; use quote::{ToTokens, format_ident, quote}; use syn::{ - Error as SynError, Expr, FnArg, Ident, ItemFn, LitInt, LitStr, Pat, PatType, Path, Stmt, Token, - Type, TypeReference, parenthesized, + DeriveInput, Error as SynError, Expr, FnArg, Ident, ItemFn, LitInt, LitStr, Pat, PatType, Path, + Stmt, Token, Type, TypeReference, parenthesized, parse::{Parse, ParseStream}, parse_macro_input, punctuated::Punctuated, }; +/// Derive structural, recursively compositional memoization behavior. +/// +/// Every field must implement `cocoindex::MemoInput`; the containing struct or +/// enum does not need to implement `serde::Serialize`. The generated identity +/// includes the type's module path and name, so moving the type to another +/// module causes a one-time cache miss. Renaming the type or a named field also +/// changes its identity. These invalidations affect cached values only and do +/// not change persisted context-key names. +#[proc_macro_derive(MemoInput)] +pub fn derive_memo_input(input: TokenStream) -> TokenStream { + let input = parse_macro_input!(input as DeriveInput); + match derive_memo_input_impl(input) { + Ok(tokens) => tokens.into(), + Err(error) => error.into_compile_error().into(), + } +} + +fn derive_memo_input_impl(input: DeriveInput) -> syn::Result { + let name = input.ident; + let mut generics = input.generics; + for ty in memo_input_field_types(&input.data) { + generics + .make_where_clause() + .predicates + .push(syn::parse_quote!(#ty: ::cocoindex::memo::MemoInput)); + } + let (impl_generics, ty_generics, where_clause) = generics.split_for_impl(); + let (write_body, state_body) = memo_input_derive_bodies(&name, &input.data)?; + + Ok(quote! { + #[::cocoindex::async_trait] + impl #impl_generics ::cocoindex::memo::MemoInput for #name #ty_generics #where_clause { + fn write_memo_key( + &self, + writer: &mut ::cocoindex::memo::MemoKeyWriter<'_>, + ) -> ::cocoindex::Result<()> { + #write_body + } + + async fn memo_state( + &self, + previous: ::core::option::Option<&::cocoindex::memo::MemoStateValue>, + ) -> ::cocoindex::Result< + ::core::option::Option<::cocoindex::memo::MemoStateDecision> + > { + #state_body + } + } + }) +} + +fn memo_input_field_types(data: &syn::Data) -> Vec<&Type> { + match data { + syn::Data::Struct(data) => data.fields.iter().map(|field| &field.ty).collect(), + syn::Data::Enum(data) => data + .variants + .iter() + .flat_map(|variant| variant.fields.iter().map(|field| &field.ty)) + .collect(), + syn::Data::Union(_) => Vec::new(), + } +} + +fn memo_input_derive_bodies( + name: &Ident, + data: &syn::Data, +) -> syn::Result<(TokenStream2, TokenStream2)> { + match data { + syn::Data::Struct(data) => Ok(memo_input_struct_bodies(name, &data.fields)), + syn::Data::Enum(data) => Ok(memo_input_enum_bodies(name, data)), + syn::Data::Union(data) => Err(SynError::new_spanned( + data.union_token, + "MemoInput cannot be derived for unions", + )), + } +} + +fn memo_input_struct_bodies(name: &Ident, fields: &syn::Fields) -> (TokenStream2, TokenStream2) { + let len = fields.len(); + let (key_writes, children): (Vec<_>, Vec<_>) = match fields { + syn::Fields::Named(fields) => fields + .named + .iter() + .map(|field| { + let ident = field.ident.as_ref().expect("named field has an identifier"); + let field_name = ident.to_string(); + ( + quote! { + writer.write(&#field_name)?; + writer.write_input(&self.#ident)?; + }, + quote! { &self.#ident as &dyn ::cocoindex::memo::MemoInput }, + ) + }) + .unzip(), + syn::Fields::Unnamed(fields) => fields + .unnamed + .iter() + .enumerate() + .map(|(index, _)| { + let index = syn::Index::from(index); + ( + quote! { writer.write_input(&self.#index)?; }, + quote! { &self.#index as &dyn ::cocoindex::memo::MemoInput }, + ) + }) + .unzip(), + syn::Fields::Unit => (Vec::new(), Vec::new()), + }; + let write_body = quote! { + writer.write(&("struct", ::core::module_path!(), ::core::stringify!(#name), #len))?; + #(#key_writes)* + Ok(()) + }; + let state_body = quote! { + let __coco_children: &[&dyn ::cocoindex::memo::MemoInput] = &[#(#children),*]; + ::cocoindex::memo::memo_state_group(__coco_children, previous, false).await + }; + (write_body, state_body) +} + +fn memo_input_enum_bodies(name: &Ident, data: &syn::DataEnum) -> (TokenStream2, TokenStream2) { + let mut write_arms = Vec::with_capacity(data.variants.len()); + let mut state_arms = Vec::with_capacity(data.variants.len()); + + for variant in &data.variants { + let variant_ident = &variant.ident; + let variant_name = variant_ident.to_string(); + let len = variant.fields.len(); + match &variant.fields { + syn::Fields::Named(fields) => { + let fields_and_bindings = fields + .named + .iter() + .map(|field| { + let field = field.ident.as_ref().expect("named field has an identifier"); + let binding = format_ident!("__coco_field_{}", field); + (field, binding) + }) + .collect::>(); + let key_writes = fields_and_bindings.iter().map(|(field, binding)| { + let field_name = field.to_string(); + quote! { + writer.write(&#field_name)?; + writer.write_input(#binding)?; + } + }); + let children = fields_and_bindings + .iter() + .map(|(_, binding)| quote! { #binding as &dyn ::cocoindex::memo::MemoInput }); + let patterns = fields_and_bindings + .iter() + .map(|(field, binding)| quote! { #field: #binding }); + write_arms.push(quote! { + Self::#variant_ident { #(#patterns),* } => { + writer.write(&( + "enum", + ::core::module_path!(), + ::core::stringify!(#name), + #variant_name, + #len, + ))?; + #(#key_writes)* + Ok(()) + } + }); + let patterns = fields_and_bindings + .iter() + .map(|(field, binding)| quote! { #field: #binding }); + state_arms.push(quote! { + Self::#variant_ident { #(#patterns),* } => { + let __coco_children: &[&dyn ::cocoindex::memo::MemoInput] = + &[#(#children),*]; + ::cocoindex::memo::memo_state_group( + __coco_children, + previous, + false, + ).await + } + }); + } + syn::Fields::Unnamed(fields) => { + let bindings = (0..fields.unnamed.len()) + .map(|index| format_ident!("__coco_field_{index}")) + .collect::>(); + let key_writes = bindings + .iter() + .map(|binding| quote! { writer.write_input(#binding)?; }); + let children = bindings + .iter() + .map(|binding| quote! { #binding as &dyn ::cocoindex::memo::MemoInput }); + write_arms.push(quote! { + Self::#variant_ident(#(#bindings),*) => { + writer.write(&( + "enum", + ::core::module_path!(), + ::core::stringify!(#name), + #variant_name, + #len, + ))?; + #(#key_writes)* + Ok(()) + } + }); + state_arms.push(quote! { + Self::#variant_ident(#(#bindings),*) => { + let __coco_children: &[&dyn ::cocoindex::memo::MemoInput] = + &[#(#children),*]; + ::cocoindex::memo::memo_state_group( + __coco_children, + previous, + false, + ).await + } + }); + } + syn::Fields::Unit => { + write_arms.push(quote! { + Self::#variant_ident => { + writer.write(&( + "enum", + ::core::module_path!(), + ::core::stringify!(#name), + #variant_name, + 0usize, + ))?; + Ok(()) + } + }); + state_arms.push(quote! { Self::#variant_ident => Ok(None) }); + } + } + } + + ( + quote! { match self { #(#write_arms),* } }, + quote! { match self { #(#state_arms),* } }, + ) +} + /// Information about a non-ctx parameter. #[derive(Clone, Debug)] struct ParamInfo { @@ -360,8 +600,9 @@ fn gen_state_collect_for_param( states_ident: &Ident, prev_states_ident: &Ident, param: &ParamInfo, + overrides: &[MemoKeyOverride], ) -> TokenStream2 { - if param.is_str_ref { + if param.is_str_ref || memo_key_override(overrides, ¶m.ident).is_some() { return quote! {}; } let ident = ¶m.ident; @@ -416,9 +657,17 @@ fn gen_memo_wrapper_block( .collect::>(); let state_collects = params .iter() - .map(|param| gen_state_collect_for_param(&states_ident, &prev_states_ident, param)) + .map(|param| { + gen_state_collect_for_param(&states_ident, &prev_states_ident, param, overrides) + }) + .collect::>(); + let state_params = params + .iter() + .filter(|param| !param.is_str_ref && memo_key_override(overrides, ¶m.ident).is_none()) + .cloned() .collect::>(); - let clone_stmts = gen_clones(params); + let state_clone_stmts = gen_clones(&state_params); + let body_clone_stmts = gen_clones(params); let memo_key = gen_function_memo_key(fn_name, hash_const_name, &key_writes); Ok(quote! {{ @@ -429,7 +678,7 @@ fn gen_memo_wrapper_block( __coco_key, #propagate_children_fn_logic, { - #(#clone_stmts)* + #(#state_clone_stmts)* move |__coco_prev_states| async move { let mut __coco_states = Vec::new(); #(#state_collects)* @@ -437,7 +686,7 @@ fn gen_memo_wrapper_block( } }, { - #(#clone_stmts)* + #(#body_clone_stmts)* #body_closure }, ) @@ -773,9 +1022,10 @@ fn expand_batching_function( /// in the key automatically. /// /// Customize individual argument keys with `memo_key(...)`. A transform -/// function replaces the argument's default representation; `skip` (or `None`) -/// excludes it entirely. A skipped parameter needs only `Any + Clone`, not -/// `Serialize`: +/// function replaces the argument's default `cocoindex::MemoInput` identity +/// and suppresses its type-owned external-state validation. `skip` (or `None`) +/// excludes the argument entirely. A skipped parameter needs only `Clone`, not +/// `MemoInput` or `Serialize`: /// ```ignore /// #[cocoindex::function(memo, memo_key(client = skip))] /// async fn fetch(ctx: &Ctx, url: &str, client: &ApiClient) -> Result { @@ -796,6 +1046,10 @@ fn expand_batching_function( /// /// let embedding: Vec = embed(ctx, text).await?; /// ``` +/// Additional parameters must implement `Serialize` because their serialized +/// values group compatible calls into a physical batch. With `memo`, they must +/// also implement `MemoInput` because they form part of each item's memo key, +/// unless `memo_key(...)` transforms or skips their memo identity. /// /// Optional `version` parameter forces cache invalidation: /// ```ignore diff --git a/rust/utils/src/fingerprint.rs b/rust/utils/src/fingerprint.rs index 892b2995f..6763c6d37 100644 --- a/rust/utils/src/fingerprint.rs +++ b/rust/utils/src/fingerprint.rs @@ -214,6 +214,12 @@ impl Serializer for &mut Fingerprinter { Ok(()) } + fn serialize_i128(self, v: i128) -> std::result::Result<(), Self::Error> { + self.write_type_tag("i16"); + self.hasher.update(v.to_le_bytes()); + Ok(()) + } + fn serialize_u8(self, v: u8) -> std::result::Result<(), Self::Error> { self.write_type_tag("u1"); self.hasher.update(v.to_le_bytes()); @@ -238,6 +244,12 @@ impl Serializer for &mut Fingerprinter { Ok(()) } + fn serialize_u128(self, v: u128) -> std::result::Result<(), Self::Error> { + self.write_type_tag("u16"); + self.hasher.update(v.to_le_bytes()); + Ok(()) + } + fn serialize_f32(self, v: f32) -> std::result::Result<(), Self::Error> { self.write_type_tag("f4"); self.hasher.update(v.to_le_bytes());