Fix BigQueryExampleGen failing without custom_config #6307

PiperOrigin-RevId: 600997877
tensorflow · Jan 24, 2024 · 3a56dbf · 3a56dbf
1 parent df5118c
commit 3a56dbf
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/RELEASE.md b/RELEASE.md
@@ -65,6 +65,7 @@
 * Fixed a synchronization bug in google_cloud_ai_platform tuner.
 * Print best tuning trials only from the chief worker of google_cloud_ai_platform tuner.
 * Add a kpf dependency in the docker-image extra packages.
+* Fix BigQueryExampleGen failing without custom_config.
 
 ## Dependency Updates
 | Package Name | Version Constraints | Previously (in `v1.14.0`) | Comments |

diff --git a/tfx/extensions/google_cloud_big_query/example_gen/executor.py b/tfx/extensions/google_cloud_big_query/example_gen/executor.py
@@ -65,11 +65,12 @@ def _BigQueryToExample(pipeline: beam.Pipeline, exec_properties: Dict[str, Any],
   """
   project = utils.parse_gcp_project(exec_properties['_beam_pipeline_args'])
   converter = _BigQueryConverter(split_pattern, project)
-  big_query_custom_config = (
-      json.loads(exec_properties['custom_config'])
-      if 'custom_config' in exec_properties
-      else None
-  )
+  big_query_custom_config = None
+  if (
+      'custom_config' in exec_properties
+      and exec_properties['custom_config'] is not None
+  ):
+    big_query_custom_config = json.loads(exec_properties['custom_config'])
 
   return (
       pipeline