bug: missing column with `SqlModel`

[zilto]

Summary

Running a dlt.hub.transformation that returns a dlt.Relation where the first row contains zeros incorrectly treats the columns as missing.

Repro

# repro.py
import dlt
from dlt.sources.rest_api import rest_api_resources
import dlthub.data_quality as dq

@dlt.source
def jaffleshop():
    jaffle_rest_resources: list = rest_api_resources({
        "client": {
            "base_url": "https://jaffle-shop.dlthub.com/api/v1",
            "paginator": {"type": "header_link"},
        },
        "resources": [  # individual resources
            "customers",
            "products",
            "orders",
        ],
        # set the time range for all resources
        "resource_defaults": {
            "endpoint": {
                "params": {
                    "start_date": "2017-01-01",
                    "end_date": "2017-01-15",
                },
            },
        },
    })
    
    return jaffle_rest_resources


@dlt.hub.transformation
def jaffle_checks(dataset: dlt.Dataset) -> dlt.Relation:
    checks = {
        "orders": [dq.checks.is_unique("id"), dq.checks.case("subtotal > 0")]
    }
    return dq.prepare_checks(dataset, checks=checks)


pipeline = dlt.pipeline("jaffle_dq", destination="duckdb")


if __name__ == "__main__":
    pipeline.run([jaffleshop()])

Output of running the code

❯ uv run python repro.py
DATE|[WARNING]|356876|137503479621440|dlt|validate.py|verify_normalized_table:100|In schema `jaffleshop`: The following columns in table `jaffle_checks` did not receive any data during this load:
  - success_count
  - success_rate

This can happen if you specify columns manually, for example, using the `merge_key`, `primary_key` or `columns` argument but they do not exist in the data.

If I change the if __name__ == "__main__": clause, I get a different result from what is materialized

if __name__ == "__main__":
    query = dq.prepare_checks(
        pipeline.dataset(),
        checks={
            "orders": [dq.checks.is_unique("id"), dq.checks.case("subtotal > 0")],
        },
    )
    print(query.arrow())

Output

❯ uv run python option1.py
pyarrow.Table
table_name: string
check_qualified_name: string
row_count: int64
success_count: int32
success_rate: double
----
table_name: [["orders","orders"]]
check_qualified_name: [["id__is_unique","subtotal__case__GT"]]
row_count: [[1570,1570]]
success_count: [[1570,1557]]
success_rate: [[1,0.9917197452229299]]

Suggested solution

I'm unsure about the origin of the bug. I wonder if we can simplify the lifecycle of dlt.Relation and SqlModel

[zilto]

I spent some time investigating the bug. Here are my cleaned up notes and takeaways.

Takeaways

• the bug is in ModelItemsNormalizer. It rewrites the SQL query of the dlt.Relation to exclude columns that are missing a data_type annotation
• the normalizer should do no query rewrite except adding _dlt_load_id and _dlt_id. All the necessary normalization is done via dlt.Relation
• the normalizer shouldn't do query validation, constraints, etc. This should be done by dlt.Relation.to_sql()

Notes

1. missing data_type

1. the initial bug / odd moment is that dlt.Relation.arrow() and pipeline.run() over a transformation yielding the same dlt.Relation gives different results
2. the pipeline.run() pointed to normalizer code in validate.py saying "did not receive any data during load"
   • side note: this message doesn't make sense for SQL queries
3. In validate.py, it checks that the dlt table schema contains "complete" columns. A complete column has name and data_type are set. In this case, data_type isn't set for the faulty columns
4. During pipeline.run(), the dlt table schema is generated via the DltResource.compute_table_schema() method.
5. In this case, it goes through the DltTransformationResource. This calls dlt.Relation.schema.
6. ❗ In the repro, we see that dq.prepare_checks(dataset, checks=checks).schema is also missing the data_type hints

This simplifies the repro and removes the @dlt.hub.transformation component.

It still doesn't explain the divergence.

2. why data_type is missing

1. dlt.Relation.schema calls lineage.compute_columns_schema() which processes sqlglot expressions (let's walk this function backwards)
2. there's no error / nothing raises in this code, but it produces "incomplete columns" in the output schema
3. ❗ In the fault case, from_sqlglot_type() returns an empty dictionary.
4. looking at from_sqlglot_type() comments indicate that the only condition for returning {} is when a column has sqlglot type UNKNOWN
   • @anuunchin found that modifying the sqlglot query improves the annotate_types() step and allows to resolve other types than UNKNOWN, thus avoiding the issue

It seems that compute_columns_schema() should have a parameter to raise if UNKNOWN types, i.e, "incomplete columns" are present in the output dlt table schema. It seems overly restrictive to require all types to be known for all cases.

This still doesn't explain the divergence

3. dlt.Relation.arrow() vs. pipeline.run()

For dlt.Relation.arrow(), the sqlglot expression is qualified, normalized, and compiled to a string. Then, the SQL string is directly passed to the dataset's destination's SQL client. For pipeline.run() there are a lot more moving pieces (assuming the same dlt.Relation)

1. Extract

1. the DltTransformationResource retrieves the dlt table schema from dlt.Relation.schema (we see the missing data_type)
2. the Extractor object does schema contract and schema evolution
3. the ModelWriter.write_data() receives the dlt.Relation and creates a .model file that contains the dlt.Relation.query_dialect and dlt.Relation.to_sql()

Note

Contrary to @dlt.resource, the relation / SQL schema is known before execution

2. Normalize

1. the .model file is loaded
2. ❗ ModelItemsNormalizer.__call__() rewrites the SQL string. The main difference is an outer wrapper that adds the _dlt_load_id. Notably, it doesn't include the faulty columns that don't have a data_type
   • This is the first source of divergence from dlt.Relation.

   SELECT 
     _dlt_subquery."table_name" AS "table_name",
     _dlt_subquery."check_qualified_name" AS  "check_qualified_name",
     _dlt_subquery."row_count" AS "row_count",
     '1764114959.071344' AS "_dlt_load_id"
   FROM ($RELATION_QUERY) AS _dlt_subquery

3. .__call__() uses schema.get_table_columns(root_table_name) in several areas, which excludes "incomplete columns" by default
4. the ModelWriter receives the edited query string and creates .model file
5. ❗validate.py reads the dlt table schema from Extract step and complains about incomplete columns. The only schema update is adding _dlt_load_id (i.e., data_type is still missing)

Note

Looking at step 2. it's obvious why some columns are missing on the destination table. The query is rewritten to exclude them

[anuunchin]

Looking at step 2. it's obvious why some columns are missing on the destination table. The query is rewritten to exclude them the query here is meant to exclude them because the normalizer is meant to follow the dlt schema -> including incomplete columns will cause all kinds of weird behaviour, such as a binder error from the destination in the test here if you add pipeline.load()

The normalizer is not including incomplete columns, because dlt can't materialize these columns in the destination if it doesn't know/can't infer the type -> which is happening here

[anuunchin]

It seems that compute_columns_schema() should have a parameter to raise if UNKNOWN types I agree with this one

[zilto]

Let's take a step back from how the current code works.

Here's my mental model:

input dataset == output dataset

1. call pipeline.run() over a resource, you get data in a dlt.Dataset. This represents real data on a destination
2. from the dataset, you create a dlt.Relation. This relation may represent existing data exactly as-is, or be an operation on top of it.
   • Key point: the relation is backed by an SQL expression (e.g., SELECT statement) which always binds to physical data on the destination
3. If I do relation.fetchall():
   1. execute the relation's SQL against the destination
   2. return the data in memory
4. If I do pipeline.run([relation]) (i.e., trying to simplify what a transformation is):
   1. normalize: wrap the query (i.e., SELECT is wrapped by INSERT FROM and adds _dlt_load_id)
   2. execute the SQL against the destination
   3. write the result to the destination

Note

The dialect, naming convention, etc. are the same for the dlt.Relation as for the pipeline.run(). We should catch things as early as possible if there are issues with the query.

input dataset != output dataset, same destination

4. If I do pipeline.run([relation])
   1. normalize: wrap the query and set output dataset name (i.e., selecting from bronze_ and writing to silver_
   2. execute the SQL against the destination
   3. write the result to the destination

Note

The dialect, naming convention, etc. are the same for the dlt.Relation as for the pipeline.run(). We should catch things as early as possible if there are issues with the query.

input dataset != output dataset, different destination

4. If I do pipeline.run([relation])
   1. execute the SQL query via relation.iter_arrow()
   2. process arrow batches like a regular @dlt.resource

Note

In this case, there's never a reason to have any normalize or query rewrite. We don't need schema inference on the dlt.Relation either because this will be derived from the pyarrow.Table

(detailed scenarios and diagram)

Moving forward

the query here is meant to exclude them because the normalizer is meant to follow the dlt schema -> including incomplete columns will cause all kinds of weird behaviour, such as a binder error from the destination in the test here if you add pipeline.load()

What I understand from this and my experience: SQL queries that can be executed directly (or via dlt.Relation) might not be executable via pipeline.run(). This gives a bad user experience because developing and testing outside the pipeline (e.g., in a notebook) != running the pipeline.

If dlt has additional constraints about the SQL for schema evolution, etc. it should communicate loud, explicitly, and early to users.

Some options:

1. let the SQL query / dlt.Relation run via pipeline.run() but disable some dlt features + user warning (e.g., schema evolution):
   • pros: dlt.Relation and pipeline.run() always match; dlt.Relation remains flexible
   • cons: break downstream assumptions for dlt, need to handle "resources" and "SQL models" in different code paths (maybe they are different enough to justify it?)
2. be more restrictive on dlt.Relation such that it always matches pipeline.run().
   • pros: can simplify normalizer and transformation code by delegating everything to dlt.Relation
   • cons: makes data exploration worse (can't run all SQL queries), unecessary coupling

[anuunchin]

Resolved upstream in dlthub