data_manipulation_ARDL_pyspark.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data manipulation for ARDL analysis\n",
    "\n",
    "In this notebook, we will use Pyspark and aggregate monthly data to quarterly data so that we can merge with GDP data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Current session configs: <tt>{'conf': {'spark.pyspark.python': 'python3', 'spark.pyspark.virtualenv.enabled': 'true', 'spark.pyspark.virtualenv.type': 'native', 'spark.pyspark.virtualenv.bin.path': '/usr/bin/virtualenv'}, 'proxyUser': 'jovyan', 'kind': 'pyspark'}</tt><br>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "No active sessions."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%configure -f\n",
    "{\n",
    "    \"conf\": {\n",
    "        \"spark.pyspark.python\": \"python3\",\n",
    "        \"spark.pyspark.virtualenv.enabled\": \"true\",\n",
    "        \"spark.pyspark.virtualenv.type\":\"native\",\n",
    "        \"spark.pyspark.virtualenv.bin.path\":\"/usr/bin/virtualenv\"\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting Spark application\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr><th>ID</th><th>YARN Application ID</th><th>Kind</th><th>State</th><th>Spark UI</th><th>Driver log</th><th>Current session?</th></tr><tr><td>0</td><td>application_1684800927608_0002</td><td>pyspark</td><td>idle</td><td><a target=\"_blank\" href=\"http://ip-172-31-17-221.ec2.internal:20888/proxy/application_1684800927608_0002/\">Link</a></td><td><a target=\"_blank\" href=\"http://ip-172-31-29-118.ec2.internal:8042/node/containerlogs/container_1684800927608_0002_01_000001/livy\">Link</a></td><td>✔</td></tr></table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SparkSession available as 'spark'.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from pyspark.sql import functions as F\n",
    "from pyspark.sql.window import Window"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load exchange rate and trade data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ex = spark.read.parquet(\"s3://trade-final-project-bucket/dataset/exchangerate2.parquet\").persist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ex = (ex.withColumnRenamed(\"Code\",\"country\")\n",
    "        .withColumnRenamed(\"Time\",\"year_month\")\n",
    "        .withColumnRenamed(\"value\",\"ex\")\n",
    "     )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- Country_name: string (nullable = true)\n",
      " |-- country: long (nullable = true)\n",
      " |-- Country_name2: string (nullable = true)\n",
      " |-- Area: string (nullable = true)\n",
      " |-- year_month: string (nullable = true)\n",
      " |-- ex: double (nullable = true)"
     ]
    }
   ],
   "source": [
    "ex.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+-------+-------------+--------------------+----------+------------------+\n",
      "|        Country_name|country|Country_name2|                Area|year_month|                ex|\n",
      "+--------------------+-------+-------------+--------------------+----------+------------------+\n",
      "|Afghanistan, Isla...|    130|  Afghanistan|                Asia|    198801|0.3092632225995748|\n",
      "|             Albania|    229|      Albania|Central_and_East_...|    198801|              null|\n",
      "+--------------------+-------+-------------+--------------------+----------+------------------+\n",
      "only showing top 2 rows"
     ]
    }
   ],
   "source": [
    "ex.show(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data = spark.read.parquet(\"s3://trade-final-project-bucket/dataset/trades.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- year_month: long (nullable = true)\n",
      " |-- export_import: long (nullable = true)\n",
      " |-- country: long (nullable = true)\n",
      " |-- hs9: long (nullable = true)\n",
      " |-- q1: long (nullable = true)\n",
      " |-- q2: long (nullable = true)\n",
      " |-- value: long (nullable = true)\n",
      " |-- hs6: string (nullable = true)\n",
      " |-- hs2: string (nullable = true)"
     ]
    }
   ],
   "source": [
    "data.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data = data.join(ex, on=['country', 'year_month'], how='inner')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data = data.withColumn('year_month_str', F.col('year_month').cast('string'))\n",
    "data = (data.withColumn('quarter', F.quarter(F.to_date('year_month_str', 'yyyyMM')))\n",
    "            .withColumn('year', F.year(F.to_date('year_month_str', 'yyyyMM')))\n",
    "       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- country: long (nullable = true)\n",
      " |-- year_month: long (nullable = true)\n",
      " |-- export_import: long (nullable = true)\n",
      " |-- hs9: long (nullable = true)\n",
      " |-- q1: long (nullable = true)\n",
      " |-- q2: long (nullable = true)\n",
      " |-- value: long (nullable = true)\n",
      " |-- hs6: string (nullable = true)\n",
      " |-- hs2: string (nullable = true)\n",
      " |-- Country_name: string (nullable = true)\n",
      " |-- Country_name2: string (nullable = true)\n",
      " |-- Area: string (nullable = true)\n",
      " |-- ex: double (nullable = true)\n",
      " |-- year_month_str: string (nullable = true)\n",
      " |-- quarter: integer (nullable = true)\n",
      " |-- year: integer (nullable = true)"
     ]
    }
   ],
   "source": [
    "data.printSchema()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Aggregate from monthly data to quarterly data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data_g = (data.groupby(\"year\", \"quarter\", \"export_import\", \"country\", \"Country_name\", \"hs2\")\n",
    "          .agg(F.mean(\"ex\").alias(\"mean_ex\"), F.sum(\"value\").alias(\"sum_value\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- year: integer (nullable = true)\n",
      " |-- quarter: integer (nullable = true)\n",
      " |-- export_import: long (nullable = true)\n",
      " |-- country: long (nullable = true)\n",
      " |-- Country_name: string (nullable = true)\n",
      " |-- hs2: string (nullable = true)\n",
      " |-- mean_ex: double (nullable = true)\n",
      " |-- sum_value: long (nullable = true)"
     ]
    }
   ],
   "source": [
    "data_g.printSchema()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load GDP data and merge with previous dataframe "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gdp = spark.read.parquet(\"s3://trade-final-project-bucket/dataset/gdp_quarter.parquet\").persist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gdp = (gdp.withColumn(\"quarter\", F.substring(F.col(\"Time\"), 2, 1).cast(\"integer\"))\n",
    "          .withColumn(\"year\", F.substring(F.col(\"Time\"), 4, 4).cast(\"integer\"))\n",
    "          .withColumnRenamed(\"value\", \"gdp\")\n",
    "      )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Row(Country_name='Argentina', Code=413, Country_name2='Argentina', Area='Middle_and_South_America', Time='Q1_1988', gdp=None, quarter=1, year=1988), Row(Country_name='Armenia, Rep. of', Code=151, Country_name2='Armenia', Area='Central_and_East_Europe_Russia', Time='Q1_1988', gdp=None, quarter=1, year=1988)]"
     ]
    }
   ],
   "source": [
    "gdp.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gdp = gdp.withColumnRenamed(\"Code\",\"country\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data_g = data_g.join(gdp, on=['country', 'quarter', 'year'], how='inner')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data_g = (data_g.withColumnRenamed(\"mean_ex\",\"ex\")\n",
    "          .withColumnRenamed(\"sum_value\",\"value\")\n",
    "         )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data_g = data_g.withColumn('year_quarter', F.concat(F.col(\"year\"), F.lit(\"_\"), F.col(\"quarter\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Row(country=205, quarter=1, year=2002, export_import=1, Country_name='United Kingdom', hs2='85', ex=0.005314876791761418, value=80384260, Country_name='United Kingdom', Country_name2='United_Kingdom', Area='Western_Europe', Time='Q1_2002', gdp=292210.0, year_quarter='2002_1'), Row(country=207, quarter=4, year=2003, export_import=1, Country_name='Netherlands, The', hs2='84', ex=None, value=146989643, Country_name='Netherlands, The', Country_name2='Netherlands', Area='Western_Europe', Time='Q4_2003', gdp=129623.89, year_quarter='2003_4')]"
     ]
    }
   ],
   "source": [
    "data_g.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- country: long (nullable = true)\n",
      " |-- quarter: integer (nullable = true)\n",
      " |-- year: integer (nullable = true)\n",
      " |-- export_import: long (nullable = true)\n",
      " |-- Country_name: string (nullable = true)\n",
      " |-- hs2: string (nullable = true)\n",
      " |-- ex: double (nullable = true)\n",
      " |-- value: long (nullable = true)\n",
      " |-- Country_name: string (nullable = true)\n",
      " |-- Country_name2: string (nullable = true)\n",
      " |-- Area: string (nullable = true)\n",
      " |-- Time: string (nullable = true)\n",
      " |-- gdp: double (nullable = true)\n",
      " |-- year_quarter: string (nullable = true)"
     ]
    }
   ],
   "source": [
    "data_g.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data_g = (data_g.withColumn(\"value_log\", F.log(F.col(\"value\")))\n",
    "                .withColumn(\"ex_log\", F.log(F.col(\"ex\")))\n",
    "                .withColumn(\"gdp_log\", F.log(F.col(\"gdp\")))\n",
    "         )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data_g = (data_g[[\"country\", \"year_quarter\", \"export_import\",  \"hs2\", \"value_log\", \"ex_log\", \"gdp_log\"]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- country: long (nullable = true)\n",
      " |-- year_quarter: string (nullable = true)\n",
      " |-- export_import: long (nullable = true)\n",
      " |-- hs2: string (nullable = true)\n",
      " |-- value_log: double (nullable = true)\n",
      " |-- ex_log: double (nullable = true)\n",
      " |-- gdp_log: double (nullable = true)"
     ]
    }
   ],
   "source": [
    "data_g.printSchema()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add lags and difference columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "windowSpec = Window.partitionBy(\"country\",\"hs2\",\"export_import\").orderBy(\"year_quarter\")\n",
    "\n",
    "# Create the lag variable\n",
    "df_with_lag = (data_g.withColumn(\"value_log_L1\", F.lag(\"value_log\").over(windowSpec))\n",
    "                 .withColumn(\"value_log_L2\", F.lag(\"value_log\",2).over(windowSpec))\n",
    "                 .withColumn(\"value_log_L3\", F.lag(\"value_log\",3).over(windowSpec))\n",
    "                 .withColumn(\"value_log_L4\", F.lag(\"value_log\",4).over(windowSpec))\n",
    "                 .withColumn(\"ex_log_L1\", F.lag(\"ex_log\").over(windowSpec))\n",
    "                 .withColumn(\"ex_log_L2\", F.lag(\"ex_log\",2).over(windowSpec))\n",
    "                 .withColumn(\"ex_log_L3\", F.lag(\"ex_log\",3).over(windowSpec))\n",
    "                 .withColumn(\"ex_log_L4\", F.lag(\"ex_log\",4).over(windowSpec))\n",
    "                 .withColumn(\"gdp_log_L1\", F.lag(\"gdp_log\").over(windowSpec))\n",
    "                 .withColumn(\"gdp_log_L2\", F.lag(\"gdp_log\",2).over(windowSpec))\n",
    "                 .withColumn(\"gdp_log_L3\", F.lag(\"gdp_log\",3).over(windowSpec))\n",
    "                 .withColumn(\"gdp_log_L4\", F.lag(\"gdp_log\",4).over(windowSpec))\n",
    "                )\n",
    "\n",
    "# Calculate the difference\n",
    "df_with_diff = (df_with_lag.withColumn(\"D_value_log_L1\", F.col(\"value_log_L1\") - F.col(\"value_log_L2\"))\n",
    "                           .withColumn(\"D_value_log_L2\", F.col(\"value_log_L2\") - F.col(\"value_log_L3\"))\n",
    "                           .withColumn(\"D_value_log_L3\", F.col(\"value_log_L3\") - F.col(\"value_log_L4\"))\n",
    "                           .withColumn(\"D_ex_log_L1\", F.col(\"ex_log_L1\") - F.col(\"ex_log_L2\"))\n",
    "                           .withColumn(\"D_ex_log_L2\", F.col(\"ex_log_L2\") - F.col(\"ex_log_L3\"))\n",
    "                           .withColumn(\"D_ex_log_L3\", F.col(\"ex_log_L3\") - F.col(\"ex_log_L4\"))\n",
    "                           .withColumn(\"D_gdp_log_L1\", F.col(\"gdp_log_L1\") - F.col(\"gdp_log_L2\"))\n",
    "                           .withColumn(\"D_gdp_log_L2\", F.col(\"gdp_log_L2\") - F.col(\"gdp_log_L3\"))\n",
    "                           .withColumn(\"D_gdp_log_L3\", F.col(\"gdp_log_L3\") - F.col(\"gdp_log_L4\"))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df_with_diff = (df_with_diff.withColumn(\"D_value_log\", F.col(\"value_log\") - F.col(\"value_log_L1\"))\n",
    "                            .withColumn(\"D_ex_log\", F.col(\"ex_log\") - F.col(\"ex_log_L1\"))\n",
    "                            .withColumn(\"D_gdp_log\", F.col(\"gdp_log\") - F.col(\"gdp_log_L1\"))\n",
    "               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------+------------+-------------+---+------------------+-------------------+-------+------------------+------------------+------------------+------------------+-------------------+-------------------+-------------------+-------------------+----------+----------+----------+----------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+------------+------------+------------+--------------------+--------------------+---------+\n",
      "|country|year_quarter|export_import|hs2|         value_log|             ex_log|gdp_log|      value_log_L1|      value_log_L2|      value_log_L3|      value_log_L4|          ex_log_L1|          ex_log_L2|          ex_log_L3|          ex_log_L4|gdp_log_L1|gdp_log_L2|gdp_log_L3|gdp_log_L4|      D_value_log_L1|      D_value_log_L2|      D_value_log_L3|         D_ex_log_L1|         D_ex_log_L2|         D_ex_log_L3|D_gdp_log_L1|D_gdp_log_L2|D_gdp_log_L3|         D_value_log|            D_ex_log|D_gdp_log|\n",
      "+-------+------------+-------------+---+------------------+-------------------+-------+------------------+------------------+------------------+------------------+-------------------+-------------------+-------------------+-------------------+----------+----------+----------+----------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+------------+------------+------------+--------------------+--------------------+---------+\n",
      "|    413|      1988_1|            1| 72| 13.78567898852514|-12.560180714380259|   null|              null|              null|              null|              null|               null|               null|               null|               null|      null|      null|      null|      null|                null|                null|                null|                null|                null|                null|        null|        null|        null|                null|                null|     null|\n",
      "|    413|      1988_2|            1| 72|14.483997662904516| -12.11304162790952|   null| 13.78567898852514|              null|              null|              null|-12.560180714380259|               null|               null|               null|      null|      null|      null|      null|                null|                null|                null|                null|                null|                null|        null|        null|        null|  0.6983186743793759|  0.4471390864707381|     null|\n",
      "|    413|      1988_3|            1| 72|15.057245801139722|-11.699285090713191|   null|14.483997662904516| 13.78567898852514|              null|              null| -12.11304162790952|-12.560180714380259|               null|               null|      null|      null|      null|      null|  0.6983186743793759|                null|                null|  0.4471390864707381|                null|                null|        null|        null|        null|  0.5732481382352059| 0.41375653719632943|     null|\n",
      "|    413|      1988_4|            1| 72|13.805675350484156| -11.50658197396841|   null|15.057245801139722|14.483997662904516| 13.78567898852514|              null|-11.699285090713191| -12.11304162790952|-12.560180714380259|               null|      null|      null|      null|      null|  0.5732481382352059|  0.6983186743793759|                null| 0.41375653719632943|  0.4471390864707381|                null|        null|        null|        null| -1.2515704506555654|  0.1927031167447808|     null|\n",
      "|    413|      1989_1|            1| 72|14.173932275405486|-11.382212164799844|   null|13.805675350484156|15.057245801139722|14.483997662904516| 13.78567898852514| -11.50658197396841|-11.699285090713191| -12.11304162790952|-12.560180714380259|      null|      null|      null|      null| -1.2515704506555654|  0.5732481382352059|  0.6983186743793759|  0.1927031167447808| 0.41375653719632943|  0.4471390864707381|        null|        null|        null| 0.36825692492132944| 0.12436980916856655|     null|\n",
      "|    413|      1989_2|            1| 72| 11.07535030294383| -9.460047958633352|   null|14.173932275405486|13.805675350484156|15.057245801139722|14.483997662904516|-11.382212164799844| -11.50658197396841|-11.699285090713191| -12.11304162790952|      null|      null|      null|      null| 0.36825692492132944| -1.2515704506555654|  0.5732481382352059| 0.12436980916856655|  0.1927031167447808| 0.41375653719632943|        null|        null|        null| -3.0985819724616555|  1.9221642061664923|     null|\n",
      "|    413|      1989_3|            1| 72|11.875371722792163| -7.755652254973471|   null| 11.07535030294383|14.173932275405486|13.805675350484156|15.057245801139722| -9.460047958633352|-11.382212164799844| -11.50658197396841|-11.699285090713191|      null|      null|      null|      null| -3.0985819724616555| 0.36825692492132944| -1.2515704506555654|  1.9221642061664923| 0.12436980916856655|  0.1927031167447808|        null|        null|        null|  0.8000214198483331|  1.7043957036598805|     null|\n",
      "|    413|      1989_4|            1| 72|12.804975078511257| -7.362421199724985|   null|11.875371722792163| 11.07535030294383|14.173932275405486|13.805675350484156| -7.755652254973471| -9.460047958633352|-11.382212164799844| -11.50658197396841|      null|      null|      null|      null|  0.8000214198483331| -3.0985819724616555| 0.36825692492132944|  1.7043957036598805|  1.9221642061664923| 0.12436980916856655|        null|        null|        null|  0.9296033557190935|  0.3932310552484859|     null|\n",
      "|    413|      1990_1|            1| 72|12.095917319787414| -6.273946127655021|   null|12.804975078511257|11.875371722792163| 11.07535030294383|14.173932275405486| -7.362421199724985| -7.755652254973471| -9.460047958633352|-11.382212164799844|      null|      null|      null|      null|  0.9296033557190935|  0.8000214198483331| -3.0985819724616555|  0.3932310552484859|  1.7043957036598805|  1.9221642061664923|        null|        null|        null| -0.7090577587238425|   1.088475072069964|     null|\n",
      "|    413|      1990_2|            1| 72|12.458875964862498| -5.734687571404053|   null|12.095917319787414|12.804975078511257|11.875371722792163| 11.07535030294383| -6.273946127655021| -7.362421199724985| -7.755652254973471| -9.460047958633352|      null|      null|      null|      null| -0.7090577587238425|  0.9296033557190935|  0.8000214198483331|   1.088475072069964|  0.3932310552484859|  1.7043957036598805|        null|        null|        null|  0.3629586450750839|  0.5392585562509682|     null|\n",
      "|    413|      1990_3|            1| 72|12.748254484349266| -5.517610509590221|   null|12.458875964862498|12.095917319787414|12.804975078511257|11.875371722792163| -5.734687571404053| -6.273946127655021| -7.362421199724985| -7.755652254973471|      null|      null|      null|      null|  0.3629586450750839| -0.7090577587238425|  0.9296033557190935|  0.5392585562509682|   1.088475072069964|  0.3932310552484859|        null|        null|        null|  0.2893785194867675| 0.21707706181383202|     null|\n",
      "|    413|      1990_4|            1| 72|12.879692910781896| -5.493871739508263|   null|12.748254484349266|12.458875964862498|12.095917319787414|12.804975078511257| -5.517610509590221| -5.734687571404053| -6.273946127655021| -7.362421199724985|      null|      null|      null|      null|  0.2893785194867675|  0.3629586450750839| -0.7090577587238425| 0.21707706181383202|  0.5392585562509682|   1.088475072069964|        null|        null|        null|  0.1314384264326307|0.023738770081958016|     null|\n",
      "|    413|      1991_1|            1| 72|11.722116125866204| -5.037448779488766|   null|12.879692910781896|12.748254484349266|12.458875964862498|12.095917319787414| -5.493871739508263| -5.517610509590221| -5.734687571404053| -6.273946127655021|      null|      null|      null|      null|  0.1314384264326307|  0.2893785194867675|  0.3629586450750839|0.023738770081958016| 0.21707706181383202|  0.5392585562509682|        null|        null|        null| -1.1575767849156922| 0.45642296001949667|     null|\n",
      "|    413|      1991_2|            1| 72|12.477270852498492| -4.941149141430503|   null|11.722116125866204|12.879692910781896|12.748254484349266|12.458875964862498| -5.037448779488766| -5.493871739508263| -5.517610509590221| -5.734687571404053|      null|      null|      null|      null| -1.1575767849156922|  0.1314384264326307|  0.2893785194867675| 0.45642296001949667|0.023738770081958016| 0.21707706181383202|        null|        null|        null|  0.7551547266322878| 0.09629963805826325|     null|\n",
      "|    413|      1991_3|            1| 72| 13.00943429224655|-4.9225321702249705|   null|12.477270852498492|11.722116125866204|12.879692910781896|12.748254484349266| -4.941149141430503| -5.037448779488766| -5.493871739508263| -5.517610509590221|      null|      null|      null|      null|  0.7551547266322878| -1.1575767849156922|  0.1314384264326307| 0.09629963805826325| 0.45642296001949667|0.023738770081958016|        null|        null|        null|  0.5321634397480572| 0.01861697120553263|     null|\n",
      "|    413|      1991_4|            1| 72|12.604912521844323|-4.8698790832509875|   null| 13.00943429224655|12.477270852498492|11.722116125866204|12.879692910781896|-4.9225321702249705| -4.941149141430503| -5.037448779488766| -5.493871739508263|      null|      null|      null|      null|  0.5321634397480572|  0.7551547266322878| -1.1575767849156922| 0.01861697120553263| 0.09629963805826325| 0.45642296001949667|        null|        null|        null|-0.40452177040222637|0.052653086973982965|     null|\n",
      "|    413|      1992_1|            1| 72|11.490822995062281| -4.856446883106402|   null|12.604912521844323| 13.00943429224655|12.477270852498492|11.722116125866204|-4.8698790832509875|-4.9225321702249705| -4.941149141430503| -5.037448779488766|      null|      null|      null|      null|-0.40452177040222637|  0.5321634397480572|  0.7551547266322878|0.052653086973982965| 0.01861697120553263| 0.09629963805826325|        null|        null|        null| -1.1140895267820419|0.013432200144585238|     null|\n",
      "|    413|      1992_2|            1| 72| 12.48274558930014|    -4.869804932483|   null|11.490822995062281|12.604912521844323| 13.00943429224655|12.477270852498492| -4.856446883106402|-4.8698790832509875|-4.9225321702249705| -4.941149141430503|      null|      null|      null|      null| -1.1140895267820419|-0.40452177040222637|  0.5321634397480572|0.013432200144585238|0.052653086973982965| 0.01861697120553263|        null|        null|        null|  0.9919225942378596|-0.01335804937659...|     null|\n",
      "|    413|      1992_3|            1| 72|13.285593853692633| -4.833428148770092|   null| 12.48274558930014|11.490822995062281|12.604912521844323| 13.00943429224655|    -4.869804932483| -4.856446883106402|-4.8698790832509875|-4.9225321702249705|      null|      null|      null|      null|  0.9919225942378596| -1.1140895267820419|-0.40452177040222637|-0.01335804937659...|0.013432200144585238|0.052653086973982965|        null|        null|        null|  0.8028482643924928|0.036376783712908534|     null|\n",
      "|    413|      1992_4|            1| 72|12.911200273080158| -4.819040771507747|   null|13.285593853692633| 12.48274558930014|11.490822995062281|12.604912521844323| -4.833428148770092|    -4.869804932483| -4.856446883106402|-4.8698790832509875|      null|      null|      null|      null|  0.8028482643924928|  0.9919225942378596| -1.1140895267820419|0.036376783712908534|-0.01335804937659...|0.013432200144585238|        null|        null|        null|-0.37439358061247496|0.014387377262345069|     null|\n",
      "+-------+------------+-------------+---+------------------+-------------------+-------+------------------+------------------+------------------+------------------+-------------------+-------------------+-------------------+-------------------+----------+----------+----------+----------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+------------+------------+------------+--------------------+--------------------+---------+\n",
      "only showing top 20 rows"
     ]
    }
   ],
   "source": [
    "df_with_diff.filter((df_with_diff.country==413)&(df_with_diff.hs2==72)&(df_with_diff.export_import==1)).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "df_with_diff.write.parquet(\"s3://trade-final-project-bucket/dataset/trades_ARDL.parquet\", mode=\"overwrite\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PySpark",
   "language": "",
   "name": "pysparkkernel"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "python",
    "version": 3
   },
   "mimetype": "text/x-python",
   "name": "pyspark",
   "pygments_lexer": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}