Fix BigQueryExampleGen failing without custom_config #6307

PiperOrigin-RevId: 600997877
tensorflow · Feb 1, 2024 · 5df14e8 · 5df14e8
1 parent 3aa2ad9
commit 5df14e8
Show file tree

Hide file tree

Showing 2 changed files with 4 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 failure 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,9 @@ 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_str := exec_properties.get('custom_config'):
+    big_query_custom_config = json.loads(custom_config_str)
 
   return (
       pipeline