wehub-resource-sync
217 lines
7.2 KiB
Python
217 lines
7.2 KiB
Python
"""Utility functions for GraphBolt."""
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import shutil
|
|
from typing import List, Union
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
from numpy.lib.format import read_array_header_1_0, read_array_header_2_0
|
|
|
|
|
|
def numpy_save_aligned(*args, **kwargs):
|
|
"""A wrapper for numpy.save(), ensures the array is stored 4KiB aligned."""
|
|
# https://github.com/numpy/numpy/blob/2093a6d5b933f812d15a3de0eafeeb23c61f948a/numpy/lib/format.py#L179
|
|
has_array_align = hasattr(np.lib.format, "ARRAY_ALIGN")
|
|
if has_array_align:
|
|
default_alignment = np.lib.format.ARRAY_ALIGN
|
|
# The maximum allowed alignment by the numpy code linked above is 4K.
|
|
# Most filesystems work with block sizes of 4K so in practice, the file
|
|
# size on the disk won't be larger.
|
|
np.lib.format.ARRAY_ALIGN = 4096
|
|
np.save(*args, **kwargs)
|
|
if has_array_align:
|
|
np.lib.format.ARRAY_ALIGN = default_alignment
|
|
|
|
|
|
def _read_torch_data(path):
|
|
return torch.load(path, weights_only=False)
|
|
|
|
|
|
def _read_numpy_data(path, in_memory=True):
|
|
if in_memory:
|
|
return torch.from_numpy(np.load(path))
|
|
return torch.as_tensor(np.load(path, mmap_mode="r+"))
|
|
|
|
|
|
def read_data(path, fmt, in_memory=True):
|
|
"""Read data from disk."""
|
|
if fmt == "torch":
|
|
return _read_torch_data(path)
|
|
elif fmt == "numpy":
|
|
return _read_numpy_data(path, in_memory=in_memory)
|
|
else:
|
|
raise RuntimeError(f"Unsupported format: {fmt}")
|
|
|
|
|
|
def save_data(data, path, fmt):
|
|
"""Save data into disk."""
|
|
# Make sure the directory exists.
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
|
|
if fmt not in ["numpy", "torch"]:
|
|
raise RuntimeError(f"Unsupported format: {fmt}")
|
|
|
|
# Perform necessary conversion.
|
|
if fmt == "numpy" and isinstance(data, torch.Tensor):
|
|
data = data.cpu().numpy()
|
|
elif fmt == "torch" and isinstance(data, np.ndarray):
|
|
data = torch.from_numpy(data).cpu()
|
|
|
|
# Save the data.
|
|
if fmt == "numpy":
|
|
if not data.flags["C_CONTIGUOUS"]:
|
|
Warning(
|
|
"The ndarray saved to disk is not contiguous, "
|
|
"so it will be copied to contiguous memory."
|
|
)
|
|
data = np.ascontiguousarray(data)
|
|
numpy_save_aligned(path, data)
|
|
elif fmt == "torch":
|
|
if not data.is_contiguous():
|
|
Warning(
|
|
"The tensor saved to disk is not contiguous, "
|
|
"so it will be copied to contiguous memory."
|
|
)
|
|
data = data.contiguous()
|
|
torch.save(data, path)
|
|
|
|
|
|
def get_npy_dim(npy_path):
|
|
"""Get the dim of numpy file."""
|
|
with open(npy_path, "rb") as f:
|
|
# For the read_array_header API provided by numpy will only read the
|
|
# length of the header, it will cause parsing failure and error if
|
|
# first 8 bytes which contains magin string and version are not read
|
|
# ahead of time. So, we need to make sure we have skipped these 8
|
|
# bytes.
|
|
f.seek(8, 0)
|
|
try:
|
|
shape, _, _ = read_array_header_1_0(f)
|
|
except ValueError:
|
|
try:
|
|
shape, _, _ = read_array_header_2_0(f)
|
|
except ValueError:
|
|
raise ValueError("Invalid file format")
|
|
|
|
return len(shape)
|
|
|
|
|
|
def _to_int32(data):
|
|
if isinstance(data, torch.Tensor):
|
|
return data.to(torch.int32)
|
|
elif isinstance(data, np.ndarray):
|
|
return data.astype(np.int32)
|
|
else:
|
|
raise TypeError(
|
|
"Unsupported input type. Please provide a torch tensor or numpy array."
|
|
)
|
|
|
|
|
|
def copy_or_convert_data(
|
|
input_path,
|
|
output_path,
|
|
input_format,
|
|
output_format="numpy",
|
|
in_memory=True,
|
|
is_feature=False,
|
|
within_int32=False,
|
|
):
|
|
"""Copy or convert the data from input_path to output_path."""
|
|
assert (
|
|
output_format == "numpy"
|
|
), "The output format of the data should be numpy."
|
|
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|
# We read the data always in case we need to cast its type.
|
|
data = read_data(input_path, input_format, in_memory)
|
|
if within_int32:
|
|
data = _to_int32(data)
|
|
if input_format == "numpy":
|
|
# If dim of the data is 1, reshape it to n * 1 and save it to output_path.
|
|
if is_feature and get_npy_dim(input_path) == 1:
|
|
data = data.reshape(-1, 1)
|
|
# If the data does not need to be modified, just copy the file.
|
|
elif not within_int32 and data.numpy().flags["C_CONTIGUOUS"]:
|
|
shutil.copyfile(input_path, output_path)
|
|
return
|
|
else:
|
|
# If dim of the data is 1, reshape it to n * 1 and save it to output_path.
|
|
if is_feature and data.dim() == 1:
|
|
data = data.reshape(-1, 1)
|
|
save_data(data, output_path, output_format)
|
|
|
|
|
|
def read_edges(dataset_dir, edge_fmt, edge_path):
|
|
"""Read egde data from numpy or csv."""
|
|
assert edge_fmt in [
|
|
"numpy",
|
|
"csv",
|
|
], f"`numpy` or `csv` is expected when reading edges but got `{edge_fmt}`."
|
|
if edge_fmt == "numpy":
|
|
edge_data = read_data(
|
|
os.path.join(dataset_dir, edge_path),
|
|
edge_fmt,
|
|
)
|
|
assert (
|
|
edge_data.shape[0] == 2 and len(edge_data.shape) == 2
|
|
), f"The shape of edges should be (2, N), but got {edge_data.shape}."
|
|
src, dst = edge_data.numpy()
|
|
else:
|
|
edge_data = pd.read_csv(
|
|
os.path.join(dataset_dir, edge_path),
|
|
names=["src", "dst"],
|
|
)
|
|
src, dst = edge_data["src"].to_numpy(), edge_data["dst"].to_numpy()
|
|
return (src, dst)
|
|
|
|
|
|
def calculate_file_hash(file_path, hash_algo="md5"):
|
|
"""Calculate the hash value of a file."""
|
|
hash_algos = ["md5", "sha1", "sha224", "sha256", "sha384", "sha512"]
|
|
if hash_algo in hash_algos:
|
|
hash_obj = getattr(hashlib, hash_algo)()
|
|
else:
|
|
raise ValueError(
|
|
f"Hash algorithm must be one of: {hash_algos}, but got `{hash_algo}`."
|
|
)
|
|
with open(file_path, "rb") as file:
|
|
for chunk in iter(lambda: file.read(4096), b""):
|
|
hash_obj.update(chunk)
|
|
return hash_obj.hexdigest()
|
|
|
|
|
|
def calculate_dir_hash(
|
|
dir_path, hash_algo="md5", ignore: Union[str, List[str]] = None
|
|
):
|
|
"""Calculte the hash values of all files under the directory."""
|
|
hashes = {}
|
|
for dirpath, _, filenames in os.walk(dir_path):
|
|
for filename in filenames:
|
|
if ignore and filename in ignore:
|
|
continue
|
|
filepath = os.path.join(dirpath, filename)
|
|
file_hash = calculate_file_hash(filepath, hash_algo=hash_algo)
|
|
hashes[filepath] = file_hash
|
|
return hashes
|
|
|
|
|
|
def check_dataset_change(dataset_dir, processed_dir):
|
|
"""Check whether dataset has been changed by checking its hash value."""
|
|
hash_value_file = "dataset_hash_value.txt"
|
|
hash_value_file_path = os.path.join(
|
|
dataset_dir, processed_dir, hash_value_file
|
|
)
|
|
if not os.path.exists(hash_value_file_path):
|
|
return True
|
|
with open(hash_value_file_path, "r") as f:
|
|
oringinal_hash_value = json.load(f)
|
|
present_hash_value = calculate_dir_hash(dataset_dir, ignore=hash_value_file)
|
|
if oringinal_hash_value == present_hash_value:
|
|
force_preprocess = False
|
|
else:
|
|
force_preprocess = True
|
|
return force_preprocess
|