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

[tomz-alt]

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:

• rust-sdk: close the ergonomics gap with the Python SDK #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)
• rust-cli: standalone cocoindex CLI for Rust projects #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.

[tomz-alt]

cc @georgeh0 , would you mind to take a look?

[badmonster0]

thanks a lot @tomz-alt ! @georgeh0 can help take a look!

[georgeh0]

Thanks @tomz-alt. Everything you proposed looks good! I looked at the existing code and examples, and noticed more potential gaps to add to the list:

1. Memoization - Memoization is not enabled for many examples / ops. e.g. text_embedding and sentence_transformer. - We should have clear way for users to customize the way to decide memoization key (like the Python SDK).

2. Module structure: currently a lot of non-core things (such as file) is at the top level mod of the crate.

3. Story for CLI support: Currently, for Python, we encourage users put environment initialization logic in a lifespan function. There's a default environment automatically initialized using it (instead of doing initialization in user-owned main function). Hence both user-owned main function and the CLI can share environments initialized in the same way. How should we support support CLI for Rust code?

Let me know if you have any thoughts. Thanks!

[ZhiHanZ]

Hi @georgeh0 , here is my thoughts:
1. Memoization — confirmed for the ops: Python's SentenceTransformerEmbedder gets batching + memoization from @coco.fn.as_async(batching=True, runner=coco.GPU) internally;
the Rust op is a plain method. It should adopt #[cocoindex::function(memo, batching)] once that lands (item 1).

2. Module structure — agreed. Regroup to mirror Python: connectors::{postgres, qdrant, …}, resources::{file, fs, id}, keep ops::*. Cheap now, expensive after release.

3. CLI — plan: one Rust cocoindex binary becomes the CLI for both SDKs.

Rust projects: pipeline logic is compiled into the user's binary, so the CLI works like cargo run — build the crate, run the user's binary, pass the subcommand through. The binary handles update/drop/--preview via SDK macros: #[cocoindex::app] and
#[cocoindex::lifespan] register apps/environments at link time (same linkme mechanism #[cocoindex::function] uses), #[cocoindex::main] dispatches CLI subcommands before falling through to the user's own main.
This also answers the environment-sharing question: link-time registration is Rust's analogue of Python's import-time registration, so the CLI and the user's main run the identical lifespan.

Python projects: same binary, distributed in the wheel ruff/uv-style (we already build with maturin). It spawns python -m cocoindex as a thin protocol child — load_user_app + the existing registries, emitting progress as JSON lines — and the ~1000-line click CLI front-end gets replaced by the Rust one, so UX parity across SDKs holds by construction.
ls/show need no user code in either language: they read the state db directly via core::inspect (both SDKs share the same core and LMDB format).

I'll follow up with a design doc under dev/plans/ covering the stdio protocol and migration. if you think that way make sense

[georgeh0]

Thanks @ZhiHanZ ! Looks good! It'll be very great as both Rust and Python can share the same CLI!

[georgeh0]

Noticed more problems on Rust SDK beyond ergonomics. Created #2326 for this.