experimental.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7cd33e77-a7a4-4e34-a124-1f7bdc0b3264",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import os\n",
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2fe8f0b8-821c-42f6-a2a9-47cc71f90cbe",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "os.environ['PYSPARK_PYTHON'] = \"/home/group15/.conda/envs/group15/bin/python\"\n",
    "os.environ['PYSPARK_DRIVER_PYTHON'] = \"/home/group15/.conda/envs/group15/bin/python\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dd8af5f-5d9c-4def-96cc-1e2be7042222",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Spark Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ff15b4a9-01bd-445a-a00a-ff9dab1c2d4d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pyspark.sql.functions as F\n",
    "import pyspark.sql.types as T\n",
    "from pyspark.ml.feature import VectorAssembler, StringIndexer, StandardScaler\n",
    "from pyspark.ml import Pipeline\n",
    "from pyspark.ml.classification import RandomForestClassifier, RandomForestClassificationModel\n",
    "from pyspark.ml.tuning import ParamGridBuilder, CrossValidator\n",
    "from pyspark.ml.evaluation import MulticlassClassificationEvaluator\n",
    "\n",
    "import json\n",
    "from itertools import chain\n",
    "import pycountry_convert as pc\n",
    "\n",
    "from twitter.config import LANG_MAP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "16bae8be-8d53-4dd9-9e96-af1d9b754522",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "23/12/24 15:46:16 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n",
      "Setting default log level to \"WARN\".\n",
      "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n",
      "23/12/24 15:46:18 WARN Utils: Service 'SparkUI' could not bind on port 4040. Attempting port 4041.\n",
      "23/12/24 15:46:18 WARN Utils: Service 'SparkUI' could not bind on port 4041. Attempting port 4042.\n"
     ]
    }
   ],
   "source": [
    "from pyspark.sql import SparkSession\n",
    "\n",
    "# config the connector jar file\n",
    "spark = (SparkSession.builder.appName(\"SimpleSparkJob\").master(\"spark://34.142.194.212:7077\")\n",
    "         .config(\"spark.jars\",\"/opt/spark/jars/gcs-connector-latest-hadoop2.jar\")\n",
    "         .config(\"spark.executor.memory\", \"1G\")  #excutor excute only 2G\n",
    "         .config(\"spark.driver.memory\",\"1G\") \n",
    "         .config(\"spark.executor.cores\",\"1\") # Cluster use only 3 cores to excute as it has 3 server\n",
    "         .config(\"spark.python.worker.memory\",\"1G\") # each worker use 1G to excute\n",
    "         .config(\"spark.driver.maxResultSize\",\"1G\") #Maximum size of result is 3G\n",
    "         .config(\"spark.kryoserializer.buffer.max\",\"1024M\")\n",
    "         .getOrCreate())\n",
    "\n",
    "# datetime migration\n",
    "spark.conf.set(\"spark.sql.legacy.timeParserPolicy\",\"LEGACY\")\n",
    "\n",
    "# config the credential to identify the google cloud hadoop file \n",
    "spark.conf.set(\"google.cloud.auth.service.account.json.keyfile\",\"/opt/bucket_connector/lucky-wall-393304-3fbad5f3943c.json\")\n",
    "spark._jsc.hadoopConfiguration().set('fs.gs.impl', 'com.google.cloud.hadoop.fs.gcs.GoogleHadoopFileSystem')\n",
    "spark._jsc.hadoopConfiguration().set('fs.gs.auth.service.account.enable', 'true')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "272608b9-2a99-4eeb-9335-5409cdde6ae1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "            <div>\n",
       "                <p><b>SparkSession - in-memory</b></p>\n",
       "                \n",
       "        <div>\n",
       "            <p><b>SparkContext</b></p>\n",
       "\n",
       "            <p><a href=\"http://bigdata-02.asia-southeast1-b.c.lucky-wall-393304.internal:4042\">Spark UI</a></p>\n",
       "\n",
       "            <dl>\n",
       "              <dt>Version</dt>\n",
       "                <dd><code>v3.5.0</code></dd>\n",
       "              <dt>Master</dt>\n",
       "                <dd><code>spark://34.142.194.212:7077</code></dd>\n",
       "              <dt>AppName</dt>\n",
       "                <dd><code>SimpleSparkJob</code></dd>\n",
       "            </dl>\n",
       "        </div>\n",
       "        \n",
       "            </div>\n",
       "        "
      ],
      "text/plain": [
       "<pyspark.sql.session.SparkSession at 0x7f52ccfd9fc0>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spark"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5985bba-0c0a-4eff-a146-a181f2e741f7",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Data Cleaning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "af5849f0-33bf-4e03-8ed8-708785617c1d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "23/12/24 15:47:27 WARN SparkStringUtils: Truncated the string representation of a plan since it was too large. This behavior can be adjusted by setting 'spark.sql.debug.maxToStringFields'.\n",
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+--------------------+--------------------+--------------------+\n",
      "|            follower|           following|               tweet|                user|\n",
      "+--------------------+--------------------+--------------------+--------------------+\n",
      "|[1700886965070209...|[241664456, 35731...|[{snscrape.module...|{snscrape.modules...|\n",
      "+--------------------+--------------------+--------------------+--------------------+\n",
      "only showing top 1 row\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# load_profile\n",
    "PROFILE_PATH = \"gs://it4043e-it5384/it4043e/it4043e_group15_problem1/output/profiles.jsonl\"\n",
    "def load_profile(fp):\n",
    "    profile_df = spark.read.json(fp)\n",
    "    return profile_df\n",
    "profile_df = load_profile(PROFILE_PATH)\n",
    "profile_df.show(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25f8edd7-9e29-4eb8-8492-39b8abc731b7",
   "metadata": {
    "tags": []
   },
   "source": [
    "## User Cleaning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "08b434b4-91e1-4911-b555-e07381af2e68",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+----+--------+--------------------+--------------------+--------------------+---------------+--------------+------------+----------+----------+-----------+--------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+--------------+--------+\n",
      "|               _type|blue|blueType|             created|    descriptionLinks|         displayname|favouritesCount|followersCount|friendsCount|        id|    id_str|listedCount|location|mediaCount|    profileBannerUrl|     profileImageUrl|protected|      rawDescription|statusesCount|                 url|      username|verified|\n",
      "+--------------------+----+--------+--------------------+--------------------+--------------------+---------------+--------------+------------+----------+----------+-----------+--------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+--------------+--------+\n",
      "|snscrape.modules....|true|    NULL|2014-10-18 15:45:...|[{https://t.co/29...|Day Trading BTC S...|             11|           868|           4|2862957343|2862957343|          2|        |      1315|https://pbs.twimg...|https://pbs.twimg...|     NULL|Making your dream...|         4140|https://twitter.c...|Day_trading247|   false|\n",
      "+--------------------+----+--------+--------------------+--------------------+--------------------+---------------+--------------+------------+----------+----------+-----------+--------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+--------------+--------+\n",
      "only showing top 1 row\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# _extract_users\n",
    "def _extract_users(profile_df):\n",
    "    user_df = profile_df.select(F.col(\"user.*\"))\n",
    "    return user_df\n",
    "\n",
    "user_df = _extract_users(profile_df)\n",
    "user_df.show(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "6c4341f7-5b91-4a2a-847f-c8f108374cd4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+----+--------+--------------------+--------------------+--------------------+---------------+--------------+------------+-------------------+-------------------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+-----------+--------+-------------+\n",
      "|               _type|blue|blueType|             created|    descriptionLinks|         displayname|favouritesCount|followersCount|friendsCount|                 id|             id_str|listedCount|     location|mediaCount|    profileBannerUrl|     profileImageUrl|protected|      rawDescription|statusesCount|                 url|   username|verified|      country|\n",
      "+--------------------+----+--------+--------------------+--------------------+--------------------+---------------+--------------+------------+-------------------+-------------------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+-----------+--------+-------------+\n",
      "|snscrape.modules....|true|    NULL|2018-05-28 21:55:...|[{https://t.co/c8...|Samuel Armes (Flo...|           3890|          2017|         465|1001220565510836224|1001220565510836224|         17|North America|       129|https://pbs.twimg...|https://pbs.twimg...|     NULL|Surf maxi 🌊 and ...|         1234|https://twitter.c...|samuelarmes|   false|United States|\n",
      "+--------------------+----+--------+--------------------+--------------------+--------------------+---------------+--------------+------------+-------------------+-------------------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+-----------+--------+-------------+\n",
      "only showing top 1 row\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# _label_users\n",
    "LOCATION_PATH = \"output/location.json\"\n",
    "MISSING_CN_TO_CNT = {\n",
    "    \"Kosovo\": \"Europe\",\n",
    "    \"Palestinian Territories\": \"Asia\",\n",
    "    \"North Pole\": \"Arctic\",\n",
    "    \"Ascension and Tristan da Cunha\": \"Africa\",\n",
    "    \"Gornja Siga\": \"Europe\"\n",
    "}\n",
    "\n",
    "def country_to_continent(country_name: str) -> str:\n",
    "    try:\n",
    "        country_alpha2 = pc.country_name_to_country_alpha2(country_name)\n",
    "        country_continent_code = pc.country_alpha2_to_continent_code(country_alpha2)\n",
    "        country_continent_name = pc.convert_continent_code_to_continent_name(country_continent_code)\n",
    "        return country_continent_name\n",
    "    except:\n",
    "        if country_name in [\"Europe\", \"Africa\", \"Antarctica\"]:\n",
    "            return country_name\n",
    "        elif country_name in MISSING_CN_TO_CNT:\n",
    "            return MISSING_CN_TO_CNT[country_name]\n",
    "        return None\n",
    "\n",
    "def _label_users(user_df):\n",
    "    with open(LOCATION_PATH, \"r\") as f:\n",
    "        location_dict = json.load(f)\n",
    "    for k, v in location_dict.items():\n",
    "        if isinstance(v , str):\n",
    "            if v == \"None\":\n",
    "                v = None\n",
    "            else:\n",
    "                v = v.strip()\n",
    "        location_dict[k] = v        \n",
    "    for unk_geo in [\"nan\", \"\"]:\n",
    "        location_dict[unk_geo] = None\n",
    "        \n",
    "    country_encode_fn = F.create_map([F.lit(x) for x in chain(*location_dict.items())])\n",
    "    user_df = user_df.withColumn(\"country\", country_encode_fn[F.col(\"location\")])\n",
    "        \n",
    "    for k, v in location_dict.items():\n",
    "        if v is not None:\n",
    "            location_dict[k] = country_to_continent(v)\n",
    "        else:\n",
    "            location_dict[k] = None\n",
    "        \n",
    "    label_encode_fn = F.create_map([F.lit(x) for x in chain(*location_dict.items())])\n",
    "    user_df = user_df.withColumn(\"location\", label_encode_fn[F.col(\"location\")])\n",
    "    \n",
    "    return user_df\n",
    "\n",
    "user_df = _label_users(user_df)\n",
    "user_df.where(F.col(\"location\").isNotNull()).show(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "5e7a9954-875b-4231-8e8c-28d3c55fe6d3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 91:>                                                         (0 + 2) / 2]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+-------+------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+\n",
      "|               _type|blue|blueType|            created|    descriptionLinks|displayname|favouritesCount|followersCount|friendsCount|user_id|id_str|listedCount|     location|mediaCount|    profileBannerUrl|     profileImageUrl|protected|      rawDescription|statusesCount|                 url|  username|verified|      country|\n",
      "+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+-------+------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+\n",
      "|snscrape.modules....|true|Business|2007-03-07 01:27:09|[{https://t.co/J0...| TechCrunch|           7394|      10249257|         454| 816653|816653|     106178|North America|     49908|https://pbs.twimg...|https://pbs.twimg...|     NULL|Technology news a...|       262860|https://twitter.c...|TechCrunch|   false|United States|\n",
      "+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+-------+------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+\n",
      "only showing top 1 row\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "# clean_users\n",
    "def clean_users(\n",
    "    profile_df\n",
    "):\n",
    "    # extract users from crawled data\n",
    "    user_df = _extract_users(profile_df)\n",
    "    \n",
    "    # consistent & unique\n",
    "    user_df = user_df \\\n",
    "                .withColumnRenamed(\"id\", \"user_id\") \\\n",
    "                .dropDuplicates(['user_id'])\n",
    "    user_df = user_df.na.drop(subset=[\"user_id\"])\n",
    "    user_df = user_df.withColumn(\"created\", F.to_timestamp(user_df.created, 'yyyy-MM-dd HH:mm:ss'))\n",
    "    \n",
    "    # handle target variable\n",
    "    user_df = _label_users(user_df)\n",
    "    \n",
    "    # drop all-NULL rows\n",
    "    user_df = user_df.na.drop(\"all\")\n",
    "    \n",
    "    return user_df\n",
    "\n",
    "user_df = clean_users(profile_df)\n",
    "user_df.show(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "566bf00f-90f0-4a1a-a13a-55d31d0fe4c8",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- _type: string (nullable = true)\n",
      " |-- blue: boolean (nullable = true)\n",
      " |-- blueType: string (nullable = true)\n",
      " |-- created: timestamp (nullable = true)\n",
      " |-- descriptionLinks: array (nullable = true)\n",
      " |    |-- element: struct (containsNull = true)\n",
      " |    |    |-- tcourl: string (nullable = true)\n",
      " |    |    |-- text: string (nullable = true)\n",
      " |    |    |-- url: string (nullable = true)\n",
      " |-- displayname: string (nullable = true)\n",
      " |-- favouritesCount: long (nullable = true)\n",
      " |-- followersCount: long (nullable = true)\n",
      " |-- friendsCount: long (nullable = true)\n",
      " |-- user_id: long (nullable = true)\n",
      " |-- id_str: string (nullable = true)\n",
      " |-- listedCount: long (nullable = true)\n",
      " |-- location: string (nullable = true)\n",
      " |-- mediaCount: long (nullable = true)\n",
      " |-- profileBannerUrl: string (nullable = true)\n",
      " |-- profileImageUrl: string (nullable = true)\n",
      " |-- protected: string (nullable = true)\n",
      " |-- rawDescription: string (nullable = true)\n",
      " |-- statusesCount: long (nullable = true)\n",
      " |-- url: string (nullable = true)\n",
      " |-- username: string (nullable = true)\n",
      " |-- verified: boolean (nullable = true)\n",
      " |-- country: string (nullable = true)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "user_df.printSchema()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e33648fa-8bea-44d9-a833-ab6f243104fd",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Tweet Cleaning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "2132aa71-2aee-413c-a5c3-4ad406983d90",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 49:>                                                         (0 + 1) / 1]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------+--------------------+--------+-------------------+-----------+--------------------+---------+-------------------+-------------------+----------------+-------------+----+---------+-----+--------------------+--------------+-----+----------+-----------+--------------------+----------+------------+--------------+--------------------+-----------+-----------------+--------------------+---------+\n",
      "|   user_id|               _type|cashtags|     conversationId|coordinates|                date| hashtags|                 id|             id_str|inReplyToTweetId|inReplyToUser|lang|likeCount|links|               media|mentionedUsers|place|quoteCount|quotedTweet|          rawContent|replyCount|retweetCount|retweetedTweet|              source|sourceLabel|        sourceUrl|                 url|viewCount|\n",
      "+----------+--------------------+--------+-------------------+-----------+--------------------+---------+-------------------+-------------------+----------------+-------------+----+---------+-----+--------------------+--------------+-----+----------+-----------+--------------------+----------+------------+--------------+--------------------+-----------+-----------------+--------------------+---------+\n",
      "|2862957343|snscrape.modules....|      []|1725519432179482712|       NULL|2023-11-17 14:21:...|[Bitcoin]|1725519432179482712|1725519432179482712|            NULL|         NULL|  en|        0|   []|{[], [{https://pb...|            []| NULL|         0|       NULL|#Bitcoin and Just...|         1|           0|          NULL|<a href=\"https://...|      IFTTT|https://ifttt.com|https://twitter.c...|       16|\n",
      "+----------+--------------------+--------+-------------------+-----------+--------------------+---------+-------------------+-------------------+----------------+-------------+----+---------+-----+--------------------+--------------+-----+----------+-----------+--------------------+----------+------------+--------------+--------------------+-----------+-----------------+--------------------+---------+\n",
      "only showing top 1 row\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "# _extract_tweets\n",
    "def _extract_tweets(profile_df):\n",
    "    tweet_df = profile_df \\\n",
    "                .select(F.col(\"user.id\").alias(\"user_id\"), \"tweet\") \\\n",
    "                .withColumn(\"tweet\", F.explode(\"tweet\")) \\\n",
    "                .select(\"user_id\", F.col(\"tweet.*\"))\n",
    "    \n",
    "    return tweet_df\n",
    "\n",
    "tweet_df = _extract_tweets(profile_df)\n",
    "tweet_df.show(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "090bc60a-3bf0-4809-a8ff-52bc7b47716b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 50:=============================>                            (1 + 1) / 2]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+--------------------+--------+--------------+-----------+--------------------+--------+----------+----------+----------------+-------------+----+---------+-----+------------+--------------+-----+----------+-----------+---------------+----------+------------+--------------+--------------------+------------------+------------------+--------------------+---------+\n",
      "|  user_id|               _type|cashtags|conversationId|coordinates|                date|hashtags|        id|    id_str|inReplyToTweetId|inReplyToUser|lang|likeCount|links|       media|mentionedUsers|place|quoteCount|quotedTweet|     rawContent|replyCount|retweetCount|retweetedTweet|              source|       sourceLabel|         sourceUrl|                 url|viewCount|\n",
      "+---------+--------------------+--------+--------------+-----------+--------------------+--------+----------+----------+----------------+-------------+----+---------+-----+------------+--------------+-----+----------+-----------+---------------+----------+------------+--------------+--------------------+------------------+------------------+--------------------+---------+\n",
      "|171446042|snscrape.modules....|      []|    1110302988|       NULL|2009-01-11 03:33:...|      []|1110302988|1110302988|            NULL|         NULL|  en|    76062|   []|{[], [], []}|            []| NULL|      7869|       NULL|Running bitcoin|      6975|       24871|          NULL|<a href=\"http://t...|Twitter Web Client|http://twitter.com|https://twitter.c...|     NULL|\n",
      "+---------+--------------------+--------+--------------+-----------+--------------------+--------+----------+----------+----------------+-------------+----+---------+-----+------------+--------------+-----+----------+-----------+---------------+----------+------------+--------------+--------------------+------------------+------------------+--------------------+---------+\n",
      "only showing top 1 row\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "# clean_tweets\n",
    "def clean_tweets(\n",
    "    profile_df\n",
    "):\n",
    "    # extract tweets from crawled data\n",
    "    tweet_df = _extract_tweets(profile_df) # <-- id -> user_id here\n",
    "    \n",
    "    # consistent & duplicate\n",
    "    tweet_df = tweet_df.dropDuplicates([\"id\"])\n",
    "    tweet_df = tweet_df.na.drop(subset=[\"id\"])\n",
    "    \n",
    "    # drop all-NULL rows\n",
    "    tweet_df = tweet_df.na.drop(\"all\")\n",
    "\n",
    "    return tweet_df\n",
    "\n",
    "tweet_df = clean_tweets(profile_df)\n",
    "tweet_df.show(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "579aeb17-3dbb-4c4c-b870-16a2c9891dd0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- user_id: long (nullable = true)\n",
      " |-- _type: string (nullable = true)\n",
      " |-- cashtags: array (nullable = true)\n",
      " |    |-- element: string (containsNull = true)\n",
      " |-- conversationId: long (nullable = true)\n",
      " |-- coordinates: string (nullable = true)\n",
      " |-- date: string (nullable = true)\n",
      " |-- hashtags: array (nullable = true)\n",
      " |    |-- element: string (containsNull = true)\n",
      " |-- id: long (nullable = true)\n",
      " |-- id_str: string (nullable = true)\n",
      " |-- inReplyToTweetId: long (nullable = true)\n",
      " |-- inReplyToUser: struct (nullable = true)\n",
      " |    |-- _type: string (nullable = true)\n",
      " |    |-- displayname: string (nullable = true)\n",
      " |    |-- id: long (nullable = true)\n",
      " |    |-- username: string (nullable = true)\n",
      " |-- lang: string (nullable = true)\n",
      " |-- likeCount: long (nullable = true)\n",
      " |-- links: array (nullable = true)\n",
      " |    |-- element: struct (containsNull = true)\n",
      " |    |    |-- tcourl: string (nullable = true)\n",
      " |    |    |-- text: string (nullable = true)\n",
      " |    |    |-- url: string (nullable = true)\n",
      " |-- media: struct (nullable = true)\n",
      " |    |-- animated: array (nullable = true)\n",
      " |    |    |-- element: struct (containsNull = true)\n",
      " |    |    |    |-- thumbnailUrl: string (nullable = true)\n",
      " |    |    |    |-- videoUrl: string (nullable = true)\n",
      " |    |-- photos: array (nullable = true)\n",
      " |    |    |-- element: struct (containsNull = true)\n",
      " |    |    |    |-- url: string (nullable = true)\n",
      " |    |-- videos: array (nullable = true)\n",
      " |    |    |-- element: struct (containsNull = true)\n",
      " |    |    |    |-- duration: long (nullable = true)\n",
      " |    |    |    |-- thumbnailUrl: string (nullable = true)\n",
      " |    |    |    |-- variants: array (nullable = true)\n",
      " |    |    |    |    |-- element: struct (containsNull = true)\n",
      " |    |    |    |    |    |-- bitrate: long (nullable = true)\n",
      " |    |    |    |    |    |-- contentType: string (nullable = true)\n",
      " |    |    |    |    |    |-- url: string (nullable = true)\n",
      " |    |    |    |-- views: string (nullable = true)\n",
      " |-- mentionedUsers: array (nullable = true)\n",
      " |    |-- element: struct (containsNull = true)\n",
      " |    |    |-- _type: string (nullable = true)\n",
      " |    |    |-- displayname: string (nullable = true)\n",
      " |    |    |-- id: long (nullable = true)\n",
      " |    |    |-- username: string (nullable = true)\n",
      " |-- place: struct (nullable = true)\n",
      " |    |-- country: string (nullable = true)\n",
      " |    |-- countryCode: string (nullable = true)\n",
      " |    |-- fullName: string (nullable = true)\n",
      " |    |-- id: string (nullable = true)\n",
      " |    |-- name: string (nullable = true)\n",
      " |    |-- type: string (nullable = true)\n",
      " |-- quoteCount: long (nullable = true)\n",
      " |-- rawContent: string (nullable = true)\n",
      " |-- replyCount: long (nullable = true)\n",
      " |-- retweetCount: long (nullable = true)\n",
      " |-- source: string (nullable = true)\n",
      " |-- sourceLabel: string (nullable = true)\n",
      " |-- sourceUrl: string (nullable = true)\n",
      " |-- url: string (nullable = true)\n",
      " |-- viewCount: long (nullable = true)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "tweet_df.drop(*[\"retweetedTweet\", \"quotedTweet\"]).printSchema()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63fc9c2f-8dc5-4fea-8bc0-6d74a5022cad",
   "metadata": {},
   "source": [
    "# Data Processing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edbec324-9022-4c15-89d4-163d3021983b",
   "metadata": {},
   "source": [
    "## Tweet Processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "18d70913-ca85-4020-8c67-44946c9a98a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 55:=============================>                            (1 + 1) / 2]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------+------------------+-----------------+------------------+----------+------------------+----+\n",
      "|user_id|        replyCount|     retweetCount|         likeCount|quoteCount|         viewCount|lang|\n",
      "+-------+------------------+-----------------+------------------+----------+------------------+----+\n",
      "| 816653|2.4761904761904763|5.380952380952381|12.666666666666666|       1.0|22106.761904761905|  en|\n",
      "+-------+------------------+-----------------+------------------+----------+------------------+----+\n",
      "only showing top 1 row\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "# preprocess_tweets\n",
    "def preprocess_tweets(\n",
    "    tweet_df,\n",
    "    keep_cols = [\"user_id\", \"replyCount\", \"retweetCount\", \n",
    "                 \"likeCount\", \"quoteCount\", \"viewCount\", \"lang\"]\n",
    "):\n",
    "    # drop unecessary fields for ML tasks\n",
    "    tweet_df = tweet_df.select(*keep_cols)\n",
    "    \n",
    "    # drop any-NULL rows\n",
    "    tweet_df = tweet_df.na.drop()\n",
    "    \n",
    "    # agg. by user\n",
    "    tweet_df = tweet_df.groupBy(\"user_id\").agg(\n",
    "        F.mean(\"replyCount\").alias(\"replyCount\"),\n",
    "        F.mean(\"retweetCount\").alias(\"retweetCount\"),\n",
    "        F.mean(\"likeCount\").alias(\"likeCount\"),\n",
    "        F.mean(\"quoteCount\").alias(\"quoteCount\"),\n",
    "        F.mean(\"viewCount\").alias(\"viewCount\"),\n",
    "        F.mode(\"lang\").alias(\"lang\")\n",
    "    )\n",
    "    \n",
    "    return tweet_df\n",
    "\n",
    "tweet_df = preprocess_tweets(tweet_df)\n",
    "tweet_df.show(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "74223ade-5189-4cab-9f02-7d3a60a68120",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- user_id: long (nullable = true)\n",
      " |-- replyCount: double (nullable = true)\n",
      " |-- retweetCount: double (nullable = true)\n",
      " |-- likeCount: double (nullable = true)\n",
      " |-- quoteCount: double (nullable = true)\n",
      " |-- viewCount: double (nullable = true)\n",
      " |-- lang: string (nullable = true)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "tweet_df.printSchema()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0786fbef-c00b-4b2c-8424-54afaf163a11",
   "metadata": {},
   "source": [
    "## User Processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "9f825b2f-5cef-4030-a894-530838980c34",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 94:=============================>                            (1 + 1) / 2]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+-------+------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+\n",
      "|               _type|blue|blueType|            created|    descriptionLinks|displayname|favouritesCount|followersCount|friendsCount|user_id|id_str|listedCount|     location|mediaCount|    profileBannerUrl|     profileImageUrl|protected|      rawDescription|statusesCount|                 url|  username|verified|      country|\n",
      "+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+-------+------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+\n",
      "|snscrape.modules....|true|Business|2007-03-07 01:27:09|[{https://t.co/J0...| TechCrunch|           7394|      10249257|         454| 816653|816653|     106178|North America|     49908|https://pbs.twimg...|https://pbs.twimg...|     NULL|Technology news a...|       262860|https://twitter.c...|TechCrunch|   false|United States|\n",
      "+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+-------+------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+\n",
      "only showing top 1 row\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "# preprocess_users\n",
    "def preprocess_users(\n",
    "    user_df,\n",
    "    neglect_cols = ['profileImageUrl','profileBannerUrl','descriptionLinks','_type', 'verified',\n",
    "                    'id_str','url','created', 'rawDescription', 'protected', 'blueType', 'displayname', 'country']\n",
    "):\n",
    "    # drop unecessary fields for ML tasks\n",
    "    #user_df = user_df.drop(*drop_cols)\n",
    "    \n",
    "    # drop rows with NaN values in columns other than \"location\"\n",
    "    process_cols = [col_name for col_name in user_df.columns if col_name not in ([\"location\"] + neglect_cols)]\n",
    "    user_df = user_df.na.drop(subset=process_cols)\n",
    "    \n",
    "    return user_df\n",
    "\n",
    "user_df = preprocess_users(user_df)\n",
    "user_df.show(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "8676c5f6-c570-4240-89b8-37da1dc5ccf1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- _type: string (nullable = true)\n",
      " |-- blue: boolean (nullable = true)\n",
      " |-- blueType: string (nullable = true)\n",
      " |-- created: timestamp (nullable = true)\n",
      " |-- descriptionLinks: array (nullable = true)\n",
      " |    |-- element: struct (containsNull = true)\n",
      " |    |    |-- tcourl: string (nullable = true)\n",
      " |    |    |-- text: string (nullable = true)\n",
      " |    |    |-- url: string (nullable = true)\n",
      " |-- displayname: string (nullable = true)\n",
      " |-- favouritesCount: long (nullable = true)\n",
      " |-- followersCount: long (nullable = true)\n",
      " |-- friendsCount: long (nullable = true)\n",
      " |-- user_id: long (nullable = true)\n",
      " |-- id_str: string (nullable = true)\n",
      " |-- listedCount: long (nullable = true)\n",
      " |-- location: string (nullable = true)\n",
      " |-- mediaCount: long (nullable = true)\n",
      " |-- profileBannerUrl: string (nullable = true)\n",
      " |-- profileImageUrl: string (nullable = true)\n",
      " |-- protected: string (nullable = true)\n",
      " |-- rawDescription: string (nullable = true)\n",
      " |-- statusesCount: long (nullable = true)\n",
      " |-- url: string (nullable = true)\n",
      " |-- username: string (nullable = true)\n",
      " |-- verified: boolean (nullable = true)\n",
      " |-- country: string (nullable = true)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "user_df.printSchema()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aadaa296-8378-4914-a814-8949894af228",
   "metadata": {},
   "source": [
    "## Data Integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "841b9a74-95c1-4393-aa0e-eb4f6ad9a5d3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+------------------+-----------------+------------------+----------+------------------+----+\n",
      "|user_id|               _type|blue|blueType|            created|    descriptionLinks|displayname|favouritesCount|followersCount|friendsCount|id_str|listedCount|     location|mediaCount|    profileBannerUrl|     profileImageUrl|protected|      rawDescription|statusesCount|                 url|  username|verified|      country|        replyCount|     retweetCount|         likeCount|quoteCount|         viewCount|lang|\n",
      "+-------+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+------------------+-----------------+------------------+----------+------------------+----+\n",
      "| 816653|snscrape.modules....|true|Business|2007-03-07 01:27:09|[{https://t.co/J0...| TechCrunch|           7394|      10249257|         454|816653|     106178|North America|     49908|https://pbs.twimg...|https://pbs.twimg...|     NULL|Technology news a...|       262860|https://twitter.c...|TechCrunch|   false|United States|2.4761904761904763|5.380952380952381|12.666666666666666|       1.0|22106.761904761905|  en|\n",
      "+-------+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+------------------+-----------------+------------------+----------+------------------+----+\n",
      "only showing top 1 row\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def merge_data(\n",
    "    df_user, df_tweet, \n",
    "    fillna_num=0, fillna_cat=\"unk\", left_join_on_cols=['user_id'],\n",
    "):\n",
    "    # perform a left join on specified columns\n",
    "    data_spark = df_user.join(df_tweet, on=left_join_on_cols, how='left')\n",
    "\n",
    "    # replace NaN values in columns for tweets\n",
    "    num_cols = [col_name for col_name in df_tweet.columns if col_name != \"location\" and col_name != \"lang\"]\n",
    "    cat_cols = [\"lang\"]\n",
    "    data_spark = data_spark.fillna(fillna_num, subset=num_cols)\n",
    "    data_spark = data_spark.fillna(fillna_cat, subset=cat_cols)\n",
    "                            \n",
    "    return data_spark\n",
    "\n",
    "data = merge_data(user_df, tweet_df)\n",
    "data.show(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "f36bed3b-bd02-496d-a88c-bed9fc9e326d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- user_id: long (nullable = true)\n",
      " |-- _type: string (nullable = true)\n",
      " |-- blue: boolean (nullable = true)\n",
      " |-- blueType: string (nullable = true)\n",
      " |-- created: timestamp (nullable = true)\n",
      " |-- descriptionLinks: array (nullable = true)\n",
      " |    |-- element: struct (containsNull = true)\n",
      " |    |    |-- tcourl: string (nullable = true)\n",
      " |    |    |-- text: string (nullable = true)\n",
      " |    |    |-- url: string (nullable = true)\n",
      " |-- displayname: string (nullable = true)\n",
      " |-- favouritesCount: long (nullable = true)\n",
      " |-- followersCount: long (nullable = true)\n",
      " |-- friendsCount: long (nullable = true)\n",
      " |-- id_str: string (nullable = true)\n",
      " |-- listedCount: long (nullable = true)\n",
      " |-- location: string (nullable = true)\n",
      " |-- mediaCount: long (nullable = true)\n",
      " |-- profileBannerUrl: string (nullable = true)\n",
      " |-- profileImageUrl: string (nullable = true)\n",
      " |-- protected: string (nullable = true)\n",
      " |-- rawDescription: string (nullable = true)\n",
      " |-- statusesCount: long (nullable = true)\n",
      " |-- url: string (nullable = true)\n",
      " |-- username: string (nullable = true)\n",
      " |-- verified: boolean (nullable = true)\n",
      " |-- country: string (nullable = true)\n",
      " |-- replyCount: double (nullable = false)\n",
      " |-- retweetCount: double (nullable = false)\n",
      " |-- likeCount: double (nullable = false)\n",
      " |-- quoteCount: double (nullable = false)\n",
      " |-- viewCount: double (nullable = false)\n",
      " |-- lang: string (nullable = false)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "data.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "fd11ba42-facd-4a29-a54d-73d6c0f48593",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "2149"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "313111d4-b120-4387-970b-e41bbfa636e3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1000"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.where(F.col(\"location\").isNotNull()).count()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b156a09-74a1-43c2-9ef4-bd21c8bb1821",
   "metadata": {},
   "source": [
    "# Data Transformation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "17886e77-5045-4e23-9fd1-f714091704b1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+------------------+-----------------+------------------+----------+------------------+----+\n",
      "|user_id|               _type|blue|blueType|            created|    descriptionLinks|displayname|favouritesCount|followersCount|friendsCount|id_str|listedCount|     location|mediaCount|    profileBannerUrl|     profileImageUrl|protected|      rawDescription|statusesCount|                 url|  username|verified|      country|        replyCount|     retweetCount|         likeCount|quoteCount|         viewCount|lang|\n",
      "+-------+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+------------------+-----------------+------------------+----------+------------------+----+\n",
      "| 816653|snscrape.modules....|true|Business|2007-03-07 01:27:09|[{https://t.co/J0...| TechCrunch|           7394|      10249257|         454|816653|     106178|North America|     49908|https://pbs.twimg...|https://pbs.twimg...|     NULL|Technology news a...|       262860|https://twitter.c...|TechCrunch|   false|United States|2.4761904761904763|5.380952380952381|12.666666666666666|       1.0|22106.761904761905|  en|\n",
      "+-------+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+------+-----------+-------------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+------------------+-----------------+------------------+----------+------------------+----+\n",
      "only showing top 1 row\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# load_data\n",
    "def extract_data(data):\n",
    "    # drop identity columns\n",
    "    #data = data.drop(\"user_id\", \"username\")\n",
    "\n",
    "    # split data into labeled and unlabeled based on the \"location\" column\n",
    "    labeled = data.filter(F.col(\"location\").isNotNull())\n",
    "    unlabeled = data.filter(F.col(\"location\").isNull())\n",
    "    \n",
    "    return labeled, unlabeled\n",
    "\n",
    "labeled, unlabeled = extract_data(data)\n",
    "labeled.show(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "baf17764-21b8-458a-b831-3557825e86a9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+-------+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+------+-----------+--------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+------------------+-----------------+------------------+----------+------------------+----+--------+------+------+------+------+------+------+\n",
      "|    all_feats_scaled|user_id|               _type|blue|blueType|            created|    descriptionLinks|displayname|favouritesCount|followersCount|friendsCount|id_str|listedCount|location|mediaCount|    profileBannerUrl|     profileImageUrl|protected|      rawDescription|statusesCount|                 url|  username|verified|      country|        replyCount|     retweetCount|         likeCount|quoteCount|         viewCount|lang|blue_tmp|lang_0|lang_1|lang_2|lang_3|lang_4|lang_5|\n",
      "+--------------------+-------+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+------+-----------+--------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+------------------+-----------------+------------------+----------+------------------+----+--------+------+------+------+------+------+------+\n",
      "|[-0.2317925306655...| 816653|snscrape.modules....|true|Business|2007-03-07 01:27:09|[{https://t.co/J0...| TechCrunch|           7394|      10249257|         454|816653|     106178|       4|     49908|https://pbs.twimg...|https://pbs.twimg...|     NULL|Technology news a...|       262860|https://twitter.c...|TechCrunch|   false|United States|2.4761904761904763|5.380952380952381|12.666666666666666|       1.0|22106.761904761905|  en|       1|     0|     0|     1|     0|     0|     1|\n",
      "+--------------------+-------+--------------------+----+--------+-------------------+--------------------+-----------+---------------+--------------+------------+------+-----------+--------+----------+--------------------+--------------------+---------+--------------------+-------------+--------------------+----------+--------+-------------+------------------+-----------------+------------------+----------+------------------+----+--------+------+------+------+------+------+------+\n",
      "only showing top 1 row\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# encode_data\n",
    "LABEL_MAP = {'Africa': 0, 'Antarctica': 1, 'Asia': 2, 'Europe': 3, 'North America': 4, 'Oceania': 5, 'South America': 6, 'Arctic': 7}\n",
    "def transform_data(\n",
    "    df,\n",
    "    neglect_cols = ['profileImageUrl','profileBannerUrl','descriptionLinks','_type', 'verified',\n",
    "                    'id_str','url','created', 'rawDescription', 'protected', 'blueType', 'displayname', 'country']\n",
    "):\n",
    "    # encode categorical features\n",
    "    bin_len = 6\n",
    "    label_map = F.create_map([F.lit(x) for x in chain(*LABEL_MAP.items())])\n",
    "    lang_map = F.create_map([F.lit(x) for x in chain(*LANG_MAP.items())])\n",
    "    df = df \\\n",
    "            .withColumn(\"location\", label_map[F.col(\"location\")]) \\\n",
    "            .withColumn(\"blue_tmp\", F.when(F.col(\"blue\") == True, 1).otherwise(0)) \\\n",
    "            .withColumn(\"lang_tmp\", lang_map[F.col(\"lang\")])\n",
    "    df = df.withColumn(\"lang_tmp\", F.lpad(F.bin(df[\"lang_tmp\"]), bin_len, '0')) # binary encoding\n",
    "    \n",
    "    split_lang = F.split(df['lang_tmp'], \"\")\n",
    "    for i in range(bin_len):\n",
    "        df = df \\\n",
    "                .withColumn(f'lang_{i}', split_lang.getItem(i).cast(T.IntegerType()))\n",
    "    df = df.drop(\"lang_tmp\")\n",
    "    \n",
    "    # enumerate types of features\n",
    "    num_cols = [name for name, datatype in df.dtypes if (not datatype.startswith('string')) and (not name.startswith('lang')) \n",
    "                and (not name.startswith('blue')) and name not in ([\"location\", \"user_id\"] + neglect_cols)]\n",
    "    cat_cols = [\"blue_tmp\"] + [name for name, datatype in df.dtypes if name.startswith('lang') and name not in (neglect_cols + [\"lang\"])]\n",
    "    \n",
    "    # define processing ops\n",
    "    features = num_cols + cat_cols\n",
    "    \n",
    "    num_assembler = VectorAssembler(inputCols=num_cols, outputCol=\"num_feats\")\n",
    "    scaler = StandardScaler(inputCol=\"num_feats\", outputCol=\"num_feats_scaled\", withStd=True, withMean=True)\n",
    "    \n",
    "    cat_assembler = VectorAssembler(inputCols=cat_cols, outputCol=\"cat_feats\", \n",
    "                                    handleInvalid =\"keep\")\n",
    "    \n",
    "    all_assembler = VectorAssembler(inputCols= [\"num_feats_scaled\", \"cat_feats\"], outputCol=\"all_feats_scaled\", \n",
    "                                    handleInvalid =\"keep\")\n",
    "\n",
    "    # assemble a pipeline\n",
    "    pipeline = Pipeline(stages=[num_assembler, scaler, cat_assembler, all_assembler])\n",
    "\n",
    "    # fit transform on pipeline\n",
    "    encode_data = pipeline.fit(df).transform(df)\n",
    "\n",
    "    # select features + targets\n",
    "    selected_cols = [\"all_feats_scaled\"] + df.columns\n",
    "    encode_data = encode_data.select(*selected_cols)\n",
    "\n",
    "    return encode_data\n",
    "\n",
    "labeled_prep = transform_data(labeled)\n",
    "labeled_prep.show(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5668555-1be7-4331-8ee5-83c1a4f1b0ec",
   "metadata": {},
   "source": [
    "# Modelling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "61c2073d-bdd2-408f-b03d-e2e7453e115b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validating model using weightedFMeasure\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "23/12/23 18:03:20 WARN DAGScheduler: Broadcasting large task binary with size 1031.1 KiB\n",
      "23/12/23 18:03:20 WARN DAGScheduler: Broadcasting large task binary with size 1311.8 KiB\n",
      "23/12/23 18:03:20 WARN DAGScheduler: Broadcasting large task binary with size 1572.8 KiB\n",
      "23/12/23 18:03:21 WARN DAGScheduler: Broadcasting large task binary with size 1794.2 KiB\n",
      "23/12/23 18:03:21 WARN DAGScheduler: Broadcasting large task binary with size 1929.8 KiB\n",
      "23/12/23 18:03:29 WARN DAGScheduler: Broadcasting large task binary with size 1677.3 KiB\n",
      "23/12/23 18:03:37 WARN DAGScheduler: Broadcasting large task binary with size 1233.5 KiB\n",
      "23/12/23 18:03:38 WARN DAGScheduler: Broadcasting large task binary with size 1483.9 KiB\n",
      "23/12/23 18:03:38 WARN DAGScheduler: Broadcasting large task binary with size 1715.1 KiB\n",
      "23/12/23 18:03:38 WARN DAGScheduler: Broadcasting large task binary with size 1878.5 KiB\n",
      "23/12/23 18:03:45 WARN DAGScheduler: Broadcasting large task binary with size 1621.0 KiB\n",
      "23/12/23 18:04:02 WARN DAGScheduler: Broadcasting large task binary with size 1256.3 KiB\n",
      "23/12/23 18:04:02 WARN DAGScheduler: Broadcasting large task binary with size 1627.0 KiB\n",
      "23/12/23 18:04:02 WARN DAGScheduler: Broadcasting large task binary with size 2021.5 KiB\n",
      "23/12/23 18:04:03 WARN DAGScheduler: Broadcasting large task binary with size 2.3 MiB\n",
      "23/12/23 18:04:03 WARN DAGScheduler: Broadcasting large task binary with size 2.6 MiB\n",
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validating model using weightedPrecision\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "23/12/23 18:04:18 WARN DAGScheduler: Broadcasting large task binary with size 1049.4 KiB\n",
      "23/12/23 18:04:18 WARN DAGScheduler: Broadcasting large task binary with size 1347.5 KiB\n",
      "23/12/23 18:04:19 WARN DAGScheduler: Broadcasting large task binary with size 1628.1 KiB\n",
      "23/12/23 18:04:19 WARN DAGScheduler: Broadcasting large task binary with size 1842.1 KiB\n",
      "23/12/23 18:04:20 WARN DAGScheduler: Broadcasting large task binary with size 1915.3 KiB\n",
      "23/12/23 18:04:26 WARN DAGScheduler: Broadcasting large task binary with size 1716.0 KiB\n",
      "23/12/23 18:04:35 WARN DAGScheduler: Broadcasting large task binary with size 1218.1 KiB\n",
      "23/12/23 18:04:36 WARN DAGScheduler: Broadcasting large task binary with size 1448.7 KiB\n",
      "23/12/23 18:04:36 WARN DAGScheduler: Broadcasting large task binary with size 1644.5 KiB\n",
      "23/12/23 18:04:36 WARN DAGScheduler: Broadcasting large task binary with size 1780.9 KiB\n",
      "23/12/23 18:04:43 WARN DAGScheduler: Broadcasting large task binary with size 1574.1 KiB\n",
      "23/12/23 18:04:59 WARN DAGScheduler: Broadcasting large task binary with size 1245.3 KiB\n",
      "23/12/23 18:04:59 WARN DAGScheduler: Broadcasting large task binary with size 1613.5 KiB\n",
      "23/12/23 18:04:59 WARN DAGScheduler: Broadcasting large task binary with size 2003.4 KiB\n",
      "23/12/23 18:05:00 WARN DAGScheduler: Broadcasting large task binary with size 2.3 MiB\n",
      "23/12/23 18:05:00 WARN DAGScheduler: Broadcasting large task binary with size 2.7 MiB\n",
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validating model using weightedRecall\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "23/12/23 18:05:14 WARN DAGScheduler: Broadcasting large task binary with size 1049.4 KiB\n",
      "23/12/23 18:05:14 WARN DAGScheduler: Broadcasting large task binary with size 1347.5 KiB\n",
      "23/12/23 18:05:15 WARN DAGScheduler: Broadcasting large task binary with size 1628.1 KiB\n",
      "23/12/23 18:05:15 WARN DAGScheduler: Broadcasting large task binary with size 1842.1 KiB\n",
      "23/12/23 18:05:16 WARN DAGScheduler: Broadcasting large task binary with size 1915.3 KiB\n",
      "23/12/23 18:05:22 WARN DAGScheduler: Broadcasting large task binary with size 1716.0 KiB\n",
      "23/12/23 18:05:30 WARN DAGScheduler: Broadcasting large task binary with size 1233.5 KiB\n",
      "23/12/23 18:05:30 WARN DAGScheduler: Broadcasting large task binary with size 1483.9 KiB\n",
      "23/12/23 18:05:30 WARN DAGScheduler: Broadcasting large task binary with size 1715.1 KiB\n",
      "23/12/23 18:05:31 WARN DAGScheduler: Broadcasting large task binary with size 1878.5 KiB\n",
      "23/12/23 18:05:37 WARN DAGScheduler: Broadcasting large task binary with size 1621.0 KiB\n",
      "23/12/23 18:05:53 WARN DAGScheduler: Broadcasting large task binary with size 1245.3 KiB\n",
      "23/12/23 18:05:54 WARN DAGScheduler: Broadcasting large task binary with size 1613.5 KiB\n",
      "23/12/23 18:05:54 WARN DAGScheduler: Broadcasting large task binary with size 2003.4 KiB\n",
      "23/12/23 18:05:55 WARN DAGScheduler: Broadcasting large task binary with size 2.3 MiB\n",
      "23/12/23 18:05:55 WARN DAGScheduler: Broadcasting large task binary with size 2.7 MiB\n",
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validating model using accuracy\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "23/12/23 18:06:10 WARN DAGScheduler: Broadcasting large task binary with size 1031.1 KiB\n",
      "23/12/23 18:06:10 WARN DAGScheduler: Broadcasting large task binary with size 1311.8 KiB\n",
      "23/12/23 18:06:10 WARN DAGScheduler: Broadcasting large task binary with size 1572.8 KiB\n",
      "23/12/23 18:06:11 WARN DAGScheduler: Broadcasting large task binary with size 1794.2 KiB\n",
      "23/12/23 18:06:11 WARN DAGScheduler: Broadcasting large task binary with size 1929.8 KiB\n",
      "23/12/23 18:06:17 WARN DAGScheduler: Broadcasting large task binary with size 1677.3 KiB\n",
      "23/12/23 18:06:25 WARN DAGScheduler: Broadcasting large task binary with size 1218.1 KiB\n",
      "23/12/23 18:06:25 WARN DAGScheduler: Broadcasting large task binary with size 1448.7 KiB\n",
      "23/12/23 18:06:26 WARN DAGScheduler: Broadcasting large task binary with size 1644.5 KiB\n",
      "23/12/23 18:06:26 WARN DAGScheduler: Broadcasting large task binary with size 1780.9 KiB\n",
      "23/12/23 18:06:32 WARN DAGScheduler: Broadcasting large task binary with size 1574.1 KiB\n",
      "23/12/23 18:06:48 WARN DAGScheduler: Broadcasting large task binary with size 1256.3 KiB\n",
      "23/12/23 18:06:48 WARN DAGScheduler: Broadcasting large task binary with size 1627.0 KiB\n",
      "23/12/23 18:06:49 WARN DAGScheduler: Broadcasting large task binary with size 2021.5 KiB\n",
      "23/12/23 18:06:49 WARN DAGScheduler: Broadcasting large task binary with size 2.3 MiB\n",
      "23/12/23 18:06:50 WARN DAGScheduler: Broadcasting large task binary with size 2.6 MiB\n",
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'weightedFMeasure': 0.9631500001432405,\n",
       " 'weightedPrecision': 0.9678116168626494,\n",
       " 'weightedRecall': 0.9701770177017702,\n",
       " 'accuracy': 0.9661966196619662}"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# initialize RandomForestClassifier\n",
    "rf_clf = RandomForestClassifier(labelCol=\"location\", featuresCol=\"all_feats_scaled\", \n",
    "                                numTrees=95, maxDepth=10, seed=42)\n",
    "scoreboard = {}\n",
    "\n",
    "# set up cross-validation\n",
    "for metric in [\"weightedFMeasure\", \"weightedPrecision\", \"weightedRecall\", \"accuracy\"]:\n",
    "    print(f\"Validating model using {metric}\")\n",
    "    paramGrid = ParamGridBuilder().build()\n",
    "    cross_validator = CrossValidator(\n",
    "        estimator=rf_clf,\n",
    "        estimatorParamMaps=paramGrid, \n",
    "        evaluator=MulticlassClassificationEvaluator(labelCol=\"location\", metricName=metric), \n",
    "        numFolds=2\n",
    "    )\n",
    "\n",
    "    # cross-validate on labeled data\n",
    "    model = cross_validator.fit(labeled_prep)\n",
    "\n",
    "    # avg. score from the cross-validated model\n",
    "    scoreboard[metric] = model.avgMetrics[0]\n",
    "    \n",
    "# show results\n",
    "scoreboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "100bfb4f-956a-40f3-9c88-38b2368cdfa0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "23/12/24 07:35:39 WARN DAGScheduler: Broadcasting large task binary with size 1396.5 KiB\n",
      "23/12/24 07:35:40 WARN DAGScheduler: Broadcasting large task binary with size 1951.0 KiB\n",
      "23/12/24 07:35:40 WARN DAGScheduler: Broadcasting large task binary with size 2.6 MiB\n",
      "23/12/24 07:35:41 WARN DAGScheduler: Broadcasting large task binary with size 3.3 MiB\n",
      "23/12/24 07:35:42 WARN DAGScheduler: Broadcasting large task binary with size 4.1 MiB\n",
      "23/12/24 07:35:56 WARN TaskSetManager: Stage 203 contains a task of very large size (1501 KiB). The maximum recommended task size is 1000 KiB.\n",
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "# re-fit model on full data\n",
    "rf_clf = RandomForestClassifier(labelCol=\"location\", featuresCol=\"all_feats_scaled\", \n",
    "                                numTrees=95, maxDepth=10, seed=42)\n",
    "rf_clf = rf_clf.fit(labeled_prep)\n",
    "\n",
    "# save the model to bucket\n",
    "rf_clf.save(\"gs://it4043e-it5384/it4043e/it4043e_group15_problem1/output/models/spark_rf-n_estimators_95-max_depth_11_noscale\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3b5add6-b24c-4119-bb15-5f516c2a9bdb",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "68ec257d-b493-4e68-bff2-3ee495e623ca",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "RandomForestClassificationModel: uid=RandomForestClassifier_a08dd87bace2, numTrees=95, numClasses=8, numFeatures=18"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load the model\n",
    "rf_clf = RandomForestClassificationModel.load(\"gs://it4043e-it5384/it4043e/it4043e_group15_problem1/output/models/spark_rf-n_estimators_95-max_depth_11_noscale\")\n",
    "rf_clf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "afbd608c-3460-4529-886b-c8dbc9142a65",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "23/12/24 16:30:50 WARN DAGScheduler: Broadcasting large task binary with size 3.4 MiB\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+-------+--------------------+-----+--------+-------------------+--------------------+------------------------+---------------+--------------+------------+-------+-----------+--------+----------+--------------------+--------------------+---------+-----------------------+-------------+--------------------+--------+--------+-------+----------+------------+---------+----------+---------+----+--------+------+------+------+------+------+------+--------------------+--------------------+----------+\n",
      "|    all_feats_scaled|user_id|               _type| blue|blueType|            created|    descriptionLinks|             displayname|favouritesCount|followersCount|friendsCount| id_str|listedCount|location|mediaCount|    profileBannerUrl|     profileImageUrl|protected|         rawDescription|statusesCount|                 url|username|verified|country|replyCount|retweetCount|likeCount|quoteCount|viewCount|lang|blue_tmp|lang_0|lang_1|lang_2|lang_3|lang_4|lang_5|       rawPrediction|         probability|prediction|\n",
      "+--------------------+-------+--------------------+-----+--------+-------------------+--------------------+------------------------+---------------+--------------+------------+-------+-----------+--------+----------+--------------------+--------------------+---------+-----------------------+-------------+--------------------+--------+--------+-------+----------+------------+---------+----------+---------+----+--------+------+------+------+------+------+------+--------------------+--------------------+----------+\n",
      "|[0.27413592883643...|6657682|snscrape.modules....|false|    NULL|2007-06-08 04:44:57|[{https://t.co/yS...|飯塚 友裕 - Web3 POD ...|          22480|          2072|        1060|6657682|         25|    NULL|       629|https://pbs.twimg...|https://pbs.twimg...|     NULL|Codablecash 代表。\\n...|        17148|https://twitter.c...|  iizuka|   false|   NULL|      2.36|       11.34|    43.22|      1.26|  3187.96|  ja|       0|     0|     1|     0|     1|     1|     1|[2.0,0.0,40.25,28...|[0.02105263157894...|       2.0|\n",
      "+--------------------+-------+--------------------+-----+--------+-------------------+--------------------+------------------------+---------------+--------------+------------+-------+-----------+--------+----------+--------------------+--------------------+---------+-----------------------+-------------+--------------------+--------+--------+-------+----------+------------+---------+----------+---------+----+--------+------+------+------+------+------+------+--------------------+--------------------+----------+\n",
      "only showing top 1 row\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Infer the unlabeled\n",
    "unlabeled_prep = transform_data(unlabeled)\n",
    "unlabeled_pred = rf_clf.transform(unlabeled_prep)\n",
    "unlabeled_pred.show(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "85b8eda0-0d33-4d2d-a51e-bb973158bd72",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- user_id: long (nullable = true)\n",
      " |-- _type: string (nullable = true)\n",
      " |-- blue: boolean (nullable = false)\n",
      " |-- blueType: string (nullable = true)\n",
      " |-- created: timestamp (nullable = true)\n",
      " |-- descriptionLinks: array (nullable = true)\n",
      " |    |-- element: struct (containsNull = true)\n",
      " |    |    |-- tcourl: string (nullable = true)\n",
      " |    |    |-- text: string (nullable = true)\n",
      " |    |    |-- url: string (nullable = true)\n",
      " |-- displayname: string (nullable = true)\n",
      " |-- favouritesCount: long (nullable = true)\n",
      " |-- followersCount: long (nullable = true)\n",
      " |-- friendsCount: long (nullable = true)\n",
      " |-- id_str: string (nullable = true)\n",
      " |-- listedCount: long (nullable = true)\n",
      " |-- mediaCount: long (nullable = true)\n",
      " |-- profileBannerUrl: string (nullable = true)\n",
      " |-- profileImageUrl: string (nullable = true)\n",
      " |-- protected: string (nullable = true)\n",
      " |-- rawDescription: string (nullable = true)\n",
      " |-- statusesCount: long (nullable = true)\n",
      " |-- url: string (nullable = true)\n",
      " |-- username: string (nullable = true)\n",
      " |-- verified: boolean (nullable = true)\n",
      " |-- country: string (nullable = true)\n",
      " |-- replyCount: double (nullable = false)\n",
      " |-- retweetCount: double (nullable = false)\n",
      " |-- likeCount: double (nullable = false)\n",
      " |-- quoteCount: double (nullable = false)\n",
      " |-- viewCount: double (nullable = false)\n",
      " |-- lang: string (nullable = false)\n",
      " |-- location: string (nullable = true)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# invert the prediction\n",
    "INV_LABEL_MAP = {v: k for k, v in LABEL_MAP.items()}\n",
    "inv_label_map = F.create_map([F.lit(x) for x in chain(*INV_LABEL_MAP.items())])\n",
    "unlabeled_show = unlabeled_pred \\\n",
    "                    .withColumn(\"prediction\", inv_label_map[F.col(\"prediction\")]) \\\n",
    "                    .select(*unlabeled.columns, \"prediction\") \\\n",
    "                    .drop(\"location\") \\\n",
    "                    .withColumnRenamed(\"prediction\", \"location\") \\\n",
    "                    .withColumn(\"blue\", F.when(F.col(\"blue\") == 1, True).otherwise(False))\n",
    "unlabeled_show.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "d0383d3f-19d5-4bdf-aa0b-dfd8dcf2e0ad",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- user_id: long (nullable = true)\n",
      " |-- _type: string (nullable = true)\n",
      " |-- blue: boolean (nullable = true)\n",
      " |-- blueType: string (nullable = true)\n",
      " |-- created: timestamp (nullable = true)\n",
      " |-- descriptionLinks: array (nullable = true)\n",
      " |    |-- element: struct (containsNull = true)\n",
      " |    |    |-- tcourl: string (nullable = true)\n",
      " |    |    |-- text: string (nullable = true)\n",
      " |    |    |-- url: string (nullable = true)\n",
      " |-- displayname: string (nullable = true)\n",
      " |-- favouritesCount: long (nullable = true)\n",
      " |-- followersCount: long (nullable = true)\n",
      " |-- friendsCount: long (nullable = true)\n",
      " |-- id_str: string (nullable = true)\n",
      " |-- listedCount: long (nullable = true)\n",
      " |-- mediaCount: long (nullable = true)\n",
      " |-- profileBannerUrl: string (nullable = true)\n",
      " |-- profileImageUrl: string (nullable = true)\n",
      " |-- protected: string (nullable = true)\n",
      " |-- rawDescription: string (nullable = true)\n",
      " |-- statusesCount: long (nullable = true)\n",
      " |-- url: string (nullable = true)\n",
      " |-- username: string (nullable = true)\n",
      " |-- verified: boolean (nullable = true)\n",
      " |-- country: string (nullable = true)\n",
      " |-- replyCount: double (nullable = false)\n",
      " |-- retweetCount: double (nullable = false)\n",
      " |-- likeCount: double (nullable = false)\n",
      " |-- quoteCount: double (nullable = false)\n",
      " |-- viewCount: double (nullable = false)\n",
      " |-- lang: string (nullable = false)\n",
      " |-- location: string (nullable = true)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "labeled.select(*unlabeled_show.columns).printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "9bd2caee-c7ce-4204-b1fe-c4379d6e500f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "2149"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# union data for report\n",
    "processed_data = labeled.select(*unlabeled_show.columns).union(unlabeled_show)\n",
    "processed_data.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb5183c0-9ad1-4e72-8264-920703c6d32d",
   "metadata": {},
   "source": [
    "# Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "a7f2497d-d04b-47e9-a812-23ed27abc271",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_3651529/791681254.py:8: DeprecationWarning: The 'http_auth' parameter is deprecated. Use 'basic_auth' or 'bearer_auth' parameters instead\n",
      "  es = Elasticsearch(['http://34.143.255.36:9200/'], http_auth=('elastic', 'elastic2023'))\n",
      "23/12/24 16:32:04 WARN DAGScheduler: Broadcasting large task binary with size 3.5 MiB\n",
      "23/12/24 16:32:05 WARN DAGScheduler: Broadcasting large task binary with size 3.5 MiB\n",
      "[Stage 217:>                                                        (0 + 1) / 1]\r"
     ]
    }
   ],
   "source": [
    "from elasticsearch import Elasticsearch\n",
    "from elasticsearch import helpers\n",
    "\n",
    "def create_es_index(\n",
    "    df, \n",
    "    idx_name\n",
    "):\n",
    "    es = Elasticsearch(['http://34.143.255.36:9200/'], http_auth=('elastic', 'elastic2023'))\n",
    "    \n",
    "    def index_data_to_elasticsearch(df):\n",
    "        actions = []\n",
    "        for row in df.rdd.toLocalIterator():\n",
    "            action = {\n",
    "                \"_index\": idx_name,\n",
    "                \"_source\": {\n",
    "                    col: row[col] for col in df.columns\n",
    "                }\n",
    "            }\n",
    "            actions.append(action)\n",
    "        helpers.bulk(es, actions)\n",
    "        \n",
    "    index_data_to_elasticsearch(df)\n",
    "\n",
    "create_es_index(processed_data, \"group15-spark-ml\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "group15",
   "language": "python",
   "name": "group15"
  },
  "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.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}