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/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/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 a67142d47..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 @@ -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( @@ -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>` + */ + 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 @@ -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} """) 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 new file mode 100644 index 000000000..832ac9017 --- /dev/null +++ b/scala_spark35/subgraph_sampler/src/test/scala/NodeAnchorBasedLinkPredictionTaskTest.scala @@ -0,0 +1,281 @@ +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("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) + + // 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..0065d6cc5 --- /dev/null +++ b/scala_spark35/subgraph_sampler/src/test/scala/SGSPureSparkV1TaskTest.scala @@ -0,0 +1,674 @@ +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.Encoder +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 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._ + // 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 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(_)) + 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 +} 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)