Feature/validate upstream tables

src/databricks/labs/dqx/check_funcs.py

@@ -2061,6 +2061,180 @@ def apply(df: DataFrame) -> DataFrame:
     return condition, apply
 
 
+@register_rule("dataset")
+def aggr_matches_dataset(
+    column: str | Column,
+    ref_table: str | None = None,
+    ref_df_name: str | None = None,
+    ref_column: str | Column | None = None,
+    aggr_type: str = "count",
+    aggr_params: dict[str, Any] | None = None,
+    group_by: list[str | Column] | None = None,
+    ref_group_by: list[str | Column] | None = None,
+    row_filter: str | None = None,
+    ref_row_filter: str | None = None,
+    abs_tolerance: float | None = None,
+    rel_tolerance: float | None = None,
+) -> tuple[Column, Callable]:
+    """
+    Build an upstream table comparison check condition and closure for dataset-level validation.
+
+    This function verifies that an aggregation on a column in the checked DataFrame matches the same
+    aggregation computed on a reference (upstream) DataFrame or table. It is commonly used to validate
+    that a row count (or other aggregate metric) in a downstream table matches its upstream source,
+    catching data loss or duplication introduced during ingestion.
+
+    Args:
+        column: Column name (str) or Column expression to aggregate in the checked DataFrame.
+            Pass *"*"* for *count(*)* over all rows.
+        ref_table: Name of the reference (upstream) table to read from the catalog.
+        ref_df_name: Name of the reference (upstream) DataFrame (used when passing DataFrames directly).
+        ref_column: Column name (str) or Column expression to aggregate in the reference DataFrame.
+            Defaults to *column* when not provided.
+        aggr_type: Aggregation type (default: 'count'). Curated types include count, sum, avg, min, max,
+            count_distinct, stddev, percentile, and more. Any Databricks built-in aggregate is supported.
+        aggr_params: Optional dict of parameters for aggregates requiring them (e.g., percentile value for
+            percentile functions, accuracy for approximate aggregates). Parameters are passed as keyword
+            arguments to the Spark function.
+        group_by: Optional list of column names or Column expressions in the checked DataFrame to compare
+            the aggregate per group instead of dataset-wide. Only simple column expressions are supported,
+            e.g. *F.col("region")*. A group present in the checked DataFrame but absent from the reference
+            is reported as a mismatch; groups present only in the reference are not surfaced. Note that
+            when *aggr_type* is a window-incompatible aggregate (e.g. *count_distinct*), the checked-side
+            grouping join is not null-safe, so a legitimately null group key may under-report.
+        ref_group_by: Optional list of group-by columns on the reference (upstream) side, matched to
+            *group_by* by position. Defaults to *group_by* when omitted. Must have the same length as
+            *group_by*. Requires *group_by* to be set.
+        row_filter: Optional SQL expression to filter rows in the checked DataFrame before aggregation.
+            Auto-injected from the check filter.
+        ref_row_filter: Optional SQL expression to filter rows in the reference DataFrame or table before
+            aggregation (e.g. to align both sides on the same date partition).
+        abs_tolerance: Values are considered equal if the absolute difference is less than or equal to the
+            tolerance. This is applicable to numeric aggregates.
+        rel_tolerance: Relative tolerance for numeric comparisons. Differences within this relative tolerance
+            are ignored. Useful if the aggregates vary in scale. Because it compares two separately-computed
+            aggregates, sum/avg over floating-point columns can differ across runs/clusters (non-associative
+            summation) even for identical data. A small rel_tolerance value is recommended for these situations.
+
+    Returns:
+        A tuple of:
+            - A Spark Column representing the condition for upstream comparison violations.
+            - A closure that applies the upstream comparison check and adds the necessary condition/metric
+              columns.
+
+    Raises:
+        MissingParameterError:
+            - if neither *ref_df_name* nor *ref_table* is provided.
+        InvalidParameterError:
+            - if both *ref_df_name* and *ref_table* are provided.
+            - if *abs_tolerance* or *rel_tolerance* is negative.
+            - if *ref_group_by* is provided without *group_by*.
+            - if *group_by* and *ref_group_by* lengths differ.
+    """
+    if ref_df_name and ref_table:
+        raise InvalidParameterError(
+            "Both 'ref_df_name' and 'ref_table' were provided. Please provide only one to avoid ambiguity."
+        )
+    if not ref_df_name and not ref_table:
+        raise MissingParameterError("Either 'ref_df_name' or 'ref_table' is required but neither was provided.")
+
+    if ref_group_by and not group_by:
+        raise InvalidParameterError(
+            "'ref_group_by' was provided without 'group_by'. Please provide 'group_by' for the checked "
+            "DataFrame as well."
+        )
+    group_by_names: list[str] | None = None
+    ref_group_by_names: list[str] | None = None
+    if group_by:
+        resolved_ref_group_by = group_by if ref_group_by is None else ref_group_by
+        if len(group_by) != len(resolved_ref_group_by):
+            raise InvalidParameterError(
+                f"'group_by' has {len(group_by)} entries but 'ref_group_by' has {len(resolved_ref_group_by)}. "
+                "Both must have the same length to allow comparison."
+            )
+        group_by_names = get_columns_as_strings(group_by, allow_simple_expressions_only=True)
+        ref_group_by_names = get_columns_as_strings(resolved_ref_group_by, allow_simple_expressions_only=True)
+
+    ref_column = column if ref_column is None else ref_column
+    _, ref_col_str, ref_col_expr = get_normalized_column_and_expr(ref_column)
+    ref_label = f"table '{ref_table}'" if ref_table else f"DataFrame '{ref_df_name}'"
+
+    unique_str = uuid.uuid4().hex  # make sure any column added to the dataframe is unique
+    ref_metric_col = f"__ref_metric_{aggr_type}_{unique_str}"
+
+    # The reference aggregate isn't known yet (it requires reading ref_table/ref_dfs, only available in
+    # apply()), so it's passed to _is_aggr_compare as the *name* of a column that will exist on df once
+    # the crossJoin/join below runs, rather than as a literal value or a Column tied to ref_df's own lineage
+    # (Spark can't resolve a Column across two unrelated DataFrames without a join).
+    condition, aggr_apply = _is_aggr_compare(
+        column=column,
+        limit=ref_metric_col,
+        aggr_type=aggr_type,
+        aggr_params=aggr_params,
+        group_by=group_by,
+        row_filter=row_filter,
+        compare_op=py_operator.ne,
+        compare_op_label=f"not equal to {ref_label} column '{ref_col_str}'",
+        compare_op_name="not_equal_to_upstream",
+        abs_tolerance=abs_tolerance,
+        rel_tolerance=rel_tolerance,
+        null_safe_limit_compare=True,
+    )
+
+    def apply(df: DataFrame, spark: SparkSession, ref_dfs: dict[str, DataFrame]) -> DataFrame:
+        """
+        Apply the upstream comparison check logic to the DataFrame.
+
+        Computes the aggregation on the reference (upstream) DataFrame or table. When *group_by* is not
+        set, the aggregate is a single scalar joined onto every row via *crossJoin*. When *group_by* is
+        set, the reference aggregate is computed per *ref_group_by* group and null-safe joined onto the
+        checked DataFrame by matching *group_by*/*ref_group_by* keys (groups missing on the reference side
+        yield a null reference metric, which is treated as a mismatch). Either way, the actual comparison
+        (including curated-aggregate type validation and null-safe matching) is delegated to
+        *_is_aggr_compare*.
+
+        Args:
+            df: The input DataFrame to validate.
+            spark: SparkSession used if reading a reference table.
+            ref_dfs: Dictionary of reference DataFrames (by name), used to resolve *ref_df_name*.
+
+        Returns:
+            The DataFrame with additional condition and metric columns for upstream comparison.
+        """
+        ref_df = _get_ref_df(ref_df_name, ref_table, ref_dfs, spark)
+        ref_filtered_df = ref_df.filter(ref_row_filter) if ref_row_filter else ref_df
+        ref_aggr_expr = _build_aggregate_expression(aggr_type, ref_col_expr, aggr_params)
+
+        if group_by_names and ref_group_by_names:
+            ref_group_cols = [F.col(col) for col in ref_group_by_names]
+            ref_metric_df = ref_filtered_df.groupBy(*ref_group_cols).agg(ref_aggr_expr.alias(ref_metric_col))
+
+            if aggr_type not in CURATED_AGGR_FUNCTIONS:
+                _validate_aggregate_return_type(ref_metric_df, aggr_type, ref_metric_col)
+
+            # groups only present in the reference are intentionally not surfaced: this check, like every
+            # other dataset check in this module, only annotates rows that exist in the checked DataFrame.
+            joined = _match_rows(
+                df.alias("df"),
+                ref_metric_df.alias("ref_df"),
+                group_by_names,
+                ref_group_by_names,
+                check_missing_records=False,
+                null_safe_row_matching=True,
+            )
+            joined = joined.select("df.*", F.col(f"ref_df.{ref_metric_col}").alias(ref_metric_col))
+            return aggr_apply(joined)
+
+        ref_metric_df = ref_filtered_df.select(ref_aggr_expr.alias(ref_metric_col)).limit(1)
+
+        if aggr_type not in CURATED_AGGR_FUNCTIONS:
+            _validate_aggregate_return_type(ref_metric_df, aggr_type, ref_metric_col)
+
+        return aggr_apply(df.crossJoin(ref_metric_df))
+
+    return condition, apply
+
+
 @register_rule("dataset")
 def compare_datasets(
     columns: list[str | Column],
@@ -3271,6 +3445,55 @@ def _build_aggregate_check_metadata(
     return name, group_by_list_str
 
 
+def _build_aggr_condition_result(
+    metric_col: Column,
+    limit_expr: Column,
+    compare_op: Callable[[Column, Column], Column],
+    abs_tolerance: float,
+    rel_tolerance: float,
+    null_safe_limit_compare: bool,
+) -> Column:
+    """
+    Compute the violation condition for an aggregate comparison, applying tolerance and/or
+    null-safe matching when the comparison is an equality check (*operator.eq*/*operator.ne*).
+
+    Args:
+        metric_col: Column holding the computed aggregate metric.
+        limit_expr: Column holding the limit to compare against.
+        compare_op: Comparison operator (e.g., operator.gt, operator.eq, operator.ne).
+        abs_tolerance: Absolute tolerance for numeric comparisons (0.0 disables it).
+        rel_tolerance: Relative tolerance for numeric comparisons (0.0 disables it).
+        null_safe_limit_compare: See *_is_aggr_compare*.
+
+    Returns:
+        A Spark Column expression that evaluates to True when the check is violated.
+    """
+    is_equality_check = compare_op in (py_operator.ne, py_operator.eq)
+
+    if (abs_tolerance > 0.0 or rel_tolerance > 0.0) and is_equality_check:
+        tolerance_match = _match_values_with_tolerance(metric_col, limit_expr, abs_tolerance, rel_tolerance)
+        # is_aggr_equal (compare_op=ne) fails when values don't match within tolerance;
+        # is_aggr_not_equal (compare_op=eq) fails when values match within tolerance.
+        condition_result = ~tolerance_match if compare_op == py_operator.ne else tolerance_match
+    else:
+        condition_result = compare_op(metric_col, limit_expr)
+
+    if null_safe_limit_compare and is_equality_check:
+        # Standard SQL NULL propagation would otherwise turn a NULL metric/limit into a NULL
+        # condition (no violation), silently hiding a mismatch. Treat two NULLs as equal and a
+        # one-sided NULL as a mismatch.
+        metric_is_null = metric_col.isNull()
+        limit_is_null = limit_expr.isNull()
+        violation_when_mismatch = compare_op == py_operator.ne
+        condition_result = (
+            F.when(metric_is_null & limit_is_null, F.lit(not violation_when_mismatch))
+            .when(metric_is_null | limit_is_null, F.lit(violation_when_mismatch))
+            .otherwise(condition_result)
+        )
+
+    return condition_result
+
+
 def _is_aggr_compare(
     column: str | Column,
     limit: int | float | Decimal | str | Column,
@@ -3283,6 +3506,7 @@ def _is_aggr_compare(
     compare_op_name: str,
     abs_tolerance: float | None = None,
     rel_tolerance: float | None = None,
+    null_safe_limit_compare: bool = False,
 ) -> tuple[Column, Callable]:
     """
     Helper to build aggregation comparison checks with a given operator.
@@ -3305,6 +3529,12 @@ def _is_aggr_compare(
         compare_op_name: Name identifier for the comparison (e.g., 'greater_than').
         abs_tolerance: Optional absolute tolerance for numeric comparisons.
         rel_tolerance: Optional relative tolerance for numeric comparisons.
+        null_safe_limit_compare: Only applies when *compare_op* is *operator.eq* or *operator.ne*. When
+            True, a NULL metric or NULL limit is never left to standard SQL NULL propagation: both NULL
+            is treated as a match (no violation), a one-sided NULL is treated as a mismatch (violation).
+            Use this when the limit can legitimately be NULL (e.g., an aggregate computed over an
+            external/reference dataset that may be empty), so a NULL limit isn't silently treated as
+            "no violation".
 
     Returns:
         A tuple of:
@@ -3355,7 +3585,15 @@ def apply(df: DataFrame) -> DataFrame:
             The DataFrame with additional condition and metric columns for aggregation validation.
         """
         filter_col = F.expr(row_filter) if row_filter else F.lit(True)
-        filtered_expr = F.when(filter_col, aggr_col_expr) if row_filter else aggr_col_expr
+        if row_filter:
+            # aggr_col_str == "*" only for count(*) over all rows (the only valid use of column="*").
+            # F.expr("*") can't be embedded as the THEN value of a CASE WHEN: Spark's star-expansion
+            # resolves it against every column in scope instead of treating it as a placeholder value.
+            # count() only cares about nullness, so a non-null literal is a safe stand-in.
+            then_expr = F.lit(1) if aggr_col_str == "*" else aggr_col_expr
+            filtered_expr = F.when(filter_col, then_expr)
+        else:
+            filtered_expr = aggr_col_expr
 
         # Build aggregation expression
         aggr_expr = _build_aggregate_expression(aggr_type, filtered_expr, aggr_params)
@@ -3397,33 +3635,24 @@ def apply(df: DataFrame) -> DataFrame:
 
             df = df.crossJoin(agg_df)  # bring the metric across all rows
 
-        # Apply tolerance-based comparison for equality checks
-        if (abs_tolerance > 0.0 or rel_tolerance > 0.0) and compare_op in (py_operator.ne, py_operator.eq):
-            tolerance_match = _match_values_with_tolerance(F.col(metric_col), limit_expr, abs_tolerance, rel_tolerance)
-
-            # Adjust based on compare_op:
-            if compare_op == py_operator.ne:
-                # is_aggr_equal case: fail when values don't match within tolerance
-                condition_result = ~tolerance_match
-            else:  # compare_op == py_operator.eq
-                # is_aggr_not_equal case: fail when values match within tolerance
-                condition_result = tolerance_match
-
-            df = df.withColumn(condition_col, condition_result)
-        else:
-            # Exact comparison or non-equality operators
-            df = df.withColumn(condition_col, compare_op(F.col(metric_col), limit_expr))
+        condition_result = _build_aggr_condition_result(
+            F.col(metric_col),
+            limit_expr,
+            compare_op,
+            abs_tolerance,
+            rel_tolerance,
+            null_safe_limit_compare,
+        )
+        df = df.withColumn(condition_col, condition_result)
 
         return df
 
-    # Get human-readable display name for aggregate function (including params if present)
-    aggr_display_name = _get_aggregate_display_name(aggr_type, aggr_params)
-
     condition = make_condition(
         condition=F.col(condition_col),
         message=F.concat_ws(
             "",
-            F.lit(f"{aggr_display_name} value "),
+            # Human-readable display name for aggregate function (including params if present)
+            F.lit(f"{_get_aggregate_display_name(aggr_type, aggr_params)} value "),
             F.col(metric_col).cast("string"),
             F.lit(f" in column '{aggr_col_str}'"),
             F.lit(f"{' per group of columns ' if group_by_list_str else ''}"),

tests/integration/test_apply_checks.py

@@ -7065,6 +7065,13 @@ def test_apply_checks_all_checks_using_classes(ws, spark):
             column="col6",
             check_func_kwargs={"window_minutes": 1, "min_records_per_window": 1, "lookback_windows": 3},
         ),
+        # aggr_matches_dataset check — row count matches the reference dataset
+        DQDatasetRule(
+            criticality="error",
+            check_func=check_funcs.aggr_matches_dataset,
+            column="*",
+            check_func_kwargs={"aggr_type": "count", "ref_df_name": "ref_df_key"},
+        ),
         # is_valid_json check
         DQRowRule(
             criticality="error",
@@ -7153,7 +7160,10 @@ def test_apply_checks_all_checks_using_classes(ws, spark):
         schema,
     )
 
-    checked = dq_engine.apply_checks(test_df, checks)
+    ref_df = test_df.withColumnRenamed("col1", "ref_col1").withColumnRenamed("col2", "ref_col2")
+    ref_dfs = {"ref_df_key": ref_df}
+
+    checked = dq_engine.apply_checks(test_df, checks, ref_dfs=ref_dfs)
 
     expected_schema = schema + REPORTING_COLUMNS
     expected = spark.createDataFrame(