feat(bigquery): non-nullable schema support for embedded fields in st…

…ruct (ibis-project#10189)
deepyaman · Sep 21, 2024 · b16db6a · b16db6a
1 parent c62efce
commit b16db6a
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/ibis/backends/bigquery/tests/unit/test_datatypes.py b/ibis/backends/bigquery/tests/unit/test_datatypes.py
@@ -24,6 +24,11 @@
         param(dt.string, "STRING", id="string"),
         param(dt.Array(dt.int64), "ARRAY<INT64>", id="array<int64>"),
         param(dt.Array(dt.string), "ARRAY<STRING>", id="array<string>"),
+        param(
+            dt.Struct({"a": dt.String(nullable=False)}),
+            "STRUCT<`a` STRING NOT NULL>",
+            id="struct<a: !string>",
+        ),
         param(
             dt.Struct.from_tuples(
                 [("a", dt.int64), ("b", dt.string), ("c", dt.Array(dt.string))]

diff --git a/ibis/backends/sql/datatypes.py b/ibis/backends/sql/datatypes.py
@@ -944,6 +944,24 @@ def _from_ibis_GeoSpatial(cls, dtype: dt.GeoSpatial) -> sge.DataType:
                 f"Current geotype: {dtype.geotype}, Current srid: {dtype.srid}"
             )
 
+    @classmethod
+    def _from_ibis_Struct(cls, dtype: dt.Struct) -> sge.DataType:
+        fields = [
+            sge.ColumnDef(
+                # always quote struct fields to allow reserved words as field names
+                this=sg.to_identifier(name, quoted=True),
+                # Bigquery supports embeddable nulls
+                kind=cls.from_ibis(field),
+                constraints=(
+                    None
+                    if field.nullable
+                    else [sge.ColumnConstraint(kind=sge.NotNullColumnConstraint())]
+                ),
+            )
+            for name, field in dtype.items()
+        ]
+        return sge.DataType(this=typecode.STRUCT, expressions=fields, nested=True)
+
 
 class BigQueryUDFType(BigQueryType):
     @classmethod