Batched Calls for Embeddings

​

Overview

Embedding models typically accept batches of inputs and return batches of vectors. When you need to embed a large dataset, sending one row at a time wastes round-trips and leaves GPU capacity idle. Calling a deployed function with F.catalog_call() inside a static resolver’s DataFrame pipeline lets Chalk parallelize the calls across rows automatically. Combined with the function’s own max_batching_size and max_buffer_duration, rows are grouped into batches before they ever reach your handler. There’s no chunk size or concurrency limit to configure by hand.

​

Define the embedding function

The remote function receives its input as a columnar batch. Declare a list[str] or pyarrow.Table parameter so it can accept multiple rows at once:

import chalkcompute

@chalkcompute.function(
    name="embed",
    image=chalkcompute.Image.base("python:3.12-slim").run_commands(
        "pip install sentence-transformers",
    ),
    cpu="2",
    memory="4Gi",
    gpu="nvidia-l4",
    max_batching_size=25,
    max_buffer_duration=2000,
    retries=3,
)
def embed(text: list[str]) -> list[list[float]]:
    from sentence_transformers import SentenceTransformer

    model = SentenceTransformer("all-MiniLM-L6-v2")
    return model.encode(text).tolist()

When max_batching_size or max_buffer_duration is set, the function queue automatically buffers incoming items and delivers them to the handler as a batch. The handler is invoked when either the buffer reaches max_batching_size items or max_buffer_duration milliseconds have elapsed, whichever comes first.

Deploy the function:

python embed.py

​

Call embed() from a DataFrame

Once embed is deployed, call it across every row of a DataFrame with F.catalog_call() inside a static resolver:

from chalk import online
from chalk import functions as F
from chalk.features import DataFrame, Vector, features, _

@features
class Document:
    id: int
    text: str
    embedding: Vector[384]

@online(static=True)
def compute_document_embeddings(
    df: DataFrame[Document.id, Document.text],
) -> DataFrame[Document.id, Document.embedding]:
    return (
        df.with_columns({
            "document.embedding": F.catalog_call("function.embed", _.text),
        })
        .select(str(Document.id), "document.embedding")
    )

Each row’s text column is passed to embed, and the result is written as the embedding feature alongside it. .with_columns() associates each computed embedding with its row by the DataFrame’s own row identity, not by call-completion order, so results line up correctly with their inputs even though embed’s internal batching and Chalk’s parallelization can complete rows out of order. Note the qualified name passed to catalog_call: it’s function.embed, not a bare function name; model.* and function.* are the two valid namespaces.

​

Ingesting embeddings from a file

To compute embeddings for documents that live in a file rather than an existing DataFrame, scan the file directly inside an @offline(static=True) resolver:

import chalkdf
from chalk import offline
from chalk import functions as F
from chalk.features import DataFrame, Vector, features, _

@features
class Document:
    id: int
    text: str
    embedding: Vector[384]

@offline(static=True)
def ingest_document_embeddings() -> DataFrame[
    Document.id,
    Document.text,
    Document.embedding,
]:
    return (
        chalkdf.DataFrame.scan(["documents.parquet"])
        .with_columns({
            "document.embedding": F.catalog_call("function.embed", _.text),
        })
        .select("document.id", "document.text", "document.embedding")
    )

This resolver reads documents.parquet, computes an embedding for every row via the same embed function, and ingests the result as Document feature data once applied and queried. There’s no separate step to manually chunk or batch the file’s rows.