From a1888791c820ed0ada3ab5f5946d564077e95698 Mon Sep 17 00:00:00 2001 From: Kyle Montemayor Date: Wed, 5 Aug 2026 16:41:15 +0000 Subject: [PATCH 1/5] Port pureSpark NALP/UDL SubgraphSampler to scala_spark35 Google blocks Dataproc image 2.0 cluster creation on 2026-08-25, so the V1 SubgraphSampler must run on image 2.2 via the scala_spark35 jar. That jar's TaskRunner threw "PureSpark SGS not supported in spark35 yet." for the non-graphdb node-anchor-based link prediction path, which is exactly what the cora_nalp/cora_udl e2e tests use. - TaskRunner: replace the throw with the NALP/UDL dispatch ported from scala/subgraph_sampler TaskRunner (UDL detection via isPos/isNegUserDefinedForCondensedEdgeType). - SGSPureSparkV1Task: restore the sampleWithReplacement feature (UDF, default params on the three sampling methods, SQL branches) that the ported tests exercise. The spark35 comment fixes are kept; this is not a wholesale copy of the spark31 file. - Port the 4 pureSpark test suites (only import-path edits plus the SupervisedNodeClassificationTask ctor difference); test assets already existed in this tree. Deliberately unported: the sample_with_replacement experimental-flag reading in the NALP/UDL task classes (documented parity gap), and any heterogeneous-graph support. sbt "subgraph_sampler/test": 25 tests, 7 suites, 0 failures. Co-Authored-By: Claude Opus 5 --- .../src/main/scala/libs/TaskRunner.scala | 26 +- .../task/pureSpark/SGSPureSparkV1Task.scala | 65 +- ...odeAnchorBasedLinkPredictionTaskTest.scala | 239 +++++++ .../test/scala/SGSPureSparkV1TaskTest.scala | 608 ++++++++++++++++++ ...SupervisedNodeClassificationTaskTest.scala | 48 ++ ...odeAnchorBasedLinkPredictionTaskTest.scala | 299 +++++++++ 6 files changed, 1273 insertions(+), 12 deletions(-) create mode 100644 scala_spark35/subgraph_sampler/src/test/scala/NodeAnchorBasedLinkPredictionTaskTest.scala create mode 100644 scala_spark35/subgraph_sampler/src/test/scala/SGSPureSparkV1TaskTest.scala create mode 100644 scala_spark35/subgraph_sampler/src/test/scala/SupervisedNodeClassificationTaskTest.scala create mode 100644 scala_spark35/subgraph_sampler/src/test/scala/UserDefinedLabelsNodeAnchorBasedLinkPredictionTaskTest.scala diff --git a/scala_spark35/subgraph_sampler/src/main/scala/libs/TaskRunner.scala b/scala_spark35/subgraph_sampler/src/main/scala/libs/TaskRunner.scala index a1a2cd789..022f44f41 100644 --- a/scala_spark35/subgraph_sampler/src/main/scala/libs/TaskRunner.scala +++ b/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 diff --git a/scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/SGSPureSparkV1Task.scala b/scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/SGSPureSparkV1Task.scala index a67142d47..7a8a2a8cf 100644 --- a/scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/SGSPureSparkV1Task.scala +++ b/scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/SGSPureSparkV1Task.scala @@ -12,6 +12,8 @@ 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 @@ -19,6 +21,7 @@ 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( diff --git a/scala_spark35/subgraph_sampler/src/test/scala/NodeAnchorBasedLinkPredictionTaskTest.scala b/scala_spark35/subgraph_sampler/src/test/scala/NodeAnchorBasedLinkPredictionTaskTest.scala new file mode 100644 index 000000000..4138e8144 --- /dev/null +++ b/scala_spark35/subgraph_sampler/src/test/scala/NodeAnchorBasedLinkPredictionTaskTest.scala @@ -0,0 +1,239 @@ +import common.test.testLibs.SharedSparkSession +import common.types.EdgeUsageType +import common.types.pb_wrappers.GbmlConfigPbWrapper +import common.types.pb_wrappers.GraphMetadataPbWrapper +import common.utils.ProtoLoader.populateProtoFromYaml +import libs.task.TaskOutputValidator +import libs.task.pureSpark.NodeAnchorBasedLinkPredictionTask +import org.apache.spark.SparkException +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row +import org.apache.spark.sql.types._ +import org.apache.spark.sql.{functions => F} +import org.scalatest.BeforeAndAfterAll +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers._ +import scalapb.spark.Implicits._ +import scalapb.spark.Implicits._ +import snapchat.research.gbml.gbml_config.GbmlConfig +import snapchat.research.gbml.graph_schema.Edge +import snapchat.research.gbml.graph_schema.Graph +import snapchat.research.gbml.graph_schema.Node +import snapchat.research.gbml.training_samples_schema.NodeAnchorBasedLinkPredictionSample + +import java.util.UUID.randomUUID + +class NodeAnchorBasedLinkPredictionTaskTest + extends AnyFunSuite + with BeforeAndAfterAll + with SharedSparkSession { + + // commmon/reused vars among tests which must be assigned in beforeAll(): + var nablpTask: NodeAnchorBasedLinkPredictionTask = _ + var gbmlConfigWrapper: GbmlConfigPbWrapper = _ + var graphMetadataPbWrapper: GraphMetadataPbWrapper = _ + + override def beforeAll(): Unit = { + super.beforeAll() + val frozenGbmlConfigUriTest = + "common/src/test/assets/subgraph_sampler/node_anchor_based_link_prediction/frozen_gbml_config.yaml" + val gbmlConfigProto = + populateProtoFromYaml[GbmlConfig](uri = frozenGbmlConfigUriTest) + gbmlConfigWrapper = GbmlConfigPbWrapper(gbmlConfigPb = gbmlConfigProto) + + val graphMetadataPbWrapper: GraphMetadataPbWrapper = GraphMetadataPbWrapper( + gbmlConfigWrapper.graphMetadataPb, + ) + + nablpTask = new NodeAnchorBasedLinkPredictionTask( + gbmlConfigWrapper = gbmlConfigWrapper, + graphMetadataPbWrapper = graphMetadataPbWrapper, + ) + + } + + def mockUnhydratedEdgeForCurrentTest: DataFrame = { + // Making scope local to avoid clash with scalapb.spark.Implicits._ + import sqlImplicits._ + // Note that below graph is bidirected + val edgeData = Seq( + (0, 1), + (0, 2), + (0, 3), + (0, 4), + (0, 5), + (0, 6), + (0, 7), + (0, 8), + (1, 2), + (1, 3), + (1, 0), + (2, 0), + (3, 0), + (4, 0), + (5, 0), + (6, 0), + (7, 0), + (8, 0), + (2, 1), + (3, 1), + ) + val unhydratedEdgeDF = edgeData.toDF("_src_node", "_dst_node") + unhydratedEdgeDF + } + + def mockSubgraphForCurrentTest: DataFrame = { + val emptyFeats = Seq.empty[Float] + val subgraphData = Seq( + Row(0, List(Row(1, 0, 0, emptyFeats)), List(Row(1, 0, Seq(0.1f))), Seq(0.0f), 0), + Row(1, List(Row(2, 1, 0, emptyFeats)), List(Row(2, 0, Seq(0.2f))), Seq(0.1f), 0), + Row(2, List(Row(1, 2, 0, emptyFeats)), List(Row(1, 0, Seq(0.1f))), Seq(0.2f), 0), + Row(3, List(Row(0, 3, 0, emptyFeats)), List(Row(0, 0, Seq(0.0f))), Seq(0.3f), 0), + Row(4, List(Row(0, 4, 0, emptyFeats)), List(Row(0, 0, Seq(0.0f))), Seq(0.4f), 0), + Row(5, List(Row(0, 5, 0, emptyFeats)), List(Row(0, 0, Seq(0.0f))), Seq(0.5f), 0), + Row(6, List(Row(0, 6, 0, emptyFeats)), List(Row(0, 0, Seq(0.0f))), Seq(0.6f), 0), + Row(7, List(Row(0, 7, 0, emptyFeats)), List(Row(0, 0, Seq(0.0f))), Seq(0.7f), 0), + Row(8, List(Row(0, 8, 0, emptyFeats)), List(Row(0, 0, Seq(0.0f))), Seq(0.8f), 0), + ) + val arrayStruct = new StructType() + .add("_root_node", IntegerType) + .add( + "_neighbor_edges", + ArrayType( + new StructType() + .add("_src_node", IntegerType) + .add("_dst_node", IntegerType) + .add("_condensed_edge_type", IntegerType) + .add("_feature_values", ArrayType(FloatType)), + ), + ) + .add( + "_neighbor_nodes", + ArrayType( + new StructType() + .add("_node_id", IntegerType) + .add("_condensed_node_type", IntegerType) + .add("_feature_values", ArrayType(FloatType)), + ), + ) + .add("_node_features", ArrayType(FloatType)) + .add("_condensed_node_type", IntegerType) + val subgraphRDD = sparkTest.sparkContext.parallelize(subgraphData) + val mockSubgraphDF = sparkTest.createDataFrame(subgraphRDD, arrayStruct) + + mockSubgraphDF + } + + test("Positive samples are valid.") { + val uniqueTestViewSuffix = "_" + randomUUID.toString.replace("-", "_") + val unhydratedEdgeDF = mockUnhydratedEdgeForCurrentTest + val unhydratedEdgeVIEW = "unhydratedEdgeDF" + uniqueTestViewSuffix + unhydratedEdgeDF.createOrReplaceTempView(unhydratedEdgeVIEW) + + val numTrainingSamples = + unhydratedEdgeDF.distinct().count().toInt // number of all nodes in graph + val numPositiveSamples = 2 + val sampledPosVIEW = nablpTask.sampleDstNodesUniformly( + numDstSamples = numPositiveSamples, + unhydratedEdgeVIEW = unhydratedEdgeVIEW, + permutationStrategy = "non-deterministic", + numTrainingSamples = numTrainingSamples, + edgeUsageType = EdgeUsageType.POS, + ) + val sampledPosDF = sparkTest.table(sampledPosVIEW) + // must choose a nodeId st number of its out-edges is > numPositiveSamples [to test randomness] + val nodeId = 0 + val posNodeList = Seq(1, 2, 3, 4, 5, 6, 7, 8) + val randomPosSamplesList = sampledPosDF + .filter(F.col("_src_node") === nodeId) + .select("_pos_dst_node") + .collect + .map(_(0)) + .toList + posNodeList should contain allElementsOf randomPosSamplesList + + } + + test("Positive node neighbors are valid with correct columns.") { + // Making scope local to avoid clash with scalapb.spark.Implicits._ + import sqlImplicits._ + val uniqueTestViewSuffix = "_" + randomUUID.toString.replace("-", "_") + val mockSubgraphDF = mockSubgraphForCurrentTest + val mockSubgraphVIEW = "mockSubgraphDF" + uniqueTestViewSuffix + mockSubgraphDF.createOrReplaceTempView(mockSubgraphVIEW) + val sampledPosData = Seq((0, 1), (0, 3), (0, 2), (1, 2), (4, 0)) + val sampledPosDF = sampledPosData.toDF("_src_node", "_pos_dst_node") + val sampledPosVIEW = "sampledPosDF" + uniqueTestViewSuffix + sampledPosDF.createOrReplaceTempView(sampledPosVIEW) + val posNeighborhoodVIEW = nablpTask.lookupDstNodeNeighborhood( + sampledDstNodesVIEW = sampledPosVIEW, + subgraphVIEW = mockSubgraphVIEW, + edgeUsageType = EdgeUsageType.POS, + ) + val posNeighborhoodDF = sparkTest.table(posNeighborhoodVIEW) + // sparkTest.table(posNeighborhoodVIEW).show() + // ensures the direction is preserved i.e. from _src_node to _pos_node not the other way + val expectedSrcNodeList = Seq(1, 4, 0) + val currentSrcNodeList = posNeighborhoodDF.select("_src_node").collect.map(_(0)).toList + expectedSrcNodeList should contain allElementsOf currentSrcNodeList + // ensures neighbors are valid + // for src node 1, pos dst node is 2 + val posNodeDstId = 2 + val srcNodeId = 1 + val expectedNeighborEdges = mockSubgraphDF + .filter(F.col("_root_node") === posNodeDstId) + .select("_neighbor_edges") + .first + .toSeq + val currentNeighborEdges = posNeighborhoodDF + .filter(F.col("_src_node") === srcNodeId) + .select("_pos_neighbor_edges") + .first + .toSeq + expectedNeighborEdges should contain allElementsOf currentNeighborEdges + } + + test("validation fails if nodes of supervision edges not present in neighborhood nodes") { + val rootNode = Node(nodeId = 0) + + // Create the nodes for neighborhood + val neighborhoodNodes = Seq( + Node(nodeId = 0), + Node(nodeId = 1), + Node(nodeId = 2), + ) + + // Create the edges for neighborhood + val neighborhoodEdges = Seq( + Edge(srcNodeId = 0, dstNodeId = 1), + Edge(srcNodeId = 1, dstNodeId = 2), + ) + + // Create the pos_edges (with one node of the edge not in neighborhood nodes) + val posEdges = Seq( + Edge(srcNodeId = 0, dstNodeId = 3), + ) + + // Create the NodeAnchorBasedLinkPredictionSample dataset + val sample = Seq( + NodeAnchorBasedLinkPredictionSample( + rootNode = Some(rootNode), + neighborhood = Some(Graph(nodes = neighborhoodNodes, edges = neighborhoodEdges)), + posEdges = posEdges, + ), + ) + + val sampleDS = sparkTest.createDataset(sample) + + // validateMainSamples method is not an action, performing collect for triggering computation + assertThrows[SparkException]( + TaskOutputValidator.validateMainSamples(sampleDS, graphMetadataPbWrapper).collect(), + ) + } + + // TODO + // ensure caching sanity, by checking pos samples and hydrated nei, in integration test + + // TODO test create node anchor sample subgraph + // check resulting subgraph pos is in nodes, hard neg is in nodes +} diff --git a/scala_spark35/subgraph_sampler/src/test/scala/SGSPureSparkV1TaskTest.scala b/scala_spark35/subgraph_sampler/src/test/scala/SGSPureSparkV1TaskTest.scala new file mode 100644 index 000000000..9fe4808f5 --- /dev/null +++ b/scala_spark35/subgraph_sampler/src/test/scala/SGSPureSparkV1TaskTest.scala @@ -0,0 +1,608 @@ +import common.test.testLibs.SharedSparkSession +import common.types.EdgeUsageType +import common.types.pb_wrappers.GbmlConfigPbWrapper +import common.utils.ProtoLoader.populateProtoFromYaml +import libs.task.pureSpark.SGSPureSparkV1Task +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row +import org.apache.spark.sql.types._ +import org.apache.spark.sql.{functions => F} +import org.scalatest.BeforeAndAfterAll +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers._ +import snapchat.research.gbml.gbml_config.GbmlConfig + +import java.util.UUID.randomUUID + +class SGSPureSparkV1TaskTest extends AnyFunSuite with BeforeAndAfterAll with SharedSparkSession { + + import sqlImplicits._ + // commmon/reused vars among tests which must be assigned in beforeAll(): + var sgsTask: SGSPureSparkV1Task = _ + var gbmlConfigWrapper: GbmlConfigPbWrapper = _ + + override def beforeAll(): Unit = { + super.beforeAll() + val frozenGbmlConfigUriTest = + "common/src/test/assets/subgraph_sampler/node_anchor_based_link_prediction/frozen_gbml_config.yaml" + val gbmlConfigProto = + populateProtoFromYaml[GbmlConfig](uri = frozenGbmlConfigUriTest) + gbmlConfigWrapper = GbmlConfigPbWrapper(gbmlConfigPb = gbmlConfigProto) + class MockSGSPureSparkV1Task( + gbmlConfigWrapper: GbmlConfigPbWrapper) + extends SGSPureSparkV1Task(gbmlConfigWrapper) { + def applyCachingToSubgraphDf( + dfVIEW: String, + withRepartition: Boolean, + ): String = ??? + def run() = ??? + } + sgsTask = new MockSGSPureSparkV1Task(gbmlConfigWrapper = gbmlConfigWrapper) + + } + + def mockUnhydratedEdgeForCurrentTest: DataFrame = { + val edgeData = Seq( + (0, 1), + (0, 2), + (0, 3), + (0, 4), + (0, 5), + (0, 6), + (0, 7), + (0, 8), + (1, 2), + (1, 3), + (1, 0), + (2, 0), + (3, 0), + (4, 0), + (5, 0), + (6, 0), + (7, 0), + (8, 0), + (2, 1), + (3, 1), + ) + + val unhydratedEdgeDF = edgeData.toDF("_src_node", "_dst_node") + unhydratedEdgeDF + } + + def mockHydratedEdgeForCurrentTest: DataFrame = { + // must have same edges as edgeData in mockUnhydratedEdgeForCurrentTest + val hydratedEdgeData = Seq( + (0, 1, Seq(0.5), 0), + (0, 2, Seq(1.0), 0), + (0, 3, Seq(1.5), 0), + (0, 4, Seq(2.0), 0), + (0, 5, Seq(2.5), 0), + (0, 6, Seq(3.0), 0), + (0, 7, Seq(3.5), 0), + (0, 8, Seq(4.0), 0), + (1, 2, Seq(1.5), 0), + (1, 3, Seq(2.0), 0), + (1, 0, Seq(0.5), 0), + (2, 0, Seq(1.0), 0), + (3, 0, Seq(1.5), 0), + (4, 0, Seq(2.0), 0), + (5, 0, Seq(2.5), 0), + (6, 0, Seq(3.0), 0), + (7, 0, Seq(3.5), 0), + (8, 0, Seq(4.0), 0), + (2, 1, Seq(1.5), 0), + (3, 1, Seq(2.0), 0), + ) + + val hydratedEdgeDF = hydratedEdgeData + .toDF("_from", "_to", "_edge_features", "_condensed_edge_type") + hydratedEdgeDF + } + + def mockHydratedNodeForCurrentTest: DataFrame = { + // must include isolated node (ie 9,10,11) + val nodeData = Seq( + (0, 0.0, 0), + (1, 0.1, 0), + (2, 0.2, 0), + (3, 0.3, 0), + (4, 0.4, 0), + (5, 0.5, 0), + (6, 0.6, 0), + (7, 0.7, 0), + (8, 0.8, 0), + (9, 0.9, 0), + (10, 0.01, 0), + (11, 0.11, 0), + ) + val hydratedNodeDF = + nodeData.toDF("_node_id", "_node_features", "_condensed_node_type") + hydratedNodeDF + } + + test("Node Dataframe with features and node type loads with correct cloumns.") { + val hydratedNodeVIEW = sgsTask.loadNodeDataframeIntoSparkSql(condensedNodeType = 0) + val hydratedNodeDF = sparkTest.table(hydratedNodeVIEW) + val expectedColSize = 3 + assert(hydratedNodeDF.columns.size == expectedColSize) + val loadedNodeDfColumns: Seq[String] = + hydratedNodeDF.columns.toSeq + assert(loadedNodeDfColumns.contains("_node_id")) + assert(loadedNodeDfColumns.contains("_node_features")) + assert(loadedNodeDfColumns.contains("_condensed_node_type")) + // need to check dtypes too? how to find field's dtype from proto easily? + // print(graph_schema.Node.scalaDescriptor.fields(0)) + } + + def is_edge_dataframe_loaded_with_correct_columns( + edgeUsageType: EdgeUsageType.EdgeUsageType, + check_self_loops: Boolean = false, + ): Unit = { + val hydratedEdgeVIEW = + sgsTask.loadEdgeDataframeIntoSparkSql(condensedEdgeType = 0, edgeUsageType = edgeUsageType) + val hydratedEdgeDF = sparkTest.table(hydratedEdgeVIEW) + val expectedColSize = 4 + assert(hydratedEdgeDF.columns.size == expectedColSize) + val hydratedEdgeDfCols: Seq[String] = hydratedEdgeDF.columns.toSeq + assert(hydratedEdgeDfCols.contains("_from")) + assert(hydratedEdgeDfCols.contains("_to")) + assert(hydratedEdgeDfCols.contains("_condensed_edge_type")) + assert(hydratedEdgeDfCols.contains("_edge_features")) + + val unhydratedEdgeVIEW = + sgsTask.loadUnhydratedEdgeDataframeIntoSparkSql(hydratedEdgeVIEW = hydratedEdgeVIEW) + val unhydratedEdgeDF = sparkTest.table(unhydratedEdgeVIEW) + val unhydratedDfExpectedColSize = 2 + assert(unhydratedEdgeDF.columns.size == unhydratedDfExpectedColSize) + val unhydratedEdgeDfCols = unhydratedEdgeDF.columns.toSeq + assert(unhydratedEdgeDfCols.contains("_src_node")) + assert(unhydratedEdgeDfCols.contains("_dst_node")) + + assert(gbmlConfigWrapper.sharedConfigPb.isGraphDirected == false) + + if (check_self_loops) { + var srcIds: List[Any] = hydratedEdgeDF.sort("_from").select("_from").collect.map(_(0)).toList + var dstIds: List[Any] = hydratedEdgeDF.sort("_to").select("_to").collect.map(_(0)).toList + srcIds shouldBe dstIds + + srcIds = unhydratedEdgeDF.sort("_src_node").select("_src_node").collect.map(_(0)).toList + dstIds = unhydratedEdgeDF.sort("_dst_node").select("_dst_node").collect.map(_(0)).toList + srcIds shouldBe dstIds + // graph should not include any self loops + val selfLoopIds = hydratedEdgeDF.filter(F.col("_from") === F.col("_to")) + assert(selfLoopIds.count() == 0) + } + + } + + test("Edge Dataframes are loaded with correct columns and bidirectionalized if required.") { + is_edge_dataframe_loaded_with_correct_columns( + edgeUsageType = EdgeUsageType.MAIN, + check_self_loops = true, + ) + is_edge_dataframe_loaded_with_correct_columns(edgeUsageType = EdgeUsageType.POS) + is_edge_dataframe_loaded_with_correct_columns(edgeUsageType = EdgeUsageType.NEG) + // TODO perhaps add asserts on the content itself + // TODO test if 3 edge dataframes are loaded with correct directions in scenarios for directed vs undirected graphs. + // val hydratedEdgeVIEW = sgsTask.loadEdgeDataframeIntoSparkSql(condensedEdgeType = 0) + } + + test("Onehop samples are valid.") { + val uniqueTestViewSuffix = "_" + randomUUID.toString.replace("-", "_") + val unhydratedEdgeDF = mockUnhydratedEdgeForCurrentTest + val unhydratedEdgeVIEW = "unhydratedEdgeDF" + uniqueTestViewSuffix + unhydratedEdgeDF.createOrReplaceTempView(unhydratedEdgeVIEW) + var numNeighborsToSample = 3 + val sampledOnehopVIEW = + sgsTask.sampleOnehopSrcNodesUniformly( + numNeighborsToSample = numNeighborsToSample, + unhydratedEdgeVIEW = unhydratedEdgeVIEW, + permutationStrategy = "non-deterministic", + ) + val sampledOnehopDF = sparkTest.table(sampledOnehopVIEW) + // must choose a nodeId st number of its in-edges is > numNeighborsToSample [to test randomness] + var nodeId = 0 + val nodeList = Seq(1, 2, 3, 4, 5, 6, 7, 8) + val randomSamplesList = sampledOnehopDF + .filter(F.col("_0_hop") === nodeId) + .select("_sampled_1_hop_arr") + .first + .getSeq[Integer](0) + nodeList should contain allElementsOf randomSamplesList + assert(randomSamplesList.length == numNeighborsToSample) + } + + test("Onehop samples with replacement are valid.") { + val uniqueTestViewSuffix = "_" + randomUUID.toString.replace("-", "_") + val unhydratedEdgeDF = mockUnhydratedEdgeForCurrentTest + val unhydratedEdgeVIEW = "unhydratedEdgeDF" + uniqueTestViewSuffix + unhydratedEdgeDF.createOrReplaceTempView(unhydratedEdgeVIEW) + var numNeighborsToSample = 10 + val sampledOnehopVIEW = + sgsTask.sampleOnehopSrcNodesUniformly( + numNeighborsToSample = numNeighborsToSample, + unhydratedEdgeVIEW = unhydratedEdgeVIEW, + permutationStrategy = "non-deterministic", + sampleWithReplacement = true, + ) + val sampledOnehopDF = sparkTest.table(sampledOnehopVIEW) + // must choose a nodeId st number of its in-edges is > numNeighborsToSample [to test randomness] + var nodeId = 0 + val nodeList = Seq(1, 2, 3, 4, 5, 6, 7, 8) + val randomSamplesList = sampledOnehopDF + .filter(F.col("_0_hop") === nodeId) + .select("_sampled_1_hop_arr") + .first + .getSeq[Integer](0) + nodeList should contain allElementsOf randomSamplesList + assert(randomSamplesList.length == numNeighborsToSample) + } + + test("Twohop samples are valid.") { + val uniqueTestViewSuffix = "_" + randomUUID.toString.replace("-", "_") + val unhydratedEdgeDF = mockUnhydratedEdgeForCurrentTest + val unhydratedEdgeVIEW = "unhydratedEdgeDF" + "_" + uniqueTestViewSuffix + unhydratedEdgeDF.createOrReplaceTempView(unhydratedEdgeVIEW) + val onehopData = Seq((1, Seq(3, 0, 2)), (3, Seq(0, 1)), (2, Seq(0, 1)), (0, Seq(1, 3, 2))) + val sampledOnehopDF = onehopData.toDF("_0_hop", "_sampled_1_hop_arr") + val sampledOnehopVIEW = "sampledOnehopDF" + uniqueTestViewSuffix + sampledOnehopDF.createOrReplaceTempView(sampledOnehopVIEW) + var numNeighborsToSample = 3 + + val sampledTwohopVIEW = sgsTask.sampleTwohopSrcNodesUniformly( + numNeighborsToSample = numNeighborsToSample, + unhydratedEdgeVIEW = unhydratedEdgeVIEW, + sampledOnehopVIEW = sampledOnehopVIEW, + permutationStrategy = "non-deterministic", + ) + val sampledTwohopDF = sparkTest.table(sampledTwohopVIEW) + // must choose a nodeId st number of its in-edges is <= numNeighborsToSample [no randomness] + var zerohopId = 1 + // must choose a nodeId st number of its in-edges is > numNeighborsToSample [to test randomness] + var onehopId = 0 + var twohopList = Seq(1, 2, 3, 4, 5, 6, 7, 8) + val randomTwohopSamplesList = sampledTwohopDF + .filter(F.col("_0_hop") === zerohopId && F.col("_1_hop") === onehopId) + .select("_sampled_2_hop_arr") + .first + .getSeq[Integer](0) + twohopList should contain allElementsOf randomTwohopSamplesList + assert(randomTwohopSamplesList.length == numNeighborsToSample) + } + + test("Twohop samples with replacement are valid.") { + val uniqueTestViewSuffix = "_" + randomUUID.toString.replace("-", "_") + val unhydratedEdgeDF = mockUnhydratedEdgeForCurrentTest + val unhydratedEdgeVIEW = "unhydratedEdgeDF" + uniqueTestViewSuffix + unhydratedEdgeDF.createOrReplaceTempView(unhydratedEdgeVIEW) + val onehopData = Seq((1, Seq(3, 0, 2)), (3, Seq(0, 1, 1)), (2, Seq(0, 1, 0)), (0, Seq(1, 3, 2))) + val sampledOnehopDF = onehopData.toDF("_0_hop", "_sampled_1_hop_arr") + val sampledOnehopVIEW = "sampledOnehopDF" + uniqueTestViewSuffix + sampledOnehopDF.createOrReplaceTempView(sampledOnehopVIEW) + var numNeighborsToSample = 10 + + val sampledTwohopVIEW = sgsTask.sampleTwohopSrcNodesUniformly( + numNeighborsToSample = numNeighborsToSample, + unhydratedEdgeVIEW = unhydratedEdgeVIEW, + sampledOnehopVIEW = sampledOnehopVIEW, + permutationStrategy = "non-deterministic", + sampleWithReplacement = true, + ) + val sampledTwohopDF = sparkTest.table(sampledTwohopVIEW) + // must choose a nodeId st number of its in-edges is <= numNeighborsToSample [no randomness] + var zerohopId = 1 + // must choose a nodeId st number of its in-edges is > numNeighborsToSample [to test randomness] + var onehopId = 0 + var twohopList = Seq(1, 2, 3, 4, 5, 6, 7, 8) + val randomTwohopSamplesList = sampledTwohopDF + .filter(F.col("_0_hop") === zerohopId && F.col("_1_hop") === onehopId) + .select("_sampled_2_hop_arr") + .first + .getSeq[Integer](0) + twohopList should contain allElementsOf randomTwohopSamplesList + assert(randomTwohopSamplesList.length == numNeighborsToSample) + } + + test("Hydrated kth hop nodes have right col and node ids") { + val uniqueTestViewSuffix = "_" + randomUUID.toString.replace("-", "_") + val edgeData = Seq( + (0, 1), + (0, 2), + (0, 3), + (1, 2), + (1, 3), + (1, 0), + (2, 0), + (3, 0), + (2, 1), + (3, 1), + ) + val unhydratedEdgeDF = edgeData.toDF("_src_node", "_dst_node") + val unhydratedEdgeVIEW = "unhydratedEdgeDF" + uniqueTestViewSuffix + unhydratedEdgeDF.createOrReplaceTempView(unhydratedEdgeVIEW) + + val nodeData = Seq( + (0, 0.0, 0), + (1, 0.1, 0), + (2, 0.2, 0), + (3, 0.3, 0), + ) + val hydratedNodeDF = + nodeData.toDF("_node_id", "_node_features", "_condensed_node_type") + val hydratedNodeVIEW = "hydratedNodeDF" + uniqueTestViewSuffix + hydratedNodeDF.createOrReplaceTempView(hydratedNodeVIEW) + + val onehopData = Seq((1, Seq(3, 0, 2)), (3, Seq(0, 1)), (2, Seq(0, 1)), (0, Seq(1, 3, 2))) + val sampledOnehopDF = onehopData.toDF("_0_hop", "_sampled_1_hop_arr") + val sampledOnehopVIEW = "sampledOnehopDF" + uniqueTestViewSuffix + sampledOnehopDF.createOrReplaceTempView(sampledOnehopVIEW) + + val twohopData = Seq( + (1, 0, Seq(1, 2, 3)), + (3, 1, Seq(0, 3, 2)), + (1, 2, Seq(0, 1)), + (1, 3, Seq(0, 1)), + (2, 1, Seq(2, 0, 3)), + (0, 1, Seq(2, 3, 0)), + (2, 0, Seq(3, 2, 1)), + (0, 2, Seq(0, 1)), + (0, 3, Seq(1, 0)), + (3, 0, Seq(2, 3, 1)), + ) + val sampledTwohopDF = + twohopData.toDF("_0_hop", "_1_hop", "_sampled_2_hop_arr") + val sampledTwohopVIEW = "sampledTwohopDF" + uniqueTestViewSuffix + sampledTwohopDF.createOrReplaceTempView(sampledTwohopVIEW) + + // must be <= max num out edges of unhydratedEdgeDF [we don't want to test randomness here] + var numNeighborsToSample = 3 + // check hydrated onehop sanity + val hydratedOnehopNodesVIEW = sgsTask.hydrateNodes( + k = 1, + hydratedNodeVIEW = hydratedNodeVIEW, + sampledKhopVIEW = sampledOnehopVIEW, + ) + val hydratedOnehopNodesDF = sparkTest.table(hydratedOnehopNodesVIEW) + val expectedOnehopColNames: Seq[String] = + Seq("_node_id", "_node_features", "_condensed_node_type", "_0_hop", "_1_hop") + val loadedOnehopColNames: Seq[String] = hydratedOnehopNodesDF.columns.toSeq + expectedOnehopColNames should contain theSameElementsAs loadedOnehopColNames + + var diffDF = hydratedOnehopNodesDF.filter(!(F.col("_node_id").contains(F.col("_1_hop")))) + assert(diffDF.count() == 0) + + // check hydrated twohop sanity + val hydratedTwohopNodesVIEW = sgsTask.hydrateNodes( + k = 2, + hydratedNodeVIEW = hydratedNodeVIEW, + sampledKhopVIEW = sampledTwohopVIEW, + ) + val hydratedTwohopNodesDF = sparkTest.table(hydratedTwohopNodesVIEW) + val expectedTwohopColNames: Seq[String] = + Seq("_node_id", "_node_features", "_condensed_node_type", "_0_hop", "_1_hop", "_2_hop") + val loadedTwohopColNames: Seq[String] = hydratedTwohopNodesDF.columns.toSeq + expectedTwohopColNames should contain theSameElementsAs loadedTwohopColNames + + diffDF = hydratedTwohopNodesDF.filter(!(F.col("_node_id").contains(F.col("_2_hop")))) + assert(diffDF.count() == 0) + + } + + test("Rooted Node Neighborhood is valid.") { + // delete all prev view names from current sparkTest session + sparkTest.sqlContext.tableNames().foreach(sparkTest.catalog.dropTempView(_)) + val uniqueTestViewSuffix = "_" + randomUUID.toString.replace("-", "_") + + val unhydratedEdgeDF = mockUnhydratedEdgeForCurrentTest + val unhydratedEdgeVIEW = "unhydratedEdgeDF" + uniqueTestViewSuffix + unhydratedEdgeDF.createOrReplaceTempView(unhydratedEdgeVIEW) + + val hydratedEdgeDF = mockHydratedEdgeForCurrentTest + val hydratedEdgeVIEW = "hydratedEdgeDF" + uniqueTestViewSuffix + hydratedEdgeDF.createOrReplaceTempView(hydratedEdgeVIEW) + + val hydratedNodeDF = mockHydratedNodeForCurrentTest + val hydratedNodeVIEW = "hydratedNodeDF" + uniqueTestViewSuffix + hydratedNodeDF.createOrReplaceTempView(hydratedNodeVIEW) + + var numNeighborsToSample = 3 + val rnnVIEW = sgsTask.createSubgraph( + numNeighborsToSample = numNeighborsToSample, + hydratedNodeVIEW = hydratedNodeVIEW, + hydratedEdgeVIEW = hydratedEdgeVIEW, + unhydratedEdgeVIEW = unhydratedEdgeVIEW, + permutationStrategy = "non-detrministic", + ) + val rnnDF = sparkTest.table(rnnVIEW) + // sparkTest.table(rnnVIEW).show(50, truncate = false) + // sparkTest.table(rnnVIEW).printSchema() + + // view names parametrized by hop number are valid + val parametrizedViews = Array( + "1hopdf", + "2hopdf", + "hydrated1hopnodesdf", + "hydrated2hopnodesdf", + "hydrated1thhopdf", + "hydrated2thhopdf", + "hydrated1hopnodesedgesdf", + "hydrated2hopnodesedgesdf", + ) + val expectedParametrizedViews: Array[String] = + parametrizedViews.map(_ + sgsTask.uniqueTempViewSuffix) + val curViews = sparkTest.sqlContext.tableNames() + curViews should contain allElementsOf expectedParametrizedViews + // edges are valid, i.e neighborhood is a valid subgraph. [ensures caching sampledDF is done correctly] + // 1. MUST choose a root node with in edges > NumNeighborsToSample + val rootNodeId = 0 + // 2. for rootNodeId, take 1hop src_node ids + val onehopSrcIds: Seq[Int] = rnnDF + .filter(F.col("_root_node") === rootNodeId) + .select( + F.filter(F.col("_neighbor_edges"), c => c.apply("_src_node") !== rootNodeId) + .getItem("_src_node"), + ) + .first + .getSeq[Int](0) + .take(numNeighborsToSample) + // 3. for each of onehopSrcIds verify that twohopDstId exists + val firstSampledSrcId = onehopSrcIds.apply(0) + var twohopDst: DataFrame = rnnDF + .filter(F.col("_root_node") === rootNodeId) + .select(F.filter(F.col("_neighbor_edges"), c => c.apply("_dst_node") === firstSampledSrcId)) + assert(twohopDst.count() == 1) + + val secondSampledSrcId = onehopSrcIds.apply(1) + twohopDst = rnnDF + .filter(F.col("_root_node") === rootNodeId) + .select(F.filter(F.col("_neighbor_edges"), c => c.apply("_dst_node") === secondSampledSrcId)) + assert(twohopDst.count() == 1) + + val thirdSampledSrcId = onehopSrcIds.apply(2) + twohopDst = rnnDF + .filter(F.col("_root_node") === rootNodeId) + .select(F.filter(F.col("_neighbor_edges"), c => c.apply("_dst_node") === thirdSampledSrcId)) + assert(twohopDst.count() == 1) + + // verify node ids in _neighbor_nodes match with node ids in _neighbor_edges + val srcIdList = rnnDF + .filter(F.col("_root_node") === rootNodeId) + .select(F.col("_neighbor_edges").getItem("_src_node")) + .first + .getSeq[Integer](0) + val dstIdList = rnnDF + .filter(F.col("_root_node") === rootNodeId) + .select(F.col("_neighbor_edges").getItem("_dst_node")) + .first + .getSeq[Integer](0) + var allIdsFromEdges: Seq[Integer] = srcIdList ++ dstIdList + allIdsFromEdges = allIdsFromEdges.distinct + val allIdsFromNodes: Seq[Integer] = rnnDF + .filter(F.col("_root_node") === rootNodeId) + .select(F.col("_neighbor_nodes").getItem("_node_id")) + .first + .getSeq[Integer](0) + allIdsFromEdges.sorted shouldBe allIdsFromNodes.sorted + + // root nodes are included in neighborhood + val rootNodeIdDF = rnnDF + .filter(F.col("_root_node") === rootNodeId) + .select( + F.filter(F.col("_neighbor_nodes"), c => c.apply("_node_id") === rootNodeId) + .alias("_curr_root_node"), + ) + val hydratedRootNodeId: Integer = rootNodeIdDF + .select(F.col("_curr_root_node").getItem("_node_id")) + .first + .getSeq[Integer](0) + .apply(0) + val hydratedRootNodeFeature: Double = rootNodeIdDF + .select(F.col("_curr_root_node").getItem("_feature_values")) + .first + .getSeq[Double](0) + .apply(0) + assert(hydratedRootNodeId == rootNodeId) + assert(hydratedRootNodeFeature == 0.0) + + } + // TODO: add unittest here for directed graphs to check if neighborless nodes are all included + test("Isolated nodes are included in inference/rooted node neighbor subgraph DF.") { + val uniqueTestViewSuffix = "_" + randomUUID.toString.replace("-", "_") + val edgeData = Seq((0, 1), (0, 2), (0, 3), (1, 3), (1, 0), (2, 0), (3, 0), (3, 1)) + val edgeRDD = sparkTest.sparkContext.parallelize(edgeData) + val unhydratedEdgeDF = sparkTest.createDataFrame(edgeRDD).toDF("_src_node", "_dst_node") + val unhydratedEdgeVIEW = "unhydratedEdgeDF" + uniqueTestViewSuffix + unhydratedEdgeDF.createOrReplaceTempView(unhydratedEdgeVIEW) + val nodeData = Seq( + (0, Seq(0.0f), 0), + (1, Seq(0.1f), 0), + (2, Seq(0.2f), 0), + (3, Seq(0.3f), 0), + (4, Seq(0.4f), 0), + (5, Seq(0.5f), 0), + ) + val nodeRDD = sparkTest.sparkContext.parallelize(nodeData) + val hydratedNodeDF = + sparkTest.createDataFrame(nodeRDD).toDF("_node_id", "_node_features", "_condensed_node_type") + val hydratedNodeVIEW = "hydratedNodeDF" + uniqueTestViewSuffix + hydratedNodeDF.createOrReplaceTempView(hydratedNodeVIEW) + + // note that subgraph data does not form an actual graph, only root_id and dtypes matter for this test + val emptyFeats = Seq.empty[Float] + val subgraphData = Seq( + Row(0, List(Row(1, 0, 0, emptyFeats)), List(Row(1, 0, Seq(0.1f))), Seq(0.0f), 0), + Row(1, List(Row(1, 0, 0, emptyFeats)), List(Row(1, 0, Seq(0.1f))), Seq(0.0f), 0), + Row(2, List(Row(1, 0, 0, emptyFeats)), List(Row(1, 0, Seq(0.1f))), Seq(0.0f), 0), + Row(3, List(Row(1, 0, 0, emptyFeats)), List(Row(1, 0, Seq(0.1f))), Seq(0.0f), 0), + ) + val arrayStruct = new StructType() + .add("_root_node", IntegerType) + .add( + "_neighbor_edges", + ArrayType( + new StructType() + .add("_src_node", IntegerType) + .add("_dst_node", IntegerType) + .add("_condensed_edge_type", IntegerType) + .add("_feature_values", ArrayType(FloatType)), + ), + ) + .add( + "_neighbor_nodes", + ArrayType( + new StructType() + .add("_node_id", IntegerType) + .add("_condensed_node_type", IntegerType) + .add("_feature_values", ArrayType(FloatType)), + ), + ) + .add("_node_features", ArrayType(FloatType)) + .add("_condensed_node_type", IntegerType) + val subgraphRDD = sparkTest.sparkContext.parallelize(subgraphData) + val mockSubgraphDF = sparkTest.createDataFrame(subgraphRDD, arrayStruct) + val mockSubgraphVIEW = "mockSubgraphDF" + uniqueTestViewSuffix + mockSubgraphDF.createOrReplaceTempView(mockSubgraphVIEW) + val subgraphWithIsolatedNodesVIEW = sgsTask.createRootedNodeNeighborhoodSubgraph( + hydratedNodeVIEW = hydratedNodeVIEW, + unhydratedEdgeVIEW = unhydratedEdgeVIEW, + subgraphVIEW = mockSubgraphVIEW, + ) + val subgraphWithIsolatedNodesDF = sparkTest.table(subgraphWithIsolatedNodesVIEW) + // subgraphWithIsolatedNodesDF.show(truncate = false) + val expectedIsolatedNodeList = Seq(4, 5) + val curIsolatedNodesList = subgraphWithIsolatedNodesDF + .filter(F.col("_neighbor_edges").isNull) + .select("_root_node") + .collect + .map(_(0)) + .toList + expectedIsolatedNodeList should contain allElementsOf curIsolatedNodesList + } + + test("sampleWithReplacementUDF returns correct number of samples") { + // Use the already registered UDF from SGSPureSparkV1Task + val sampleWithReplacementUDF = sgsTask.sampleWithReplacementUDF + + // Create a DataFrame with sample data + val data = Seq( + (Seq(1, 2, 3, 4, 5), 10), + (Seq(6, 7, 8, 9, 10), 2), + (Seq.empty[Int], 3), + (null, 3), + ).toDF("array", "numSamples") + + // Apply the UDF + val resultDF = + data.withColumn("samples", sampleWithReplacementUDF(F.col("array"), F.col("numSamples"))) + + // Collect the results + val results = resultDF.collect() + + // Write assertions + assert(results(0).getAs[Seq[Int]]("samples").length == 10) + assert(results(1).getAs[Seq[Int]]("samples").length == 2) + assert(results(2).getAs[Seq[Int]]("samples").isEmpty) + assert(results(3).getAs[Seq[Int]]("samples").isEmpty) + } + +} diff --git a/scala_spark35/subgraph_sampler/src/test/scala/SupervisedNodeClassificationTaskTest.scala b/scala_spark35/subgraph_sampler/src/test/scala/SupervisedNodeClassificationTaskTest.scala new file mode 100644 index 000000000..e4992f83b --- /dev/null +++ b/scala_spark35/subgraph_sampler/src/test/scala/SupervisedNodeClassificationTaskTest.scala @@ -0,0 +1,48 @@ +import common.test.testLibs.SharedSparkSession +import common.types.pb_wrappers.GbmlConfigPbWrapper +import common.types.pb_wrappers.GraphMetadataPbWrapper +import common.utils.ProtoLoader.populateProtoFromYaml +import libs.task.pureSpark.SupervisedNodeClassificationTask +import org.scalatest.BeforeAndAfterAll +import org.scalatest.funsuite.AnyFunSuite +import snapchat.research.gbml.gbml_config.GbmlConfig + +class SupervisedNodeClassificationTaskTest + extends AnyFunSuite + with BeforeAndAfterAll + with SharedSparkSession { + + // commmon/reused vars among tests which must be assigned in beforeAll(): + var sncTask: SupervisedNodeClassificationTask = _ + var gbmlConfigWrapper: GbmlConfigPbWrapper = _ + + override def beforeAll(): Unit = { + super.beforeAll() + val frozenGbmlConfigUriTest = + "common/src/test/assets/subgraph_sampler/supervised_node_classification/frozen_gbml_config.yaml" + val gbmlConfigProto = + populateProtoFromYaml[GbmlConfig](uri = frozenGbmlConfigUriTest) + gbmlConfigWrapper = GbmlConfigPbWrapper(gbmlConfigPb = gbmlConfigProto) + + val graphMetadataPbWrapper: GraphMetadataPbWrapper = GraphMetadataPbWrapper( + gbmlConfigWrapper.graphMetadataPb, + ) + + sncTask = new SupervisedNodeClassificationTask( + gbmlConfigWrapper = gbmlConfigWrapper, + ) + + } + + test("Node label dataframe loads with correct columns.") { + val nodeLabelVIEW = sncTask.loadNodeLabelDataframeIntoSparkSql(condensedNodeType = 0) + val nodeLabelDF = sparkTest.table(nodeLabelVIEW) + val expectedColSize = 3 + assert(nodeLabelDF.columns.size == expectedColSize) + val loadedNodeLabelDfColumns: Seq[String] = nodeLabelDF.columns.toSeq + assert(loadedNodeLabelDfColumns.contains("_node_id")) + assert(loadedNodeLabelDfColumns.contains("_label_key")) + assert(loadedNodeLabelDfColumns.contains("_label_type")) + + } +} diff --git a/scala_spark35/subgraph_sampler/src/test/scala/UserDefinedLabelsNodeAnchorBasedLinkPredictionTaskTest.scala b/scala_spark35/subgraph_sampler/src/test/scala/UserDefinedLabelsNodeAnchorBasedLinkPredictionTaskTest.scala new file mode 100644 index 000000000..87f0ed9ac --- /dev/null +++ b/scala_spark35/subgraph_sampler/src/test/scala/UserDefinedLabelsNodeAnchorBasedLinkPredictionTaskTest.scala @@ -0,0 +1,299 @@ +import common.test.testLibs.SharedSparkSession +import common.types.EdgeUsageType +import common.types.pb_wrappers.GbmlConfigPbWrapper +import common.types.pb_wrappers.GraphMetadataPbWrapper +import common.utils.ProtoLoader.populateProtoFromYaml +import libs.task.pureSpark.UserDefinedLabelsNodeAnchorBasedLinkPredictionTask +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row +import org.apache.spark.sql.types._ +import org.apache.spark.sql.{functions => F} +import org.scalatest.BeforeAndAfterAll +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.matchers.should.Matchers._ +import snapchat.research.gbml.gbml_config.GbmlConfig + +import java.util.UUID.randomUUID + +class UserDefinedLabelsNodeAnchorBasedLinkPredictionTaskTest + extends AnyFunSuite + with BeforeAndAfterAll + with SharedSparkSession { + + // commmon/reused vars among tests which must be assigned in beforeAll(): + var nablpTask: UserDefinedLabelsNodeAnchorBasedLinkPredictionTask = _ + var gbmlConfigWrapper: GbmlConfigPbWrapper = _ + var graphMetadataPbWrapper: GraphMetadataPbWrapper = _ + + override def beforeAll(): Unit = { + super.beforeAll() + val frozenGbmlConfigUriTest = + "common/src/test/assets/subgraph_sampler/node_anchor_based_link_prediction/frozen_gbml_config.yaml" + val gbmlConfigProto = + populateProtoFromYaml[GbmlConfig](uri = frozenGbmlConfigUriTest) + gbmlConfigWrapper = GbmlConfigPbWrapper(gbmlConfigPb = gbmlConfigProto) + + val graphMetadataPbWrapper: GraphMetadataPbWrapper = GraphMetadataPbWrapper( + gbmlConfigWrapper.graphMetadataPb, + ) + + nablpTask = new UserDefinedLabelsNodeAnchorBasedLinkPredictionTask( + gbmlConfigWrapper = gbmlConfigWrapper, + graphMetadataPbWrapper = graphMetadataPbWrapper, + isPosUserDefined = true, + isNegUserDefined = true, + ) + + } + + def mockNumberOfNodesForCurrentTest: Int = { + val numNodes: Int = 9 + numNodes + } + + def mockUnhydratedEdgeForCurrentTest: DataFrame = { + // Making scope local to avoid clash with scalapb.spark.Implicits._ + import sqlImplicits._ + // Note that below graph is bidirected + val edgeData = Seq( + (0, 1), + (0, 2), + (0, 3), + (0, 4), + (0, 5), + (0, 6), + (0, 7), + (0, 8), + (1, 2), + (1, 3), + (1, 0), + (2, 0), + (3, 0), + (4, 0), + (5, 0), + (6, 0), + (7, 0), + (8, 0), + (2, 1), + (3, 1), + ) + val unhydratedEdgeDF = edgeData.toDF("_src_node", "_dst_node") + unhydratedEdgeDF + } + + def mockUnhydratedUserDefinedPosEdgesForCurrentTest: DataFrame = { + // Making scope local to avoid clash with scalapb.spark.Implicits._ + import sqlImplicits._ + val userDefinedPosEdgeData = Seq( + (0, 1), + (0, 3), + (0, 2), + (1, 2), + (8, 0), // src only + (3, 1), + ) + val userDefinePosEdgeDF = userDefinedPosEdgeData.toDF("_src_node", "_dst_node") + userDefinePosEdgeDF + } + + def mockUnhydratedUserDefinedNegEdgesForCurrentTest: DataFrame = { + // Making scope local to avoid clash with scalapb.spark.Implicits._ + import sqlImplicits._ + val userDefinedNegEdgeData = Seq( + (2, 4), + (3, 2), + (6, 7), + (2, 7), + (8, 4), + (2, 6), + ) + val userDefineNegEdgeDF = userDefinedNegEdgeData.toDF("_src_node", "_dst_node") + userDefineNegEdgeDF + } + + def mockHydratedNodeForCurrentTest: DataFrame = { + // Making scope local to avoid clash with scalapb.spark.Implicits._ + import sqlImplicits._ + // must include isolated node (ie 9,10,11) + val nodeData = Seq( + (0, Seq(0.0f), 0), + (1, Seq(0.1f), 0), + (2, Seq(0.2f), 0), + (3, Seq(0.3f), 0), + (4, Seq(0.4f), 0), + (5, Seq(0.5f), 0), + (6, Seq(0.6f), 0), + (7, Seq(0.7f), 0), + (8, Seq(0.8f), 0), + (9, Seq(0.9f), 0), + (10, Seq(1.0f), 0), + (11, Seq(1.1f), 0), + ) + val hydratedNodeDF = + nodeData.toDF("_node_id", "_node_features", "_condensed_node_type") + hydratedNodeDF + } + + def mockSubgraphForCurrentTest: DataFrame = { + val emptyFeats = Seq.empty[Float] + val subgraphData = Seq( + Row(0, List(Row(1, 0, 0, emptyFeats)), List(Row(1, 0, Seq(0.1f))), Seq(0.0f), 0), + Row(1, List(Row(2, 1, 0, emptyFeats)), List(Row(2, 0, Seq(0.2f))), Seq(0.1f), 0), + Row(2, List(Row(1, 2, 0, emptyFeats)), List(Row(1, 0, Seq(0.1f))), Seq(0.2f), 0), + Row(3, List(Row(0, 3, 0, emptyFeats)), List(Row(0, 0, Seq(0.0f))), Seq(0.3f), 0), + Row(4, List(Row(0, 4, 0, emptyFeats)), List(Row(0, 0, Seq(0.0f))), Seq(0.4f), 0), + Row(5, List(Row(0, 5, 0, emptyFeats)), List(Row(0, 0, Seq(0.0f))), Seq(0.5f), 0), + Row(6, List(Row(0, 6, 0, emptyFeats)), List(Row(0, 0, Seq(0.0f))), Seq(0.6f), 0), + Row(7, List(Row(0, 7, 0, emptyFeats)), List(Row(0, 0, Seq(0.0f))), Seq(0.7f), 0), + Row(8, List(Row(8, 0, 0, emptyFeats)), List(Row(0, 0, Seq(0.0f))), Seq(0.8f), 0), + ) + val arrayStruct = new StructType() + .add("_root_node", IntegerType) + .add( + "_neighbor_edges", + ArrayType( + new StructType() + .add("_src_node", IntegerType) + .add("_dst_node", IntegerType) + .add("_condensed_edge_type", IntegerType) + .add("_feature_values", ArrayType(FloatType)), + ), + ) + .add( + "_neighbor_nodes", + ArrayType( + new StructType() + .add("_node_id", IntegerType) + .add("_condensed_node_type", IntegerType) + .add("_feature_values", ArrayType(FloatType)), + ), + ) + .add("_node_features", ArrayType(FloatType)) + .add("_condensed_node_type", IntegerType) + val subgraphRDD = sparkTest.sparkContext.parallelize(subgraphData) + val mockSubgraphDF = sparkTest.createDataFrame(subgraphRDD, arrayStruct) + + mockSubgraphDF + } + + def are_user_defined_samples_valid(edgeUsageType: EdgeUsageType.EdgeUsageType): Unit = { + val uniqueTestViewSuffix = "_" + randomUUID.toString.replace("-", "_") + val unhydratedEdgeDF = + if (edgeUsageType == EdgeUsageType.POS) mockUnhydratedUserDefinedPosEdgesForCurrentTest + else mockUnhydratedUserDefinedNegEdgesForCurrentTest + val unhydratedEdgeVIEW = s"unhydrated${edgeUsageType}EdgeDF" + uniqueTestViewSuffix + unhydratedEdgeDF.createOrReplaceTempView(unhydratedEdgeVIEW) + + val numTrainingSamples = mockNumberOfNodesForCurrentTest // number of all nodes in graph + val numSamples = 2 + val sampledDstVIEW = nablpTask.sampleDstNodesUniformly( + numDstSamples = numSamples, + unhydratedEdgeVIEW = unhydratedEdgeVIEW, + permutationStrategy = "non-deterministic", + numTrainingSamples = numTrainingSamples, + edgeUsageType = edgeUsageType, + ) + val sampledDstDF = sparkTest.table(sampledDstVIEW) + // must choose a nodeId st number of its out-edges is > numSamples [to test randomness] + val nodeId = if (edgeUsageType == EdgeUsageType.POS) 0 else 2 + val sampleNodeList = if (edgeUsageType == EdgeUsageType.POS) Seq(1, 2, 3) else Seq(4, 6, 7) + val randomSamplesList = sampledDstDF + .filter(F.col("_src_node") === nodeId) + .select(f"_${edgeUsageType}_dst_node") + .collect + .map(_(0)) + .toList + sampleNodeList should contain allElementsOf randomSamplesList + } + + // test correctness of sampleDstNodesUniformly function for UDL + test("User Defined Positive & Negative samples are valid") { + are_user_defined_samples_valid(EdgeUsageType.POS) + are_user_defined_samples_valid(EdgeUsageType.NEG) + } + + // test correctness of lookupDstNodeNeighborhood function for UDL + test("Negative neighbors are valid with correct columns") { + // Making scope local to avoid clash with scalapb.spark.Implicits._ + val uniqueTestViewSuffix = "_" + randomUUID.toString.replace("-", "_") + val mockSubgraphDF = mockSubgraphForCurrentTest + val mockSubgraphVIEW = "mockSubgraphDF" + uniqueTestViewSuffix + mockSubgraphDF.createOrReplaceTempView(mockSubgraphVIEW) + val negEdgesDF = mockUnhydratedUserDefinedNegEdgesForCurrentTest + val sampledNegDF = + negEdgesDF.select(F.col("_src_node"), F.col("_dst_node").alias("_neg_dst_node")) + val sampledNegVIEW = "sampledNegDF" + uniqueTestViewSuffix + sampledNegDF.createOrReplaceTempView(sampledNegVIEW) + val negNeighborhoodVIEW = nablpTask.lookupDstNodeNeighborhood( + sampledDstNodesVIEW = sampledNegVIEW, + subgraphVIEW = mockSubgraphVIEW, + edgeUsageType = EdgeUsageType.NEG, + ) + val negNeighborhoodDF = sparkTest.table(negNeighborhoodVIEW) + // ensures the direction is preserved i.e. from _src_node to _pos_node not the other way + val expectedSrcNodeList = Seq(2, 3, 6, 8) + val currentSrcNodeList = negNeighborhoodDF.select("_src_node").collect.map(_(0)).toList + expectedSrcNodeList should contain allElementsOf currentSrcNodeList + // ensures neighbors are valid + // for src node 1, pos dst node is 2 + val negNodeDstId = 7 + val srcNodeId = 6 + + val expectedNeighborEdges = mockSubgraphDF + .filter(F.col("_root_node") === negNodeDstId) + .select("_neighbor_edges") + .first + .toSeq + val currentNeighborEdges = negNeighborhoodDF + .filter(F.col("_src_node") === srcNodeId) + .select("_neg_neighbor_edges") + .first + .toSeq + expectedNeighborEdges should contain allElementsOf currentNeighborEdges + + } + + test("Src Only nodes added to reference subgraph") { + // Making scope local to avoid clash with scalapb.spark.Implicits._ + val uniqueTestViewSuffix = "_" + randomUUID.toString.replace("-", "_") + + // initialize assets + val unhydratedEdgeDF = mockUnhydratedEdgeForCurrentTest + val unhydratedEdgeVIEW = "unhydratedEdgeDF" + uniqueTestViewSuffix + unhydratedEdgeDF.createOrReplaceTempView(unhydratedEdgeVIEW) + + val unhydratedUserDefinedPosEdgesDF = mockUnhydratedUserDefinedPosEdgesForCurrentTest + val unhydratedUserDefinedPosEdgesVIEW = "unhydratedUserDefinedPosEdgesDF" + uniqueTestViewSuffix + unhydratedUserDefinedPosEdgesDF.createOrReplaceTempView(unhydratedUserDefinedPosEdgesVIEW) + + val mockSubgraphDF = mockSubgraphForCurrentTest + val mockSubgraphVIEW = "mockSubgraphDF" + uniqueTestViewSuffix + mockSubgraphDF.createOrReplaceTempView(mockSubgraphVIEW) + + val mockHydratedNodeDF = mockHydratedNodeForCurrentTest + val mockHydratedNodeVIEW = "mockHydratedNodeDF" + uniqueTestViewSuffix + mockHydratedNodeDF.createOrReplaceTempView(mockHydratedNodeVIEW) + + // run function + val subgraphWithUdlPosSrcOnlyView = nablpTask.addUserDefSrcOnlyNodesToRNNSubgraph( + unhydratedMainEdgeVIEW = unhydratedEdgeVIEW, + unhydratedUserDefEdgeVIEW = unhydratedUserDefinedPosEdgesVIEW, + referenceSubgraphVIEW = mockSubgraphVIEW, + hydratedNodeVIEW = mockHydratedNodeVIEW, + ) + + // check outputs + val userDefinedPosSrcNode = + unhydratedUserDefinedPosEdgesDF.select("_src_node").collect.map(_(0)).toList + val subgraphWithUdlPosSrcOnlyRootNodes = + sparkTest.table(subgraphWithUdlPosSrcOnlyView).select("_root_node").collect.map(_(0)).toList + + subgraphWithUdlPosSrcOnlyRootNodes should contain allElementsOf userDefinedPosSrcNode + } + + // TODO + // ensure caching sanity, by checking pos samples and hydrated nei, in integration test + + // TODO(yliu2) maybe can create test to check each sample subgraph output + // check resulting subgraph pos is in nodes, hard neg is in nodes +} From c8fc28efd01c686355c677b8a86d1222d1f8a083 Mon Sep 17 00:00:00 2001 From: Kyle Montemayor Date: Wed, 5 Aug 2026 16:41:27 +0000 Subject: [PATCH 2/5] Default SubgraphSampler to the Spark 3.5 runner Google blocks Dataproc image 2.0 cluster creation on 2026-08-25; the use_spark35=False path creates 2.0.47 clusters. Flip the use_spark35_runner experimental-flag default to "True" so default-configured pipelines (including the cora e2e tests) move to Dataproc 2.2 now, while an explicit "False" remains a rollback escape hatch until the 2.0 branch is deleted after the cutoff. test_subgraph_sampler_for_spark now exercises the new default: the spark35 sidecar jar is uploaded and passed alongside the 3.5 tfrecord jar, and the cluster is created with use_spark35=True. Co-Authored-By: Claude Opus 5 --- gigl/src/subgraph_sampler/subgraph_sampler.py | 5 ++++- .../subgraph_sampler/subgraph_sampler_test.py | 21 +++++++++++-------- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/gigl/src/subgraph_sampler/subgraph_sampler.py b/gigl/src/subgraph_sampler/subgraph_sampler.py index 8b7c140d1..359418e8a 100644 --- a/gigl/src/subgraph_sampler/subgraph_sampler.py +++ b/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" ) ) ) diff --git a/tests/unit/src/subgraph_sampler/subgraph_sampler_test.py b/tests/unit/src/subgraph_sampler/subgraph_sampler_test.py index cbdf7b7b4..a30ee5836 100644 --- a/tests/unit/src/subgraph_sampler/subgraph_sampler_test.py +++ b/tests/unit/src/subgraph_sampler/subgraph_sampler_test.py @@ -8,10 +8,7 @@ import gigl.env.dep_constants as dep_constants import gigl.src.common.constants.gcs as gcs_constants from gigl.common import GcsUri, LocalUri, UriFactory -from gigl.common.constants import ( - SPARK_31_TFRECORD_JAR_GCS_PATH, - SPARK_35_TFRECORD_JAR_GCS_PATH, -) +from gigl.common.constants import SPARK_35_TFRECORD_JAR_GCS_PATH from gigl.src.common.types import AppliedTaskIdentifier from gigl.src.common.utils import metrics_service_provider from gigl.src.subgraph_sampler import subgraph_sampler @@ -148,9 +145,7 @@ def test_subgraph_sampler_for_spark( applied_task_identifier=self.task_identifier, resource_config_uri=LocalUri(self.resource_config_path_local_path), task_config_uri=LocalUri(self.gbml_config_path_local_path), - additional_spark35_jar_file_uris=[ - LocalUri("/does/not/exist/should/not/be/passed/in") - ], + additional_spark35_jar_file_uris=[LocalUri(self.sidecar_jar_local_path)], ) subgraph_sampler_root = gcs_constants.get_subgraph_sampler_root_dir( applied_task_identifier=self.task_identifier @@ -193,6 +188,10 @@ def test_subgraph_sampler_for_spark( subgraph_sampler_root, "subgraph_sampler.jar", ), + LocalUri(self.sidecar_jar_local_path): GcsUri.join( + subgraph_sampler_root, + "sidecar.jar", + ), }, ) @@ -205,9 +204,13 @@ def test_subgraph_sampler_for_spark( max_job_duration=ANY, runtime_args=ANY, extra_jar_file_uris=[ - SPARK_31_TFRECORD_JAR_GCS_PATH, + GcsUri.join( + subgraph_sampler_root, + "sidecar.jar", + ).uri, + SPARK_35_TFRECORD_JAR_GCS_PATH, ], - use_spark35=False, + use_spark35=True, ) @patch(_INGESTOR_FQN) From f7834784093d1b6403e1823b02467689976e7f5d Mon Sep 17 00:00:00 2001 From: Kyle Montemayor Date: Wed, 5 Aug 2026 16:41:45 +0000 Subject: [PATCH 3/5] Bump Dataproc 2.2 image to 2.2.85 2.2.19 is over a year stale; 2.2.85 (2026-06-30, Spark 3.5.3) is the newest 2.2.x listed on the Dataproc release page. Separate commit because this line also moves SplitGenerator, which already runs on 2.2 unconditionally, so the bump can be reverted independently of the SubgraphSampler migration. Co-Authored-By: Claude Opus 5 --- gigl/src/common/utils/spark_job_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gigl/src/common/utils/spark_job_manager.py b/gigl/src/common/utils/spark_job_manager.py index f72db92ec..bb743e0d2 100644 --- a/gigl/src/common/utils/spark_job_manager.py +++ b/gigl/src/common/utils/spark_job_manager.py @@ -96,7 +96,7 @@ def create_dataproc_cluster( image_version: str if use_spark35: # https://cloud.google.com/dataproc/docs/concepts/versioning/dataproc-release-2.2 - image_version = "2.2.19-ubuntu22" + image_version = "2.2.85-ubuntu22" gce_cluster_config = GceClusterConfig( service_account=cluster_init_data.service_account, service_account_scopes=[ From 9a97cd23e778850847802d4a088aa0de650ba4b4 Mon Sep 17 00:00:00 2001 From: Kyle Montemayor Date: Wed, 5 Aug 2026 22:34:47 +0000 Subject: [PATCH 4/5] Revert "Bump Dataproc 2.2 image to 2.2.85" This reverts commit f78347840e2c99b40a0b28cbc4b19b48eec9c1a6. Image 2.2.85 breaks the SubgraphSampler Spark job. Every V1 e2e pipeline failed within minutes of the cluster coming up, all with the same error: Exception in thread "main" java.lang.NoClassDefFoundError: io/grpc/Context dblp_nalp fails too, and it already ran the spark35 jar on 2.2.19 before this branch, so the sub-minor is the variable, not the ported pureSpark code. 2.2.85 evidently ships a different gRPC classpath than 2.2.19. 2.2.19 stays on a supported minor, and escaping image 2.0 before it becomes uncreatable on 2026-08-25 does not require the newest sub-minor. Bumping it is worth doing separately, once the classpath change is understood. Co-Authored-By: Claude Opus 5 --- gigl/src/common/utils/spark_job_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gigl/src/common/utils/spark_job_manager.py b/gigl/src/common/utils/spark_job_manager.py index bb743e0d2..f72db92ec 100644 --- a/gigl/src/common/utils/spark_job_manager.py +++ b/gigl/src/common/utils/spark_job_manager.py @@ -96,7 +96,7 @@ def create_dataproc_cluster( image_version: str if use_spark35: # https://cloud.google.com/dataproc/docs/concepts/versioning/dataproc-release-2.2 - image_version = "2.2.85-ubuntu22" + image_version = "2.2.19-ubuntu22" gce_cluster_config = GceClusterConfig( service_account=cluster_init_data.service_account, service_account_scopes=[ From 08c777e948b321f07c0985d54bcde735398abc68 Mon Sep 17 00:00:00 2001 From: Kyle Montemayor Date: Thu, 6 Aug 2026 01:42:37 +0000 Subject: [PATCH 5/5] Rename nested array-element fields to proto names in pureSpark SGS casts The three pureSpark `castTo*ProtoSchema` functions renamed only top-level columns to their proto names. Array-of-struct columns were passed through carrying the sampling pipeline's internal `_`-prefixed element field names (`_node_id`, `_src_node`, ...). The always-empty `neg_edges` / `hard_neg_edges` were emitted as bare `ARRAY()` literals. Both were latent bugs that `sparksql31-scalapb0_11 1.0.0` hid. In 1.0.0, `FromCatalystHelpers.fieldFromCatalyst` wrapped repeated fields in `MapObjects(lambda, input, protoSql.singularDataType(fd))`. The lambda variable carried the proto's *declared* struct type, so the by-name field lookup inside the lambda resolved against the declared names and compiled down to ordinal access. At runtime the data was therefore read positionally and the wrong element names never mattered. Field order happens to match the proto at every creation site, so Spark 3.1 produced correct protos regardless. `sparksql35-scalapb0_11 1.0.4` builds the deserializer with `UnresolvedMapObjects` instead, deferring element typing to the analyzer. Spark's `ResolveDeserializer` binds the lambda variable to the *actual* element type and resolves fields by name, which fails at analysis time: [FIELD_NOT_FOUND] No such struct field `node_id` in `_node_id`, `_condensed_node_type`, `_feature_values` A bare `ARRAY()` is `ARRAY`, and with no declared element type the analyzer has nothing to bind the lambda variable to there either: [INVALID_EXTRACT_BASE_FIELD_TYPE] Can't extract a value from "lambdavariable(MapObject, NullType, false, 347)" ... but got "VOID" This is why cora_nalp, cora_udl and cora_snc all die in SubgraphSampler on Dataproc 2.2 while dblp_nalp passes -- dblp takes the GraphDB path, whose SQL already aliases nested fields to the proto names. Downgrading is not an option: only 1.0.4 and 1.0.5 exist for the `sparksql35` artifact, and both postdate the change. Rewrite the element field names to the proto names at the cast boundary with `transform`, and give the empty edge arrays an explicit element type. The rewrite lives in three shared helpers on `SGSPureSparkV1Task` so the five call sites across the three cast functions cannot drift apart. `transform` is NULL-safe, so the UNION branches that emit `NULL` for their neighbor arrays still deserialize to empty sequences. Creation sites and intermediate joins keep their `_`-prefixed convention. The legacy `scala/` (Spark 3.1) tree is deliberately left alone: 1.0.0's positional semantics make the same code correct there, and it is being retired. No pureSpark test previously crossed the `castTo*` + `.as[proto]` boundary, which is why this reached production, so add two regression tests. The RootedNodeNeighborhood one uses the production loaders against the checked-in TFRecord assets rather than this suite's mock fixtures, since those mocks use scalar/double features that the proto encoder rejects for unrelated reasons. Co-Authored-By: Claude Opus 5 --- ...odeAnchorBasedLinkPredictionBaseTask.scala | 20 +++--- .../task/pureSpark/SGSPureSparkV1Task.scala | 49 +++++++++++++- .../SupervisedNodeClassificationTask.scala | 4 +- ...odeAnchorBasedLinkPredictionTaskTest.scala | 42 ++++++++++++ .../test/scala/SGSPureSparkV1TaskTest.scala | 66 +++++++++++++++++++ 5 files changed, 167 insertions(+), 14 deletions(-) diff --git a/scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/NodeAnchorBasedLinkPredictionBaseTask.scala b/scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/NodeAnchorBasedLinkPredictionBaseTask.scala index af33dc7d7..fcfd28e4c 100644 --- a/scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/NodeAnchorBasedLinkPredictionBaseTask.scala +++ b/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} """) diff --git a/scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/SGSPureSparkV1Task.scala b/scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/SGSPureSparkV1Task.scala index 7a8a2a8cf..bf80c2a9b 100644 --- a/scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/SGSPureSparkV1Task.scala +++ b/scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/SGSPureSparkV1Task.scala @@ -1017,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>` + */ + 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>` + */ + 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`. 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>>)" + def castToRootedNodeNeighborhoodProtoSchema(dfVIEW: String): String = { /** Creates desirable schema for RootedNodeNeighborhood defined in training_samples_schema.proto. Returns rnnVIEW @@ -1030,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} """) diff --git a/scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/SupervisedNodeClassificationTask.scala b/scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/SupervisedNodeClassificationTask.scala index 46fdaa1a0..c205e206a 100644 --- a/scala_spark35/subgraph_sampler/src/main/scala/libs/task/pureSpark/SupervisedNodeClassificationTask.scala +++ b/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 diff --git a/scala_spark35/subgraph_sampler/src/test/scala/NodeAnchorBasedLinkPredictionTaskTest.scala b/scala_spark35/subgraph_sampler/src/test/scala/NodeAnchorBasedLinkPredictionTaskTest.scala index 4138e8144..832ac9017 100644 --- a/scala_spark35/subgraph_sampler/src/test/scala/NodeAnchorBasedLinkPredictionTaskTest.scala +++ b/scala_spark35/subgraph_sampler/src/test/scala/NodeAnchorBasedLinkPredictionTaskTest.scala @@ -193,6 +193,48 @@ class NodeAnchorBasedLinkPredictionTaskTest expectedNeighborEdges should contain allElementsOf currentNeighborEdges } + test("Training samples conform to the NodeAnchorBasedLinkPredictionSample proto schema.") { + // Regression test for the Spark 3.5 / sparksql-scalapb 1.0.4 upgrade. Two distinct failures are + // covered here, both of which only appear once repeated proto fields are resolved by name: + // 1. array-of-struct elements carrying `_`-prefixed field names instead of the proto's names, + // which fails with `FIELD_NOT_FOUND`, and + // 2. the always-empty `neg_edges` / `hard_neg_edges` literals, which are `ARRAY` unless + // explicitly CAST, and fail with `INVALID_EXTRACT_BASE_FIELD_TYPE ... got "VOID"`. + val uniqueTestViewSuffix = "_" + randomUUID.toString.replace("-", "_") + val mockSubgraphDF = mockSubgraphForCurrentTest + val mockSubgraphVIEW = "mockSubgraphDF" + uniqueTestViewSuffix + mockSubgraphDF.createOrReplaceTempView(mockSubgraphVIEW) + // Reuse the neighbor edges as the supervision edges; only their schema matters here. + val trainingSubgraphDF = sparkTest.sql(s""" + SELECT + _root_node, _condensed_node_type, _node_features, _neighbor_nodes, _neighbor_edges, + _neighbor_edges AS _pos_hydrated_edges, + _neighbor_edges AS _neg_hydrated_edges + FROM ${mockSubgraphVIEW} + """) + val trainingSubgraphVIEW = "trainingSubgraphDF" + uniqueTestViewSuffix + trainingSubgraphDF.createOrReplaceTempView(trainingSubgraphVIEW) + + for (hasHardNegs <- Seq(false, true)) { + val nablpWithSchemaVIEW = nablpTask.castToTrainingSampleProtoSchema( + dfVIEW = trainingSubgraphVIEW, + hasHardNegs = hasHardNegs, + ) + val samples = + sparkTest.table(nablpWithSchemaVIEW).as[NodeAnchorBasedLinkPredictionSample].collect() + assert(samples.length == mockSubgraphDF.count()) + // `neg_edges` is always the empty literal, for both branches. + assert(samples.forall(_.negEdges.isEmpty)) + assert(samples.forall(_.posEdges.nonEmpty)) + assert(samples.forall(_.hardNegEdges.nonEmpty) == hasHardNegs) + // Field values must survive the rename, not just the field names. + val sample = samples.head + sample.posEdges.map(_.srcNodeId) should contain allElementsOf + sample.neighborhood.get.edges.map(_.srcNodeId) + sample.neighborhood.get.nodes.map(_.nodeId) shouldNot be(empty) + } + } + test("validation fails if nodes of supervision edges not present in neighborhood nodes") { val rootNode = Node(nodeId = 0) diff --git a/scala_spark35/subgraph_sampler/src/test/scala/SGSPureSparkV1TaskTest.scala b/scala_spark35/subgraph_sampler/src/test/scala/SGSPureSparkV1TaskTest.scala index 9fe4808f5..0065d6cc5 100644 --- a/scala_spark35/subgraph_sampler/src/test/scala/SGSPureSparkV1TaskTest.scala +++ b/scala_spark35/subgraph_sampler/src/test/scala/SGSPureSparkV1TaskTest.scala @@ -4,6 +4,7 @@ import common.types.pb_wrappers.GbmlConfigPbWrapper import common.utils.ProtoLoader.populateProtoFromYaml import libs.task.pureSpark.SGSPureSparkV1Task import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Encoder import org.apache.spark.sql.Row import org.apache.spark.sql.types._ import org.apache.spark.sql.{functions => F} @@ -11,9 +12,22 @@ import org.scalatest.BeforeAndAfterAll import org.scalatest.funsuite.AnyFunSuite import org.scalatest.matchers.should.Matchers._ import snapchat.research.gbml.gbml_config.GbmlConfig +import snapchat.research.gbml.training_samples_schema.RootedNodeNeighborhood import java.util.UUID.randomUUID +/** Holds the scalapb proto encoders for this suite. + * + * They live outside the suite because `SGSPureSparkV1TaskTest` imports `sqlImplicits._` at class + * level to build mock DataFrames from Scala Seqs, and `sqlImplicits`' generic product encoder + * takes precedence over the scalapb one for proto case classes. + */ +private object ProtoEncoders { + import scalapb.spark.Implicits._ + val rootedNodeNeighborhood: Encoder[RootedNodeNeighborhood] = + implicitly[Encoder[RootedNodeNeighborhood]] +} + class SGSPureSparkV1TaskTest extends AnyFunSuite with BeforeAndAfterAll with SharedSparkSession { import sqlImplicits._ @@ -388,6 +402,58 @@ class SGSPureSparkV1TaskTest extends AnyFunSuite with BeforeAndAfterAll with Sha } + test("Rooted Node Neighborhood conforms to the RootedNodeNeighborhood proto schema.") { + // Regression test for the Spark 3.5 / sparksql-scalapb 1.0.4 upgrade. 1.0.4 resolves repeated + // proto fields by name rather than by position, so the `_`-prefixed field names that the + // sampling pipeline uses inside array-of-struct columns have to be rewritten to the proto's + // names before `.as[RootedNodeNeighborhood]`. Without the rewrite this fails at analysis time + // with `FIELD_NOT_FOUND: No such struct field node_id in _node_id, ...`. + // + // The production loaders are used deliberately: they read the checked-in TFRecord assets and + // therefore produce the `array` feature types that real preprocessed data has. The + // mock fixtures in this suite use scalar/double features, which the proto encoder rejects for + // unrelated reasons. + sparkTest.sqlContext.tableNames().foreach(sparkTest.catalog.dropTempView(_)) + + val hydratedNodeVIEW = sgsTask.loadNodeDataframeIntoSparkSql(condensedNodeType = 0) + val hydratedEdgeVIEW = sgsTask.loadEdgeDataframeIntoSparkSql(condensedEdgeType = 0) + val unhydratedEdgeVIEW = sgsTask.loadUnhydratedEdgeDataframeIntoSparkSql(hydratedEdgeVIEW) + + val subgraphVIEW = sgsTask.createSubgraph( + hydratedNodeVIEW = hydratedNodeVIEW, + hydratedEdgeVIEW = hydratedEdgeVIEW, + unhydratedEdgeVIEW = unhydratedEdgeVIEW, + numNeighborsToSample = 3, + permutationStrategy = "non-deterministic", + ) + val rnnVIEW = sgsTask.createRootedNodeNeighborhoodSubgraph( + hydratedNodeVIEW = hydratedNodeVIEW, + unhydratedEdgeVIEW = unhydratedEdgeVIEW, + subgraphVIEW = subgraphVIEW, + ) + // Neighborless root nodes get a NULL `_neighbor_edges`; the cast has to keep passing those + // through so they deserialize to an empty `edges` sequence. + val numRootNodesWithoutEdges = + sparkTest.table(rnnVIEW).filter(F.col("_neighbor_edges").isNull).count() + assert(numRootNodesWithoutEdges > 0) + + val rnnWithSchemaVIEW = sgsTask.castToRootedNodeNeighborhoodProtoSchema(dfVIEW = rnnVIEW) + val samples = sparkTest + .table(rnnWithSchemaVIEW) + .as[RootedNodeNeighborhood](ProtoEncoders.rootedNodeNeighborhood) + .collect() + + assert(samples.length == sparkTest.table(rnnVIEW).count()) + assert(samples.count(_.neighborhood.get.edges.isEmpty) == numRootNodesWithoutEdges) + // Field values must survive the rename, not just the field names. + val sampleWithEdges = samples.filter(_.neighborhood.get.edges.nonEmpty).head + val neighborhood = sampleWithEdges.neighborhood.get + neighborhood.nodes.map(_.nodeId) should contain(sampleWithEdges.rootNode.get.nodeId) + val nodeIdsFromEdges = + neighborhood.edges.flatMap(edge => Seq(edge.srcNodeId, edge.dstNodeId)).distinct + neighborhood.nodes.map(_.nodeId) should contain allElementsOf nodeIdsFromEdges + } + test("Rooted Node Neighborhood is valid.") { // delete all prev view names from current sparkTest session sparkTest.sqlContext.tableNames().foreach(sparkTest.catalog.dropTempView(_))