
# What if a 9B model could beat Sonnet at SQL?

I was building a data agent. The kind of thing where you type a question in plain English and it writes SQL, runs it, and shows you a chart. The demo worked great. Single question, single query, clean result. Ship it.

Then I tried using it the way a real analyst would.

```text
Show revenue by country.
Only France.
Now break it down by month.
That looks empty. Try the country code instead.
```

Four turns. By the third one, the model had forgotten what we were looking at. By the fourth, it was writing SQL against a table it had already abandoned. The demo was a lie. Not because the model couldn't write SQL — it clearly could. But because it couldn't hold a conversation about data.

That failure sent me down a rabbit hole.

## The benchmark gap nobody talks about

When you hear that frontier models score 85% on text-to-SQL, that number almost certainly comes from a single-turn benchmark. One question, one answer, fresh context. [Spider](https://yale-lily.github.io/spider) and [BIRD](https://arxiv.org/abs/2305.03111) are the gold standards, and they've driven real progress.

But I kept pulling at the thread: what happens when we test models the way analysts actually work? Not one question at a time, but in sequences. Follow-ups. Refinements. "Wait, that doesn't look right — try it this way."

That's where [SParC](https://arxiv.org/abs/1906.02285) and [CoSQL](https://arxiv.org/abs/1909.05378) come in. SParC tests whether a model can handle sequences of related questions over databases it has never seen. CoSQL goes further — it's actual dialogue, with ambiguity, clarification requests, and all the messiness of a real conversation between an analyst and a SQL assistant.

And then there's [BIRD-Interact](https://arxiv.org/abs/2510.05318), which argues that even those benchmarks are too polite. Real assistants need to recover from execution errors, ask for clarification, and deal with the fact that their own earlier answers might have been wrong.

Here's the thing that got me: **frontier models that look dominant on single-turn benchmarks get dramatically less impressive on multi-turn ones.** The gap isn't small. It's the kind of gap that makes you wonder whether we've been measuring the right thing.

So I started a thought experiment. If multi-turn is where the big models stumble, and if the problem is more about *behavior* than raw intelligence — keeping context, carrying filters, resolving references — could a small model be fine-tuned to close that gap?

Could a 9B model running on my desk beat Sonnet at the thing that actually matters?

## Setting up the experiment

I grabbed [Qwen 3.5 9B](https://huggingface.co/Qwen) as my base. Why Qwen? It runs on consumer hardware, Unsloth has good support for it, and the instruction-tuned variant already knows what SQL looks like. The training ran on a single RTX 5090 with bf16 LoRA — rank 32, alpha 64, all the standard attention and MLP projections. Nothing exotic. About four hours of compute for each run.

For evaluation, I set up a CoSQL proxy: 100 turns spread across 32 real dialogs. The blog results I'm reporting use a clean subset of 43 follow-up turns from 12 of those dialogs, chosen because every run used the same rows, the same scorer, and the same protocol.

One crucial design choice: I used **generated-history rollout**, not teacher-forced history. In teacher-forced evaluation, the model gets perfect prior turns — it sees what the ideal previous SQL *should have been*. That's a useful diagnostic, but it's a fantasy. In practice, your model sees its own earlier output. If turn 2 was wrong, turn 3 has to deal with that. Generated-history makes the numbers uglier, but it's the only honest way to test a real assistant.

For the frontier baseline, I ran Claude Sonnet 4.6 through OpenRouter on the exact same rows, same scorer, same protocol. Same rules for both sides.

## Five ideas, one clear signal

I didn't just try one training strategy. I was exploring a research question, and that meant testing different hypotheses about *what* the model needs to learn.

[DIN-SQL](https://arxiv.org/abs/2304.11015) had shown that decomposing the text-to-SQL problem into subtasks could help. That felt right intuitively — writing SQL isn't one skill. It's figuring out which tables to use, understanding what values look like in the database, grasping what the user actually wants, generating the code, and catching mistakes. All tangled together. The question was which piece matters most for multi-turn conversations.

Here's what I tried:

**Direct SQL** — the control group. Take the user's question and the database schema, train the model to produce SQL directly. No intermediate steps, no clever representations. Just question in, query out. If this doesn't improve over the base model, nothing else matters.

**Semantic decomposition** — before writing any SQL, first have the model describe *what the user is asking for* in structured terms: which entity, which metric, which filter, which grouping. The idea is that understanding "break it down by month" depends on first understanding what "it" refers to from the previous turn. That's a meaning problem before it's a SQL problem. The training data included a Cube-inspired semantic model (entities, dimensions, measures, join hints) derived from the database schema, so the model could learn to map user language to governed business concepts before touching SQL.

**Metric DSL** — a small intermediate language for business metrics. Instead of going straight to `SUM(orders.amount)`, the model would first emit something like `MEASURE(revenue) BY customer_country`, and a compiler would expand that into the actual SQL. Attractive in theory, because when an analyst says "show me revenue," they're thinking about a business metric, not a column name. The model should preserve that intent and let something else handle the SQL plumbing.

**Behavior recovery** — in real use, the model will sometimes write bad SQL. The next turn has to deal with that. So instead of always training on perfect conversation histories, I generated training examples where earlier turns contained the model's *actual* (sometimes wrong) SQL output. The idea: teach the model to recognize when a previous query returned empty results or wrong data, and adjust. Show it the messy states it'll actually face in production.

**Schema selection** — this one was a deliberate cheat, designed to answer a diagnostic question. I took the correct reference SQL for each question, parsed it to extract exactly which tables, columns, and joins were needed, then fed those as hints in the prompt before asking the model to write SQL. Obviously you can't do this in production — you'd need to already know the answer to give the hint. But it tells you something critical: *if the model already knew where to look in the database, how good would its SQL be?*

The results told a clear story.

## What the numbers say

| Run | Value accuracy | Strict accuracy | What it means |
| --- | ---: | ---: | --- |
| Claude Sonnet 4.6 (via OpenRouter) | `0.674` | `0.395` | The frontier. Same questions, same scorer. |
| Qwen 9B + schema hints from correct answer | `0.651` | `0.465` | Cheating — told which tables and columns to use. Not deployable, but revealing. |
| Qwen 9B + semantic training (50 steps) | `0.581` | `0.326` | Best result from a fair, deployable local model. |
| Qwen 9B + direct SQL training (50 steps) | `0.581` | `0.302` | Same value accuracy, slightly less precise output shape. |
| Qwen 9B out of the box | `0.558` | `0.302` | No training, just the base model with a SQL prompt. |
| Qwen 9B + recovery training | `0.558` | `0.302` | Trained on messy histories. Didn't help yet. |

Before reading into these, a note on the metrics. *Value accuracy* asks: did the SQL return the right data? It's lenient about column aliases — if you called the column `total` instead of `revenue` but got the right numbers, that's a pass. *Strict accuracy* also cares about output shape. Both matter, but value accuracy is the more meaningful measure of whether the model is actually useful.

Some of these results are what you'd expect. Some are not.

**The clean local lift is real but small.** The best fine-tuned adapters (semantic and direct SQL) both reach `0.581` value accuracy, up from `0.558` for raw Qwen. That's a 2.3 percentage point improvement. Not nothing — it proves the model can learn from the task-specific data — but it doesn't close the gap to Sonnet's `0.674`.

**Behavior recovery was a disappointment.** Conceptually, it's the right idea. In practice, it tied raw Qwen. The training data didn't teach the useful repair moves yet. That's a data quality problem, not a conceptual one, but it means recovery training is still a hypothesis, not a result.

**The schema hints result is the real story.** When you cheat and tell the model which tables and columns it needs, the local model jumps to `0.651` — within two points of Sonnet on value accuracy and actually *ahead* on strict accuracy (`0.465` vs `0.395`). You can't deploy this — the hints come from parsing the correct answer, which you obviously don't have at inference time. But it's the kind of result that changes the direction of your research.

## What the schema hints result actually means

Here's what the cheat does concretely. The evaluation code takes the reference SQL (the correct answer), parses it with [sqlglot](https://github.com/tobymao/sqlglot), and extracts the relevant tables, the specific columns used, the join conditions, the projection shape, and grouping. It then injects those as hints into the prompt: "you need the `orders` table joined to `customers` on `customer_id`, selecting `amount` and `country`." The model still has to write the actual SQL, but it no longer has to *figure out where to look*.

Most wrong SQL is plausible SQL. The model doesn't usually produce garbage — it produces a reasonable-looking query against the wrong table, with the wrong join, using a value format that doesn't match what's in the database. "France" instead of "FR". `customer_orders` instead of `orders`. A `LEFT JOIN` where an `INNER JOIN` was needed.

When you remove that confusion — when the model knows where to look — a 9B local model gets within striking distance of the frontier. The remaining errors are about value grounding (the database stores "FR" but the user said "France"), execution edge cases, and grain issues (accidentally duplicating rows through a many-to-many join). The table-and-column selection problem is the bottleneck, not raw SQL writing ability.

This completely reframes the question. I started asking "can a small model write SQL as well as Sonnet?" and ended up learning that "can a small model pick the right tables?" is the question that actually matters.

## The honest evidence boundary

I want to be precise about what I can and can't claim from this work.

**I can say:** The best clean local fine-tunes slightly beat raw Qwen on the same 43 generated-history rows. Fine-tuning works. The model learns something real from multi-turn SQL data.

**I can say:** Sonnet 4.6 remains ahead on the clean comparison. The frontier hasn't been reached.

**I can say:** When the model is told which tables and columns to use (by peeking at the answer), it nearly closes the gap to Sonnet. That tells us where the real problem is.

**I cannot say:** A local 9B model beats Sonnet. Not yet. Not on this data.

**I cannot say:** The schema-hints setup is deployable. You can't peek at the answer in production. It's a ceiling test that shows where to push.

**I cannot say:** Recovery training is proven. It tied the baseline. The idea might be right, but the current training data didn't teach it.

This is a 43-row generated-history slice from CoSQL. It's reproducible, controlled, and honest. It is not a final benchmark claim.

## So what was the point?

The thought experiment started with "what if a 9B model could beat Sonnet at SQL?" and the honest answer is: not yet, not cleanly. But the experiment was worth running because of what it revealed about *where* the gap actually lives.

The gap isn't in SQL syntax. Raw Qwen already produces syntactically valid SQL almost every time — `1.000` syntax accuracy across the board. The gap isn't even in understanding English questions — the model generally knows what you're asking.

The gap is in navigating the database: knowing which tables to query, which columns to select, how to join them, and what the stored values actually look like. And in multi-turn conversations, that problem compounds. You're not just picking the right tables for one question. You're carrying forward context from previous turns, resolving references like "break *that* down by month," and adapting when a filter that worked on turn 2 needs a different column on turn 4.

That's a learnable problem. It's not a matter of needing a smarter model — it's a matter of needing better training data about which tables matter, how values are stored, and what carries forward across turns. The schema-hints experiment proved that the 9B model *can* do this work when it knows where to look. The next step is teaching it to figure that out on its own.

## What comes next

The next experiment should teach context selection directly:

1. Which tables matter for this question?
2. Which columns express the requested metric or dimension?
3. Which stored values match the user's wording?
4. What should carry forward from the previous turn?
5. What changed in the follow-up: filter, grouping, metric, or repair?
6. When execution returns nothing, what's the next reasonable attempt?

That's the bridge from a small CoSQL checkpoint toward [BIRD-Interact](https://arxiv.org/abs/2510.05318)-style evaluation. First make the local protocol stable. Then scale the benchmark. Then ask whether the local model can close the hosted gap.

The thought experiment isn't over. It just got more specific.

## Code

The full fine-tuning repo: [`github.com/xdanny/multiturn-sql-finetuning`](https://github.com/xdanny/multiturn-sql-finetuning).

## Sources

- [BIRD paper](https://arxiv.org/abs/2305.03111)
- [SParC paper](https://arxiv.org/abs/1906.02285)
- [CoSQL paper](https://arxiv.org/abs/1909.05378)
- [BIRD-Interact paper](https://arxiv.org/abs/2510.05318)
- [DIN-SQL paper](https://arxiv.org/abs/2304.11015)
- [Qwen model family](https://huggingface.co/Qwen)
- [OpenRouter](https://openrouter.ai/)


What if a 9B model could beat Sonnet at SQL?


# RAG Is Dead. The Engineering Isn't.

"RAG is dead." I kept seeing this take, and after building a retrieval system from scratch, I think they're half right. The term deserved to die. It always collapsed three distinct engineering problems — finding information, assembling context, and generating answers — into one buzzword that made it sound like a solved pattern. Stick an embedding model in front of a vector database, pipe the results into a prompt, done. RAG.

But the retrieval problem didn't go away just because we stopped calling it RAG. If anything, killing the term forced a more honest question: what are you *actually* building when you build this thing?

Jeff Huber, CEO of Chroma, put it sharply in a [Latent Space interview](https://www.latent.space/p/chroma):

> "Context engineering is the job of figuring out what should be in the context window at any given LLM generation step."

That framing clicked for me. Not because it's a better label — although it is — but because it decomposes the problem into pieces you can reason about independently. Retrieval. Filtering. Re-ranking. Assembly. Memory. Evaluation. Each one has its own failure modes, its own tuning knobs, its own way of going wrong.

So I built a code search system to see which of those pieces actually matters. The answer surprised me.

**Full tutorial**: [learn-strands/rag-chatbot](https://github.com/xdanny/learn-strands/tree/main/rag-chatbot)

---

## What I built

The system indexes the [Strands SDK](https://strandsagents.com) codebase — Python files, markdown docs, examples — and lets you ask questions about the framework in natural language. A chatbot that answers questions about Strands by searching its own source code.

The stack:
- **FAISS** for dense vector search (local embeddings, no API calls)
- **Strands SDK** for agent orchestration (tools, multi-agent coordination)
- **mem0** for conversation memory
- **Gemini 2.5 Flash Lite** for generation

Nothing exotic. I wanted to see how far you get with simple, well-composed primitives before reaching for anything fancy.

---

## Retrieval was the straightforward part

I expected retrieval to be where I'd spend most of my time. It wasn't. FAISS with `IndexFlatL2` and sentence-transformers gives you exact search with 100% recall, and for a codebase under 10K chunks, it's fast enough that you don't need approximate indices:

```python
class FAISSVectorStore:
    """FAISS-based vector store with local embeddings."""

    def __init__(self, local_model: str = "all-MiniLM-L6-v2", dimension: int = 384):
        self.embedder = SentenceTransformer(local_model)
        self.dimension = dimension
        self.index = faiss.IndexFlatL2(dimension)
        self.documents = []

    def add_documents(self, documents: List[Dict[str, Any]], batch_size: int = 32):
        """Add documents with batched embedding generation."""
        texts = [doc.get('text', '') for doc in documents]

        all_embeddings = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i + batch_size]
            embeddings = self.embedder.encode(batch, show_progress_bar=False)
            all_embeddings.append(embeddings)

        embeddings_array = np.vstack(all_embeddings).astype('float32')
        faiss.normalize_L2(embeddings_array)

        self.index.add(embeddings_array)
        self.documents.extend(documents)

    def search(self, query: str, k: int = 5, threshold: float = 0.3) -> List[Dict[str, Any]]:
        """Dense vector search with cosine similarity."""
        if self.index.ntotal == 0:
            return []

        query_embedding = self.embedder.encode([query]).astype('float32')
        faiss.normalize_L2(query_embedding)

        distances, indices = self.index.search(query_embedding, min(k, self.index.ntotal))
        similarities = 1 - (distances[0] / 2)  # L2 on normalized vectors → cosine

        results = []
        for idx, similarity in zip(indices[0], similarities):
            if similarity >= threshold and idx < len(self.documents):
                doc = self.documents[int(idx)].copy()
                doc['similarity'] = float(similarity)
                results.append(doc)

        return results
```

L2 normalization lets you get cosine similarity from Euclidean distance. The 0.3 threshold is aggressive but intentional — I'd rather get fewer, better chunks than flood the context with marginal matches.

The search itself becomes a Strands tool that the agent can call explicitly:

```python
@tool
def search_strands_sdk(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
    """Semantic search over Strands SDK codebase."""
    results = vector_store.search(query=query, k=max_results, threshold=0.3)

    return [{
        'text': r['text'],
        'similarity': r['similarity'],
        'file_path': r.get('file_path', 'unknown'),
        'file_type': r.get('file_type', 'unknown')
    } for r in results]


@tool
def search_by_file_type(query: str, file_type: str, max_results: int = 5) -> List[Dict[str, Any]]:
    """Combine dense search with metadata filtering."""
    all_results = vector_store.search(query=query, k=max_results * 3, threshold=0.3)
    filtered = [r for r in all_results if r.get('file_type', '') == file_type][:max_results]

    return [{
        'text': r['text'],
        'similarity': r['similarity'],
        'file_path': r.get('file_path', 'unknown')
    } for r in filtered]
```

These aren't buried in a pipeline. They're named tools the agent invokes explicitly — `search_strands_sdk` for broad semantic queries, `search_by_file_type` when you need to narrow by file extension. Huber calls this "naming the primitives," and he's right that it matters: when your retrieval is an explicit, composable step rather than a hidden subroutine, you can reason about it, debug it, and swap it out.

This whole layer took maybe a day to get working well. I kept waiting for the hard part. It wasn't here.

---

## The surprise: accurate retrieval saves money

Here's what I didn't expect. When first-stage retrieval is precise — when you go from 6,000 chunks to 20 good ones instead of 200 mediocre ones — the downstream agent burns dramatically fewer tokens.

Think about what happens when retrieval is sloppy. The agent gets a pile of vaguely relevant chunks, can't find a clear answer, and starts exploring. It might call the search tool again with a rephrased query. It might ask the LLM to reason through ambiguous context. It might generate a hedged, uncertain answer that prompts the user to ask a follow-up. Every one of those extra steps costs tokens.

With tight retrieval — good embeddings, reasonable thresholds, metadata filtering — the agent gets what it needs on the first call. It reads five chunks, finds the answer, cites the source, done. The retrieval precision translates directly into agent efficiency.

Huber frames this as "win the first stage":

> "Using signals like vector search, like full text search, like metadata filtering... to go from 10,000 down to 300."

He's talking about recall — making sure the good chunks are in the candidate set. But what I found is that the precision side matters just as much, maybe more, when an agent is consuming the results. An agent with 20 excellent chunks behaves completely differently from an agent with 200 okay chunks. The first one answers. The second one wanders.

---

## Where I fell short: re-ranking

Here's where I have to be honest about a gap.

Huber is clear about re-ranking:

> "Using an LLM as a re-ranker and brute forcing from 300 down to 30, I've seen now emerging... way more cost effective than a lot of people realize."

That's a *separate pass*. You take your 300 first-stage candidates, send them through a cross-encoder or an LLM scoring step, produce a ranked list, take the top 30, and *then* generate your answer from those 30. The re-ranking step reduces what enters the generation context.

My system doesn't do this. What I built was a response specialist that receives all retrieved chunks in a single prompt and generates from them:

```python
@tool
def response_specialist_tool(query: str, context: str) -> str:
    """Generate response from retrieved context."""
    agent = Agent(
        system_prompt="""You are a response generation specialist for Strands SDK queries.

Generate answers based ONLY on provided context.

Guidelines:
1. PRIORITIZE the most relevant chunks from context
2. CITE sources with file paths and line numbers
3. COMBINE information from multiple chunks coherently
4. If context is insufficient, say so clearly
5. Provide runnable code snippets when possible""",
        tools=[use_llm],
        model=gemini_model
    )

    response = agent(f"""User Question: {query}

Retrieved Context:
{context}

Generate a comprehensive answer using ONLY the context above.""")

    return str(response)
```

I initially told myself that the LLM's attention mechanism was doing "implicit re-ranking" — it focuses on relevant chunks and ignores the rest. But that's not what re-ranking is. Everything is still in the context window. Every token still counts against the context budget. The model might attend more to good chunks, but the bad ones are still there, still contributing to what Huber calls context rot:

> "As you use more and more tokens, the model can pay attention to less and then also can reason sort of less effectively."

So I'm eating the cost of stuffing noise into the context and hoping the model copes. For a small codebase with tight first-stage retrieval (where I'm sending 5-10 chunks, not 300), this works well enough. But it won't scale. A proper re-ranking step — score the candidates, take the top N, discard the rest — is on the list.

---

## The hard part: assembly and memory

Retrieval was a day of work. Assembly took three.

Assembly is the problem of taking retrieved chunks — disconnected fragments from different files, different sections of documentation, different code examples — and composing them into context that the model can reason about coherently. It's not just concatenation. The order matters. The framing matters. Whether you include file paths and line numbers matters. Whether you show the chunk verbatim or summarize it matters.

And then there's conversation memory. In a multi-turn chat, the model needs context from previous turns, but you can't just append the full conversation history. That grows without bound and directly triggers context rot.

This is where mem0 earned its place in the stack:

```python
mem0_config = {
    "vector_store": {
        "provider": "qdrant",
        "config": {
            "collection_name": "strands_chat",
            "embedding_model_dims": 384,
            "path": ":memory:"
        }
    },
    "embedder": {
        "provider": "huggingface",
        "config": {"model": "all-MiniLM-L6-v2"}
    }
}

memory = Memory.from_config(mem0_config)


@tool
def remember_conversation(user_message: str, assistant_response: str, user_id: str = "user") -> str:
    """Extract and store salient facts from the conversation."""
    memory.add(
        f"User asked: {user_message}\nAssistant responded: {assistant_response}",
        user_id=user_id
    )
    return f"Stored conversation in memory for {user_id}"


@tool
def recall_conversation(query: str = "", user_id: str = "user") -> str:
    """Retrieve relevant conversation history — not the full transcript."""
    if not query:
        query = "recent conversation history"

    results = memory.search(query, user_id=user_id, limit=5)

    if not results or 'results' not in results or not results['results']:
        return "No previous conversation history found."

    history = []
    for item in results['results']:
        if 'memory' in item:
            history.append(item['memory'])

    return "\n\n".join(history) if history else "No relevant conversation history found."
```

Instead of appending full dialogue turns (which bloats the context window with every exchange), mem0 extracts salient facts and stores them as searchable memories. When the agent needs conversation context, it retrieves *relevant* memories, not the entire transcript. Huber's principle — "tight, structured contexts beat maximal windows" — is exactly right, and memory is where most systems violate it first.

---

## Putting it together

The orchestrator ties retrieval, response generation, and memory into a single agent:

```python
rag_chatbot = Agent(
    system_prompt="""You are an intelligent assistant with expertise in the Strands SDK.

WORKFLOW:
1. RETRIEVE: Use retrieval_specialist to find relevant docs and code
2. RECALL: Use recall_conversation for relevant conversation context
3. RESPOND: Use response_specialist to generate a cited answer
4. REMEMBER: Use remember_conversation to store salient facts

Always retrieve context before answering technical questions.
Prefer code examples from the actual Strands SDK codebase.""",
    tools=[
        retrieval_specialist,
        response_specialist_tool,
        remember_conversation,
        recall_conversation,
        use_llm
    ],
    model=gemini_model
)
```

This is what context engineering looks like in practice: not a monolithic pipeline but composable steps — retrieval, memory, assembly, generation — each one an explicit tool the orchestrator invokes as needed. When you ask "How do I create an agent with custom tools?", the orchestrator:

1. Calls `retrieval_specialist` → searches the codebase, filters by relevance
2. Calls `recall_conversation` → pulls relevant memories from earlier turns
3. Passes both to `response_specialist` → generates an answer citing specific files
4. Calls `remember_conversation` → stores the key facts for future turns

Every step is visible. Every step is debuggable. And when something breaks — when the answer is wrong, when the citations are off — you can tell *which* step failed instead of staring at a monolithic "RAG pipeline" wondering where the problem lives.

**Full working tutorial** (self-contained Jupyter notebook, one-command setup): [learn-strands/rag-chatbot](https://github.com/xdanny/learn-strands/tree/main/rag-chatbot)

---

## What I'd measure next

Huber's most practical recommendation is one I haven't implemented yet: golden datasets.

> "People should be creating small golden data sets of what queries they want to work and what chunks should return... quantitatively evaluate what matters."

The idea is dead simple — spend an evening creating labeled query-chunk pairs:

```json
[
  {
    "query": "How do I create an agent with custom tools?",
    "expected_chunks": [
      "data/strands-sdk/examples/custom_tools.py",
      "data/strands-sdk/docs/agents.md"
    ],
    "expected_concepts": ["@tool decorator", "Agent class initialization"]
  }
]
```

Then wire Recall@10 into CI and fail the build if it drops below a threshold. Without this, retrieval quality degrades silently. You change an embedding model, update a chunking strategy, re-index the codebase, and never notice that three important queries stopped returning the right files. Golden datasets are the eval loop that keeps the system honest.

---

## What "RAG is dead" actually means

I started this project because everyone was saying RAG is dead, and I wanted to know what comes after it. The answer, it turns out, is the same engineering work — but decomposed into pieces that are honest about where the difficulty lives.

Retrieval was the easy part. FAISS, sentence-transformers, a threshold, some metadata filters. A day of work, and it was good enough. The real surprise was how much *accurate* retrieval matters downstream — not just for answer quality, but for token cost. A precise first stage means the agent stops wandering and starts answering.

Assembly was the hard part. Taking disconnected chunks from different files and composing them into context the model can reason about. Managing conversation memory without bloating the context window. Deciding what to include, what to summarize, and what to discard. This is where I spent most of my time and where the system still has the most room to improve.

And re-ranking — the step between retrieval and assembly — is where I cut a corner I shouldn't have. For a small codebase it's survivable, but it's the next thing to fix.

That's context engineering. Not a pipeline. Not a pattern you install from a template. A set of specific engineering decisions about what enters the context window and what doesn't, made independently at each stage, debugged independently when they break.

RAG is dead. Good riddance. The work was always more interesting than the acronym.

---

## Sources

- [Latent Space: Jeff Huber on Context Engineering](https://www.latent.space/p/chroma)
- [learn-strands Tutorial](https://github.com/xdanny/learn-strands/tree/main/rag-chatbot)
- [FAISS](https://github.com/facebookresearch/faiss)
- [Strands SDK](https://strandsagents.com)
- [mem0](https://mem0.ai)


RAG Is Dead. The Engineering Isn't.


# Building data quality checks in your pySpark data pipelines

Data quality is a rather critical part of any production data pipeline. In order to provide accurate SLA metrics
and to ensure that the data is correct, it is important to have a way to validate the data and report the metrics
for further analysis. In this post, we will look at how to build data quality checks in your pySpark data pipelines.

## Exploring Delta Live Tables

Delta Live Tables is a new feature in Databricks that allows users to build reliable data pipelines with built-in
data quality metrics and monitoring. It is a new abstraction on top of Delta Lake that allows users to query the
data using streaming live tables. The data is updated in real-time as the underlying data changes. What caught my eye
was the data quality capabilities that the users can specify on dataset level. Using python decorators we can specify
@expect_all, @expect_all_or_drop, and @expect_all_or_fail expectations that accept a python dictionary as an argument,
where the key is the expectation name and the value is the expectation constraint. Example:

```python
@dlt.expect("valid timestamp", "col("timestamp") > '2012-01-01'")
@dlt.expect_or_drop("valid_current_page", "current_page_id IS NOT NULL AND current_page_title IS NOT NULL")
@dlt.expect_or_fail("valid_count", "count > 0")
```

Metrics of clean records and failed records are automatically collected and stored in the Delta Live Table metadata, so
the users can set up alerts and monitor the data quality of their pipelines.

Delta Live Tables however still face quite some limitations and are not yet ready for production use. Some limitations
include:
1. The data quality checks are only available for streaming live tables, not for batch tables. We can still create
streaming tables from batch tables, but if the version of your data is changing the pipeline will fail.
2. Lack of testing capabilities. There is no way to test the data quality checks in a local environment because
dlt package is available only in Databricks runtime.
3. Lack of documentation. The documentation is very limited and it is not clear how to use the data quality checks.
Currently only python and SQL API are supported.
4. Setting up DLT job doesn't support all the parameters that are available in the Databricks job.

## Building your own data quality checks as python decorators

In order to overcome the limitations of Delta Live Tables, we can build our own data quality checks as python decorators.
The idea is to create a decorator that will accept a python list of arguments, which will be constraints that we will
apply for a determined column. We will collect all the necessary metrics and store them as part of the Delta Lake
metadata.

We will start by building two simple conditions for our data, uniqueness and filtering based on a condition.

```python
from abc import ABC, abstractmethod

class ColumnCondition(ABC):
    @abstractmethod
    def get_cols(self):
        pass


class UniqueCondition(ColumnCondition):

    def __init__(self, col):
        self.col = col

    def get_cols(self):
        return self.col


class FilterCondition(ColumnCondition):

    def __init__(self, left_col, right_col):
        self.left_col = left_col
        self.right_col = right_col

    def get_cols(self):
        return self.left_col, self.right_col
```

For ease of use, let's define a couple of factory functions that will return the right condition:

```python
def is_not_null(col):
    return FilterCondition(col + " is not null", col + " is null")

def is_unique(col):
    return UniqueCondition(col)
```

The main idea is to use a function as a decorator argument using a certain column, which will return a condition object.
We can use the condition object to pattern match and apply the specific function depending on the condition type.
We start by creating a simple python decorator using functools wraps:

```python
    def expect_or_drop(self, conditions: List[FilterCondition]):
        def decorator(function):
            @wraps(function)
            def wrapper(*args, **kwargs):
                retval = function(*args, **kwargs)
                # apply conditions
                return retval
            return wrapper
        return decorator
```

We will create an Expectations class that will contain all the data quality checks. The rsd arguments represents the
maximum relative standard deviation allowed for the approx_count_distinct_functions. Read more
[here](https://spark.apache.org/docs/3.1.2/api/python/reference/api/pyspark.sql.functions.approx_count_distinct.html).

```python

class Expectations:

    def __init__(self, spark: SparkSession, rsd=0.05):
        self.spark = spark
        self.schema = StructType([StructField("condition", StringType(), True),
                                  StructField("dropped_records", IntegerType(), True),
                                  StructField("clean_records", IntegerType(), True)])
        emptyRDD = spark.sparkContext.emptyRDD()
        self.metrics = spark.createDataFrame(emptyRDD, schema=self.schema)
        self.rsd = rsd
```

The metrics dataframe will contain the metrics for each data quality check. We can proceed to create our filtering
and uniqueness checks:

```python

    def apply_condition(self, dataframe, condition):
        if isinstance(condition, FilterCondition):
            return self.filter_condition(dataframe, condition.get_cols())
        elif isinstance(condition, UniqueCondition):
            return self.is_unique_extend(dataframe, condition.get_cols())
        return dataframe

    def filter_condition(self, dataframe: DataFrame, left_right) -> DataFrame:
        left, right = left_right
        total_records = dataframe.count()
        dropped_records = dataframe.filter(right).count()
        df = self.spark.createDataFrame([(left, dropped_records, (total_records - dropped_records))], schema=self.schema)
        self.metrics = self.metrics.unionAll(df)
        return dataframe.filter(left)

    def is_unique_extend(self, dataframe: DataFrame, col) -> DataFrame:
        total_records = dataframe.select(F.col(col)).count()
        distinct_records = dataframe.select(F.approx_count_distinct(col, self.rsd)).collect()[0][0]
        dropped_records = total_records - distinct_records
        df = self.spark.createDataFrame([(col + " is unique", dropped_records, distinct_records)], schema=self.schema)
        self.metrics = self.metrics.unionAll(df)
        return dataframe.dropDuplicates([col])
```

In order to apply the conditions we will use the apply_condition function to every condition in the list. In order to
do that we will use the functools reduce function as foldLeft:

```python
   foldl = lambda func, acc, xs: reduce(func, xs, acc)
   @wraps(function)
   def wrapper(*args, **kwargs):
        retval = function(*args, **kwargs)
        return foldl(self.apply_condition, retval, conditions)
```

Wrapping it all together:

```python

from abc import ABC, abstractmethod
from functools import reduce, wraps
from typing import List

import pyspark.sql.functions as F
from pyspark.sql import DataFrame, SparkSession
from pyspark.sql.types import IntegerType, StringType, StructField, StructType

foldl = lambda func, acc, xs: reduce(func, xs, acc)


class ColumnCondition(ABC):
    @abstractmethod
    def get_cols(self):
        pass


class UniqueCondition(ColumnCondition):

    def __init__(self, col):
        self.col = col

    def get_cols(self):
        return self.col


class FilterCondition(ColumnCondition):

    def __init__(self, left_col, right_col):
        self.left_col = left_col
        self.right_col = right_col

    def get_cols(self):
        return self.left_col, self.right_col


def is_not_null(col):
    return FilterCondition(col + " is not null", col + " is null")


def is_unique(col):
    return UniqueCondition(col)


class Expectations:

    def __init__(self, spark: SparkSession, rsd=0.05):
        self.spark = spark
        self.schema = StructType([StructField("condition", StringType(), True),
                                  StructField("dropped_records", IntegerType(), True),
                                  StructField("clean_records", IntegerType(), True)])
        emptyRDD = spark.sparkContext.emptyRDD()
        self.metrics = spark.createDataFrame(emptyRDD, schema=self.schema)
        self.rsd = rsd

    def expect_or_drop(self, conditions: List[FilterCondition]):
        def decorator(function):
            @wraps(function)
            def wrapper(*args, **kwargs):
                retval = function(*args, **kwargs)
                return foldl(self.apply_condition, retval, conditions)
            return wrapper
        return decorator

    def apply_condition(self, dataframe, condition):
        if isinstance(condition, FilterCondition):
            return self.filter_condition(dataframe, condition.get_cols())
        elif isinstance(condition, UniqueCondition):
            return self.is_unique_extend(dataframe, condition.get_cols())
        return dataframe

    def filter_condition(self, dataframe: DataFrame, left_right) -> DataFrame:
        left, right = left_right
        total_records = dataframe.count()
        dropped_records = dataframe.filter(right).count()
        df = self.spark.createDataFrame([(left, dropped_records, (total_records - dropped_records))], schema=self.schema)
        self.metrics = self.metrics.unionAll(df)
        return dataframe.filter(left)

    def is_unique_extend(self, dataframe: DataFrame, col) -> DataFrame:
        total_records = dataframe.select(F.col(col)).count()
        distinct_records = dataframe.select(F.approx_count_distinct(col, self.rsd)).collect()[0][0]
        dropped_records = total_records - distinct_records
        df = self.spark.createDataFrame([(col + " is unique", dropped_records, distinct_records)], schema=self.schema)
        self.metrics = self.metrics.unionAll(df)
        return dataframe.dropDuplicates([col])
```

Let's see the decorator in action:

```python
    @expectation.expect_or_drop([is_not_null("row"), is_unique("row")])
    def read_dataframe(df):
        return df

    result = read_dataframe(df_1)
    print(result.collect())
    print(expectation.metrics.collect())
```

The console will print:
```bash
[Row(row='row1', row_number=1)]
[Row(condition='row is not null', dropped_records=1, clean_records=2), Row(condition='row is unique',
dropped_records=1, clean_records=1)]
```

We can extend and add some plotting functions as well to our Expectations class.

```python
    def plot_pie_with_total(self, figsize=(10, 10)):
        labels = ["clean_records", "dropped_records"]
        df = self.metrics.toPandas()
        size = len(df.index)
        if size == 1:
            fig, axs = plt.subplots(1)
            axs.pie(df.iloc[0][labels], labels=labels, autopct='%1.1f%%')
            axs.set_title(df.iloc[0]["condition"])
        else:
            fig, axs = plt.subplots(1, size, figsize=figsize)
            for i in range(size):
                axs[i].pie(df.iloc[i][labels], labels=labels, autopct='%1.1f%%')
                axs[i].set_title(df.iloc[i]["condition"])
        plt.show()
```

That will generate a pie chart with the metrics for every condition:
![Image Name](/assets/blog/data-quality-pyspark/data_quality.png)

Happy coding and stay safe!


Building data quality checks in your pySpark data pipelines


# Integration Testing for your Databricks CI/CD Data Pipelines with Microsoft Nutter

In this blogpost we will continue our journey of testing our Data Pipelines. If you haven't checked out
the first post, [make sure you do](/post/automate-all-your-pyspark-unit-test-with-hypothesis). As assessed in the first post, writing unit tests in pyspark can be challenging.
Especially when we go to the world of notebooks, setting up your testing infrastructure can be challenging.
Most often the data scientists/data engineers in your organization will prefer working on and developing their
Data Pipelines in a dedicated environment like Databricks or JupyterLab. In order to enable cross-team
collaboration, peer-review and enforce best testing practices, the notebooks will have to be version controlled.
In Databricks, for example, the preferred way to work on notebooks is to use
[Databricks Repos](https://docs.databricks.com/repos/index.html). However, because Databricks doesn't have
support for pytest, you cannot write a separate notebook for tests. The data engineer will be writing
pytest's using an IDE as part of the notebook pull request, the execution part being left to the CI/CD pipeline
in a mocked Spark environment. This ensures that our code is determinstic and our functions are correct, but
it does not necessarily translate to correct functioning in a production-like environment.

## Microsoft Nutter - Testing Framework for Databricks

To overcome the limits described above, a popular framework open-sourced by Microsoft is
[Nutter](https://github.com/microsoft/nutter). The library has two main components: the runner and the CLI.
We will be exploring in this blogpost only the runner, leaving the CLI as an exercise for the readers. The
runner runtime is a module that can be used once you install the library in your notebook. In Databricks
we can use the following command, that will install nutter as well as hypothesis and tinsel libraries used to
generate our datasets:

```python
%pip install -U nutter hypothesis tinsel
```

As an example, we will be using the same function as in the first series blogpost to illustrate the example
in a separate notebook:

```python
from pyspark.sql import DataFrame
import pyspark.sql.functions as F

def test_dataframe_transformations(df: DataFrame) -> DataFrame:
    temp_df = df.select('*', F.explode_outer('order_items'))\
        .select(F.col('created_at'), F.col('col.item_value').alias('item_value'))\
        .withColumn('day', F.to_date('created_at'))\
        .withColumn('week', F.weekofyear('created_at'))

    weekly_revenue_df = temp_df.groupBy('week').agg(
        F.sum('item_value').alias('weekly_revenue')
    )

    return temp_df.groupBy('week', 'day').agg(
        F.sum('item_value').alias('daily_revenue')
    ).join(weekly_revenue_df, on='week', how='inner')
```

Now, in our testing notebook, we can import the needed functions with the following Databricks command:

```python
%run /{notebook_to_test_path}
```

In order to setup our test datasets, we will be using the same setup as in the previous blogpost:

```python
import hypothesis.strategies as st
from hypothesis import given, settings, HealthCheck

settings.register_profile(
"my_profile",
max_examples=50,
deadline=60 * 1000,  # Allow 1 min per example (deadline is specified in milliseconds)
suppress_health_check=(HealthCheck.too_slow, HealthCheck.data_too_large),
)
```

```python
from datetime import datetime, timedelta
from typing import List, NamedTuple
from tinsel import struct, transform

@struct
class OrderItem(NamedTuple):
item_id: str
item_value: float

@struct
class Order(NamedTuple):
order_id: str
created_at: datetime
updated_at: datetime
customer_id: str
order_items: List[OrderItem]

week_sunday = (datetime.now() - timedelta(days = datetime.now().weekday() + 1))
time_window = (datetime.now() - timedelta(days = datetime.now().weekday() + 7))

def order_strategy():
return st.lists(st.builds(Order,
order_id=st.text(min_size=5),
created_at=st.datetimes(min_value=time_window, max_value=week_sunday),
updated_at=st.datetimes(min_value=time_window, max_value=week_sunday),
customer_id=st.text(min_size=5),
order_items=st.lists(st.builds(OrderItem,
item_id=st.text(min_size=5),
item_value=st.floats(min_value=0, max_value=100)), min_size=1)), min_size=5)
```

We are now ready to create our first tests with Microsoft Nutter. The test framework follows the following
conventions for the test fixtures:

```
before_(testname) - (optional) - if provided, is run prior to the 'run_' method. This method can be used to setup any test pre-conditions

run_(testname) - (optional) - if provider, is run after 'before_' if before was provided, otherwise run first. This method is typically used to run the notebook under test

assertion_(testname) (required) - run after 'run_', if run was provided. This method typically contains the test assertions

after_(testname) (optional) - if provided, run after 'assertion_'. This method typically is used to clean up any test data used by the test
```

In the case where we want to test full notebooks, we can setup for example our static dataframe from a file
in the before function, run the ETL transformations, assert that the created delta table corresponds to our
expectations and use the after in order to clean up all the test resources we might have spinned up. In this
blogpost, however, we will only be using the assertion as we are generating the data on the fly. Our test
fixture then would look like this:

```python
from runtime.nutterfixture import NutterFixture, tag

default_timeout = 600

class Test1Fixture(NutterFixture):
def init(self):
NutterFixture.init(self)

@given(data=order_strategy())
@settings(settings.load_profile("my_profile"))
def test_samples(self, data):
  df = spark.createDataFrame(data=data, schema=transform(Order), verifySchema=False)
  collected_list = test_dataframe_transformations(df).collect()
  for row in collected_list:
    assert row['weekly_revenue'] >= row['daily_revenue']

def assertion_test_transform_data(self):
    self.test_samples()

def after_test_transform_data(self):
    print('done')
```

We can now run our test notebook, and the results will look like this:

```
Notebook exited:
Notebook: N/A - Lifecycle State: N/A, Result: N/A
Run Page URL: N/A

PASSING TESTS

test_transform_data (31.414987 seconds)
```

This notebook can be added as part of your CI/CD pipeline, and we will explore in the next blogpost how.
We can write and run unit/integration/end2end tests using Nutter and export the results via CSV or JUnit.
Happy Coding!

### References

[Microsoft Nutter](https://github.com/microsoft/nutter)


Integration Testing for your Databricks CI/CD Data Pipelines with Microsoft Nutter


# Automate all your PySpark Unit Test with Hypothesis!

Unit testing is often regarded as a main pillar of testing your software applications, and it usually involves
testing a single/unit component and ensuring that it covers all the edge cases the software developer can think of.
The Functional Programming world, specifically Haskell with QuickCheck and Scala with ScalaCheck, have introduced the
idea of testing based on property specifications and automatic test data generation in order to complement
the traditional unit testing approach. The idea is that we define a property that specifies the program behaviour and
the data is automatically generated for us to overcome all the test scenarios, using shrinking in order to simplify the
verification of the test case.

## PySpark property-based testing

Testing your PySpark notebooks can be quite challenging by itself, the data infrastructure itself has often quite
limited support for pytest and adding such tests to your existing data pipelines adds quite some complexity.
Additionally, testing your ETL transformation function can be quite cumbersome as it involves generating data,
maintaining it, or simply skipping the unit tests in favor of simple integrity tests. In this blog post we will see
how we can easily generate data from dataclasses and think of unit testing in terms of properties rather than data
integrity. We will start with Hypothesis, an excellent Python library built for property-based testing.

We can begin by specifying some simple dataclasses we will be using in order to showcase the potential of the library:

```python
from datetime import datetime, timedelta
from typing import List, NamedTuple
from tinsel import struct, transform

@struct
class OrderItem(NamedTuple):
  item_id: str
  item_value: float

@struct
class Order(NamedTuple):
  order_id: str
  created_at: datetime
  updated_at: datetime
  customer_id: str
  order_items: List[OrderItem]
```

Our transaction dataset will be composed by orders with a set of order items for each order. The dataclasses are defined
as NamedTuples because PySpark has limited support for dataclasses. We will be working with a secondary python library
called Tinsel, which allows us to easily convert our dataclass schema to a Spark DataFrame schema.

At this point we can leverage the Hypothesis library to specify arbitrary generators we will be using for our dataclasses:

```python
import hypothesis.strategies as st
from hypothesis import given, settings, HealthCheck

week_sunday = (datetime.now() - timedelta(days = datetime.now().weekday() + 1))
time_window = (datetime.now() - timedelta(days = datetime.now().weekday() + 7))

def order_strategy():
    return st.lists(st.builds(Order,
                     order_id=st.text(min_size=5),
                     created_at=st.datetimes(min_value=time_window, max_value=week_sunday),
                     updated_at=st.datetimes(min_value=time_window, max_value=week_sunday),
                     customer_id=st.text(min_size=5),
                     order_items=st.lists(st.builds(OrderItem,
                                           item_id=st.text(min_size=5),
                                           item_value=st.floats(min_value=0, max_value=100)), min_size=1)), min_size=5)
```

In the above code snippet I have specified some upper and down boundaries for the datetimes I will be using in my test
dataset. Likewise for the strings, floats and lists I have used different specifiers. For additional custom generator
strategy consult the excellent Hypothesis documentation.

Now that we have our custom generators we can specify our example transformation function we would like to test:

```python
from pyspark.sql import DataFrame
import pyspark.sql.functions as F

def test_dataframe_transformations(df: DataFrame) -> DataFrame:
    temp_df = df.select('*', F.explode_outer('order_items'))\
        .select(F.col('created_at'), F.col('col.item_value').alias('item_value'))\
        .withColumn('day', F.to_date('created_at'))\
        .withColumn('week', F.weekofyear('created_at'))

    weekly_revenue_df = temp_df.groupBy('week').agg(
        F.sum('item_value').alias('weekly_revenue')
    )

    return temp_df.groupBy('week', 'day').agg(
        F.sum('item_value').alias('daily_revenue')
    ).join(weekly_revenue_df, on='week', how='inner')
```
The property I am interested to test in the above code snippet is the revenue itself. Logically the daily revenue should
always be smaller or equal than the weekly revenue, unless I have an undetected bug in my code.

Now, in order to be able to use pyspark in your unit testing you need to setup a global SparkSession:

```python
import unittest

from pyspark.sql import SparkSession

class PySparkTestCase(unittest.TestCase):
    """Set-up of global test SparkSession"""

    @classmethod
    def setUpClass(cls):
        cls.spark = (SparkSession
                     .builder
                     .master("local[1]")
                     .appName("PySpark unit test")
                     .getOrCreate())

    @classmethod
    def tearDownClass(cls):
        cls.spark.stop()
```

We can setup different profiles for our Hypothesis generators, for example the amount of examples or the deadline:

```python
settings.register_profile(
    "my_profile",
    max_examples=50,
    deadline=60 * 1000,  # Allow 1 min per example (deadline is specified in milliseconds)
    suppress_health_check=(HealthCheck.too_slow, HealthCheck.data_too_large),
)
```

Using our generators, settings and SparkSession we can now write our first unit test:

```python
class SimpleTestCase(PySparkTestCase):

    def test_spark(self):
        @given(data=order_strategy())
        @settings(settings.load_profile("my_profile"))
        def test_samples(data):
            df = self.spark.createDataFrame(data=data, schema=schema, verifySchema=False)
            collected_list = test_dataframe_transformations(df).collect()

            for row in collected_list:
                assert row['weekly_revenue'] >= row['daily_revenue']

        test_samples()
```

This code snippet will feed all the generated datasets to the transformation function as a Spark DataFrame, collect
the results and check the property we were after all along.

In the next blog post we will tackle the orchestration of unit tests and the addition to your Databricks CI/CD pipeline
using [Microsoft's Nutter](https://github.com/microsoft/nutter)! Stay Tuned!

#### The code is available on [github](https://github.com/xdanny/property-testing-pyspark)

### References
[Hypothesis Docs](https://hypothesis.readthedocs.io/en/latest/)
[Tinsel PySpark schema converter](https://github.com/benchsci/tinsel)
[Property Based Testing with QuickCheck](https://typeable.io/blog/2021-08-09-pbt)


Automate all your PySpark Unit Test with Hypothesis!


# Improve your PySpark ETL's performance by providing explicit schema

Have you ever stumbled upon a Spark ETL and you were left wondering how a simple loading of a
dataset can take hours, even though the filtered dataset you are specifying is relatively small?
While there can be multiple reasons for the ETL being slow, from the cloud provider to wrong cluster Spark
configuration of executors, we will focus in this blogpost on optimizing dataframe reads for json and csv datasets.

Going through the [Spark Scala source code](https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/streaming/DataStreamReader.scala#L213)
we can already understand one of the reasons of our query being slow. When reading json datasets Spark will go
through all the json files once to infer the schema before loading our datasets. In case of a heavily distributed
dataset across multiple json files, this can be a source of considerable time spent by Spark Executors. We can
reference the [Spark Documentation](https://spark.apache.org/docs/latest/sql-data-sources-json.html) in order
to better understand what kind of properties are set by default when reading json datasets. A quick fix to speed
our query is to set a smaller sampling ratio for Spark schema inference, by default this value is set to `1.0`.
A better long term fix for your ETL code, especially if the data infrastructure provides some schema guarantees,
is to provide Spark Schema when reading your datasets.

## How to provide a Spark schema when querying json files
### #1 Provide a typed PySpark Schema

As an example we will use an order event that we might receive via our event-bus. This event will have some custom
metadata as well as the order itself with custom order items. In this case our PySpark schema would look like this:

```python
from pyspark.sql.types import *

schema = StructType([
                StructField("metadata",StringType(),False),
                StructField("order",
                            StructType([
                                StructField("order_id",StringType(),False),
                                StructField("created_at",TimestampType(),True),
                                StructField("updated_at",TimestampType(),True),
                                StructField("customer_id",StringType(),False),
                                StructField("order_items",
                                            ArrayType(
                                                StructType([
                                                    StructField("item_id",StringType(),False),
                                                    StructField("item_value",DoubleType(),False)]),False),False)]),False)])

```

We would read the dataframe with the following PySpark command:
```python
df = spark.read.json("our/path", schema=schema)
```

With this configuration spark will read the dataset directly without trying to infer the schema. This setup
can feel quite cumbersome as the data engineer needs to work with Spark types directly, and the definition
of the struct can be quite verbose.

### #2 Provide PySpark schema as Python dataclasses

We can improve the example above using the library `tinsel` that we have leveraged in previous blogposts.
To install the library in your notebook simply run:
```%pip install tinsel```
Using tinsel transformer we can write our schema as dataclasses as in the following example:
```python
from typing import List, NamedTuple, Optional
from tinsel import struct, transform

@struct
class OrderItem(NamedTuple):
  item_id: str
  item_value: float

@struct
class Order(NamedTuple):
  order_id: str
  created_at: Optional[datetime]
  updated_at: Optional[datetime]
  customer_id: str
  order_items: List[OrderItem]

@struct
class Event(NamedTuple):
    metadata: str
    order: Order

schema = transform(Event)
df = spark.read.json("our/path", schema=schema)
```

### #3 Save and load json schema for PySpark dataframes

Writing our schema as python dataclasses is already an excellent step forward, however this might not always
be the right solution. Maintaining schema and schema migration can be quite challenging, and the software
developers might opt on using version control to specify the schemas as yaml or json. With PySpark we can
load the schema specified as json as a static resource, for example from S3. Using the example above we
can generate the json schema:
```python
df.schema.json()
```
Which would print our schema:
```python
json_schema =
"""
    {
  "fields": [
    {
      "metadata": {},
      "name": "metadata",
      "nullable": false,
      "type": "string"
    },
    {
      "metadata": {},
      "name": "order",
      "nullable": false,
      "type": {
        "fields": [
          {
            "metadata": {},
            "name": "order_id",
            "nullable": false,
            "type": "string"
          },
          {
            "metadata": {},
            "name": "created_at",
            "nullable": true,
            "type": "timestamp"
          },
          {
            "metadata": {},
            "name": "updated_at",
            "nullable": true,
            "type": "timestamp"
          },
          {
            "metadata": {},
            "name": "customer_id",
            "nullable": false,
            "type": "string"
          },
          {
            "metadata": {},
            "name": "order_items",
            "nullable": false,
            "type": {
              "containsNull": false,
              "elementType": {
                "fields": [
                  {
                    "metadata": {},
                    "name": "item_id",
                    "nullable": false,
                    "type": "string"
                  },
                  {
                    "metadata": {},
                    "name": "item_value",
                    "nullable": false,
                    "type": "double"
                  }
                ],
                "type": "struct"
              },
              "type": "array"
            }
          }
        ],
        "type": "struct"
      }
    }
  ],
  "type": "struct"
}"""
```

The schema can now be loaded using the following command:
```python
import json

new_schema = StructType.fromJson(json.loads(json_schema))
df = spark.read.json("our/path", schema=new_schema)
```


schema1

Improve your PySpark ETL's performance by providing explicit schema

Recommended

What if a 9B model could beat Sonnet at SQL?

RAG Is Dead. The Engineering Isn't.

Building data quality checks in your pySpark data pipelines

Integration Testing for your Databricks CI/CD Data Pipelines with Microsoft Nutter

Automate all your PySpark Unit Test with Hypothesis!

Follow new posts