wehub-resource-sync
178 lines
6.2 KiB
Python
178 lines
6.2 KiB
Python
import os
|
|
import tempfile
|
|
from collections import namedtuple
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from distpartitioning import array_readwriter, constants
|
|
from distpartitioning.parmetis_preprocess import gen_edge_files
|
|
from distpartitioning.utils import generate_roundrobin_read_list
|
|
from numpy.testing import assert_array_equal
|
|
|
|
NODE_TYPE = "n1"
|
|
EDGE_TYPE = f"{NODE_TYPE}:e1:{NODE_TYPE}"
|
|
|
|
|
|
def _read_file(fname, fmt_name, fmt_delimiter):
|
|
"""Read a file
|
|
|
|
Parameters:
|
|
-----------
|
|
fname : string
|
|
filename of the input file to read
|
|
fmt_name : string
|
|
specifying whether it is a csv or a parquet file
|
|
fmt_delimiter : string
|
|
string specifying the delimiter used in the input file
|
|
"""
|
|
reader_fmt_meta = {
|
|
"name": fmt_name,
|
|
}
|
|
if fmt_name == constants.STR_CSV:
|
|
reader_fmt_meta["delimiter"] = fmt_delimiter
|
|
data_df = array_readwriter.get_array_parser(**reader_fmt_meta).read(fname)
|
|
return data_df
|
|
|
|
|
|
def _get_test_data(edges_dir, num_chunks, edge_fmt, edge_fmt_del):
|
|
"""Creates unit test input which are a set of edge files
|
|
in the following format "src_node_id<delimiter>dst_node_id"
|
|
|
|
Parameters:
|
|
-----------
|
|
edges_dir : str
|
|
folder where edge files are stored
|
|
num_chunks : int
|
|
no. of files to create for each edge type
|
|
edge_fmt : str, optional
|
|
to specify whether this file is csv or parquet
|
|
edge_fmt_del : str optional
|
|
delimiter to use in the edges file
|
|
|
|
Returns:
|
|
--------
|
|
dict :
|
|
dictionary created which represents the schema used for
|
|
creating the input dataset
|
|
"""
|
|
schema = {}
|
|
schema["num_nodes_per_type"] = [10]
|
|
schema["edge_type"] = [EDGE_TYPE]
|
|
schema["node_type"] = [NODE_TYPE]
|
|
|
|
edges = {}
|
|
edges[EDGE_TYPE] = {}
|
|
edges[EDGE_TYPE]["format"] = {}
|
|
edges[EDGE_TYPE]["format"]["name"] = edge_fmt
|
|
edges[EDGE_TYPE]["format"]["delimiter"] = edge_fmt_del
|
|
|
|
os.makedirs(edges_dir, exist_ok=True)
|
|
fmt_meta = {"name": edge_fmt}
|
|
if edge_fmt == "csv":
|
|
fmt_meta["delimiter"] = edge_fmt_del
|
|
|
|
edge_files = []
|
|
for idx in range(num_chunks):
|
|
path = os.path.join(edges_dir, f"test_file_{idx}.{fmt_meta['name']}")
|
|
array_parser = array_readwriter.get_array_parser(**fmt_meta)
|
|
edge_data = (
|
|
np.array([np.arange(10), np.arange(10)]).reshape(10, 2) + 10 * idx
|
|
)
|
|
array_parser.write(path, edge_data)
|
|
|
|
edge_files.append(path)
|
|
|
|
edges[EDGE_TYPE]["data"] = edge_files
|
|
schema["edges"] = edges
|
|
|
|
return schema
|
|
|
|
|
|
@pytest.mark.parametrize("num_chunks, num_parts", [[4, 1], [4, 2], [4, 4]])
|
|
@pytest.mark.parametrize("edges_fmt", ["csv", "parquet"])
|
|
@pytest.mark.parametrize("edges_delimiter", [" ", ","])
|
|
def test_gen_edge_files(num_chunks, num_parts, edges_fmt, edges_delimiter):
|
|
"""Unit test case for the function
|
|
tools/distpartitioning/parmetis_preprocess.py::gen_edge_files
|
|
|
|
Parameters:
|
|
-----------
|
|
num_chunks : int
|
|
no. of chunks the input graph needs to be partititioned into
|
|
num_parts : int
|
|
no. of partitions
|
|
edges_fmt : string
|
|
specifying the storage format for the edge files
|
|
edges_delimiter : string
|
|
specifying the delimiter used in the edge files
|
|
"""
|
|
# Create the input dataset
|
|
with tempfile.TemporaryDirectory() as root_dir:
|
|
|
|
# Create expected environment for test
|
|
input_dir = os.path.join(root_dir, "chunked-data")
|
|
output_dir = os.path.join(root_dir, "preproc_dir")
|
|
|
|
# Mock a parser object
|
|
fn_params = namedtuple("fn_params", "input_dir output_dir num_parts")
|
|
fn_params.input_dir = input_dir
|
|
fn_params.output_dir = output_dir
|
|
fn_params.num_parts = num_parts
|
|
|
|
# Create test files and get corresponding file schema
|
|
schema_map = _get_test_data(
|
|
input_dir, num_chunks, edges_fmt, edges_delimiter
|
|
)
|
|
edges_file_list = schema_map["edges"][EDGE_TYPE]["data"]
|
|
# This is breaking encapsulation, but no other good way to get file list
|
|
rank_assignments = generate_roundrobin_read_list(
|
|
len(edges_file_list), num_parts
|
|
)
|
|
|
|
# Get the global node id offsets for each node type
|
|
# There is only one node-type in the test graph
|
|
# which range from 0 thru 9.
|
|
ntype_gnid_offset = {}
|
|
ntype_gnid_offset[NODE_TYPE] = np.array([0, 10 * num_chunks]).reshape(
|
|
1, 2
|
|
)
|
|
|
|
# Iterate over no. of partitions
|
|
for rank in range(num_parts):
|
|
actual_results = gen_edge_files(rank, schema_map, fn_params)
|
|
|
|
# Get the original files
|
|
original_files = [
|
|
edges_file_list[file_idx] for file_idx in rank_assignments[rank]
|
|
]
|
|
|
|
# Validate the results with the baseline results
|
|
# Test 1. no. of files should have the same count per rank
|
|
assert len(original_files) == len(actual_results)
|
|
assert len(actual_results) > 0
|
|
|
|
# Test 2. Check the contents of each file and verify the
|
|
# file contents match with the expected results.
|
|
for actual_fname, original_fname in zip(
|
|
actual_results, original_files
|
|
):
|
|
# Check the actual file exists
|
|
assert os.path.isfile(actual_fname)
|
|
# Read both files and compare the edges
|
|
# Here note that the src and dst end points are global_node_ids
|
|
actual_data = _read_file(actual_fname, "csv", " ")
|
|
expected_data = _read_file(
|
|
original_fname, edges_fmt, edges_delimiter
|
|
)
|
|
|
|
# Subtract the global node id offsets, so that we get type node ids
|
|
# In the current unit test case, the graph has only one node-type.
|
|
# and this means that type-node-ids are same as the global-node-ids.
|
|
# Below two lines will take take into effect when the graphs have
|
|
# more than one node type.
|
|
actual_data[:, 0] -= ntype_gnid_offset[NODE_TYPE][0, 0]
|
|
actual_data[:, 1] -= ntype_gnid_offset[NODE_TYPE][0, 0]
|
|
|
|
# Verify that the contents are equal
|
|
assert_array_equal(expected_data, actual_data)
|