Compare ARRAY & STRUCT types in BigQuery (simplistically)

Author: Sergey Vasilyev

data_diff/sqeleton/abcs/database_types.py

@@ -13,6 +13,7 @@
 DbTime = datetime
 
 
+@dataclass
 class ColType:
     supported = True
 
@@ -143,6 +144,18 @@ def __post_init__(self):
         assert self.precision == 0
 
 
+@dataclass
+class Array(ColType):
+    item_type: ColType
+
+
+@dataclass
+class Struct(ColType):
+    # TODO: Later, add the exact definition of the struct so that we could hash & compare individual fields.
+    #   Meanwhile, we rely on JSON comparison, so the internal structure does not matter.
+    pass
+
+
 @dataclass
 class UnknownColType(ColType):
     text: str
@@ -221,6 +234,10 @@ def parse_type(
     ) -> ColType:
         "Parse type info as returned by the database"
 
+    @abstractmethod
+    def to_comparable(self, value: str, coltype: ColType) -> str:
+        """Ensure that the expression is comparable in ``IS DISTINCT FROM``."""
+
 
 from typing import TypeVar, Generic
 

data_diff/sqeleton/abcs/mixins.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from .database_types import TemporalType, FractionalType, ColType_UUID, Boolean, ColType, String_UUID
+from .database_types import Array, Struct, TemporalType, FractionalType, ColType_UUID, Boolean, ColType, String_UUID
 from .compiler import Compilable
 
 
@@ -8,6 +8,11 @@ class AbstractMixin(ABC):
 
 
 class AbstractMixin_NormalizeValue(AbstractMixin):
+
+    @abstractmethod
+    def to_comparable(self, value: str, coltype: ColType) -> str:
+        """Ensure that the expression is comparable in ``IS DISTINCT FROM``."""
+
     @abstractmethod
     def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
         """Creates an SQL expression, that converts 'value' to a normalized timestamp.
@@ -43,6 +48,14 @@ def normalize_boolean(self, value: str, _coltype: Boolean) -> str:
         """Creates an SQL expression, that converts 'value' to either '0' or '1'."""
         return self.to_string(value)
 
+    def normalize_array(self, value: str, _coltype: Array) -> str:
+        """Creates an SQL expression, that serialized an array into a JSON string."""
+        return self.to_string(value)
+
+    def normalize_struct(self, value: str, _coltype: Struct) -> str:
+        """Creates an SQL expression, that serialized a struct into a JSON string."""
+        return self.to_string(value)
+
     def normalize_uuid(self, value: str, coltype: ColType_UUID) -> str:
         """Creates an SQL expression, that strips uuids of artifacts like whitespace."""
         if isinstance(coltype, String_UUID):
@@ -73,6 +86,10 @@ def normalize_value_by_type(self, value: str, coltype: ColType) -> str:
             return self.normalize_uuid(value, coltype)
         elif isinstance(coltype, Boolean):
             return self.normalize_boolean(value, coltype)
+        elif isinstance(coltype, Array):
+            return self.normalize_array(value, coltype)
+        elif isinstance(coltype, Struct):
+            return self.normalize_struct(value, coltype)
         return self.to_string(value)
 
 

data_diff/sqeleton/databases/base.py

@@ -164,6 +164,10 @@ def concat(self, items: List[str]) -> str:
         joined_exprs = ", ".join(items)
         return f"concat({joined_exprs})"
 
+    def to_comparable(self, value: str, coltype: ColType) -> str:
+        """Ensure that the expression is comparable in ``IS DISTINCT FROM``."""
+        return value
+
     def is_distinct_from(self, a: str, b: str) -> str:
         return f"{a} is distinct from {b}"
 
@@ -228,7 +232,7 @@ def parse_type(
         """ """
 
         cls = self._parse_type_repr(type_repr)
-        if not cls:
+        if cls is None:
             return UnknownColType(type_repr)
 
         if issubclass(cls, TemporalType):

data_diff/sqeleton/databases/bigquery.py

@@ -1,6 +1,8 @@
-from typing import List, Union
+import re
+from typing import Any, List, Union
 from ..abcs.database_types import (
-    Timestamp,
+    Array,
+    ColType, Struct, Timestamp,
     Datetime,
     Integer,
     Decimal,
@@ -10,6 +12,7 @@
     FractionalType,
     TemporalType,
     Boolean,
+    UnknownColType,
 )
 from ..abcs.mixins import (
     AbstractMixin_MD5,
@@ -36,6 +39,7 @@ def md5_as_int(self, s: str) -> str:
 
 
 class Mixin_NormalizeValue(AbstractMixin_NormalizeValue):
+
     def normalize_timestamp(self, value: str, coltype: TemporalType) -> str:
         if coltype.rounds:
             timestamp = f"timestamp_micros(cast(round(unix_micros(cast({value} as timestamp))/1000000, {coltype.precision})*1000000 as int))"
@@ -57,6 +61,20 @@ def normalize_number(self, value: str, coltype: FractionalType) -> str:
     def normalize_boolean(self, value: str, _coltype: Boolean) -> str:
         return self.to_string(f"cast({value} as int)")
 
+    def normalize_array(self, value: str, _coltype: Array) -> str:
+        # BigQuery is unable to compare arrays & structs with ==/!=/distinct from, e.g.:
+        #   Got error: 400 Grouping is not defined for arguments of type ARRAY<INT64> at …
+        # So we do the best effort and compare it as strings, hoping that the JSON forms
+        # match on both sides: i.e. have properly ordered keys, same spacing, same quotes, etc.
+        return f"to_json_string({value})"
+
+    def normalize_struct(self, value: str, _coltype: Struct) -> str:
+        # BigQuery is unable to compare arrays & structs with ==/!=/distinct from, e.g.:
+        #   Got error: 400 Grouping is not defined for arguments of type ARRAY<INT64> at …
+        # So we do the best effort and compare it as strings, hoping that the JSON forms
+        # match on both sides: i.e. have properly ordered keys, same spacing, same quotes, etc.
+        return f"to_json_string({value})"
+
 
 class Mixin_Schema(AbstractMixin_Schema):
     def list_tables(self, table_schema: str, like: Compilable = None) -> Compilable:
@@ -117,6 +135,8 @@ class Dialect(BaseDialect, Mixin_Schema):
         # Boolean
         "BOOL": Boolean,
     }
+    TYPE_ARRAY_RE = re.compile(r'ARRAY<(.+)>')
+    TYPE_STRUCT_RE = re.compile(r'STRUCT<(.+)>')
     MIXINS = {Mixin_Schema, Mixin_MD5, Mixin_NormalizeValue, Mixin_TimeTravel, Mixin_RandomSample}
 
     def random(self) -> str:
@@ -134,6 +154,35 @@ def type_repr(self, t) -> str:
         except KeyError:
             return super().type_repr(t)
 
+    def parse_type(
+        self,
+        table_path: DbPath,
+        col_name: str,
+        type_repr: str,
+        *args: Any,  # pass-through args
+        **kwargs: Any,  # pass-through args
+    ) -> ColType:
+        col_type = super().parse_type(table_path, col_name, type_repr, *args, **kwargs)
+        if isinstance(col_type, UnknownColType):
+
+            m = self.TYPE_ARRAY_RE.fullmatch(type_repr)
+            if m:
+                item_type = self.parse_type(table_path, col_name, m.group(1), *args, **kwargs)
+                col_type = Array(item_type=item_type)
+
+            m = self.TYPE_STRUCT_RE.fullmatch(type_repr)
+            if m:
+                col_type = Struct()
+
+        return col_type
+
+    def to_comparable(self, value: str, coltype: ColType) -> str:
+        """Ensure that the expression is comparable in ``IS DISTINCT FROM``."""
+        if isinstance(coltype, (Array, Struct)):
+            return self.normalize_value_by_type(value, coltype)
+        else:
+            return super().to_comparable(value, coltype)
+
     def set_timezone_to_utc(self) -> str:
         raise NotImplementedError()
 

data_diff/sqeleton/queries/ast_classes.py

@@ -352,7 +352,9 @@ class IsDistinctFrom(ExprNode, LazyOps):
     type = bool
 
     def compile(self, c: Compiler) -> str:
-        return c.dialect.is_distinct_from(c.compile(self.a), c.compile(self.b))
+        a = c.dialect.to_comparable(c.compile(self.a), self.a.type)
+        b = c.dialect.to_comparable(c.compile(self.b), self.b.type)
+        return c.dialect.is_distinct_from(a, b)
 
 
 @dataclass(eq=False, order=False)