Jchmura/edge feature quantization

gigl/common/data/dataloaders.py

@@ -398,16 +398,6 @@ def load_as_torch_tensors(
                 feature_spec_dict[entity_key] = tf.io.FixedLenFeature(
                     shape=[], dtype=tf.int64
                 )
-            if (
-                packed_feature_key is not None
-                and packed_feature_key not in feature_spec_dict
-            ):
-                logger.info(
-                    f"Injecting packed feature key {packed_feature_key} into feature spec dictionary with value `tf.io.FixedLenFeature(shape=[], dtype=tf.string)`"
-                )
-                feature_spec_dict[packed_feature_key] = tf.io.FixedLenFeature(
-                    shape=[], dtype=tf.string
-                )
         else:
             id_concat_axis = 1
             proccess_id_tensor = lambda t: tf.stack(
@@ -433,6 +423,17 @@ def load_as_torch_tensors(
                     shape=[], dtype=tf.int64
                 )
 
+        if (
+            packed_feature_key is not None
+            and packed_feature_key not in feature_spec_dict
+        ):
+            logger.info(
+                f"Injecting packed feature key {packed_feature_key} into feature spec dictionary with value `tf.io.FixedLenFeature(shape=[], dtype=tf.string)`"
+            )
+            feature_spec_dict[packed_feature_key] = tf.io.FixedLenFeature(
+                shape=[], dtype=tf.string
+            )
+
         uris = self._partition_children_uris(
             serialized_tf_record_info.tfrecord_uri_prefix,
             serialized_tf_record_info.tfrecord_uri_pattern,

gigl/common/data/load_torch_tensors.py

@@ -1,7 +1,7 @@
 import time
 import traceback
-from dataclasses import dataclass
-from typing import MutableMapping, Optional, Union
+from dataclasses import dataclass, replace
+from typing import MutableMapping, Optional, Union, cast
 
 import torch
 import torch.multiprocessing as mp
@@ -119,6 +119,134 @@ class SerializedGraphMetadata:
     node_quantization_metadata: Optional[
         Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]]
     ] = None
+    edge_quantization_metadata: Optional[
+        Union[FeatureQuantizationMetadata, dict[EdgeType, FeatureQuantizationMetadata]]
+    ] = None
+
+
+def _validate_weight_edge_feature_name(
+    edge_entity_info: Union[
+        SerializedTFRecordInfo, dict[EdgeType, SerializedTFRecordInfo]
+    ],
+    weight_edge_feat_name: Optional[Union[str, dict[EdgeType, str]]],
+) -> None:
+    if weight_edge_feat_name is None:
+        return
+
+    configured_weights: list[tuple[EdgeType, str, SerializedTFRecordInfo]]
+    if isinstance(edge_entity_info, SerializedTFRecordInfo):
+        if not isinstance(weight_edge_feat_name, str):
+            raise ValueError("weight_edge_feat_name must be str for homogeneous graph")
+        edge_type = DEFAULT_HOMOGENEOUS_EDGE_TYPE
+        configured_weights = [(edge_type, weight_edge_feat_name, edge_entity_info)]
+    else:
+        if isinstance(weight_edge_feat_name, str):
+            if len(edge_entity_info) != 1:
+                raise ValueError(
+                    "weight_edge_feat_name must be dict[EdgeType, str] for heterogeneous graph with multiple edge types"
+                )
+            edge_type, serialized_info = next(iter(edge_entity_info.items()))
+            configured_weights = [(edge_type, weight_edge_feat_name, serialized_info)]
+        else:
+            unknown_edge_types = set(weight_edge_feat_name) - set(edge_entity_info)
+            if unknown_edge_types:
+                raise ValueError(
+                    f"weight_edge_feat_name contains unknown edge types: {unknown_edge_types}"
+                )
+            configured_weights = [
+                (edge_type, feature_name, edge_entity_info[edge_type])
+                for edge_type, feature_name in weight_edge_feat_name.items()
+            ]
+
+    for edge_type, feature_name, serialized_info in configured_weights:
+        if feature_name not in serialized_info.feature_keys:
+            raise ValueError(
+                f"Sampling-weight field '{feature_name}' for edge type {edge_type} must be an unquantized raw edge feature."
+            )
+
+
+def remove_sampling_weight_from_edge_quantization_metadata(
+    serialized_graph_metadata: SerializedGraphMetadata,
+    weight_edge_feat_name: Optional[Union[str, dict[EdgeType, str]]],
+) -> Optional[
+    Union[FeatureQuantizationMetadata, dict[EdgeType, FeatureQuantizationMetadata]]
+]:
+    """Remove separately stored sampling weights from edge reconstruction metadata.
+
+    TFRecord loading removes the sampling-weight column from raw edge features
+    before registering it with the weighted sampler. The resulting metadata
+    must describe the remaining model features so batch reconstruction scatters
+    raw and dequantized columns into the correct positions.
+
+    Args:
+        serialized_graph_metadata: Serialized edge schema and quantization metadata.
+        weight_edge_feat_name: Raw scalar feature configured as sampling weights.
+
+    Returns:
+        Quantization metadata for the model-facing edge features.
+    """
+    quantization_metadata = serialized_graph_metadata.edge_quantization_metadata
+    if quantization_metadata is None or weight_edge_feat_name is None:
+        return quantization_metadata
+
+    if isinstance(serialized_graph_metadata.edge_entity_info, SerializedTFRecordInfo):
+        assert isinstance(quantization_metadata, FeatureQuantizationMetadata)
+        assert isinstance(weight_edge_feat_name, str)
+        edge_info_by_type: dict[EdgeType, SerializedTFRecordInfo] = {
+            DEFAULT_HOMOGENEOUS_EDGE_TYPE: serialized_graph_metadata.edge_entity_info
+        }
+        metadata_by_type: dict[EdgeType, FeatureQuantizationMetadata] = {
+            DEFAULT_HOMOGENEOUS_EDGE_TYPE: quantization_metadata
+        }
+        weight_by_type: dict[EdgeType, str] = {
+            DEFAULT_HOMOGENEOUS_EDGE_TYPE: weight_edge_feat_name
+        }
+        is_homogeneous = True
+    else:
+        assert isinstance(quantization_metadata, dict)
+        edge_info_by_type: dict[EdgeType, SerializedTFRecordInfo] = (
+            serialized_graph_metadata.edge_entity_info
+        )
+        metadata_by_type: dict[EdgeType, FeatureQuantizationMetadata] = cast(
+            dict[EdgeType, FeatureQuantizationMetadata], quantization_metadata
+        )
+        if isinstance(weight_edge_feat_name, str):
+            edge_type = next(iter(edge_info_by_type))
+            weight_by_type: dict[EdgeType, str] = {edge_type: weight_edge_feat_name}
+        else:
+            weight_by_type: dict[EdgeType, str] = weight_edge_feat_name
+        is_homogeneous = False
+
+    adjusted_metadata: dict[EdgeType, FeatureQuantizationMetadata] = {}
+    for edge_type, metadata in metadata_by_type.items():
+        weight_feature_name = weight_by_type.get(edge_type)
+        if weight_feature_name is None:
+            adjusted_metadata[edge_type] = metadata
+            continue
+
+        edge_info = edge_info_by_type[edge_type]
+        raw_column_offset = 0
+        for feature_name in edge_info.feature_keys:
+            if feature_name == weight_feature_name:
+                break
+            feature_spec = edge_info.feature_spec[feature_name]
+            raw_column_offset += feature_spec.shape[-1] if feature_spec.shape else 1
+        weight_logical_index = metadata.raw_feature_indices[raw_column_offset]
+        adjusted_quantized_feature_indices = tuple(
+            quantized_feature_index - 1
+            if quantized_feature_index > weight_logical_index
+            else quantized_feature_index
+            for quantized_feature_index in metadata.quantized_feature_indices
+        )
+        adjusted_metadata[edge_type] = replace(
+            metadata,
+            feature_dim=metadata.feature_dim - 1,
+            quantized_feature_indices=adjusted_quantized_feature_indices,
+        )
+
+    if is_homogeneous:
+        return adjusted_metadata[DEFAULT_HOMOGENEOUS_EDGE_TYPE]
+    return adjusted_metadata
 
 
 def _data_loading_process(
@@ -199,14 +327,6 @@ def _data_loading_process(
                 raise NotImplementedError(
                     "Label keys are not supported for edge entities"
                 )
-            if (
-                serialized_entity_tf_record_info.packed_feature_key is not None
-                and not serialized_entity_tf_record_info.is_node_entity
-            ):
-                # TODO(quantization): Support feature quantization for edge features.
-                raise NotImplementedError(
-                    "Packed feature keys are not supported for edge entities"
-                )
             loaded_entity = tf_record_dataloader.load_as_torch_tensors(
                 serialized_tf_record_info=serialized_entity_tf_record_info,
                 tf_dataset_options=tf_dataset_options,
@@ -396,6 +516,11 @@ def load_torch_tensors_from_tf_record(
         loaded_graph_tensors (LoadedGraphTensors): Unpartitioned Graph Tensors
     """
 
+    _validate_weight_edge_feature_name(
+        edge_entity_info=serialized_graph_metadata.edge_entity_info,
+        weight_edge_feat_name=weight_edge_feat_name,
+    )
+
     logger.info(f"Rank {rank} starting loading torch tensors from serialized info ...")
     start_time = time.time()
 
@@ -525,6 +650,9 @@ def load_torch_tensors_from_tf_record(
 
     edge_index = edge_output_dict[_ID_FMT.format(entity=_EDGE_KEY)]
     edge_features = edge_output_dict.get(_FEATURE_FMT.format(entity=_EDGE_KEY), None)
+    edge_quantized_features = edge_output_dict.get(
+        _PACKED_FEATURE_FMT.format(entity=_EDGE_KEY), None
+    )
     edge_weights = edge_output_dict.get(_EDGE_WEIGHTS_KEY, None)
 
     positive_labels = edge_output_dict.get(
@@ -552,6 +680,7 @@ def load_torch_tensors_from_tf_record(
         node_labels=node_labels,
         edge_index=edge_index,
         edge_features=edge_features,
+        edge_quantized_features=edge_quantized_features,
         positive_label=positive_labels,
         negative_label=negative_labels,
         edge_weights=edge_weights,

gigl/common/utils/feature_quantization/README.md

@@ -22,11 +22,11 @@ this as a useful tradeoff for GiGL.
 The built-in flow is:
 
 1. The data preprocessor computes feature summary statistics offline.
-2. The preprocessor quantizes selected scalar feature columns with NumPy.
+2. The preprocessor quantizes selected scalar node or main-edge feature columns with NumPy.
 3. The packed `uint8` feature sidecar is written to TFRecords.
 4. Distributed dataset construction partitions and samples the packed bytes.
 5. The dataloader collate path dequantizes sampled packed features with Torch.
-6. Dequantized columns are scattered back into the logical `x` feature matrix.
+6. Dequantized columns are scattered back into the logical `x` or `edge_attr` feature matrix.
 
 The NumPy/Torch split is intentional:
 

gigl/distributed/base_dist_loader.py

@@ -246,6 +246,7 @@ def __init__(
         self._node_feature_info = dataset_schema.node_feature_info
         self._edge_feature_info = dataset_schema.edge_feature_info
         self._node_quantization_metadata = dataset_schema.node_quantization_metadata
+        self._edge_quantization_metadata = dataset_schema.edge_quantization_metadata
 
         self._sampler_options = sampler_options
         self._non_blocking_transfers = non_blocking_transfers
@@ -435,7 +436,10 @@ def create_sampling_config(
             batch_size=batch_size,
             shuffle=shuffle,
             drop_last=drop_last,
-            with_edge=dataset_schema.edge_feature_info is not None,
+            with_edge=(
+                dataset_schema.edge_feature_info is not None
+                or dataset_schema.edge_quantization_metadata is not None
+            ),
             collect_features=True,
             with_neg=False,
             with_weight=with_weight,

gigl/distributed/base_sampler.py

@@ -19,6 +19,7 @@
 
 from gigl.common.logger import Logger
 from gigl.distributed.sampler import (
+    EDGE_PACKED_FEATURES_METADATA_KEY,
     NEGATIVE_LABEL_METADATA_KEY,
     NODE_PACKED_FEATURES_METADATA_KEY,
     POSITIVE_LABEL_METADATA_KEY,
@@ -116,6 +117,7 @@ def __init__(self, *args, **kwargs) -> None:
         self._sampling_error_sent: bool = False
 
         self.dist_node_quantized_feature: Optional[DistFeature] = None
+        self.dist_edge_quantized_feature: Optional[DistFeature] = None
         if (
             self.collect_features
             and data is not None
@@ -132,6 +134,20 @@ def __init__(self, *args, **kwargs) -> None:
                 rpc_router=self.rpc_router,
                 device=self.device,
             )
+        if (
+            self.collect_features
+            and data is not None
+            and getattr(data, "edge_quantized_features", None) is not None
+        ):
+            self.dist_edge_quantized_feature = DistFeature(
+                data.num_partitions,
+                data.partition_idx,
+                data.edge_quantized_features,
+                data.edge_pb,
+                local_only=False,
+                rpc_router=self.rpc_router,
+                device=self.device,
+            )
 
     def _prepare_sample_loop_inputs(
         self,
@@ -436,6 +452,8 @@ async def _collate_fn(
                         )
             if self.dist_edge_feature is not None and self.with_edge:
                 for etype in self.edge_types:
+                    if etype not in self.dist_edge_feature.local_feature:
+                        continue
                     if self.edge_dir == "in":
                         eids = result_map.get(
                             f"{as_str(reverse_edge_type(etype))}.eids", None
@@ -451,6 +469,26 @@ async def _collate_fn(
                         futs[result_key] = wrap_torch_future(
                             self.dist_edge_feature.async_get(eids, etype)
                         )
+            if self.dist_edge_quantized_feature is not None and self.with_edge:
+                for etype in self.edge_types:
+                    # Like node features, an edge partition book covers every
+                    # edge type while a feature store may register only some.
+                    if etype not in self.dist_edge_quantized_feature.local_feature:
+                        continue
+                    result_edge_type = (
+                        reverse_edge_type(etype) if self.edge_dir == "in" else etype
+                    )
+                    eids = result_map.get(f"{as_str(result_edge_type)}.eids")
+                    if eids is not None:
+                        eids = eids.to(torch.long)
+                        # GLT maps incoming wire edge types back to the dataset edge
+                        # type during collation. Metadata bypasses that mapping, so its
+                        # transport key must already match the final output store.
+                        futs[f"#META.{EDGE_PACKED_FEATURES_METADATA_KEY}.{etype}"] = (
+                            wrap_torch_future(
+                                self.dist_edge_quantized_feature.async_get(eids, etype)
+                            )
+                        )
             if output.batch is not None:
                 for ntype, batch in output.batch.items():
                     result_map[f"{as_str(ntype)}.batch"] = batch
@@ -490,6 +528,10 @@ async def _collate_fn(
                 futs["efeats"] = wrap_torch_future(
                     self.dist_edge_feature.async_get(eids)
                 )
+            if self.dist_edge_quantized_feature is not None:
+                futs[f"#META.{EDGE_PACKED_FEATURES_METADATA_KEY}"] = wrap_torch_future(
+                    self.dist_edge_quantized_feature.async_get(result_map["eids"])
+                )
             if output.batch is not None:
                 result_map["batch"] = output.batch
 

gigl/distributed/dataset_factory.py

@@ -25,6 +25,7 @@
     SerializedGraphMetadata,
     TFDatasetOptions,
     load_torch_tensors_from_tf_record,
+    remove_sampling_weight_from_edge_quantization_metadata,
 )
 from gigl.common.logger import Logger
 from gigl.common.utils.decorator import tf_on_cpu
@@ -194,6 +195,10 @@ def _load_and_build_partitioned_dataset(
         partitioner.register_edge_features(
             edge_features=loaded_graph_tensors.edge_features
         )
+    if loaded_graph_tensors.edge_quantized_features is not None:
+        partitioner.register_edge_quantized_features(
+            edge_quantized_features=loaded_graph_tensors.edge_quantized_features
+        )
     if loaded_graph_tensors.positive_label is not None:
         partitioner.register_labels(
             label_edge_index=loaded_graph_tensors.positive_label, is_positive=True
@@ -212,6 +217,7 @@ def _load_and_build_partitioned_dataset(
         loaded_graph_tensors.node_quantized_features,
         loaded_graph_tensors.edge_index,
         loaded_graph_tensors.edge_features,
+        loaded_graph_tensors.edge_quantized_features,
         loaded_graph_tensors.edge_weights,
         loaded_graph_tensors.positive_label,
         loaded_graph_tensors.negative_label,
@@ -227,6 +233,10 @@ def _load_and_build_partitioned_dataset(
         world_size=world_size,
         edge_dir=edge_dir,
         node_quantization_metadata=serialized_graph_metadata.node_quantization_metadata,
+        edge_quantization_metadata=remove_sampling_weight_from_edge_quantization_metadata(
+            serialized_graph_metadata=serialized_graph_metadata,
+            weight_edge_feat_name=weight_edge_feat_name,
+        ),
     )
 
     dataset.build(