diff --git a/benchmark/in_process_io.py b/benchmark/in_process_io.py index f4c67c3..d8ad619 100644 --- a/benchmark/in_process_io.py +++ b/benchmark/in_process_io.py @@ -125,7 +125,7 @@ def run(self): graph=self._graph, schema=self._schema, path=os.path.join(self.work_dir, f"WriteGFGraphInMemory"), - max_num_shards=20, + num_shards=20, ) @override diff --git a/dgf/src/io/graph_in_memory.py b/dgf/src/io/graph_in_memory.py index 8b06786..34552ef 100644 --- a/dgf/src/io/graph_in_memory.py +++ b/dgf/src/io/graph_in_memory.py @@ -493,7 +493,7 @@ def write_graph( schema: schema_lib.GraphSchema, path: str, verbose: bool = False, - max_num_shards: Optional[int] = None, + num_shards: Optional[int] = None, compression: str = "snappy", container: ( str | gf_metadata_lib.Container @@ -511,8 +511,8 @@ def write_graph( schema: The schema of the graph. path: The path to the GF Graph directory. verbose: If True, print progress information. - max_num_shards: If provided, limits the maximum number of shards used when - writing the Parquet files for each node and edge set. + num_shards: If provided, the number of shards used when writing the files + for each node and edge set. If None, the number of shards is estimated. compression: Compression algorithm for Parquet files. container: The file format for writing data blocks. """ @@ -553,9 +553,12 @@ def write_graph( if verbose: log.info("Writing nodeset %s to %s", nodeset_name, node_dir) - num_shards, _ = shard_lib.estimate_num_node_shards(node_set.num_nodes or 0) - if max_num_shards is not None: - num_shards = min(num_shards, max_num_shards) + if num_shards is not None: + effective_num_shards = num_shards + else: + effective_num_shards, _ = shard_lib.estimate_num_node_shards( + node_set.num_nodes or 0 + ) features_to_write = node_set.features if features_to_write is None: @@ -566,7 +569,7 @@ def write_graph( nodeset_name, node_dir, nodeset_schema.features, - num_shards, + effective_num_shards, verbose, compression=compression, container_type=container, @@ -609,16 +612,19 @@ def write_graph( edgeset_schema.target ].features[target_primary_key] - num_shards, _ = shard_lib.estimate_num_edge_shards(edge_set.num_edges()) - if max_num_shards is not None: - num_shards = min(num_shards, max_num_shards) + if num_shards is not None: + effective_num_shards = num_shards + else: + effective_num_shards, _ = shard_lib.estimate_num_edge_shards( + edge_set.num_edges() + ) _write_container( features_to_write, edgeset_name, edge_dir, features_schema, - num_shards, + effective_num_shards, verbose, compression=compression, container_type=container, diff --git a/dgf/src/io/graph_in_memory_test.py b/dgf/src/io/graph_in_memory_test.py index ccd8bbd..035cbc6 100644 --- a/dgf/src/io/graph_in_memory_test.py +++ b/dgf/src/io/graph_in_memory_test.py @@ -14,6 +14,7 @@ import os import tempfile +from unittest import mock from absl.testing import absltest from absl.testing import parameterized from dgf.src.data import distributed_graph @@ -182,12 +183,10 @@ def test_write_and_read_graph_timestamp(self): self.assertIsNone(loaded_graph2.timestamp) @parameterized.product( - max_num_shards=[1, 2, 3], + num_shards=[1, 2, 3], container=["PARQUET", "TF_RECORD", "RECORDIO"], ) - def test_write_and_read_sharded_graph( - self, max_num_shards: int, container: str - ): + def test_write_and_read_sharded_graph(self, num_shards: int, container: str): with tempfile.TemporaryDirectory() as tmpdir: output_path = os.path.join(tmpdir, "sharded_gf_graph") num_nodes = 2500 @@ -255,7 +254,7 @@ def test_write_and_read_sharded_graph( graph, schema, output_path, - max_num_shards=max_num_shards, + num_shards=num_shards, container=container, ) loaded_graph, loaded_schema = gf_graph_in_memory.read_graph(output_path) @@ -270,7 +269,7 @@ def test_write_and_read_sharded_graph( extension = gf_graph_in_memory.get_extension( gf_metadata_lib.Container(container) ) - expected_num_shards = min(3, max_num_shards) + expected_num_shards = num_shards expected_files = ["/schema.json", "/metadata.json"] for s in range(expected_num_shards): expected_files.append( @@ -288,6 +287,47 @@ def test_write_and_read_sharded_graph( ) self.assertSameElements(sorted(actual_files), sorted(expected_files)) + def test_write_graph_num_shards_skips_estimate(self): + with tempfile.TemporaryDirectory() as tmpdir: + output_path = os.path.join(tmpdir, "output_gf_graph") + in_memory_graph = gen_test_graph.generate_in_memory_graph( + node_ids=True, edge_ids=True + ) + schema = gen_test_graph.generate_schema( + node_ids=True, edge_ids=True, semantic=True + ) + + with mock.patch.object( + gf_graph_in_memory.shard_lib, + "estimate_num_node_shards", + autospec=True, + ) as mock_node_shards, mock.patch.object( + gf_graph_in_memory.shard_lib, + "estimate_num_edge_shards", + autospec=True, + ) as mock_edge_shards: + gf_graph_in_memory.write_graph( + in_memory_graph, schema, output_path, num_shards=2 + ) + mock_node_shards.assert_not_called() + mock_edge_shards.assert_not_called() + + output_path2 = os.path.join(tmpdir, "output_gf_graph2") + with mock.patch.object( + gf_graph_in_memory.shard_lib, + "estimate_num_node_shards", + return_value=(1, 10000), + ) as mock_node_shards, mock.patch.object( + gf_graph_in_memory.shard_lib, + "estimate_num_edge_shards", + return_value=(1, 10000), + ) as mock_edge_shards: + gf_graph_in_memory.write_graph( + in_memory_graph, schema, output_path2, num_shards=None + ) + self.assertTrue(mock_node_shards.called) + self.assertTrue(mock_edge_shards.called) + def _canonicalize_graph( graph: in_memory_graph_lib.InMemoryGraph, diff --git a/dgf/src/util/shard.py b/dgf/src/util/shard.py index 293228e..48c459a 100644 --- a/dgf/src/util/shard.py +++ b/dgf/src/util/shard.py @@ -16,13 +16,15 @@ from __future__ import annotations +import math import re from typing import List, Optional, Tuple from dgf.src.util.weak_dep.weak_dep_tensorflow import tf -NUM_NODES_PER_SHARD = 1000 -NUM_EDGES_PER_SHARD = 1000 +MAX_NUM_SHARDS = 100 +NUM_NODES_PER_SHARD = 10000 +NUM_EDGES_PER_SHARD = 10000 def expand_output_paths(path: str, num_shards: Optional[int]) -> List[str]: @@ -271,13 +273,11 @@ def estimate_num_node_shards(num_nodes: int) -> tuple[int, int]: Returns: The estimated number of shards, and the number of nodes per shard. """ - num_shards = (num_nodes // NUM_NODES_PER_SHARD) + ( - num_nodes % NUM_NODES_PER_SHARD > 0 - ) - if num_shards <= 100: + num_shards = math.ceil(num_nodes / NUM_NODES_PER_SHARD) + if num_shards <= MAX_NUM_SHARDS: return num_shards, NUM_NODES_PER_SHARD else: - return 100, num_nodes // 100 + (num_nodes % 100 > 0) + return MAX_NUM_SHARDS, math.ceil(num_nodes / MAX_NUM_SHARDS) def estimate_num_edge_shards(num_edges: int) -> tuple[int, int]: @@ -291,10 +291,8 @@ def estimate_num_edge_shards(num_edges: int) -> tuple[int, int]: Returns: The estimated number of shards, and the number of edges per shard. """ - num_shards = (num_edges // NUM_EDGES_PER_SHARD) + ( - num_edges % NUM_EDGES_PER_SHARD > 0 - ) - if num_shards <= 100: + num_shards = math.ceil(num_edges / NUM_EDGES_PER_SHARD) + if num_shards <= MAX_NUM_SHARDS: return num_shards, NUM_EDGES_PER_SHARD else: - return 100, num_edges // 100 + (num_edges % 100 > 0) + return MAX_NUM_SHARDS, math.ceil(num_edges / MAX_NUM_SHARDS) diff --git a/dgf/src/util/shard_test.py b/dgf/src/util/shard_test.py index 48b5c6d..c23fb40 100644 --- a/dgf/src/util/shard_test.py +++ b/dgf/src/util/shard_test.py @@ -180,13 +180,13 @@ def test_expand_output_paths(self): shard.expand_output_paths("m*.ext", num_shards=2) @parameterized.parameters( - (100, 1, 1000), - (1000, 1, 1000), - (1001, 2, 1000), - (1000 * 100, 100, 1000), - (1000 * 100 + 1, 100, 1001), - (1000 * 200, 100, 2000), - (1000 * 200 + 1, 100, 2001), + (100, 1, 10000), + (10000, 1, 10000), + (10001, 2, 10000), + (10000 * 100, 100, 10000), + (10000 * 100 + 1, 100, 10001), + (10000 * 200, 100, 20000), + (10000 * 200 + 1, 100, 20001), ) def test_estimate_num_node_shards( self, num_nodes, expected_num_shards, expected_num_nodes_per_shard @@ -197,20 +197,20 @@ def test_estimate_num_node_shards( ) @parameterized.parameters( - (100, 1, 1000), - (1000, 1, 1000), - (1001, 2, 1000), - (1000 * 100, 100, 1000), - (1000 * 100 + 1, 100, 1001), - (1000 * 200, 100, 2000), - (1000 * 200 + 1, 100, 2001), + (100, 1, 10000), + (10000, 1, 10000), + (10001, 2, 10000), + (10000 * 100, 100, 10000), + (10000 * 100 + 1, 100, 10001), + (10000 * 200, 100, 20000), + (10000 * 200 + 1, 100, 20001), ) def test_estimate_num_edge_shards( - self, num_edges, expected_num_shards, expected_num_edeg_per_shard + self, num_edges, expected_num_shards, expected_num_edges_per_shard ): self.assertEqual( shard.estimate_num_edge_shards(num_edges), - (expected_num_shards, expected_num_edeg_per_shard), + (expected_num_shards, expected_num_edges_per_shard), )