Feature: Categorical values distribution - #1450
Conversation
mwojtyczka
left a comment
There was a problem hiding this comment.
Automated code review (high-effort recall pass) for is_in_distribution. 6 findings inline, ranked most-severe first. Correctness items #1–#3 are marked plausible — verified by reading the sibling dataset-check closures rather than executing Spark locally.
This review was generated with assistance from Claude Code.
| aggr_type: avg | ||
| lookback_num_intervals: 2 | ||
| warmup_num_intervals: 2 No newline at end of file | ||
| warmup_num_intervals: 2 |
There was a problem hiding this comment.
[test-coverage] Fixture uses distance: 1.0 (the TVD maximum), so it can never fail.
is_violation = tvd > 1.0 is unreachable (TVD is bounded by 1). The 'apply every check' regression fixture therefore only verifies the check runs without error, never that a real distribution mismatch is caught end-to-end through the YAML/metadata path. A regression that inverts the comparison or mis-buckets residual mass would still pass. Consider a second fixture entry with a distribution + small distance that is expected to flag.
|
Hello, @mwojtyczka! Thanks for a review. The previous feedback has been addressed. Would it be possible to have another round? |
mwojtyczka
left a comment
There was a problem hiding this comment.
Follow-up review pass
Resolved (verified against HEAD — threads now closed): the earlier .collect()/streaming (L599), driver-baked verdict (L646), df.schema[expr] KeyError (L565), implicit int upcast (L640), and missing -> None (L691) concerns are all addressed by the refactor to lazy conditional aggregation + explicit key casts. Marked those five threads resolved.
Still open: the YAML fixture thread (all_dataset_checks.yaml:267) is only partially addressed — the new second entry (TVD=0.5, distance=1.0) guards against a comparison inversion, but no fixture entry actually flags a mismatch end-to-end (both entries use distance: 1.0 and pass by design). Left open so you can decide whether positive-detection belongs here or in the integration test.
New findings (inline): two correctness edge cases (0.0-probability key under impute=False; is_close abs_tol=0 at distance=0) plus three minor items.
| ) | ||
| is_violation_column = ( | ||
| F.when(total_column == 0, F.lit(False)) | ||
| .when(has_missing_column, F.lit(True)) |
There was a problem hiding this comment.
[correctness] impute=False flags a 0.0-probability key as 'missing'. has_missing_column is true whenever any expected key has a per-key count of 0, so a key whose expected probability is exactly 0.0 (which validation permits) is reported as a violation when it is legitimately absent. Example: distribution={A: 0.0, B: 1.0}, data all B -> 'missing expected keys [A]', yet with impute=True the same data yields TVD=0 and passes. The two modes disagree on identical data. Consider excluding expected-probability-0 keys from the missing-key set.
| # given distance" already promises the boundary passes, so equality-within-precision must | ||
| # not fire the check. | ||
| distance_lit = F.lit(distance) | ||
| is_close_col = F.abs(tvd_col - distance_lit) <= F.lit(1e-9) * F.greatest(F.abs(tvd_col), F.abs(distance_lit)) |
There was a problem hiding this comment.
[correctness] is_close uses abs_tol=0, leaving the distance=0 boundary unprotected against float noise. The tolerance is 1e-9 * max(|tvd|, |distance|), which collapses to ~0 when distance=0 and tvd is ~1e-16. For three equal 1/3 categories with a 1:1:1 sample, expected_residual = 1 - fsum(1/3,1/3,1/3) = 1.1e-16 (L737) gives tvd = 5.5e-17; the tolerance is ~5.5e-26, so ~is_close is True and tvd > 0 -> every row is flagged despite a mathematically perfect match. distance=0 is exercised throughout the tests. math.isclose's default abs_tol=0 is unsafe against zero — add a small absolute floor to the tolerance or special-case distance=0.
| else: | ||
| # Missing keys are those whose per-key count is zero; sort for stable output. | ||
| missing_key_names = [ | ||
| F.when(F.col(alias) == 0, F.lit(repr(key))) for key, alias in zip(expected_keys, per_key_aliases) |
There was a problem hiding this comment.
[minor] With case_sensitive=False, the 'missing keys' message echoes the case-normalized key. expected_keys here comes from normalized_distribution (lowercased at L641-642), so repr(key) renders 'gold' rather than the user's 'Gold' — the violation text reports a key the user never wrote. Carry the original key alongside the normalized one for the message.
|
|
||
| is_string_column = isinstance(column_type, (types.StringType, types.CharType)) | ||
| group_expr = col_expr | ||
| normalized_distribution: dict[Any, float] = distribution |
There was a problem hiding this comment.
[conventions] Any in new type hints. dict[Any, float] here (and list[Any] / dict[Any, float] on the helper signatures) — AGENTS.md Type Hints: 'Unknown types -> object or Protocol, not Any ... New code outside [anomaly/] must not introduce it.' The supported key types are a bounded set; object or a bounded union satisfies the rule with identical behavior.
| @@ -15,6 +16,7 @@ | |||
| import pyspark.sql.functions as F | |||
| from pyspark.sql import types | |||
| from pyspark.sql import Column, DataFrame, SparkSession | |||
| from pyspark.sql.types import DataType | |||
There was a problem hiding this comment.
[minor] Redundant import. from pyspark.sql.types import DataType duplicates the already-imported types module used everywhere else in this file (types.StringType, etc.). _is_in_distribution_get_data_type can return types.DataType; dropping this line keeps the file's single-import convention.
mwojtyczka
left a comment
There was a problem hiding this comment.
Left some additional comments, previous comments resolved except one only partilally addressed.
Changes
New check: Categorical values distribution check at dataset level using TDV measurement.
Linked issues
Resolves #1344
Tests
Documentation and Demos