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)


Automate all your PySpark Unit Test with Hypothesis!


# Can a local 9B model compete with state-of-the-art systems for data analysis?

This is the core question driving this project. I'm not interested in seeing if a small model can generate toy SQL queries or look impressive in a carefully curated demo. The real test is whether a model running on a single local workstation can handle the heavy lifting we usually reserve for frontier API models. Can it read a complex database schema, carry context across multiple follow-up questions, nail the right joins, normalize messy user inputs, and actually return trustworthy data?

I honestly don't know the answer yet, which is exactly why I'm building this out.

I chose SQL generation as the proving ground because it eliminates the wiggle room you get with natural language tasks. A SQL query either executes and returns the correct rows, or it fails. Multi-turn SQL pushes this further by mimicking how real analysts work: a user asks a question, gets a result, immediately asks for a breakdown or filter, and expects the system to remember everything they just talked about.

This first post isn't about claiming a victory lap for a local Qwen 3.5 9B model against GPT-4o, and I won't be dropping a final BIRD-Interact score just yet. Instead, this post is about establishing the engineering pipeline needed to make that comparison honest. We have to walk before we can run: preparing the multi-turn data, fine-tuning LoRA adapters, serving them via vLLM, executing the SQL locally against SQLite, and—crucially—isolating genuine progress from accidental data leakage.

Here is how I'm breaking down the journey:

| Post | The Focus |
| --- | --- |
| **1. Benchmark Setup** | Setting up a local evaluation loop we can trust. (This post) |
| **2. Local Training and Serving** | The nuts and bolts of the RTX 5090, Unsloth, and vLLM pipeline. |
| **3. The Dataset Bridge** | Reconciling BIRD-Interact, BIRD mini-dev, CoSQL, SParC, and synthetic SQL. |
| **4. Analyzing Fine-tuning Runs** | Which training tweaks actually move the needle? |
| **5. Deconstructing Failures** | Why do schema linking, joins, and dialogue state break down in practice? |
| **6. Building a Planner System** | Moving from an "oracle" cheat to a production-ready query planner. |

## Why one-shot SQL isn't enough

Single-turn text-to-SQL is already a notoriously difficult problem. The model has to ingest a schema, parse a standalone question, and write a flawless query:

```text
Which airlines fly to Boston?
```

But multi-turn SQL fundamentally changes the game. Real users rarely restate their entire intent. They speak in shorthand:

```text
Which airlines fly to Boston?
Only show the ones from France.
```

That second question is entirely dependent on the first. For an agent to succeed here, it has to carry forward the route target and the airline entity, append a country filter, select the correct tables, and preserve the original output shape. If it drops the ball on any of these, it might generate perfectly valid SQL that answers the completely wrong question.

When building this out, I ran into several concrete failure modes:
- **Filter amnesia**: The model forgets a constraint established in a previous turn.
- **Table misattribution**: Applying the correct value to the wrong table entirely.
- **Join fan-out**: Selecting a join path that inadvertently multiplies the rows.
- **Projection drift**: Returning a raw count when the user specifically asked for names.
- **Value grounding failures**: Querying for a string that simply doesn't exist in the database (e.g. `'France'` vs `'FR'`).
- **History misalignment**: Resolving a follow-up query against the wrong turn in a long conversation.

This is why we can't treat multi-turn SQL as just "single-turn SQL plus a chat history string." The model has to actively *plan* before it generates code.

## Isolating the datasets

I'm pulling from several SQL datasets, but tossing them all into one massive fine-tuning pile is a bad idea. Each dataset serves a specific engineering purpose:

| Dataset | Role in this project |
| --- | --- |
| **BIRD-Interact** | Our north star. This is the interactive benchmark closest to real-world data analysis. |
| **BIRD mini-dev** | Single-turn BIRD-style rows, great for testing our execution harness but not a substitute for dialogue. |
| **CoSQL** | Our current local proxy. It provides multi-turn dialogues over Spider-style databases and runs easily against SQLite. |
| **SParC** | Provides related, context-dependent SQL data to bolster the training signal. |
| **Synthetic SQL** | Single-turn SQL paired with complex schemas to prevent the model from overfitting to simple database layouts. |

For the benchmarks reported below, I consistently evaluated against the exact same proxy slice: **100 turns across 32 CoSQL dialogues**. 

I deliberately kept this slice small so I could iterate quickly while still exposing the system to complex follow-up failures. I also used "teacher-forced" dialogue history—meaning when the model is evaluated on turn 3, it's given the correct, ground-truth SQL from turns 1 and 2, rather than its own potentially broken history. This tells us whether the model is capable of utilizing clean context, deferring the harder problem of error recovery for later.

## Testing the fine-tuning loop

The fine-tuning runs I'm highlighting here are small, exploratory LoRA experiments, not final training runs. I used balanced slices of 192 and 384 rows from our dataset mix just to verify that the training loop could actually move the needle.

Starting out, the base Qwen 3.5 9B model achieved a meager `0.370` strict execution accuracy. After a 100-step LoRA fine-tune, that jumped to `0.530`.

| Run | Training rows | Steps | Strict execution accuracy |
| --- | ---: | ---: | ---: |
| Base `unsloth/Qwen3.5-9B` | 0 | 0 | 0.370 |
| 50-step LoRA | 192 | 50 | 0.420 |
| 100-step LoRA | 192 | 100 | 0.530 |
| 384-row LoRA | 384 | 50 | 0.420 |

While a `0.530` accuracy is a real improvement, it's completely unviable for production. If a model fails on nearly half of its turns, the user's conversation is going to derail constantly. However, these numbers gave me exactly what I needed: proof that the training and evaluation loop was sensitive. The fact that the scores shifted gave me a reliable signal to start debugging.

## Moving beyond strict scoring

I quickly realized that strict execution scoring was masking part of the picture. The scorer would brutally fail a query simply because the model chose a different output alias than the reference SQL, even if the returned rows and values were identical.

To fix this, I implemented a **value-aware execution scorer**. It still rigorously checks row values, duplicate counts, and order-sensitive outputs, but it forgives harmless cosmetic differences in column labels.

Under this fairer metric, my best non-oracle run maxed out at **`0.640`**. This is the number I consider our true production-style proxy ceiling. It represents a run using only the inputs a real system would have access to: the question, conversation history, schema context, and standard prompt instructions.

<figure>
<svg viewBox="0 0 820 360" role="img" aria-labelledby="accuracy-title accuracy-desc" style="max-width:100%;height:auto;border:1px solid #ddd;background:#fff">
  <title id="accuracy-title">Proxy SQL accuracy on the fixed CoSQL slice</title>
  <desc id="accuracy-desc">Bar chart showing base strict accuracy 0.370, 100-step LoRA strict accuracy 0.530, best non-oracle value accuracy 0.640, oracle prompt value accuracy 0.850, and oracle-trained value accuracy 0.890.</desc>
  <rect x="0" y="0" width="820" height="360" fill="#fff"/>
  <text x="36" y="38" font-family="Arial, sans-serif" font-size="22" font-weight="700" fill="#111">The useful result is the gap, not the top number</text>
  <text x="36" y="63" font-family="Arial, sans-serif" font-size="13" fill="#555">Fixed 100-turn CoSQL proxy. The last two bars use gold SQL-derived planning labels.</text>
  <line x1="86" y1="288" x2="782" y2="288" stroke="#222"/>
  <line x1="86" y1="88" x2="86" y2="288" stroke="#222"/>
  <line x1="86" y1="248" x2="782" y2="248" stroke="#eee"/>
  <line x1="86" y1="208" x2="782" y2="208" stroke="#eee"/>
  <line x1="86" y1="168" x2="782" y2="168" stroke="#eee"/>
  <line x1="86" y1="128" x2="782" y2="128" stroke="#eee"/>
  <rect x="120" y="214" width="82" height="74" fill="#222"/>
  <rect x="252" y="182" width="82" height="106" fill="#555"/>
  <rect x="384" y="160" width="82" height="128" fill="#888"/>
  <rect x="516" y="118" width="82" height="170" fill="#ccc"/>
  <rect x="648" y="110" width="82" height="178" fill="#e5e5e5" stroke="#222"/>
  <text x="136" y="206" font-family="Arial, sans-serif" font-size="14" font-weight="700" fill="#111">0.370</text>
  <text x="268" y="174" font-family="Arial, sans-serif" font-size="14" font-weight="700" fill="#111">0.530</text>
  <text x="400" y="152" font-family="Arial, sans-serif" font-size="14" font-weight="700" fill="#111">0.640</text>
  <text x="532" y="110" font-family="Arial, sans-serif" font-size="14" font-weight="700" fill="#111">0.850</text>
  <text x="664" y="102" font-family="Arial, sans-serif" font-size="14" font-weight="700" fill="#111">0.890</text>
  <text x="111" y="313" font-family="Arial, sans-serif" font-size="12" fill="#222">Base</text>
  <text x="233" y="313" font-family="Arial, sans-serif" font-size="12" fill="#222">100-step</text>
  <text x="371" y="313" font-family="Arial, sans-serif" font-size="12" fill="#222">Best real</text>
  <text x="507" y="313" font-family="Arial, sans-serif" font-size="12" fill="#222">Oracle</text>
  <text x="640" y="313" font-family="Arial, sans-serif" font-size="12" fill="#222">Oracle train</text>
  <text x="108" y="331" font-family="Arial, sans-serif" font-size="11" fill="#666">strict</text>
  <text x="239" y="331" font-family="Arial, sans-serif" font-size="11" fill="#666">strict</text>
  <text x="370" y="331" font-family="Arial, sans-serif" font-size="11" fill="#666">value</text>
  <text x="502" y="331" font-family="Arial, sans-serif" font-size="11" fill="#666">value, gold plan</text>
  <text x="637" y="331" font-family="Arial, sans-serif" font-size="11" fill="#666">value, gold plan</text>
</svg>
<figcaption>The first three bars are production-style proxy runs. The last two are diagnostics that use labels derived from the reference SQL.</figcaption>
</figure>

## The ceiling test: Diagnosing the bottleneck

Once I hit that `0.640` ceiling, I decided to run a diagnostic test that you would never be allowed to run in production. 

I parsed the ground-truth reference SQL for every turn and extracted the exact tables, columns, joins, query skeleton, projection shape, and duplicate-row policy. I then fed these extracted labels directly to the model as planning hints.

This is what I call an **oracle run**—we are intentionally leaking the answer key. It's completely useless for production, but it's an incredible diagnostic tool. It isolates a very specific question: *If the schema planning was already perfect, how well could the model actually write the SQL?*

The results were eye-opening. A prompt-only oracle run hit `0.850` value accuracy, and a model specifically trained on that oracle-labelled format reached a massive **`0.890`**.

This doesn't mean the project is solved, but it proves something vital: the SQL generator itself is actually quite robust when it isn't burdened with schema selection and projection logic. The core missing piece isn't just better prompting; it's a dedicated planner that can make those structural decisions *without* peeking at the answer key.

## Formalizing the planner

Based on this insight, I restructured the pipeline to make the planner an explicit, measurable entity. Rather than asking the model to just "think step by step" in the prompt, we enforce a structured planning phase before any SQL is written.

Each evaluated turn now carries:
- A `gold_plan`, extracted from the reference SQL and used purely for training/diagnostics.
- A `predicted_plan`, generated from scratch without seeing the reference SQL.
- Distinct scoring metrics for how well the planner chose tables, columns, joins, aggregations, and grouping.

To establish a baseline, I threw together a quick, naive lexical planner. It simply reads the visible schema, looks at the user's question, and attempts to guess the tables and columns using raw string matching.

As expected, it struggles:

| Planner source | Turns | Dialogs | Oracle prompt rows | Macro score | Table F1 | Column F1 | Skeleton F1 |
| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
| `lexical_schema_baseline` | 100 | 32 | 0 | 0.571 | 0.599 | 0.117 | 0.648 |

<figure>
<svg viewBox="0 0 820 330" role="img" aria-labelledby="planner-title planner-desc" style="max-width:100%;height:auto;border:1px solid #ddd;background:#fff">
  <title id="planner-title">Planner baseline metrics on the fixed CoSQL slice</title>
  <desc id="planner-desc">Bar chart showing macro planner score 0.571, table F1 0.599, column F1 0.117, skeleton F1 0.648, join F1 0.790, group-by F1 0.830, and selected-count match 0.180.</desc>
  <rect x="0" y="0" width="820" height="330" fill="#fff"/>
  <text x="36" y="38" font-family="Arial, sans-serif" font-size="22" font-weight="700" fill="#111">The planner can now fail in specific ways</text>
  <text x="36" y="63" font-family="Arial, sans-serif" font-size="13" fill="#555">Lexical non-oracle planner. 100 CoSQL turns, 32 dialogs, zero oracle prompt rows.</text>
  <line x1="80" y1="270" x2="780" y2="270" stroke="#222"/>
  <line x1="80" y1="82" x2="80" y2="270" stroke="#222"/>
  <line x1="80" y1="232" x2="780" y2="232" stroke="#eee"/>
  <line x1="80" y1="194" x2="780" y2="194" stroke="#eee"/>
  <line x1="80" y1="157" x2="780" y2="157" stroke="#eee"/>
  <line x1="80" y1="119" x2="780" y2="119" stroke="#eee"/>
  <rect x="114" y="163" width="58" height="107" fill="#222"/>
  <rect x="211" y="157" width="58" height="113" fill="#444"/>
  <rect x="308" y="248" width="58" height="22" fill="#bdbdbd"/>
  <rect x="405" y="148" width="58" height="122" fill="#666"/>
  <rect x="502" y="122" width="58" height="148" fill="#888"/>
  <rect x="599" y="114" width="58" height="156" fill="#aaa"/>
  <rect x="696" y="236" width="58" height="34" fill="#d9d9d9" stroke="#222"/>
  <text x="116" y="154" font-family="Arial, sans-serif" font-size="13" font-weight="700" fill="#111">0.571</text>
  <text x="213" y="148" font-family="Arial, sans-serif" font-size="13" font-weight="700" fill="#111">0.599</text>
  <text x="310" y="239" font-family="Arial, sans-serif" font-size="13" font-weight="700" fill="#111">0.117</text>
  <text x="407" y="139" font-family="Arial, sans-serif" font-size="13" font-weight="700" fill="#111">0.648</text>
  <text x="504" y="113" font-family="Arial, sans-serif" font-size="13" font-weight="700" fill="#111">0.790</text>
  <text x="601" y="105" font-family="Arial, sans-serif" font-size="13" font-weight="700" fill="#111">0.830</text>
  <text x="698" y="227" font-family="Arial, sans-serif" font-size="13" font-weight="700" fill="#111">0.180</text>
  <text x="101" y="294" font-family="Arial, sans-serif" font-size="12" fill="#222">Macro</text>
  <text x="203" y="294" font-family="Arial, sans-serif" font-size="12" fill="#222">Table</text>
  <text x="294" y="294" font-family="Arial, sans-serif" font-size="12" fill="#222">Column</text>
  <text x="388" y="294" font-family="Arial, sans-serif" font-size="12" fill="#222">Skeleton</text>
  <text x="499" y="294" font-family="Arial, sans-serif" font-size="12" fill="#222">Join</text>
  <text x="587" y="294" font-family="Arial, sans-serif" font-size="12" fill="#222">Group-by</text>
  <text x="681" y="294" font-family="Arial, sans-serif" font-size="12" fill="#222">Count</text>
</svg>
<figcaption>The planner sometimes finds the rough table area, but column selection is still incredibly weak. This highlights exactly where our next engineering push needs to be.</figcaption>
</figure>

This decoupling is a huge win for observability. Now, if the planner picks the wrong table, we flag it as a planning failure. If the planner is spot-on but the SQL syntax is mangled, that's a generation failure. If both are perfect but execution still fails, we know we have an issue with value normalization or duplicate rows. By splitting these concerns, we avoid collapsing every error into a useless bucket of "bad SQL."

## The road ahead

With this groundwork laid, the path forward for our data-analysis claim is clear:

1. Continue iterating on the CoSQL proxy slice for tight local feedback loops.
2. Upgrade the naive planner to a robust, non-oracle system that can accurately predict tables, columns, joins, and value maps.
3. Feed the predicted plans into the SQL generator and formally score planning distinct from execution.
4. Establish baselines with heavy, hosted models on this exact same proxy slice.
5. Finally, migrate this entire evaluation contract over to the rigorous BIRD-Interact benchmark.

At this stage, the data engineering effort is just as critical as the model fine-tuning. Up next, I'll be building out a value index, entity-resolution labels, join/fan-out fixtures, and a schema validator to catch bad column references before they trigger generic SQL failures.

While the local 9B model definitely improved through fine-tuning, the real victory of this first phase is that our evaluation loop successfully highlighted what the system *cannot* do yet. 

To summarize our honest, current state:
- **`0.640`** value accuracy is our production-ready proxy baseline.
- **`0.890`** value accuracy is our oracle-assisted diagnostic ceiling.
- The planner baseline is finally measurable, but needs significant work.
- The definitive BIRD-Interact and hosted-model showdowns are up next.

A local 9B model might very well become a powerhouse for serious data analysis, but we have to build the evidence methodically: perfect the proxy, nail the planner, establish the baselines, and then tackle the final benchmark.

### Reference Material

- [Qwen 3.5 9B model card](https://huggingface.co/Qwen/Qwen3.5-9B)
- [BIRD-Interact project page](https://bird-interact.github.io/)
- [BIRD-Interact paper](https://arxiv.org/abs/2510.05318)
- [CoSQL project page](https://yale-lily.github.io/cosql)
- [CoSQL paper](https://arxiv.org/abs/1909.05378)
- [SParC paper](https://arxiv.org/abs/1906.02285)
- [Cube semantic layer docs](https://docs.cube.dev/docs/introduction)
- [dbt Semantic Layer docs](https://docs.getdbt.com/docs/build/semantic-models)


cosql1

Can a local 9B model compete with state-of-the-art systems for data analysis?

Recommended

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