From 6c774776db5e016d597e582c7006ba8d27006f9d Mon Sep 17 00:00:00 2001 From: Daniel Baumann Date: Thu, 26 Oct 2023 19:21:50 +0200 Subject: Adding upstream version 18.17.0. Signed-off-by: Daniel Baumann --- tests/test_optimizer.py | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) (limited to 'tests/test_optimizer.py') 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") + 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") + self.assertEqual(expression.type.sql(), "ARRAY") + + 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>") + self.assertEqual(expression.this.type.sql(), "ARRAY") + 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>") + + 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": { -- cgit v1.2.3