Skip to content

fix: make kl_divergence and chi_square audits run on DuckDB and Trino #6129

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shinichitazawa wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: make kl_divergence and chi_square audits run on DuckDB and Trino #6129

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 22 additions & 56 deletions sqlmesh/core/audit/builtin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -532,84 +532,44 @@ def create_non_blocking_copy(audit: Audit) -> Audit:
WITH
table_a AS (
SELECT
@source_column,
@source_column AS x,
COUNT(*) AS num_rows
FROM @this_model
WHERE @condition
GROUP BY @source_column
),
table_b AS (
SELECT
@target_column,
@target_column AS x,
COUNT(*) AS num_rows
FROM @this_model
WHERE @condition
GROUP BY @target_column
),
table_a_with_p AS (
SELECT
@source_column,
num_rows,
x,
num_rows / SUM(num_rows) OVER () AS p
FROM table_a
),
table_b_with_q AS (
SELECT
@target_column,
num_rows,
x,
num_rows / SUM(num_rows) OVER () AS q
FROM table_b
),
table_a_with_q AS (
-- Values that appear in only one column are skipped, as their terms are undefined
divergence AS (
SELECT
@source_column,
num_rows,
p,
COALESCE(q, 0) AS q
SUM((p - q) * LN(p / q)) AS kl_divergence
FROM table_a_with_p
LEFT JOIN table_b_with_q USING (@source_column)
),
table_b_with_p AS (
SELECT
@target_column,
num_rows,
q,
COALESCE(p, 0) AS p
FROM table_b_with_q
LEFT JOIN table_a_with_p USING (@target_column)
),
table_a_with_kl AS (
SELECT
@source_column,
num_rows,
p,
q,
p * LOG(p / NULLIF(q, 0)) AS kl
FROM table_a_with_q
),
table_b_with_kl AS (
SELECT
@target_column,
num_rows,
p,
q,
q * LOG(q / NULLIF(p, 0)) AS kl
FROM table_b_with_p
),
unioned AS (
SELECT *
FROM table_a_with_kl
UNION ALL
SELECT *
FROM table_b_with_kl
INNER JOIN table_b_with_q
ON table_a_with_p.x = table_b_with_q.x
)
SELECT
@source_column,
@target_column,
SUM(kl) AS kl_divergence
FROM unioned
GROUP BY @source_column, @target_column
HAVING kl_divergence > @threshold
kl_divergence
FROM divergence
WHERE kl_divergence > @threshold
""",
)

Expand Down Expand Up @@ -646,17 +606,23 @@ def create_non_blocking_copy(audit: Audit) -> Audit:
(SELECT COUNT(*) FROM samples) as g_t -- g_t is the grand total
FROM samples AS r
GROUP BY x_a, x_b
),
statistic AS (
SELECT
(@c - 1) * (@r - 1) as degrees_of_freedom,
SUM((observed - @E) * (observed - @E) / @E) as chi_square
FROM contingency_table
)

SELECT
(@c - 1) * (@r - 1) as degrees_of_freedom,
SUM((observed - @E) * (observed - @E) / @E) as chi_square
FROM contingency_table
degrees_of_freedom,
chi_square
FROM statistic
-- H0: the two variables are independent
-- H1: the two variables are dependent
-- if chi_square > critical_value, reject H0
-- if chi_square <= critical_value, fail to reject H0
HAVING NOT @IF(@dependent, chi_square > @critical_value, chi_square <= @critical_value)
WHERE NOT @IF(@dependent, chi_square > @critical_value, chi_square <= @critical_value)
""",
)

Expand Down
15 changes: 14 additions & 1 deletion tests/core/test_audit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -651,7 +651,20 @@ def test_chi_square_audit(model: Model):
)
assert (
rendered_query.sql()
== """WITH "samples" AS (SELECT "a" AS "x_a", "b" AS "x_b" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_0" WHERE (NOT "a" IS NULL AND NOT "b" IS NULL) AND TRUE), "contingency_table" AS (SELECT "x_a", "x_b", COUNT(*) AS "observed", (SELECT COUNT(*) FROM "samples" AS "t" WHERE "r"."x_a" = "t"."x_a") AS "tot_a", (SELECT COUNT(*) FROM "samples" AS "t" WHERE "r"."x_b" = "t"."x_b") AS "tot_b", (SELECT COUNT(*) FROM "samples" AS "samples") AS "g_t" /* g_t is the grand total */ FROM "samples" AS "r" GROUP BY "x_a", "x_b") SELECT ((SELECT COUNT(DISTINCT "x_a") FROM "contingency_table" AS "contingency_table") - 1) * ((SELECT COUNT(DISTINCT "x_b") FROM "contingency_table" AS "contingency_table") - 1) AS "degrees_of_freedom", SUM(("observed" - ("tot_a" * "tot_b" / "g_t")) * ("observed" - ("tot_a" * "tot_b" / "g_t")) / ("tot_a" * "tot_b" / "g_t")) AS "chi_square" FROM "contingency_table" AS "contingency_table" /* H0: the two variables are independent */ /* H1: the two variables are dependent */ /* if chi_square > critical_value, reject H0 */ /* if chi_square <= critical_value, fail to reject H0 */ HAVING NOT "chi_square" > 9.48773"""
== """WITH "samples" AS (SELECT "a" AS "x_a", "b" AS "x_b" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_0" WHERE (NOT "a" IS NULL AND NOT "b" IS NULL) AND TRUE), "contingency_table" AS (SELECT "x_a", "x_b", COUNT(*) AS "observed", (SELECT COUNT(*) FROM "samples" AS "t" WHERE "r"."x_a" = "t"."x_a") AS "tot_a", (SELECT COUNT(*) FROM "samples" AS "t" WHERE "r"."x_b" = "t"."x_b") AS "tot_b", (SELECT COUNT(*) FROM "samples" AS "samples") AS "g_t" /* g_t is the grand total */ FROM "samples" AS "r" GROUP BY "x_a", "x_b"), "statistic" AS (SELECT ((SELECT COUNT(DISTINCT "x_a") FROM "contingency_table" AS "contingency_table") - 1) * ((SELECT COUNT(DISTINCT "x_b") FROM "contingency_table" AS "contingency_table") - 1) AS "degrees_of_freedom", SUM(("observed" - ("tot_a" * "tot_b" / "g_t")) * ("observed" - ("tot_a" * "tot_b" / "g_t")) / ("tot_a" * "tot_b" / "g_t")) AS "chi_square" FROM "contingency_table" AS "contingency_table") SELECT "degrees_of_freedom" AS "degrees_of_freedom", "chi_square" AS "chi_square" FROM "statistic" AS "statistic" /* H0: the two variables are independent */ /* H1: the two variables are dependent */ /* if chi_square > critical_value, reject H0 */ /* if chi_square <= critical_value, fail to reject H0 */ WHERE NOT "chi_square" > 9.48773"""
)


def test_kl_divergence_audit(model: Model):
rendered_query = model.render_audit_query(
builtin.kl_divergence_audit,
source_column=exp.to_column("a"),
target_column=exp.to_column("b"),
threshold=exp.convert(0.1),
)
assert (
rendered_query.sql()
== """WITH "table_a" AS (SELECT "a" AS "x", COUNT(*) AS "num_rows" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_0" WHERE TRUE GROUP BY "a"), "table_b" AS (SELECT "b" AS "x", COUNT(*) AS "num_rows" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_1" WHERE TRUE GROUP BY "b"), "table_a_with_p" AS (SELECT "x", "num_rows" / SUM("num_rows") OVER () AS "p" FROM "table_a" AS "table_a"), "table_b_with_q" AS (SELECT "x", "num_rows" / SUM("num_rows") OVER () AS "q" FROM "table_b" AS "table_b"), "divergence" /* Values that appear in only one column are skipped, as their terms are undefined */ AS (SELECT SUM(("p" - "q") * LN("p" / "q")) AS "kl_divergence" FROM "table_a_with_p" AS "table_a_with_p" INNER JOIN "table_b_with_q" AS "table_b_with_q" ON "table_a_with_p"."x" = "table_b_with_q"."x") SELECT "kl_divergence" AS "kl_divergence" FROM "divergence" AS "divergence" WHERE "kl_divergence" > 0.1"""
)


Expand Down