pudl_diff.rows ¶

def compare_rows_with_pk(
    left: pl.LazyFrame,
    right: pl.LazyFrame,
    pk_cols: Sequence[str],
    *,
    rtol: float = 1e-5,
    atol: float = 1e-8,
) -> KeyedRowDiff:
    """Compare two tables sharing a primary key.

    Args:
        left: The "left" table to compare. Must have the same columns as
            `right`, with matching dtypes for the primary key columns.
        right: The "right" table to compare against `left`.
        pk_cols: Names of the primary key columns shared by both tables, e.g.
            from `PudlDiffDataset.primary_key(table_name)`.
        rtol: Relative tolerance used to treat two floating point values as
            equal, matching `numpy.isclose()`'s default.
        atol: Absolute tolerance used to treat two floating point values as
            equal, matching `numpy.isclose()`'s default.
    """
    schema = left.collect_schema()
    columns = list(schema.keys())
    non_pk_cols = [name for name in schema if name not in pk_cols]
    left_cols = [f"{c}_left" for c in non_pk_cols]
    right_cols = [f"{c}_right" for c in non_pk_cols]

    spill_dir = SpillDir()
    delta = _diff_by_key_hash(
        left,
        right,
        schema,
        pk_cols,
        rtol,
        atol,
        spill_dir,
        value_columns=non_pk_cols,
        multiset=False,
    )

    # Only keys present on both sides whose non-key values hash differently can
    # have changed, so compare just those rows, exactly and with tolerance.
    on_both_sides = (pl.col(_LEFT_COUNT_COL) > 0) & (pl.col(_RIGHT_COUNT_COL) > 0)
    left_candidates = _spill(
        _semi_rows(delta.left_keyed, delta.affected, on_both_sides).select(columns),
        spill_dir,
        "left_candidates",
    )
    right_candidates = _spill(
        _semi_rows(delta.right_keyed, delta.affected, on_both_sides).select(columns),
        spill_dir,
        "right_candidates",
    )
    left_renamed = left_candidates.rename(
        dict(zip(non_pk_cols, left_cols, strict=True))
    )
    right_renamed = right_candidates.rename(
        dict(zip(non_pk_cols, right_cols, strict=True))
    )
    joined = left_renamed.join(right_renamed, on=list(pk_cols), how="inner")
    diff_flag_cols = [f"_pudl_diff_flag_{c}" for c in non_pk_cols]
    joined = joined.with_columns(
        [
            _values_differ(lc, rc, schema[c], rtol, atol).alias(flag_col)
            for c, lc, rc, flag_col in zip(
                non_pk_cols, left_cols, right_cols, diff_flag_cols, strict=True
            )
        ]
    )

    # Keep only the rows that really changed (both sides' values and the
    # per-column flags). Everything below - the counts and both output frames -
    # is derived from this small spilled result.
    any_diff = pl.any_horizontal(diff_flag_cols) if diff_flag_cols else pl.lit(False)
    changed = _spill(joined.filter(any_diff), spill_dir, "changed")

    counts = (
        changed.select([pl.col(fc).sum() for fc in diff_flag_cols]).collect(
            engine="streaming"
        )
        if non_pk_cols
        else pl.DataFrame()
    )
    column_changes = {
        c: count
        for c, fc in zip(non_pk_cols, diff_flag_cols, strict=True)
        if (count := counts[fc][0]) > 0
    }

    changed_left = (
        changed.select([*pk_cols, *left_cols])
        .rename(dict(zip(left_cols, non_pk_cols, strict=True)))
        .select(columns)
    )
    changed_right = (
        changed.select([*pk_cols, *right_cols])
        .rename(dict(zip(right_cols, non_pk_cols, strict=True)))
        .select(columns)
    )

    return KeyedRowDiff(
        pk_diff=delta.row_set_diff,
        column_changes=column_changes,
        changed_row_count=count_rows(changed),
        changed_left=changed_left,
        changed_right=changed_right,
    )

def compare_rows_without_pk(
    left: pl.LazyFrame,
    right: pl.LazyFrame,
    *,
    rtol: float = 1e-5,
    atol: float = 1e-8,
) -> RowSetDiff:
    """Compare two tables with no primary key by taking the symmetric difference of rows.

    The comparison is of multisets: a row that appears a different number of times
    in the two tables counts as a difference, and the surplus copies are reported
    as only in the table that has more of them.

    Args:
        left: The "left" table to compare. Must have the same columns as
            `right`, with the same dtypes.
        right: The "right" table to compare against `left`.
        rtol: Relative tolerance used to treat two floating point values as
            equal, matching `numpy.isclose()`'s default. Set to `0` along
            with `atol` to require exact float equality.
        atol: Absolute tolerance used to treat two floating point values as
            equal, matching `numpy.isclose()`'s default.

    Raises:
        polars.exceptions.SchemaError: if a column's dtype differs between
            `left` and `right`.
    """
    schema = left.collect_schema()
    spill_dir = SpillDir()
    return _diff_by_key_hash(
        left, right, schema, schema.keys(), rtol, atol, spill_dir, multiset=True
    ).row_set_diff