Skip to content

Commit 4635131

Browse files
sfc-gh-yixiesfc-gh-yixie
committed
Add tests in simplifier
1 parent 7f5c5de commit 4635131

File tree

1 file changed

+79
-0
lines changed

1 file changed

+79
-0
lines changed

‎tests/integ/test_simplifier_suite.py‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1649,6 +1649,22 @@ def test_chained_sort(session):
16491649
.filter(col("A") > 2),
16501650
'SELECT "A", "B", 12 :: INT AS "TWELVE" FROM ( SELECT $1 AS "A", $2 AS "B" FROM VALUES (1 :: INT, -2 :: INT), (3 :: INT, -4 :: INT) ) WHERE (("A" > 1{POSTFIX}) AND ("A" > 2{POSTFIX}))',
16511651
),
1652+
# Flattened if the dropped columns are not used in filter
1653+
(
1654+
lambda df: df.filter(col("A") >= 1)
1655+
.select(col("A").alias("C"), col("B").alias("D"))
1656+
.filter(col("C") > 2)
1657+
.select(col("C")),
1658+
'SELECT "A" AS "C" FROM (SELECT $1 AS "A", $2 AS "B" FROM VALUES (1 :: INT, -2 :: INT), (3 :: INT, -4 :: INT)) WHERE (("A" >= 1{POSTFIX}) AND ("C" > 2{POSTFIX}))',
1659+
),
1660+
# Flattened if the dropped columns are not in the filter clause's dependent columns
1661+
(
1662+
lambda df: df.filter(col("A") >= 1)
1663+
.select(col("A").alias("C"), col("B").alias("D"))
1664+
.filter((col("C") + 1) > 2)
1665+
.select(col("C")),
1666+
'SELECT "A" AS "C" FROM (SELECT $1 AS "A", $2 AS "B" FROM VALUES (1 :: INT, -2 :: INT), (3 :: INT, -4 :: INT)) WHERE (("A" >= 1{POSTFIX}) AND (("C" + 1{POSTFIX}) > 2{POSTFIX}))',
1667+
),
16521668
# Not fully flattened, since col("A") > 1 and col("A") > 2 are referring to different columns
16531669
(
16541670
lambda df: df.filter(col("A") > 1)
@@ -1671,6 +1687,29 @@ def test_chained_sort(session):
16711687
lambda df: df.filter(col("$1") > 1).select(col("B"), col("A")),
16721688
'SELECT "B", "A" FROM ( SELECT * FROM ( SELECT "A", "B" FROM ( SELECT $1 AS "A", $2 AS "B" FROM VALUES (1 :: INT, -2 :: INT), (3 :: INT, -4 :: INT) ) ) WHERE ("$1" > 1{POSTFIX}) )',
16731689
),
1690+
# Not flattened if a dropped column is used in the filter clause
1691+
(
1692+
lambda df: df.filter(col("A") >= 1)
1693+
.select(col("A"), col("B").alias("D"))
1694+
.filter(col("D") > -3)
1695+
.select(col("A").alias("E")),
1696+
'SELECT "A" AS "E" FROM (SELECT "A", "B" AS "D" FROM (SELECT $1 AS "A", $2 AS "B" FROM VALUES (1 :: INT, -2 :: INT), (3 :: INT, -4 :: INT)) WHERE (("A" >= 1{POSTFIX}) AND ("D" > -3{POSTFIX})))',
1697+
),
1698+
# Not flattened if a dropped column is used in the select clause's dependent columns
1699+
(
1700+
lambda df: df.filter(col("A") >= 1)
1701+
.select(col("A"), col("B").alias("D"))
1702+
.filter((col("D") - 1) > -4)
1703+
.select((col("A") + 1).alias("E")),
1704+
'SELECT ("A" + 1{POSTFIX}) AS "E" FROM (SELECT "A", "B" AS "D" FROM (SELECT $1 AS "A", $2 AS "B" FROM VALUES (1 :: INT, -2 :: INT), (3 :: INT, -4 :: INT)) WHERE (("A" >= 1{POSTFIX}) AND (("D" - 1{POSTFIX}) > -4{POSTFIX})))',
1705+
),
1706+
# Not flattened if a dropped column that was changed expression is used in the select clause's dependent columns
1707+
(
1708+
lambda df: df.select(col("A"), (col("B") + 1).alias("B"))
1709+
.filter((col("B") - 1) > -4)
1710+
.select((col("A") + 1).alias("E")),
1711+
'SELECT ("A" + 1{POSTFIX}) AS "E" FROM (SELECT "A", ("B" + 1{POSTFIX}) AS "B" FROM (SELECT $1 AS "A", $2 AS "B" FROM VALUES (1 :: INT, -2 :: INT), (3 :: INT, -4 :: INT))) WHERE (("B" - 1{POSTFIX}) > -4{POSTFIX})',
1712+
),
16741713
],
16751714
)
16761715
def test_select_after_filter(setup_reduce_cast, session, operation, simplified_query):
@@ -1741,6 +1780,46 @@ def test_select_after_filter(setup_reduce_cast, session, operation, simplified_q
17411780
'SELECT "A" FROM ( SELECT $1 AS "A", $2 AS "B" FROM VALUES (1 :: INT, -2 :: INT), (3 :: INT, -4 :: INT) ) ORDER BY "B" ASC NULLS FIRST, "A" ASC NULLS FIRST',
17421781
True,
17431782
),
1783+
# Flattened if the dropped columns are not used in filter
1784+
(
1785+
lambda df: df.select(col("A").alias("C"), col("B").alias("D"))
1786+
.order_by(col("C"))
1787+
.select(col("C")),
1788+
'SELECT "A" AS "C" FROM (SELECT $1 AS "A", $2 AS "B" FROM VALUES (1 :: INT, -2 :: INT), (3 :: INT, -4 :: INT)) ORDER BY "C" ASC NULLS FIRST',
1789+
True,
1790+
),
1791+
# Flattened if the dropped columns are not in the order by clause's dependent columns
1792+
(
1793+
lambda df: df.select(col("A").alias("C"), col("B").alias("D"))
1794+
.order_by(col("C") + 1)
1795+
.select(col("C")),
1796+
'SELECT "A" AS "C" FROM (SELECT $1 AS "A", $2 AS "B" FROM VALUES (1 :: INT, -2 :: INT), (3 :: INT, -4 :: INT)) ORDER BY ("C" + 1{POSTFIX}) ASC NULLS FIRST',
1797+
True,
1798+
),
1799+
# Not flattened if a dropped new column is used in the order by clause
1800+
(
1801+
lambda df: df.select(col("A"), col("B").alias("D"))
1802+
.order_by(col("D"))
1803+
.select(col("A").alias("E")),
1804+
'SELECT "A" AS "E" FROM (SELECT "A", "B" AS "D" FROM (SELECT $1 AS "A", $2 AS "B" FROM VALUES (1 :: INT, -2 :: INT), (3 :: INT, -4 :: INT)) ORDER BY "D" ASC NULLS FIRST)',
1805+
True,
1806+
),
1807+
# Not flattened if a dropped new column is used in the order by clause's dependent columns
1808+
(
1809+
lambda df: df.select(col("A"), col("B").alias("D"))
1810+
.order_by(col("D") - 1)
1811+
.select((col("A") + 1).alias("E")),
1812+
'SELECT ("A" + 1{POSTFIX}) AS "E" FROM (SELECT "A", "B" AS "D" FROM (SELECT $1 AS "A", $2 AS "B" FROM VALUES (1 :: INT, -2 :: INT), (3 :: INT, -4 :: INT)) ORDER BY ("D" - 1{POSTFIX}) ASC NULLS FIRST)',
1813+
True,
1814+
),
1815+
# Not flattened if a dropped column that was changed expression is used in the select clause's dependent columns
1816+
(
1817+
lambda df: df.select(col("A"), (col("B") + 1).alias("B"))
1818+
.filter((col("B") - 1) > -4)
1819+
.select((col("A") + 1).alias("E")),
1820+
'SELECT ("A" + 1{POSTFIX}) AS "E" FROM (SELECT "A", ("B" + 1{POSTFIX}) AS "B" FROM (SELECT $1 AS "A", $2 AS "B" FROM VALUES (1 :: INT, -2 :: INT), (3 :: INT, -4 :: INT))) WHERE (("B" - 1{POSTFIX}) > -4{POSTFIX})',
1821+
True,
1822+
),
17441823
],
17451824
)
17461825
def test_select_after_orderby(

0 commit comments

Comments
 (0)