Configure arguments in assets decorator that could be used in asset_checks.

[marcilj]

I found it hard to understand why the asset check feature is so limitating.
In my mind, most (90%) of the checks for all assets are similar and then you can create some that target something very unique that you want to validate for a specific asset. But in general, once I would build a asset checks I'll probably apply it to all my assets. For example, nulls, duplicates, difference between previous value, quantity exceeds xyz, increase from a median value, etc.

It would be nice if we could configure fields close to the data in our @asset decorator that would be given to an asset checks.

What I mean by this is instead of generating a list of things and providing that list to a factory like describe here .
Same thing for the factory pattern that you have defined in the doc

def make_checks(
    check_blobs: Sequence[Mapping[str, str]]
) -> Sequence[AssetChecksDefinition]:
    checks = []
    for check_blob in check_blobs:

        @asset_check(name=check_blob["name"], asset=check_blob["asset"])
        def _check(context):
            db_connection = ...
            rows = db_connection.execute(check_blob["sql"])
            return AssetCheckResult(
                passed=len(rows) == 0, metadata={"num_rows": len(rows)}
            )

        checks.append(_check)

    return checks


check_blobs = [
    {
        "name": "orders_id_has_no_nulls",
        "asset": "orders",
        "sql": "select * from orders where order_id is null",
    },
    {
        "name": "items_id_has_no_nulls",
        "asset": "items",
        "sql": "select * from items where item_id is null",
    },
]

We could instead configure the list of checks with the necessary parameters in the asset decorator.

@asset(
    op_tags={"kind": "python"},
    io_manager_key="warehouse_io_manager",
    required_resource_keys={"snowflake_client", "aws_secrets"},
    metadata={
        "priority": "high",
        "partition_expr": "time",
    },
    run_checks=[
        is_unique_key_unique(column_to_validate="unique_id"),
        contain_nulls(column_to_validate="time"),
    ]

)
def my_asset(context: OpExecutionContext):

@asset_check()
def is_unique_key_unique(asset: df.DataFrame, column_to_validate: str):
    num_duplicates = asset.duplicated(subset=[column_to_validate]).values.sum()

    return AssetCheckResult(
        passed=bool(num_duplicates == 0),
        check_name="is_unique_key_unique",
        metadata={
            "num_duplicates": int(num_duplicates),
        },
        severity=AssetCheckSeverity.ERROR,
    )


def count_nulls(df: df.DataFrame, column: str):
    return df[f"{column}"].isna().sum()

@asset_check()
def have_nulls(asset: df.DataFrame, column_to_validate: str):
    num_nulls = count_nulls(asset, column_to_validate)

    return AssetCheckResult(
        passed=bool(num_nulls == 0),
        check_name="have_nulls",
        metadata={
            "num_nulls: int(num_nulls),
        },
        severity=AssetCheckSeverity.ERROR,
    )

Or something like a dict or key (That represent the asset check name to run) and the value that represent the parameter to look at.

@asset(
    run_checks={
        is_unique_key_unique: "unique_id",
        contain_nulls: "time",
    }
)

This would be similar to the feature that Dagster already have for dagster called "partition_expr" where you can define the field that would be used for the partition in the metadata dict.

    metadata={
        "partition_expr": "partition_timestamp",
    },

For most of the use-case that I found with the asset check feature, having them configure directly in the function doesn't really make sense, as it will grow larger and large each time I want to validate something.

https://docs.dagster.io/concepts/assets/asset-checks#checks-that-execute-in-the-same-op-that-materializes-the-asset

@asset(check_specs=[AssetCheckSpec(name="orders_id_has_no_nulls", asset="orders")])
def orders(context: AssetExecutionContext):
    orders_df = pd.DataFrame({"order_id": [1, 2], "item_id": [432, 878]})

    # save the output and indicate that it's been saved
    orders_df.to_csv("orders")
    yield Output(value=None)

    # check nulls
    num_null_order_ids = orders_df["order_id"].isna().sum()
    yield AssetCheckResult(
        passed=bool(num_null_order_ids == 0),
    )

    # check duplicate
    num_duplicates = orders_df.duplicated(subset=["order_id"]).values.sum()

    yield AssetCheckResult(
        passed=bool(num_duplicates == 0),
        check_name="is_unique_key_unique",
        metadata={
            "num_duplicates": int(num_duplicates),
        },
        severity=AssetCheckSeverity.ERROR,
    )

    # check something else
    yield AssetCheckResult()

    # check something else again
    yield AssetCheckResult()

Lastly, the factory solution works but isn't really scallable IMO, since we would endup with a big list that we need to keep around and manage instead of specifying the checks at the asset levels.

If I have 100 assets where I want to validate nulls for, I would need to generate a list of 100 elements in the check_blobs.

Instead configuring the different elements closer to the asset would make more sense IMO.

Since in the configs, the only thing that changes are the asset name that could be provided directly by the function name decorated with @asset and the sql query.

    {
        "name": "orders_id_has_no_nulls",
        "asset": "orders",
        "sql": "select * from orders where order_id is null",
    },
    {
        "name": "items_id_has_no_nulls",
        "asset": "items",
        "sql": "select * from items where item_id is null",
    },

In that example, the table and where condition are different. But since we can leverage the IO Manager instead of writing the query directly, we could replace the sql query by a parameter in that case the name of the columns and pass that as a parameter to the @asset_check function.

I hope I was mostly clear in my proposition, don't hesitate to ask question or tell me which feature that's already available that I missed.

Have a great day.

[johannkm]

Thank you for this write up!

The factory you propose is similar to something we considered initially when designing checks. We liked how easy it was to add checks across your assets, but we were worried about parameterizing. The current factory pattern is maximally flexible for how you build your check based on whatever input you chose. With the other approach we'd have to provide the the apis to parameterize them, and they could get more complicated as we improve checks (e.g. by adding support for checks on multiple assets).

I wonder if there's a way to achieve some of the ergonomics you want with the current approach. Re:

configuring the different elements closer to the asset would make more sense

I'm curious how you'd feel about something like this:

def create_is_key_unique_check(target_asset: AssetsDefinition, column_to_validate: str):
      
  @asset_check(asset=target_asset, name=f"{column_to_validate}_is_unique")
  def is_key_unique(asset_to_check: df.DataFrame):
    num_duplicates = asset_to_check.duplicated(subset=[column_to_validate]).values.sum()
    return AssetCheckResult(passed=bool(num_duplicates == 0))

  return is_key_unique

where you'd use it next to each asset:

@asset
def my_asset(): ...

check_my_asset_id = create_is_key_unique_check(target_asset=my_asset, column_to_validate='id')

@asset
def my_other_asset(): ...

check_my_other_asset_id = create_is_key_unique_check(target_asset=my_other_asset, column_to_validate='id')
check_my_other_asset_other_column = create_is_key_unique_check(target_asset=my_other_asset, column_to_validate='other_column')

and load both of these with

defs = Definitions(
  assets=load_assets_from_current_module(),
  asset_checks=load_asset_checks_from_current_module(),
)

[marcilj]

The solution you gave me is quite interesting.

That would allow me to configure the assets in a simple way, closer to the asset, and I could even create groups of checks the same way and apply a bunch at the same time.

Thank you for your quick response, but I like this option.

I'll try it.

Thank you very much.

[bsarden]

I gave this a shot this morning as well, but for some reason the asset_check instances are not being picked up by load_asset_checks_from_current_module()

from dagster import AssetsDefinition, asset_check, AssetCheckResult, AssetCheckSeverity, asset

import polars as pl

def create_is_key_unique_check(target_asset: AssetsDefinition, column_to_validate: str):

    @asset_check(asset=target_asset)
    def _check(asset_to_check: pl.DataFrame):
        breakpoint()
        num_duplicates = len(asset_to_check.select(
          pl.col(column_to_validate).is_duplicated().alias("is_duplicated")
        ).filter(
          pl.col("is_duplicated") == True
        ))
        return AssetCheckResult(
          passed=bool(num_duplicates == 0),
          check_name="is_key_unique",
          metadata={
            "num_duplicates": num_duplicates,
          },
          severity=AssetCheckSeverity.ERROR,
        )
    print(f"Created asset check")

    return _check
    
@asset
def my_asset() -> pl.DataFrame:
    return pl.DataFrame(
        { "my_col": [1, 2, 3, 4, 5] }
    )
 
 create_is_key_unique_check(target_asset='my_asset', column_to_validate='my_col') 

But the assets are not discovered.

from dagster import load_asset_checks_from_package_module
import my_module

my_asset_checks = load_asset_checks_from_package_module(package_module=my_module)

# my_asset_checks = []

The assets are properly discovered when using the regular @asset_check decorator without the factory abstraction. Maybe this is a PEBCAK issue? Thanks in advance!

[johannkm]

@bsarden this was an issue in my example. For load_asset_checks_from_* to pick up the check it needs to be assigned to a variable (the name doesn't matter). I've updated the example

[johannkm]

Also updated to set @asset_check(name=... to avoid conflicts