BigQuery IO cleanup

dgf/src/io/gcp/bigquery_graph.py

@@ -129,7 +129,6 @@ def _graph_element_table(
 
 def metadata_to_schema(
     bigquery_graph_metadata: bigquery_graph_metadata_lib.BigQueryGraphMetadata,
-    combine_as_json: bool,
 ) -> schema_lib.GraphSchema:
   """Converts BigQuery graph metadata to a GraphFlow schema."""
 
@@ -139,7 +138,7 @@ def metadata_to_schema(
         features=gcp_common_lib.infer_feature_set_schema(
             _graph_element_table(node_table.label_and_properties),  # pyrefly: ignore[bad-argument-type]
             node_table.key_columns,
-            combine_as_json,
+            combine_as_json=False,
         )
     )
 
@@ -151,7 +150,7 @@ def metadata_to_schema(
         features=gcp_common_lib.infer_feature_set_schema(
             _graph_element_table(edge_table.label_and_properties),  # pyrefly: ignore[bad-argument-type]
             [],  # Add support for edge ids,
-            combine_as_json,
+            combine_as_json=False,
         ),
     )
 
@@ -210,7 +209,6 @@ def read_bigquery_graph_schema(
     dataset: str,
     graph: str,
     *,
-    combine_as_json: bool = False,
     verbose: Union[int, bool] = True,
 ):
   """Reads the schema of a BigQuery graph into a GF schema.
@@ -231,7 +229,6 @@ def read_bigquery_graph_schema(
     project: The GCP project ID of the BigQuery Graph.
     dataset: The BQ dataset ID of the BigQuery Graph.
     graph: The ID of the BigQuery Graph.
-    combine_as_json: Whether to combine the features as JSON.
     verbose: Amount of verbose (0: no, 1 or true: a little, 2: a lot).
 
   Returns:
@@ -241,7 +238,7 @@ def read_bigquery_graph_schema(
   metadata = get_metadata(project, dataset, graph)
   if verbose >= 2:
     log.info("meta-data:\n%s", metadata)
-  return metadata_to_schema(metadata, combine_as_json)
+  return metadata_to_schema(metadata)
 
 
 def read_bigquery_graph(
@@ -251,7 +248,7 @@ def read_bigquery_graph(
     *,
     schema: Optional[schema_lib.GraphSchema] = None,
     work_dir: str,
-    combine_as_json: bool = False,
+    remove_dangling_edges: bool = False,
     max_workers: int = 10,
     verbose: Union[int, bool] = True,
 ) -> Tuple[in_memory_graph_lib.InMemoryGraph, schema_lib.GraphSchema]:
@@ -273,8 +270,12 @@ def read_bigquery_graph(
     project: The Google Cloud project ID.
     dataset: The BigQuery dataset ID.
     graph: The BigQuery graph ID.
+    schema: Optional GraphFlow schema. If None, the schema is inferred from
+      BigQuery graph metadata.
     work_dir: The working directory to use for the temporary storage.
-    combine_as_json: Whether to combine the features as JSON.
+    remove_dangling_edges: If False (default), fails if an edge is dangling
+      (i.e., it refers to non-existing nodes). If True, dangling edges are
+      removed.
     max_workers: The maximum number of workers to use for parallel processing.
     verbose: Amount of verbose (0: no, 1 or true: a little, 2: a lot).
 
@@ -294,14 +295,15 @@ def read_bigquery_graph(
         project=project,
         dataset=dataset,
         graph=graph,
-        combine_as_json=combine_as_json,
         max_workers=max_workers,
         verbose=verbose,
         schema=schema,
     )
 
     in_memory_graph, in_memory_schema = gf_graph_in_memory_lib.read_graph(
-        work_dir, verbose=verbose >= 1
+        work_dir,
+        remove_dangling_edges=remove_dangling_edges,
+        verbose=verbose >= 1,
     )
   finally:
     try:
@@ -319,7 +321,6 @@ def export_bigquery_to_disk(
     graph: str,
     *,
     schema: Optional[schema_lib.GraphSchema] = None,
-    combine_as_json: bool = False,
     max_workers: int = 10,
     verbose: Union[int, bool] = True,
 ):
@@ -342,7 +343,8 @@ def export_bigquery_to_disk(
     project: The Google Cloud project ID.
     dataset: The BigQuery dataset ID.
     graph: The BigQuery graph ID.
-    combine_as_json: Whether to combine the features as JSON.
+    schema: Optional GraphFlow schema. If None, the schema is inferred from
+      BigQuery graph metadata.
     max_workers: The maximum number of workers to use for parallel processing.
     verbose: Amount of verbose (0: no, 1 or true: a little, 2: a lot).
 
@@ -361,7 +363,7 @@ def export_bigquery_to_disk(
   # TODO(b/328622124): Add feature_shapes, feature_semantics, and
   # num_categorical_values to the function signature and pass them here.
   if schema is None:
-    schema = metadata_to_schema(metadata, combine_as_json)
+    schema = metadata_to_schema(metadata)
 
   if verbose >= 2:
     log.info("%s", print_schema_lib.print_schema(schema, return_output=True))

dgf/src/io/gcp/bigquery_graph_beam.py

@@ -69,7 +69,6 @@ def distributed_read_beam(
     dataset_id: str,
     graph_id: str,
     p: beam.pvalue.PBegin,
-    combine_as_json: bool = False,
 ) -> distributed_graph_lib.Graph:
   """Read BigQuery Graph via Beam and return a distributed GraphFlow graph.
 
@@ -78,13 +77,12 @@ def distributed_read_beam(
     dataset_id: The BQ dataset ID of the BigQuery Graph.
     graph_id: The ID of the BigQuery Graph.
     p: The Beam pipeline.
-    combine_as_json: Whether to combine the features as JSON.
 
   Returns:
     A distributed GraphFlow graph.
   """
   metadata = bigquery_graph.get_metadata(project_id, dataset_id, graph_id)
-  schema = bigquery_graph.metadata_to_schema(metadata, combine_as_json)
+  schema = bigquery_graph.metadata_to_schema(metadata)
 
   node_sets = {}
   for node_table in metadata.node_tables:

dgf/src/io/gcp/bigquery_graph_test.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for bigquery_graph."""
+"""Tests for BigQuery Graph reading, schema generation, and export."""
 
 import copy
 from unittest import mock
@@ -195,10 +195,7 @@ def test_graph_schema(self):
     metadata = bigquery_graph_metadata_lib.BigQueryGraphMetadata.from_dict(  # pyrefly: ignore[missing-attribute]
         infoschema_query_response_json
     )
-    schema = bigquery_graph.metadata_to_schema(
-        metadata,
-        combine_as_json=False,
-    )
+    schema = bigquery_graph.metadata_to_schema(metadata)
     self.assertIn("nodes", schema.node_sets)
     self.assertIn("biggraphs-poc.ogbn_arxiv_2.edges", schema.edge_sets)
 
@@ -258,6 +255,85 @@ def test_read_bigquery_graph(
         )
     )
     self.assertEqual(mock_read_graph.call_args[1]["verbose"], 1)
+    self.assertFalse(mock_read_graph.call_args[1]["remove_dangling_edges"])
+
+  @mock.patch("dgf.src.io.gcp.bigquery_graph.get_metadata")
+  @mock.patch("dgf.src.io.gcp.bigquery_graph._execute_query")
+  @mock.patch("dgf.src.io.gcp.parquet_export.create_export_sql")
+  @mock.patch("dgf.src.io.schema.write_schema")
+  @mock.patch("dgf.src.util.filesystem.open_write")
+  @mock.patch("dgf.src.io.graph_in_memory.read_graph")
+  def test_read_bigquery_graph_remove_dangling_edges(
+      self,
+      mock_read_graph,
+      mock_open_write,
+      mock_write_schema,
+      mock_create_export_sql,
+      mock_execute_query,
+      mock_load_metadata,
+  ):
+    del mock_open_write, mock_write_schema, mock_create_export_sql  # Unused.
+    metadata = bigquery_graph_metadata_lib.BigQueryGraphMetadata.from_dict(  # pyrefly: ignore[missing-attribute]
+        infoschema_query_response_json
+    )
+    mock_load_metadata.return_value = metadata
+    mock_execute_query.return_value = mock.Mock()
+    mock_read_graph.return_value = (mock.Mock(), mock.Mock())
+
+    bigquery_graph.read_bigquery_graph(
+        "project",
+        "dataset",
+        "graph",
+        work_dir="gs://bucket/prefix",
+        remove_dangling_edges=True,
+    )
+
+    self.assertTrue(mock_read_graph.call_args[1]["remove_dangling_edges"])
+
+  @mock.patch("dgf.src.io.gcp.bigquery_graph.get_metadata")
+  def test_read_bigquery_graph_schema(self, mock_load_metadata):
+    metadata = bigquery_graph_metadata_lib.BigQueryGraphMetadata.from_dict(  # pyrefly: ignore[missing-attribute]
+        infoschema_query_response_json
+    )
+    mock_load_metadata.return_value = metadata
+
+    schema = bigquery_graph.read_bigquery_graph_schema(
+        "project", "dataset", "graph"
+    )
+
+    self.assertIn("nodes", schema.node_sets)
+    self.assertIn("biggraphs-poc.ogbn_arxiv_2.edges", schema.edge_sets)
+    mock_load_metadata.assert_called_once_with("project", "dataset", "graph")
+
+  @mock.patch("dgf.src.io.gcp.bigquery_graph.get_metadata")
+  @mock.patch("dgf.src.io.gcp.bigquery_graph._execute_query")
+  @mock.patch("dgf.src.io.gcp.parquet_export.create_export_sql")
+  @mock.patch("dgf.src.io.schema.write_schema")
+  @mock.patch("dgf.src.util.filesystem.open_write")
+  def test_export_bigquery_to_disk(
+      self,
+      mock_open_write,
+      mock_write_schema,
+      mock_create_export_sql,
+      mock_execute_query,
+      mock_load_metadata,
+  ):
+    metadata = bigquery_graph_metadata_lib.BigQueryGraphMetadata.from_dict(  # pyrefly: ignore[missing-attribute]
+        infoschema_query_response_json
+    )
+    mock_load_metadata.return_value = metadata
+    mock_execute_query.return_value = mock.Mock()
+    mock_create_export_sql.return_value = "EXPORT DATA SQL"
+
+    bigquery_graph.export_bigquery_to_disk(
+        "gs://bucket/prefix", "project", "dataset", "graph"
+    )
+
+    mock_load_metadata.assert_called_once_with("project", "dataset", "graph")
+    self.assertEqual(mock_execute_query.call_count, 2)
+    mock_create_export_sql.assert_called()
+    mock_write_schema.assert_called_once()
+    mock_open_write.assert_called_once()
 
 
 if __name__ == "__main__":