[GSProcessing] Enforce re-order for node label processing during clas…

…sification (#1136)

*Issue #, if available:*

Fixes #1135 
Fixes #1138

*Description of changes:*

* We guarantee ordering for node label classification by ordering the
transformed label DF after processing by the NODE_INT_MAPPING id, and
doing the same for masks.
* Because there is no guarantee for order when writing to Parquet from
Spark _even for ordered Spark DataFrames_, we collect the labels and
masks to a Pandas DF on the Spark leader, and write that using pyarrow.
* NOTE: DistPart requires at least N files to be present for every
node/edge type, where N is the number of requested partitions. So when
writing back the masks and label we still use the number of incoming
partitions to create as many files.


By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

graphstorm-processing/graphstorm_processing/config/label_config_base.py

@@ -18,6 +18,8 @@
 import logging
 from typing import Any, Dict, Optional
 
+from graphstorm_processing.constants import VALID_TASK_TYPES
+
 
 class LabelConfig(abc.ABC):
     """Basic class for label config"""
@@ -55,6 +57,9 @@ def __init__(self, config_dict: Dict[str, Any]):
             self._mask_field_names = None
 
     def _sanity_check(self):
+        assert (
+            self._task_type in VALID_TASK_TYPES
+        ), f"Invalid task type {self._task_type}, must be one of {VALID_TASK_TYPES}"
         if self._label_column == "":
             assert self._task_type == "link_prediction", (
                 "When no label column is specified, the task type must be link_prediction, "
@@ -83,6 +88,25 @@ def _sanity_check(self):
             assert all(isinstance(x, str) for x in self._mask_field_names)
             assert len(self._mask_field_names) == 3
 
+    def __repr__(self) -> str:
+        """Formal object representation for debugging"""
+        return (
+            f"{self.__class__.__name__}(label_column={self._label_column!r}, "
+            f"task_type={self._task_type!r}, separator={self._separator!r}, "
+            f"multilabel={self._multilabel!r}, split={self._split!r}, "
+            f"custom_split_filenames={self._custom_split_filenames!r}, "
+            f"mask_field_names={self._mask_field_names!r})"
+        )
+
+    def __str__(self) -> str:
+        """Informal object representation for readability"""
+        task_desc = f"{self._task_type} task"
+        if self._label_column:
+            task_desc += f" on column '{self._label_column}'"
+        if self._multilabel:
+            task_desc += f" (multilabel with separator '{self._separator}')"
+        return task_desc
+
     @property
     def label_column(self) -> str:
         """The name of the column storing the target label property value."""

graphstorm-processing/graphstorm_processing/constants.py

@@ -66,6 +66,16 @@
 NODE_MAPPING_STR = "orig"
 NODE_MAPPING_INT = "new"
 
+################# Reserved columns  ################
+DATA_SPLIT_SET_MASK_COL = "GSP-SAMPLE-SET-MASK"
+
+################# Supported task types  ##############
+VALID_TASK_TYPES = {
+    "classification",
+    "regression",
+    "link_prediction",
+}
+
 
 ################# Supported execution envs  ##############
 class ExecutionEnv(Enum):

graphstorm-processing/graphstorm_processing/data_transformations/dist_label_loader.py

@@ -16,9 +16,10 @@
 
 from dataclasses import dataclass
 from math import fsum
+from typing import Optional
 
 from pyspark.sql import DataFrame, SparkSession
-from pyspark.sql.types import FloatType
+from pyspark.sql.types import NumericType
 
 from graphstorm_processing.config.label_config_base import LabelConfig
 from graphstorm_processing.data_transformations.dist_transformations import (
@@ -100,11 +101,33 @@ class DistLabelLoader:
         The SparkSession to use for processing.
     """
 
-    def __init__(self, label_config: LabelConfig, spark: SparkSession) -> None:
+    def __init__(
+        self, label_config: LabelConfig, spark: SparkSession, order_col: Optional[str] = None
+    ) -> None:
         self.label_config = label_config
         self.label_column = label_config.label_column
         self.spark = spark
         self.label_map: dict[str, int] = {}
+        self.order_col = order_col
+
+    def __str__(self) -> str:
+        """Informal object representation for readability"""
+        return (
+            f"DistLabelLoader(label_column='{self.label_column}', "
+            f"task_type='{self.label_config.task_type}', "
+            f"multilabel={self.label_config.multilabel}, "
+            f"order_col={self.order_col!r})"
+        )
+
+    def __repr__(self) -> str:
+        """Formal object representation for debugging"""
+        return (
+            f"DistLabelLoader("
+            f"label_config={self.label_config!r}, "
+            f"spark={self.spark!r}, "
+            f"order_col={self.order_col!r}, "
+            f"label_map={self.label_map!r})"
+        )
 
     def process_label(self, input_df: DataFrame) -> DataFrame:
         """Transforms the label column in the input DataFrame to conform to GraphStorm expectations.
@@ -134,23 +157,32 @@ def process_label(self, input_df: DataFrame) -> DataFrame:
         label_type = input_df.schema[self.label_column].dataType
 
         if self.label_config.task_type == "classification":
+            assert self.order_col, f"{self.order_col} must be provided for classification tasks"
             if self.label_config.multilabel:
                 assert self.label_config.separator
                 label_transformer = DistMultiLabelTransformation(
                     [self.label_config.label_column], self.label_config.separator
                 )
             else:
                 label_transformer = DistSingleLabelTransformation(
-                    [self.label_config.label_column], self.spark
+                    [self.label_config.label_column],
+                    self.spark,
                 )
 
-            transformed_label = label_transformer.apply(input_df).select(self.label_column)
+            transformed_label = label_transformer.apply(input_df)
+            if self.order_col:
+                assert self.order_col in transformed_label.columns, (
+                    f"{self.order_col=} needs to be part of transformed "
+                    f"label DF, got {transformed_label.columns=}"
+                )
+                transformed_label = transformed_label.sort(self.order_col).cache()
+
             self.label_map = label_transformer.value_map
             return transformed_label
         elif self.label_config.task_type == "regression":
-            if not isinstance(label_type, FloatType):
+            if not isinstance(label_type, NumericType):
                 raise RuntimeError(
-                    "Data type for regression should be FloatType, "
+                    "Data type for regression should be a NumericType, "
                     f"got {label_type} for {self.label_column}"
                 )
             return input_df.select(self.label_column)

graphstorm-processing/graphstorm_processing/data_transformations/dist_transformations/dist_category_transformation.py

@@ -313,8 +313,12 @@ def __init__(self, cols: Sequence[str], separator: str) -> None:
     def get_transformation_name() -> str:
         return "DistMultiCategoryTransformation"
 
-    def apply(self, input_df: DataFrame) -> DataFrame:
+    def apply(self, input_df: DataFrame, return_all_cols: bool = False) -> DataFrame:
         col_datatype = input_df.schema[self.multi_column].dataType
+        if return_all_cols:
+            original_cols = {*input_df.columns} - {self.multi_column}
+        else:
+            original_cols = {}
         is_array_col = False
         if col_datatype.typeName() == "array":
             assert isinstance(col_datatype, ArrayType)
@@ -326,13 +330,19 @@ def apply(self, input_df: DataFrame) -> DataFrame:
 
             is_array_col = True
 
+        # Parquet input might come with arrays already, CSV will need splitting
         if is_array_col:
-            list_df = input_df.select(self.multi_column).alias(self.multi_column)
+            multi_column = F.col(self.multi_column)
         else:
-            list_df = input_df.select(
-                F.split(F.col(self.multi_column), self.separator).alias(self.multi_column)
+            multi_column = F.split(F.col(self.multi_column), self.separator).alias(
+                self.multi_column
             )
 
+        list_df = input_df.select(
+            multi_column,
+            *original_cols,
+        )
+
         distinct_category_counts = (
             list_df.withColumn(SINGLE_CATEGORY_COL, F.explode(F.col(self.multi_column)))
             .groupBy(SINGLE_CATEGORY_COL)

graphstorm-processing/graphstorm_processing/data_transformations/dist_transformations/dist_label_transformation.py

@@ -45,6 +45,7 @@ def __init__(self, cols: Sequence[str], spark: SparkSession) -> None:
 
     def apply(self, input_df: DataFrame) -> DataFrame:
         assert self.spark
+        original_cols = {*input_df.columns} - {self.label_column}
         processed_col_name = self.label_column + "_processed"
 
         str_indexer = StringIndexer(
@@ -63,13 +64,15 @@ def apply(self, input_df: DataFrame) -> DataFrame:
 
         # Labels that were missing and were assigned the value numLabels by the StringIndexer
         # are converted to None
-        long_class_label = indexed_df.select(F.col(self.label_column).cast("long")).select(
+        long_class_label = indexed_df.select(
             F.when(
                 F.col(self.label_column) == len(str_indexer_model.labelsArray[0]),  # type: ignore
                 F.lit(None),
             )
             .otherwise(F.col(self.label_column))
-            .alias(self.label_column)
+            .cast("long")
+            .alias(self.label_column),
+            *original_cols,
         )
 
         # Get a mapping from original label to encoded value
@@ -112,7 +115,7 @@ def __init__(self, cols: Sequence[str], separator: str) -> None:
         super().__init__(cols, separator)
         self.label_column = cols[0]
 
-    def apply(self, input_df: DataFrame) -> DataFrame:
-        multi_cat_df = super().apply(input_df)
+    def apply(self, input_df: DataFrame, return_all_cols=True) -> DataFrame:
+        multi_cat_df = super().apply(input_df, return_all_cols=return_all_cols)
 
         return multi_cat_df

graphstorm-processing/graphstorm_processing/distributed_executor.py

@@ -264,6 +264,7 @@ def __init__(
             output_prefix=self.output_prefix,
             precomputed_transformations=self.precomputed_transformations,
         )
+
         self.loader = DistHeterogeneousGraphLoader(
             self.spark,
             loader_config,
@@ -287,17 +288,18 @@ def _upload_output_files(self, loader: DistHeterogeneousGraphLoader, force=False
             bucket, s3_prefix = s3_utils.extract_bucket_and_key(self.output_prefix)
             s3 = boto3.resource("s3")
 
-            output_files = os.listdir(loader.output_path)
+            output_files = os.listdir(loader.local_meta_output_path)
             for output_file in output_files:
                 s3.meta.client.upload_file(
-                    f"{os.path.join(loader.output_path, output_file)}",
+                    f"{os.path.join(loader.local_meta_output_path, output_file)}",
                     bucket,
                     f"{s3_prefix}/{output_file}",
                 )
 
     def run(self) -> None:
         """
-        Executes the Spark processing job.
+        Executes the Spark processing job, optional repartition job, and uploads any metadata files
+        if needed.
         """
         logging.info("Performing data processing with PySpark...")
 
@@ -355,7 +357,7 @@ def run(self) -> None:
         # If any of the metadata modification took place, write an updated metadata file
         if updated_metadata:
             updated_meta_path = os.path.join(
-                self.loader.output_path, "updated_row_counts_metadata.json"
+                self.loader.local_meta_output_path, "updated_row_counts_metadata.json"
             )
             with open(
                 updated_meta_path,