Repository: dmlc/dgl
Pull request: [Feature] Add parquet support for node/edge features in chunked data
## Description
Support node/edge features with parquet format in chunked data when doing data pre-processing. The provided parquet files are assumed to store only one NdArray, and its dim is default to 2 if not explicitly specified in the file. 

## Checklist
Please feel free to remove inapplicable items for your PR.
- [x] The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
- [x] I've leverage the [tools](https://docs.google.com/document/d/1iHyj7zlmygKSk5gBPsqIqL5ASPzJSPREaNT_QdsiYA4/edit) to beautify the python and c++ code.
- [x] The PR is complete and small, read the [Google eng practice (CL equals to PR)](https://google.github.io/eng-practices/review/developer/small-cls.html) to understand more about small PR. In DGL, we consider PRs with less than 200 lines of core code change are small (example, test and documentation could be exempted).
- [x] All changes have test coverage
- [x] Code is well-documented
- [ ] To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change
- [x] Related issue is referred in this PR
- [ ] If the PR is for a new model/paper, I've updated the example index [here](../examples/README.md).

## Changes
<!-- You could use following template
- [ ] Feature1, tests, (and when applicable, API doc)
- [ ] Feature2, tests, (and when applicable, API doc)
-->


Problem statement:
[Feature request] tools/distpartitioning/data_proc_pipeline.py need to support loading edge or node features in Parquet format.
## 🚀 Feature
The data_proc_pipeline.py, which combines ParMetis output with node and edge features, should support loading edge and node features stored in Parquet format. This is important as it is common to use Spark to preprocess the graph data.

## Motivation

<!-- Please outline the motivation for the proposal. Is your feature request
related to a problem? e.g., I'm always frustrated when [...]. If this is
related to another GitHub issue, please link here too -->

## Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

## Pitch

<!-- A clear and concise description of what you want to happen. -->

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->

Hints from issue discussion:

Commit message: add e2e test for parquet

Patch to review:
diff --git a/tests/tools/create_chunked_dataset.py b/tests/tools/create_chunked_dataset.py
index c5021d6ff1a2..95d581739b47 100644
--- a/tests/tools/create_chunked_dataset.py
+++ b/tests/tools/create_chunked_dataset.py
@@ -15,7 +15,7 @@
 
 
 def create_chunked_dataset(
-    root_dir, num_chunks, include_masks=False
+    root_dir, num_chunks, include_masks=False, data_fmt='numpy'
 ):
     """
     This function creates a sample dataset, based on MAG240 dataset.
@@ -233,6 +233,7 @@ def rand_edges(num_src, num_dst, num_edges):
         edge_data,
         num_chunks=num_chunks,
         output_path=output_dir,
+        data_fmt=data_fmt,
     )
     print('Done with creating chunked graph')
 
diff --git a/tests/tools/test_array_readwriter.py b/tests/tools/test_array_readwriter.py
new file mode 100644
index 000000000000..c61ce2553760
--- /dev/null
+++ b/tests/tools/test_array_readwriter.py
@@ -0,0 +1,23 @@
+import tempfile
+import os
+import pytest
+import numpy as np
+from utils import array_readwriter
+
+@pytest.mark.parametrize("shape", [[500], [300,10], [200,5,5], [100,5,5,5]])
+@pytest.mark.parametrize("format", ['numpy', 'parquet'])
+def test_array_readwriter(format, shape):
+    original_arr = np.random.rand(*shape)
+    fmt_meta = {"name": format}
+    
+    with tempfile.TemporaryDirectory() as test_dir:
+        path = os.path.join(test_dir, f'nodes.{format}')
+        array_readwriter.get_array_parser(
+            **fmt_meta
+        ).write(path, original_arr)
+        arr = array_readwriter.get_array_parser(
+            **fmt_meta
+        ).read(path)
+
+        assert original_arr.shape == arr.shape
+        assert np.array_equal(original_arr, arr)
diff --git a/tests/tools/test_dist_part.py b/tests/tools/test_dist_part.py
index dae61667b74a..56a413a39364 100644
--- a/tests/tools/test_dist_part.py
+++ b/tests/tools/test_dist_part.py
@@ -8,6 +8,7 @@
 import torch
 from chunk_graph import chunk_graph
 from create_chunked_dataset import create_chunked_dataset
+from utils import array_readwriter
 
 import dgl
 from dgl.data.utils import load_graphs, load_tensors
@@ -80,11 +81,12 @@ def _verify_graph_feats(
 
 
 @pytest.mark.parametrize("num_chunks", [1, 8])
-def test_chunk_graph(num_chunks):
+@pytest.mark.parametrize("data_fmt", ['numpy', 'parquet'])
+def test_chunk_graph(num_chunks, data_fmt):
 
     with tempfile.TemporaryDirectory() as root_dir:
 
-        g = create_chunked_dataset(root_dir, num_chunks)
+        g = create_chunked_dataset(root_dir, num_chunks, data_fmt=data_fmt)
 
         # check metadata.json
         output_dir = os.path.join(root_dir, "chunked-data")
@@ -111,12 +113,16 @@ def test_chunk_graph(num_chunks):
                     assert isinstance(int(num2), int)
 
         # check node/edge_data
+        suffix = 'npy' if data_fmt=='numpy' else 'parquet'
+        reader_fmt_meta = {"name": data_fmt}
         def test_data(sub_dir, feat, expected_data, expected_shape):
             data = []
             for i in range(num_chunks):
-                fname = os.path.join(sub_dir, f'{feat}-{i}.npy')
+                fname = os.path.join(sub_dir, f'{feat}-{i}.{suffix}')
                 assert os.path.isfile(fname)
-                feat_array = np.load(fname)
+                feat_array =  array_readwriter.get_array_parser(
+                            **reader_fmt_meta
+                        ).read(fname)
                 assert feat_array.shape[0] == expected_shape
                 data.append(feat_array)
             data = np.concatenate(data, 0)
@@ -136,14 +142,14 @@ def test_data(sub_dir, feat, expected_data, expected_shape):
                 test_data(sub_dir, feat, data, g.num_edges(c_etype) // num_chunks)
 
 
-def _test_pipeline(num_chunks, num_parts, graph_formats=None):
+def _test_pipeline(num_chunks, num_parts, graph_formats=None, data_fmt='numpy'):
     if num_chunks < num_parts:
         # num_parts should less/equal than num_chunks
         return
 
     with tempfile.TemporaryDirectory() as root_dir:
 
-        g = create_chunked_dataset(root_dir, num_chunks)
+        g = create_chunked_dataset(root_dir, num_chunks, data_fmt=data_fmt)
 
         # Step1: graph partition
         in_dir = os.path.join(root_dir, "chunked-data")
@@ -211,8 +217,9 @@ def read_orig_ids(fname):
 
 @pytest.mark.parametrize("num_chunks", [1, 3, 4, 8])
 @pytest.mark.parametrize("num_parts", [1, 3, 4, 8])
-def test_pipeline_basics(num_chunks, num_parts):
-    _test_pipeline(num_chunks, num_parts)
+@pytest.mark.parametrize("data_fmt", ['numpy', 'parquet'])
+def test_pipeline_basics(num_chunks, num_parts, data_fmt):
+    _test_pipeline(num_chunks, num_parts, data_fmt=data_fmt)
 
 
 @pytest.mark.parametrize(
diff --git a/tools/chunk_graph.py b/tools/chunk_graph.py
index afbf35ef2f75..821f3bbd8046 100644
--- a/tools/chunk_graph.py
+++ b/tools/chunk_graph.py
@@ -25,7 +25,7 @@ def chunk_numpy_array(arr, fmt_meta, chunk_sizes, path_fmt):
     return paths
 
 
-def _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path):
+def _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path, data_fmt):
     # First deal with ndata and edata that are homogeneous (i.e. not a dict-of-dict)
     if len(g.ntypes) == 1 and not isinstance(
         next(iter(ndata_paths.values())), dict
@@ -94,6 +94,8 @@ def _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path):
             metadata["edges"][etypestr] = edges_meta
 
     # Chunk node data
+    reader_fmt_meta, writer_fmt_meta = {"name": "numpy"}, {"name": data_fmt}
+    file_suffix = 'npy' if data_fmt == 'numpy' else 'parquet'
     metadata["node_data"] = {}
     with setdir("node_data"):
         for ntype, ndata_per_type in ndata_paths.items():
@@ -104,7 +106,6 @@ def _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path):
                         "Chunking node data for type %s key %s" % (ntype, key)
                     )
                     ndata_key_meta = {}
-                    reader_fmt_meta = writer_fmt_meta = {"name": "numpy"}
                     arr = array_readwriter.get_array_parser(
                         **reader_fmt_meta
                     ).read(path)
@@ -113,7 +114,7 @@ def _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path):
                         arr,
                         writer_fmt_meta,
                         num_nodes_per_chunk_dict[ntype],
-                        key + "-%d.npy",
+                        key + "-%d." + file_suffix,
                     )
                     ndata_meta[key] = ndata_key_meta
 
@@ -131,7 +132,6 @@ def _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path):
                         % (etypestr, key)
                     )
                     edata_key_meta = {}
-                    reader_fmt_meta = writer_fmt_meta = {"name": "numpy"}
                     arr = array_readwriter.get_array_parser(
                         **reader_fmt_meta
                     ).read(path)
@@ -141,7 +141,7 @@ def _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path):
                         arr,
                         writer_fmt_meta,
                         num_edges_per_chunk_dict[etype],
-                        key + "-%d.npy",
+                        key + "-%d." + file_suffix,
                     )
                     edata_meta[key] = edata_key_meta
 
@@ -153,7 +153,7 @@ def _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path):
     logging.info("Saved metadata in %s" % os.path.abspath(metadata_path))
 
 
-def chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path):
+def chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path, data_fmt='numpy'):
     """
     Split the graph into multiple chunks.
 
@@ -184,7 +184,7 @@ def chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path):
         for key in edata.keys():
             edata[key] = os.path.abspath(edata[key])
     with setdir(output_path):
-        _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path)
+        _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, output_path, data_fmt)
 
 
 if __name__ == "__main__":
diff --git a/tools/distpartitioning/constants.py b/tools/distpartitioning/constants.py
index db84dbc87b3e..40b37e20dbed 100644
--- a/tools/distpartitioning/constants.py
+++ b/tools/distpartitioning/constants.py
@@ -30,6 +30,7 @@
 STR_EDGE_DATA = "edge_data"
 
 STR_NUMPY = "numpy"
+STR_PARQUET = "parquet"
 STR_CSV = "csv"
 STR_NAME = "name"
 
diff --git a/tools/distpartitioning/convert_partition.py b/tools/distpartitioning/convert_partition.py
index 37d7f227c6cf..7ec93ccfdab7 100644
--- a/tools/distpartitioning/convert_partition.py
+++ b/tools/distpartitioning/convert_partition.py
@@ -13,7 +13,7 @@
 from pyarrow import csv
 
 import constants
-from utils import get_idranges, memory_snapshot, read_json
+from data_utils import get_idranges, memory_snapshot, read_json
 from dgl.distributed.partition import (
     RESERVED_FIELD_DTYPE,
     _etype_str_to_tuple,
diff --git a/tools/distpartitioning/data_shuffle.py b/tools/distpartitioning/data_shuffle.py
index 501a945baab0..b464d3bfca04 100644
--- a/tools/distpartitioning/data_shuffle.py
+++ b/tools/distpartitioning/data_shuffle.py
@@ -20,7 +20,7 @@
                        assign_shuffle_global_nids_nodes,
                        lookup_shuffle_global_nids_edges)
 from gloo_wrapper import allgather_sizes, alltoallv_cpu, gather_metadata_json
-from utils import (augment_edge_data, get_edge_types, get_etype_featnames,
+from data_utils import (augment_edge_data, get_edge_types, get_etype_featnames,
                    get_gnid_range_map, get_idranges, get_node_types,
                    get_ntype_featnames, memory_snapshot, read_json,
                    read_ntype_partition_files, write_dgl_objects,
diff --git a/tools/distpartitioning/utils.py b/tools/distpartitioning/data_utils.py
similarity index 100%
rename from tools/distpartitioning/utils.py
rename to tools/distpartitioning/data_utils.py
diff --git a/tools/distpartitioning/dataset_utils.py b/tools/distpartitioning/dataset_utils.py
index 90bc7675c588..e43d221fe9ac 100644
--- a/tools/distpartitioning/dataset_utils.py
+++ b/tools/distpartitioning/dataset_utils.py
@@ -5,9 +5,10 @@
 import pyarrow
 import torch
 from pyarrow import csv
+from utils import array_readwriter
 
 import constants
-from utils import get_idranges
+from data_utils import get_idranges
 
 
 def get_dataset(input_dir, graph_name, rank, world_size, schema_map):
@@ -104,7 +105,7 @@ def get_dataset(input_dir, graph_name, rank, world_size, schema_map):
     to mention all the files present for this given node feature.
 
     Data read from each of the node features file is a multi-dimensional tensor data and is read
-    in numpy format, which is also the storage format of node features on the permanent storage.
+    in numpy or parquet format, which is also the storage format of node features on the permanent storage.
     '''
 
     #iterate over the "node_data" dictionary in the schema_map
@@ -117,7 +118,8 @@ def get_dataset(input_dir, graph_name, rank, world_size, schema_map):
             #where key: feature_name, value: dictionary in which keys are "format", "data"
             node_feature_tids[ntype_name] = []
             for feat_name, feat_data in ntype_feature_data.items():
-                assert feat_data[constants.STR_FORMAT][constants.STR_NAME] == constants.STR_NUMPY
+                assert feat_data[constants.STR_FORMAT][constants.STR_NAME] in [constants.STR_NUMPY, constants.STR_PARQUET]
+                reader_fmt_meta = {"name": feat_data[constants.STR_FORMAT][constants.STR_NAME]}
                 num_chunks = len(feat_data[constants.STR_DATA])
                 read_list = np.array_split(np.arange(num_chunks), world_size)
                 nfeat = []
@@ -126,10 +128,13 @@ def get_dataset(input_dir, graph_name, rank, world_size, schema_map):
                     if not os.path.isabs(nfeat_file):
                         nfeat_file = os.path.join(input_dir, nfeat_file)
                     logging.info(f'Loading node feature[{feat_name}] of ntype[{ntype_name}] from {nfeat_file}')
-                    nfeat.append(np.load(nfeat_file))
+                    nfeat.append(
+                        array_readwriter.get_array_parser(
+                            **reader_fmt_meta
+                        ).read(nfeat_file)
+                    )
                 nfeat = np.concatenate(nfeat)
                 node_features[ntype_name + '/' + feat_name] = torch.from_numpy(nfeat)
-
                 node_feature_tids[ntype_name].append([feat_name, -1, -1])
 
     '''
@@ -224,7 +229,8 @@ def get_dataset(input_dir, graph_name, rank, world_size, schema_map):
             #where key: feature_name, value: dictionary in which keys are "format", "data"
             edge_feature_tids[etype_name] = []
             for feat_name, feat_data in etype_feature_data.items():
-                assert feat_data[constants.STR_FORMAT][constants.STR_NAME] == constants.STR_NUMPY
+                assert feat_data[constants.STR_FORMAT][constants.STR_NAME] in [constants.STR_NUMPY, constants.STR_PARQUET]
+                reader_fmt_meta = {"name": feat_data[constants.STR_FORMAT][constants.STR_NAME]}
                 num_chunks = len(feat_data[constants.STR_DATA])
                 read_list = np.array_split(np.arange(num_chunks), world_size)
                 efeat = []
@@ -235,7 +241,11 @@ def get_dataset(input_dir, graph_name, rank, world_size, schema_map):
                     logging.info(
                         f'Loading edge feature[{feat_name}] of etype[{etype_name}] from {efeat_file}'
                     )
-                    efeat.append(np.load(efeat_file))
+                    efeat.append(
+                        array_readwriter.get_array_parser(
+                            **reader_fmt_meta
+                        ).read(efeat_file)
+                    )
                 efeat = np.concatenate(efeat)
                 edge_features[etype_name + '/' + feat_name] = torch.from_numpy(efeat)
 
diff --git a/tools/distpartitioning/globalids.py b/tools/distpartitioning/globalids.py
index f48d2534b7fb..a447bda9a017 100644
--- a/tools/distpartitioning/globalids.py
+++ b/tools/distpartitioning/globalids.py
@@ -7,7 +7,7 @@
 import constants
 from dist_lookup import DistLookupService
 from gloo_wrapper import allgather_sizes, alltoallv_cpu
-from utils import memory_snapshot
+from data_utils import memory_snapshot
 
 
 def get_shuffle_global_nids(rank, world_size, global_nids_ranks, node_data):
diff --git a/tools/distpartitioning/parmetis_postprocess.py b/tools/distpartitioning/parmetis_postprocess.py
index 5b7447ec7159..f241ab1e7913 100644
--- a/tools/distpartitioning/parmetis_postprocess.py
+++ b/tools/distpartitioning/parmetis_postprocess.py
@@ -10,7 +10,7 @@
 import pyarrow.csv as csv
 
 import constants
-from utils import get_idranges, get_node_types, read_json
+from data_utils import get_idranges, get_node_types, read_json
 
 
 def post_process(params):
diff --git a/tools/distpartitioning/parmetis_preprocess.py b/tools/distpartitioning/parmetis_preprocess.py
index 44350460c4e4..8c34385e6f9f 100644
--- a/tools/distpartitioning/parmetis_preprocess.py
+++ b/tools/distpartitioning/parmetis_preprocess.py
@@ -11,7 +11,7 @@
 import torch.distributed as dist
 
 import constants
-from utils import get_idranges, get_node_types, read_json
+from data_utils import get_idranges, get_node_types, read_json
 
 
 def get_proc_info():
diff --git a/tools/distpartitioning/parmetis_wrapper.py b/tools/distpartitioning/parmetis_wrapper.py
index 521fada400d0..6ae905cf1782 100644
--- a/tools/distpartitioning/parmetis_wrapper.py
+++ b/tools/distpartitioning/parmetis_wrapper.py
@@ -6,7 +6,7 @@
 from pathlib import Path
 
 import constants
-from utils import read_json
+from data_utils import read_json
 
 
 def check_dependencies():
diff --git a/tools/utils/array_readwriter/__init__.py b/tools/utils/array_readwriter/__init__.py
index 6b8514c0cc46..3cce32874b98 100644
--- a/tools/utils/array_readwriter/__init__.py
+++ b/tools/utils/array_readwriter/__init__.py
@@ -1,2 +1,2 @@
-from . import csv, numpy_array
+from . import csv, numpy_array, parquet
 from .registry import get_array_parser, register_array_parser
diff --git a/tools/utils/array_readwriter/parquet.py b/tools/utils/array_readwriter/parquet.py
new file mode 100644
index 000000000000..05f70283e54d
--- /dev/null
+++ b/tools/utils/array_readwriter/parquet.py
@@ -0,0 +1,36 @@
+import logging
+
+import pandas as pd
+import pyarrow
+import pyarrow.parquet
+
+from .registry import register_array_parser
+
+
+@register_array_parser("parquet")
+class ParquetArrayParser(object):
+    def __init__(self):
+        pass
+
+    def read(self, path):
+        logging.info("Reading from %s using parquet format" % path)
+        metadata = pyarrow.parquet.read_metadata(path)
+        metadata = metadata.schema.to_arrow_schema().metadata
+        # As parquet data are tabularized, we assume the dim of ndarray is 2.
+        # If not, it should be explictly specified in the file as metadata.
+        shape = metadata.get(b'shape', None)
+        table = pyarrow.parquet.read_table(path, memory_map=True)
+        logging.info("Done reading from %s" % path)
+        arr = table.to_pandas().to_numpy()
+        shape = tuple(eval(shape.decode())) if shape else arr.shape
+        return arr.reshape(shape)
+
+    def write(self, path, arr):
+        logging.info("Writing to %s using parquet format" % path)
+        shape = arr.shape
+        if len(shape) > 2:
+            arr = arr.reshape(shape[0], -1)
+        table = pyarrow.Table.from_pandas(pd.DataFrame(arr))
+        table = table.replace_schema_metadata({"shape": str(shape)})
+        pyarrow.parquet.write_table(table, path)
+        logging.info("Done writing to %s" % path)


Code review comments:
diff_hunk: @@ -126,10 +128,13 @@ def get_dataset(input_dir, graph_name, rank, world_size, schema_map):
                     if not os.path.isabs(nfeat_file):
                         nfeat_file = os.path.join(input_dir, nfeat_file)
                     logging.info(f'Loading node feature[{feat_name}] of ntype[{ntype_name}] from {nfeat_file}')
-                    nfeat.append(np.load(nfeat_file))
+                    nfeat.append(
+                        array_readwriter.get_array_parser(
path: tools/distpartitioning/dataset_utils.py
text: @user1:
could we have a wrapper that could help create `fmt_meta` for the caller? such as `utils.read(nfeat_file, fmt=feat_data[][])`?