record_linkage_match_key.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "record_linkage_match_key.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyOGJ+AmoQn5CVyv7/lBvNpR",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Brent-Morrison/Misc_scripts/blob/master/record_linkage_match_key.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rW9bfj8zsMuA"
      },
      "source": [
        "## Introduction\r\n",
        "\r\n",
        "This notebook will document a methodology for deriving a match key for records identified as duplicates.  The challenge is in situations where more than two records are matched.  I.e., where record linking suggests three records belong to the one entity.\r\n",
        "\r\n",
        "For example, consider that record 1 is matched to record 2, and record 2 is matched to record 3.\r\n",
        "\r\n",
        "This will be implemented in Spark 2.4.7.\r\n",
        "<br>\r\n",
        "<br>\r\n",
        "### Set up"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-V47Worzrw7y"
      },
      "source": [
        "\r\n",
        "\r\n",
        "```\r\n",
        "# This is formatted as code\r\n",
        "```\r\n",
        "\r\n",
        "#### PySpark installations"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "n-vzqUSrry5Y",
        "outputId": "226b4566-264b-4c31-f80f-4f367b068ddc"
      },
      "source": [
        "!apt-get install openjdk-8-jdk-headless -qq > /dev/null\r\n",
        "!wget -q https://www-eu.apache.org/dist/spark/spark-2.4.7/spark-2.4.7-bin-hadoop2.7.tgz\r\n",
        "!tar xf spark-2.4.7-bin-hadoop2.7.tgz\r\n",
        "!pip install -q findspark\r\n",
        "!pip install pyspark\r\n",
        "!pip install pyarrow"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: pyspark in /usr/local/lib/python3.6/dist-packages (3.0.1)\n",
            "Requirement already satisfied: py4j==0.10.9 in /usr/local/lib/python3.6/dist-packages (from pyspark) (0.10.9)\n",
            "Requirement already satisfied: pyarrow in /usr/local/lib/python3.6/dist-packages (0.14.1)\n",
            "Requirement already satisfied: six>=1.0.0 in /usr/local/lib/python3.6/dist-packages (from pyarrow) (1.15.0)\n",
            "Requirement already satisfied: numpy>=1.14 in /usr/local/lib/python3.6/dist-packages (from pyarrow) (1.19.4)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5jbA6sUZr1qv"
      },
      "source": [
        "Set the environment variables so that Colab can find Spark"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TUMD0EHDr2O-"
      },
      "source": [
        "import os\r\n",
        "os.environ[\"JAVA_HOME\"] = \"/usr/lib/jvm/java-8-openjdk-amd64\"\r\n",
        "os.environ[\"SPARK_HOME\"] = \"/content/spark-2.4.7-bin-hadoop2.7\""
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gTSYzb29r4dz"
      },
      "source": [
        "Add PySpark to sys.path\r\n",
        "\r\n",
        "PySpark isn't on sys.path by default, but that doesn't mean it can't be used as a regular library. You can address this by either symlinking pyspark into your site-packages, or adding pyspark to sys.path at runtime. [findspark](https://github.com/minrk/findspark) does the latter."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HxCizwptr996"
      },
      "source": [
        "import findspark\r\n",
        "findspark.init()"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g7itCha-sCU1"
      },
      "source": [
        "Create the Spark session"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Dj2ZJUZQsC_a"
      },
      "source": [
        "from pyspark.sql import SparkSession\r\n",
        "spark = SparkSession.builder.master(\"local[*]\").getOrCreate()"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iu9hSKf77om1"
      },
      "source": [
        "Prior to using PySpark we need the required classses from the PySpark sql module."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lMkJ3-3t7ryd"
      },
      "source": [
        "from pyspark.sql import functions as F\r\n",
        "from pyspark.sql import Window as W\r\n",
        "from pyspark.sql import Column as C\r\n",
        "from pyspark.sql import GroupedData as G\r\n",
        "from pyspark.sql import DataFrame\r\n",
        "from pyspark.sql.types import *"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9NvxwOs4tcL5"
      },
      "source": [
        "### Solution\r\n",
        "\r\n",
        "#### Create mock data in a PySpark data frame"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wImmT2TRtfP6",
        "outputId": "ccd7ca59-14d6-4b4b-bf32-bccc64a910d3"
      },
      "source": [
        "m1 = spark.createDataFrame(\r\n",
        "  [\r\n",
        "    (1,2)\r\n",
        "    ,(2,3)\r\n",
        "    ,(2,3) # There may be dupes here if different tests trigger the same combination\r\n",
        "    ,(4,5)\r\n",
        "    ,(6,7)\r\n",
        "    ,(7,8)\r\n",
        "    ,(8,9)\r\n",
        "  ],\r\n",
        "  ['ecid_1', 'ecid_2', ] \r\n",
        ")\r\n",
        "\r\n",
        "m1.show()\r\n",
        "m1.printSchema()"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "+------+------+\n",
            "|ecid_1|ecid_2|\n",
            "+------+------+\n",
            "|     1|     2|\n",
            "|     2|     3|\n",
            "|     2|     3|\n",
            "|     4|     5|\n",
            "|     6|     7|\n",
            "|     7|     8|\n",
            "|     8|     9|\n",
            "+------+------+\n",
            "\n",
            "root\n",
            " |-- ecid_1: long (nullable = true)\n",
            " |-- ecid_2: long (nullable = true)\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V8TEK_8Ptk4f"
      },
      "source": [
        "### Create the match key  \r\n",
        "\r\n",
        "This is done with an array type column."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "w4UjavTfpLBi",
        "outputId": "5cb7c46a-ef25-419e-e2cd-b4e648ac6c6d"
      },
      "source": [
        "m2 = m1.withColumn('match_key1', F.array(F.col('ecid_1'),F.col('ecid_2')))\r\n",
        "m2.show(truncate=False)"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "+------+------+----------+\n",
            "|ecid_1|ecid_2|match_key1|\n",
            "+------+------+----------+\n",
            "|1     |2     |[1, 2]    |\n",
            "|2     |3     |[2, 3]    |\n",
            "|2     |3     |[2, 3]    |\n",
            "|4     |5     |[4, 5]    |\n",
            "|6     |7     |[6, 7]    |\n",
            "|7     |8     |[7, 8]    |\n",
            "|8     |9     |[8, 9]    |\n",
            "+------+------+----------+\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gfgQH4Vgw7Ad"
      },
      "source": [
        "#### Union"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ozZZuOOuw8kJ",
        "outputId": "315ca2c6-c3e7-44f0-f891-510735f49712"
      },
      "source": [
        "m3 = m2.select(F.col('ecid_1'),F.col('match_key1')).unionAll(m2.select(F.col('ecid_2'),F.col('match_key1')))\r\n",
        "m3.show(truncate=False)"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "+------+----------+\n",
            "|ecid_1|match_key1|\n",
            "+------+----------+\n",
            "|1     |[1, 2]    |\n",
            "|2     |[2, 3]    |\n",
            "|2     |[2, 3]    |\n",
            "|4     |[4, 5]    |\n",
            "|6     |[6, 7]    |\n",
            "|7     |[7, 8]    |\n",
            "|8     |[8, 9]    |\n",
            "|2     |[1, 2]    |\n",
            "|3     |[2, 3]    |\n",
            "|3     |[2, 3]    |\n",
            "|5     |[4, 5]    |\n",
            "|7     |[6, 7]    |\n",
            "|8     |[7, 8]    |\n",
            "|9     |[8, 9]    |\n",
            "+------+----------+\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "An5cTw1Dp1f4"
      },
      "source": [
        "#### Explode"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6Cmfwrsl0xHi",
        "outputId": "ac90b184-b1fb-4d36-95b1-476b6b45ab95"
      },
      "source": [
        "m4 = m3.withColumn('match_key1', F.explode('match_key1'))\r\n",
        "m4.show(truncate=False)"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "+------+----------+\n",
            "|ecid_1|match_key1|\n",
            "+------+----------+\n",
            "|1     |1         |\n",
            "|1     |2         |\n",
            "|2     |2         |\n",
            "|2     |3         |\n",
            "|2     |2         |\n",
            "|2     |3         |\n",
            "|4     |4         |\n",
            "|4     |5         |\n",
            "|6     |6         |\n",
            "|6     |7         |\n",
            "|7     |7         |\n",
            "|7     |8         |\n",
            "|8     |8         |\n",
            "|8     |9         |\n",
            "|2     |1         |\n",
            "|2     |2         |\n",
            "|3     |2         |\n",
            "|3     |3         |\n",
            "|3     |2         |\n",
            "|3     |3         |\n",
            "+------+----------+\n",
            "only showing top 20 rows\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5P_rT8cSwYQd"
      },
      "source": [
        "#### Collect to array"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5u_B8SoEqD4D",
        "outputId": "1b797292-95a3-4d0f-8e30-f857c97eeb60"
      },
      "source": [
        "m5 = (m4\r\n",
        "  .groupBy('ecid_1').agg(F.collect_list(F.col('match_key1')).alias('match_key1'))\r\n",
        "  .withColumn('match_key1', F.array_distinct('match_key1'))\r\n",
        "  .sort(F.col('ecid_1'))\r\n",
        "  )\r\n",
        "m5.show(truncate=False)"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "+------+----------+\n",
            "|ecid_1|match_key1|\n",
            "+------+----------+\n",
            "|1     |[1, 2]    |\n",
            "|2     |[2, 3, 1] |\n",
            "|3     |[2, 3]    |\n",
            "|4     |[4, 5]    |\n",
            "|5     |[4, 5]    |\n",
            "|6     |[6, 7]    |\n",
            "|7     |[7, 8, 6] |\n",
            "|8     |[8, 9, 7] |\n",
            "|9     |[8, 9]    |\n",
            "+------+----------+\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "brx8XbCkwmXh"
      },
      "source": [
        "#### Cross join"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Jtlhzbh-wqc_",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5cdf9565-a3d6-4bbf-cac1-6abdacdae015"
      },
      "source": [
        "m6 = m5.crossJoin(m5.withColumnRenamed('ecid_1','ecid_2').withColumnRenamed('match_key1','match_key2'))\r\n",
        "m6.show(truncate=False)"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "+------+----------+------+----------+\n",
            "|ecid_1|match_key1|ecid_2|match_key2|\n",
            "+------+----------+------+----------+\n",
            "|1     |[1, 2]    |1     |[1, 2]    |\n",
            "|1     |[1, 2]    |2     |[2, 3, 1] |\n",
            "|1     |[1, 2]    |3     |[2, 3]    |\n",
            "|1     |[1, 2]    |4     |[4, 5]    |\n",
            "|1     |[1, 2]    |5     |[4, 5]    |\n",
            "|1     |[1, 2]    |6     |[6, 7]    |\n",
            "|1     |[1, 2]    |7     |[7, 8, 6] |\n",
            "|1     |[1, 2]    |8     |[8, 9, 7] |\n",
            "|1     |[1, 2]    |9     |[8, 9]    |\n",
            "|2     |[2, 3, 1] |1     |[1, 2]    |\n",
            "|2     |[2, 3, 1] |2     |[2, 3, 1] |\n",
            "|2     |[2, 3, 1] |3     |[2, 3]    |\n",
            "|2     |[2, 3, 1] |4     |[4, 5]    |\n",
            "|2     |[2, 3, 1] |5     |[4, 5]    |\n",
            "|2     |[2, 3, 1] |6     |[6, 7]    |\n",
            "|2     |[2, 3, 1] |7     |[7, 8, 6] |\n",
            "|2     |[2, 3, 1] |8     |[8, 9, 7] |\n",
            "|2     |[2, 3, 1] |9     |[8, 9]    |\n",
            "|3     |[2, 3]    |1     |[1, 2]    |\n",
            "|3     |[2, 3]    |2     |[2, 3, 1] |\n",
            "+------+----------+------+----------+\n",
            "only showing top 20 rows\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wkh49hdwqOfj"
      },
      "source": [
        "#### Filter for overlaps"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "An7PRd0rzBgf",
        "outputId": "b26c95c0-14fc-4103-d935-d804adfe0f71"
      },
      "source": [
        "m7 = (m6\r\n",
        "  .withColumn('ovelap', F.arrays_overlap('match_key1','match_key2'))\r\n",
        "  .filter(F.col('ovelap') == True)\r\n",
        "  .withColumn('match_key',F.array_union(F.col('match_key1'),F.col('match_key2')))\r\n",
        "  .withColumn('match_key', F.array_sort('match_key'))\r\n",
        "  .withColumn('match_key_size',F.size(F.col('match_key')))\r\n",
        ")\r\n",
        "m7.show(truncate=False)"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "+------+----------+------+----------+------+------------+--------------+\n",
            "|ecid_1|match_key1|ecid_2|match_key2|ovelap|match_key   |match_key_size|\n",
            "+------+----------+------+----------+------+------------+--------------+\n",
            "|1     |[1, 2]    |1     |[1, 2]    |true  |[1, 2]      |2             |\n",
            "|1     |[1, 2]    |2     |[2, 3, 1] |true  |[1, 2, 3]   |3             |\n",
            "|1     |[1, 2]    |3     |[2, 3]    |true  |[1, 2, 3]   |3             |\n",
            "|2     |[2, 3, 1] |1     |[1, 2]    |true  |[1, 2, 3]   |3             |\n",
            "|2     |[2, 3, 1] |2     |[2, 3, 1] |true  |[1, 2, 3]   |3             |\n",
            "|2     |[2, 3, 1] |3     |[2, 3]    |true  |[1, 2, 3]   |3             |\n",
            "|3     |[2, 3]    |1     |[1, 2]    |true  |[1, 2, 3]   |3             |\n",
            "|3     |[2, 3]    |2     |[2, 3, 1] |true  |[1, 2, 3]   |3             |\n",
            "|3     |[2, 3]    |3     |[2, 3]    |true  |[2, 3]      |2             |\n",
            "|4     |[4, 5]    |4     |[4, 5]    |true  |[4, 5]      |2             |\n",
            "|4     |[4, 5]    |5     |[4, 5]    |true  |[4, 5]      |2             |\n",
            "|5     |[4, 5]    |4     |[4, 5]    |true  |[4, 5]      |2             |\n",
            "|5     |[4, 5]    |5     |[4, 5]    |true  |[4, 5]      |2             |\n",
            "|6     |[6, 7]    |6     |[6, 7]    |true  |[6, 7]      |2             |\n",
            "|6     |[6, 7]    |7     |[7, 8, 6] |true  |[6, 7, 8]   |3             |\n",
            "|6     |[6, 7]    |8     |[8, 9, 7] |true  |[6, 7, 8, 9]|4             |\n",
            "|7     |[7, 8, 6] |6     |[6, 7]    |true  |[6, 7, 8]   |3             |\n",
            "|7     |[7, 8, 6] |7     |[7, 8, 6] |true  |[6, 7, 8]   |3             |\n",
            "|7     |[7, 8, 6] |8     |[8, 9, 7] |true  |[6, 7, 8, 9]|4             |\n",
            "|7     |[7, 8, 6] |9     |[8, 9]    |true  |[6, 7, 8, 9]|4             |\n",
            "+------+----------+------+----------+------+------------+--------------+\n",
            "only showing top 20 rows\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lAIA6e1lAh7z"
      },
      "source": [
        "#### Filter to derive final table\r\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9i7nuBAL068U",
        "outputId": "6ddf710e-45cc-4cb0-f20d-920d074d7c61"
      },
      "source": [
        "w1 = W.partitionBy('ecid_1','match_key_size').orderBy('ecid_2')\r\n",
        "w2 = W.partitionBy('ecid_1')\r\n",
        "\r\n",
        "m8 = (m7\r\n",
        "  .withColumn('rank', F.rank().over(w1))\r\n",
        "  .withColumn('max_size', F.max(F.col('match_key_size')).over(w2))\r\n",
        "  .filter(\r\n",
        "    (F.col('match_key_size') == F.col('max_size')) &\r\n",
        "    (F.col('rank') == F.lit(1))\r\n",
        "    )\r\n",
        "  .select(F.col('ecid_1').alias('ecid'),F.col('match_key'))\r\n",
        "  .sort(F.col('ecid_1'))#,F.col('rank'))\r\n",
        ") \r\n",
        "m8.show(truncate=False)"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "+----+------------+\n",
            "|ecid|match_key   |\n",
            "+----+------------+\n",
            "|1   |[1, 2, 3]   |\n",
            "|2   |[1, 2, 3]   |\n",
            "|3   |[1, 2, 3]   |\n",
            "|4   |[4, 5]      |\n",
            "|5   |[4, 5]      |\n",
            "|6   |[6, 7, 8, 9]|\n",
            "|7   |[6, 7, 8, 9]|\n",
            "|8   |[6, 7, 8, 9]|\n",
            "|9   |[6, 7, 8, 9]|\n",
            "+----+------------+\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w82tkt4F0lDY"
      },
      "source": [
        "### In one step \r\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "05xs79Ylktpj",
        "outputId": "882396f2-9ea0-4dde-eca7-360010234b9a"
      },
      "source": [
        "m9 = spark.createDataFrame(\r\n",
        "  [\r\n",
        "    (1,2)\r\n",
        "    ,(2,3)\r\n",
        "    ,(2,3) # There may be dupes here if different tests trigger the same combination\r\n",
        "    ,(4,5)\r\n",
        "    ,(6,7)\r\n",
        "    ,(7,8)\r\n",
        "    ,(8,9)\r\n",
        "  ],\r\n",
        "  ['ecid_1','ecid_2',] \r\n",
        ")\r\n",
        "\r\n",
        "m9.show(truncate=False)"
      ],
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "+------+------+\n",
            "|ecid_1|ecid_2|\n",
            "+------+------+\n",
            "|1     |2     |\n",
            "|2     |3     |\n",
            "|2     |3     |\n",
            "|4     |5     |\n",
            "|6     |7     |\n",
            "|7     |8     |\n",
            "|8     |9     |\n",
            "+------+------+\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NGTwKbiAvT4I",
        "outputId": "ad2ba1f0-9eb3-46e5-cf69-d2d456f9f798"
      },
      "source": [
        "m10 = (\r\n",
        "  m9.withColumn('match_key1', F.array(F.col('ecid_1'),F.col('ecid_2'))).select(F.col('ecid_1'),F.col('match_key1'))\r\n",
        "  .unionAll(\r\n",
        "    m9.withColumn('match_key1', F.array(F.col('ecid_1'),F.col('ecid_2'))).select(F.col('ecid_2'),F.col('match_key1'))\r\n",
        "  )\r\n",
        "  .withColumn('match_key1', F.explode('match_key1'))\r\n",
        "  .groupBy('ecid_1').agg(F.collect_list(F.col('match_key1')).alias('match_key1'))\r\n",
        "  .withColumn('match_key1', F.array_distinct('match_key1'))\r\n",
        "  \r\n",
        "  .crossJoin(\r\n",
        "    m9.withColumn('match_key1', F.array(F.col('ecid_1'),F.col('ecid_2'))).select(F.col('ecid_1'),F.col('match_key1'))\r\n",
        "    .unionAll(\r\n",
        "      m9.withColumn('match_key1', F.array(F.col('ecid_1'),F.col('ecid_2'))).select(F.col('ecid_2'),F.col('match_key1'))\r\n",
        "      )\r\n",
        "    .withColumn('match_key1', F.explode('match_key1'))\r\n",
        "    .groupBy('ecid_1').agg(F.collect_list(F.col('match_key1')).alias('match_key1'))\r\n",
        "    .withColumn('match_key1', F.array_distinct('match_key1')).withColumnRenamed('ecid_1','ecid_2').withColumnRenamed('match_key1','match_key2')\r\n",
        "    )\r\n",
        "  \r\n",
        "  .withColumn('overlap', F.arrays_overlap('match_key1','match_key2'))\r\n",
        "  .filter(F.col('overlap') == True)\r\n",
        "  .withColumn('match_key',F.array_union(F.col('match_key1'),F.col('match_key2')))\r\n",
        "  .withColumn('match_key', F.array_sort('match_key'))\r\n",
        "  .withColumn('match_key_size',F.size(F.col('match_key')))\r\n",
        "  .withColumn('rank', F.rank().over(W.partitionBy('ecid_1','match_key_size').orderBy('ecid_2')))\r\n",
        "  .withColumn('max_size', F.max(F.col('match_key_size')).over(W.partitionBy('ecid_1')))\r\n",
        "  .filter(\r\n",
        "    (F.col('match_key_size') == F.col('max_size')) &\r\n",
        "    (F.col('rank') == F.lit(1))\r\n",
        "    )\r\n",
        "  .select(F.col('ecid_1').alias('ecid'),F.col('match_key'))\r\n",
        "  .sort(F.col('ecid_1'))\r\n",
        ")\r\n",
        "\r\n",
        "m10.show(truncate=False)"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "+----+------------+\n",
            "|ecid|match_key   |\n",
            "+----+------------+\n",
            "|1   |[1, 2, 3]   |\n",
            "|2   |[1, 2, 3]   |\n",
            "|3   |[1, 2, 3]   |\n",
            "|4   |[4, 5]      |\n",
            "|5   |[4, 5]      |\n",
            "|6   |[6, 7, 8, 9]|\n",
            "|7   |[6, 7, 8, 9]|\n",
            "|8   |[6, 7, 8, 9]|\n",
            "|9   |[6, 7, 8, 9]|\n",
            "+----+------------+\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Lak4saus95w"
      },
      "source": [
        "Note that the array type column can be converted to a string with the ```concat_ws``` function.\r\n",
        "\r\n",
        "This may be required if writing to a csv file or querying via an engine that does not support the array type.  \r\n",
        "  \r\n",
        "Impala for example (Hive is OK)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "--EqWUfCtHf3",
        "outputId": "ad39857b-7de9-40bc-e83d-94d90d53dcc9"
      },
      "source": [
        "m11 = m10.withColumn('match_key_string', F.concat_ws('_',F.col('match_key')))\r\n",
        "\r\n",
        "m11.show()\r\n",
        "m11.printSchema()"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "+----+------------+----------------+\n",
            "|ecid|   match_key|match_key_string|\n",
            "+----+------------+----------------+\n",
            "|   1|   [1, 2, 3]|           1_2_3|\n",
            "|   2|   [1, 2, 3]|           1_2_3|\n",
            "|   3|   [1, 2, 3]|           1_2_3|\n",
            "|   4|      [4, 5]|             4_5|\n",
            "|   5|      [4, 5]|             4_5|\n",
            "|   6|[6, 7, 8, 9]|         6_7_8_9|\n",
            "|   7|[6, 7, 8, 9]|         6_7_8_9|\n",
            "|   8|[6, 7, 8, 9]|         6_7_8_9|\n",
            "|   9|[6, 7, 8, 9]|         6_7_8_9|\n",
            "+----+------------+----------------+\n",
            "\n",
            "root\n",
            " |-- ecid: long (nullable = true)\n",
            " |-- match_key: array (nullable = true)\n",
            " |    |-- element: long (containsNull = true)\n",
            " |-- match_key_string: string (nullable = false)\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}