Putting LLMs inside data pipelines without breaking them
Lessons from wiring LLMs into Spark and Databricks pipelines: where they belong, where they don't, and how to keep a probabilistic step inside a deterministic system.

Samith Deshai Siddo
Forward Deployed AI Engineer, Data Color AI
Data pipelines are built on a promise: the same input produces the same output, every time. LLMs break that promise by design. The same prompt can return a different answer tomorrow, and a model upgrade can shift every answer at once.
Even so, some of the most useful LLM work I have done has been inside data pipelines rather than in chat interfaces. At Data Color AI I built a data quality framework on Databricks where an LLM writes validation rules, Spark runs them, and MLflow tracks how both behave over time. Getting there meant learning where a probabilistic step belongs in a deterministic system, and how to fence it in.
Where LLMs belong
The test I use is simple. If a competent engineer could write the logic as code, write it as code. Use an LLM when the logic depends on reading and interpreting messy human input.
Every item on the right can be done by an LLM, and teams do try it, usually because it works in a notebook demo with twenty rows. At twenty million rows it becomes slow, expensive, and wrong in ways that are hard to find.
The pattern: generate once, execute deterministically
The most valuable pattern I have found is to keep the LLM out of the per-row path entirely. Instead of asking a model to judge each row, ask it once to produce an artifact such as a rule, a mapping or a SQL expression. A person reviews that artifact, and the engine runs it at scale.
For data quality this works very well. The model gets a profile of a column: its type, null rate, distinct count, a sample of values, the most common patterns and what the column means in the business glossary. It proposes expectations. A steward approves, edits or rejects each one. Approved rules become versioned code that Spark evaluates on every run, with no model call at run time.
column tax_id
rule matches ^\d{2}-\d{7}$
when country = 'US'
“98.6% of US rows match the EIN pattern; the rest look like typos.”
Data steward
rule v3 · prompt v12
4,912,007 rows
No LLM calls at run time
{
"rule_id": "dq.customer.tax_id.us_ein_format",
"column": "tax_id",
"check": "regex_match",
"params": { "pattern": "^\\d{2}-\\d{7}$" },
"filter": "country = 'US'",
"severity": "warn",
"rationale": "98.6% of US rows match the EIN format; non-matching values appear to be typos or SSNs.",
"provenance": {
"model": "anthropic.claude-sonnet-5",
"prompt_version": "dq-proposer@12",
"profile_run_id": "a41f0c2e",
"approved_by": "steward:mkim"
}
}The provenance block is not decoration. When a rule starts failing six months later, the first question is why it exists. The answer should be one lookup away: which model proposed it, from which prompt version, based on which profile, and who approved it.
This pattern turns an LLM problem into an ordinary software problem. Rules can be unit tested, diffed and rolled back, and the pipeline stays reproducible because the thing that runs is plain code.
When you do need per-row inference
Sometimes the value is in the row itself: classifying a free-text expense description, pulling fields out of an attached contract, or normalizing a supplier name that appears in forty spellings. Then the LLM has to run per row, and five things matter.
1. Batch, and cache by input hash
Real datasets repeat themselves. The same product description appears thousands of times. Hash the normalized input together with the prompt version and model ID, and check a cache table before calling the model. On reference data, that often removes most of the calls. Send what remains in batches to reduce per-request overhead, and use the provider's batch API when latency does not matter.
2. Structured output, validated
Never parse free text. Give the model a JSON schema and use the provider's structured output or tool calling mode, whether that is Bedrock, Vertex AI or Azure OpenAI. Then validate the result again on your side, because a valid shape can still contain a category that does not exist in your taxonomy. Rows that fail validation go to a quarantine table, not into Silver.
3. Confidence thresholds and a review queue
Every prediction carries a confidence, either from the model or from a secondary check such as agreement between two prompts. Above the threshold, the row moves on. Below it, the row goes to a review queue where a person decides, and their decisions become labeled data for evaluation.
Invoice #88213 – 40 cartons nitrile gloves
Medical supplies
Medical supplies
Svc agreement, annual HVAC maint.
Facilities
Facilities
Misc – see attached
Unknown
Unknown
AWS mkt. subscription Q3
Software
Software
Reimb. – J. Ortiz travel Dallas
Travel
Travel
import json, hashlib
import pandas as pd
from pydantic import BaseModel, ValidationError
PROMPT_VERSION = "expense-classifier@7"
MODEL_ID = "anthropic.claude-haiku-4-5" # pinned, never "latest"
CATEGORIES = {"Medical supplies", "Facilities", "Software", "Travel", "Unknown"}
class Label(BaseModel):
category: str
confidence: float
def key(text: str) -> str:
norm = " ".join(text.lower().split())
return hashlib.sha256(f"{MODEL_ID}|{PROMPT_VERSION}|{norm}".encode()).hexdigest()
def classify_batches(batches):
for pdf in batches:
out = []
for chunk in (pdf[i:i + 25] for i in range(0, len(pdf), 25)):
raw = llm_json(MODEL_ID, PROMPT_VERSION, chunk["description"].tolist()) # one call per 25 rows
for row, item in zip(chunk.itertuples(), raw):
try:
lbl = Label(**item)
ok = lbl.category in CATEGORIES
except ValidationError:
lbl, ok = Label(category="Unknown", confidence=0.0), False
out.append((row.id, row.cache_key, lbl.category, lbl.confidence, ok,
MODEL_ID, PROMPT_VERSION, json.dumps(item)))
yield pd.DataFrame(out, columns=["id", "cache_key", "category", "confidence", "valid",
"model_id", "prompt_version", "raw_output"])
todo = (df.withColumn("cache_key", key_udf("description"))
.join(cache, "cache_key", "left_anti")) # only uncached inputs hit the model
labeled = todo.repartition(32).mapInPandas(classify_batches, schema=OUT_SCHEMA)
labeled.filter("valid AND confidence >= 0.8").write.mode("append").saveAsTable("silver.expenses_labeled")
labeled.filter("NOT valid OR confidence < 0.8").write.mode("append").saveAsTable("ops.review_queue")A few choices in there are deliberate. The repartition controls concurrency, and with it how hard you hit rate limits. The model output is stored raw next to the parsed value, so a parsing bug can be fixed without paying for inference again. Every row records the model and prompt version that produced it.
4. Idempotency and reproducibility
A rerun of yesterday's job should produce yesterday's output. With LLMs that only holds if you pin the model version, version your prompts like code, set temperature to zero, and read from the cache before calling the model. When you do change the model or prompt, treat it as a migration: run the new version on a sample, compare, then backfill on purpose.
Fragile
latest, prompt inlined in a notebook, output overwritten in place, no record of which run produced which label.Reproducible
model_id and prompt_version columns.5. Evaluate continuously
A golden set of a few hundred labeled rows, drawn partly from the review queue, is the most useful asset in the project. Every prompt or model change runs against it before deployment. I log those runs to MLflow with the prompt version as a parameter and accuracy, per-class precision, validity rate and cost per thousand rows as metrics. In production, the same metrics are computed on each batch, along with the share of rows sent to review. A sudden jump in review volume is usually the first sign that the input data has changed.
Lineage
Data teams already care about lineage, and LLM steps have to join that graph. For every derived value, you should be able to answer three questions: which input rows produced it, which model and prompt version were used, and whether a person reviewed it. On Databricks, storing those as columns and registering the prompt as an MLflow artifact gets you most of the way. If a downstream dashboard looks wrong, you can trace a number back to a specific prompt revision.
What I would tell a team starting out
- Look for places where people currently read text and make a judgment. That is where LLMs pay off.
- Prefer generating rules and mappings over per-row inference. Execution stays cheap and testable.
- When you do run per row, cache, batch, validate, and route low-confidence output to people.
- Pin everything, version prompts, and store raw outputs.
- Build the golden set before you build the pipeline.
None of this is exotic. It is the same discipline data engineering has always applied to external dependencies, applied to a dependency that happens to be a language model.