tests for unary

Author: geooo109

sqlglot/optimizer/annotate_types.py

@@ -450,13 +450,16 @@ def _annotate_binary(self, expression: B) -> B:
         if isinstance(expression, (exp.Connector, exp.Predicate)):
             self._set_type(expression, exp.DataType.Type.BOOLEAN)
             expr_type = expression.type
-            left_is_nonnull = left_type and (left_type.args.get("nonnull") is True)
-            right_is_nonnull = right_type and (right_type.args.get("nonnull") is True)
-            if expr_type:
-                expr_type.set(
-                    "nonnull",
-                    isinstance(expression, exp.Is) or (left_is_nonnull and right_is_nonnull),
+            if expr_type and (
+                isinstance(expression, exp.Is)
+                or (
+                    left_type
+                    and (left_type.args.get("nonnull") is True)
+                    and right_type
+                    and (right_type.args.get("nonnull") is True)
                 )
+            ):
+                expr_type.set("nonnull", True)
         elif (left_type_this, right_type_this) in self.binary_coercions:
             self._set_type(
                 expression, self.binary_coercions[(left_type_this, right_type_this)](left, right)

tests/test_optimizer.py

@@ -1727,7 +1727,7 @@ def test_nonnull_annotation(self):
                     ):
                         sql = f"SELECT {sql_predicate} FROM foo"
                         query = parse_one(sql)
-                        annotated = annotate_types(query)
+                        annotated = annotate_types(query, schema=schema)
                         assert annotated.selects[0].type == exp.DataType.build(
                             "BOOLEAN", nonnull=nonnull
                         )
@@ -1738,5 +1738,15 @@ def test_nonnull_annotation(self):
                 ):
                     sql = f"SELECT {sql_predicate} FROM foo"
                     query = parse_one(sql)
-                    annotated = annotate_types(query)
+                    annotated = annotate_types(query, schema=schema)
                     assert annotated.selects[0].type == exp.DataType.build("BOOLEAN", nonnull=True)
+
+        for unary, unary_type in (("NOT", "BOOLEAN"), ("-", "INT")):
+            for value, nonnull in (("1", True), ("foo.id", False)):
+                with self.subTest(f"Test NULL propagation for unary: {unary} with value: {value}"):
+                    sql = f"SELECT {unary} {value} FROM foo"
+                    query = parse_one(sql)
+                    annotated = annotate_types(query, schema=schema)
+                    assert annotated.selects[0].type == exp.DataType.build(
+                        unary_type, nonnull=nonnull
+                    )