examples/spark_engine.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a54fe998",
   "metadata": {},
   "source": [
    "![image](https://user-images.githubusercontent.com/1651790/221876073-61ef4edb-adcd-4f10-b3fc-8ddc24918ea1.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f46fdd40",
   "metadata": {},
   "outputs": [],
   "source": [
    "# install ng_ai in the first run\n",
    "%pip install ng_ai"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5b4e4143",
   "metadata": {},
   "source": [
    "## AI Suite Spark Engine Examples\n",
    "### read data with spark engine, scan mode"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f17abcf8",
   "metadata": {},
   "source": [
    "In this example, we are leveraging the Spark Engine of NebulaGraph AI Suite, with the Storage Scan mode.\n",
    "\n",
    "#### Step 1, get dataframe by scanning the Graph\n",
    "\n",
    "We will scan all edge in type `follow` first as dataframe: `df`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e158440f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r\n",
      "[Stage 0:>                                                          (0 + 1) / 1]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+---------+-----+------+\n",
      "|   _srcId|   _dstId|_rank|degree|\n",
      "+---------+---------+-----+------+\n",
      "|player105|player100|    0|    70|\n",
      "|player105|player104|    0|    83|\n",
      "+---------+---------+-----+------+\n",
      "only showing top 2 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r\n",
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "from ng_ai import NebulaReader\n",
    "\n",
    "# read data with spark engine, scan mode\n",
    "reader = NebulaReader(engine=\"spark\")\n",
    "reader.scan(edge=\"follow\", props=\"degree\")\n",
    "df = reader.read()\n",
    "df.show(2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3617de5f",
   "metadata": {},
   "source": [
    "#### Step 2, run Pagerank Algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "90069aaf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+\n",
      "|       id|\n",
      "+---------+\n",
      "|player108|\n",
      "|player129|\n",
      "|player120|\n",
      "|player103|\n",
      "|player128|\n",
      "|player148|\n",
      "|player117|\n",
      "|player139|\n",
      "|player140|\n",
      "|player134|\n",
      "|player149|\n",
      "|player150|\n",
      "|player125|\n",
      "|player137|\n",
      "|player143|\n",
      "|player101|\n",
      "|player141|\n",
      "|player144|\n",
      "|player102|\n",
      "|player121|\n",
      "+---------+\n",
      "only showing top 20 rows\n",
      "\n",
      "+-----+------+-----------+----------+\n",
      "|_rank|degree|     _srcId|    _dstId|\n",
      "+-----+------+-----------+----------+\n",
      "|    0|    90|34359738371|         3|\n",
      "|    0|    90|25769803786|         6|\n",
      "|    0|    90|34359738369|         6|\n",
      "|    0|    80| 8589934596|         2|\n",
      "|    0|    99|25769803784|         2|\n",
      "|    0|    90|25769803777|         2|\n",
      "|    0|    90|          1|         4|\n",
      "|    0|    90|17179869189|         4|\n",
      "|    0|    90|          4|         1|\n",
      "|    0|    10| 8589934598|         1|\n",
      "|    0|    90|17179869189|         1|\n",
      "|    0|    80| 8589934598|         5|\n",
      "|    0|    85|25769803786|         5|\n",
      "|    0|    70|34359738373|         5|\n",
      "|    0|    95|17179869185|8589934597|\n",
      "|    0|    95|25769803778|8589934597|\n",
      "|    0|    99|25769803784|8589934597|\n",
      "|    0|    90|34359738368|8589934597|\n",
      "|    0|    85|          5|8589934598|\n",
      "|    0|    90| 8589934596|8589934598|\n",
      "+-----+------+-----------+----------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "23/03/01 13:17:34 WARN BlockManager: Block rdd_75_2 already exists on this machine; not re-adding it\n",
      "23/03/01 13:17:34 WARN BlockManager: Block rdd_75_3 already exists on this machine; not re-adding it\n",
      "23/03/01 13:17:34 WARN BlockManager: Block rdd_75_1 already exists on this machine; not re-adding it\n"
     ]
    }
   ],
   "source": [
    "pr_result = df.algo.pagerank(reset_prob=0.15, max_iter=10)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "66e70ca0",
   "metadata": {},
   "source": [
    "#### Step 3, check results of the algorithm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "abbce2fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+-------------------+\n",
      "|      _id|           pagerank|\n",
      "+---------+-------------------+\n",
      "|player133|0.18601069183310504|\n",
      "|player126|0.18601069183310504|\n",
      "|player130|  1.240071278887367|\n",
      "|player108|0.18601069183310504|\n",
      "|player102| 1.6602373739502536|\n",
      "+---------+-------------------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "pr_result.show(5)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3eb228f8",
   "metadata": {},
   "source": [
    "------------------\n",
    "### read data with spark engine, query mode\n",
    "\n",
    "In this example, we are leveraging the Spark Engine of NebulaGraph DI Suite, with the Graph Query mode.\n",
    "\n",
    "#### Step 1, get dataframe by querying the Graph with a Cypher\n",
    "\n",
    "We will query 100000 edges in type `follow` as a dataframe: `df`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e44ac3e9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+---------+-----+------+\n",
      "|   _srcId|   _dstId|_rank|degree|\n",
      "+---------+---------+-----+------+\n",
      "|player102|player100|    0|    75|\n",
      "|player102|player101|    0|    75|\n",
      "+---------+---------+-----+------+\n",
      "only showing top 2 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from ng_ai import NebulaReader\n",
    "\n",
    "# read data with spark engine, query mode\n",
    "reader = NebulaReader(engine=\"spark\")\n",
    "query = \"\"\"\n",
    "    MATCH ()-[e:follow]->()\n",
    "    RETURN e LIMIT 100000\n",
    "\"\"\"\n",
    "reader.query(query=query, edge=\"follow\", props=\"degree\")\n",
    "df = reader.read()  # this will take some time\n",
    "df.show(2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "49becbdb",
   "metadata": {},
   "source": [
    "#### Step 2, run Conncted Components Algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cfbcda82",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+\n",
      "|       id|\n",
      "+---------+\n",
      "|player129|\n",
      "|player120|\n",
      "|player148|\n",
      "|player103|\n",
      "|player128|\n",
      "|player108|\n",
      "|player117|\n",
      "|player150|\n",
      "|player125|\n",
      "|player137|\n",
      "|player139|\n",
      "|player140|\n",
      "|player134|\n",
      "|player149|\n",
      "|player102|\n",
      "|player135|\n",
      "|player147|\n",
      "|player121|\n",
      "|player143|\n",
      "|player101|\n",
      "+---------+\n",
      "only showing top 20 rows\n",
      "\n",
      "+-----+------+-----------+----------+\n",
      "|_rank|degree|     _srcId|    _dstId|\n",
      "+-----+------+-----------+----------+\n",
      "|    0|    90|34359738368|         3|\n",
      "|    0|    90|25769803781|         6|\n",
      "|    0|    90|34359738371|         6|\n",
      "|    0|    80| 8589934592|         1|\n",
      "|    0|    99|25769803779|         1|\n",
      "|    0|    90|25769803784|         1|\n",
      "|    0|    90|          0|         4|\n",
      "|    0|    90|17179869187|         4|\n",
      "|    0|    90|          4|         0|\n",
      "|    0|    10| 8589934594|         0|\n",
      "|    0|    90|17179869187|         0|\n",
      "|    0|    80| 8589934594|         2|\n",
      "|    0|    85|25769803781|         2|\n",
      "|    0|    70|34359738370|         2|\n",
      "|    0|    95|17179869189|8589934593|\n",
      "|    0|    95|25769803785|8589934593|\n",
      "|    0|    99|25769803779|8589934593|\n",
      "|    0|    90|34359738372|8589934593|\n",
      "|    0|    85|          2|8589934594|\n",
      "|    0|    90| 8589934592|8589934594|\n",
      "+-----+------+-----------+----------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "cc_result = df.algo.connected_components(max_iter=10)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "38181d45",
   "metadata": {},
   "source": [
    "#### Step 3, check results of the algorithm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bed14375",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+---------+\n",
      "|      _id|       cc|\n",
      "+---------+---------+\n",
      "|player115|player129|\n",
      "|player113|player129|\n",
      "|player100|player129|\n",
      "|player129|player129|\n",
      "|player137|player129|\n",
      "+---------+---------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "cc_result.show(5)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3d088006",
   "metadata": {},
   "source": [
    "### Write back algo result to NebulaGraph as TAG\n",
    "\n",
    "Assume that we have a Spark DataFrame `df_result` computed with `df.algo.label_propagation()` with the following schema:\n",
    "\n",
    "```python\n",
    "df_result.printSchema()\n",
    "# result:\n",
    "root\n",
    " |-- _id: string (nullable = false)\n",
    " |-- lpa: string (nullable = false)\n",
    "```\n",
    "\n",
    "There are two columns, `_id` is the vid, and `lpa` is the label propagation detected cluster id, let's write them back to tag: label_propagation(cluster_id). So we create a TAG `label_propagation` in NebulaGraph on same space with the following schema:\n",
    "\n",
    "```ngql\n",
    "CREATE TAG IF NOT EXISTS label_propagation (\n",
    "    cluster_id string NOT NULL\n",
    ");\n",
    "```\n",
    "\n",
    "Then, we could write the label propagation result to NebulaGraph, map the column `lpa` to `cluster_id`:\n",
    "```python\n",
    "properties = {\n",
    "    \"lpa\": \"cluster_id\"\n",
    "}\n",
    "```\n",
    "And pass it to NebulaWriter in `spark` engine and `nebulagraph_vertex` sink"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6b43261f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+\n",
      "|       id|\n",
      "+---------+\n",
      "|player129|\n",
      "|player120|\n",
      "|player148|\n",
      "|player103|\n",
      "|player128|\n",
      "|player108|\n",
      "|player117|\n",
      "|player150|\n",
      "|player125|\n",
      "|player137|\n",
      "|player139|\n",
      "|player140|\n",
      "|player134|\n",
      "|player149|\n",
      "|player102|\n",
      "|player135|\n",
      "|player147|\n",
      "|player121|\n",
      "|player143|\n",
      "|player101|\n",
      "+---------+\n",
      "only showing top 20 rows\n",
      "\n",
      "+-----+------+----------+-----------+\n",
      "|_rank|degree|    _srcId|     _dstId|\n",
      "+-----+------+----------+-----------+\n",
      "|    0|    70|         3|17179869184|\n",
      "|    0|    80|         5|25769803785|\n",
      "|    0|    80|         5|17179869189|\n",
      "|    0|    80|         1|25769803784|\n",
      "|    0|    90|         4|25769803778|\n",
      "|    0|    90|         4|17179869187|\n",
      "|    0|    90|         4|          0|\n",
      "|    0|    90|         0|25769803778|\n",
      "|    0|    90|         0|17179869187|\n",
      "|    0|    90|         0|          4|\n",
      "|    0|    80|         2|34359738370|\n",
      "|    0|    90|         2|25769803781|\n",
      "|    0|    85|         2| 8589934594|\n",
      "|    0|    90|8589934593|25769803785|\n",
      "|    0|    -1|8589934597|17179869187|\n",
      "|    0|    10|8589934594|          0|\n",
      "|    0|    80|8589934594|25769803781|\n",
      "|    0|    80|8589934594|          2|\n",
      "|    0|    99|8589934595|25769803776|\n",
      "|    0|    90|8589934596|34359738368|\n",
      "+-----+------+----------+-----------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "23/03/01 13:18:01 WARN CacheManager: Asked to cache already cached data.\n",
      "23/03/01 13:18:01 WARN CacheManager: Asked to cache already cached data.\n"
     ]
    }
   ],
   "source": [
    "# Run label Propagation Algorithm\n",
    "df_result = df.algo.label_propagation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a6e011db",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- _id: string (nullable = false)\n",
      " |-- lpa: string (nullable = false)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# check the result schema\n",
    "df_result.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c5bbf9e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ng_ai import NebulaWriter\n",
    "from ng_ai.config import NebulaGraphConfig\n",
    "\n",
    "config = NebulaGraphConfig()\n",
    "writer = NebulaWriter(\n",
    "    data=df_result, sink=\"nebulagraph_vertex\", config=config, engine=\"spark\"\n",
    ")\n",
    "\n",
    "# map column louvain into property cluster_id\n",
    "properties = {\"lpa\": \"cluster_id\"}\n",
    "\n",
    "writer.set_options(\n",
    "    space=\"basketballplayer\",\n",
    "    tag=\"label_propagation\",\n",
    "    vid_field=\"_id\",\n",
    "    properties=properties,\n",
    "    batch_size=256,\n",
    "    write_mode=\"insert\",\n",
    ")\n",
    "# write back to NebulaGraph\n",
    "writer.write()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9da30271",
   "metadata": {},
   "source": [
    "Then we could query the result in NebulaGraph:\n",
    "\n",
    "```cypher\n",
    "MATCH (v:label_propagation)\n",
    "RETURN id(v), v.label_propagation.cluster_id LIMIT 10;\n",
    "```\n",
    "Result:\n",
    "\n",
    "```cypher\n",
    "(root@nebula) [basketballplayer]> \"\"\"\n",
    "                               -> MATCH (v:label_propagation)\n",
    "                               -> RETURN id(v), v.label_propagation.cluster_id LIMIT 10;\n",
    "                               -> \"\"\"\n",
    "+-------------+--------------------------------+\n",
    "| id(v)       | v.label_propagation.cluster_id |\n",
    "+-------------+--------------------------------+\n",
    "| \"player103\" | \"player101\"                    |\n",
    "| \"player113\" | \"player129\"                    |\n",
    "| \"player121\" | \"player129\"                    |\n",
    "| \"player128\" | \"player129\"                    |\n",
    "| \"player130\" | \"player130\"                    |\n",
    "| \"player136\" | \"player136\"                    |\n",
    "| \"player127\" | \"player137\"                    |\n",
    "| \"player135\" | \"player101\"                    |\n",
    "| \"player147\" | \"player148\"                    |\n",
    "| \"player148\" | \"player148\"                    |\n",
    "+-------------+--------------------------------+\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "670bd34b",
   "metadata": {},
   "source": [
    "### Result being written as edge\n",
    "\n",
    "Similar to TAG, we first need to ensure to create schema first\n",
    "\n",
    "```\n",
    "CREATE EDGE jaccard_similarity(similarity double);\n",
    "```\n",
    "\n",
    "Then we run a algorithm writting results to edge:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "493cc969",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run Jaccard Algorithm\n",
    "df_result = df.algo.jaccard()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b9e76ef2",
   "metadata": {},
   "source": [
    "Then let's write the result to NebulaGraph, map the column `similarity` to `similarity`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e91af7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "writer = NebulaWriter(\n",
    "    data=df_result, sink=\"nebulagraph_edge\", config=config, engine=\"spark\"\n",
    ")\n",
    "\n",
    "# map column louvain into property cluster_id\n",
    "properties = {\"similarity\": \"similarity\"}\n",
    "\n",
    "writer.set_options(\n",
    "    space=\"basketballplayer\",\n",
    "    edge_type=\"jaccard_similarity\",\n",
    "    src_id=\"srcId\",\n",
    "    dst_id=\"dstId\",\n",
    "    src_id_policy=\"\",\n",
    "    dst_id_policy=\"\",\n",
    "    properties=properties,\n",
    "    batch_size=256,\n",
    "    write_mode=\"insert\",\n",
    ")\n",
    "\n",
    "# write back to NebulaGraph\n",
    "writer.write()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8165fec9",
   "metadata": {},
   "source": [
    "Check result:\n",
    "\n",
    "```\n",
    "(root@nebula) [basketballplayer]> MATCH ()-[e:jaccard_similarity]->() RETURN e LIMIT 3\n",
    "+-------------------------------------------------------------------------------------+\n",
    "| e                                                                                   |\n",
    "+-------------------------------------------------------------------------------------+\n",
    "| [:jaccard_similarity \"player102\"->\"player100\" @0 {similarity: 0.07692307692307687}] |\n",
    "| [:jaccard_similarity \"player102\"->\"player101\" @0 {similarity: 0.11111111111111116}] |\n",
    "| [:jaccard_similarity \"player102\"->\"player104\" @0 {similarity: 0.33333333333333326}] |\n",
    "+-------------------------------------------------------------------------------------+\n",
    "Got 3 rows (time spent 39.984ms/44.574542ms)\n",
    "\n",
    "Wed, 06 Sep 2023 13:04:38 CST\n",
    "\n",
    "(root@nebula) [basketballplayer]>\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5bcb02e2",
   "metadata": {},
   "source": [
    "## How to run other algorithm examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff5a866d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# lpa_result  = df.algo.label_propagation()\n",
    "# louvain_result = df.algo.louvain()\n",
    "# k_core_result = df.algo.k_core()\n",
    "# degree_statics_result = df.algo.degree_statics()\n",
    "# betweenness_centrality_result = df.algo.betweenness_centrality()\n",
    "# clustering_coefficient_result = df.algo.clustering_coefficient()\n",
    "# bfs_result = df.algo.bfs()\n",
    "# hanp_result = df.algo.hanp()\n",
    "# jaccard_result = df.algo.jaccard()\n",
    "# strong_connected_components_result = df.algo.strong_connected_components()\n",
    "# triangle_count_result = df.algo.triangle_count()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}