
# Can a local 9B model compete on multi-turn data analysis?

Single-turn text-to-SQL has become a much more serious benchmark. The
[BIRD benchmark](https://bird-bench.github.io/) pushed the field toward
real databases, messy schemas, and executable SQL instead of toy examples. On
that kind of one-question task, a strong model can often read a schema, read
one question, and produce a useful query with zero-shot prompting.

That is impressive. It is also not the same thing as analysis.

Analysis usually sounds more like this:

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

None of those follow-ups is hard because of SQL syntax alone. The system has to
remember the metric, carry filters forward, normalize values, change grain,
preserve semantic intent, notice empty results, and recover. A query can execute
cleanly and still be the wrong analytical move.

That is the question behind this project:

> Can a small specialized local model learn the behaviors and semantic concepts
> needed to challenge hosted state-of-the-art systems on multi-turn data
> analysis?

I do not have that hosted comparison yet. The current public claim is narrower:
on a fixed local proxy slice, the best non-oracle run improved value accuracy
over the base model. A local-vs-hosted claim would need the same rows, same
scorer, same oracle boundary, hosted and local manifests, generated-history
rollout, latency and cost accounting, and a transfer run on an interactive
benchmark such as [BIRD-Interact](https://arxiv.org/abs/2510.05318).

## Attached Code

The attached code lives in
[`xdanny/multiturn-sql-finetuning`](https://github.com/xdanny/multiturn-sql-finetuning).
The runnable lab is published at
[`/labs/local-multiturn-sql-finetuning/`](/labs/local-multiturn-sql-finetuning/),
with the
[Marimo source](https://github.com/xdanny/multiturn-sql-finetuning/blob/main/notebooks/labs/local_multiturn_sql_lab.py)
in the repository.

The lab is a compact code inset, not the main argument and not a setup guide. It
runs on CPU, shows the tiny failure case, and leaves the full endpoint/training
path in the repository.

## The Benchmarks In Plain English

The benchmark names are easy to blur together, but they test different things.

[BIRD](https://bird-bench.github.io/) is a single-turn text-to-SQL benchmark:
the model sees a database schema and one natural-language question, then writes
one executable SQL query. It is useful because it uses more realistic databases
than older toy datasets, but it does not test whether the model can carry an
analysis through follow-up turns.

[CoSQL](https://arxiv.org/abs/1909.05378) is conversational. A user asks a
sequence of related database questions, and each turn is still grounded in SQL.
That makes it a practical proxy for this project: small enough to run locally,
but close enough to expose context carryover failures.

[SParC](https://arxiv.org/abs/1906.02285) is also context-dependent text-to-SQL.
It focuses on coherent question sequences across databases. I treat it as a
context signal, not as the final target.

[BIRD-Interact](https://bird-interact.github.io/) is closer to the end goal. It
turns text-to-SQL into an interactive setting with a database, extra knowledge,
tools, and a user simulator. That is the kind of benchmark where recovery,
clarification, and repeated analysis matter.

Synthetic schema-rich SQL is not a benchmark claim at all. It is a lab bench:
small hand-built cases for value normalization, entity resolution, fanout,
bridge tables, governed metrics, and empty-result recovery.

## What Breaks After One Query

The first mistake is assuming that single-turn SQL success transfers cleanly to
multi-turn analysis. It does not fail in one generic way.

The main failure modes are:

- state carryover: remembering that the second question still depends on the
  first one;
- value grounding: mapping a phrase like "France" to the stored value `FR`;
- grain shift: moving from country totals to monthly totals without losing the
  metric;
- metric intent: preserving that revenue means the governed revenue metric, not
  any convenient `SUM`;
- result-aware recovery: changing behavior after an empty or suspicious result;
- generated-history drift: surviving the model's own earlier mistakes, not only
  clean reference history.

That is why the data sources are split. CoSQL is the fast proxy. SParC helps
reason about context-dependent SQL. BIRD mini-dev is useful for single-turn
execution-harness checks. BIRD-Interact is the eventual interactive target.
Synthetic fixtures isolate rare failures before spending endpoint time.

The point is to stop treating all SQL data as one pile. Each dataset answers a
different engineering question.

## What The Current Loop Proves

The current loop is deliberately small: 100 CoSQL turns from 32 dialogs, one
local endpoint path, and one SQLite scorer. It uses teacher-forced history,
which means later turns receive clean previous reference SQL rather than the
model's own earlier SQL. That makes this a controlled proxy, not a full
interactive rollout.

The scoring also matters. Strict execution accuracy asks whether the query
returns the expected result under the scorer. Value accuracy is more forgiving
about harmless output labels and focuses on returned values. That distinction
changed the interpretation: some queries were right in substance but failed
because an alias name differed.

On that slice, base
[Qwen 3.5 9B](https://huggingface.co/Qwen/Qwen3.5-9B) scored `0.370` strict
execution accuracy. A 100-step LoRA moved strict execution to `0.530`. After
fixing the scorer to compare result values rather than harmless output labels,
the best non-oracle proxy reached `0.640` value accuracy.

![Proxy SQL accuracy ladder generated from the claim ledger](/assets/blog/local-multiturn-sql-finetuning/evidence/accuracy-ladder.svg)

That is progress, not a conclusion. The remaining misses are not just "bad
SQL." They split across state, value grounding, schema alignment, projection,
grain, semantic intent, and recovery.

The schema validator made one class more concrete. It catches unknown columns,
wrong-table columns, and ambiguous unqualified columns before execution. That
turns some misses into repairable labels instead of vague execution failures.

## Which Target Is Actually Ready?

Direct SQL supervised fine-tuning is the control arm. It asks the model to emit
SQL directly and learns from examples. It is the first thing to beat, not the
only thing worth training.

The serious method comparison is between several targets:

- direct SQL: generate the final SQL immediately;
- planner-first SQL: predict tables, columns, joins, filters, and output shape
  before generating SQL;
- semantic-layer state: recover entities, values, grain, and governed metrics
  before SQL, similar in spirit to semantic layers in
  [Cube](https://docs.cube.dev/docs/introduction) or
  [dbt](https://docs.getdbt.com/docs/build/semantic-models);
- `MEASURE()`-preserving DSL: use a small domain-specific language, or DSL, to
  keep metric intent explicit before compiling to SQL;
- behavior/recovery tuning: teach the model what to do after empty results,
  failed executions, or suspicious intermediate answers.

Each target has to win against the right control on the same rows, same scorer,
same prompt boundary, and same oracle policy. Current evidence says direct SQL
is rankable as the control. Planner, metric DSL, recovery, semantic-layer, and
hosted comparisons each still need a specific missing artifact before they can
be promoted as better methods.

[DSPy](https://arxiv.org/abs/2310.03714) fits here, but only if it optimizes a
scoreable intermediate program. A prompt search over final SQL wording is weak
evidence. A DSPy search over a planner, value resolver, or semantic DSL program
is better because the output can be judged before SQL generation. The
[DSPy optimizer docs](https://dspy.ai/learn/optimization/optimizers/) frame
that the same way: optimize a program against a metric, not just a prettier
prompt.

## SQL Is Too Late For Some Failures

If a user says "France" and the database stores `FR`, a syntactically valid SQL
query can return no rows. If a follow-up changes country totals into monthly
totals, the model has to change grain without losing the metric. If an earlier
result was empty, continuing blindly is not analysis.

That is why the repository now treats data artifacts as part of the model
problem. The important artifacts are:

- value grounding labels, derived from gold SQL, for supervision and
  evaluation;
- a value index, built from database contents, for production-like lookup;
- schema diagnostics, so wrong-table and ambiguous-column errors are visible
  before execution;
- synthetic fixtures for aliases, casing, abbreviations, Roman numerals, dates,
  bridge-table fanout, `MEASURE()` preservation, and empty-result recovery.

The gap between value labels and the value index is the semantic-layer problem
in miniature: the value may exist in the database, but the user's phrase may not
match the stored literal.

## The Missing Plan

The oracle result is the most useful and most dangerous part of the project. In
this post, oracle means information extracted from the reference SQL. That is
allowed for diagnosis, but it is not allowed in a production system because it
looks at the answer.

When the model was given tables, columns, joins, projection shape, skeleton, and
duplicate policy derived from the reference SQL, value accuracy rose to `0.850`
with prompt-only oracle hints and `0.890` with an oracle-trained format.

That does not mean the system works. It means the missing plan is valuable.

A production system has to predict that plan without reading the answer. The
repo now carries `gold_plan` for scoring and `predicted_plan` for non-oracle
generation. The first lexical planner is intentionally simple.

![Planner baseline generated from the planning support code](/assets/blog/local-multiturn-sql-finetuning/evidence/planner-baseline.svg)

The next gate is clear: row pairing is ready, but the planner is too weak on
columns and projection shape to promote predicted-planner SQL as a win. That is
a useful failure. It says the next endpoint run should not be "try a bigger
prompt." It should be "improve the non-oracle planner, then compare
predicted-planner SQL against direct SQL on identical rows."

## Metric Intent Needs Its Own Target

For metric-heavy analysis, raw SQL may be the wrong first language. If the user
is asking for revenue, the system should preserve the governed metric as long as
possible. A DSL, or domain-specific language, is just a small intermediate
language for the task:

```text
MEASURE(revenue) BY customer_country
WHERE customer_country = 'FR'
ORDER BY MEASURE(revenue) DESC
LIMIT 5
```

Then the semantic model can compile the SQL:

```sql
SELECT customers.country AS customer_country, SUM(orders.amount) AS revenue
FROM orders
JOIN customers ON orders.customer_id = customers.id
WHERE customers.country = 'FR'
GROUP BY customers.country
ORDER BY revenue DESC
LIMIT 5
```

This is not a win yet. It is a better-shaped experiment. The DSL target only
matters if it parses, preserves governed `MEASURE()` intent, compiles through a
semantic model, and beats direct SQL on matching metric-heavy rows.

## What The Next Run Has To Prove

The current answer is modest: a small local model moved on a fixed proxy, and
the oracle ceiling says intermediate planning is high leverage. It does not yet
prove that a local 9B model beats hosted state-of-the-art systems.

The order I would trust is:

1. Improve the non-oracle planner and score planner F1 directly.
2. Run predicted-planner SQL against the direct SQL control.
3. Score semantic-layer artifacts for values, entities, grain, and governed
   measures before using them as prompt context.
4. Evaluate DSL-first generation on metric-heavy rows and compare compiled SQL
   against direct SQL.
5. Run generated-history rollout, not only teacher-forced history.
6. Add hosted baselines on the same rows with the same scorer, latency, and cost
   accounting.
7. Transfer the protocol to BIRD-Interact.

So the project is no longer just "fine-tune a SQL model." It is a set of
competing fine-tuning hypotheses:

- direct SQL SFT may be enough for some distribution shift;
- planner-first may help when schema linking and projection shape are the
  bottleneck;
- semantic-layer tuning may help when values, entities, grain, and governed
  metrics are the bottleneck;
- DSL-first may help when preserving `MEASURE()` is more important than emitting
  SQL immediately;
- behavior/recovery tuning may help only when generated-history rollout exposes
  the model's own previous mistakes.

That is the bar for the next version. A local 9B model might become competitive
on serious multi-turn analysis, but the project has to earn that claim one
artifact at a time.

### Sources Used In The Text

- [Qwen 3.5 9B model card](https://huggingface.co/Qwen/Qwen3.5-9B)
- [BIRD benchmark project page](https://bird-bench.github.io/)
- [BIRD paper: Can LLM Already Serve as a Database Interface?](https://papers.nips.cc/paper_files/paper/2023/file/83fc8fab1710363050bbd1d4b8cc0021-Paper-Datasets_and_Benchmarks.pdf)
- [BIRD-Interact project page](https://bird-interact.github.io/)
- [BIRD-Interact paper](https://arxiv.org/abs/2510.05318)
- [CoSQL paper](https://arxiv.org/abs/1909.05378)
- [SParC paper](https://arxiv.org/abs/1906.02285)
- [DSPy paper](https://arxiv.org/abs/2310.03714)
- [DSPy optimizer docs](https://dspy.ai/learn/optimization/optimizers/)
- [Cube semantic layer docs](https://docs.cube.dev/docs/introduction)
- [dbt Semantic Layer docs](https://docs.getdbt.com/docs/build/semantic-models)
- [marimo docs](https://docs.marimo.io/)


Can a local 9B model compete on multi-turn data analysis?

Context Engineering with FAISS and Strands: Beyond the RAG Hype


# 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


# 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](https://unpackingdata.com/property-testing-pyspark). 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)


Welcome to Unpacking Data. Exploring data engineering, AI engineering, and building robust systems at scale.

Featured Post

Can a local 9B model compete on multi-turn data analysis?

Recent Posts

Context Engineering with FAISS and Strands: Beyond the RAG Hype

Building data quality checks in your pySpark data pipelines

Improve your PySpark ETL's performance by providing explicit schema

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

Automate all your PySpark Unit Test with Hypothesis!

Subscribe to new posts