From fecc9540b90882f7718e0da09977233eba5942fc Mon Sep 17 00:00:00 2001 From: Mathieu Guillame-Bert Date: Thu, 24 Sep 2026 08:50:57 -0700 Subject: [PATCH] Add tutorials for offline distributed samplers. PiperOrigin-RevId: 987536318 --- .../gcp_offline_distributed_sampler.ipynb | 1886 +++++ doc/docs/tutorial/sampler.ipynb | 6180 ++++++++--------- doc/mkdocs.yml | 3 +- 3 files changed, 4967 insertions(+), 3102 deletions(-) create mode 100644 doc/docs/tutorial/gcp_offline_distributed_sampler.ipynb diff --git a/doc/docs/tutorial/gcp_offline_distributed_sampler.ipynb b/doc/docs/tutorial/gcp_offline_distributed_sampler.ipynb new file mode 100644 index 0000000..a632283 --- /dev/null +++ b/doc/docs/tutorial/gcp_offline_distributed_sampler.ipynb @@ -0,0 +1,1886 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "52d5d84b" + }, + "source": [ + "## Offline Distributed Sampler (GCP)\n", + "\n", + "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/distributed_graph_flow/blob/main/doc/docs/tutorial/gcp_offline_distributed_sampler.ipynb)\n", + "\n", + "This tutorial shows how to generate graph samples with the offline distributed\n", + "sampler on Google Cloud (GCP). The sampler runs as an Apache Beam pipeline on\n", + "Dataflow, launched through a Vertex AI custom job. This option is well suited\n", + "for large graphs (e.g., more than 100 million nodes). For smaller graphs, the\n", + "in-process sampler (see the \"In-process Sampler\" tutorial) is simpler to use.\n", + "\n", + "In this tutorial, you will:\n", + "\n", + "1. Copy a toy graph (OGB MAG) to a Cloud Storage bucket.\n", + "2. Generate graph samples with the offline distributed sampler.\n", + "3. Read and plot a generated sample.\n", + "4. Train a node classification model on the samples.\n", + "\n", + "**Note:** This tutorial is based on the\n", + "`create_graph_samples_offline_distributed_gcp.py` example." + ], + "id": "52d5d84b" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "EsL3MiYR9nD6" + }, + "source": [ + "" + ], + "id": "EsL3MiYR9nD6" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "D1I1ijy08t0v" + }, + "source": [ + "## Installing GF\n", + "\n", + "The sampling runs on GCP, so it does not require a local accelerator. However,\n", + "the last section of this tutorial trains a model: make sure your machine has a\n", + "GPU or TPU, otherwise training will be slow. If you are using Google Colab, you\n", + "can get one for free: go to Edit > Notebook settings and select a GPU or TPU as\n", + "the hardware accelerator." + ], + "id": "D1I1ijy08t0v" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:37.856620Z", + "iopub.status.busy": "2026-09-23T18:32:37.856434Z", + "iopub.status.idle": "2026-09-23T18:32:38.101388Z", + "shell.execute_reply": "2026-09-23T18:32:38.100901Z" + }, + "id": "eASOvqWCLNMQ" + }, + "outputs": [], + "source": [ + "# Install DGF (Distributed Graph Flow), the Vertex AI SDK, and OGB (for the toy\n", + "# dataset).\n", + "!pip install dgf ogb google-cloud-aiplatform -U" + ], + "id": "eASOvqWCLNMQ" + }, + { + "cell_type": "markdown", + "metadata": { + "id": "71e9cf1c" + }, + "source": [ + "## Importing libraries" + ], + "id": "71e9cf1c" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:38.144052Z", + "iopub.status.busy": "2026-09-23T18:32:38.143850Z", + "iopub.status.idle": "2026-09-23T18:32:42.505988Z", + "shell.execute_reply": "2026-09-23T18:32:42.505456Z" + }, + "id": "f5431c44" + }, + "outputs": [], + "source": [ + "import dgf # Import Graph Flow" + ], + "id": "f5431c44" + }, + { + "cell_type": "markdown", + "source": [ + "## Setup\n", + "\n", + "In this example, we generate samples from the OGB MAG dataset. This dataset is\n", + "small, but it is sufficient to demonstrate the sampler. To prepare, we\n", + "authenticate with GCP and copy OGB MAG to a Cloud Storage bucket.\n", + "\n", + "**Requirements:** You need a GCP project with billing enabled, and with the\n", + "Vertex AI, Dataflow, Compute Engine, and Cloud Storage APIs enabled. For\n", + "example:\n", + "\n", + "```shell\n", + "gcloud services enable aiplatform.googleapis.com dataflow.googleapis.com \\\n", + " compute.googleapis.com storage.googleapis.com\n", + "```" + ], + "metadata": { + "id": "0g2tYOlRcCTq" + }, + "id": "0g2tYOlRcCTq" + }, + { + "cell_type": "code", + "source": [ + "# Authenticate with your Google account (in Colab).\n", + "from google.colab import auth\n", + "\n", + "auth.authenticate_user()\n", + "\n", + "# Outside of Colab, use instead:\n", + "# !gcloud auth login --no-launch-browser\n", + "# !gcloud auth application-default login --no-launch-browser" + ], + "metadata": { + "id": "lLHF9EzbcEYA" + }, + "id": "lLHF9EzbcEYA", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Set the GCP project. Replace `graphflow-experiments-49784` with your own\n", + "project ID, here and in the rest of the tutorial:" + ], + "metadata": { + "id": "5LPKZC4bwqs7" + }, + "id": "5LPKZC4bwqs7" + }, + { + "cell_type": "code", + "source": [ + "!gcloud config set project \"graphflow-experiments-49784\"" + ], + "metadata": { + "id": "bva_tdZacGzX" + }, + "id": "bva_tdZacGzX", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Create a Cloud Storage bucket to store the graph and the generated samples.\n", + "Bucket names are globally unique: replace `dgf-tutorial-bucket` with a name of\n", + "your own, here and in the rest of the tutorial." + ], + "metadata": { + "id": "YpaEG65twpam" + }, + "id": "YpaEG65twpam" + }, + { + "cell_type": "code", + "source": [ + "!gcloud storage buckets create gs://dgf-tutorial-bucket --location=us-central1" + ], + "metadata": { + "id": "AshsKDoPciLx" + }, + "id": "AshsKDoPciLx", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "DGF provides direct access to OGB MAG. Let's download it and copy it to our\n", + "bucket:" + ], + "metadata": { + "id": "cV-UP_ctcYCp" + }, + "id": "cV-UP_ctcYCp" + }, + { + "cell_type": "code", + "source": [ + "graph, schema = dgf.io.fetch_ogb_graph(\"mag\")\n", + "dgf.io.write_graph(graph, schema, \"gs://dgf-tutorial-bucket/graph\")" + ], + "metadata": { + "id": "0uI54W3mdAlo" + }, + "id": "0uI54W3mdAlo", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Let's check that the graph files are there:" + ], + "metadata": { + "id": "cRtmA0-mdRZX" + }, + "id": "cRtmA0-mdRZX" + }, + { + "cell_type": "code", + "source": [ + "!gcloud storage ls gs://dgf-tutorial-bucket/graph" + ], + "metadata": { + "id": "RB_fj5EVdTjg", + "executionInfo": { + "status": "ok", + "timestamp": 1790242397276, + "user_tz": -120, + "elapsed": 2347, + "user": { + "displayName": "", + "userId": "" + } + }, + "outputId": "f0fcc452-bc36-423e-ef64-ddce4077654a" + }, + "id": "RB_fj5EVdTjg", + "execution_count": null, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "gs://dgf-tutorial-bucket/graph/edgesets_$folder$\n", + "gs://dgf-tutorial-bucket/graph/metadata.json\n", + "gs://dgf-tutorial-bucket/graph/nodesets_$folder$\n", + "gs://dgf-tutorial-bucket/graph/schema.json\n", + "gs://dgf-tutorial-bucket/graph/edgesets/\n", + "gs://dgf-tutorial-bucket/graph/nodesets/\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "source": [ + "It is always a good idea to check the schema of a graph:" + ], + "metadata": { + "id": "DxoYuQgBeoI5" + }, + "id": "DxoYuQgBeoI5" + }, + { + "cell_type": "code", + "source": [ + "dgf.print.schema(\n", + " dgf.io.read_schema(\"gs://dgf-tutorial-bucket/graph/schema.json\")\n", + ")" + ], + "metadata": { + "id": "51l31yqseqUv", + "executionInfo": { + "status": "ok", + "timestamp": 1790242435792, + "user_tz": -120, + "elapsed": 2064, + "user": { + "displayName": "", + "userId": "" + } + }, + "outputId": "1d68288f-d473-455e-8e6b-35c3c51197a9" + }, + "id": "51l31yqseqUv", + "execution_count": null, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Graph Schema:\n", + "\n", + "Node Sets:\n", + " author:\n", + " | Feature | Format | Semantic | Shape | Detail |\n", + " |-----------|----------|------------|---------|----------|\n", + " | #id | BYTES | PRIMARY_ID | None | |\n", + "\n", + " field_of_study:\n", + " | Feature | Format | Semantic | Shape | Detail |\n", + " |-----------|----------|------------|---------|----------|\n", + " | #id | BYTES | PRIMARY_ID | None | |\n", + "\n", + " institution:\n", + " | Feature | Format | Semantic | Shape | Detail |\n", + " |-----------|----------|------------|---------|----------|\n", + " | #id | BYTES | PRIMARY_ID | None | |\n", + "\n", + " paper:\n", + " | Feature | Format | Semantic | Shape | Detail |\n", + " |-----------|------------|-------------|---------|----------|\n", + " | #id | BYTES | PRIMARY_ID | None | |\n", + " | #split | BYTES | CATEGORICAL | None | |\n", + " | feat | FLOAT_32 | EMBEDDING | (128,) | |\n", + " | labels | INTEGER_64 | CATEGORICAL | None | |\n", + " | year | INTEGER_64 | NUMERICAL | None | |\n", + "\n", + "\n", + "Edge Sets:\n", + " affiliated_with: (Source: author, Target: institution)\n", + " (No features)\n", + "\n", + " cites: (Source: paper, Target: paper)\n", + " (No features)\n", + "\n", + " has_topic: (Source: paper, Target: field_of_study)\n", + " (No features)\n", + "\n", + " writes: (Source: author, Target: paper)\n", + " (No features)\n", + "\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "source": [ + "## Run the sampler" + ], + "metadata": { + "id": "Mx6JdDoNddsG" + }, + "id": "Mx6JdDoNddsG" + }, + { + "cell_type": "markdown", + "source": [ + "The sampler runs as the default Compute Engine service account of your\n", + "project. Give this service account read and write access to the bucket:" + ], + "metadata": { + "id": "jkPVxQkawoFA" + }, + "id": "jkPVxQkawoFA" + }, + { + "cell_type": "code", + "source": [ + "!gcloud storage buckets add-iam-policy-binding \"gs://dgf-tutorial-bucket\" \\\n", + " --member=\"serviceAccount:$(gcloud projects describe graphflow-experiments-49784 --format='value(projectNumber)')-compute@developer.gserviceaccount.com\" \\\n", + " --role=\"roles/storage.objectAdmin\"" + ], + "metadata": { + "id": "tJUi-idsv2V6" + }, + "id": "tJUi-idsv2V6", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "The sampling plan defines how the graph is traversed to build each sample. In\n", + "most cases, a `SimpleSamplingConfig` is sufficient: it specifies the seed\n", + "nodeset (i.e., the nodes to generate samples for), the number of hops, and the\n", + "number of neighbors to sample at each hop (the hop width). For finer control\n", + "(e.g., configuring each hop individually), use a `SamplingPlan` instead." + ], + "metadata": { + "id": "rkTFs8aQw4vp" + }, + "id": "rkTFs8aQw4vp" + }, + { + "cell_type": "code", + "source": [ + "plan = dgf.sampling.SimpleSamplingConfig(\n", + " seed_nodeset=\"paper\",\n", + " num_hops=2,\n", + " hop_width=5,\n", + ")" + ], + "metadata": { + "id": "5DVVvI_ldejd" + }, + "id": "5DVVvI_ldejd", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Then, run the sampler. With `blocking=True`, the function waits until the\n", + "samples are generated and logs the progress of the job:" + ], + "metadata": { + "id": "pP0lvq3oeFDx" + }, + "id": "pP0lvq3oeFDx" + }, + { + "cell_type": "code", + "source": [ + "dgf.sampling.offline_distributed_sampler_gcp(\n", + " input_path=\"gs://dgf-tutorial-bucket/graph\",\n", + " output_path=\"gs://dgf-tutorial-bucket/samples\",\n", + " project=\"graphflow-experiments-49784\",\n", + " plan=plan,\n", + " num_workers=5, # Increase for larger graphs.\n", + " num_seeds=10000, # Generate 10k samples from randomly selected seed nodes.\n", + " blocking=True, # Wait until the samples are generated.\n", + ")" + ], + "metadata": { + "id": "XotLW7vCeIh8", + "executionInfo": { + "status": "ok", + "timestamp": 1790243624192, + "user_tz": -120, + "elapsed": 542006, + "user": { + "displayName": "", + "userId": "" + } + }, + "outputId": "244e2d35-f16a-4dff-a5c3-cef9bd3fae31" + }, + "id": "XotLW7vCeIh8", + "execution_count": null, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Reading graph schema from gs://dgf-tutorial-bucket/graph/schema.json\n", + "Writing sampling plan to gs://dgf-tutorial-bucket/samples/sampling_config.json\n", + "Submitting Vertex AI CustomJob 'dgf-distributed-sampler-glassbox-20260924114442'...\n", + "Creating CustomJob\n", + "CustomJob created. Resource name: projects/569105783505/locations/us-central1/customJobs/3683233532072689664\n", + "To use this CustomJob in another session:\n", + "custom_job = aiplatform.CustomJob.get('projects/569105783505/locations/us-central1/customJobs/3683233532072689664')\n", + "View Custom Job:\n", + "https://console.cloud.google.com/agent-platform/locations/us-central1/training/3683233532072689664?project=569105783505\n", + "Vertex AI CustomJob created: projects/569105783505/locations/us-central1/customJobs/3683233532072689664\n", + " Cloud Console: https://console.cloud.google.com/vertex-ai/locations/us-central1/training/3683233532072689664?project=graphflow-experiments-49784\n", + " Cloud Logging: https://console.cloud.google.com/logs/viewer?project=graphflow-experiments-49784&resource=ml_job%2Fjob_id%2F3683233532072689664\n" + ] + }, + { + "output_type": "stream", + "name": "stderr", + "text": [ + "\rDistributed Sampler [1/4: Job submitted]: 0%| | 0/4 [00:00 \n", + "resource name: projects/569105783505/locations/us-central1/customJobs/3683233532072689664" + ] + }, + "metadata": {}, + "execution_count": 15 + } + ] + }, + { + "cell_type": "markdown", + "source": [ + "**Remarks:**\n", + "\n", + "- MAG is a small dataset: most of the time is spent setting up the job.\n", + "- After a few seconds, the job appears in the\n", + " [Dataflow console](https://console.cloud.google.com/dataflow/jobs).\n", + "- Without `num_seeds`, one sample is generated for each node of the seed\n", + " nodeset. Use `num_samples_per_seed` to generate multiple samples per seed\n", + " node, or `input_seeds` / `filter_seed_node` to control which nodes are used\n", + " as seeds.\n", + "\n", + "Let's list the generated sample files:" + ], + "metadata": { + "id": "TD0IfCMLegCv" + }, + "id": "TD0IfCMLegCv" + }, + { + "cell_type": "code", + "source": [ + "!gcloud storage ls gs://dgf-tutorial-bucket/samples" + ], + "metadata": { + "id": "aIfUscnMeigA", + "executionInfo": { + "status": "ok", + "timestamp": 1790245835293, + "user_tz": -120, + "elapsed": 3180, + "user": { + "displayName": "", + "userId": "" + } + }, + "outputId": "1fb5a5b7-2994-43ce-ce13-350b9a005fa7" + }, + "id": "aIfUscnMeigA", + "execution_count": null, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "gs://dgf-tutorial-bucket/samples/samples-00000-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00001-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00002-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00003-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00004-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00005-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00006-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00007-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00008-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00009-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00010-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00011-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00012-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00013-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00014-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00015-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00016-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00017-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00018-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/samples-00019-of-00020.tfrecord.gz\n", + "gs://dgf-tutorial-bucket/samples/sampling_config.json\n", + "gs://dgf-tutorial-bucket/samples/schema.json\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "source": [ + "Samples are stored as TF-GNN graphs in sharded TFRecord files, next to their\n", + "schema. Let's read the first sample:" + ], + "metadata": { + "id": "8CxqjeoJejdv" + }, + "id": "8CxqjeoJejdv" + }, + { + "cell_type": "code", + "source": [ + "sample_schema = dgf.io.read_schema(\n", + " \"gs://dgf-tutorial-bucket/samples/schema.json\"\n", + ")\n", + "sample_generator = dgf.io.read_tfgnn_graphs(\n", + " \"gs://dgf-tutorial-bucket/samples/samples@*.tfrecord.gz\", sample_schema\n", + ")\n", + "sample = next(iter(sample_generator))" + ], + "metadata": { + "id": "IBCvp7kCfGMF" + }, + "id": "IBCvp7kCfGMF", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [ + "Let's plot it. `features=False` hides the features, which makes the topology\n", + "of the sample easier to see:" + ], + "metadata": { + "id": "Pum40NA5-CMc" + }, + "id": "Pum40NA5-CMc" + }, + { + "cell_type": "code", + "source": [ + "dgf.plot.plot_graph(sample, sample_schema, features=False)" + ], + "metadata": { + "colab": { + "height": 1000 + }, + "id": "fJhNCszigwmJ", + "executionInfo": { + "status": "ok", + "timestamp": 1790245882343, + "user_tz": -120, + "elapsed": 99, + "user": { + "displayName": "", + "userId": "" + } + }, + "outputId": "0e14525b-d5fe-4422-9f84-ee3a823d48fb" + }, + "id": "fJhNCszigwmJ", + "execution_count": null, + "outputs": [ + { + "output_type": "execute_result", + "data": { + "image/svg+xml": "\n\n\n\n\n\n\n\n\nauthor_0\n\nauthor_0\n\n\n\npaper_12\n\npaper_12\n\n\n\nauthor_0->paper_12\n\n\nwrites\n\n\n\nauthor_1\n\nauthor_1\n\n\n\nauthor_1->paper_12\n\n\nwrites\n\n\n\nauthor_2\n\nauthor_2\n\n\n\ninstitution_0\n\ninstitution_0\n\n\n\nauthor_2->institution_0\n\n\naffiliated_with\n\n\n\npaper_0\n\npaper_0\n\n\n\nauthor_2->paper_0\n\n\nwrites\n\n\n\npaper_1\n\npaper_1\n\n\n\nauthor_2->paper_1\n\n\nwrites\n\n\n\npaper_19\n\npaper_19\n\n\n\nauthor_2->paper_19\n\n\nwrites\n\n\n\npaper_23\n\npaper_23\n\n\n\nauthor_2->paper_23\n\n\nwrites\n\n\n\npaper_27\n\npaper_27\n\n\n\nauthor_2->paper_27\n\n\nwrites\n\n\n\nauthor_3\n\nauthor_3\n\n\n\nauthor_3->institution_0\n\n\naffiliated_with\n\n\n\nauthor_3->paper_0\n\n\nwrites\n\n\n\nauthor_3->paper_19\n\n\nwrites\n\n\n\nauthor_4\n\nauthor_4\n\n\n\nauthor_4->institution_0\n\n\naffiliated_with\n\n\n\nauthor_4->paper_0\n\n\nwrites\n\n\n\nauthor_4->paper_19\n\n\nwrites\n\n\n\nauthor_5\n\nauthor_5\n\n\n\nauthor_5->paper_12\n\n\nwrites\n\n\n\nauthor_6\n\nauthor_6\n\n\n\nauthor_6->paper_12\n\n\nwrites\n\n\n\nauthor_7\n\nauthor_7\n\n\n\nauthor_7->paper_12\n\n\nwrites\n\n\n\nfield_of_study_0\n\nfield_of_study_0\n\n\n\nfield_of_study_1\n\nfield_of_study_1\n\n\n\nfield_of_study_2\n\nfield_of_study_2\n\n\n\nfield_of_study_3\n\nfield_of_study_3\n\n\n\nfield_of_study_4\n\nfield_of_study_4\n\n\n\nfield_of_study_5\n\nfield_of_study_5\n\n\n\nfield_of_study_6\n\nfield_of_study_6\n\n\n\nfield_of_study_7\n\nfield_of_study_7\n\n\n\nfield_of_study_8\n\nfield_of_study_8\n\n\n\npaper_0->field_of_study_0\n\n\nhas_topic\n\n\n\npaper_0->field_of_study_2\n\n\nhas_topic\n\n\n\npaper_0->field_of_study_3\n\n\nhas_topic\n\n\n\npaper_0->field_of_study_4\n\n\nhas_topic\n\n\n\npaper_0->field_of_study_6\n\n\nhas_topic\n\n\n\npaper_0->paper_12\n\n\ncites\n\n\n\npaper_2\n\npaper_2\n\n\n\npaper_2->field_of_study_4\n\n\nhas_topic\n\n\n\npaper_3\n\npaper_3\n\n\n\npaper_3->field_of_study_4\n\n\nhas_topic\n\n\n\npaper_4\n\npaper_4\n\n\n\npaper_4->field_of_study_0\n\n\nhas_topic\n\n\n\npaper_5\n\npaper_5\n\n\n\npaper_5->field_of_study_2\n\n\nhas_topic\n\n\n\npaper_6\n\npaper_6\n\n\n\npaper_6->field_of_study_4\n\n\nhas_topic\n\n\n\npaper_7\n\npaper_7\n\n\n\npaper_7->field_of_study_6\n\n\nhas_topic\n\n\n\npaper_8\n\npaper_8\n\n\n\npaper_8->field_of_study_6\n\n\nhas_topic\n\n\n\npaper_9\n\npaper_9\n\n\n\npaper_9->field_of_study_0\n\n\nhas_topic\n\n\n\npaper_10\n\npaper_10\n\n\n\npaper_10->field_of_study_2\n\n\nhas_topic\n\n\n\npaper_11\n\npaper_11\n\n\n\npaper_11->field_of_study_2\n\n\nhas_topic\n\n\n\npaper_12->field_of_study_0\n\n\nhas_topic\n\n\n\npaper_12->field_of_study_1\n\n\nhas_topic\n\n\n\npaper_12->field_of_study_5\n\n\nhas_topic\n\n\n\npaper_12->field_of_study_7\n\n\nhas_topic\n\n\n\npaper_12->field_of_study_8\n\n\nhas_topic\n\n\n\npaper_13\n\npaper_13\n\n\n\npaper_13->field_of_study_4\n\n\nhas_topic\n\n\n\npaper_14\n\npaper_14\n\n\n\npaper_14->field_of_study_2\n\n\nhas_topic\n\n\n\npaper_15\n\npaper_15\n\n\n\npaper_15->field_of_study_0\n\n\nhas_topic\n\n\n\npaper_16\n\npaper_16\n\n\n\npaper_16->field_of_study_6\n\n\nhas_topic\n\n\n\npaper_17\n\npaper_17\n\n\n\npaper_17->field_of_study_2\n\n\nhas_topic\n\n\n\npaper_18\n\npaper_18\n\n\n\npaper_18->paper_12\n\n\ncites\n\n\n\npaper_19->paper_12\n\n\ncites\n\n\n\npaper_20\n\npaper_20\n\n\n\npaper_20->field_of_study_4\n\n\nhas_topic\n\n\n\npaper_21\n\npaper_21\n\n\n\npaper_21->field_of_study_6\n\n\nhas_topic\n\n\n\npaper_22\n\npaper_22\n\n\n\npaper_22->field_of_study_0\n\n\nhas_topic\n\n\n\npaper_24\n\npaper_24\n\n\n\npaper_24->field_of_study_0\n\n\nhas_topic\n\n\n\npaper_25\n\npaper_25\n\n\n\npaper_25->field_of_study_6\n\n\nhas_topic\n\n\n\npaper_26\n\npaper_26\n\n\n\npaper_26->paper_12\n\n\ncites\n\n\n\n", + "text/plain": [ + "" + ] + }, + "metadata": {}, + "execution_count": 20 + } + ] + }, + { + "cell_type": "markdown", + "source": [ + "The graph samples can be used directly to train a model with the simple API:\n", + "\n", + "**Note:** Reduce `num_train_steps` if you don't have a GPU/TPU." + ], + "metadata": { + "id": "WwQSMWPShIPk" + }, + "id": "WwQSMWPShIPk" + }, + { + "cell_type": "code", + "source": [ + "model = dgf.learning.train_node_model(\n", + " graph=\"gs://dgf-tutorial-bucket/samples/samples@*.tfrecord.gz\",\n", + " schema=sample_schema,\n", + " target_column=\"labels\",\n", + " target_nodeset=\"paper\",\n", + " num_train_steps=1000,\n", + " valid_every_n_steps=100,\n", + " verbose=1,\n", + ")" + ], + "metadata": { + "id": "QBuA4zekhLZW", + "executionInfo": { + "status": "ok", + "timestamp": 1790246307673, + "user_tz": -120, + "elapsed": 103320, + "user": { + "displayName": "", + "userId": "" + } + }, + "outputId": "a285a5e2-eaf1-4c4b-cf44-5e93fe22832d" + }, + "id": "QBuA4zekhLZW", + "execution_count": null, + "outputs": [ + { + "output_type": "stream", + "name": "stderr", + "text": [ + "[Warning] The model is trained on graph samples and no `sampling_plan` is provided. This is fine, but it is recommended to provide the `sampling_plan` used to generate those samples to make the model easier to use later (e.g. to generate predictions from a full graph).\n" + ] + }, + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Preparing dataset\n" + ] + }, + { + "output_type": "stream", + "name": "stderr", + "text": [ + "[Warning] No normalizer created for node set 'author', feature '#id'.\n", + "[Warning] No normalizer created for node set 'field_of_study', feature '#id'.\n", + "[Warning] No normalizer created for node set 'institution', feature '#id'.\n", + "[Warning] No normalizer created for node set 'paper', feature '#id'.\n" + ] + }, + { + "output_type": "stream", + "name": "stdout", + "text": [ + "Preparing dataset finished in 18.57 seconds\n", + "Caching validation dataset\n", + "Caching validation dataset finished in 9.65 seconds\n", + "Number of cache validation batches: 313\n", + "Training model\n", + "Generate first batch to initialize model\n", + "Create model variables\n", + "...Tracing model\n", + "Create model variables finished in 7.36 seconds\n", + "Will validate model every 100 step(s)\n", + "Will checkpoint model every 1000 step(s)\n", + "Start training. The first two steps are generally slow.\n" + ] + }, + { + "output_type": "stream", + "name": "stderr", + "text": [ + "\rTraining: 0%| | 0/1000 [00:00" + ], + "text/html": [ + "
\n", + "\n", + "\n", + "\n", + "
\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
\n", + "
\n", + "

Node prediction model: Predict the value of a node feature.

\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
TypeNode prediction model - Predict the value of a node feature.
Target nodesetpaper
Target columnlabels
Number of label classes349
\n", + "
\n", + "
\n", + "\n", + " \n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + " \n", + "
WARNINGThe model is trained on graph samples and no `sampling_plan` is provided. This is fine, but it is recommended to provide the `sampling_plan` used to generate those samples to make the model easier to use later (e.g. to generate predictions from a full graph).
WARNINGNo normalizer created for node set 'author', feature '#id'.
WARNINGNo normalizer created for node set 'field_of_study', feature '#id'.
WARNINGNo normalizer created for node set 'institution', feature '#id'.
WARNINGNo normalizer created for node set 'paper', feature '#id'.
\n", + "
\n", + "
\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
Accuracy0.3168
Num Examples10000
AUC0.959368428779122
Per Class Metrics [349]:
    \n", + "
  • Class 0: AUC=0.9503, PR-AUC=0.1430
    • \n", + "
  • Class 1: AUC=0.9674, PR-AUC=0.4592
    • \n", + "
  • Class 2: AUC=0.9013, PR-AUC=0.0009
    • \n", + "
  • ... (344 omitted) ...
    • \n", + "
  • Class 347: AUC=0.9923, PR-AUC=0.5048
    • \n", + "
  • Class 348: AUC=0.9759, PR-AUC=0.0057
    • \n", + "
\n", + "

*Showing plots for the first 20 classes out of 349 total classes.

\n", + "\n", + " \n", + " \n", + " \n", + " \n", + "
\n", + "
\n", + "
\n", + "\n", + " \n", + "\n", + " \n", + " \n", + " \n", + "\n", + " \n", + " \n", + "
Number of training seed nodesNone
Number of validation seed nodesNone
Training duration1m 30s
Number of training steps (final model)1000
\n", + "
\n", + " \n", + " \n", + " \n", + " \n", + "
\n", + "
\n", + "

Note: The logs for the first training step are not shown.

\n", + "
\n", + "
\n", + "\n", + " \n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + " \n", + "
num_sampling_hops2
sampling_width15
num_layers2
batch_size32
max_training_time_secondsNone
num_train_steps1000
random_seed42
node_embedding_dim128
learning_rate0.001
opt_weight_decay0.0001
dropout0.1
message_pooling'sum'
architecture<Architecture.HETEROGENEOUS_MESSAGE_PASSING: 'HETEROGENEOUS_MESSAGE_PASSING'>
early_stopping{'patience': 5, 'min_improvement': 1e-06}
\n", + "
\n", + "
Raw schema plot\n", + "\n", + "Raw schema textual\n", + "
Node Sets:\n",
+              "  author:\n",
+              "    | Feature   | Format   | Semantic   | Shape   | Detail   |\n",
+              "    |-----------|----------|------------|---------|----------|\n",
+              "    | #id       | BYTES    | PRIMARY_ID | None    |          |\n",
+              "\n",
+              "  field_of_study:\n",
+              "    | Feature   | Format   | Semantic   | Shape   | Detail   |\n",
+              "    |-----------|----------|------------|---------|----------|\n",
+              "    | #id       | BYTES    | PRIMARY_ID | None    |          |\n",
+              "\n",
+              "  institution:\n",
+              "    | Feature   | Format   | Semantic   | Shape   | Detail   |\n",
+              "    |-----------|----------|------------|---------|----------|\n",
+              "    | #id       | BYTES    | PRIMARY_ID | None    |          |\n",
+              "\n",
+              "  paper:\n",
+              "    | Feature   | Format     | Semantic    | Shape   | Detail   |\n",
+              "    |-----------|------------|-------------|---------|----------|\n",
+              "    | #id       | BYTES      | PRIMARY_ID  | None    |          |\n",
+              "    | #split    | BYTES      | CATEGORICAL | None    |          |\n",
+              "    | feat      | FLOAT_32   | EMBEDDING   | (128,)  |          |\n",
+              "    | labels    | INTEGER_64 | CATEGORICAL | None    |          |\n",
+              "    | year      | INTEGER_64 | NUMERICAL   | None    |          |\n",
+              "\n",
+              "\n",
+              "Edge Sets:\n",
+              "  affiliated_with: (Source: author, Target: institution)\n",
+              "    (No features)\n",
+              "\n",
+              "  cites: (Source: paper, Target: paper)\n",
+              "    (No features)\n",
+              "\n",
+              "  has_topic: (Source: paper, Target: field_of_study)\n",
+              "    (No features)\n",
+              "\n",
+              "  writes: (Source: author, Target: paper)\n",
+              "    (No features)\n",
+              "
\n", + "Normalized schema plot\n", + "\n", + "Normalized schema textual\n", + "
Node Sets:\n",
+              "  author:\n",
+              "    (No features)\n",
+              "\n",
+              "  field_of_study:\n",
+              "    (No features)\n",
+              "\n",
+              "  institution:\n",
+              "    (No features)\n",
+              "\n",
+              "  paper:\n",
+              "    | Feature            | Format     | Semantic    | Shape   | Detail       |\n",
+              "    |--------------------|------------|-------------|---------|--------------|\n",
+              "    | #split_INDEX       | INTEGER_64 | CATEGORICAL | ()      | #num.cat:4   |\n",
+              "    | feat               | FLOAT_32   | EMBEDDING   | (128,)  |              |\n",
+              "    | labels             | INTEGER_64 | CATEGORICAL | None    | #num.cat:349 |\n",
+              "    | year_SOFT_QUANTILE | FLOAT_32   | EMBEDDING   | ()      |              |\n",
+              "\n",
+              "\n",
+              "Edge Sets:\n",
+              "  affiliated_with: (Source: author, Target: institution)\n",
+              "    (No features)\n",
+              "\n",
+              "  cites: (Source: paper, Target: paper)\n",
+              "    (No features)\n",
+              "\n",
+              "  has_topic: (Source: paper, Target: field_of_study)\n",
+              "    (No features)\n",
+              "\n",
+              "  writes: (Source: author, Target: paper)\n",
+              "    (No features)\n",
+              "
\n", + "
\n", + "
Default feature statistics\n", + "
GraphFeatureStatistics:\n",
+              "  Node Sets (4):\n",
+              "    'author':\n",
+              "      '#id': count=213142, min=nan, max=nan\n",
+              "    'field_of_study':\n",
+              "      '#id': count=237143, min=nan, max=nan\n",
+              "    'institution':\n",
+              "      '#id': count=32329, min=nan, max=nan\n",
+              "    'paper':\n",
+              "      '#id': count=736157, min=nan, max=nan\n",
+              "      '#split': count=736157, min=nan, max=nan, dictionary=(3)['train': 617839, 'valid': 68072, 'test': 50246]\n",
+              "      'feat': count=736157, min=nan, max=nan\n",
+              "      'labels': count=736157, min=0.0000, max=348.0000\n",
+              "      'year': count=736157, min=2010.0000, max=2019.0000, quantiles=(100)[2010.0000, 2010.0000, 2010.0000, ..., 2019.0000, 2019.0000, 2019.0000]\n",
+              "
\n", + "
\n", + "
Default sampling plan\n", + "

The sampling plan is not available: the model was trained on already sampled graph samples, and no sampling plan was provided.

\n", + "
\n", + "
Model Structure\n", + "
EmbedGraph(cat-embedding=64)\n",
+              "Dense(128)\n",
+              "Activation(silu)\n",
+              "Norm(layer_norm)\n",
+              "Graph Convolution Block x2:\n",
+              "    X = ...\n",
+              "    MPNN:\n",
+              "      Message:\n",
+              "        Dense(128)\n",
+              "        Activation(silu)\n",
+              "        Dense(128)\n",
+              "      Update:\n",
+              "        Dense(128)\n",
+              "        Activation(silu)\n",
+              "        Dropout(0.1)\n",
+              "        Dense(128)\n",
+              "    Residual(X)\n",
+              "    # Post MPNN\n",
+              "    X = ...\n",
+              "    Norm(rms_norm)\n",
+              "    Dense(512)\n",
+              "    Activation(silu)\n",
+              "    Dense(128)\n",
+              "    Dropout(0.1)\n",
+              "    Residual(X)\n",
+              "Identity\n",
+              "Dense(349) # Classification head
\n", + "Model Weights\n", + "
{'float32': 2312413}
\n", + "
Default padding\n", + "
Node Sets:\n",
+              "  author: 939 nodes\n",
+              "  field_of_study: 995 nodes\n",
+              "  institution: 167 nodes\n",
+              "  paper: 3125 nodes\n",
+              "\n",
+              "Edge Sets:\n",
+              "  affiliated_with: 249 edges\n",
+              "  cites: 2115 edges\n",
+              "  has_topic: 2332 edges\n",
+              "  writes: 1516 edges
\n", + "
\n", + "\n", + "\n", + "\n", + "
" + ] + }, + "metadata": {}, + "execution_count": 28 + } + ] + }, + { + "cell_type": "markdown", + "source": [ + "## Clean up\n", + "\n", + "To avoid storage costs, delete the bucket and its content when you are done:\n", + "\n", + "```shell\n", + "gcloud storage rm --recursive gs://dgf-tutorial-bucket\n", + "```" + ], + "metadata": { + "id": "cleanupBucket01" + }, + "id": "cleanupBucket01" + } + ], + "metadata": { + "colab": { + "last_runtime": { + "build_target": "", + "kind": "local" + }, + "provenance": [ + { + "file_id": "1hMr9gK0hhKcCtIs1Pro3_ksOnYcdAL2P", + "timestamp": 1771583153488 + } + ], + "toc_visible": true, + "views": { + "output_only": { + "cells": [ + { + "id": "52d5d84b" + }, + { + "id": "EsL3MiYR9nD6" + }, + { + "id": "D1I1ijy08t0v" + }, + { + "id": "eASOvqWCLNMQ" + }, + { + "id": "71e9cf1c" + }, + { + "id": "f5431c44" + }, + { + "id": "7eb84356" + }, + { + "id": "7f9272a6" + }, + { + "id": "d0fb3022" + }, + { + "id": "2896f2d1" + }, + { + "id": "fba58d46" + }, + { + "id": "4ee7378e" + }, + { + "id": "daee962f" + }, + { + "id": "a98cd4ee" + }, + { + "id": "2d0902f0" + }, + { + "id": "5178bcff" + }, + { + "id": "1f9373f2" + }, + { + "id": "07fe45e1" + }, + { + "id": "8cea9faa" + }, + { + "id": "d67499e6" + }, + { + "id": "3a98d6ce" + }, + { + "id": "31410d56" + }, + { + "id": "3efc55a0" + }, + { + "id": "c93aafb9" + }, + { + "id": "e34b4b65" + }, + { + "id": "O5j_22f5HJYn" + }, + { + "id": "LGu-U8SkHTMV" + }, + { + "id": "Yxpn64EQb_0R" + }, + { + "id": "OrG4j5_OZ4lB" + }, + { + "id": "grTlK4Lnbvoq" + }, + { + "id": "868fTc8JbkzT" + }, + { + "id": "dK98Ep4Mcnya" + }, + { + "id": "5c5d4YfWcpar" + }, + { + "id": "bfbc2230" + }, + { + "id": "NvZ9fNQ7LxSs" + }, + { + "id": "5ad2a4f9" + }, + { + "id": "c6d421dc" + }, + { + "id": "909d6d86" + }, + { + "id": "91611d90" + }, + { + "id": "5cb0d579" + }, + { + "id": "a35660d2" + }, + { + "id": "ef4b8f03" + }, + { + "id": "4151f762" + }, + { + "id": "71fc0067" + }, + { + "id": "f4a3e3de" + }, + { + "id": "ecbaa932" + }, + { + "id": "70dcd856" + }, + { + "id": "b39436b6" + }, + { + "id": "d75df1fd" + } + ], + "hide_code": true + } + } + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.14.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/doc/docs/tutorial/sampler.ipynb b/doc/docs/tutorial/sampler.ipynb index afecf7f..d32e517 100644 --- a/doc/docs/tutorial/sampler.ipynb +++ b/doc/docs/tutorial/sampler.ipynb @@ -1,3174 +1,3152 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "id": "52d5d84b" - }, - "source": [ - "## Sampler\n", - "\n", - "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/distributed_graph_flow/blob/main/doc/docs/tutorial/sampler.ipynb)\n", - "\n", - "This tutorial shows the different options to generate graph samples.\n", - "\n", - "You'll learn how to use the in-memory sampler and the semi-distributed Beam sampler." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "EsL3MiYR9nD6" - }, - "source": [ - "" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "D1I1ijy08t0v" - }, - "source": [ - "## Installing GF\n", - "\n", - "Make sure your machine has a GPU or TPU, otherwise training is going to take forever.\n", - "If you are using Google Colab, you can get one for free. Just go to Edit > Notebook settings and select your hardware accelerator as TPU or GPU." - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:37.856620Z", - "iopub.status.busy": "2026-09-23T18:32:37.856434Z", - "iopub.status.idle": "2026-09-23T18:32:38.101388Z", - "shell.execute_reply": "2026-09-23T18:32:38.100901Z" - }, - "executionInfo": { - "elapsed": 286, - "status": "ok", - "timestamp": 1790188358142.653, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 - }, - "id": "eASOvqWCLNMQ" - }, - "outputs": [], - "source": [ - "# Install DGF (Distributed Graph Flow) and OGB (for the toy dataset).\n", - "!pip install dgf ogb -U" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "71e9cf1c" - }, - "source": [ - "## Importing libraries" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:38.144052Z", - "iopub.status.busy": "2026-09-23T18:32:38.143850Z", - "iopub.status.idle": "2026-09-23T18:32:42.505988Z", - "shell.execute_reply": "2026-09-23T18:32:42.505456Z" - }, - "executionInfo": { - "elapsed": 4364, - "status": "ok", - "timestamp": 1790188362507.4407, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 - }, - "id": "f5431c44" - }, - "outputs": [], - "source": [ - "import dgf # Import Graph Flow\n", - "import numpy as np\n", - "import copy" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "7eb84356" - }, - "source": [ - "## Load some data\n", - "\n", - "Let's start by loading a toy dataset." - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:42.508824Z", - "iopub.status.busy": "2026-09-23T18:32:42.508200Z", - "iopub.status.idle": "2026-09-23T18:32:42.589533Z", - "shell.execute_reply": "2026-09-23T18:32:42.589133Z" + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## In-process Sampler\n", + "\n", + "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/distributed_graph_flow/blob/main/doc/docs/tutorial/sampler.ipynb)\n", + "\n", + "This tutorial explains how to generate graph samples in-process. This approach works best when training on graphs with fewer than a few hundred million nodes.\n", + "\n", + "**Note:** When using the simple API, graph samples are generated automatically during training, so you do not need to create them beforehand." + ], + "id": "52d5d84b" }, - "executionInfo": { - "elapsed": 83, - "status": "ok", - "timestamp": 1790188362590.8496, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "" + ], + "id": "EsL3MiYR9nD6" }, - "id": "7f9272a6", - "outputId": "e216c215-39c3-4e80-a711-5d4c78928767" - }, - "outputs": [ { - "name": "stdout", - "output_type": "stream", - "text": [ - "Caching arxiv graph at /tmp/gf_fetch/arxiv.cache\n", - "OGB dependency not available. Downloading graph from CNS.\n" - ] - } - ], - "source": [ - "# Download the Arxiv graph from the OGB repo.\n", - "graph, schema = dgf.io.fetch_ogb_graph(\"arxiv\")" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:42.591860Z", - "iopub.status.busy": "2026-09-23T18:32:42.591647Z", - "iopub.status.idle": "2026-09-23T18:32:42.595939Z", - "shell.execute_reply": "2026-09-23T18:32:42.595387Z" + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Installing GF\n", + "\n", + "Make sure your machine has a GPU or TPU, otherwise training is going to take forever.\n", + "If you are using Google Colab, you can get one for free. Just go to Edit > Notebook settings and select your hardware accelerator as TPU or GPU." + ], + "id": "D1I1ijy08t0v" }, - "executionInfo": { - "elapsed": 5, - "status": "ok", - "timestamp": 1790188362596.8052, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "code", + "execution_count": 2, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:37.856620Z", + "iopub.status.busy": "2026-09-23T18:32:37.856434Z", + "iopub.status.idle": "2026-09-23T18:32:38.101388Z", + "shell.execute_reply": "2026-09-23T18:32:38.100901Z" + }, + "executionInfo": { + "elapsed": 286, + "status": "ok", + "timestamp": 1790188358142.653, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + } + }, + "outputs": [], + "source": [ + "# Install DGF (Distributed Graph Flow) and OGB (for the toy dataset).\n", + "!pip install dgf ogb -U" + ], + "id": "eASOvqWCLNMQ" }, - "id": "d0fb3022", - "outputId": "c177f5ab-8b16-4c6a-fd61-c9fa5cbfcc7e" - }, - "outputs": [ { - "name": "stdout", - "output_type": "stream", - "text": [ - "Graph Schema:\n", - "\n", - "Node Sets:\n", - " nodes:\n", - " | Feature | Format | Semantic | Shape | Detail |\n", - " |-----------|------------|-------------|---------|-------------|\n", - " | #id | BYTES | PRIMARY_ID | None | |\n", - " | #split | BYTES | CATEGORICAL | None | |\n", - " | feat | FLOAT_32 | EMBEDDING | (128,) | |\n", - " | labels | INTEGER_64 | CATEGORICAL | None | #num.cat:40 |\n", - " | year | INTEGER_64 | NUMERICAL | None | |\n", - "\n", - "\n", - "Edge Sets:\n", - " edges: (Source: nodes, Target: nodes)\n", - " (No features)\n", - "\n" - ] - } - ], - "source": [ - "dgf.print.schema(schema)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "2896f2d1" - }, - "source": [ - "## In-process Sampler\n", - "\n", - "GF has an in-memory graph sampler that implements GraphSAGE and other graph\n", - "sampling algorithms.\n", - "\n", - "Let's initialize a sampler." - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:42.597618Z", - "iopub.status.busy": "2026-09-23T18:32:42.597405Z", - "iopub.status.idle": "2026-09-23T18:32:42.801351Z", - "shell.execute_reply": "2026-09-23T18:32:42.800746Z" + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Importing libraries" + ], + "id": "71e9cf1c" }, - "executionInfo": { - "elapsed": 205, - "status": "ok", - "timestamp": 1790188362802.2468, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:38.144052Z", + "iopub.status.busy": "2026-09-23T18:32:38.143850Z", + "iopub.status.idle": "2026-09-23T18:32:42.505988Z", + "shell.execute_reply": "2026-09-23T18:32:42.505456Z" + }, + "executionInfo": { + "elapsed": 4364, + "status": "ok", + "timestamp": 1790188362507.4407, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + } + }, + "outputs": [], + "source": [ + "import copy\n", + "import dgf # Import Graph Flow\n", + "import numpy as np" + ], + "id": "f5431c44" }, - "id": "fba58d46" - }, - "outputs": [], - "source": [ - "# Create a sampler\n", - "sampler_config = dgf.sampling.SimpleSamplingConfig(\n", - " # Start the expansion at the \"nodes\" nodeset.\n", - " # Note: This graph only has one nodeset.\n", - " seed_nodeset=\"nodes\",\n", - " # Maximum distances to consider.\n", - " num_hops=2,\n", - " # How many neighbors we consider at each hop.\n", - " hop_width=2,\n", - " # Follow the edges on both directions.\n", - " reverse=True,\n", - ")\n", - "\n", - "sampler = dgf.sampling.create_sampler(\n", - " graph=graph,\n", - " schema=schema,\n", - " plan=sampler_config,\n", - " num_threads=5,\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "4ee7378e" - }, - "source": [ - "We can now generate and plot a graph sample." - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "metadata": { - "colab": { - "height": 596 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Load some data\n", + "\n", + "Let's start by loading a toy dataset." + ], + "id": "7eb84356" }, - "execution": { - "iopub.execute_input": "2026-09-23T18:32:42.803404Z", - "iopub.status.busy": "2026-09-23T18:32:42.802955Z", - "iopub.status.idle": "2026-09-23T18:32:42.846004Z", - "shell.execute_reply": "2026-09-23T18:32:42.845528Z" + { + "cell_type": "code", + "execution_count": 4, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:42.508824Z", + "iopub.status.busy": "2026-09-23T18:32:42.508200Z", + "iopub.status.idle": "2026-09-23T18:32:42.589533Z", + "shell.execute_reply": "2026-09-23T18:32:42.589133Z" + }, + "executionInfo": { + "elapsed": 83, + "status": "ok", + "timestamp": 1790188362590.8496, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + }, + "outputId": "e216c215-39c3-4e80-a711-5d4c78928767" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Caching arxiv graph at /tmp/gf_fetch/arxiv.cache\n", + "OGB dependency not available. Downloading graph from CNS.\n" + ] + } + ], + "source": [ + "# Download the Arxiv graph from the OGB repo.\n", + "graph, schema = dgf.io.fetch_ogb_graph(\"arxiv\")" + ], + "id": "7f9272a6" }, - "executionInfo": { - "elapsed": 44, - "status": "ok", - "timestamp": 1790188362847.2637, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "code", + "execution_count": 5, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:42.591860Z", + "iopub.status.busy": "2026-09-23T18:32:42.591647Z", + "iopub.status.idle": "2026-09-23T18:32:42.595939Z", + "shell.execute_reply": "2026-09-23T18:32:42.595387Z" + }, + "executionInfo": { + "elapsed": 5, + "status": "ok", + "timestamp": 1790188362596.8052, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + }, + "outputId": "c177f5ab-8b16-4c6a-fd61-c9fa5cbfcc7e" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Graph Schema:\n", + "\n", + "Node Sets:\n", + " nodes:\n", + " | Feature | Format | Semantic | Shape | Detail |\n", + " |-----------|------------|-------------|---------|-------------|\n", + " | #id | BYTES | PRIMARY_ID | None | |\n", + " | #split | BYTES | CATEGORICAL | None | |\n", + " | feat | FLOAT_32 | EMBEDDING | (128,) | |\n", + " | labels | INTEGER_64 | CATEGORICAL | None | #num.cat:40 |\n", + " | year | INTEGER_64 | NUMERICAL | None | |\n", + "\n", + "\n", + "Edge Sets:\n", + " edges: (Source: nodes, Target: nodes)\n", + " (No features)\n", + "\n" + ] + } + ], + "source": [ + "dgf.print.schema(schema)" + ], + "id": "d0fb3022" }, - "id": "daee962f", - "outputId": "b9b56a2a-1070-4cc1-e895-667861003202" - }, - "outputs": [ { - "data": { - "image/svg+xml": [ - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "nodes_0\n", - "\n", - "nodes_0\n", - "\n", - "\n", - "\n", - "nodes_1\n", - "\n", - "nodes_1\n", - "\n", - "\n", - "\n", - "nodes_0->nodes_1\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_6\n", - "\n", - "nodes_6\n", - "\n", - "\n", - "\n", - "nodes_0->nodes_6\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_2\n", - "\n", - "nodes_2\n", - "\n", - "\n", - "\n", - "nodes_1->nodes_2\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_3\n", - "\n", - "nodes_3\n", - "\n", - "\n", - "\n", - "nodes_1->nodes_3\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_4\n", - "\n", - "nodes_4\n", - "\n", - "\n", - "\n", - "nodes_4->nodes_1\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_5\n", - "\n", - "nodes_5\n", - "\n", - "\n", - "\n", - "nodes_5->nodes_1\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_7\n", - "\n", - "nodes_7\n", - "\n", - "\n", - "\n", - "nodes_7->nodes_6\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_8\n", - "\n", - "nodes_8\n", - "\n", - "\n", - "\n", - "nodes_8->nodes_6\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_9\n", - "\n", - "nodes_9\n", - "\n", - "\n", - "\n", - "nodes_9->nodes_0\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_10\n", - "\n", - "nodes_10\n", - "\n", - "\n", - "\n", - "nodes_9->nodes_10\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_11\n", - "\n", - "nodes_11\n", - "\n", - "\n", - "\n", - "nodes_9->nodes_11\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_12\n", - "\n", - "nodes_12\n", - "\n", - "\n", - "\n", - "nodes_12->nodes_9\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_13\n", - "\n", - "nodes_13\n", - "\n", - "\n", - "\n", - "nodes_13->nodes_9\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_14\n", - "\n", - "nodes_14\n", - "\n", - "\n", - "\n", - "nodes_14->nodes_0\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_15\n", - "\n", - "nodes_15\n", - "\n", - "\n", - "\n", - "nodes_14->nodes_15\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_16\n", - "\n", - "nodes_16\n", - "\n", - "\n", - "\n", - "nodes_14->nodes_16\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_17\n", - "\n", - "nodes_17\n", - "\n", - "\n", - "\n", - "nodes_17->nodes_14\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_18\n", - "\n", - "nodes_18\n", - "\n", - "\n", - "\n", - "nodes_18->nodes_14\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n" - ], - "text/plain": [ - "" - ] - }, - "execution_count": 7, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "sample = sampler.sample(seed_node_idxs=0)\n", - "dgf.plot.plot_graph(sample, schema, features=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "a98cd4ee" - }, - "source": [ - "**Remarks:**\n", - "\n", - "- `features=False` removes the features from the plot. It is great for\n", - " visualizing the topology of the graph.\n", - "- In a sample, the first node (`nodes_0`) is always the one the sample was\n", - " generated for (specified with `seed_node_idxs=0`, a.k.a. the seed node).\n", - "- Sampling one seed node at a time is not very efficient; in practice, it is\n", - " better to provide multiple seed nodes, e.g., `seed_node_idxs=[0,1,2]`.\n", - "- Since `seed_node_idxs` is an integer, `sample` returns a single sample. If\n", - " `seed_node_idxs` were a list, `sample` would return a list of samples." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "2d0902f0" - }, - "source": [ - "While `SimpleSamplingConfig` offers basic options, you can use a `SamplingPlan`\n", - "for more granular control (e.g., managing individual hops).\n", - "\n", - "To see what a SamplingPlan looks like, use the\n", - "`simple_sampling_config_to_sampling_plan` method to convert your simple\n", - "configuration into a full plan." - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:42.848313Z", - "iopub.status.busy": "2026-09-23T18:32:42.848026Z", - "iopub.status.idle": "2026-09-23T18:32:42.851406Z", - "shell.execute_reply": "2026-09-23T18:32:42.851045Z" + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## In-process Sampler\n", + "\n", + "GF has an in-memory graph sampler that implements GraphSAGE and other graph\n", + "sampling algorithms.\n", + "\n", + "Let's initialize a sampler." + ], + "id": "2896f2d1" }, - "executionInfo": { - "elapsed": 5, - "status": "ok", - "timestamp": 1790188362852.3203, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "code", + "execution_count": 6, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:42.597618Z", + "iopub.status.busy": "2026-09-23T18:32:42.597405Z", + "iopub.status.idle": "2026-09-23T18:32:42.801351Z", + "shell.execute_reply": "2026-09-23T18:32:42.800746Z" + }, + "executionInfo": { + "elapsed": 205, + "status": "ok", + "timestamp": 1790188362802.2468, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + } + }, + "outputs": [], + "source": [ + "# Create a sampler\n", + "sampler_config = dgf.sampling.SimpleSamplingConfig(\n", + " # Start the expansion at the \"nodes\" nodeset.\n", + " # Note: This graph only has one nodeset.\n", + " seed_nodeset=\"nodes\",\n", + " # Maximum distances to consider.\n", + " num_hops=2,\n", + " # How many neighbors we consider at each hop.\n", + " hop_width=2,\n", + " # Follow the edges on both directions.\n", + " reverse=True,\n", + ")\n", + "\n", + "sampler = dgf.sampling.create_sampler(\n", + " graph=graph,\n", + " schema=schema,\n", + " plan=sampler_config,\n", + " num_threads=5,\n", + ")" + ], + "id": "fba58d46" }, - "id": "5178bcff", - "outputId": "2c017bbd-2d89-4d09-f109-aa0de4d96bb2" - }, - "outputs": [ { - "data": { - "text/plain": [ - "SamplingPlan(root=PlanNode(nodeset='nodes', children=[PlanEdge(edgeset='edges', reversed=False, node=PlanNode(nodeset='nodes', children=[PlanEdge(edgeset='edges', reversed=False, node=PlanNode(nodeset='nodes', children=[]), hop_width=2), PlanEdge(edgeset='edges', reversed=True, node=PlanNode(nodeset='nodes', children=[]), hop_width=2)]), hop_width=2), PlanEdge(edgeset='edges', reversed=True, node=PlanNode(nodeset='nodes', children=[PlanEdge(edgeset='edges', reversed=False, node=PlanNode(nodeset='nodes', children=[]), hop_width=2), PlanEdge(edgeset='edges', reversed=True, node=PlanNode(nodeset='nodes', children=[]), hop_width=2)]), hop_width=2)]), with_replacement=False, temporal_sampling=False, multi_visit=True, max_timeseries_len=32, propagate_timestamp_to_edges=True)" - ] - }, - "execution_count": 8, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "sampling_plan = dgf.sampling.simple_sampling_config_to_sampling_plan(\n", - " sampler_config, schema=schema\n", - ")\n", - "sampling_plan" - ] - }, - { - "cell_type": "code", - "execution_count": 9, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:42.853147Z", - "iopub.status.busy": "2026-09-23T18:32:42.852869Z", - "iopub.status.idle": "2026-09-23T18:32:42.855544Z", - "shell.execute_reply": "2026-09-23T18:32:42.855087Z" + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can now generate and plot a graph sample." + ], + "id": "4ee7378e" }, - "executionInfo": { - "elapsed": 4, - "status": "ok", - "timestamp": 1790188362856.455, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "code", + "execution_count": 7, + "metadata": { + "colab": { + "height": 596 + }, + "execution": { + "iopub.execute_input": "2026-09-23T18:32:42.803404Z", + "iopub.status.busy": "2026-09-23T18:32:42.802955Z", + "iopub.status.idle": "2026-09-23T18:32:42.846004Z", + "shell.execute_reply": "2026-09-23T18:32:42.845528Z" + }, + "executionInfo": { + "elapsed": 44, + "status": "ok", + "timestamp": 1790188362847.2637, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + }, + "outputId": "b9b56a2a-1070-4cc1-e895-667861003202" + }, + "outputs": [ + { + "data": { + "image/svg+xml": [ + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "nodes_0\n", + "\n", + "nodes_0\n", + "\n", + "\n", + "\n", + "nodes_1\n", + "\n", + "nodes_1\n", + "\n", + "\n", + "\n", + "nodes_0->nodes_1\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_6\n", + "\n", + "nodes_6\n", + "\n", + "\n", + "\n", + "nodes_0->nodes_6\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_2\n", + "\n", + "nodes_2\n", + "\n", + "\n", + "\n", + "nodes_1->nodes_2\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_3\n", + "\n", + "nodes_3\n", + "\n", + "\n", + "\n", + "nodes_1->nodes_3\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_4\n", + "\n", + "nodes_4\n", + "\n", + "\n", + "\n", + "nodes_4->nodes_1\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_5\n", + "\n", + "nodes_5\n", + "\n", + "\n", + "\n", + "nodes_5->nodes_1\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_7\n", + "\n", + "nodes_7\n", + "\n", + "\n", + "\n", + "nodes_7->nodes_6\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_8\n", + "\n", + "nodes_8\n", + "\n", + "\n", + "\n", + "nodes_8->nodes_6\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_9\n", + "\n", + "nodes_9\n", + "\n", + "\n", + "\n", + "nodes_9->nodes_0\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_10\n", + "\n", + "nodes_10\n", + "\n", + "\n", + "\n", + "nodes_9->nodes_10\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_11\n", + "\n", + "nodes_11\n", + "\n", + "\n", + "\n", + "nodes_9->nodes_11\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_12\n", + "\n", + "nodes_12\n", + "\n", + "\n", + "\n", + "nodes_12->nodes_9\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_13\n", + "\n", + "nodes_13\n", + "\n", + "\n", + "\n", + "nodes_13->nodes_9\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_14\n", + "\n", + "nodes_14\n", + "\n", + "\n", + "\n", + "nodes_14->nodes_0\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_15\n", + "\n", + "nodes_15\n", + "\n", + "\n", + "\n", + "nodes_14->nodes_15\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_16\n", + "\n", + "nodes_16\n", + "\n", + "\n", + "\n", + "nodes_14->nodes_16\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_17\n", + "\n", + "nodes_17\n", + "\n", + "\n", + "\n", + "nodes_17->nodes_14\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_18\n", + "\n", + "nodes_18\n", + "\n", + "\n", + "\n", + "nodes_18->nodes_14\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sample = sampler.sample(seed_node_idxs=0)\n", + "dgf.plot.plot_graph(sample, schema, features=False)" + ], + "id": "daee962f" }, - "id": "1f9373f2", - "outputId": "77a7549e-8b9c-4bc2-ae41-b7c140f38b3b" - }, - "outputs": [ { - "name": "stdout", - "output_type": "stream", - "text": [ - "Sampling Plan:\n", - "\n", - "Root: nodes\n", - "├── edges [width=2] ➔ nodes\n", - "│ ├── edges [width=2] ➔ nodes\n", - "│ └── edges (reversed) [width=2] ➔ nodes\n", - "└── edges (reversed) [width=2] ➔ nodes\n", - " ├── edges [width=2] ➔ nodes\n", - " └── edges (reversed) [width=2] ➔ nodes\n" - ] - } - ], - "source": [ - "dgf.print.sampling_plan(sampling_plan)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "07fe45e1" - }, - "source": [ - "To be more efficient, let's sample multiple graphs at the same time:" - ] - }, - { - "cell_type": "code", - "execution_count": 10, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:42.857234Z", - "iopub.status.busy": "2026-09-23T18:32:42.857006Z", - "iopub.status.idle": "2026-09-23T18:32:42.864538Z", - "shell.execute_reply": "2026-09-23T18:32:42.864122Z" + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Remarks:**\n", + "\n", + "- `features=False` removes the features from the plot. It is great for\n", + " visualizing the topology of the graph.\n", + "- In a sample, the first node (`nodes_0`) is always the one the sample was\n", + " generated for (specified with `seed_node_idxs=0`, a.k.a. the seed node).\n", + "- Sampling one seed node at a time is not very efficient; in practice, it is\n", + " better to provide multiple seed nodes, e.g., `seed_node_idxs=[0,1,2]`.\n", + "- Since `seed_node_idxs` is an integer, `sample` returns a single sample. If\n", + " `seed_node_idxs` were a list, `sample` would return a list of samples." + ], + "id": "a98cd4ee" }, - "executionInfo": { - "elapsed": 8, - "status": "ok", - "timestamp": 1790188362865.455, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "While `SimpleSamplingConfig` offers basic options, you can use a `SamplingPlan`\n", + "for more granular control (e.g., managing individual hops).\n", + "\n", + "To see what a SamplingPlan looks like, use the\n", + "`simple_sampling_config_to_sampling_plan` method to convert your simple\n", + "configuration into a full plan." + ], + "id": "2d0902f0" }, - "id": "8cea9faa", - "outputId": "e2d60567-92ac-4c42-e04f-519e0481e04d" - }, - "outputs": [ { - "name": "stdout", - "output_type": "stream", - "text": [ - "Number of nodes in the first graph: 15\n", - "Number of nodes in the second graph: 6\n" - ] - } - ], - "source": [ - "samples = sampler.sample(seed_node_idxs=[0, 1])\n", - "print(\n", - " \"Number of nodes in the first graph:\",\n", - " samples[0].node_sets[\"nodes\"].num_nodes,\n", - ")\n", - "print(\n", - " \"Number of nodes in the second graph:\",\n", - " samples[1].node_sets[\"nodes\"].num_nodes,\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "d67499e6" - }, - "source": [ - "Many GF functions use iterators / generators. Let's wrap our sampler into a\n", - "generator and show some of those functions:" - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:42.866345Z", - "iopub.status.busy": "2026-09-23T18:32:42.866157Z", - "iopub.status.idle": "2026-09-23T18:32:42.902122Z", - "shell.execute_reply": "2026-09-23T18:32:42.901663Z" + "cell_type": "code", + "execution_count": 8, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:42.848313Z", + "iopub.status.busy": "2026-09-23T18:32:42.848026Z", + "iopub.status.idle": "2026-09-23T18:32:42.851406Z", + "shell.execute_reply": "2026-09-23T18:32:42.851045Z" + }, + "executionInfo": { + "elapsed": 5, + "status": "ok", + "timestamp": 1790188362852.3203, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + }, + "outputId": "2c017bbd-2d89-4d09-f109-aa0de4d96bb2" + }, + "outputs": [ + { + "data": { + "text/plain": [ + "SamplingPlan(root=PlanNode(nodeset='nodes', children=[PlanEdge(edgeset='edges', reversed=False, node=PlanNode(nodeset='nodes', children=[PlanEdge(edgeset='edges', reversed=False, node=PlanNode(nodeset='nodes', children=[]), hop_width=2), PlanEdge(edgeset='edges', reversed=True, node=PlanNode(nodeset='nodes', children=[]), hop_width=2)]), hop_width=2), PlanEdge(edgeset='edges', reversed=True, node=PlanNode(nodeset='nodes', children=[PlanEdge(edgeset='edges', reversed=False, node=PlanNode(nodeset='nodes', children=[]), hop_width=2), PlanEdge(edgeset='edges', reversed=True, node=PlanNode(nodeset='nodes', children=[]), hop_width=2)]), hop_width=2)]), with_replacement=False, temporal_sampling=False, multi_visit=True, max_timeseries_len=32, propagate_timestamp_to_edges=True)" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sampling_plan = dgf.sampling.simple_sampling_config_to_sampling_plan(\n", + " sampler_config, schema=schema\n", + ")\n", + "sampling_plan" + ], + "id": "5178bcff" }, - "executionInfo": { - "elapsed": 38, - "status": "ok", - "timestamp": 1790188362903.595, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "code", + "execution_count": 9, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:42.853147Z", + "iopub.status.busy": "2026-09-23T18:32:42.852869Z", + "iopub.status.idle": "2026-09-23T18:32:42.855544Z", + "shell.execute_reply": "2026-09-23T18:32:42.855087Z" + }, + "executionInfo": { + "elapsed": 4, + "status": "ok", + "timestamp": 1790188362856.455, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + }, + "outputId": "77a7549e-8b9c-4bc2-ae41-b7c140f38b3b" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Sampling Plan:\n", + "\n", + "Root: nodes\n", + "├── edges [width=2] ➔ nodes\n", + "│ ├── edges [width=2] ➔ nodes\n", + "│ └── edges (reversed) [width=2] ➔ nodes\n", + "└── edges (reversed) [width=2] ➔ nodes\n", + " ├── edges [width=2] ➔ nodes\n", + " └── edges (reversed) [width=2] ➔ nodes\n" + ] + } + ], + "source": [ + "dgf.print.sampling_plan(sampling_plan)" + ], + "id": "1f9373f2" }, - "id": "3a98d6ce", - "outputId": "da141109-12d0-4d1b-d91b-4ba1eabf0946" - }, - "outputs": [ { - "name": "stdout", - "output_type": "stream", - "text": [ - "........................................Done generating\n" - ] - } - ], - "source": [ - "# Create a generator of graph samples.\n", - "def sample_generator(num_batches: int = 10, batch_size: int = 4):\n", - " num_nodes = graph.node_sets[\"nodes\"].num_nodes\n", - " for _ in range(num_batches):\n", - " seed_node_idxs = np.random.choice(num_nodes, size=batch_size, replace=False)\n", - " samples = sampler.sample(seed_node_idxs)\n", - " for sample in samples:\n", - " yield sample\n", - "\n", - "\n", - "# Test the generator.\n", - "for sample in sample_generator():\n", - " print(\".\", end=\"\")\n", - "print(\"Done generating\")" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "31410d56" - }, - "source": [ - "For example, the `dgf.io.write_tfgnn_graphs` function takes a graph generator,\n", - "and saves the values to a TF-GNN Graph record." - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:42.904429Z", - "iopub.status.busy": "2026-09-23T18:32:42.904255Z", - "iopub.status.idle": "2026-09-23T18:32:43.223100Z", - "shell.execute_reply": "2026-09-23T18:32:43.222625Z" + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To be more efficient, let's sample multiple graphs at the same time:" + ], + "id": "07fe45e1" }, - "executionInfo": { - "elapsed": 320, - "status": "ok", - "timestamp": 1790188363224.457, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "code", + "execution_count": 10, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:42.857234Z", + "iopub.status.busy": "2026-09-23T18:32:42.857006Z", + "iopub.status.idle": "2026-09-23T18:32:42.864538Z", + "shell.execute_reply": "2026-09-23T18:32:42.864122Z" + }, + "executionInfo": { + "elapsed": 8, + "status": "ok", + "timestamp": 1790188362865.455, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + }, + "outputId": "e2d60567-92ac-4c42-e04f-519e0481e04d" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Number of nodes in the first graph: 15\n", + "Number of nodes in the second graph: 6\n" + ] + } + ], + "source": [ + "samples = sampler.sample(seed_node_idxs=[0, 1])\n", + "print(\n", + " \"Number of nodes in the first graph:\",\n", + " samples[0].node_sets[\"nodes\"].num_nodes,\n", + ")\n", + "print(\n", + " \"Number of nodes in the second graph:\",\n", + " samples[1].node_sets[\"nodes\"].num_nodes,\n", + ")" + ], + "id": "8cea9faa" }, - "id": "3efc55a0" - }, - "outputs": [], - "source": [ - "# Write the graph samples to disk.\n", - "dgf.io.write_tfgnn_graphs(\n", - " graphs=sample_generator(),\n", - " schema=schema,\n", - " path=\"/tmp/samples_graph.tfrecord.gz\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "c93aafb9" - }, - "source": [ - "When learning GNNs, you will likely batch multiple graphs together. The `merge`\n", - "method takes a list of graphs, and returns a single merged graph." - ] - }, - { - "cell_type": "code", - "execution_count": 13, - "metadata": { - "colab": { - "height": 852 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Many GF functions use iterators / generators. Let's wrap our sampler into a\n", + "generator and show some of those functions:" + ], + "id": "d67499e6" }, - "execution": { - "iopub.execute_input": "2026-09-23T18:32:43.225613Z", - "iopub.status.busy": "2026-09-23T18:32:43.225126Z", - "iopub.status.idle": "2026-09-23T18:32:43.266862Z", - "shell.execute_reply": "2026-09-23T18:32:43.266297Z" + { + "cell_type": "code", + "execution_count": 11, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:42.866345Z", + "iopub.status.busy": "2026-09-23T18:32:42.866157Z", + "iopub.status.idle": "2026-09-23T18:32:42.902122Z", + "shell.execute_reply": "2026-09-23T18:32:42.901663Z" + }, + "executionInfo": { + "elapsed": 38, + "status": "ok", + "timestamp": 1790188362903.595, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + }, + "outputId": "da141109-12d0-4d1b-d91b-4ba1eabf0946" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "........................................Done generating\n" + ] + } + ], + "source": [ + "# Create a generator of graph samples.\n", + "def sample_generator(num_batches: int = 10, batch_size: int = 4):\n", + " num_nodes = graph.node_sets[\"nodes\"].num_nodes\n", + " for _ in range(num_batches):\n", + " seed_node_idxs = np.random.choice(num_nodes, size=batch_size, replace=False)\n", + " samples = sampler.sample(seed_node_idxs)\n", + " for sample in samples:\n", + " yield sample\n", + "\n", + "\n", + "# Test the generator.\n", + "for sample in sample_generator():\n", + " print(\".\", end=\"\")\n", + "print(\"Done generating\")" + ], + "id": "3a98d6ce" }, - "executionInfo": { - "elapsed": 43, - "status": "ok", - "timestamp": 1790188363267.97, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For example, the `dgf.io.write_tfgnn_graphs` function takes a graph generator,\n", + "and saves the values to a TF-GNN Graph record." + ], + "id": "31410d56" }, - "id": "e34b4b65", - "outputId": "38f50aa2-d177-46d6-f3b6-4cdb4a414307" - }, - "outputs": [ { - "data": { - "image/svg+xml": [ - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "nodes_0\n", - "\n", - "nodes_0\n", - "\n", - "\n", - "\n", - "nodes_1\n", - "\n", - "nodes_1\n", - "\n", - "\n", - "\n", - "nodes_0->nodes_1\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_6\n", - "\n", - "nodes_6\n", - "\n", - "\n", - "\n", - "nodes_0->nodes_6\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_2\n", - "\n", - "nodes_2\n", - "\n", - "\n", - "\n", - "nodes_1->nodes_2\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_3\n", - "\n", - "nodes_3\n", - "\n", - "\n", - "\n", - "nodes_1->nodes_3\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_4\n", - "\n", - "nodes_4\n", - "\n", - "\n", - "\n", - "nodes_4->nodes_1\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_5\n", - "\n", - "nodes_5\n", - "\n", - "\n", - "\n", - "nodes_5->nodes_1\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_7\n", - "\n", - "nodes_7\n", - "\n", - "\n", - "\n", - "nodes_7->nodes_6\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_8\n", - "\n", - "nodes_8\n", - "\n", - "\n", - "\n", - "nodes_8->nodes_6\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_9\n", - "\n", - "nodes_9\n", - "\n", - "\n", - "\n", - "nodes_9->nodes_0\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_10\n", - "\n", - "nodes_10\n", - "\n", - "\n", - "\n", - "nodes_9->nodes_10\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_11\n", - "\n", - "nodes_11\n", - "\n", - "\n", - "\n", - "nodes_9->nodes_11\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_12\n", - "\n", - "nodes_12\n", - "\n", - "\n", - "\n", - "nodes_12->nodes_9\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_13\n", - "\n", - "nodes_13\n", - "\n", - "\n", - "\n", - "nodes_13->nodes_9\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_14\n", - "\n", - "nodes_14\n", - "\n", - "\n", - "\n", - "nodes_14->nodes_0\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_15\n", - "\n", - "nodes_15\n", - "\n", - "\n", - "\n", - "nodes_14->nodes_15\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_16\n", - "\n", - "nodes_16\n", - "\n", - "\n", - "\n", - "nodes_16->nodes_14\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_17\n", - "\n", - "nodes_17\n", - "\n", - "\n", - "\n", - "nodes_17->nodes_14\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_18\n", - "\n", - "nodes_18\n", - "\n", - "\n", - "\n", - "nodes_19\n", - "\n", - "nodes_19\n", - "\n", - "\n", - "\n", - "nodes_18->nodes_19\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_20\n", - "\n", - "nodes_20\n", - "\n", - "\n", - "\n", - "nodes_19->nodes_20\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_21\n", - "\n", - "nodes_21\n", - "\n", - "\n", - "\n", - "nodes_21->nodes_19\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_22\n", - "\n", - "nodes_22\n", - "\n", - "\n", - "\n", - "nodes_22->nodes_19\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_23\n", - "\n", - "nodes_23\n", - "\n", - "\n", - "\n", - "nodes_23->nodes_18\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_23->nodes_19\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n" - ], - "text/plain": [ - "" - ] - }, - "execution_count": 13, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "samples = sampler.sample(seed_node_idxs=[0, 1])\n", - "merged_graph, offset = dgf.transform.GraphMerger(schema=schema, padding=None)(samples)\n", - "dgf.plot.plot_graph(merged_graph, schema, features=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "O5j_22f5HJYn" - }, - "source": [ - "The indices of the original graph nodes in the merged graph are available with `offset`:" - ] - }, - { - "cell_type": "code", - "execution_count": 14, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:43.268897Z", - "iopub.status.busy": "2026-09-23T18:32:43.268674Z", - "iopub.status.idle": "2026-09-23T18:32:43.271866Z", - "shell.execute_reply": "2026-09-23T18:32:43.271521Z" + "cell_type": "code", + "execution_count": 12, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:42.904429Z", + "iopub.status.busy": "2026-09-23T18:32:42.904255Z", + "iopub.status.idle": "2026-09-23T18:32:43.223100Z", + "shell.execute_reply": "2026-09-23T18:32:43.222625Z" + }, + "executionInfo": { + "elapsed": 320, + "status": "ok", + "timestamp": 1790188363224.457, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + } + }, + "outputs": [], + "source": [ + "# Write the graph samples to disk.\n", + "dgf.io.write_tfgnn_graphs(\n", + " graphs=sample_generator(),\n", + " schema=schema,\n", + " path=\"/tmp/samples_graph.tfrecord.gz\",\n", + ")" + ], + "id": "3efc55a0" }, - "executionInfo": { - "elapsed": 4, - "status": "ok", - "timestamp": 1790188363272.7253, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "When learning GNNs, you will likely batch multiple graphs together. The `merge`\n", + "method takes a list of graphs, and returns a single merged graph." + ], + "id": "c93aafb9" }, - "id": "LGu-U8SkHTMV" - }, - "outputs": [ { - "data": { - "text/plain": [ - "{'nodes': array([ 0, 18, 24])}" - ] - }, - "execution_count": 14, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "offset" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "Yxpn64EQb_0R" - }, - "source": [ - "**Remarks:**\n", - "\n", - "- The last value of the `offset` is a sentinel reporting the number of nodes. Call `merge_graph` with `sentinel_offset=False` to remove it.\n" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "OrG4j5_OZ4lB" - }, - "source": [ - "## Subgraph extraction\n", - "\n", - "The `sample` method builds graph samples by randomly traversing edges and aggregating all visited edges and nodes. In this section, we will show the `subgraph` method that extracts all the nodes and edges in a certain radius.\n", - "\n", - "By default, `subgraph` returns a single subgraph containing all the nodes and edges at a distance less than or equal to all the seed nodes.\n" - ] - }, - { - "cell_type": "code", - "execution_count": 15, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:43.273623Z", - "iopub.status.busy": "2026-09-23T18:32:43.273293Z", - "iopub.status.idle": "2026-09-23T18:32:43.316332Z", - "shell.execute_reply": "2026-09-23T18:32:43.315862Z" + "cell_type": "code", + "execution_count": 13, + "metadata": { + "colab": { + "height": 852 + }, + "execution": { + "iopub.execute_input": "2026-09-23T18:32:43.225613Z", + "iopub.status.busy": "2026-09-23T18:32:43.225126Z", + "iopub.status.idle": "2026-09-23T18:32:43.266862Z", + "shell.execute_reply": "2026-09-23T18:32:43.266297Z" + }, + "executionInfo": { + "elapsed": 43, + "status": "ok", + "timestamp": 1790188363267.97, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + }, + "outputId": "38f50aa2-d177-46d6-f3b6-4cdb4a414307" + }, + "outputs": [ + { + "data": { + "image/svg+xml": [ + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "nodes_0\n", + "\n", + "nodes_0\n", + "\n", + "\n", + "\n", + "nodes_1\n", + "\n", + "nodes_1\n", + "\n", + "\n", + "\n", + "nodes_0->nodes_1\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_6\n", + "\n", + "nodes_6\n", + "\n", + "\n", + "\n", + "nodes_0->nodes_6\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_2\n", + "\n", + "nodes_2\n", + "\n", + "\n", + "\n", + "nodes_1->nodes_2\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_3\n", + "\n", + "nodes_3\n", + "\n", + "\n", + "\n", + "nodes_1->nodes_3\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_4\n", + "\n", + "nodes_4\n", + "\n", + "\n", + "\n", + "nodes_4->nodes_1\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_5\n", + "\n", + "nodes_5\n", + "\n", + "\n", + "\n", + "nodes_5->nodes_1\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_7\n", + "\n", + "nodes_7\n", + "\n", + "\n", + "\n", + "nodes_7->nodes_6\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_8\n", + "\n", + "nodes_8\n", + "\n", + "\n", + "\n", + "nodes_8->nodes_6\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_9\n", + "\n", + "nodes_9\n", + "\n", + "\n", + "\n", + "nodes_9->nodes_0\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_10\n", + "\n", + "nodes_10\n", + "\n", + "\n", + "\n", + "nodes_9->nodes_10\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_11\n", + "\n", + "nodes_11\n", + "\n", + "\n", + "\n", + "nodes_9->nodes_11\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_12\n", + "\n", + "nodes_12\n", + "\n", + "\n", + "\n", + "nodes_12->nodes_9\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_13\n", + "\n", + "nodes_13\n", + "\n", + "\n", + "\n", + "nodes_13->nodes_9\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_14\n", + "\n", + "nodes_14\n", + "\n", + "\n", + "\n", + "nodes_14->nodes_0\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_15\n", + "\n", + "nodes_15\n", + "\n", + "\n", + "\n", + "nodes_14->nodes_15\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_16\n", + "\n", + "nodes_16\n", + "\n", + "\n", + "\n", + "nodes_16->nodes_14\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_17\n", + "\n", + "nodes_17\n", + "\n", + "\n", + "\n", + "nodes_17->nodes_14\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_18\n", + "\n", + "nodes_18\n", + "\n", + "\n", + "\n", + "nodes_19\n", + "\n", + "nodes_19\n", + "\n", + "\n", + "\n", + "nodes_18->nodes_19\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_20\n", + "\n", + "nodes_20\n", + "\n", + "\n", + "\n", + "nodes_19->nodes_20\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_21\n", + "\n", + "nodes_21\n", + "\n", + "\n", + "\n", + "nodes_21->nodes_19\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_22\n", + "\n", + "nodes_22\n", + "\n", + "\n", + "\n", + "nodes_22->nodes_19\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_23\n", + "\n", + "nodes_23\n", + "\n", + "\n", + "\n", + "nodes_23->nodes_18\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_23->nodes_19\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "samples = sampler.sample(seed_node_idxs=[0, 1])\n", + "merged_graph, offset = dgf.transform.GraphMerger(schema=schema, padding=None)(\n", + " samples\n", + ")\n", + "dgf.plot.plot_graph(merged_graph, schema, features=False)" + ], + "id": "e34b4b65" }, - "executionInfo": { - "elapsed": 44, - "status": "ok", - "timestamp": 1790188363317.5933, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The indices of the original graph nodes in the merged graph are available with `offset`:" + ], + "id": "O5j_22f5HJYn" }, - "id": "grTlK4Lnbvoq" - }, - "outputs": [ { - "data": { - "image/svg+xml": [ - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "nodes_0\n", - "\n", - "nodes_0\n", - "\n", - "\n", - "\n", - "nodes_2\n", - "\n", - "nodes_2\n", - "\n", - "\n", - "\n", - "nodes_0->nodes_2\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_3\n", - "\n", - "nodes_3\n", - "\n", - "\n", - "\n", - "nodes_0->nodes_3\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_1\n", - "\n", - "nodes_1\n", - "\n", - "\n", - "\n", - "nodes_17\n", - "\n", - "nodes_17\n", - "\n", - "\n", - "\n", - "nodes_1->nodes_17\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_22\n", - "\n", - "nodes_22\n", - "\n", - "\n", - "\n", - "nodes_1->nodes_22\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_4\n", - "\n", - "nodes_4\n", - "\n", - "\n", - "\n", - "nodes_3->nodes_4\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_5\n", - "\n", - "nodes_5\n", - "\n", - "\n", - "\n", - "nodes_3->nodes_5\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_6\n", - "\n", - "nodes_6\n", - "\n", - "\n", - "\n", - "nodes_6->nodes_3\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_7\n", - "\n", - "nodes_7\n", - "\n", - "\n", - "\n", - "nodes_7->nodes_3\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_8\n", - "\n", - "nodes_8\n", - "\n", - "\n", - "\n", - "nodes_8->nodes_2\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_9\n", - "\n", - "nodes_9\n", - "\n", - "\n", - "\n", - "nodes_9->nodes_2\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_10\n", - "\n", - "nodes_10\n", - "\n", - "\n", - "\n", - "nodes_10->nodes_0\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_11\n", - "\n", - "nodes_11\n", - "\n", - "\n", - "\n", - "nodes_10->nodes_11\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_12\n", - "\n", - "nodes_12\n", - "\n", - "\n", - "\n", - "nodes_12->nodes_0\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_12->nodes_2\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_13\n", - "\n", - "nodes_13\n", - "\n", - "\n", - "\n", - "nodes_12->nodes_13\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_14\n", - "\n", - "nodes_14\n", - "\n", - "\n", - "\n", - "nodes_12->nodes_14\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_15\n", - "\n", - "nodes_15\n", - "\n", - "\n", - "\n", - "nodes_15->nodes_12\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_16\n", - "\n", - "nodes_16\n", - "\n", - "\n", - "\n", - "nodes_16->nodes_12\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_18\n", - "\n", - "nodes_18\n", - "\n", - "\n", - "\n", - "nodes_17->nodes_18\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_19\n", - "\n", - "nodes_19\n", - "\n", - "\n", - "\n", - "nodes_17->nodes_19\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_20\n", - "\n", - "nodes_20\n", - "\n", - "\n", - "\n", - "nodes_20->nodes_17\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_21\n", - "\n", - "nodes_21\n", - "\n", - "\n", - "\n", - "nodes_21->nodes_17\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_23\n", - "\n", - "nodes_23\n", - "\n", - "\n", - "\n", - "nodes_23->nodes_22\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_24\n", - "\n", - "nodes_24\n", - "\n", - "\n", - "\n", - "nodes_24->nodes_22\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_25\n", - "\n", - "nodes_25\n", - "\n", - "\n", - "\n", - "nodes_25->nodes_1\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_26\n", - "\n", - "nodes_26\n", - "\n", - "\n", - "\n", - "nodes_25->nodes_26\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_27\n", - "\n", - "nodes_27\n", - "\n", - "\n", - "\n", - "nodes_25->nodes_27\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_28\n", - "\n", - "nodes_28\n", - "\n", - "\n", - "\n", - "nodes_28->nodes_2\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_29\n", - "\n", - "nodes_29\n", - "\n", - "\n", - "\n", - "nodes_28->nodes_29\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_30\n", - "\n", - "nodes_30\n", - "\n", - "\n", - "\n", - "nodes_30->nodes_28\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_31\n", - "\n", - "nodes_31\n", - "\n", - "\n", - "\n", - "nodes_31->nodes_28\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n" - ], - "text/plain": [ - "" - ] - }, - "execution_count": 15, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "sample = sampler.subgraph(seed_node_idxs=[0, 150780, 93487])\n", - "dgf.plot.plot_graph(sample, schema, features=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "868fTc8JbkzT" - }, - "source": [ - "**Remarks**:\n", - "\n", - "- Notice that all the nodes are connected. This is because the selected seed nodes (0, 150780, 93487)are close.\n", - "- By default, `subgraph` runs on a single thread. However, it is fully thread-safe: You can call it in parallel using Python's multi-threading.\n", - "\n", - "`multisubgraph` is related to `subgraph`: Instead of returning a single graph, it returns a separate sub-graph around each seed node:\n" - ] - }, - { - "cell_type": "code", - "execution_count": 16, - "metadata": { - "colab": { - "height": 1000 + "cell_type": "code", + "execution_count": 14, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:43.268897Z", + "iopub.status.busy": "2026-09-23T18:32:43.268674Z", + "iopub.status.idle": "2026-09-23T18:32:43.271866Z", + "shell.execute_reply": "2026-09-23T18:32:43.271521Z" + }, + "executionInfo": { + "elapsed": 4, + "status": "ok", + "timestamp": 1790188363272.7253, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "{'nodes': array([ 0, 18, 24])}" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "offset" + ], + "id": "LGu-U8SkHTMV" }, - "execution": { - "iopub.execute_input": "2026-09-23T18:32:43.318562Z", - "iopub.status.busy": "2026-09-23T18:32:43.318280Z", - "iopub.status.idle": "2026-09-23T18:32:43.365244Z", - "shell.execute_reply": "2026-09-23T18:32:43.364671Z" + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Remarks:**\n", + "\n", + "- The last value of the `offset` is a sentinel reporting the number of nodes. Call `merge_graph` with `sentinel_offset=False` to remove it.\n" + ], + "id": "Yxpn64EQb_0R" }, - "executionInfo": { - "elapsed": 48, - "status": "ok", - "timestamp": 1790188363366.272, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Subgraph extraction\n", + "\n", + "The `sample` method builds graph samples by randomly traversing edges and aggregating all visited edges and nodes. In this section, we will show the `subgraph` method that extracts all the nodes and edges in a certain radius.\n", + "\n", + "By default, `subgraph` returns a single subgraph containing all the nodes and edges at a distance less than or equal to all the seed nodes.\n" + ], + "id": "OrG4j5_OZ4lB" }, - "id": "dK98Ep4Mcnya", - "outputId": "a67644aa-b0b5-49d4-d040-59a3263aeca1" - }, - "outputs": [ { - "data": { - "image/svg+xml": [ - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "nodes_0\n", - "\n", - "nodes_0\n", - "\n", - "\n", - "\n", - "nodes_1\n", - "\n", - "nodes_1\n", - "\n", - "\n", - "\n", - "nodes_0->nodes_1\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_6\n", - "\n", - "nodes_6\n", - "\n", - "\n", - "\n", - "nodes_0->nodes_6\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_2\n", - "\n", - "nodes_2\n", - "\n", - "\n", - "\n", - "nodes_1->nodes_2\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_3\n", - "\n", - "nodes_3\n", - "\n", - "\n", - "\n", - "nodes_1->nodes_3\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_4\n", - "\n", - "nodes_4\n", - "\n", - "\n", - "\n", - "nodes_4->nodes_1\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_5\n", - "\n", - "nodes_5\n", - "\n", - "\n", - "\n", - "nodes_5->nodes_1\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_7\n", - "\n", - "nodes_7\n", - "\n", - "\n", - "\n", - "nodes_7->nodes_6\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_8\n", - "\n", - "nodes_8\n", - "\n", - "\n", - "\n", - "nodes_8->nodes_6\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_9\n", - "\n", - "nodes_9\n", - "\n", - "\n", - "\n", - "nodes_9->nodes_0\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_10\n", - "\n", - "nodes_10\n", - "\n", - "\n", - "\n", - "nodes_9->nodes_10\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_11\n", - "\n", - "nodes_11\n", - "\n", - "\n", - "\n", - "nodes_9->nodes_11\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_12\n", - "\n", - "nodes_12\n", - "\n", - "\n", - "\n", - "nodes_12->nodes_9\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_13\n", - "\n", - "nodes_13\n", - "\n", - "\n", - "\n", - "nodes_13->nodes_0\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_14\n", - "\n", - "nodes_14\n", - "\n", - "\n", - "\n", - "nodes_13->nodes_14\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_15\n", - "\n", - "nodes_15\n", - "\n", - "\n", - "\n", - "nodes_13->nodes_15\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_16\n", - "\n", - "nodes_16\n", - "\n", - "\n", - "\n", - "nodes_16->nodes_13\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_17\n", - "\n", - "nodes_17\n", - "\n", - "\n", - "\n", - "nodes_17->nodes_13\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_18\n", - "\n", - "nodes_18\n", - "\n", - "\n", - "\n", - "nodes_19\n", - "\n", - "nodes_19\n", - "\n", - "\n", - "\n", - "nodes_18->nodes_19\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_24\n", - "\n", - "nodes_24\n", - "\n", - "\n", - "\n", - "nodes_18->nodes_24\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_20\n", - "\n", - "nodes_20\n", - "\n", - "\n", - "\n", - "nodes_19->nodes_20\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_21\n", - "\n", - "nodes_21\n", - "\n", - "\n", - "\n", - "nodes_19->nodes_21\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_22\n", - "\n", - "nodes_22\n", - "\n", - "\n", - "\n", - "nodes_22->nodes_19\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_23\n", - "\n", - "nodes_23\n", - "\n", - "\n", - "\n", - "nodes_23->nodes_19\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_25\n", - "\n", - "nodes_25\n", - "\n", - "\n", - "\n", - "nodes_24->nodes_25\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_26\n", - "\n", - "nodes_26\n", - "\n", - "\n", - "\n", - "nodes_24->nodes_26\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_27\n", - "\n", - "nodes_27\n", - "\n", - "\n", - "\n", - "nodes_27->nodes_24\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_28\n", - "\n", - "nodes_28\n", - "\n", - "\n", - "\n", - "nodes_28->nodes_24\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_29\n", - "\n", - "nodes_29\n", - "\n", - "\n", - "\n", - "nodes_29->nodes_18\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_30\n", - "\n", - "nodes_30\n", - "\n", - "\n", - "\n", - "nodes_29->nodes_30\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_31\n", - "\n", - "nodes_31\n", - "\n", - "\n", - "\n", - "nodes_29->nodes_31\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_32\n", - "\n", - "nodes_32\n", - "\n", - "\n", - "\n", - "nodes_33\n", - "\n", - "nodes_33\n", - "\n", - "\n", - "\n", - "nodes_33->nodes_32\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_34\n", - "\n", - "nodes_34\n", - "\n", - "\n", - "\n", - "nodes_33->nodes_34\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_35\n", - "\n", - "nodes_35\n", - "\n", - "\n", - "\n", - "nodes_33->nodes_35\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_36\n", - "\n", - "nodes_36\n", - "\n", - "\n", - "\n", - "nodes_36->nodes_33\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_37\n", - "\n", - "nodes_37\n", - "\n", - "\n", - "\n", - "nodes_37->nodes_33\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_38\n", - "\n", - "nodes_38\n", - "\n", - "\n", - "\n", - "nodes_38->nodes_32\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_39\n", - "\n", - "nodes_39\n", - "\n", - "\n", - "\n", - "nodes_38->nodes_39\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_40\n", - "\n", - "nodes_40\n", - "\n", - "\n", - "\n", - "nodes_38->nodes_40\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_41\n", - "\n", - "nodes_41\n", - "\n", - "\n", - "\n", - "nodes_41->nodes_38\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_42\n", - "\n", - "nodes_42\n", - "\n", - "\n", - "\n", - "nodes_42->nodes_38\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n" - ], - "text/plain": [ - "" - ] - }, - "execution_count": 16, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "samples = sampler.multisubgraph(seed_node_idxs=[0, 150780, 93487])\n", - "\n", - "# Grouping the graphs to create a single plot.\n", - "merge_graph, offset = dgf.transform.GraphMerger(schema=schema, padding=None)(samples)\n", - "dgf.plot.plot_graph(merge_graph, schema, features=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "5c5d4YfWcpar" - }, - "source": [ - "\n", - "**Remark:**\n", - "\n", - "- Notice the 3 independent subgraphs.\n", - "- `multisubgraph` is multi-threaded (one thread per subgraphs)." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "bfbc2230" - }, - "source": [ - "## Temporal masking\n", - "\n", - "Some applications require sampling graphs with temporal masking. This means only\n", - "sampling nodes and edges with a timestamp before a specific value, where the\n", - "limit changes for each sample.\n", - "\n", - "Timestamps should be indicated either on nodes or edges with \"is_creation_time\" equal to true.\n", - "\n", - "**Note:** Edge timestamps are more powerful but less common. Internally, GF converts node timestamps into edge timestamps with the `dgf.transform.propagate_timestamp_to_edges` method.\n", - "\n", - "The following example shows how to create a graph with node timestamps, convert\n", - "them to edge timestamps, and perform temporal sampling.\n", - "\n", - "Let's Arxiv graph contains `year` features on the nodes that we can use as a\n", - "timestamp, and propagate to the edges." - ] - }, - { - "cell_type": "code", - "execution_count": 17, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:43.367494Z", - "iopub.status.busy": "2026-09-23T18:32:43.367068Z", - "iopub.status.idle": "2026-09-23T18:32:43.372803Z", - "shell.execute_reply": "2026-09-23T18:32:43.372352Z" + "cell_type": "code", + "execution_count": 15, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:43.273623Z", + "iopub.status.busy": "2026-09-23T18:32:43.273293Z", + "iopub.status.idle": "2026-09-23T18:32:43.316332Z", + "shell.execute_reply": "2026-09-23T18:32:43.315862Z" + }, + "executionInfo": { + "elapsed": 44, + "status": "ok", + "timestamp": 1790188363317.5933, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + } + }, + "outputs": [ + { + "data": { + "image/svg+xml": [ + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "nodes_0\n", + "\n", + "nodes_0\n", + "\n", + "\n", + "\n", + "nodes_2\n", + "\n", + "nodes_2\n", + "\n", + "\n", + "\n", + "nodes_0->nodes_2\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_3\n", + "\n", + "nodes_3\n", + "\n", + "\n", + "\n", + "nodes_0->nodes_3\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_1\n", + "\n", + "nodes_1\n", + "\n", + "\n", + "\n", + "nodes_17\n", + "\n", + "nodes_17\n", + "\n", + "\n", + "\n", + "nodes_1->nodes_17\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_22\n", + "\n", + "nodes_22\n", + "\n", + "\n", + "\n", + "nodes_1->nodes_22\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_4\n", + "\n", + "nodes_4\n", + "\n", + "\n", + "\n", + "nodes_3->nodes_4\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_5\n", + "\n", + "nodes_5\n", + "\n", + "\n", + "\n", + "nodes_3->nodes_5\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_6\n", + "\n", + "nodes_6\n", + "\n", + "\n", + "\n", + "nodes_6->nodes_3\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_7\n", + "\n", + "nodes_7\n", + "\n", + "\n", + "\n", + "nodes_7->nodes_3\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_8\n", + "\n", + "nodes_8\n", + "\n", + "\n", + "\n", + "nodes_8->nodes_2\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_9\n", + "\n", + "nodes_9\n", + "\n", + "\n", + "\n", + "nodes_9->nodes_2\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_10\n", + "\n", + "nodes_10\n", + "\n", + "\n", + "\n", + "nodes_10->nodes_0\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_11\n", + "\n", + "nodes_11\n", + "\n", + "\n", + "\n", + "nodes_10->nodes_11\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_12\n", + "\n", + "nodes_12\n", + "\n", + "\n", + "\n", + "nodes_12->nodes_0\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_12->nodes_2\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_13\n", + "\n", + "nodes_13\n", + "\n", + "\n", + "\n", + "nodes_12->nodes_13\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_14\n", + "\n", + "nodes_14\n", + "\n", + "\n", + "\n", + "nodes_12->nodes_14\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_15\n", + "\n", + "nodes_15\n", + "\n", + "\n", + "\n", + "nodes_15->nodes_12\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_16\n", + "\n", + "nodes_16\n", + "\n", + "\n", + "\n", + "nodes_16->nodes_12\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_18\n", + "\n", + "nodes_18\n", + "\n", + "\n", + "\n", + "nodes_17->nodes_18\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_19\n", + "\n", + "nodes_19\n", + "\n", + "\n", + "\n", + "nodes_17->nodes_19\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_20\n", + "\n", + "nodes_20\n", + "\n", + "\n", + "\n", + "nodes_20->nodes_17\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_21\n", + "\n", + "nodes_21\n", + "\n", + "\n", + "\n", + "nodes_21->nodes_17\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_23\n", + "\n", + "nodes_23\n", + "\n", + "\n", + "\n", + "nodes_23->nodes_22\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_24\n", + "\n", + "nodes_24\n", + "\n", + "\n", + "\n", + "nodes_24->nodes_22\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_25\n", + "\n", + "nodes_25\n", + "\n", + "\n", + "\n", + "nodes_25->nodes_1\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_26\n", + "\n", + "nodes_26\n", + "\n", + "\n", + "\n", + "nodes_25->nodes_26\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_27\n", + "\n", + "nodes_27\n", + "\n", + "\n", + "\n", + "nodes_25->nodes_27\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_28\n", + "\n", + "nodes_28\n", + "\n", + "\n", + "\n", + "nodes_28->nodes_2\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_29\n", + "\n", + "nodes_29\n", + "\n", + "\n", + "\n", + "nodes_28->nodes_29\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_30\n", + "\n", + "nodes_30\n", + "\n", + "\n", + "\n", + "nodes_30->nodes_28\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_31\n", + "\n", + "nodes_31\n", + "\n", + "\n", + "\n", + "nodes_31->nodes_28\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sample = sampler.subgraph(seed_node_idxs=[0, 150780, 93487])\n", + "dgf.plot.plot_graph(sample, schema, features=False)" + ], + "id": "grTlK4Lnbvoq" }, - "executionInfo": { - "elapsed": 7, - "status": "ok", - "timestamp": 1790188363373.767, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Remarks**:\n", + "\n", + "- Notice that all the nodes are connected. This is because the selected seed nodes (0, 150780, 93487)are close.\n", + "- By default, `subgraph` runs on a single thread. However, it is fully thread-safe: You can call it in parallel using Python's multi-threading.\n", + "\n", + "`multisubgraph` is related to `subgraph`: Instead of returning a single graph, it returns a separate sub-graph around each seed node:\n" + ], + "id": "868fTc8JbkzT" }, - "id": "NvZ9fNQ7LxSs", - "outputId": "f833e47c-3754-44a8-bdb8-f4b0ff6fd1e7" - }, - "outputs": [ { - "name": "stdout", - "output_type": "stream", - "text": [ - "Graph Schema:\n", - "\n", - "Node Sets:\n", - " nodes:\n", - " | Feature | Format | Semantic | Shape | Detail |\n", - " |-----------|------------|-------------|---------|-------------|\n", - " | #id | BYTES | PRIMARY_ID | None | |\n", - " | #split | BYTES | CATEGORICAL | None | |\n", - " | feat | FLOAT_32 | EMBEDDING | (128,) | |\n", - " | labels | INTEGER_64 | CATEGORICAL | None | #num.cat:40 |\n", - " | year | INTEGER_64 | NUMERICAL | None | creation |\n", - "\n", - "\n", - "Edge Sets:\n", - " edges: (Source: nodes, Target: nodes)\n", - " (No features)\n", - "\n" - ] - } - ], - "source": [ - "time_schema = copy.deepcopy(schema)\n", - "time_schema.node_sets[\"nodes\"].features[\"year\"].is_creation_time = True\n", - "dgf.print.schema(time_schema)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "5ad2a4f9" - }, - "source": [ - "The following example shows how to create a graph with node timestamps, convert\n", - "them to edge timestamps, and perform temporal sampling." - ] - }, - { - "cell_type": "code", - "execution_count": 18, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:43.374679Z", - "iopub.status.busy": "2026-09-23T18:32:43.374451Z", - "iopub.status.idle": "2026-09-23T18:32:43.627765Z", - "shell.execute_reply": "2026-09-23T18:32:43.627273Z" + "cell_type": "code", + "execution_count": 16, + "metadata": { + "colab": { + "height": 1000 + }, + "execution": { + "iopub.execute_input": "2026-09-23T18:32:43.318562Z", + "iopub.status.busy": "2026-09-23T18:32:43.318280Z", + "iopub.status.idle": "2026-09-23T18:32:43.365244Z", + "shell.execute_reply": "2026-09-23T18:32:43.364671Z" + }, + "executionInfo": { + "elapsed": 48, + "status": "ok", + "timestamp": 1790188363366.272, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + }, + "outputId": "a67644aa-b0b5-49d4-d040-59a3263aeca1" + }, + "outputs": [ + { + "data": { + "image/svg+xml": [ + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "nodes_0\n", + "\n", + "nodes_0\n", + "\n", + "\n", + "\n", + "nodes_1\n", + "\n", + "nodes_1\n", + "\n", + "\n", + "\n", + "nodes_0->nodes_1\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_6\n", + "\n", + "nodes_6\n", + "\n", + "\n", + "\n", + "nodes_0->nodes_6\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_2\n", + "\n", + "nodes_2\n", + "\n", + "\n", + "\n", + "nodes_1->nodes_2\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_3\n", + "\n", + "nodes_3\n", + "\n", + "\n", + "\n", + "nodes_1->nodes_3\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_4\n", + "\n", + "nodes_4\n", + "\n", + "\n", + "\n", + "nodes_4->nodes_1\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_5\n", + "\n", + "nodes_5\n", + "\n", + "\n", + "\n", + "nodes_5->nodes_1\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_7\n", + "\n", + "nodes_7\n", + "\n", + "\n", + "\n", + "nodes_7->nodes_6\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_8\n", + "\n", + "nodes_8\n", + "\n", + "\n", + "\n", + "nodes_8->nodes_6\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_9\n", + "\n", + "nodes_9\n", + "\n", + "\n", + "\n", + "nodes_9->nodes_0\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_10\n", + "\n", + "nodes_10\n", + "\n", + "\n", + "\n", + "nodes_9->nodes_10\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_11\n", + "\n", + "nodes_11\n", + "\n", + "\n", + "\n", + "nodes_9->nodes_11\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_12\n", + "\n", + "nodes_12\n", + "\n", + "\n", + "\n", + "nodes_12->nodes_9\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_13\n", + "\n", + "nodes_13\n", + "\n", + "\n", + "\n", + "nodes_13->nodes_0\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_14\n", + "\n", + "nodes_14\n", + "\n", + "\n", + "\n", + "nodes_13->nodes_14\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_15\n", + "\n", + "nodes_15\n", + "\n", + "\n", + "\n", + "nodes_13->nodes_15\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_16\n", + "\n", + "nodes_16\n", + "\n", + "\n", + "\n", + "nodes_16->nodes_13\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_17\n", + "\n", + "nodes_17\n", + "\n", + "\n", + "\n", + "nodes_17->nodes_13\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_18\n", + "\n", + "nodes_18\n", + "\n", + "\n", + "\n", + "nodes_19\n", + "\n", + "nodes_19\n", + "\n", + "\n", + "\n", + "nodes_18->nodes_19\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_24\n", + "\n", + "nodes_24\n", + "\n", + "\n", + "\n", + "nodes_18->nodes_24\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_20\n", + "\n", + "nodes_20\n", + "\n", + "\n", + "\n", + "nodes_19->nodes_20\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_21\n", + "\n", + "nodes_21\n", + "\n", + "\n", + "\n", + "nodes_19->nodes_21\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_22\n", + "\n", + "nodes_22\n", + "\n", + "\n", + "\n", + "nodes_22->nodes_19\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_23\n", + "\n", + "nodes_23\n", + "\n", + "\n", + "\n", + "nodes_23->nodes_19\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_25\n", + "\n", + "nodes_25\n", + "\n", + "\n", + "\n", + "nodes_24->nodes_25\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_26\n", + "\n", + "nodes_26\n", + "\n", + "\n", + "\n", + "nodes_24->nodes_26\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_27\n", + "\n", + "nodes_27\n", + "\n", + "\n", + "\n", + "nodes_27->nodes_24\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_28\n", + "\n", + "nodes_28\n", + "\n", + "\n", + "\n", + "nodes_28->nodes_24\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_29\n", + "\n", + "nodes_29\n", + "\n", + "\n", + "\n", + "nodes_29->nodes_18\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_30\n", + "\n", + "nodes_30\n", + "\n", + "\n", + "\n", + "nodes_29->nodes_30\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_31\n", + "\n", + "nodes_31\n", + "\n", + "\n", + "\n", + "nodes_29->nodes_31\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_32\n", + "\n", + "nodes_32\n", + "\n", + "\n", + "\n", + "nodes_33\n", + "\n", + "nodes_33\n", + "\n", + "\n", + "\n", + "nodes_33->nodes_32\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_34\n", + "\n", + "nodes_34\n", + "\n", + "\n", + "\n", + "nodes_33->nodes_34\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_35\n", + "\n", + "nodes_35\n", + "\n", + "\n", + "\n", + "nodes_33->nodes_35\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_36\n", + "\n", + "nodes_36\n", + "\n", + "\n", + "\n", + "nodes_36->nodes_33\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_37\n", + "\n", + "nodes_37\n", + "\n", + "\n", + "\n", + "nodes_37->nodes_33\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_38\n", + "\n", + "nodes_38\n", + "\n", + "\n", + "\n", + "nodes_38->nodes_32\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_39\n", + "\n", + "nodes_39\n", + "\n", + "\n", + "\n", + "nodes_38->nodes_39\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_40\n", + "\n", + "nodes_40\n", + "\n", + "\n", + "\n", + "nodes_38->nodes_40\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_41\n", + "\n", + "nodes_41\n", + "\n", + "\n", + "\n", + "nodes_41->nodes_38\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_42\n", + "\n", + "nodes_42\n", + "\n", + "\n", + "\n", + "nodes_42->nodes_38\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "samples = sampler.multisubgraph(seed_node_idxs=[0, 150780, 93487])\n", + "\n", + "# Grouping the graphs to create a single plot.\n", + "merge_graph, offset = dgf.transform.GraphMerger(schema=schema, padding=None)(\n", + " samples\n", + ")\n", + "dgf.plot.plot_graph(merge_graph, schema, features=False)" + ], + "id": "dK98Ep4Mcnya" }, - "executionInfo": { - "elapsed": 255, - "status": "ok", - "timestamp": 1790188363628.9556, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "**Remark:**\n", + "\n", + "- Notice the 3 independent subgraphs.\n", + "- `multisubgraph` is multi-threaded (one thread per subgraphs)." + ], + "id": "5c5d4YfWcpar" }, - "id": "c6d421dc" - }, - "outputs": [], - "source": [ - "# Create a sampler\n", - "time_sampler_config = dgf.sampling.SimpleSamplingConfig(\n", - " seed_nodeset=\"nodes\",\n", - " num_hops=2,\n", - " hop_width=2,\n", - " reverse=True,\n", - " # This is the new part: the creation time of the nodes and of the edges is\n", - " # inferred from the \"is_creation_time\" features of the schema.\n", - " temporal_sampling=True,\n", - ")\n", - "\n", - "time_sampler = dgf.sampling.create_sampler(\n", - " graph=graph,\n", - " schema=time_schema,\n", - " plan=time_sampler_config,\n", - " num_threads=5,\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "909d6d86" - }, - "source": [ - "Let's look at the year of the first node:" - ] - }, - { - "cell_type": "code", - "execution_count": 19, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:43.629903Z", - "iopub.status.busy": "2026-09-23T18:32:43.629602Z", - "iopub.status.idle": "2026-09-23T18:32:43.632730Z", - "shell.execute_reply": "2026-09-23T18:32:43.632437Z" + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Temporal masking\n", + "\n", + "Some applications require sampling graphs with temporal masking. This means only\n", + "sampling nodes and edges with a timestamp before a specific value, where the\n", + "limit changes for each sample.\n", + "\n", + "Timestamps should be indicated either on nodes or edges with \"is_creation_time\" equal to true.\n", + "\n", + "**Note:** Edge timestamps are more powerful but less common. Internally, GF converts node timestamps into edge timestamps with the `dgf.transform.propagate_timestamp_to_edges` method.\n", + "\n", + "The following example shows how to create a graph with node timestamps, convert\n", + "them to edge timestamps, and perform temporal sampling.\n", + "\n", + "Let's Arxiv graph contains `year` features on the nodes that we can use as a\n", + "timestamp, and propagate to the edges." + ], + "id": "bfbc2230" }, - "executionInfo": { - "elapsed": 4, - "status": "ok", - "timestamp": 1790188363633.6614, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "code", + "execution_count": 17, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:43.367494Z", + "iopub.status.busy": "2026-09-23T18:32:43.367068Z", + "iopub.status.idle": "2026-09-23T18:32:43.372803Z", + "shell.execute_reply": "2026-09-23T18:32:43.372352Z" + }, + "executionInfo": { + "elapsed": 7, + "status": "ok", + "timestamp": 1790188363373.767, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + }, + "outputId": "f833e47c-3754-44a8-bdb8-f4b0ff6fd1e7" + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Graph Schema:\n", + "\n", + "Node Sets:\n", + " nodes:\n", + " | Feature | Format | Semantic | Shape | Detail |\n", + " |-----------|------------|-------------|---------|-------------|\n", + " | #id | BYTES | PRIMARY_ID | None | |\n", + " | #split | BYTES | CATEGORICAL | None | |\n", + " | feat | FLOAT_32 | EMBEDDING | (128,) | |\n", + " | labels | INTEGER_64 | CATEGORICAL | None | #num.cat:40 |\n", + " | year | INTEGER_64 | NUMERICAL | None | creation |\n", + "\n", + "\n", + "Edge Sets:\n", + " edges: (Source: nodes, Target: nodes)\n", + " (No features)\n", + "\n" + ] + } + ], + "source": [ + "time_schema = copy.deepcopy(schema)\n", + "time_schema.node_sets[\"nodes\"].features[\"year\"].is_creation_time = True\n", + "dgf.print.schema(time_schema)" + ], + "id": "NvZ9fNQ7LxSs" }, - "id": "91611d90", - "outputId": "055c8769-b8dd-4897-aba5-1d9acea0d769" - }, - "outputs": [ { - "data": { - "text/plain": [ - "np.int64(2013)" - ] - }, - "execution_count": 19, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "graph.node_sets[\"nodes\"].features[\"year\"][0]" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "5cb0d579" - }, - "source": [ - "So, let's create a graph sample by only considering edges prior (non-strict)\n", - "to 2013." - ] - }, - { - "cell_type": "code", - "execution_count": 20, - "metadata": { - "colab": { - "height": 395 + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The following example shows how to create a graph with node timestamps, convert\n", + "them to edge timestamps, and perform temporal sampling." + ], + "id": "5ad2a4f9" }, - "execution": { - "iopub.execute_input": "2026-09-23T18:32:43.634340Z", - "iopub.status.busy": "2026-09-23T18:32:43.634171Z", - "iopub.status.idle": "2026-09-23T18:32:43.678957Z", - "shell.execute_reply": "2026-09-23T18:32:43.678428Z" + { + "cell_type": "code", + "execution_count": 18, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:43.374679Z", + "iopub.status.busy": "2026-09-23T18:32:43.374451Z", + "iopub.status.idle": "2026-09-23T18:32:43.627765Z", + "shell.execute_reply": "2026-09-23T18:32:43.627273Z" + }, + "executionInfo": { + "elapsed": 255, + "status": "ok", + "timestamp": 1790188363628.9556, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + } + }, + "outputs": [], + "source": [ + "# Create a sampler\n", + "time_sampler_config = dgf.sampling.SimpleSamplingConfig(\n", + " seed_nodeset=\"nodes\",\n", + " num_hops=2,\n", + " hop_width=2,\n", + " reverse=True,\n", + " # This is the new part: the creation time of the nodes and of the edges is\n", + " # inferred from the \"is_creation_time\" features of the schema.\n", + " temporal_sampling=True,\n", + ")\n", + "\n", + "time_sampler = dgf.sampling.create_sampler(\n", + " graph=graph,\n", + " schema=time_schema,\n", + " plan=time_sampler_config,\n", + " num_threads=5,\n", + ")" + ], + "id": "c6d421dc" }, - "executionInfo": { - "elapsed": 46, - "status": "ok", - "timestamp": 1790188363680.2434, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Let's look at the year of the first node:" + ], + "id": "909d6d86" }, - "id": "a35660d2", - "outputId": "16028134-3b41-4be1-db5d-a239de4912d0" - }, - "outputs": [ { - "data": { - "image/svg+xml": [ - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "nodes_0\n", - "\n", - "nodes_0\n", - "#id: b'n0'\n", - "#split: b'train'\n", - "feat: [-0.057943 -0.05253  -0.072603 -0.026555  0.130435[...]\n", - "labels: 4\n", - "year: 2013\n", - "\n", - "\n", - "\n", - "nodes_1\n", - "\n", - "nodes_1\n", - "#id: b'n52893'\n", - "#split: b'train'\n", - "feat: [-0.055733 -0.031606 -0.292581 -0.054655  0.088294[...]\n", - "labels: 24\n", - "year: 2010\n", - "\n", - "\n", - "\n", - "nodes_0->nodes_1\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_4\n", - "\n", - "nodes_4\n", - "#id: b'n93487'\n", - "#split: b'train'\n", - "feat: [-0.215197  0.051735 -0.050193 -0.02174   0.12137 [...]\n", - "labels: 24\n", - "year: 2012\n", - "\n", - "\n", - "\n", - "nodes_0->nodes_4\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_2\n", - "\n", - "nodes_2\n", - "#id: b'n14528'\n", - "#split: b'train'\n", - "feat: [ 6.38240e-02 -4.81730e-02 -2.85770e-01 -1.85823e-[...]\n", - "labels: 24\n", - "year: 2010\n", - "\n", - "\n", - "\n", - "nodes_1->nodes_2\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_3\n", - "\n", - "nodes_3\n", - "#id: b'n71730'\n", - "#split: b'train'\n", - "feat: [-0.115429  0.011448 -0.260941  0.03005   0.145451[...]\n", - "labels: 24\n", - "year: 2009\n", - "\n", - "\n", - "\n", - "nodes_1->nodes_3\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n", - "nodes_5\n", - "\n", - "nodes_5\n", - "#id: b'n67258'\n", - "#split: b'train'\n", - "feat: [-0.149148 -0.08377  -0.191802  0.092441 -0.005551[...]\n", - "labels: 24\n", - "year: 2013\n", - "\n", - "\n", - "\n", - "nodes_5->nodes_4\n", - "\n", - "\n", - "edges\n", - "\n", - "\n", - "\n" - ], - "text/plain": [ - "" - ] - }, - "execution_count": 20, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "sample = time_sampler.sample(seed_node_idxs=0, seed_timestamps=2013)\n", - "dgf.plot.plot_graph(sample, time_schema, features=True)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "ef4b8f03" - }, - "source": [ - "## Semi-distributed sampler (v2)\n", - "\n", - "The semi-distributed sampler is a simple way to distribute the computation of\n", - "graph samples over multiple machines using Apache Beam. It is not a fully\n", - "distributed sampler though: Each worker will load the full graph topology (i.e.,\n", - "the edges; but not the features) in memory.\n", - "\n", - "**Note:** For a full example, check\n", - "`dgf/examples/create_graph_samples_semi_distributed_v2.py`.\n", - "\n", - "Let's configure and run the semi-distributed graph sampler:" - ] - }, - { - "cell_type": "code", - "execution_count": 21, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:43.681333Z", - "iopub.status.busy": "2026-09-23T18:32:43.681009Z", - "iopub.status.idle": "2026-09-23T18:32:43.906957Z", - "shell.execute_reply": "2026-09-23T18:32:43.906597Z" + "cell_type": "code", + "execution_count": 19, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:43.629903Z", + "iopub.status.busy": "2026-09-23T18:32:43.629602Z", + "iopub.status.idle": "2026-09-23T18:32:43.632730Z", + "shell.execute_reply": "2026-09-23T18:32:43.632437Z" + }, + "executionInfo": { + "elapsed": 4, + "status": "ok", + "timestamp": 1790188363633.6614, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + }, + "outputId": "055c8769-b8dd-4897-aba5-1d9acea0d769" + }, + "outputs": [ + { + "data": { + "text/plain": [ + "np.int64(2013)" + ] + }, + "execution_count": 19, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "graph.node_sets[\"nodes\"].features[\"year\"][0]" + ], + "id": "91611d90" }, - "executionInfo": { - "elapsed": 227, - "status": "ok", - "timestamp": 1790188363908.234, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "So, let's create a graph sample by only considering edges prior (non-strict)\n", + "to 2013." + ], + "id": "5cb0d579" }, - "id": "4151f762" - }, - "outputs": [], - "source": [ - "from absl import flags\n", - "import apache_beam as beam\n", - "from apache_beam.options import pipeline_options\n", - "from google3.pipeline.flume.py import runner as flume_runner\n", - "\n", - "# For the user: Comment / uncomment on of the block.\n", - "# Note: Defining the \"flume_exec_mode\" in the \"PipelineOptions\" does not work.\n", - "\n", - "# Run the execution in-process. Great for debugging / iteration on small data.\n", - "# ===\n", - "flags.FLAGS.flume_exec_mode = \"IN_PROCESS\"\n", - "options = pipeline_options.PipelineOptions()\n", - "\n", - "# Run the execution on Borg. Great for large data.\n", - "# ===\n", - "# flags.FLAGS.flume_exec_mode = \"BORG\"\n", - "# options = pipeline_options.PipelineOptions(\n", - "# flume_borg_accounting_charged_user_name=\"simple-ml-accounting\",\n", - "# flume_borg_cells=\"is\",\n", - "# flume_use_batch_scheduler=True,\n", - "# flume_batch_scheduler_strategy=\"RUN_SOON\",\n", - "# )" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "71fc0067" - }, - "source": [ - "Save our graph to disk. This will be the input of the sampler." - ] - }, - { - "cell_type": "code", - "execution_count": 22, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:43.909135Z", - "iopub.status.busy": "2026-09-23T18:32:43.908954Z", - "iopub.status.idle": "2026-09-23T18:32:44.288082Z", - "shell.execute_reply": "2026-09-23T18:32:44.287545Z" + { + "cell_type": "code", + "execution_count": 20, + "metadata": { + "colab": { + "height": 395 + }, + "execution": { + "iopub.execute_input": "2026-09-23T18:32:43.634340Z", + "iopub.status.busy": "2026-09-23T18:32:43.634171Z", + "iopub.status.idle": "2026-09-23T18:32:43.678957Z", + "shell.execute_reply": "2026-09-23T18:32:43.678428Z" + }, + "executionInfo": { + "elapsed": 46, + "status": "ok", + "timestamp": 1790188363680.2434, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + }, + "outputId": "16028134-3b41-4be1-db5d-a239de4912d0" + }, + "outputs": [ + { + "data": { + "image/svg+xml": [ + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "nodes_0\n", + "\n", + "nodes_0\n", + "#id: b'n0'\n", + "#split: b'train'\n", + "feat: [-0.057943 -0.05253  -0.072603 -0.026555  0.130435[...]\n", + "labels: 4\n", + "year: 2013\n", + "\n", + "\n", + "\n", + "nodes_1\n", + "\n", + "nodes_1\n", + "#id: b'n52893'\n", + "#split: b'train'\n", + "feat: [-0.055733 -0.031606 -0.292581 -0.054655  0.088294[...]\n", + "labels: 24\n", + "year: 2010\n", + "\n", + "\n", + "\n", + "nodes_0->nodes_1\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_4\n", + "\n", + "nodes_4\n", + "#id: b'n93487'\n", + "#split: b'train'\n", + "feat: [-0.215197  0.051735 -0.050193 -0.02174   0.12137 [...]\n", + "labels: 24\n", + "year: 2012\n", + "\n", + "\n", + "\n", + "nodes_0->nodes_4\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_2\n", + "\n", + "nodes_2\n", + "#id: b'n14528'\n", + "#split: b'train'\n", + "feat: [ 6.38240e-02 -4.81730e-02 -2.85770e-01 -1.85823e-[...]\n", + "labels: 24\n", + "year: 2010\n", + "\n", + "\n", + "\n", + "nodes_1->nodes_2\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_3\n", + "\n", + "nodes_3\n", + "#id: b'n71730'\n", + "#split: b'train'\n", + "feat: [-0.115429  0.011448 -0.260941  0.03005   0.145451[...]\n", + "labels: 24\n", + "year: 2009\n", + "\n", + "\n", + "\n", + "nodes_1->nodes_3\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n", + "nodes_5\n", + "\n", + "nodes_5\n", + "#id: b'n67258'\n", + "#split: b'train'\n", + "feat: [-0.149148 -0.08377  -0.191802  0.092441 -0.005551[...]\n", + "labels: 24\n", + "year: 2013\n", + "\n", + "\n", + "\n", + "nodes_5->nodes_4\n", + "\n", + "\n", + "edges\n", + "\n", + "\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "execution_count": 20, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sample = time_sampler.sample(seed_node_idxs=0, seed_timestamps=2013)\n", + "dgf.plot.plot_graph(sample, time_schema, features=True)" + ], + "id": "a35660d2" }, - "executionInfo": { - "elapsed": 380, - "status": "ok", - "timestamp": 1790188364289.1096, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Semi-distributed sampler (v2)\n", + "\n", + "The semi-distributed sampler is a simple way to distribute the computation of\n", + "graph samples over multiple machines using Apache Beam. It is not a fully\n", + "distributed sampler though: Each worker will load the full graph topology (i.e.,\n", + "the edges; but not the features) in memory.\n", + "\n", + "**Note:** For a full example, check\n", + "`dgf/examples/create_graph_samples_semi_distributed_v2.py`.\n", + "\n", + "Let's configure and run the semi-distributed graph sampler:" + ], + "id": "ef4b8f03" }, - "id": "f4a3e3de" - }, - "outputs": [ { - "name": "stdout", - "output_type": "stream", - "text": [ - "Graph written from memory in 0.38s\n" - ] - } - ], - "source": [ - "graph_path = \"/tmp/my_graph\"\n", - "dgf.io.write_graph(graph, schema, path=graph_path)" - ] - }, - { - "cell_type": "code", - "execution_count": 23, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:32:44.289991Z", - "iopub.status.busy": "2026-09-23T18:32:44.289809Z", - "iopub.status.idle": "2026-09-23T18:34:42.743424Z", - "shell.execute_reply": "2026-09-23T18:34:42.742728Z" + "cell_type": "code", + "execution_count": 21, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:43.681333Z", + "iopub.status.busy": "2026-09-23T18:32:43.681009Z", + "iopub.status.idle": "2026-09-23T18:32:43.906957Z", + "shell.execute_reply": "2026-09-23T18:32:43.906597Z" + }, + "executionInfo": { + "elapsed": 227, + "status": "ok", + "timestamp": 1790188363908.234, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + } + }, + "outputs": [], + "source": [ + "from absl import flags\n", + "import apache_beam as beam\n", + "from apache_beam.options import pipeline_options\n", + "from google3.pipeline.flume.py import runner as flume_runner\n", + "\n", + "# For the user: Comment / uncomment on of the block.\n", + "# Note: Defining the \"flume_exec_mode\" in the \"PipelineOptions\" does not work.\n", + "\n", + "# Run the execution in-process. Great for debugging / iteration on small data.\n", + "# ===\n", + "flags.FLAGS.flume_exec_mode = \"IN_PROCESS\"\n", + "options = pipeline_options.PipelineOptions()\n", + "\n", + "# Run the execution on Borg. Great for large data.\n", + "# ===\n", + "# flags.FLAGS.flume_exec_mode = \"BORG\"\n", + "# options = pipeline_options.PipelineOptions(\n", + "# flume_borg_accounting_charged_user_name=\"simple-ml-accounting\",\n", + "# flume_borg_cells=\"is\",\n", + "# flume_use_batch_scheduler=True,\n", + "# flume_batch_scheduler_strategy=\"RUN_SOON\",\n", + "# )" + ], + "id": "4151f762" }, - "executionInfo": { - "elapsed": 118455, - "status": "ok", - "timestamp": 1790188482744.6162, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Save our graph to disk. This will be the input of the sampler." + ], + "id": "71fc0067" }, - "id": "ecbaa932" - }, - "outputs": [ { - "data": { - "application/javascript": [ - "\n", - " if (typeof window.interactive_beam_jquery == 'undefined') {\n", - " var jqueryScript = document.createElement('script');\n", - " jqueryScript.src = 'https://code.jquery.com/jquery-3.4.1.slim.min.js';\n", - " jqueryScript.type = 'text/javascript';\n", - " jqueryScript.onload = function() {\n", - " var datatableScript = document.createElement('script');\n", - " datatableScript.src = 'https://cdn.datatables.net/1.10.20/js/jquery.dataTables.min.js';\n", - " datatableScript.type = 'text/javascript';\n", - " datatableScript.onload = function() {\n", - " window.interactive_beam_jquery = jQuery.noConflict(true);\n", - " window.interactive_beam_jquery(document).ready(function($){\n", - " \n", - " });\n", - " }\n", - " document.head.appendChild(datatableScript);\n", - " };\n", - " document.head.appendChild(jqueryScript);\n", - " } else {\n", - " window.interactive_beam_jquery(document).ready(function($){\n", - " \n", - " });\n", - " }" - ] - }, - "metadata": {}, - "output_type": "display_data" + "cell_type": "code", + "execution_count": 22, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:43.909135Z", + "iopub.status.busy": "2026-09-23T18:32:43.908954Z", + "iopub.status.idle": "2026-09-23T18:32:44.288082Z", + "shell.execute_reply": "2026-09-23T18:32:44.287545Z" + }, + "executionInfo": { + "elapsed": 380, + "status": "ok", + "timestamp": 1790188364289.1096, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Graph written from memory in 0.38s\n" + ] + } + ], + "source": [ + "graph_path = \"/tmp/my_graph\"\n", + "dgf.io.write_graph(graph, schema, path=graph_path)" + ], + "id": "f4a3e3de" }, { - "name": "stdout", - "output_type": "stream", - "text": [ - "Using override schema\n", - "Reading 1 nodeset(s), 1 edgeset(s), and 5 feature(s)\n", - "Reading metadata from /tmp/my_graph\n", - "Reading nodeset nodes from /tmp/my_graph\n", - ".concatenating values\n", - "Reading edgeset edges from /tmp/my_graph\n", - ".concatenating values\n", - "Graph read in memory in 0.52s\n" - ] - } - ], - "source": [ - "with beam.Pipeline(runner=flume_runner.FlumeRunner(), options=options) as root:\n", - "\n", - " # Read only the node and the #id feature from the graph. Those\n", - " # is the nodes to seed.\n", - " seed_graph = dgf.beam.io.read_graph(\n", - " root,\n", - " graph_path,\n", - " schema_filter=dgf.data.GraphSchemaFilter(\n", - " nodeset_fn=lambda key, sch: key == \"nodes\",\n", - " edgeset_fn=lambda key, sch: False,\n", - " feature_fn=lambda key, sch: key == \"#id\",\n", - " ),\n", - " )\n", - " seed_node_ids = dgf.beam.sampling.extract_nodes_ids(seed_graph, \"nodes\")\n", - "\n", - " # Alternatively, you can do:\n", - " # seed_node_ids = beam.Create([b\"\", b\"\"])\n", - "\n", - " # Randomly select 50 seed nodes.\n", - " seed_node_ids = (\n", - " seed_node_ids\n", - " | \"Sample seeds\" >> beam.combiners.Sample.FixedSizeGlobally(50)\n", - " | beam.FlatMap(lambda xs: xs)\n", - " )\n", - "\n", - " # Generate samples\n", - " samples, output_schema = dgf.beam.sampling.semi_distributed_sampler_v2(\n", - " graph_path=graph_path,\n", - " plan=sampler_config,\n", - " seeds=seed_node_ids,\n", - " num_threads=20,\n", - " beam_feature_collection=False,\n", - " )\n", - "\n", - " # Save the samples to disk\n", - " dgf.beam.io.write_tfgnn_graphs(samples, \"/tmp/samples@*\", output_schema)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "70dcd856" - }, - "source": [ - "**Remark:**\n", - "\n", - "- The semi-distributed sampler (v2) takes as input a graph stored on disk. It\n", - " is today, the fastest option (faster than loading the input graph from a\n", - " beam pcollection like the v1).\n", - "- The `num_threads` option control the number of threads for each worker.\n", - "- Setting `beam_feature_collection=True` tells the sampler to generate\n", - " topologies first and then collect feature values using Apache Beam. This\n", - " approach scales effectively because workers only store the graph topology in\n", - " memory, though it results in slower overall performance. Alternatively,\n", - " `beam_feature_collection=False` triggers feature collection at the same time\n", - " the sampler generates the topology. While this method is faster, it requires\n", - " workers to load all feature values into memory, which limits its ability to\n", - " scale.\n", - "- Note that all the beam method start with `dgf.beam.*`." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "b39436b6" - }, - "source": [ - "We can now load and plot the graph samples." - ] - }, - { - "cell_type": "code", - "execution_count": 24, - "metadata": { - "execution": { - "iopub.execute_input": "2026-09-23T18:34:42.745544Z", - "iopub.status.busy": "2026-09-23T18:34:42.745327Z", - "iopub.status.idle": "2026-09-23T18:34:47.479165Z", - "shell.execute_reply": "2026-09-23T18:34:47.478595Z" + "cell_type": "code", + "execution_count": 23, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:32:44.289991Z", + "iopub.status.busy": "2026-09-23T18:32:44.289809Z", + "iopub.status.idle": "2026-09-23T18:34:42.743424Z", + "shell.execute_reply": "2026-09-23T18:34:42.742728Z" + }, + "executionInfo": { + "elapsed": 118455, + "status": "ok", + "timestamp": 1790188482744.6162, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + } + }, + "outputs": [ + { + "data": { + "application/javascript": [ + "\n", + " if (typeof window.interactive_beam_jquery == 'undefined') {\n", + " var jqueryScript = document.createElement('script');\n", + " jqueryScript.src = 'https://code.jquery.com/jquery-3.4.1.slim.min.js';\n", + " jqueryScript.type = 'text/javascript';\n", + " jqueryScript.onload = function() {\n", + " var datatableScript = document.createElement('script');\n", + " datatableScript.src = 'https://cdn.datatables.net/1.10.20/js/jquery.dataTables.min.js';\n", + " datatableScript.type = 'text/javascript';\n", + " datatableScript.onload = function() {\n", + " window.interactive_beam_jquery = jQuery.noConflict(true);\n", + " window.interactive_beam_jquery(document).ready(function($){\n", + " \n", + " });\n", + " }\n", + " document.head.appendChild(datatableScript);\n", + " };\n", + " document.head.appendChild(jqueryScript);\n", + " } else {\n", + " window.interactive_beam_jquery(document).ready(function($){\n", + " \n", + " });\n", + " }" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Using override schema\n", + "Reading 1 nodeset(s), 1 edgeset(s), and 5 feature(s)\n", + "Reading metadata from /tmp/my_graph\n", + "Reading nodeset nodes from /tmp/my_graph\n", + ".concatenating values\n", + "Reading edgeset edges from /tmp/my_graph\n", + ".concatenating values\n", + "Graph read in memory in 0.52s\n" + ] + } + ], + "source": [ + "with beam.Pipeline(runner=flume_runner.FlumeRunner(), options=options) as root:\n", + "\n", + " # Read only the node and the #id feature from the graph. Those\n", + " # is the nodes to seed.\n", + " seed_graph = dgf.beam.io.read_graph(\n", + " root,\n", + " graph_path,\n", + " schema_filter=dgf.data.GraphSchemaFilter(\n", + " nodeset_fn=lambda key, sch: key == \"nodes\",\n", + " edgeset_fn=lambda key, sch: False,\n", + " feature_fn=lambda key, sch: key == \"#id\",\n", + " ),\n", + " )\n", + " seed_node_ids = dgf.beam.sampling.extract_nodes_ids(seed_graph, \"nodes\")\n", + "\n", + " # Alternatively, you can do:\n", + " # seed_node_ids = beam.Create([b\"\", b\"\"])\n", + "\n", + " # Randomly select 50 seed nodes.\n", + " seed_node_ids = (\n", + " seed_node_ids\n", + " | \"Sample seeds\" >> beam.combiners.Sample.FixedSizeGlobally(50)\n", + " | beam.FlatMap(lambda xs: xs)\n", + " )\n", + "\n", + " # Generate samples\n", + " samples, output_schema = dgf.beam.sampling.semi_distributed_sampler_v2(\n", + " graph_path=graph_path,\n", + " plan=sampler_config,\n", + " seeds=seed_node_ids,\n", + " num_threads=20,\n", + " beam_feature_collection=False,\n", + " )\n", + "\n", + " # Save the samples to disk\n", + " dgf.beam.io.write_tfgnn_graphs(samples, \"/tmp/samples@*\", output_schema)" + ], + "id": "ecbaa932" }, - "executionInfo": { - "elapsed": 4735, - "status": "ok", - "timestamp": 1790188487480.6619, - "user": { - "user_id": "gbm" - }, - "user_tz": -120 + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Remark:**\n", + "\n", + "- The semi-distributed sampler (v2) takes as input a graph stored on disk. It\n", + " is today, the fastest option (faster than loading the input graph from a\n", + " beam pcollection like the v1).\n", + "- The `num_threads` option control the number of threads for each worker.\n", + "- Setting `beam_feature_collection=True` tells the sampler to generate\n", + " topologies first and then collect feature values using Apache Beam. This\n", + " approach scales effectively because workers only store the graph topology in\n", + " memory, though it results in slower overall performance. Alternatively,\n", + " `beam_feature_collection=False` triggers feature collection at the same time\n", + " the sampler generates the topology. While this method is faster, it requires\n", + " workers to load all feature values into memory, which limits its ability to\n", + " scale.\n", + "- Note that all the beam method start with `dgf.beam.*`." + ], + "id": "70dcd856" }, - "id": "d75df1fd" - }, - "outputs": [ { - "name": "stdout", - "output_type": "stream", - "text": [ - "Found 50 graphs\n" - ] - } - ], - "source": [ - "graphs = list(dgf.io.read_tfgnn_graphs(path=\"/tmp/samples@*\", schema=schema))\n", - "print(f\"Found {len(graphs)} graphs\")" - ] - } - ], - "metadata": { - "colab": { - "last_runtime": { - "build_target": "", - "kind": "local" - }, - "provenance": [ + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can now load and plot the graph samples." + ], + "id": "b39436b6" + }, { - "file_id": "1hMr9gK0hhKcCtIs1Pro3_ksOnYcdAL2P", - "timestamp": 1771583153488 - } - ], - "toc_visible": true, - "views": { - "output_only": { - "cells": [ - { - "id": "52d5d84b" - }, - { - "id": "EsL3MiYR9nD6" - }, - { - "id": "D1I1ijy08t0v" - }, - { - "id": "eASOvqWCLNMQ" - }, - { - "id": "71e9cf1c" - }, - { - "id": "f5431c44" - }, - { - "id": "7eb84356" - }, - { - "id": "7f9272a6" - }, - { - "id": "d0fb3022" - }, - { - "id": "2896f2d1" - }, - { - "id": "fba58d46" - }, - { - "id": "4ee7378e" - }, - { - "id": "daee962f" - }, - { - "id": "a98cd4ee" - }, - { - "id": "2d0902f0" - }, - { - "id": "5178bcff" - }, - { - "id": "1f9373f2" - }, - { - "id": "07fe45e1" - }, - { - "id": "8cea9faa" - }, - { - "id": "d67499e6" - }, - { - "id": "3a98d6ce" - }, - { - "id": "31410d56" - }, - { - "id": "3efc55a0" - }, - { - "id": "c93aafb9" - }, - { - "id": "e34b4b65" - }, - { - "id": "O5j_22f5HJYn" + "cell_type": "code", + "execution_count": 24, + "metadata": { + "execution": { + "iopub.execute_input": "2026-09-23T18:34:42.745544Z", + "iopub.status.busy": "2026-09-23T18:34:42.745327Z", + "iopub.status.idle": "2026-09-23T18:34:47.479165Z", + "shell.execute_reply": "2026-09-23T18:34:47.478595Z" + }, + "executionInfo": { + "elapsed": 4735, + "status": "ok", + "timestamp": 1790188487480.6619, + "user": { + "user_id": "gbm" + }, + "user_tz": -120 + } }, - { - "id": "LGu-U8SkHTMV" - }, - { - "id": "Yxpn64EQb_0R" - }, - { - "id": "OrG4j5_OZ4lB" - }, - { - "id": "grTlK4Lnbvoq" - }, - { - "id": "868fTc8JbkzT" - }, - { - "id": "dK98Ep4Mcnya" - }, - { - "id": "5c5d4YfWcpar" - }, - { - "id": "bfbc2230" - }, - { - "id": "NvZ9fNQ7LxSs" - }, - { - "id": "5ad2a4f9" - }, - { - "id": "c6d421dc" - }, - { - "id": "909d6d86" - }, - { - "id": "91611d90" - }, - { - "id": "5cb0d579" - }, - { - "id": "a35660d2" - }, - { - "id": "ef4b8f03" - }, - { - "id": "4151f762" - }, - { - "id": "71fc0067" - }, - { - "id": "f4a3e3de" - }, - { - "id": "ecbaa932" - }, - { - "id": "70dcd856" - }, - { - "id": "b39436b6" + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Found 50 graphs\n" + ] + } + ], + "source": [ + "graphs = list(dgf.io.read_tfgnn_graphs(path=\"/tmp/samples@*\", schema=schema))\n", + "print(f\"Found {len(graphs)} graphs\")" + ], + "id": "d75df1fd" + } + ], + "metadata": { + "colab": { + "last_runtime": { + "build_target": "", + "kind": "local" }, - { - "id": "d75df1fd" + "provenance": [ + { + "file_id": "1hMr9gK0hhKcCtIs1Pro3_ksOnYcdAL2P", + "timestamp": 1771583153488 + } + ], + "toc_visible": true, + "views": { + "output_only": { + "cells": [ + { + "id": "52d5d84b" + }, + { + "id": "EsL3MiYR9nD6" + }, + { + "id": "D1I1ijy08t0v" + }, + { + "id": "eASOvqWCLNMQ" + }, + { + "id": "71e9cf1c" + }, + { + "id": "f5431c44" + }, + { + "id": "7eb84356" + }, + { + "id": "7f9272a6" + }, + { + "id": "d0fb3022" + }, + { + "id": "2896f2d1" + }, + { + "id": "fba58d46" + }, + { + "id": "4ee7378e" + }, + { + "id": "daee962f" + }, + { + "id": "a98cd4ee" + }, + { + "id": "2d0902f0" + }, + { + "id": "5178bcff" + }, + { + "id": "1f9373f2" + }, + { + "id": "07fe45e1" + }, + { + "id": "8cea9faa" + }, + { + "id": "d67499e6" + }, + { + "id": "3a98d6ce" + }, + { + "id": "31410d56" + }, + { + "id": "3efc55a0" + }, + { + "id": "c93aafb9" + }, + { + "id": "e34b4b65" + }, + { + "id": "O5j_22f5HJYn" + }, + { + "id": "LGu-U8SkHTMV" + }, + { + "id": "Yxpn64EQb_0R" + }, + { + "id": "OrG4j5_OZ4lB" + }, + { + "id": "grTlK4Lnbvoq" + }, + { + "id": "868fTc8JbkzT" + }, + { + "id": "dK98Ep4Mcnya" + }, + { + "id": "5c5d4YfWcpar" + }, + { + "id": "bfbc2230" + }, + { + "id": "NvZ9fNQ7LxSs" + }, + { + "id": "5ad2a4f9" + }, + { + "id": "c6d421dc" + }, + { + "id": "909d6d86" + }, + { + "id": "91611d90" + }, + { + "id": "5cb0d579" + }, + { + "id": "a35660d2" + }, + { + "id": "ef4b8f03" + }, + { + "id": "4151f762" + }, + { + "id": "71fc0067" + }, + { + "id": "f4a3e3de" + }, + { + "id": "ecbaa932" + }, + { + "id": "70dcd856" + }, + { + "id": "b39436b6" + }, + { + "id": "d75df1fd" + } + ], + "hide_code": true + } } - ], - "hide_code": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.14.5" } - } - }, - "kernelspec": { - "display_name": "Python 3", - "name": "python3" }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.14.5" - } - }, - "nbformat": 4, - "nbformat_minor": 0 + "nbformat": 4, + "nbformat_minor": 5 } diff --git a/doc/mkdocs.yml b/doc/mkdocs.yml index 1eaaaeb..b4f3024 100644 --- a/doc/mkdocs.yml +++ b/doc/mkdocs.yml @@ -78,7 +78,8 @@ nav: - Timeseries: tutorial/timeseries_node_prediction.ipynb - Advanced API: - In-memory graph: tutorial/in_memory_graph.ipynb - - Sampler: tutorial/sampler.ipynb + - In-process Sampler: tutorial/sampler.ipynb + - Offline Distributed Sampler (GCP): tutorial/gcp_offline_distributed_sampler.ipynb - Normaliers: tutorial/normalizer.ipynb - Guides: - Graph file format: file_formats.md