[Dist] Add support for Parquet-formatted edges files, remove some assumptions...

[Dist] Add support for Parquet-formatted edges files, remove some assumptions on edge file number. (#5051)

* [Dist] Add support for Parquet-formatted edges files, remove some assumptions on edge file number.

* [Dist] Add parquet edges option to unit tests.
Co-authored-by: xiang song(charlie.song) <classicxsong@gmail.com>

tests/tools/test_dist_part.py

@@ -5,6 +5,7 @@ import tempfile
 import numpy as np
 import pytest
 import torch
+import pyarrow.parquet as pq
 from utils import create_chunked_dataset
 
 from distpartitioning import array_readwriter
@@ -80,6 +81,7 @@ def _verify_graph_feats(
 def _test_chunk_graph(
     num_chunks,
     data_fmt = 'numpy',
+    edges_fmt = 'csv',
     num_chunks_nodes = None,
     num_chunks_edges = None,
     num_chunks_node_data = None,
@@ -88,7 +90,7 @@ def _test_chunk_graph(
     with tempfile.TemporaryDirectory() as root_dir:
 
         g = create_chunked_dataset(root_dir, num_chunks,
-                data_fmt=data_fmt,
+                data_fmt=data_fmt, edges_fmt=edges_fmt,
                 num_chunks_nodes=num_chunks_nodes,
                 num_chunks_edges=num_chunks_edges,
                 num_chunks_node_data=num_chunks_node_data,
@@ -117,11 +119,17 @@ def _test_chunk_graph(
                     output_edge_index_dir, f'{c_etype_str}{i}.txt'
                 )
                 assert os.path.isfile(fname)
-                with open(fname, "r") as f:
-                    header = f.readline()
-                    num1, num2 = header.rstrip().split(" ")
-                    assert isinstance(int(num1), int)
-                    assert isinstance(int(num2), int)
+                if edges_fmt == 'csv':
+                    with open(fname, "r") as f:
+                        header = f.readline()
+                        num1, num2 = header.rstrip().split(" ")
+                        assert isinstance(int(num1), int)
+                        assert isinstance(int(num2), int)
+                elif edges_fmt == 'parquet':
+                    metadata = pq.read_metadata(fname)
+                    assert metadata.num_columns == 2
+                else:
+                    assert False, f"Invalid edges_fmt: {edges_fmt}"
 
         # check node/edge_data
         suffix = 'npy' if data_fmt=='numpy' else 'parquet'
@@ -179,8 +187,9 @@ def _test_chunk_graph(
 
 @pytest.mark.parametrize("num_chunks", [1, 8])
 @pytest.mark.parametrize("data_fmt", ['numpy', 'parquet'])
-def test_chunk_graph_basics(num_chunks, data_fmt):
-    _test_chunk_graph(num_chunks, data_fmt=data_fmt)
+@pytest.mark.parametrize("edges_fmt", ['csv', 'parquet'])
+def test_chunk_graph_basics(num_chunks, data_fmt, edges_fmt):
+    _test_chunk_graph(num_chunks, data_fmt=data_fmt, edges_fmt=edges_fmt)
 
 
 @pytest.mark.parametrize(

tests/tools/utils.py

@@ -76,7 +76,7 @@ def _initialize_num_chunks(g, num_chunks, kwargs=None):
 
 
 def _chunk_graph(
-    g, name, ndata_paths, edata_paths, num_chunks, data_fmt, **kwargs
+    g, name, ndata_paths, edata_paths, num_chunks, data_fmt, edges_format, **kwargs
 ):
     # First deal with ndata and edata that are homogeneous
     # (i.e. not a dict-of-dict)
@@ -138,7 +138,12 @@ def _chunk_graph(
             etypestr = etypestrs[etype]
             logging.info("Chunking edge index for %s" % etypestr)
             edges_meta = {}
-            fmt_meta = {"name": "csv", "delimiter": " "}
+            if edges_format == 'csv':
+                fmt_meta = {"name": edges_format, "delimiter": " "}
+            elif edges_format == 'parquet':
+                fmt_meta = {"name": edges_format}
+            else:
+                raise RuntimeError(f"Invalid edges_fmt: {edges_format}")
             edges_meta["format"] = fmt_meta
 
             srcdst = torch.stack(g.edges(etype=etype), 1)
@@ -244,6 +249,7 @@ def chunk_graph(
     num_chunks,
     output_path,
     data_fmt="numpy",
+    edges_fmt='csv',
     **kwargs,
 ):
     """
@@ -281,11 +287,12 @@ def chunk_graph(
             edata[key] = os.path.abspath(edata[key])
     with setdir(output_path):
         _chunk_graph(
-            g, name, ndata_paths, edata_paths, num_chunks, data_fmt, **kwargs
+            g, name, ndata_paths, edata_paths, num_chunks, data_fmt, edges_fmt, **kwargs
         )
 
 
-def create_chunked_dataset(root_dir, num_chunks, data_fmt="numpy", **kwargs):
+def create_chunked_dataset(
+    root_dir, num_chunks, data_fmt="numpy", edges_fmt='csv', **kwargs):
     """
     This function creates a sample dataset, based on MAG240 dataset.
 
@@ -523,6 +530,7 @@ def create_chunked_dataset(root_dir, num_chunks, data_fmt="numpy", **kwargs):
         num_chunks=num_chunks,
         output_path=output_dir,
         data_fmt=data_fmt,
+        edges_fmt=edges_fmt,
         **kwargs,
     )
     print("Done with creating chunked graph")

tools/distpartitioning/constants.py

@@ -25,6 +25,7 @@ STR_EDGE_TYPE = "edge_type"
 STR_NUM_EDGES_PER_CHUNK = "num_edges_per_chunk"
 STR_EDGES = "edges"
 STR_FORMAT = "format"
+STR_FORMAT_DELIMITER = "delimiter"
 STR_DATA = "data"
 STR_NODE_DATA = "node_data"
 STR_EDGE_DATA = "edge_data"

tools/distpartitioning/dataset_utils.py

@@ -4,6 +4,7 @@ import gc
 
 import numpy as np
 import pyarrow
+import pyarrow.parquet as pq
 import torch
 import torch.distributed as dist
 
@@ -501,8 +502,6 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
         edge_datadict[col] = []
 
     for etype_name, etype_info in edge_data.items():
-        assert etype_info[constants.STR_FORMAT][constants.STR_NAME] == constants.STR_CSV
-
         edge_info = etype_info[constants.STR_DATA]
 
         #edgetype strings are in canonical format, src_node_type:edge_type:dst_node_type
@@ -528,16 +527,24 @@ def get_dataset(input_dir, graph_name, rank, world_size, num_parts, schema_map):
                 edge_file = os.path.join(input_dir, edge_file)
             logging.info(f'Loading edges of etype[{etype_name}] from {edge_file}')
 
-            read_options=pyarrow.csv.ReadOptions(use_threads=True, block_size=4096, autogenerate_column_names=True)
-            parse_options=pyarrow.csv.ParseOptions(delimiter=' ')
-            with pyarrow.csv.open_csv(edge_file, read_options=read_options, parse_options=parse_options) as reader:
-                for next_chunk in reader:
-                    if next_chunk is None:
-                        break
-
-                    next_table = pyarrow.Table.from_batches([next_chunk])
-                    src_ids.append(next_table['f0'].to_numpy())
-                    dst_ids.append(next_table['f1'].to_numpy())
+            if etype_info[constants.STR_FORMAT][constants.STR_NAME] == constants.STR_CSV:
+                read_options=pyarrow.csv.ReadOptions(use_threads=True, block_size=4096, autogenerate_column_names=True)
+                parse_options=pyarrow.csv.ParseOptions(delimiter=' ')
+                with pyarrow.csv.open_csv(edge_file, read_options=read_options, parse_options=parse_options) as reader:
+                    for next_chunk in reader:
+                        if next_chunk is None:
+                            break
+
+                        next_table = pyarrow.Table.from_batches([next_chunk])
+                        src_ids.append(next_table['f0'].to_numpy())
+                        dst_ids.append(next_table['f1'].to_numpy())
+            elif etype_info[constants.STR_FORMAT][constants.STR_NAME] == constants.STR_PARQUET:
+                data_df = pq.read_table(edge_file)
+                data_df = data_df.rename_columns(["f0", "f1"])
+                src_ids.append(data_df['f0'].to_numpy())
+                dst_ids.append(data_df['f1'].to_numpy())
+            else:
+                raise ValueError(f'Unknown edge format {etype_info[constants.STR_FORMAT][constants.STR_NAME]} for edge type {etype_name}')
 
         src_ids = np.concatenate(src_ids)
         dst_ids = np.concatenate(dst_ids)

tools/distpartitioning/parmetis_preprocess.py

@@ -7,6 +7,7 @@ from pathlib import Path
 import numpy as np
 import pyarrow
 import pyarrow.csv as csv
+import pyarrow.parquet as pq
 import torch
 import torch.distributed as dist
 
@@ -77,7 +78,8 @@ def gen_edge_files(schema_map, output):
     num_parts = len(schema_map[constants.STR_NUM_EDGES_PER_CHUNK][0])
     for etype_name, etype_info in edge_data.items():
 
-        edge_info = etype_info[constants.STR_DATA]
+        edges_format = etype_info[constants.STR_FORMAT][constants.STR_NAME]
+        edge_data_files = etype_info[constants.STR_DATA]
 
         # ``edgetype`` strings are in canonical format, src_node_type:edge_type:dst_node_type
         tokens = etype_name.split(":")
@@ -87,29 +89,42 @@ def gen_edge_files(schema_map, output):
         rel_name = tokens[1]
         dst_ntype_name = tokens[2]
 
-        data_df = csv.read_csv(
-            edge_info[rank],
-            read_options=pyarrow.csv.ReadOptions(
-                autogenerate_column_names=True
-            ),
-            parse_options=pyarrow.csv.ParseOptions(delimiter=" "),
-        )
-        data_f0 = data_df["f0"].to_numpy()
-        data_f1 = data_df["f1"].to_numpy()
+        def convert_to_numpy_and_write_back(data_df):
+            data_f0 = data_df["f0"].to_numpy()
+            data_f1 = data_df["f1"].to_numpy()
 
-        global_src_id = data_f0 + ntype_gnid_offset[src_ntype_name][0, 0]
-        global_dst_id = data_f1 + ntype_gnid_offset[dst_ntype_name][0, 0]
-        cols = [global_src_id, global_dst_id]
-        col_names = ["global_src_id", "global_dst_id"]
+            global_src_id = data_f0 + ntype_gnid_offset[src_ntype_name][0, 0]
+            global_dst_id = data_f1 + ntype_gnid_offset[dst_ntype_name][0, 0]
+            cols = [global_src_id, global_dst_id]
+            col_names = ["global_src_id", "global_dst_id"]
 
-        out_file = edge_info[rank].split("/")[-1]
-        out_file = os.path.join(outdir, "edges_{}".format(out_file))
-        options = csv.WriteOptions(include_header=False, delimiter=" ")
-        options.delimiter = " "
+            out_file = edge_data_files[rank].split("/")[-1]
+            out_file = os.path.join(outdir, "edges_{}".format(out_file))
 
-        csv.write_csv(
-            pyarrow.Table.from_arrays(cols, names=col_names), out_file, options
-        )
+            # TODO(thvasilo): We should support writing to the same format as the input
+            options = csv.WriteOptions(include_header=False, delimiter=" ")
+            options.delimiter = " "
+            csv.write_csv(
+                pyarrow.Table.from_arrays(cols, names=col_names), out_file, options
+            )
+            return out_file
+
+        if edges_format == constants.STR_CSV:
+            delimiter = etype_info[constants.STR_FORMAT][constants.STR_FORMAT_DELIMITER]
+            data_df = csv.read_csv(
+                edge_data_files[rank],
+                read_options=pyarrow.csv.ReadOptions(
+                    autogenerate_column_names=True
+                ),
+                parse_options=pyarrow.csv.ParseOptions(delimiter=delimiter),
+            )
+        elif edges_format == constants.STR_PARQUET:
+            data_df = pq.read_table(edge_data_files[rank])
+            data_df = data_df.rename_columns(["f0", "f1"])
+        else:
+            raise NotImplementedError(f"Unknown edge format {edges_format}")
+
+        out_file = convert_to_numpy_and_write_back(data_df)
         edge_files.append(out_file)
 
     return edge_files
@@ -277,7 +292,8 @@ def gen_parmetis_input_args(params, schema_map):
     """
 
     num_nodes_per_chunk = schema_map[constants.STR_NUM_NODES_PER_CHUNK]
-    num_parts = len(num_nodes_per_chunk[0])
+    # TODO: This makes the assumption that all node files have the same number of chunks
+    num_node_parts = len(num_nodes_per_chunk[0])
     ntypes_ntypeid_map, ntypes, ntid_ntype_map = get_node_types(schema_map)
     type_nid_dict, ntype_gnid_offset = get_idranges(
         schema_map[constants.STR_NODE_TYPE],
@@ -319,7 +335,7 @@ def gen_parmetis_input_args(params, schema_map):
     os.makedirs(outdir, exist_ok=True)
     for ntype_id, ntype_name in ntid_ntype_map.items():
         global_nid_offset = ntype_gnid_offset[ntype_name][0, 0]
-        for r in range(num_parts):
+        for r in range(num_node_parts):
             type_start, type_end = (
                 type_nid_dict[ntype_name][r][0],
                 type_nid_dict[ntype_name][r][1],
@@ -345,9 +361,9 @@ def gen_parmetis_input_args(params, schema_map):
     edge_data = schema_map[constants.STR_EDGES]
     edge_files = []
     for etype_name, etype_info in edge_data.items():
-        edge_info = etype_info[constants.STR_DATA]
-        for r in range(num_parts):
-            out_file = edge_info[r].split("/")[-1]
+        edge_data_files = etype_info[constants.STR_DATA]
+        for edge_file_path in edge_data_files:
+            out_file = os.path.basename(edge_file_path)
             out_file = os.path.join(outdir, "edges_{}".format(out_file))
             edge_files.append(out_file)