diff options
Diffstat (limited to 'tests/test_optimizer.py')
| -rw-r--r-- | tests/test_optimizer.py | 30 |
1 files changed, 30 insertions, 0 deletions
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index 9b68c78..3f96a0e 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -229,6 +229,15 @@ class TestOptimizer(unittest.TestCase): def test_qualify_columns(self, logger): self.assertEqual( optimizer.qualify_columns.qualify_columns( + parse_one("WITH x AS (SELECT a FROM db.y) SELECT * FROM db.x"), + schema={"db": {"x": {"z": "int"}, "y": {"a": "int"}}}, + expand_stars=False, + ).sql(), + "WITH x AS (SELECT y.a AS a FROM db.y) SELECT * FROM db.x", + ) + + self.assertEqual( + optimizer.qualify_columns.qualify_columns( parse_one("WITH x AS (SELECT a FROM db.y) SELECT z FROM db.x"), schema={"db": {"x": {"z": "int"}, "y": {"a": "int"}}}, infer_schema=False, @@ -285,6 +294,11 @@ class TestOptimizer(unittest.TestCase): schema = MappingSchema(self.schema, {"x": {"a"}, "y": {"b"}, "z": {"b"}}) self.check_file("qualify_columns__with_invisible", qualify_columns, schema=schema) + def test_pushdown_cte_alias_columns(self): + self.check_file( + "pushdown_cte_alias_columns", optimizer.qualify_columns.pushdown_cte_alias_columns + ) + def test_qualify_columns__invalid(self): for sql in load_sql_fixtures("optimizer/qualify_columns__invalid.sql"): with self.subTest(sql): @@ -534,6 +548,22 @@ FROM READ_CSV('tests/fixtures/optimizer/tpc-h/nation.csv.gz', 'delimiter', '|') level="warning", ) + def test_struct_type_annotation(self): + tests = { + "SELECT STRUCT(1 AS col)": "STRUCT<col INT>", + "SELECT STRUCT(1 AS col, 2.5 AS row)": "STRUCT<col INT, row DOUBLE>", + "SELECT STRUCT(1)": "STRUCT<INT>", + "SELECT STRUCT(1 AS col, 2.5 AS row, struct(3.5 AS inner_col, 4 AS inner_row) AS nested_struct)": "STRUCT<col INT, row DOUBLE, nested_struct STRUCT<inner_col DOUBLE, inner_row INT>>", + "SELECT STRUCT(1 AS col, 2.5, ARRAY[1, 2, 3] AS nested_array, 'foo')": "STRUCT<col INT, DOUBLE, nested_array ARRAY<INT>, VARCHAR>", + "SELECT STRUCT(1, 2.5, 'bar')": "STRUCT<INT, DOUBLE, VARCHAR>", + 'SELECT STRUCT(1 AS "CaseSensitive")': 'STRUCT<"CaseSensitive" INT>', + } + + for sql, target_type in tests.items(): + with self.subTest(sql): + expression = annotate_types(parse_one(sql)) + assert expression.expressions[0].is_type(target_type) + def test_literal_type_annotation(self): tests = { "SELECT 5": exp.DataType.Type.INT, |