diff --git a/dgf/src/io/gcp/bigquery_graph.py b/dgf/src/io/gcp/bigquery_graph.py index c2d9eea..2d337e4 100644 --- a/dgf/src/io/gcp/bigquery_graph.py +++ b/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)) diff --git a/dgf/src/io/gcp/bigquery_graph_beam.py b/dgf/src/io/gcp/bigquery_graph_beam.py index 926a9f7..bd4aa90 100644 --- a/dgf/src/io/gcp/bigquery_graph_beam.py +++ b/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: diff --git a/dgf/src/io/gcp/bigquery_graph_test.py b/dgf/src/io/gcp/bigquery_graph_test.py index f013f23..ec9491f 100644 --- a/dgf/src/io/gcp/bigquery_graph_test.py +++ b/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__":