Adding handling of Nulled lists to beam_row_from_dict (#33830)

* Adding handling of Nulled lists to beam_row_from_dict * Update sdks/python/apache_beam/io/gcp/bigquery_tools.py Co-authored-by: Ahmed Abualsaud <[email protected]> --------- Co-authored-by: Ahmed Abualsaud <[email protected]>
apache · Feb 3, 2025 · c7b2695 · c7b2695
1 parent df13ffe
commit c7b2695
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
@@ -1606,11 +1606,13 @@ def beam_row_from_dict(row: dict, schema):
     # This requires that each row has all the fields in the schema.
     # However, it's possible that some nullable fields don't appear in the row.
     # For this case, we create the field with a `None` value
-    if name not in row and mode == "NULLABLE":
+    # None is also set when a repeated field is missing as BigQuery
+    # converts Null Repeated fields to empty lists
+    if name not in row and mode != "REQUIRED":
       row[name] = None
 
     value = row[name]
-    if type in ["RECORD", "STRUCT"]:
+    if type in ["RECORD", "STRUCT"] and value:
       # if this is a list of records, we create a list of Beam Rows
       if mode == "REPEATED":
         list_of_beam_rows = []

diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
@@ -867,7 +867,15 @@ def test_dict_to_beam_row_all_types_repeated(self):
     self.assertEqual(expected_beam_row, beam_row_from_dict(dict_row, schema))
 
   def test_dict_to_beam_row_all_types_nullable(self):
-    schema = {"fields": self.get_schema_fields_with_mode("nullable")}
+    schema_fields_with_nested = [{
+        "name": "nested_record",
+        "type": "record",
+        "mode": "repeated",
+        "fields": self.get_schema_fields_with_mode("nullable")
+    }]
+    schema_fields_with_nested.extend(
+        self.get_schema_fields_with_mode("nullable"))
+    schema = {"fields": schema_fields_with_nested}
     dict_row = {k: None for k in self.DICT_ROW}
 
     # input dict row with missing nullable fields should still yield a full
@@ -876,6 +884,7 @@ def test_dict_to_beam_row_all_types_nullable(self):
     del dict_row['bool']
 
     expected_beam_row = beam.Row(
+        nested_record=None,
         str=None,
         bool=None,
         bytes=None,