Run SubgraphSampler on Dataproc 2.2 / Spark 3.5 before image 2.0 becomes uncreatable

gigl/src/subgraph_sampler/subgraph_sampler.py

@@ -97,10 +97,13 @@ def run(
                 "graph_db_config"
             )
         )
+        # Default to the Spark 3.5 runner (Dataproc 2.2): Google blocks cluster creation on
+        # Dataproc image 2.0 starting 2026-08-25. Setting the `use_spark35_runner`
+        # experimental flag to "False" remains a temporary escape hatch until then.
         use_spark35: bool = bool(
             strtobool(
                 gbml_config_pb_wrapper.dataset_config.subgraph_sampler_config.experimental_flags.get(
-                    "use_spark35_runner", "False"
+                    "use_spark35_runner", "True"
                 )
             )
         )

scala_spark35/subgraph_sampler/src/main/scala/libs/TaskRunner.scala

@@ -6,7 +6,9 @@ import common.types.pb_wrappers.TaskMetadataPbWrapper
 import common.types.pb_wrappers.TaskMetadataType
 import libs.task.SubgraphSamplerTask
 import libs.task.graphdb.GraphDBNodeAnchorBasedLinkPredictionTask
+import libs.task.pureSpark.NodeAnchorBasedLinkPredictionTask
 import libs.task.pureSpark.SupervisedNodeClassificationTask
+import libs.task.pureSpark.UserDefinedLabelsNodeAnchorBasedLinkPredictionTask
 import libs.task.pureSparkV2.EgoNetGeneration
 
 object TaskRunner {
@@ -67,9 +69,27 @@ object TaskRunner {
             giglResourceConfigWrapper = giglResourceConfigWrapper,
           )
         } else {
-          throw new Exception(
-            "PureSpark SGS not supported in spark35 yet.",
-          )
+          // TODO for heterogeneous will have to update this to be more flexible to detect UDL for different edge types
+          val defaultCondensedEdgeType: Int =
+            gbmlConfigWrapper.preprocessedMetadataWrapper.preprocessedMetadataPb.condensedEdgeTypeToPreprocessedMetadata.keysIterator
+              .next()
+          val isPosUserDefined: Boolean = !gbmlConfigWrapper.preprocessedMetadataWrapper
+            .isPosUserDefinedForCondensedEdgeType(condensedEdgeType = defaultCondensedEdgeType)
+          val isNegUserDefined: Boolean = !gbmlConfigWrapper.preprocessedMetadataWrapper
+            .isNegUserDefinedForCondensedEdgeType(condensedEdgeType = defaultCondensedEdgeType)
+          if (isPosUserDefined || isNegUserDefined) {
+            new UserDefinedLabelsNodeAnchorBasedLinkPredictionTask(
+              gbmlConfigWrapper = gbmlConfigWrapper,
+              graphMetadataPbWrapper = gbmlConfigWrapper.graphMetadataPbWrapper,
+              isPosUserDefined = isPosUserDefined,
+              isNegUserDefined = isNegUserDefined,
+            )
+          } else {
+            new NodeAnchorBasedLinkPredictionTask(
+              gbmlConfigWrapper = gbmlConfigWrapper,
+              graphMetadataPbWrapper = gbmlConfigWrapper.graphMetadataPbWrapper,
+            )
+          }
         }
       } else {
         // TODO: (svij) Note this implementation is copy from spark 3.2 implementation but untested

scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/NodeAnchorBasedLinkPredictionBaseTask.scala

@@ -396,11 +396,11 @@ abstract class NodeAnchorBasedLinkPredictionBaseTask(
         STRUCT( _root_node AS node_id,
           _condensed_node_type AS condensed_node_type,
           _node_features AS feature_values ) AS root_node,
-        ARRAY() AS hard_neg_edges,
-        _pos_hydrated_edges AS pos_edges,
-        ARRAY() AS neg_edges,
-        STRUCT( _neighbor_nodes AS nodes,
-          _neighbor_edges AS edges ) AS neighborhood
+        ${emptyEdgeArrayWithProtoSchema} AS hard_neg_edges,
+        ${castEdgeArrayToProtoSchema("_pos_hydrated_edges")} AS pos_edges,
+        ${emptyEdgeArrayWithProtoSchema} AS neg_edges,
+        STRUCT( ${castNodeArrayToProtoSchema("_neighbor_nodes")} AS nodes,
+          ${castEdgeArrayToProtoSchema("_neighbor_edges")} AS edges ) AS neighborhood
       FROM
         ${dfVIEW}
                 """)
@@ -410,11 +410,11 @@ abstract class NodeAnchorBasedLinkPredictionBaseTask(
         STRUCT( _root_node AS node_id,
           _condensed_node_type AS condensed_node_type,
           _node_features AS feature_values ) AS root_node,
-        _neg_hydrated_edges AS hard_neg_edges,
-        _pos_hydrated_edges AS pos_edges,
-        ARRAY() AS neg_edges,
-        STRUCT( _neighbor_nodes AS nodes,
-          _neighbor_edges AS edges ) AS neighborhood
+        ${castEdgeArrayToProtoSchema("_neg_hydrated_edges")} AS hard_neg_edges,
+        ${castEdgeArrayToProtoSchema("_pos_hydrated_edges")} AS pos_edges,
+        ${emptyEdgeArrayWithProtoSchema} AS neg_edges,
+        STRUCT( ${castNodeArrayToProtoSchema("_neighbor_nodes")} AS nodes,
+          ${castEdgeArrayToProtoSchema("_neighbor_edges")} AS edges ) AS neighborhood
       FROM
         ${dfVIEW}
                 """)

scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/SGSPureSparkV1Task.scala

@@ -12,13 +12,16 @@ import libs.task.SamplingStrategy.shuffleBasedUniformPermutation
 import libs.task.SubgraphSamplerTask
 import org.apache.spark.sql.Column
 import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.expressions.UserDefinedFunction
+import org.apache.spark.sql.functions.udf
 import org.apache.spark.sql.types.IntegerType
 import org.apache.spark.sql.{functions => F}
 import org.apache.spark.storage.StorageLevel
 import snapchat.research.gbml.preprocessed_metadata.PreprocessedMetadata
 
 import java.util.UUID.randomUUID
 import scala.collection.mutable.ListBuffer
+import scala.util.Random
 
 abstract class SGSPureSparkV1Task(
   gbmlConfigWrapper: GbmlConfigPbWrapper)
@@ -36,6 +39,16 @@ abstract class SGSPureSparkV1Task(
 
   val uniqueTempViewSuffix: String = generateUniqueSuffix
 
+  val sampleWithReplacementUDF: UserDefinedFunction = udf((array: Seq[Int], numSamples: Int) => {
+    if (array == null || array.isEmpty) {
+      Seq.empty[Int]
+    } else {
+      val random = new Random()
+      (1 to numSamples).map(_ => array(random.nextInt(array.length)))
+    }
+  })
+  spark.udf.register("sampleWithReplacementUDF", sampleWithReplacementUDF)
+
   def loadNodeDataframeIntoSparkSql(condensedNodeType: Int): String = {
 
     /** For a given condensed_node_type, loads the node feature/type dataframe with columns:
@@ -302,6 +315,7 @@ abstract class SGSPureSparkV1Task(
     numNeighborsToSample: Int,
     unhydratedEdgeVIEW: String,
     permutationStrategy: String,
+    sampleWithReplacement: Boolean = false,
   ): String = {
 
     /** 1. for each dst node, take all in-edges as onehop array 2. randomly shuffle onehop array
@@ -339,13 +353,27 @@ abstract class SGSPureSparkV1Task(
     val permutedOnehopArrayVIEW = "permutedOnehopArrayDF" + uniqueTempViewSuffix
     permutedOnehopArrayDF.createOrReplaceTempView(permutedOnehopArrayVIEW)
 
-    val sampledOnehopDF: DataFrame = spark.sql(f"""
-          SELECT
-            _dst_node AS _0_hop,
-            slice(_shuffled_1_hop_arr, 1, ${numNeighborsToSample}) AS _sampled_1_hop_arr
-          FROM
-            ${permutedOnehopArrayVIEW}
-          """)
+    val sampledOnehopDF: DataFrame = if (sampleWithReplacement) {
+      spark.sql(
+        f"""
+        SELECT
+          _dst_node AS _0_hop,
+          sampleWithReplacementUDF(_shuffled_1_hop_arr, ${numNeighborsToSample}) AS _sampled_1_hop_arr
+        FROM
+          ${permutedOnehopArrayVIEW}
+         """,
+      )
+    } else {
+      spark.sql(
+        f"""
+        SELECT
+          _dst_node AS _0_hop,
+          slice(_shuffled_1_hop_arr, 1, ${numNeighborsToSample}) AS _sampled_1_hop_arr
+        FROM
+          ${permutedOnehopArrayVIEW}
+        """,
+      )
+    }
 
     // @spark: critical NOTE: cache is necessary here, to break parallelism and enforce random shuffle/sampling happens once, to circumvent NON-determinism in F.shuffle. Otherwise, for all calls to sampledOnehopDF downstream, this stage will run in parallel and mess up onehop samples!
     // @spark: maybe in future we wanna try caching the exploded verison of below DF.
@@ -365,6 +393,7 @@ abstract class SGSPureSparkV1Task(
     unhydratedEdgeVIEW: String,
     sampledOnehopVIEW: String,
     permutationStrategy: String,
+    sampleWithReplacement: Boolean = false,
   ): String = {
 
     /**   1. uses onehop nodes in sampledOnehopVIEW as reference to obtain twohop neighbors for each
@@ -431,14 +460,29 @@ abstract class SGSPureSparkV1Task(
     val permutedTwohopArrayVIEW = "permutedTwohopArrayDF" + uniqueTempViewSuffix
     permutedTwohopArrayDF.createOrReplaceTempView(permutedTwohopArrayVIEW)
     // @spark: Since twohop df is called only once, no need to take care of NON determinism in the shuffle/sampling (for k hop any k-1 df must be cached)
-    val sampledTwohopDF: DataFrame = spark.sql(f"""
+    val sampledTwohopDF: DataFrame = if (sampleWithReplacement) {
+      spark.sql(
+        f"""
+        SELECT
+          _0_hop,
+          _1_hop,
+          sampleWithReplacementUDF(_shuffled_2_hop_arr, ${numNeighborsToSample}) AS _sampled_2_hop_arr
+        FROM
+          ${permutedTwohopArrayVIEW}
+         """,
+      )
+    } else {
+      spark.sql(
+        f"""
         SELECT
           _0_hop,
           _1_hop,
           slice(_shuffled_2_hop_arr ,1 , ${numNeighborsToSample}) AS _sampled_2_hop_arr
         FROM
           ${permutedTwohopArrayVIEW}
-        """)
+        """,
+      )
+    }
     val sampledTwohopVIEW = "sampledTwohopDF" + uniqueTempViewSuffix
     sampledTwohopDF.createOrReplaceTempView(sampledTwohopVIEW)
     //
@@ -631,6 +675,7 @@ abstract class SGSPureSparkV1Task(
     unhydratedEdgeVIEW: String,
     numNeighborsToSample: Int,
     permutationStrategy: String,
+    sampleWithReplacement: Boolean = false,
   ): String = {
 
     /** Adds root node features to hydrated neighborhood and creates final subgraphDF for each root
@@ -650,12 +695,14 @@ abstract class SGSPureSparkV1Task(
       numNeighborsToSample = numNeighborsToSample,
       unhydratedEdgeVIEW = unhydratedEdgeVIEW,
       permutationStrategy = permutationStrategy,
+      sampleWithReplacement = sampleWithReplacement,
     )
     val sampledTwohopVIEW = sampleTwohopSrcNodesUniformly(
       numNeighborsToSample = numNeighborsToSample,
       unhydratedEdgeVIEW = unhydratedEdgeVIEW,
       sampledOnehopVIEW = sampledOnehopVIEW,
       permutationStrategy = permutationStrategy,
+      sampleWithReplacement = sampleWithReplacement,
     )
     // hydrate onehop neighbors
     val hydratedOnehopNeighborsVIEW = createKthHydratedNeighborhood(
@@ -970,6 +1017,51 @@ abstract class SGSPureSparkV1Task(
     subgraphsWithNeighborlessNodesVIEW
   }
 
+  /** SQL fragment renaming a node array's *element* fields to the `Node` proto field names.
+    *
+    * Everywhere inside the sampling pipeline, array-of-struct elements carry `_`-prefixed field
+    * names (`_node_id`, `_condensed_node_type`, `_feature_values`). The proto expects
+    * `node_id`, `condensed_node_type`, `feature_values`, and sparksql-scalapb resolves repeated
+    * message fields **by name**, so the element fields have to be renamed before `.as[proto]`.
+    *
+    * `transform` is NULL-safe (`transform(NULL, ...)` is NULL), which matters because several UNION
+    * branches emit `NULL` for their neighbor arrays.
+    *
+    * @param nodeArrayColumnName column holding `ARRAY<STRUCT<_node_id, _condensed_node_type, _feature_values>>`
+    */
+  protected def castNodeArrayToProtoSchema(nodeArrayColumnName: String): String =
+    s"""transform(${nodeArrayColumnName}, node -> struct(
+       |    node._node_id AS node_id,
+       |    node._condensed_node_type AS condensed_node_type,
+       |    node._feature_values AS feature_values
+       |  ))""".stripMargin
+
+  /** SQL fragment renaming an edge array's *element* fields to the `Edge` proto field names.
+    *
+    * Same by-name-resolution reasoning as [[castNodeArrayToProtoSchema]]. Note the source struct
+    * calls the endpoints `_src_node` / `_dst_node` while the proto calls them `src_node_id` /
+    * `dst_node_id`.
+    *
+    * @param edgeArrayColumnName column holding `ARRAY<STRUCT<_src_node, _dst_node, _condensed_edge_type, _feature_values>>`
+    */
+  protected def castEdgeArrayToProtoSchema(edgeArrayColumnName: String): String =
+    s"""transform(${edgeArrayColumnName}, edge -> struct(
+       |    edge._src_node AS src_node_id,
+       |    edge._dst_node AS dst_node_id,
+       |    edge._condensed_edge_type AS condensed_edge_type,
+       |    edge._feature_values AS feature_values
+       |  ))""".stripMargin
+
+  /** SQL fragment for an always-empty edge array with an explicit `Edge` proto element type.
+    *
+    * A bare `ARRAY()` literal is `ARRAY<VOID>`. sparksql-scalapb hands the element type to the
+    * analyzer rather than declaring it, so a VOID element type fails with
+    * `INVALID_EXTRACT_BASE_FIELD_TYPE` instead of deserializing to an empty `Seq`.
+    */
+  protected val emptyEdgeArrayWithProtoSchema: String =
+    "CAST(ARRAY() AS ARRAY<STRUCT<src_node_id: INT, dst_node_id: INT, " +
+      "condensed_edge_type: INT, feature_values: ARRAY<FLOAT>>>)"
+
   def castToRootedNodeNeighborhoodProtoSchema(dfVIEW: String): String = {
 
     /** Creates desirable schema for RootedNodeNeighborhood defined in training_samples_schema.proto. Returns rnnVIEW
@@ -983,8 +1075,8 @@ abstract class SGSPureSparkV1Task(
                     _node_features AS feature_values
                   ) AS root_node,
                   struct(
-                    _neighbor_nodes AS nodes,
-                    _neighbor_edges AS edges
+                    ${castNodeArrayToProtoSchema("_neighbor_nodes")} AS nodes,
+                    ${castEdgeArrayToProtoSchema("_neighbor_edges")} AS edges
                   ) AS neighborhood
                 FROM ${dfVIEW}
                 """)

scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/SupervisedNodeClassificationTask.scala

@@ -323,8 +323,8 @@ class SupervisedNodeClassificationTask(
           STRUCT( _root_node AS node_id,
             _condensed_node_type AS condensed_node_type,
             _node_features AS feature_values ) AS root_node,
-          STRUCT( _neighbor_nodes AS nodes,
-            _neighbor_edges AS edges ) AS neighborhood,
+          STRUCT( ${castNodeArrayToProtoSchema("_neighbor_nodes")} AS nodes,
+            ${castEdgeArrayToProtoSchema("_neighbor_edges")} AS edges ) AS neighborhood,
           ARRAY(STRUCT( _label_type AS label_type,
               CAST(_label AS INTEGER) AS label) ) AS root_node_labels
         FROM