add meta

Author: geooo109

sqlglot/dialects/clickhouse.py

@@ -162,7 +162,7 @@ def _timestrtotime_sql(self: ClickHouse.Generator, expression: exp.TimeStrToTime
     datatype = exp.DataType.build(
         exp.DataType.Type.DATETIME64,
         expressions=[exp.DataTypeParam(this=exp.Literal.number(6)), *expressions],
-        nonnull=True,
+        nullable=False,
     )
 
     return self.sql(exp.cast(ts, datatype, dialect=self.dialect))
@@ -606,13 +606,13 @@ def _parse_types(
             dtype = super()._parse_types(
                 check_func=check_func, schema=schema, allow_identifiers=allow_identifiers
             )
-            if isinstance(dtype, exp.DataType) and dtype.args.get("nonnull") is not False:
+            if isinstance(dtype, exp.DataType) and dtype.args.get("nullable") is not True:
                 # Mark every type as non-nullable which is ClickHouse's default, unless it's
                 # already marked as nullable. This marker helps us transpile types from other
                 # dialects to ClickHouse, so that we can e.g. produce `CAST(x AS Nullable(String))`
                 # from `CAST(x AS TEXT)`. If there is a `NULL` value in `x`, the former would
                 # fail in ClickHouse without the `Nullable` type constructor.
-                dtype.set("nonnull", True)
+                dtype.set("nullable", False)
 
             return dtype
 
@@ -1237,7 +1237,7 @@ def trycast_sql(self, expression: exp.TryCast) -> str:
             dtype = expression.to
             if not dtype.is_type(*self.NON_NULLABLE_TYPES, check_nullable=True):
                 # Casting x into Nullable(T) appears to behave similarly to TRY_CAST(x AS T)
-                dtype.set("nonnull", False)
+                dtype.set("nullable", True)
 
             return super().cast_sql(expression)
 
@@ -1294,9 +1294,9 @@ def datatype_sql(self, expression: exp.DataType) -> str:
             #   String or FixedString (possibly LowCardinality) or UUID or IPv6"
             # - It's not a composite type, e.g. `Nullable(Array(...))` is not a valid type
             parent = expression.parent
-            nonnull = expression.args.get("nonnull")
-            if (
-                nonnull is not True
+            nullable = expression.args.get("nullable")
+            if nullable is True or (
+                nullable is None
                 and not (
                     isinstance(parent, exp.DataType)
                     and parent.is_type(exp.DataType.Type.MAP, check_nullable=True)

sqlglot/expressions.py

@@ -4607,8 +4607,8 @@ def name(self) -> str:
         return self.this.name
 
 
-# The `nonnull` arg is helpful when transpiling types from other dialects to ClickHouse, which
-# assumes non-nullable types by default. Values `None` and `False` mean the type is nullable.
+# The `nullable` arg is helpful when transpiling types from other dialects to ClickHouse, which
+# assumes non-nullable types by default. Values `None` and `True` mean the type is nullable.
 class DataType(Expression):
     arg_types = {
         "this": True,
@@ -4617,7 +4617,7 @@ class DataType(Expression):
         "values": False,
         "prefix": False,
         "kind": False,
-        "nonnull": False,
+        "nullable": False,
     }
 
     class Type(AutoName):
@@ -4901,23 +4901,20 @@ def is_type(self, *dtypes: DATA_TYPE, check_nullable: bool = False) -> bool:
         Returns:
             True, if and only if there is a type in `dtypes` which is equal to this DataType.
         """
-        self_nonnull = self.args.get("nonnull")
+        self_is_nullable = self.args.get("nullable")
         for dtype in dtypes:
             other_type = DataType.build(dtype, copy=False, udt=True)
-            other_nonnull = other_type.args.get("nonnull")
-
+            other_is_nullable = other_type.args.get("nullable")
             if (
                 other_type.expressions
+                or (check_nullable and (self_is_nullable or other_is_nullable))
                 or self.this == DataType.Type.USERDEFINED
                 or other_type.this == DataType.Type.USERDEFINED
             ):
                 matches = self == other_type
             else:
                 matches = self.this == other_type.this
 
-            if matches and check_nullable:
-                matches = self_nonnull == other_nonnull
-
             if matches:
                 return True
         return False

sqlglot/optimizer/annotate_types.py

@@ -208,16 +208,9 @@ def _set_type(
         prev_type = expression.type
         expression_id = id(expression)
 
-        expr_type = target_type or exp.DataType.Type.UNKNOWN  # type: ignore
-        expression.type = expr_type
+        expression.type = target_type or exp.DataType.Type.UNKNOWN  # type: ignore
         self._visited.add(expression_id)
 
-        if isinstance(expr_type, exp.DataType):
-            if isinstance(expression, exp.Null) or (
-                isinstance(expression, exp.Column) and expr_type.args.get("nonnull") is not True
-            ):
-                expr_type.set("nonnull", False)
-
         if (
             not self._supports_null_type
             and t.cast(exp.DataType, expression.type).this == exp.DataType.Type.NULL
@@ -444,28 +437,20 @@ def _annotate_binary(self, expression: B) -> B:
             self._set_type(expression, None)
             return expression
 
-        left_type, right_type = left.type, right.type
-        left_type_this, right_type_this = left_type.this, right_type.this  # type: ignore
+        left_type, right_type = left.type.this, right.type.this  # type: ignore
 
         if isinstance(expression, (exp.Connector, exp.Predicate)):
             self._set_type(expression, exp.DataType.Type.BOOLEAN)
-            expr_type = expression.type
-            if expr_type and (
+            if (
                 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)
-                )
+                or left.meta.get("nullable") is False
+                or right.meta.get("nullable") is False
             ):
-                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)
-            )
+                expression.meta["nullable"] = False
+        elif (left_type, right_type) in self.binary_coercions:
+            self._set_type(expression, self.binary_coercions[(left_type, right_type)](left, right))
         else:
-            self._set_type(expression, self._maybe_coerce(left_type_this, right_type_this))
+            self._set_type(expression, self._maybe_coerce(left_type, right_type))
 
         return expression
 
@@ -477,12 +462,8 @@ def _annotate_unary(self, expression: E) -> E:
         else:
             self._set_type(expression, expression.this.type)
 
-        if (
-            (this_type := expression.this.type)
-            and this_type.args.get("nonnull")
-            and (expr_type := expression.type)
-        ):
-            expr_type.set("nonnull", True)
+        if expression.this.meta.get("nullable") is False:
+            expression.meta["nullable"] = False
 
         return expression
 
@@ -494,10 +475,7 @@ def _annotate_literal(self, expression: exp.Literal) -> exp.Literal:
         else:
             self._set_type(expression, exp.DataType.Type.DOUBLE)
 
-        if not isinstance(expression.parent, (exp.Alias, exp.ColumnDef, exp.PropertyEQ)) and (
-            expr_type := expression.type
-        ):
-            expr_type.set("nonnull", True)
+        expression.meta["nullable"] = False
 
         return expression
 

sqlglot/optimizer/simplify.py

@@ -402,7 +402,7 @@ def remove_complements(expression, root=True):
         ops = set(expression.flatten())
         for op in ops:
             if isinstance(op, exp.Not) and op.this in ops:
-                if (expr_type := expression.type) and (expr_type.args.get("nonnull")):
+                if expression.meta.get("nullable") is False:
                     return exp.false() if isinstance(expression, exp.And) else exp.true()
 
     return expression

sqlglot/parser.py

@@ -5472,7 +5472,7 @@ def _parse_types(
                 )
                 if type_token == TokenType.NULLABLE and len(expressions) == 1:
                     this = expressions[0]
-                    this.set("nonnull", False)
+                    this.set("nullable", True)
                     self._match_r_paren()
                     return this
             elif type_token in self.ENUM_TYPE_TOKENS:

tests/test_optimizer.py

@@ -1700,46 +1700,36 @@ def test_annotate_object_construct(self):
             annotated.selects[0].type.sql("snowflake"), 'OBJECT("foo" VARCHAR, "a b" VARCHAR)'
         )
 
-    def test_nonnull_annotation(self):
-        for literal_sql, literal_type in (("1", "INT"), ("'foo'", "VARCHAR"), ("2.5", "DOUBLE")):
+    def test_nullable_annotation(self):
+        for literal_sql in ("1", "'foo'", "2.5"):
             with self.subTest(f"Test NULL annotation for literal: {literal_sql}"):
                 sql = f"SELECT {literal_sql}"
                 query = parse_one(sql)
                 annotated = annotate_types(query)
-                assert annotated.selects[0].type == exp.DataType.build(literal_type, nonnull=True)
+                assert annotated.selects[0].meta.get("nullable") is False
 
         schema = {"foo": {"id": "INT"}}
-        sql = "SELECT foo.id FROM foo"
-        query = parse_one(sql)
-        annotated = annotate_types(query, schema=schema)
-        self.assertTrue(
-            annotated.selects[0].type.is_type(
-                exp.DataType.build("INT", nonnull=False), check_nullable=True
-            )
-        )
 
         for predicate in (">", "<", ">=", "<=", "=", "!=", "<>", "LIKE", "NOT LIKE"):
-            for operand, nonnull in (("1", True), ("foo.id", False)):
+            for operand, nullable in (("1", False), ("foo.id", None)):
                 sql_predicate = f"{operand} {predicate} {operand}"
                 with self.subTest(f"Test NULL propagation for predicate: {predicate}"):
                     sql = f"SELECT {sql_predicate} FROM foo"
                     query = parse_one(sql)
                     annotated = annotate_types(query, schema=schema)
-                    assert annotated.selects[0].type == exp.DataType.build(
-                        "BOOLEAN", nonnull=nonnull
-                    )
+                    assert annotated.selects[0].meta.get("nullable") is nullable
 
         for predicate in ("IS NULL", "IS NOT NULL"):
             sql_predicate = f"foo.id {predicate}"
             with self.subTest(f"Test NULL propagation for predicate: {predicate}"):
                 sql = f"SELECT {sql_predicate} FROM foo"
                 query = parse_one(sql)
                 annotated = annotate_types(query, schema=schema)
-                assert annotated.selects[0].type == exp.DataType.build("BOOLEAN", nonnull=True)
+                assert annotated.selects[0].meta.get("nullable") is False
 
         for connector in ("AND", "OR"):
             for predicate in (">", "<", ">=", "<=", "=", "!=", "<>", "LIKE", "NOT LIKE"):
-                for operand, nonnull in (("1", True), ("foo.id", False)):
+                for operand, nullable in (("1", False), ("foo.id", None)):
                     sql_predicate = f"({operand} {predicate} {operand})"
                     sql_connector = f"{sql_predicate} {connector} {sql_predicate}"
                     with self.subTest(
@@ -1748,16 +1738,12 @@ def test_nonnull_annotation(self):
                         sql = f"SELECT {sql_connector} FROM foo"
                         query = parse_one(sql)
                         annotated = annotate_types(query, schema=schema)
-                        assert annotated.selects[0].type == exp.DataType.build(
-                            "BOOLEAN", nonnull=nonnull
-                        )
+                        assert annotated.selects[0].meta.get("nullable") is nullable
 
-        for unary, unary_type in (("NOT", "BOOLEAN"), ("-", "INT")):
-            for value, nonnull in (("1", True), ("foo.id", False)):
+        for unary in ("NOT", "-"):
+            for value, nullable in (("1", False), ("foo.id", None)):
                 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
-                    )
+                    assert annotated.selects[0].meta.get("nullable") is nullable