Change how the user can control the number of shards when writting a graph.

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

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,

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,

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)

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),
     )