Feature/validate upstream tables by STEFANOVIVAS · Pull Request #1309 · databrickslabs/dqx

STEFANOVIVAS · 2026-07-07T02:52:13Z

Changes

Added new check function: aggr_matches_dataset— a new dataset-level check that compares an aggregate metric (row count by default, or any curated/built-in aggregate) computed on the checked DataFrame against the same aggregate computed on a reference (upstream) table or DataFrame, with optional abs_tolerance/rel_tolerance and independent row_filter/ref_row_filter on each side.

Linked issues

Resolves #1046

Tests

Documentation and Demos

added/updated demos
added/updated docs
added/updated agent skills

CLAassistant · 2026-07-07T02:52:21Z

All committers have signed the CLA.

Copilot

Pull request overview

This PR introduces a new dataset-level quality check, validate_upstream_table, to validate ingestion correctness by comparing an aggregate metric (row count by default, or other curated/built-in aggregates) computed on the checked DataFrame against the same metric computed on an upstream reference (DataFrame via ref_df_name or Unity Catalog table via ref_table), with optional absolute/relative tolerances and independent row filters.

Changes:

Added validate_upstream_table to check_funcs.py, including parameter validation, aggregate evaluation on both sides, and tolerance/null-safe comparison logic.
Added unit and integration tests covering parameter validation, tolerance behavior, row filtering, '*'/count(*) usage, and ref-table vs ref-DataFrame paths.
Documented the new check in the quality checks reference, including YAML and Python usage examples.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
`src/databricks/labs/dqx/check_funcs.py`	Adds the new dataset-level upstream aggregate comparison check implementation.
`tests/unit/test_dataset_checks.py`	Adds rule-build-time parameter validation tests for the new check.
`tests/integration/test_dataset_checks.py`	Adds end-to-end Spark integration tests for the new check’s runtime behavior.
`docs/dqx/docs/reference/quality_checks.mdx`	Documents the new check and adds usage examples.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

ghanse

Headed in the right direction. Left some comments. Do you also want to add a similar check for a custom SQL query?

mwojtyczka

Review of validate_upstream_table. Useful check and largely correct; main themes are: it reimplements existing aggregate machinery (_is_aggr_compare + _get_ref_df) rather than composing it, and a few edges around filter validation, the empty-upstream message, and the exact-equality default. Inline below.

Cleared: rel_tolerance has no div-by-zero (guarded by max()); the apply closure captures spark/ref_dfs correctly (standard dataset-check pattern); .limit(1) matches _is_aggr_compare; behavioral coverage lives in the integration tests (dataset checks can't be unit-tested without Spark).

mwojtyczka

Going in the right direction, left some comments

…ng functions as _is_aggr_compare and _get_ref_df, avoiding re-implementing existing logic

…dataset

mwojtyczka

LGTM. Just one request to rename the function to aggr_matches_dataset to make it more generic and update docs and docs strings accordingly

STEFANOVIVAS · 2026-07-16T16:51:17Z

LGTM. Just one request to rename the function to aggr_matches_dataset to make it more generic and update docs and docs strings accordingly

Hi, @mwojtyczka! One thing to point out: Should we implement a group by clause in this check? _is_aggr_compare function already let us apply it in the input Dataframe. Could make sense because in many situations we need to make a per-group reconciliation (date, region,source_system,etc).

mwojtyczka · 2026-07-16T17:01:01Z

LGTM. Just one request to rename the function to aggr_matches_dataset to make it more generic and update docs and docs strings accordingly

Hi, @mwojtyczka! One thing to point out: Should we implement a group by clause in this check? _is_aggr_compare function already let us apply it in the input Dataframe. Could make sense because in many situations we need to make a per-group reconciliation (date, region,source_system,etc).

yes makes sense

mwojtyczka · 2026-07-22T17:28:57Z

@STEFANOVIVAS pls also add the check to tests/resources/all_dataset_checks.yaml, and adjust the relevant tests. We also need a performance test so that we can validate regressions in the future

mwojtyczka · 2026-07-22T17:50:33Z

Code review — `aggr_matches_dataset`

Overall this is a high-quality, well-tested addition (23 integration tests + unit exception cases, careful null-safety handling, behavior-preserving refactor of _is_aggr_compare). No blocking bugs found in the happy paths. Notes below are minor correctness edge-cases and test-coverage gaps.

Note: the non-null-safe grouped-join issue I initially flagged for count_distinct + group_by is already being fixed in #1356 (it replaces the df.join(agg_df, on=join_cols, how="left") in _is_aggr_compare with a null-safe helper). Once #1356 lands, aggr_matches_dataset's grouped count_distinct path inherits the fix — so it's excluded here.

1. Not wired into the "all dataset checks" coverage (test-coverage)

tests/resources/all_dataset_checks.yaml (+ the all-checks integration tests) — aggr_matches_dataset isn't included, unlike sibling dataset checks (compare_datasets, foreign_key, …). The declarative metadata / YAML round-trip for this check is therefore unexercised.

2. No `aggr_matches_dataset` test for `count_distinct` with `group_by` (test-coverage)

tests/integration/test_dataset_checks.py — the check is tested with count_distinct (no group_by) and with group_by (count/sum), but not both together. #1356 adds a grouped-count_distinct regression for the base is_aggr_* path; a small aggr_matches_dataset-specific case (ideally with a NULL group key, to lock in the #1356 fix for this check) would close the loop.

mwojtyczka

Code looks good, left comments regading tests

- Wire aggr_matches_dataset into the all-dataset-checks coverage: add it to all_dataset_checks.yaml and both all-checks integration tests (test_apply_checks_all_checks_as_yaml already loads the YAML; test_apply_checks_all_checks_using_classes gains the class-based check plus ref_dfs), exercising the declarative/YAML round-trip like the sibling dataset checks. - Add count_distinct + group_by integration tests (match, mismatch, and a NULL group key) to lock in grouped distinct-count behaviour for this check; the NULL-key case asserts flagging behaviour rather than exact message text since the window-incompatible join is not null-safe (documented). - Add performance benchmarks: dataset-wide count comparison and the count_distinct + group_by two-stage join path. Co-authored-by: Isaac

mwojtyczka · 2026-07-23T10:08:00Z

Code review — aggr_matches_dataset

Overall this is a high-quality, well-tested addition (23 integration tests + unit exception cases, careful null-safety handling, behavior-preserving refactor of _is_aggr_compare). No blocking bugs found in the happy paths. Notes below are minor correctness edge-cases and test-coverage gaps.

Note: the non-null-safe grouped-join issue I initially flagged for count_distinct + group_by is already being fixed in #1356 (it replaces the df.join(agg_df, on=join_cols, how="left") in _is_aggr_compare with a null-safe helper). Once #1356 lands, aggr_matches_dataset's grouped count_distinct path inherits the fix — so it's excluded here.

1. Not wired into the "all dataset checks" coverage (test-coverage)

tests/resources/all_dataset_checks.yaml (+ the all-checks integration tests) — aggr_matches_dataset isn't included, unlike sibling dataset checks (compare_datasets, foreign_key, …). The declarative metadata / YAML round-trip for this check is therefore unexercised.

2. No aggr_matches_dataset test for count_distinct with group_by (test-coverage)

tests/integration/test_dataset_checks.py — the check is tested with count_distinct (no group_by) and with group_by (count/sum), but not both together. #1356 adds a grouped-count_distinct regression for the base is_aggr_* path; a small aggr_matches_dataset-specific case (ideally with a NULL group key, to lock in the #1356 fix for this check) would close the loop.

I fixed these

mwojtyczka

LGTM

mwojtyczka

LGTM

# Conflicts: # src/databricks/labs/dqx/check_funcs.py

The MCP suite (tests/integration_mcp) only ran on PRs touching mcp-server/, src/databricks/labs/dqx/, or the mcp workflow — so nightly never exercised it. Add an mcp-tests job to nightly.yml mirroring the mcp.yml PR job (setup-env, prebuild-wheel, install Databricks CLI, acceptance harness with the integration_mcp codegen), minus the PR-only fork gate since nightly is schedule-triggered. No coverage upload (the MCP server/runner execute remotely). Co-authored-by: Isaac

mwojtyczka · 2026-07-23T14:41:05Z

Heads-up: I pushed a small CI addition to this branch (96ef0896) — an mcp-tests job in .github/workflows/nightly.yml so the MCP integration suite runs on the nightly schedule.

It's unrelated to the aggr_matches_dataset change here — it's a gap we missed when the MCP server landed: mcp.yml only triggers on PRs touching mcp-server/, so nightly never exercised the suite. Bundling the fix here for expediency.

implemented requested changes

STEFANOVIVAS added 2 commits July 6, 2026 23:22

create validate-upstream_table check and apply tests

8999fe2

update test_dataset_check.py file

133e537

STEFANOVIVAS requested a review from a team as a code owner July 7, 2026 02:52

STEFANOVIVAS requested review from pratikk-databricks and removed request for a team July 7, 2026 02:52

Merge branch 'main' into feature/validate_upstream_tables

0e78738

ghanse requested a review from Copilot July 7, 2026 17:49

Copilot started reviewing on behalf of ghanse July 7, 2026 17:50 View session

Copilot AI reviewed Jul 7, 2026

View reviewed changes

Comment thread src/databricks/labs/dqx/check_funcs.py Outdated

ghanse previously requested changes Jul 7, 2026

View reviewed changes

Comment thread src/databricks/labs/dqx/check_funcs.py Outdated

Comment thread src/databricks/labs/dqx/check_funcs.py Outdated

ghanse added under-review This PR is currently being reviewed by one of DQX maintainers. needs-changes Changes required after review labels Jul 7, 2026

mwojtyczka reviewed Jul 12, 2026

View reviewed changes

Comment thread docs/dqx/docs/reference/quality_checks.mdx Outdated

mwojtyczka reviewed Jul 12, 2026

View reviewed changes

Comment thread src/databricks/labs/dqx/check_funcs.py Outdated

mwojtyczka reviewed Jul 12, 2026

View reviewed changes

Comment thread src/databricks/labs/dqx/check_funcs.py Outdated

mwojtyczka requested changes Jul 12, 2026

View reviewed changes

STEFANOVIVAS added 3 commits July 13, 2026 22:07

refactoring validate_upstream_table check to take advantage of existi…

b7d488b

…ng functions as _is_aggr_compare and _get_ref_df, avoiding re-implementing existing logic

renaming check from validate_upstream_table to aggr_matches_upstream_…

57bcb91

…dataset

Add Pylint minor modifications

019f136

STEFANOVIVAS requested review from ghanse and mwojtyczka July 14, 2026 01:10

Merge branch 'main' into feature/validate_upstream_tables

a140cee

mwojtyczka approved these changes Jul 16, 2026

View reviewed changes

mwojtyczka and others added 3 commits July 16, 2026 23:51

Merge branch 'main' into feature/validate_upstream_tables

3f6ec68

Merge branch 'main' into feature/validate_upstream_tables

7f8c094

Add group by clause to aggr_matches_upstream_dataset

145684c

mwojtyczka requested changes Jul 22, 2026

View reviewed changes

mwojtyczka added 2 commits July 23, 2026 00:53

Merge branch 'main' into feature/validate_upstream_tables

11d68e8

mwojtyczka approved these changes Jul 23, 2026

View reviewed changes

mwojtyczka added Approved to Merge When PR is reviewed and approved. To be merged once all tests pass and removed under-review This PR is currently being reviewed by one of DQX maintainers. labels Jul 23, 2026

Merge remote-tracking branch 'origin/main' into HEAD

c1166dc

# Conflicts: # src/databricks/labs/dqx/check_funcs.py

mwojtyczka mentioned this pull request Jul 23, 2026

Fork test: PR #1309 - Feature/validate upstream tables #1367

Closed

mwojtyczka had a problem deploying to tool July 23, 2026 11:01 — with GitHub Actions Failure

mwojtyczka temporarily deployed to tool July 23, 2026 11:01 — with GitHub Actions Inactive

mwojtyczka temporarily deployed to tool July 23, 2026 14:40 — with GitHub Actions Inactive

mwojtyczka merged commit e9d2cca into databrickslabs:main Jul 23, 2026
42 checks passed

Uh oh!

Conversation

STEFANOVIVAS commented Jul 7, 2026 • edited by mwojtyczka Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

Linked issues

Tests

Documentation and Demos

Uh oh!

CLAassistant commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

ghanse left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mwojtyczka left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mwojtyczka left a comment

Choose a reason for hiding this comment

Uh oh!

mwojtyczka left a comment

Choose a reason for hiding this comment

Uh oh!

STEFANOVIVAS commented Jul 16, 2026

Uh oh!

mwojtyczka commented Jul 16, 2026

Uh oh!

mwojtyczka commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mwojtyczka commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code review — aggr_matches_dataset

1. Not wired into the "all dataset checks" coverage (test-coverage)

2. No aggr_matches_dataset test for count_distinct with group_by (test-coverage)

Uh oh!

mwojtyczka left a comment

Choose a reason for hiding this comment

Uh oh!

mwojtyczka commented Jul 23, 2026

Code review — aggr_matches_dataset

1. Not wired into the "all dataset checks" coverage (test-coverage)

2. No aggr_matches_dataset test for count_distinct with group_by (test-coverage)

Uh oh!

mwojtyczka left a comment

Choose a reason for hiding this comment

Uh oh!

mwojtyczka left a comment

Choose a reason for hiding this comment

Uh oh!

mwojtyczka commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

STEFANOVIVAS commented Jul 7, 2026 •

edited by mwojtyczka

Loading

CLAassistant commented Jul 7, 2026 •

edited

Loading

mwojtyczka commented Jul 22, 2026 •

edited

Loading

mwojtyczka commented Jul 22, 2026 •

edited

Loading

Code review — `aggr_matches_dataset`

2. No `aggr_matches_dataset` test for `count_distinct` with `group_by` (test-coverage)

Code review — `aggr_matches_dataset`

2. No `aggr_matches_dataset` test for `count_distinct` with `group_by` (test-coverage)

mwojtyczka commented Jul 23, 2026 •

edited

Loading