diff options
Diffstat (limited to 'tests/test_optimizer.py')
| -rw-r--r-- | tests/test_optimizer.py | 41 |
1 files changed, 41 insertions, 0 deletions
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index c43a84e..8f5dd08 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -550,6 +550,47 @@ FROM READ_CSV('tests/fixtures/optimizer/tpc-h/nation.csv.gz', 'delimiter', '|') self.assertEqual(expression.right.this.left.type.this, exp.DataType.Type.INT) self.assertEqual(expression.right.this.right.type.this, exp.DataType.Type.INT) + def test_bracket_annotation(self): + expression = annotate_types(parse_one("SELECT A[:]")).expressions[0] + + self.assertEqual(expression.type.this, exp.DataType.Type.UNKNOWN) + self.assertEqual(expression.expressions[0].type.this, exp.DataType.Type.UNKNOWN) + + expression = annotate_types(parse_one("SELECT ARRAY[1, 2, 3][1]")).expressions[0] + self.assertEqual(expression.this.type.sql(), "ARRAY<INT>") + self.assertEqual(expression.type.this, exp.DataType.Type.INT) + + expression = annotate_types(parse_one("SELECT ARRAY[1, 2, 3][1 : 2]")).expressions[0] + self.assertEqual(expression.this.type.sql(), "ARRAY<INT>") + self.assertEqual(expression.type.sql(), "ARRAY<INT>") + + expression = annotate_types( + parse_one("SELECT ARRAY[ARRAY[1], ARRAY[2], ARRAY[3]][1][2]") + ).expressions[0] + self.assertEqual(expression.this.this.type.sql(), "ARRAY<ARRAY<INT>>") + self.assertEqual(expression.this.type.sql(), "ARRAY<INT>") + self.assertEqual(expression.type.this, exp.DataType.Type.INT) + + expression = annotate_types( + parse_one("SELECT ARRAY[ARRAY[1], ARRAY[2], ARRAY[3]][1:2]") + ).expressions[0] + self.assertEqual(expression.type.sql(), "ARRAY<ARRAY<INT>>") + + expression = annotate_types(parse_one("MAP(1.0, 2, '2', 3.0)['2']", read="spark")) + self.assertEqual(expression.type.this, exp.DataType.Type.DOUBLE) + + expression = annotate_types(parse_one("MAP(1.0, 2, x, 3.0)[2]", read="spark")) + self.assertEqual(expression.type.this, exp.DataType.Type.UNKNOWN) + + expression = annotate_types(parse_one("MAP(ARRAY(1.0, x), ARRAY(2, 3.0))[x]")) + self.assertEqual(expression.type.this, exp.DataType.Type.DOUBLE) + + expression = annotate_types( + parse_one("SELECT MAP(1.0, 2, 2, t.y)[2] FROM t", read="spark"), + schema={"t": {"y": "int"}}, + ).expressions[0] + self.assertEqual(expression.type.this, exp.DataType.Type.INT) + def test_interval_math_annotation(self): schema = { "x": { |