Feature/categorical values distribution - #1450
Conversation
mwojtyczka
left a comment
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.
| for key, alias in zip(expected_keys_ordered, per_key_aliases) | ||
| ] | ||
| aggregations.append(F.count(group_expr).alias("__total")) | ||
| aggregate_row = filtered.agg(*aggregations).collect()[0] |
There was a problem hiding this comment.
[correctness] Eager .collect() in the apply() closure breaks streaming.
Every other dataset check (is_unique, has_no_outliers, aggr_*, foreign_key) keeps its closure lazy — only transformations, returning df.withColumn(...). is_in_distribution is the only one that calls filtered.agg(...).collect() inside apply, which the executor invokes during apply_checks.
Failure: apply_checks() on a streaming DataFrame reaches closure_func(df) → .collect(), which raises AnalysisException ('Queries with streaming sources must be executed with writeStream.start()'). In batch it also forces an extra eager job mid-plan. Consider a lazy formulation (broadcast-join the expected distribution and compute TVD via aggregation expressions).
| message_col = f"__message_{col_str_norm}_{unique_str}" | ||
|
|
||
| def apply(df: DataFrame) -> DataFrame: | ||
| column_type = df.schema[col_expr_str].dataType |
There was a problem hiding this comment.
[correctness] df.schema[col_expr_str] KeyErrors for a Column expression input.
col_expr_str is the rendered expression (e.g. upper(v)), not a schema field name. For is_in_distribution(F.upper(F.col('v')), ...), df.schema['upper(v)'] raises KeyError since the schema only has field v. The signature/docstring advertise column expressions, but only F.col("value") (name == field) is tested. Same pattern exists in has_no_outliers, so it's a shared latent gap rather than net-new — worth a test and/or resolving the type off the expression.
| # nudges the computed tvd just above distance (e.g. abs(0.15-0.2) == 0.05000000000000002). | ||
| # The docstring's "less than or equal to the given distance" already promises the | ||
| # boundary passes, so equality-within-precision must not fire the check. | ||
| is_violation = tvd > distance and not math.isclose(tvd, distance) |
There was a problem hiding this comment.
[correctness] int keys vs byte/short/date columns rely on implicit Spark upcasting.
For a byte column with {1: 0.75, 2: 0.25}, F.lit(1) is an INT literal compared to a BYTE column. Spark upcasts so the tests pass, but there's no explicit cast of the literal to the column type. Residual/total bucketing depends on this equality holding for every supported small-int/date type; a type where Spark equality semantics differ (or a future type addition) would silently miscount into the residual bucket rather than erroring. An explicit cast of the key literal to the column type would make this robust.
| f"which exceeds the allowed distance={distance}." | ||
| ) | ||
|
|
||
| return df.withColumn(condition_col, F.lit(is_violation)).withColumn(message_col, F.lit(message_text)) |
There was a problem hiding this comment.
[altitude] Verdict baked as F.lit(python_bool) from the driver-side collect result.
Because the TVD is computed on the driver from collected counts, is_violation is a Python bool turned into F.lit, and the whole dataset shares a single plan-baked verdict. This is a design consequence of the .collect() approach (finding on line 599) and is what prevents the check from participating in a lazy/streaming plan. Reformulating to compute TVD with aggregation expressions would remove both the collect and the literal-baking.
|
|
||
| def _validate_distribution_values( | ||
| distribution: dict[bool, float] | dict[str, float] | dict[int, float] | dict[datetime.date, float] | ||
| ): |
There was a problem hiding this comment.
[conventions] Missing -> None return annotations.
AGENTS.md (Type Hints): "Every parameter and return value must be annotated. Enforced by mypy (make lint)." _validate_distribution_values, _validate_distribution_keys_collision, _validate_distribution_keys_types, and _validate_distribution_items end in ): with no -> None. _validate_distribution and _validate_distance in the same block do annotate -> None, so these four are inconsistent.
| 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.
mwojtyczka
added
the
under-review
2 participants
Changes
#1344
Linked issues
Resolves #..
Tests
Documentation and Demos