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

[mkolodner-sc]

Scope of work done

• Adds ability for TFRecordDataloader to load labels as part of the features and corresponding tests
• Adds function for separating labels from features and corresponding tests
• Updates dataset factory to call this function to separate node labels if they are present

Explanation of Node Classification Labels

1. As input to our node classification pipeline, we should have prepared bigquery node feature tables which have columns corresponding to the node classification label.

Taken from external-snap-ci-github-gigl.public_gigl.cora_homogeneous_supervised_node_classification_edge_features_paper_nodes_2024-07-15--21-30-07-UTC

2. In our Data Preprocessor config, labels are specified [1] in the get_nodes_preprocessing_spec function.
3. After running data preprocessor, we have generated TFRecords with the node label as a feature. We can see the preprocessed metadata output at gs://public-gigl/mocked_assets/2024-07-15--21-30-07-UTC/cora_homogeneous_supervised_node_classification/data_preprocess/preprocessed_metadata.yaml:

...
    - f1430
    - f1431
    - f1432
    labelKeys:
    - node_label
    nodeIdKey: node_id
    schemaUri: gs://public-gigl/mocked_assets/2024-07-15--21-30-07-UTC/cora_homogeneous_supervised_node_classification/data_preprocess/node_features_dir/paper/schema.pbtxt
    tfrecordUriPrefix: gs://public-gigl/mocked_assets/2024-07-15--21-30-07-UTC/cora_homogeneous_supervised_node_classification/data_preprocess/node_features_dir/paper/features/

4. We load the preprocessed metadata into a PreprocessedMetadataPbWrapper. This creates and isolates fields we'll need for loading the TFRecords like the FeatureSchemaDict, FeatureDim, LabelKeys, etc. As of this PR, the FeatureSchemaDict includes the feature schemas for both the feature keys and label keys.
5. We pass this information to the SerializedTFRecordInfo.
6. We load the node features and labels, with the labels being appended as an additional feature to the normal feature tensor.
7. We call our partitioner and partition the "node features" across all machines. The benefit of having labels as part of this process is that 1) there are less distributed calls and 2) there are guarantees that the mapping of global_id to local_id (id2idx) is correct for both node features and node labels.
8. We extract the labels from the partitioned node features with _get_labels_from_features.
9. We register those labels to the dataset in dataset.build() (TODO in the next PR)

Where is the documentation for this feature?: N/A

Did you add automated tests or write a test plan?

Updated Changelog.md? NO

Ready for code review?: NO

[yliu2-sc]

I'm ok with this implementation. Only nit is that if we can store labels separately from features.

[svij-sc]

Looks okay to me.

[mkolodner-sc]

/unit_test

[mkolodner-sc]

/e2e_test

[github-actions]

GiGL Automation

@ 06:44:04UTC : 🔄 Unit Test started.

@ 07:20:10UTC : ❌ Workflow failed.
Please check the logs for more details.

[github-actions]

GiGL Automation

@ 06:44:05UTC : 🔄 E2E Test started.

@ 08:08:05UTC : ✅ Workflow completed successfully.

[mkolodner-sc]

/unit_test

[github-actions]

GiGL Automation

@ 08:16:53UTC : 🔄 Unit Test started.

@ 08:53:30UTC : ✅ Workflow completed successfully.