Add and test BQ boolean support.

PiperOrigin-RevId: 481679118

Author: brandonmayer

tensorflow_gnn/data/unigraph.py

@@ -610,8 +610,9 @@ def row_to_keyed_example(self, row: Mapping[str, Any]) -> Any:
       example_feature = example.features.feature[tf_feature_name]
       if feature.dtype == tf.float32.as_datatype_enum:
         example_feature.float_list.value.append(feature_value)
-      elif feature.dtype == tf.int64.as_datatype_enum:
-        example_feature.int64_list.value.append(feature_value)
+      elif (feature.dtype == tf.int64.as_datatype_enum or
+            feature.dtype == tf.bool.as_datatype_enum):
+        example_feature.int64_list.value.append(int(feature_value))
       elif feature.dtype == tf.string.as_datatype_enum:
         example_feature.bytes_list.value.append(
             row.get(feature_name, "").encode("utf-8"))

tensorflow_gnn/data/unigraph_test.py

@@ -188,18 +188,24 @@ def test_bigquery_row_to_keyed_example_node_set(self):
             dtype: DT_INT64
           }
         }
-        features: {
+        features {
           key: "float_feature"
           value {
               dtype: DT_FLOAT
           }
         }
-        features: {
+        features {
             key: "string_feature"
             value {
                 dtype: DT_STRING
             }
         }
+        features {
+            key: "bool_feature"
+            value {
+                dtype: DT_BOOL
+            }
+        }
         metadata {
             bigquery {
                 table_spec {
@@ -224,12 +230,14 @@ def fake_bq_reader(**unused_kwargs):
           "id": "id1",
           "string_feature": "a",
           "float_feature": 1.0,
-          "int_feature": 2
+          "int_feature": 2,
+          "bool_feature": True
       }, {
           "id": "id2",
           "string_feature": "b",
           "int_feature": 3,
-          "float_feature": 4.0
+          "float_feature": 4.0,
+          "bool_feature": False
       }])
 
     with test_pipeline.TestPipeline() as pipeline:
@@ -277,6 +285,14 @@ def fake_bq_reader(**unused_kwargs):
                                       }
                                     }
                                   }
+                                  feature {
+                                      key: "bool_feature"
+                                      value {
+                                        int64_list {
+                                          value: 1
+                                        }
+                                      }
+                                  }
                               }""", tf.train.Example())),
                          (b"id2",
                           text_format.Parse(
@@ -313,6 +329,14 @@ def fake_bq_reader(**unused_kwargs):
                                       }
                                     }
                                   }
+                                  feature {
+                                      key: "bool_feature"
+                                      value {
+                                          int64_list {
+                                              value: 0
+                                          }
+                                      }
+                                  }
                                 }""", tf.train.Example()))]))
 
   def test_bigquery_row_to_keyed_example_edge_set(self):