Rust SDK ergonomics: close the remaining macro/API gap with the Python SDK

Rust SDK ergonomics: close the remaining macro/API gap with the Python SDK

Problem

The Rust SDK (rust/sdk/cocoindex) requires users to understand more low-level machinery than the Python SDK for the same pipeline. Comparing examples/rust/text_embedding with examples/text_embedding, a Rust user today hand-writes things Python handles with one decorator flag or one line: batching wiring, memo-key hash constants, ContextKey statics, and stringly-typed table schemas.

Most of the foundation already exists — #[cocoindex::function] (with memo, memo_key, version, logic_tracking), use_mount!, mount_each!, and #[derive(SchemaFields)] cover the core of Python's @coco.fn / mount ergonomics. The gap is a handful of specific holes, plus stale docs that make the SDK look more low-level than it is.

Gaps

1. Batching macro is documented but not implemented

rust/sdk/SHOWCASE.md specifies #[cocoindex::function(batching)] and #[cocoindex::function(memo, batching)], but the macro parser (rust/sdk/cocoindex_macros/src/lib.rs) only accepts memo, memo_key(...), version = N, and logic_tracking = "...".

Today's equivalent requires hand-wiring the generated hash constant into a static (rust/sdk/cocoindex/src/batched.rs):

#[cocoindex::function]
async fn embed_batch(texts: Vec<String>) -> Result<Vec<Vec<f32>>> { ... }

static EMBED: LazyLock<Batched<String, Vec<f32>>> =
    LazyLock::new(|| Batched::new(embed_batch, __COCO_FN_HASH_EMBED_BATCH));
// call site: EMBED.call(&ctx, text)

Python: @coco.fn(batching=True).

2. The macro already covers non-serializable resources — but tests/docs steer users to manual ctx.memo instead

#[cocoindex::function(memo)] already handles the "memoized function that uses a client/pool" case, because (a) the memo body closure receives an owned Ctx (cached_by_fingerprint_with_state, rust/sdk/cocoindex/src/memo.rs:126), so ctx.get_key(...) / ctx.get_or_err::<T>() work inside the body, and (b) memo_key(param = skip) params only need Any + Clone, not Serialize (collect_memo_arg_state<T: Any>, memo.rs:254).

But the doc comment in rust/sdk/cocoindex/tests/pipeline.rs (~2691) calls hand-rolled ctx.memo(&(__COCO_FN_HASH_ANALYZE, input), ...) "the realistic pattern" — steering users to the one API where they must hand-wire hash constants and where forgetting the constant silently serves stale results after a code edit (manual ctx.memo closures are not logic-tracked). Python users never see fingerprints at all.

3. ContextKey declaration ceremony

Python:

PG_DB = coco.ContextKey[asyncpg.Pool]("text_embedding_db")

Rust (examples/rust/text_embedding/src/main.rs):

static DB: LazyLock<ContextKey<postgres::Database>> = LazyLock::new(|| {
    ContextKey::new_with_state("text_embedding_db", |db: &postgres::Database| {
        db.state_id().to_string()
    })
});

Duplicate key names also panic process-wide (rust/sdk/cocoindex/src/ctx.rs), which the static ritual exists to avoid.

4. #[derive(SchemaFields)] exists but connectors and examples don't use it

TableSchema::from_row::<T>() is implemented only for Doris and SQLite. Postgres, LanceDB, Qdrant, and Turbopuffer still require hand-written column strings (ColumnDef::new("bigint"), ColumnDef::new("vector(384)")), and no examples/rust/ project uses the derive — even though it was built as the analogue of Python's TableSchema.from_class.

5. Docs describe an SDK that doesn't exist

• rust/sdk/SHOWCASE.md documents #[cocoindex::function(batching)], ctx.write_file, ctx.batch, and a sync App::open — none exist in that form (real: Batched, DirTarget, async App::open / App::open_blocking).
• The docs site has zero Rust SDK pages; the only accurate reference is rust/sdk/cocoindex/tests/pipeline.rs.

A customer evaluating the Rust SDK reads a stale pitch doc, hits compile errors, and falls back to reading tests — which is where the "you must understand the low-level internals" impression comes from.

Workstreams (tracking issues)

Split into two tracking issues, sized for modular review:

• #2278 — rust-sdk: close the ergonomics gap (module regroup → batching macro → context_key! → memo guidance fix → from_row for remaining connectors → ops::sentence_transformers adoption → docs truth pass)
• #2279 — rust-cli: standalone cocoindex CLI (#[app] / #[lifespan] / #[main] registration + build-and-exec stdio protocol; phase 2 later ships the binary in the Python wheels and replaces the click front-end so one CLI serves both SDKs)

API sketches (target UX vs Python vs today)

Batching

# Python
@coco.fn(memo=True, batching=True, max_batch_size=32)
async def embed(texts: list[str]) -> list[NDArray]:
    return await coco.use_context(EMBEDDER).embed_batch(texts)

vec = await embed(chunk.text)   # called with a single item

// Rust today
#[cocoindex::function]
async fn embed_batch(texts: Vec<String>) -> Result<Vec<Vec<f32>>> { ... }
static EMBED: LazyLock<Batched<String, Vec<f32>>> =
    LazyLock::new(|| Batched::new(embed_batch, __COCO_FN_HASH_EMBED_BATCH));
let vec = EMBED.call(&ctx, text).await?;

// Rust proposed — same semantics as Python: declared Vec -> Vec, called item -> item
#[cocoindex::function(memo, batching, max_batch_size = 32)]
async fn embed(ctx: &Ctx, texts: Vec<String>) -> Result<Vec<Vec<f32>>> {
    ctx.get_key(&EMBEDDER)?.embed_batch(texts).await
}
let vec: Vec<f32> = embed(&ctx, text).await?;

The macro generates the Batched static and rewrites the callable signature (Vec<String> -> Vec<Vec<f32>> body, String -> Vec<f32> call), exactly mirroring Python's list[T] -> list[U] declared / T -> U called contract. Batched stays public as the escape hatch (dynamic batch functions, custom dispatch).

Context keys

# Python
PG_DB = coco.ContextKey[asyncpg.Pool]("text_embedding_db")
EMBEDDER = coco.ContextKey[SentenceTransformerEmbedder]("embedder", detect_change=True)

// Rust today
static DB: LazyLock<ContextKey<postgres::Database>> = LazyLock::new(|| {
    ContextKey::new_with_state("text_embedding_db", |db: &postgres::Database| {
        db.state_id().to_string()
    })
});

// Rust proposed — one line per form, mirroring new / new_detect_change / new_with_state
cocoindex::context_key!(static CONFIG: AppConfig = "app_config");
cocoindex::context_key!(static EMBEDDER: SentenceTransformerEmbedder = "embedder", detect_change);
cocoindex::context_key!(static DB: postgres::Database = "text_embedding_db", state = Database::state_id);

Memoization with resources (no new API — documentation fix)

// What tests/pipeline.rs currently calls "the realistic pattern":
#[cocoindex::function]
async fn analyze(ctx: &Ctx, file: &FileEntry) -> Result<Info> {
    let client = ctx.get_or_err::<Client>()?.clone();
    let content = file.content_str()?;
    ctx.memo(&(__COCO_FN_HASH_ANALYZE, file.fingerprint()), move |_ctx| async move {
        client.call(&content).await
    }).await
}

// What already works and should be documented as the default:
#[cocoindex::function(memo)]
async fn analyze(ctx: &Ctx, file: &FileEntry) -> Result<Info> {
    let client = ctx.get_or_err::<Client>()?.clone();   // memo body receives Ctx
    client.call(&file.content_str()?).await
}

Table schema from row struct

# Python
@dataclass
class DocEmbedding:
    id: int
    filename: str
    embedding: Annotated[NDArray, EMBEDDER]   # dim inferred from provided embedder

schema = await postgres.TableSchema.from_class(DocEmbedding, primary_key=["id"])

// Rust today (postgres)
postgres::TableSchema::new(
    [("id", ColumnDef::new("bigint")), ("filename", ColumnDef::new("text")),
     ("embedding", ColumnDef::new(format!("vector({EMBED_DIM})")))],
    ["id"],
)

// Rust proposed — already works for sqlite/doris; wire into remaining connectors
#[derive(Clone, Serialize, SchemaFields)]
struct DocEmbedding {
    id: i64,
    filename: String,
    #[coco(vector = 384)]
    embedding: Vec<f32>,
}
let schema = postgres::TableSchema::from_row::<DocEmbedding>(["id"])?
    .with_vector_dim("embedding", embedder.dim());   // optional runtime override ≈ Annotated[NDArray, EMBEDDER]

Redundancy review of the combined surface

• (memo, batching) vs Batched: not redundant — the struct remains the escape hatch; the macro is the default spelling.
• context_key! vs type-keyed provide::<T>(): not redundant — type-keyed injection has no change detection, no state fn, and allows one value per type. Docs should state when to use which.
• memo! vs #[function(memo, memo_key(...))]: redundant for whole functions (see above); only block-level memoization justifies it, hence demoted to optional.
• Mount spellings (ctx.scope / ctx.mount_each vs use_mount! / mount_each!): not redundant — the methods take explicit keys and skip the component-memo fingerprint fast-path; the macros auto-derive subpaths and fingerprint args. Docs should present the macros as the default and the methods as the explicit-control variant.

Non-goals

• Removing explicit &Ctx threading — a stated design choice ("explicit &Ctx instead of hidden globals").
• Hiding Serialize + DeserializeOwned + Send + 'static bounds on memo/mount boundaries — inherent to Rust; Python pays the equivalent cost invisibly via pickle fallbacks.
• Macro-izing target-connector authoring — one TargetHandler trait + closure-built sinks is already roughly as compact as Python's protocol.