examples.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c5a4658f",
   "metadata": {},
   "source": [
    "# Pyspark Helper Functions\n",
    "### [For less verbose and foolproof operations]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5ebecf66",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    from pyspark import SparkConf\n",
    "except ImportError:\n",
    "    ! pip install pyspark==3.2.1\n",
    "\n",
    "from pyspark import SparkConf\n",
    "from pyspark.sql import SparkSession, types as st\n",
    "from IPython.display import HTML\n",
    "\n",
    "import spark.helpers as sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "61340588",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup Spark\n",
    "\n",
    "conf = SparkConf().setMaster(\"local[1]\").setAppName(\"examples\")\n",
    "spark = SparkSession.builder.config(conf=conf).getOrCreate()\n",
    "spark.sparkContext.setLogLevel('ERROR')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9d699692",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div style='float:left'><h4>Dataset 1:</h3><table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>x1</th>\n",
       "      <th>x2</th>\n",
       "      <th>x3</th>\n",
       "      <th>x4</th>\n",
       "      <th>x5</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>A</td>\n",
       "      <td>J</td>\n",
       "      <td>734</td>\n",
       "      <td>499</td>\n",
       "      <td>595.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>B</td>\n",
       "      <td>J</td>\n",
       "      <td>357</td>\n",
       "      <td>202</td>\n",
       "      <td>525.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>C</td>\n",
       "      <td>H</td>\n",
       "      <td>864</td>\n",
       "      <td>568</td>\n",
       "      <td>433.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>D</td>\n",
       "      <td>J</td>\n",
       "      <td>530</td>\n",
       "      <td>703</td>\n",
       "      <td>112.3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>E</td>\n",
       "      <td>H</td>\n",
       "      <td>61</td>\n",
       "      <td>521</td>\n",
       "      <td>906.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>F</td>\n",
       "      <td>H</td>\n",
       "      <td>482</td>\n",
       "      <td>496</td>\n",
       "      <td>13.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>G</td>\n",
       "      <td>A</td>\n",
       "      <td>350</td>\n",
       "      <td>279</td>\n",
       "      <td>941.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>H</td>\n",
       "      <td>C</td>\n",
       "      <td>171</td>\n",
       "      <td>267</td>\n",
       "      <td>423.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>I</td>\n",
       "      <td>C</td>\n",
       "      <td>755</td>\n",
       "      <td>133</td>\n",
       "      <td>600.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>J</td>\n",
       "      <td>A</td>\n",
       "      <td>228</td>\n",
       "      <td>765</td>\n",
       "      <td>7.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table></div><div style='float:left; margin-left:50px;'><h4>Dataset 2:</h3><table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>x1</th>\n",
       "      <th>x3</th>\n",
       "      <th>x4</th>\n",
       "      <th>x6</th>\n",
       "      <th>x7</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>W</td>\n",
       "      <td>K</td>\n",
       "      <td>391</td>\n",
       "      <td>140</td>\n",
       "      <td>872.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>X</td>\n",
       "      <td>G</td>\n",
       "      <td>88</td>\n",
       "      <td>483</td>\n",
       "      <td>707.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Y</td>\n",
       "      <td>M</td>\n",
       "      <td>144</td>\n",
       "      <td>476</td>\n",
       "      <td>714.3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Z</td>\n",
       "      <td>J</td>\n",
       "      <td>896</td>\n",
       "      <td>68</td>\n",
       "      <td>902.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>A</td>\n",
       "      <td>O</td>\n",
       "      <td>946</td>\n",
       "      <td>187</td>\n",
       "      <td>431.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>B</td>\n",
       "      <td>P</td>\n",
       "      <td>692</td>\n",
       "      <td>523</td>\n",
       "      <td>503.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>C</td>\n",
       "      <td>Q</td>\n",
       "      <td>550</td>\n",
       "      <td>988</td>\n",
       "      <td>181.05</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>D</td>\n",
       "      <td>R</td>\n",
       "      <td>50</td>\n",
       "      <td>419</td>\n",
       "      <td>42.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>E</td>\n",
       "      <td>S</td>\n",
       "      <td>824</td>\n",
       "      <td>805</td>\n",
       "      <td>558.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>F</td>\n",
       "      <td>T</td>\n",
       "      <td>69</td>\n",
       "      <td>722</td>\n",
       "      <td>721.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load example datasets\n",
    "\n",
    "dataframe_1 = spark.read.options(header=True).csv(\"./data/dataset_1.csv\")\n",
    "dataframe_2 = spark.read.options(header=True).csv(\"./data/dataset_2.csv\")\n",
    "html = (\n",
    "    \"<div style='float:left'><h4>Dataset 1:</h3>\" +\n",
    "    dataframe_1.toPandas().to_html() + \n",
    "    \"</div><div style='float:left; margin-left:50px;'><h4>Dataset 2:</h3>\" +\n",
    "    dataframe_2.toPandas().to_html() +\n",
    "    \"</div>\"\n",
    ")\n",
    "HTML(html)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96265d5c",
   "metadata": {},
   "source": [
    "## 1. Pandas-like group by"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7be19620",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A  =>  2\n",
      "C  =>  2\n",
      "H  =>  3\n",
      "J  =>  3\n"
     ]
    }
   ],
   "source": [
    "for group, data in sh.group_iterator(dataframe_1, \"x2\"):\n",
    "    print(group, \" => \", data.toPandas().shape[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18894f4f",
   "metadata": {},
   "source": [
    "### [Multiple columns group by]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "858bc27c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('A', 'J')  =>  1\n",
      "('B', 'J')  =>  1\n",
      "('C', 'H')  =>  1\n",
      "('D', 'J')  =>  1\n",
      "('E', 'H')  =>  1\n",
      "('F', 'H')  =>  1\n",
      "('G', 'A')  =>  1\n",
      "('H', 'C')  =>  1\n",
      "('I', 'C')  =>  1\n",
      "('J', 'A')  =>  1\n"
     ]
    }
   ],
   "source": [
    "for group, data in sh.group_iterator(dataframe_1, [\"x1\", \"x2\"]):\n",
    "    print(group, \" => \", data.toPandas().shape[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e009be9a",
   "metadata": {},
   "source": [
    "## 2. Bulk-change schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ea6cae6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "before = [(x[\"name\"], x[\"type\"]) for x in dataframe_1.schema.jsonValue()[\"fields\"]]\n",
    "\n",
    "schema = {\n",
    "    \"x2\": st.IntegerType(),\n",
    "    \"x5\": st.FloatType(),\n",
    "}\n",
    "new_dataframe = sh.change_schema(dataframe_1, schema)\n",
    "\n",
    "after = [(x[\"name\"], x[\"type\"]) for x in new_dataframe.schema.jsonValue()[\"fields\"]]\n",
    "check = [\n",
    "    ('x1', 'string'),\n",
    "    ('x2', 'integer'),\n",
    "    ('x3', 'string'),\n",
    "    ('x4', 'string'),\n",
    "    ('x5', 'float')\n",
    "]\n",
    "\n",
    "assert before != after\n",
    "assert after == check"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e43941e",
   "metadata": {},
   "source": [
    "## 3. Improved joins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0381789c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>x1</th>\n",
       "      <th>x2</th>\n",
       "      <th>x3</th>\n",
       "      <th>x4</th>\n",
       "      <th>x5</th>\n",
       "      <th>x6</th>\n",
       "      <th>x7</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>A</td>\n",
       "      <td>A</td>\n",
       "      <td>O</td>\n",
       "      <td>946</td>\n",
       "      <td>7.0</td>\n",
       "      <td>187</td>\n",
       "      <td>431.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>A</td>\n",
       "      <td>A</td>\n",
       "      <td>O</td>\n",
       "      <td>946</td>\n",
       "      <td>941.0</td>\n",
       "      <td>187</td>\n",
       "      <td>431.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>C</td>\n",
       "      <td>C</td>\n",
       "      <td>Q</td>\n",
       "      <td>550</td>\n",
       "      <td>600.0</td>\n",
       "      <td>988</td>\n",
       "      <td>181.05</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>C</td>\n",
       "      <td>C</td>\n",
       "      <td>Q</td>\n",
       "      <td>550</td>\n",
       "      <td>423.0</td>\n",
       "      <td>988</td>\n",
       "      <td>181.05</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  x1 x2 x3   x4     x5   x6      x7\n",
       "0  A  A  O  946    7.0  187   431.0\n",
       "1  A  A  O  946  941.0  187   431.0\n",
       "2  C  C  Q  550  600.0  988  181.05\n",
       "3  C  C  Q  550  423.0  988  181.05"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "joined = sh.join(dataframe_1.select(\"x2\", \"x5\"), dataframe_2, sh.JoinStatement(\"x2\", \"x1\"))\n",
    "joined.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac8321a6",
   "metadata": {},
   "source": [
    "### [When there are overlapping columns]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f3bdc599",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error raised as expected: \n",
      "\n",
      "Overlapping columns found in the dataframes: ['x1', 'x3', 'x4']\n",
      "Please provide the `overwrite_strategy` argument therefore, to select a selection strategy:\n",
      "\t* \"left\": Use all the intersecting columns from the left dataframe\n",
      "\t* \"right\": Use all the intersecting columns from the right dataframe\n",
      "\t* [[\"x_in_left\", \"y_in_left\"], [\"z_in_right\"]]: Provide column names for both\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>x1</th>\n",
       "      <th>x2</th>\n",
       "      <th>x3</th>\n",
       "      <th>x4</th>\n",
       "      <th>x5</th>\n",
       "      <th>x6</th>\n",
       "      <th>x7</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>A</td>\n",
       "      <td>J</td>\n",
       "      <td>734</td>\n",
       "      <td>499</td>\n",
       "      <td>595.0</td>\n",
       "      <td>187</td>\n",
       "      <td>431.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>B</td>\n",
       "      <td>J</td>\n",
       "      <td>357</td>\n",
       "      <td>202</td>\n",
       "      <td>525.0</td>\n",
       "      <td>523</td>\n",
       "      <td>503.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>C</td>\n",
       "      <td>H</td>\n",
       "      <td>864</td>\n",
       "      <td>568</td>\n",
       "      <td>433.5</td>\n",
       "      <td>988</td>\n",
       "      <td>181.05</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>D</td>\n",
       "      <td>J</td>\n",
       "      <td>530</td>\n",
       "      <td>703</td>\n",
       "      <td>112.3</td>\n",
       "      <td>419</td>\n",
       "      <td>42.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>E</td>\n",
       "      <td>H</td>\n",
       "      <td>61</td>\n",
       "      <td>521</td>\n",
       "      <td>906.0</td>\n",
       "      <td>805</td>\n",
       "      <td>558.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>F</td>\n",
       "      <td>H</td>\n",
       "      <td>482</td>\n",
       "      <td>496</td>\n",
       "      <td>13.0</td>\n",
       "      <td>722</td>\n",
       "      <td>721.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  x1 x2   x3   x4     x5   x6      x7\n",
       "0  A  J  734  499  595.0  187   431.0\n",
       "1  B  J  357  202  525.0  523   503.5\n",
       "2  C  H  864  568  433.5  988  181.05\n",
       "3  D  J  530  703  112.3  419    42.0\n",
       "4  E  H   61  521  906.0  805   558.2\n",
       "5  F  H  482  496   13.0  722   721.0"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "try:\n",
    "    joined = sh.join(dataframe_1, dataframe_2, sh.JoinStatement(\"x1\"))\n",
    "except ValueError as error:\n",
    "    print(f\"Error raised as expected: {error}\")\n",
    "    joined = sh.join(dataframe_1, dataframe_2, sh.JoinStatement(\"x1\"), overwrite_strategy=\"left\")\n",
    "joined.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37c9173c",
   "metadata": {},
   "source": [
    "### [Keeping the duplicate columns from the right dataframe]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "96dfff0a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>x1</th>\n",
       "      <th>x2</th>\n",
       "      <th>x3</th>\n",
       "      <th>x4</th>\n",
       "      <th>x5</th>\n",
       "      <th>x6</th>\n",
       "      <th>x7</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>A</td>\n",
       "      <td>J</td>\n",
       "      <td>O</td>\n",
       "      <td>946</td>\n",
       "      <td>595.0</td>\n",
       "      <td>187</td>\n",
       "      <td>431.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>B</td>\n",
       "      <td>J</td>\n",
       "      <td>P</td>\n",
       "      <td>692</td>\n",
       "      <td>525.0</td>\n",
       "      <td>523</td>\n",
       "      <td>503.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>C</td>\n",
       "      <td>H</td>\n",
       "      <td>Q</td>\n",
       "      <td>550</td>\n",
       "      <td>433.5</td>\n",
       "      <td>988</td>\n",
       "      <td>181.05</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>D</td>\n",
       "      <td>J</td>\n",
       "      <td>R</td>\n",
       "      <td>50</td>\n",
       "      <td>112.3</td>\n",
       "      <td>419</td>\n",
       "      <td>42.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>E</td>\n",
       "      <td>H</td>\n",
       "      <td>S</td>\n",
       "      <td>824</td>\n",
       "      <td>906.0</td>\n",
       "      <td>805</td>\n",
       "      <td>558.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>F</td>\n",
       "      <td>H</td>\n",
       "      <td>T</td>\n",
       "      <td>69</td>\n",
       "      <td>13.0</td>\n",
       "      <td>722</td>\n",
       "      <td>721.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  x1 x2 x3   x4     x5   x6      x7\n",
       "0  A  J  O  946  595.0  187   431.0\n",
       "1  B  J  P  692  525.0  523   503.5\n",
       "2  C  H  Q  550  433.5  988  181.05\n",
       "3  D  J  R   50  112.3  419    42.0\n",
       "4  E  H  S  824  906.0  805   558.2\n",
       "5  F  H  T   69   13.0  722   721.0"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "joined = sh.join(dataframe_1, dataframe_2, sh.JoinStatement(\"x1\"), overwrite_strategy=\"right\")\n",
    "joined.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15c5a9c2",
   "metadata": {},
   "source": [
    "### [Keeping the duplicate columns from both]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e7a32f49",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>x1</th>\n",
       "      <th>x2</th>\n",
       "      <th>x3</th>\n",
       "      <th>x4</th>\n",
       "      <th>x5</th>\n",
       "      <th>x6</th>\n",
       "      <th>x7</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>A</td>\n",
       "      <td>J</td>\n",
       "      <td>734</td>\n",
       "      <td>946</td>\n",
       "      <td>595.0</td>\n",
       "      <td>187</td>\n",
       "      <td>431.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>B</td>\n",
       "      <td>J</td>\n",
       "      <td>357</td>\n",
       "      <td>692</td>\n",
       "      <td>525.0</td>\n",
       "      <td>523</td>\n",
       "      <td>503.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>C</td>\n",
       "      <td>H</td>\n",
       "      <td>864</td>\n",
       "      <td>550</td>\n",
       "      <td>433.5</td>\n",
       "      <td>988</td>\n",
       "      <td>181.05</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>D</td>\n",
       "      <td>J</td>\n",
       "      <td>530</td>\n",
       "      <td>50</td>\n",
       "      <td>112.3</td>\n",
       "      <td>419</td>\n",
       "      <td>42.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>E</td>\n",
       "      <td>H</td>\n",
       "      <td>61</td>\n",
       "      <td>824</td>\n",
       "      <td>906.0</td>\n",
       "      <td>805</td>\n",
       "      <td>558.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>F</td>\n",
       "      <td>H</td>\n",
       "      <td>482</td>\n",
       "      <td>69</td>\n",
       "      <td>13.0</td>\n",
       "      <td>722</td>\n",
       "      <td>721.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  x1 x2   x3   x4     x5   x6      x7\n",
       "0  A  J  734  946  595.0  187   431.0\n",
       "1  B  J  357  692  525.0  523   503.5\n",
       "2  C  H  864  550  433.5  988  181.05\n",
       "3  D  J  530   50  112.3  419    42.0\n",
       "4  E  H   61  824  906.0  805   558.2\n",
       "5  F  H  482   69   13.0  722   721.0"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "joined = sh.join(\n",
    "    dataframe_1, dataframe_2, sh.JoinStatement(\"x1\"), \n",
    "    overwrite_strategy=[[\"x1\", \"x3\"], [\"x4\"]]\n",
    ")\n",
    "joined.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24d273a5",
   "metadata": {},
   "source": [
    "### [Complex join]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "708ce08f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>x1</th>\n",
       "      <th>x2</th>\n",
       "      <th>x3</th>\n",
       "      <th>x4</th>\n",
       "      <th>x5</th>\n",
       "      <th>x6</th>\n",
       "      <th>x7</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>A</td>\n",
       "      <td>J</td>\n",
       "      <td>734</td>\n",
       "      <td>499</td>\n",
       "      <td>595.0</td>\n",
       "      <td>187</td>\n",
       "      <td>431.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>B</td>\n",
       "      <td>J</td>\n",
       "      <td>357</td>\n",
       "      <td>202</td>\n",
       "      <td>525.0</td>\n",
       "      <td>523</td>\n",
       "      <td>503.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>C</td>\n",
       "      <td>H</td>\n",
       "      <td>864</td>\n",
       "      <td>568</td>\n",
       "      <td>433.5</td>\n",
       "      <td>988</td>\n",
       "      <td>181.05</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>D</td>\n",
       "      <td>J</td>\n",
       "      <td>530</td>\n",
       "      <td>703</td>\n",
       "      <td>112.3</td>\n",
       "      <td>419</td>\n",
       "      <td>42.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>E</td>\n",
       "      <td>H</td>\n",
       "      <td>61</td>\n",
       "      <td>521</td>\n",
       "      <td>906.0</td>\n",
       "      <td>805</td>\n",
       "      <td>558.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>F</td>\n",
       "      <td>H</td>\n",
       "      <td>482</td>\n",
       "      <td>496</td>\n",
       "      <td>13.0</td>\n",
       "      <td>722</td>\n",
       "      <td>721.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>G</td>\n",
       "      <td>A</td>\n",
       "      <td>350</td>\n",
       "      <td>279</td>\n",
       "      <td>941.0</td>\n",
       "      <td>483</td>\n",
       "      <td>707.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>J</td>\n",
       "      <td>A</td>\n",
       "      <td>228</td>\n",
       "      <td>765</td>\n",
       "      <td>7.0</td>\n",
       "      <td>68</td>\n",
       "      <td>902.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  x1 x2   x3   x4     x5   x6      x7\n",
       "0  A  J  734  499  595.0  187   431.0\n",
       "1  B  J  357  202  525.0  523   503.5\n",
       "2  C  H  864  568  433.5  988  181.05\n",
       "3  D  J  530  703  112.3  419    42.0\n",
       "4  E  H   61  521  906.0  805   558.2\n",
       "5  F  H  482  496   13.0  722   721.0\n",
       "6  G  A  350  279  941.0  483   707.1\n",
       "7  J  A  228  765    7.0   68   902.0"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x1_x1 = sh.JoinStatement(\"x1\")\n",
    "x1_x3 = sh.JoinStatement(\"x1\", \"x3\")\n",
    "statement = sh.JoinStatement(x1_x1, x1_x3, \"or\")\n",
    "joined = sh.join(dataframe_1, dataframe_2, statement, overwrite_strategy=\"left\")\n",
    "joined.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9167b546",
   "metadata": {},
   "source": [
    "### [Further nested joins are not supported]\n",
    "(Perform sequential joins instead)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "217a0f06",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error raised as expected: [Recursive JoinStatement not implemented]\n"
     ]
    }
   ],
   "source": [
    "x1_x1 = sh.JoinStatement(\"x1\")\n",
    "x1_x2 = sh.JoinStatement(\"x1\", \"x3\")\n",
    "statement = sh.JoinStatement(x1_x1, x1_x2, \"or\")\n",
    "statement_complex = sh.JoinStatement(statement, statement, \"and\")\n",
    "try:\n",
    "    joined = sh.join(dataframe_1, dataframe_2, statement_complex, overwrite_strategy=\"left\")\n",
    "except NotImplementedError as error:\n",
    "    print(f\"Error raised as expected: [{error}]\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}