
# Why I Compiled My Time Series Forecasting Engine to WebAssembly

If you build tools for enterprises, you quickly learn a painful truth: **handling raw data can be messy.** 

The moment you ask a financial analyst, a logistics manager, or an operations lead to upload their historical sales or inventory CSV to your servers, you trigger a compliance gauntlet. InfoSec questionnaires, data processing agreements, GDPR/CCPA audits, and the constant fear of data leaks. 

But what if you didn't have to touch their raw data at all?

When I built [WaySightAI](https://github.com/xdanny/WaysightAI), I set a strict constraint: **raw time-series data must never leave the user's browser.** The backend should only handle authentication, saved metadata (project names, configurations), usage limits, and high-level AI orchestration. The actual heavy lifting—parsing the CSV, cleaning the data, running stationarity tests, and computing forecasts—had to happen locally.

To do this, I wrote my core time-series math in Rust and compiled it to WebAssembly (WASM). Here is how I structured the hybrid architecture and what I learned along the way.

---

## The Monorepo Architecture

WaySightAI is organized as a monorepo to maintain a clear boundary between the UI, the mathematical engine, and the management layer:

```text
packages/
  web/        React 18 + TypeScript + Vite (renders the workflow UI)
  wasm/       Rust forecasting engine compiled with wasm-pack
  api/        FastAPI backend (SQLAlchemy, Clerk authentication, Redis, Stripe hooks)
```

The forecasting path follows a strict client-side flow:
1. The user uploads a CSV in the React web application.
2. The browser parses the CSV locally and lets the user select date and target columns.
3. The React app loads the `wasm` package dynamically and transfers the numeric vectors to the WASM memory buffer.
4. The Rust/WASM engine processes the calculations and returns JSON containing historical fitted values, forecast points, confidence intervals, and diagnostics.
5. The React app draws the charts using local data.

---

## Writing the Core Math in Rust

Rust is an ideal language for time-series forecasting. It gives us deterministic memory management, zero-cost abstractions, and an ecosystem of libraries like `ndarray` and `statrs` for numerical analysis. 

For example, when a user uploads raw data, it often contains missing values or extreme anomalies (outliers) that would break classical statistical models like ARIMA. In my Rust core (`packages/wasm/src/preprocessing.rs`), we implemented fast IQR-based outlier detection that runs in microseconds:

```rust
use serde::{Deserialize, Serialize};
use wasm_bindgen::prelude::*;

#[derive(Serialize, Deserialize)]
pub struct OutlierIndices {
    pub indices: Vec<usize>,
    pub z_scores: Vec<f64>,
}

/// Detect outliers using the IQR method
#[wasm_bindgen]
pub fn detect_outliers_iqr(values: &[f64]) -> Result<JsValue, JsValue> {
    let mut valid_values: Vec<(usize, f64)> = values
        .iter()
        .enumerate()
        .filter(|(_, &v)| v.is_finite())
        .map(|(i, &v)| (i, v))
        .collect();

    if valid_values.len() < 4 {
        return Ok(serde_wasm_bindgen::to_value(&OutlierIndices {
            indices: vec![],
            z_scores: vec![],
        })?);
    }

    // Sort to find quartiles
    valid_values.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());

    let q1_idx = valid_values.len() / 4;
    let q3_idx = (valid_values.len() * 3) / 4;
    let q1 = valid_values[q1_idx].1;
    let q3 = valid_values[q3_idx].1;
    let iqr = q3 - q1;

    let lower_bound = q1 - 1.5 * iqr;
    let upper_bound = q3 + 1.5 * iqr;

    let mean: f64 = valid_values.iter().map(|(_, v)| v).sum::<f64>() / valid_values.len() as f64;
    let variance: f64 = valid_values
        .iter()
        .map(|(_, v)| (v - mean).powi(2))
        .sum::<f64>()
        / valid_values.len() as f64;
    let std_dev = variance.sqrt();

    let mut outlier_indices = Vec::new();
    let mut z_scores = Vec::new();

    for (idx, &val) in values.iter().enumerate() {
        if val.is_finite() && (val < lower_bound || val > upper_bound) {
            outlier_indices.push(idx);
            z_scores.push((val - mean) / std_dev);
        }
    }

    Ok(serde_wasm_bindgen::to_value(&OutlierIndices {
        indices: outlier_indices,
        z_scores,
    })?)
}
```

By compiling this function via `wasm-pack`, JavaScript can invoke it directly, passing array buffers and receiving the results back as native JS values via `serde_wasm_bindgen`.

---

## Bridging Rust and React in the Browser

To make loading the compiled WASM file seamless, I built a TypeScript wrapper class (`WasmEngine`) that handles dynamic imports, module initialization, and the type conversion between standard JavaScript arrays and WebAssembly-friendly typed arrays (`Float64Array`).

Here is a simplified version of the loader wrapper:

```typescript
import type {
  WaySightAIWasm,
  ForecastConfig,
  ForecastResult,
  DataStats
} from '../../types/wasm';

export class WasmEngine {
  private wasmModule: WaySightAIWasm | null = null;
  private isInitialized = false;
  private loadingPromise: Promise<void> | null = null;

  /**
   * Dynamically import and initialize the WASM module
   */
  async load(): Promise<void> {
    if (this.loadingPromise) return this.loadingPromise;
    if (this.isInitialized && this.wasmModule) return Promise.resolve();

    this.loadingPromise = (async () => {
      try {
        // Dynamically import the JS entry point generated by wasm-pack
        const wasmModule = await import('/pkg/waysightai_wasm.js');

        // Initialize the WebAssembly binary instance
        await wasmModule.default();
        wasmModule.init();

        this.wasmModule = wasmModule as unknown as WaySightAIWasm;
        this.isInitialized = true;
        console.log('[WasmEngine] WASM module loaded successfully');
      } catch (error) {
        console.error('[WasmEngine] Failed to load WASM module:', error);
        throw error;
      }
    })();

    await this.loadingPromise;
    this.loadingPromise = null;
  }

  private ensureInitialized(): WaySightAIWasm {
    if (!this.isInitialized || !this.wasmModule) {
      throw new Error('WASM module not initialized. Call load() first.');
    }
    return this.wasmModule;
  }

  /**
   * Run a local forecast job
   */
  async runForecast(
    timestamps: number[],
    values: number[],
    isUnixMs: boolean,
    config: ForecastConfig
  ): Promise<ForecastResult> {
    const wasm = this.ensureInitialized();
    const configJson = JSON.stringify(config);
    
    // Call the exported Rust function
    const resultJson = wasm.run_forecast(timestamps, values, isUnixMs, configJson);
    return JSON.parse(resultJson);
  }
}
```

Because WebAssembly compilation produces a `.wasm` file alongside a `.js` wrapper, Vite allows us to load this dynamically on demand. This keeps the initial page bundle small, downloading the 1.2MB WASM forecasting package only when the user reaches the active editor dashboard.

---

## Performance: Browser vs. Server

Running ARIMA models (which rely on MLE optimization) and triple exponential smoothing (Holt-Winters) in a single-threaded browser environment might sound slow. However, my benchmarks showed surprising results:

| Dataset Size (Rows) | Operation | Execution Location | Duration (ms) |
| :--- | :--- | :--- | :--- |
| 500 rows | Stationarity Tests (ADF + KPSS) | Browser (WASM) | 8ms |
| 500 rows | ARIMA(1,1,1) Parameter Tuning | Browser (WASM) | 45ms |
| 5,000 rows | Outlier Detection & Imputation | Browser (WASM) | 12ms |
| 5,000 rows | Holt-Winters Forecast (Horizon=30) | Browser (WASM) | 185ms |

For typical business dashboards where datasets range from 100 to 10,000 rows, the execution is effectively instantaneous. In contrast, sending this data over an API to a Python service, waiting for it to parse, compute, serialize, and return would easily take 800ms to 2.5 seconds—even before accounting for server cold starts or database query times.

Furthermore, running the computations locally means we can build **interactive visual sliders**. If a user wants to test different seasonal parameters (alpha, beta, gamma), they can slide the control in React, and the chart updates in near real-time (under 50ms) as the local WASM engine re-calculates the parameters.

---

## The Security and Cost Tradeoffs

Going fully client-side with WASM is not a silver bullet. It introduces key trade-offs:

1. **WASM Overhead**: The first load requires fetching the WASM binary. For users on slow connections, this adds a couple of seconds of initialization delay.
2. **IP Exposure**: Because WebAssembly runs on the client, the compiled code is technically accessible. If you have proprietary forecasting models that you want to keep secret, running them in WASM is a risk, as the assembly can be reverse-engineered. For standard algorithms (ARIMA, Simple/Double/Triple Exponential Smoothing), this is not a concern.
3. **CPU Constraints**: If a user uploads a dataset with millions of rows, the browser will lag, and you risk hitting browser memory limits. For massive scale, you still need an orchestrator like AWS SageMaker.

However, the cost savings are massive. Since my servers don't perform any math, my FastAPI backend runs on a cheap container instance. I can handle thousands of concurrent users without paying for expensive auto-scaling CPU nodes or GPU workers.

## Conclusion

Compiling my forecasting math to WebAssembly turned out to be the single best architectural decision I made. It allowed me to promise absolute data privacy to users, gave sub-100ms forecast updates, and kept server infrastructure costs virtually zero.

If you are building data-heavy SaaS tools, don't default to sending every CSV to a server. Let Rust and WebAssembly do the heavy lifting where the data already lives: right in the browser.


Why I Compiled My Time Series Forecasting Engine to WebAssembly


# Orchestrating MLOps with Bedrock AgentCore and SageMaker

When you build machine learning pipelines, you realize quickly that the modeling itself is only a small fraction of the effort. The real work is in the plumbing: cleaning datasets, exploring stationarity and seasonality, engineering lagging or rolling window features, running hyperparameter search, tracking jobs, building reports, and deploying endpoints.

Traditionally, data scientists write messy Jupyter notebooks to do this manually. But what if an **autonomous AI agent** could handle this entire lifecycle?

I set out to build exactly that: a production-ready forecasting system that takes a raw CSV in S3, performs advanced statistical analysis, tunes hyperparameters, trains ARIMA/LSTM models, and deploys a live endpoint—all managed by an LLM-driven orchestrator.

Here is how I built a three-layer architecture combining **AWS Bedrock AgentCore**, **Code Interpreter** sandboxes, and **AWS SageMaker**.

---

## The Three-Layer Architecture

An autonomous system needs to split concerns between reasoning, executing code, and running heavy machine learning jobs. I designed a three-layer architecture:

```mermaid
graph TD
    subgraph Layer1["Layer 1: Orchestration (AgentCore)"]
        Agent["Intelligent Agent<br/>(16 Tools)"]
        Boto3["AWS SDK (boto3)"]
    end

    subgraph Layer2["Layer 2: Data Science (Code Interpreter)"]
        Sandbox["Isolated Python Sandbox"]
        EDA["EDA / Statistical Tests"]
        Features["Feature Engineering"]
        Plotly["Plotly HTML Reports"]
    end

    subgraph Layer3["Layer 3: Scalable ML (SageMaker)"]
        Tuning["Bayesian HPO Jobs"]
        Training["Scale Model Training"]
        Endpoint["Production REST Endpoint"]
    end

    Agent --> Sandbox
    Agent --> Boto3
    Boto3 --> Tuning
    Boto3 --> Training
    Boto3 --> Endpoint
```

Let's break down how these layers interact.

---

## Layer 1: Orchestration with AWS Bedrock AgentCore

The orchestrator is built using [AWS Bedrock AgentCore](https://aws.amazon.com/bedrock/) (packaged via the Strands library). The agent has access to 16 specialized tools ranging from statistical analysis to SageMaker deployment commands.

In `agent.py`, the agent is initialized with a list of Python tool references:

```python
from strands import Agent
from agents.advanced_eda_agent import run_advanced_eda
from agents.intelligent_feature_engineering_agent import recommend_features, create_features
from agents.sagemaker_simple import (
    create_sagemaker_training_job,
    get_training_job_status,
    deploy_sagemaker_model,
    invoke_sagemaker_endpoint
)

agent = Agent(
    name="IntelligentForecastingAgent",
    description="Intelligent time series forecasting with ARIMA and LSTM model comparison.",
    tools=[
        run_advanced_eda,
        recommend_features,
        create_features,
        create_sagemaker_training_job,
        get_training_job_status,
        deploy_sagemaker_model,
        invoke_sagemaker_endpoint,
        # ... other tools
    ]
)
```

The agent uses these tools to execute a 7-step forecasting workflow:
1. **Advanced EDA**: Runs ADF/KPSS tests and seasonal decomposition.
2. **Feature Recommendations**: The agent examines stats and recommends features.
3. **Feature Engineering**: Generates rolling stats, lag columns, and calendar features.
4. **Bayesian Tuning**: Launches SageMaker HPO to find optimal hyperparameters.
5. **Model Training**: Trains classical models (ARIMA) and deep learning models (LSTM).
6. **Comparison**: Evaluates metrics on hold-out test sets.
7. **Report Generation**: Deploys the best model and outputs a presigned URL to an interactive Plotly report.

---

## Layer 2: Safe Execution via Code Interpreter

A major challenge for LLM agents is handling arbitrary code execution. I cannot run unverified code written by an LLM directly on my application server or inside my main database container.

My solution is to offload all data science computations (EDA, feature engineering, Plotly charting) to an isolated **Code Interpreter** sandbox. This sandbox has pre-installed scientific libraries (`pandas`, `scipy`, `statsmodels`, `plotly`) but is locked down with limited CPU, memory, and network permissions.

Here is the design pattern I used for running our Advanced EDA tool:

```python
@tool
def run_advanced_eda(dataset_s3_path: str, time_column: str = None, value_column: str = None) -> str:
    """
    Runs stationarity tests (ADF, KPSS), seasonal decomposition, and ACF/PACF analysis.
    """
    # 1. Parse the target S3 paths
    bucket, key = parse_s3_uri(dataset_s3_path)
    
    # 2. Build the Python script to run in the sandbox
    code = f'''
import pandas as pd
import numpy as np
import json
import subprocess
from statsmodels.tsa.stattools import adfuller, kpss
from statsmodels.tsa.seasonal import seasonal_decompose

# Downloader helper - Sandbox has pre-configured read-only AWS CLI credentials
subprocess.run(['aws', 's3', 'cp', 's3://{bucket}/{key}', '/tmp/data.csv'], check=True)
df = pd.read_csv('/tmp/data.csv')

# ... Auto-detect time/value columns and clean data ...
series = df[value_col].dropna()

# Run ADF and KPSS tests
adf_stat, adf_pvalue, _, _, _, _ = adfuller(series)
kpss_stat, kpss_pvalue, _, _ = kpss(series, regression='c')

# Return results as JSON
print("===JSON_START===")
print(json.dumps({{
    "adf": {{"statistic": adf_stat, "p_value": adf_pvalue}},
    "kpss": {{"statistic": kpss_stat, "p_value": kpss_pvalue}}
}}))
print("===JSON_END===")
'''
    # 3. Send script to the Code Interpreter sandbox and parse stdout
    stdout = sandbox_client.execute(code)
    return extract_json_from_stdout(stdout)
```

By packaging Python code as a text block and passing it to the sandbox, we keep our agent's host server secure. The sandbox pulls the data from S3, computes stats, and returns clean, structured JSON back to the agent.

---

## Layer 3: Scalable ML Training and Deployment on SageMaker

While Code Interpreter is great for light calculations, it lacks the memory and compute capacity to train deep learning models (like LSTMs) or run large-scale hyperparameter searches. For that, I delegate to **AWS SageMaker**.

When the agent reaches the model training step, it uses the AWS SDK (`boto3`) to launch real SageMaker HPO (Hyperparameter Optimization) jobs using preconfigured Scikit-Learn or PyTorch ECR containers:

```python
import boto3
import json

sagemaker = boto3.client('sagemaker')

@tool
def create_sagemaker_training_job(
    job_name: str,
    dataset_s3_path: str,
    role_arn: str
) -> str:
    """
    Launches a Scikit-Learn ECR container on SageMaker for training.
    """
    try:
        response = sagemaker.create_training_job(
            TrainingJobName=job_name,
            AlgorithmSpecification={
                'TrainingImage': '683313688378.dkr.ecr.us-east-1.amazonaws.com/sagemaker-scikit-learn:1.2-1-cpu-py3',
                'TrainingInputMode': 'File'
            },
            RoleArn=role_arn,
            InputDataConfig=[{
                'ChannelName': 'train',
                'DataSource': {
                    'S3DataSource': {
                        'S3DataType': 'S3Prefix',
                        'S3Uri': dataset_s3_path,
                    }
                },
                'ContentType': 'text/csv'
            }],
            OutputDataConfig={
                'S3OutputPath': f's3://{BUCKET}/sagemaker/models/'
            },
            ResourceConfig={
                'InstanceType': 'ml.m5.xlarge',
                'InstanceCount': 1,
                'VolumeSizeInGB': 10
            },
            HyperParameters={
                'p': '2',
                'd': '1',
                'q': '3',
                'seasonal-p': '1',
            },
            StoppingCondition={'MaxRuntimeInSeconds': 3600}
        )
        return json.dumps({'success': True, 'job_arn': response['TrainingJobArn']})
    except Exception as e:
        return json.dumps({'success': False, 'error': str(e)})
```

Because SageMaker jobs run asynchronously and can take 5 to 60 minutes, the agent uses a polling loop tool (`get_training_job_status`) to track the execution. Once the job succeeds, the agent triggers `deploy_sagemaker_model` to package the model artifact (`model.tar.gz`) from S3 and host it behind a real-time SageMaker endpoint.

---

## The Two-Phase Data Pattern for Reports

Generating and serving rich reports presents a unique security challenge. The Code Interpreter sandbox (which generates the Plotly interactive HTML files) should **never** have write access to my production S3 buckets, nor should it hold long-term AWS credentials.

To solve this, I implemented a **Two-Phase Data Handling** pattern:

```mermaid
sequenceDiagram
    participant Agent as AgentCore (Has S3 Write Credentials)
    participant CI as Code Interpreter Sandbox (No S3 Credentials)
    participant S3 as AWS S3

    Agent->>CI: Execute charting script
    CI->>CI: Generate Plotly chart & serialize to HTML
    CI-->>Agent: Print HTML string to stdout
    Agent->>Agent: Extract HTML string from stdout
    Agent->>S3: Upload HTML content using boto3
    S3-->>Agent: Return secure presigned URL
```

By keeping the AWS write credentials strictly in the Orchestration Layer (AgentCore) and treating the Code Interpreter as a pure calculation engine, I maintain least-privilege security across my cloud infrastructure.

---

## Conclusion & Lessons Learned

Building an autonomous ML pipeline taught me several key lessons:

1. **Decouple CPU Profiles**: Keep light calculations (cleaning, simple charts) inside rapid, cheap sandboxes, and reserve expensive cloud computing instances (SageMaker ml.m5/ml.p3 nodes) for heavy model training.
2. **Handle Asynchrony Gracefully**: LLMs are naturally sequential. To prevent the agent from locking up during a 30-minute SageMaker job, use explicit state-machine trackers and periodic polling tools.
3. **Sandbox Everything**: Never let an LLM-generated script run natively. Use a secure container runtime with strict network isolation.

By combining AWS Bedrock AgentCore's reasoning capabilities with SageMaker's massive scaling power, I successfully took a standard time-series pipeline and turned it into a fully autonomous MLOps platform.


Orchestrating MLOps with Bedrock AgentCore and SageMaker


# 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.


# 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)
```


Improve your PySpark ETL's performance by providing explicit schema


# 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


# 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!


data quality1 posts

Building data quality checks in your pySpark data pipelines

Recommended

Why I Compiled My Time Series Forecasting Engine to WebAssembly

Orchestrating MLOps with Bedrock AgentCore and SageMaker

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