Add ability to load, partition, and separate node labels with features

CHANGELOG.md

@@ -12,6 +12,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 - Added support for NodeAnchorSplitter
 
+- Add support for loading, partitioning, and separate node labels from features
+
 ### Changed
 
 ### Deprecated

python/gigl/common/data/dataloaders.py

@@ -1,6 +1,6 @@
 import time
 from copy import deepcopy
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from functools import partial
 from typing import Callable, Optional, Sequence, Tuple, Union
 
@@ -29,12 +29,17 @@ class SerializedTFRecordInfo:
     tfrecord_uri_prefix: Uri
     # Feature names to load for the current entity
     feature_keys: Sequence[str]
-    # a dict of feature name -> FeatureSpec (eg. FixedLenFeature, VarlenFeature, SparseFeature, RaggedFeature). If entity keys are not present, we insert them during tensor loading
+    # A dict of feature name -> FeatureSpec (eg. FixedLenFeature, VarlenFeature, SparseFeature, RaggedFeature).
+    # If entity keys are not present, we insert them during tensor loading. For example, if the FeatureSpecDict
+    # doesn't have the "node_id" identifier, we populate the feature_spce with a FixedLenFeature with shape=[], dtype=tf.int64.
+    # Note that entity label keys should also be included in the feature_spec if they are present.
     feature_spec: FeatureSpecDict
     # Feature dimension of current entity
     feature_dim: int
     # Entity ID Key for current entity. If this is a Node Entity, this must be a string. If this is an edge entity, this must be a Tuple[str, str] for the source and destination ids.
     entity_key: Union[str, Tuple[str, str]]
+    # Name of the label columns for the current entity, defaults to an empty list.
+    label_keys: Sequence[str] = field(default_factory=list)
     # The regex pattern to match the TFRecord files at the specified prefix
     tfrecord_uri_pattern: str = ".*-of-.*\.tfrecord(\.gz)?$"
 
@@ -79,23 +84,31 @@ class TFDatasetOptions:
 def _concatenate_features_by_names(
     feature_key_to_tf_tensor: dict[str, tf.Tensor],
     feature_keys: Sequence[str],
+    label_keys: Sequence[str],
 ) -> tf.Tensor:
     """
     Concatenates feature tensors in the order specified by feature names.
+    Also concatenates labels to the end of the feature list if they are present using the corresponding label key
 
     It is assumed that feature_names is a subset of the keys in feature_name_to_tf_tensor.
 
     Args:
         feature_key_to_tf_tensor (dict[str, tf.Tensor]): A dictionary mapping feature names to their corresponding tf tensors.
-        feature_keys (list[str]): A list of feature names specifying the order in which tensors should be concatenated.
+        feature_keys (Sequence[str]): A list of feature names specifying the order in which tensors should be concatenated.
+        label_keys (Sequence[str]): Name of the label columns for the current entity.
 
     Returns:
-        tf.Tensor: A concatenated tensor of the features in the specified order.
+        tf.Tensor: A concatenated tensor of the features in the specified order, with the labels being concatenated at the end if it exists
     """
 
     features: list[tf.Tensor] = []
 
-    for feature_key in feature_keys:
+    feature_iterable = list(feature_keys)
+
+    for label_key in label_keys:
+        feature_iterable.append(label_key)
+
+    for feature_key in feature_iterable:
         tensor = feature_key_to_tf_tensor[feature_key]
 
         # TODO(kmonte, xgao, zfan): We will need to add support for this if we're trying to scale up.
@@ -298,6 +311,7 @@ def load_as_torch_tensors(
         """
         entity_key = serialized_tf_record_info.entity_key
         feature_keys = serialized_tf_record_info.feature_keys
+        label_keys = serialized_tf_record_info.label_keys
 
         # We make a deep copy of the feature spec dict so that future modifications don't redirect to the input
 
@@ -356,11 +370,16 @@ def load_as_torch_tensors(
                 if entity_type == FeatureTypes.NODE
                 else torch.empty(2, 0)
             )
-            empty_feature = (
-                torch.empty(0, serialized_tf_record_info.feature_dim)
-                if feature_keys
-                else None
-            )
+            if label_keys and feature_keys:
+                empty_feature = torch.empty(
+                    0, serialized_tf_record_info.feature_dim + len(label_keys)
+                )
+            elif label_keys and not feature_keys:
+                empty_feature = torch.empty(0, len(label_keys))
+            elif not label_keys and feature_keys:
+                empty_feature = torch.empty(0, serialized_tf_record_info.feature_dim)
+            else:
+                empty_feature = None
             return empty_entity, empty_feature
 
         dataset = TFRecordDataLoader._build_dataset_for_uris(
@@ -375,9 +394,9 @@ def load_as_torch_tensors(
         feature_tensors = []
         for idx, batch in enumerate(dataset):
             id_tensors.append(proccess_id_tensor(batch))
-            if feature_keys:
+            if feature_keys or label_keys:
                 feature_tensors.append(
-                    _concatenate_features_by_names(batch, feature_keys)
+                    _concatenate_features_by_names(batch, feature_keys, label_keys)
                 )
             num_entities_processed += (
                 id_tensors[-1].shape[0]

python/gigl/distributed/dataset_factory.py

@@ -2,8 +2,9 @@
 DatasetFactory is responsible for building and returning a DistLinkPredictionDataset class or subclass. It does this by spawning a
 process which initializes rpc + worker group, loads and builds a partitioned dataset, and shuts down the rpc + worker group.
 """
+import gc
 import time
-from collections import abc
+from collections.abc import Mapping
 from distutils.util import strtobool
 from typing import Literal, MutableMapping, Optional, Tuple, Type, Union
 
@@ -19,7 +20,7 @@
 )
 
 from gigl.common import Uri, UriFactory
-from gigl.common.data.dataloaders import TFRecordDataLoader
+from gigl.common.data.dataloaders import SerializedTFRecordInfo, TFRecordDataLoader
 from gigl.common.data.load_torch_tensors import (
     SerializedGraphMetadata,
     TFDatasetOptions,
@@ -40,9 +41,9 @@
 from gigl.distributed.utils.serialized_graph_metadata_translator import (
     convert_pb_to_serialized_graph_metadata,
 )
-from gigl.src.common.types.graph_data import EdgeType
+from gigl.src.common.types.graph_data import EdgeType, NodeType
 from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper
-from gigl.types.graph import DEFAULT_HOMOGENEOUS_EDGE_TYPE
+from gigl.types.graph import DEFAULT_HOMOGENEOUS_EDGE_TYPE, FeaturePartitionData
 from gigl.utils.data_splitters import (
     HashedNodeAnchorLinkSplitter,
     NodeAnchorLinkSplitter,
@@ -52,6 +53,39 @@
 logger = Logger()
 
 
+def _get_labels_from_features(
+    feature_and_label_tensor: torch.Tensor, label_dim: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Given a combined tensor of features and labels, returns the features and labels separately.
+    Args:
+        feature_and_label_tensor (torch.Tensor): Tensor of features and labels
+        label_dim (int): Dimension of the labels
+    Returns:
+        feature_tensor (torch.Tensor): Tensor of features
+        label_tensor (torch.Tensor): Tensor of labels
+    """
+
+    if len(feature_and_label_tensor.shape) != 2:
+        raise ValueError(
+            f"Expected tensor to be 2D for extracting labels, but got shape {feature_and_label_tensor.shape}"
+        )
+
+    _, feature_and_label_dim = feature_and_label_tensor.shape
+
+    if label_dim > feature_and_label_dim:
+        raise ValueError(
+            f"Got invalid label dim {label_dim} for extracting labels from tensor of shape {feature_and_label_tensor.shape}"
+        )
+
+    feature_dim = feature_and_label_dim - label_dim
+
+    return (
+        feature_and_label_tensor[:, :feature_dim],
+        feature_and_label_tensor[:, feature_dim:],
+    )
+
+
 @tf_on_cpu
 def _load_and_build_partitioned_dataset(
     serialized_graph_metadata: SerializedGraphMetadata,
@@ -112,7 +146,7 @@ def _load_and_build_partitioned_dataset(
                 "Cannot have loaded positive and negative labels when attempting to select self-supervised positive edges from edge index."
             )
         positive_label_edges: Union[torch.Tensor, dict[EdgeType, torch.Tensor]]
-        if isinstance(loaded_graph_tensors.edge_index, abc.Mapping):
+        if isinstance(loaded_graph_tensors.edge_index, Mapping):
             # This assert is required while `select_ssl_positive_label_edges` exists out of any splitter. Once this is in transductive splitter,
             # we can remove this assert.
             assert isinstance(
@@ -192,6 +226,56 @@ def _load_and_build_partitioned_dataset(
 
     partition_output = partitioner.partition()
 
+    node_labels: Optional[Union[torch.Tensor, dict[NodeType, torch.Tensor]]] = None
+    if isinstance(partition_output.partitioned_node_features, Mapping):
+        node_labels = {}
+        for (
+            node_type,
+            node_feature,
+        ) in partition_output.partitioned_node_features.items():
+            # serialized_graph_metadata can be heterogeneous for a heterogeneous node features
+            if isinstance(serialized_graph_metadata.node_entity_info, Mapping):
+                label_dim = len(
+                    serialized_graph_metadata.node_entity_info[node_type].label_keys
+                )
+            # serialized_graph_metadata can be homogeneous for a heterogeneous node features,
+            # since we inject positive and negative label types as edges
+            else:
+                label_dim = len(serialized_graph_metadata.node_entity_info.label_keys)
+            if label_dim > 0:
+                node_features, node_labels[node_type] = _get_labels_from_features(
+                    node_feature.feats, label_dim=label_dim
+                )
+                partition_output.partitioned_node_features[
+                    node_type
+                ] = FeaturePartitionData(feats=node_features, ids=node_feature.ids)
+                del node_feature
+                gc.collect()
+
+    elif isinstance(partition_output.partitioned_node_features, FeaturePartitionData):
+        # serialized graph metadata must be homogeneous if partitioned node features is homogeneous
+        if not isinstance(
+            serialized_graph_metadata.node_entity_info, SerializedTFRecordInfo
+        ):
+            raise ValueError(
+                f"Expected partitioned node features to be type SerializedTFRecordInfo, got {type(partition_output.partitioned_node_features)}"
+            )
+        label_dim = len(serialized_graph_metadata.node_entity_info.label_keys)
+        if label_dim > 0:
+            node_features, node_labels = _get_labels_from_features(
+                partition_output.partitioned_node_features.feats, label_dim=label_dim
+            )
+            partition_output.partitioned_node_features = FeaturePartitionData(
+                feats=node_features, ids=partition_output.partitioned_node_features.ids
+            )
+            gc.collect()
+    else:
+        raise ValueError(
+            f"Expected to have partitioned node features if labels are present, but got node features {partition_output.partitioned_node_features}"
+        )
+
+    # TODO (mkolodner-sc): Add node labels to the dataset
+
     logger.info(
         f"Initializing DistLinkPredictionDataset instance with edge direction {edge_dir}"
     )

python/gigl/src/common/types/pb_wrappers/preprocessed_metadata.py

@@ -68,13 +68,14 @@ def __post_init__(self):
             condensed_node_type_to_feature_dim_map[
                 CondensedNodeType(condensed_node_type)
             ] = node_metadata.feature_dim
-            node_feature_keys = list(node_metadata.feature_keys)
+            node_feature_keys = sorted(list(node_metadata.feature_keys))
+            label_keys = sorted(list(node_metadata.label_keys))
             node_feature_schema = self.__build_feature_schema(
                 schema_uri=UriFactory.create_uri(node_metadata.schema_uri),
                 transform_fn_assets_uri=UriFactory.create_uri(
                     node_metadata.transform_fn_assets_uri
                 ),
-                feature_keys=node_feature_keys,
+                feature_keys=node_feature_keys + label_keys,
             )
             condensed_node_type_to_feature_schema_map[
                 CondensedNodeType(condensed_node_type)
@@ -380,12 +381,10 @@ def condensed_node_type_to_feature_keys_map(
             dict[CondensedNodeType, list[str]]: A mapping which stores the feature keys of each CondensedNodeTypes
         """
         return {
-            condensed_node_type: list(
-                self.condensed_node_type_to_feature_schema_map[
-                    condensed_node_type
-                ].feature_spec.keys()
+            CondensedNodeType(condensed_node_type): sorted(
+                list(preprocessed_metadata.feature_keys)
             )
-            for condensed_node_type in self.condensed_node_type_to_feature_schema_map
+            for condensed_node_type, preprocessed_metadata in self.preprocessed_metadata_pb.condensed_node_type_to_preprocessed_metadata.items()
         }
 
     @property