assigment.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Assigment problem\n",
    "\n",
    "In this notebook we have a practical example where we count pairs of elements that are identified by id_1 and id_2. The problem consists in assigning the elements 1 to 1 minimizing the distance between the pairs.\n",
    "\n",
    "**Not all combinations have values, it is considered that if they have missing values, the records are very different.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyspark.sql.functions as F\n",
    "from pyspark.sql.types import StructType\n",
    "from scipy.optimize import linear_sum_assignment\n",
    "from pyspark.sql import Window\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+---------------+---------+---------------+-----+\n",
      "|id1_x|id2_x|id1_y|id2_y|distance|\n",
      "+---------+---------------+---------+---------------+-----+\n",
      "|        _|              a|        a|              _|  9.1|\n",
      "|        _|              a|        b|              _|  5.1|\n",
      "|        _|              b|        a|              _|  7.1|\n",
      "|        _|              b|        c|              _|  1.1|\n",
      "|        _|              c|        d|              _|  8.1|\n",
      "|        _|              c|        b|              _|  7.1|\n",
      "|        _|              d|        d|              _|  1.1|\n",
      "|        _|              d|        a|              _|  3.1|\n",
      "+---------+---------------+---------+---------------+-----+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df = spark.createDataFrame(\n",
    "    [\n",
    "        ('_', 'a', 'a', '_', 9.1),\n",
    "        ('_', 'a', 'b', '_', 5.1),\n",
    "        ('_', 'b', 'a', '_', 7.1),\n",
    "        ('_', 'b', 'c', '_', 1.1),\n",
    "        ('_', 'c', 'd', '_', 8.1),\n",
    "        ('_', 'c', 'b', '_', 7.1),\n",
    "        ('_', 'd', 'd', '_', 1.1),\n",
    "        ('_', 'd', 'a', '_', 3.1),\n",
    "    ],\n",
    "    [\"id1_x\", \"id2_x\", \"id1_y\", \"id2_y\", \"distance\"]\n",
    ")\n",
    "\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+---------------+---------+---------------+-----+-----+\n",
      "|id1_x|id2_x|id1_y|id2_y|distance|row_id|\n",
      "+---------+---------------+---------+---------------+-----+-----+\n",
      "|        _|              a|        a|              _|  9.1|    1|\n",
      "|        _|              a|        b|              _|  5.1|    1|\n",
      "|        _|              b|        a|              _|  7.1|    2|\n",
      "|        _|              b|        c|              _|  1.1|    2|\n",
      "|        _|              c|        d|              _|  8.1|    3|\n",
      "|        _|              c|        b|              _|  7.1|    3|\n",
      "|        _|              d|        d|              _|  1.1|    4|\n",
      "|        _|              d|        a|              _|  3.1|    4|\n",
      "+---------+---------------+---------+---------------+-----+-----+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# create row id to identify each combination\n",
    "w = Window.orderBy(['id1_x', 'id2_x'])\n",
    "\n",
    "df = df.withColumn('row_id', F.dense_rank().over(w))\n",
    "\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+---------------+-----+\n",
      "|id1_x|id2_x|row_id|\n",
      "+---------+---------------+-----+\n",
      "|        _|              a|    1|\n",
      "|        _|              b|    2|\n",
      "|        _|              c|    3|\n",
      "|        _|              d|    4|\n",
      "+---------+---------------+-----+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df3 = df.select('id1_x','id2_x','row_id').dropDuplicates()\n",
    "\n",
    "df3.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+---------------+----+----+----+----+\n",
      "|id1_y|id2_y|   1|   2|   3|   4|\n",
      "+---------+---------------+----+----+----+----+\n",
      "|        a|              _| 9.1| 7.1|null| 3.1|\n",
      "|        b|              _| 5.1|null| 7.1|null|\n",
      "|        c|              _|null| 1.1|null|null|\n",
      "|        d|              _|null|null| 8.1| 1.1|\n",
      "+---------+---------------+----+----+----+----+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# create cost matrix with all combinations\n",
    "df1 = df.groupby('id1_y', 'id2_y').pivot('row_id').agg(F.first('distance'))\n",
    "\n",
    "df1.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define assigment algorithm as pandas udf fucntion\n",
    "schema = StructType.fromJson(df1.schema.jsonValue()).add('assigned', 'string')\n",
    "\n",
    "def hungarian_algorithm(df_matrix, fill_values=1e3):\n",
    "    matrix = df_matrix.iloc[:,2:]\n",
    "    cols = matrix.columns\n",
    "    n_row, n_col = matrix.shape\n",
    "    cost_matrix = matrix.fillna(fill_values).values\n",
    "\n",
    "    if n_row > n_col:\n",
    "        cost_matrix = np.hstack((cost_matrix, np.full((n_row, n_row - n_col), fill_values)))\n",
    "    elif n_col > n_row:\n",
    "        cost_matrix = np.vstack((cost_matrix, np.full((n_col - n_row , n_col), fill_values)))\n",
    "\n",
    "    _, idx = linear_sum_assignment(cost_matrix)\n",
    "\n",
    "    return df_matrix.assign(assigned=['%s' % (i + 1) for i in idx][:n_row])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "find_assigned = F.pandas_udf(lambda x: hungarian_algorithm(x), schema, F.PandasUDFType.GROUPED_MAP)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+---------------+----+----+----+----+--------+\n",
      "|id1_y|id2_y|   1|   2|   3|   4|assigned|\n",
      "+---------+---------------+----+----+----+----+--------+\n",
      "|        a|              _| 9.1| 7.1|null| 3.1|       4|\n",
      "|        b|              _| 5.1|null| 7.1|null|       1|\n",
      "|        c|              _|null| 1.1|null|null|       2|\n",
      "|        d|              _|null|null| 8.1| 1.1|       3|\n",
      "+---------+---------------+----+----+----+----+--------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df2 = df1.groupby('id2_y').apply(find_assigned)\n",
    "\n",
    "df2.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1, 2, 3, 4]"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get column range\n",
    "columns_ = df2.columns\n",
    "columns_range = columns_[2:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+---------------+--------+-----+-----+--------+\n",
      "|id1_y|id2_y|assigned|row_id|distance|is_match|\n",
      "+---------+---------------+--------+-----+-----+--------+\n",
      "|        a|              _|       4|    1|  9.1|   false|\n",
      "|        a|              _|       4|    2|  7.1|   false|\n",
      "|        a|              _|       4|    3| null|   false|\n",
      "|        a|              _|       4|    4|  3.1|    true|\n",
      "|        b|              _|       1|    1|  5.1|    true|\n",
      "|        b|              _|       1|    2| null|   false|\n",
      "|        b|              _|       1|    3|  7.1|   false|\n",
      "|        b|              _|       1|    4| null|   false|\n",
      "|        c|              _|       2|    1| null|   false|\n",
      "|        c|              _|       2|    2|  1.1|    true|\n",
      "|        c|              _|       2|    3| null|   false|\n",
      "|        c|              _|       2|    4| null|   false|\n",
      "|        d|              _|       3|    1| null|   false|\n",
      "|        d|              _|       3|    2| null|   false|\n",
      "|        d|              _|       3|    3|  8.1|    true|\n",
      "|        d|              _|       3|    4|  1.1|   false|\n",
      "+---------+---------------+--------+-----+-----+--------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_new = df2.selectExpr(\n",
    "    \"id1_y\",\n",
    "    \"id2_y\",\n",
    "    'assigned',\n",
    "    'stack({},{}) as (row_id, distance)'.format(len(columns_range),\n",
    "                                                ','.join(\"int('{0}'), `{0}`\".format(c) for c in columns_range))\n",
    ") \\\n",
    ".withColumn('is_match', F.expr(\"assigned=row_id\"))\n",
    "\n",
    "df_new.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use cases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+---------------+---------+---------------+-----+\n",
      "|id1_x|id2_x|id1_y|id2_y|distance|\n",
      "+---------+---------------+---------+---------------+-----+\n",
      "|        _|              a|        a|              _|  9.1|\n",
      "|        _|              a|        b|              _|  5.1|\n",
      "|        _|              b|        a|              _|  7.1|\n",
      "|        _|              b|        c|              _|  1.1|\n",
      "|        _|              c|        d|              _|  8.1|\n",
      "|        _|              c|        b|              _|  7.1|\n",
      "|        _|              d|        d|              _|  1.1|\n",
      "|        _|              d|        a|              _|  3.1|\n",
      "+---------+---------------+---------+---------------+-----+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df = spark.createDataFrame(\n",
    "    [\n",
    "        ('_', 'a', 'a', '_', 9.1),\n",
    "        ('_', 'a', 'b', '_', 5.1),\n",
    "        ('_', 'b', 'a', '_', 7.1),\n",
    "        ('_', 'b', 'c', '_', 1.1),\n",
    "        ('_', 'c', 'd', '_', 8.1),\n",
    "        ('_', 'c', 'b', '_', 7.1),\n",
    "        ('_', 'd', 'd', '_', 1.1),\n",
    "        ('_', 'd', 'a', '_', 3.1),\n",
    "    ],\n",
    "    [\"id1_x\", \"id2_x\", \"id1_y\", \"id2_y\", \"distance\"]\n",
    ")\n",
    "\n",
    "df_.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+---------------+---------+---------------+-----+\n",
      "|id1_x|id2_x|id1_y|id2_y|distance|\n",
      "+---------+---------------+---------+---------------+-----+\n",
      "|        _|              a|        a|              _|  9.1|\n",
      "|        _|              a|        b|              _|  5.1|\n",
      "|        _|              b|        a|              _|  7.1|\n",
      "|        _|              b|        c|              _|  1.1|\n",
      "|        _|              c|        d|              _|  8.1|\n",
      "|        _|              c|        b|              _|  7.1|\n",
      "+---------+---------------+---------+---------------+-----+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# more rows\n",
    "df_1 = spark.createDataFrame(\n",
    "    [\n",
    "        ('_', 'a', 'a', '_', 9.1),\n",
    "        ('_', 'a', 'b', '_', 5.1),\n",
    "        ('_', 'b', 'a', '_', 7.1),\n",
    "        ('_', 'b', 'c', '_', 1.1),\n",
    "        ('_', 'c', 'd', '_', 8.1),\n",
    "        ('_', 'c', 'b', '_', 7.1),\n",
    "    ],\n",
    "    [\"id1_x\", \"id2_x\", \"id1_y\", \"id2_y\", \"distance\"]\n",
    ")\n",
    "\n",
    "df_1.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------+---------------+---------+---------------+-----+\n",
      "|id1_x|id2_x|id1_y|id2_y|distance|\n",
      "+---------+---------------+---------+---------------+-----+\n",
      "|        _|              a|        a|              _|  9.1|\n",
      "|        _|              a|        b|              _|  5.1|\n",
      "|        _|              a|        c|              _|  7.1|\n",
      "|        _|              b|        a|              _|  1.1|\n",
      "|        _|              b|        b|              _|  8.1|\n",
      "|        _|              b|        c|              _|  7.1|\n",
      "+---------+---------------+---------+---------------+-----+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# more cols\n",
    "df_2 = spark.createDataFrame(\n",
    "    [\n",
    "        ('_', 'a', 'a', '_', 9.1),\n",
    "        ('_', 'a', 'b', '_', 5.1),\n",
    "        ('_', 'a', 'c', '_', 7.1),\n",
    "        ('_', 'b', 'a', '_', 1.1),\n",
    "        ('_', 'b', 'b', '_', 8.1),\n",
    "        ('_', 'b', 'c', '_', 7.1),\n",
    "    ],\n",
    "    [\"id1_x\", \"id2_x\", \"id1_y\", \"id2_y\", \"distance\"]\n",
    ")\n",
    "\n",
    "df_2.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hungarian_algorithm(df_matrix, fill_values=1e3):\n",
    "    matrix = df_matrix.iloc[:,2:]\n",
    "    cols = matrix.columns\n",
    "    n_row, n_col = matrix.shape\n",
    "    cost_matrix = matrix.fillna(fill_values).values\n",
    "\n",
    "    if n_row > n_col:\n",
    "        cost_matrix = np.hstack((cost_matrix, np.full((n_row, n_row - n_col), fill_values)))\n",
    "    elif n_col > n_row:\n",
    "        cost_matrix = np.vstack((cost_matrix, np.full((n_col - n_row , n_col), fill_values)))\n",
    "\n",
    "    _, idx = linear_sum_assignment(cost_matrix)\n",
    "\n",
    "    return df_matrix.assign(assigned=['%s' % (i + 1) for i in idx][:n_row])\n",
    "\n",
    "def solve_assigment(df):\n",
    "    n_rows = df_.drop_duplicates(['id1_x', 'id2_x']).count()\n",
    "    w = Window.orderBy(['id1_x', 'id2_x'])\n",
    "\n",
    "    # get row_id\n",
    "    df = df.withColumn('row_id', F.dense_rank().over(w))\n",
    "\n",
    "    # drop duplicates\n",
    "    df3 = df.select('id1_x','id2_x','row_id').dropDuplicates()\n",
    "\n",
    "    # create pivot table\n",
    "    df1 = df.groupby('id1_y', 'id2_y').pivot('row_id').agg(F.first('distance'))\n",
    "\n",
    "    # apply hungarian algorithm\n",
    "    schema = StructType.fromJson(df1.schema.jsonValue()).add('assigned', 'string')\n",
    "    find_assigned = F.pandas_udf(lambda x: __find_assigned(x), schema, F.PandasUDFType.GROUPED_MAP)\n",
    "    df2 = df1.groupby('id2_y').apply(hungarian_algorithm)\n",
    "\n",
    "    # get column range\n",
    "    columns_ = df2.columns\n",
    "    columns_range = columns_[2:-1]\n",
    "\n",
    "    df_new = df2.selectExpr(\n",
    "        \"id1_y\",\n",
    "        \"id2_y\",\n",
    "        'assigned',\n",
    "        'stack({},{}) as (row_id, distance)'.format(len(columns_range),\n",
    "                                                    ','.join(\"int('{0}'), `{0}`\".format(c) for c in columns_range))\n",
    "    ) \\\n",
    "    .withColumn('is_match', F.expr(\"assigned=row_id\"))\n",
    "\n",
    "\n",
    "    return df_new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n",
      "+---------+---------------+---------+---------------+-----+-----+\n",
      "|id1_x|id2_x|id1_y|id2_y|distance|row_id|\n",
      "+---------+---------------+---------+---------------+-----+-----+\n",
      "|        _|              a|        a|              _|  9.1|    1|\n",
      "|        _|              a|        b|              _|  5.1|    1|\n",
      "|        _|              b|        a|              _|  7.1|    2|\n",
      "|        _|              b|        c|              _|  1.1|    2|\n",
      "|        _|              c|        d|              _|  8.1|    3|\n",
      "|        _|              c|        b|              _|  7.1|    3|\n",
      "|        _|              d|        d|              _|  1.1|    4|\n",
      "|        _|              d|        a|              _|  3.1|    4|\n",
      "+---------+---------------+---------+---------------+-----+-----+\n",
      "\n",
      "+---------+---------------+-----+\n",
      "|id1_x|id2_x|row_id|\n",
      "+---------+---------------+-----+\n",
      "|        _|              a|    1|\n",
      "|        _|              b|    2|\n",
      "|        _|              c|    3|\n",
      "|        _|              d|    4|\n",
      "+---------+---------------+-----+\n",
      "\n",
      "+---------+---------------+----+----+----+----+\n",
      "|id1_y|id2_y|   1|   2|   3|   4|\n",
      "+---------+---------------+----+----+----+----+\n",
      "|        a|              _| 9.1| 7.1|null| 3.1|\n",
      "|        b|              _| 5.1|null| 7.1|null|\n",
      "|        c|              _|null| 1.1|null|null|\n",
      "|        d|              _|null|null| 8.1| 1.1|\n",
      "+---------+---------------+----+----+----+----+\n",
      "\n",
      "+---------+---------------+----+----+----+----+--------+\n",
      "|id1_y|id2_y|   1|   2|   3|   4|assigned|\n",
      "+---------+---------------+----+----+----+----+--------+\n",
      "|        a|              _| 9.1| 7.1|null| 3.1|       4|\n",
      "|        b|              _| 5.1|null| 7.1|null|       1|\n",
      "|        c|              _|null| 1.1|null|null|       2|\n",
      "|        d|              _|null|null| 8.1| 1.1|       3|\n",
      "+---------+---------------+----+----+----+----+--------+\n",
      "\n",
      "['1', '2', '3', '4']\n",
      "+---------+---------------+--------+-----+-----+--------+\n",
      "|id1_y|id2_y|assigned|row_id|distance|is_match|\n",
      "+---------+---------------+--------+-----+-----+--------+\n",
      "|        a|              _|       4|    1|  9.1|   false|\n",
      "|        a|              _|       4|    2|  7.1|   false|\n",
      "|        a|              _|       4|    3| null|   false|\n",
      "|        a|              _|       4|    4|  3.1|    true|\n",
      "|        b|              _|       1|    1|  5.1|    true|\n",
      "|        b|              _|       1|    2| null|   false|\n",
      "|        b|              _|       1|    3|  7.1|   false|\n",
      "|        b|              _|       1|    4| null|   false|\n",
      "|        c|              _|       2|    1| null|   false|\n",
      "|        c|              _|       2|    2|  1.1|    true|\n",
      "|        c|              _|       2|    3| null|   false|\n",
      "|        c|              _|       2|    4| null|   false|\n",
      "|        d|              _|       3|    1| null|   false|\n",
      "|        d|              _|       3|    2| null|   false|\n",
      "|        d|              _|       3|    3|  8.1|    true|\n",
      "|        d|              _|       3|    4|  1.1|   false|\n",
      "+---------+---------------+--------+-----+-----+--------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_1_assigned = emparejar(df)\n",
    "df_1_assigned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n",
      "+---------+---------------+---------+---------------+-----+-----+\n",
      "|id1_x|id2_x|id1_y|id2_y|distance|row_id|\n",
      "+---------+---------------+---------+---------------+-----+-----+\n",
      "|        _|              a|        a|              _|  9.1|    1|\n",
      "|        _|              a|        b|              _|  5.1|    1|\n",
      "|        _|              b|        a|              _|  7.1|    2|\n",
      "|        _|              b|        c|              _|  1.1|    2|\n",
      "|        _|              c|        d|              _|  8.1|    3|\n",
      "|        _|              c|        b|              _|  7.1|    3|\n",
      "+---------+---------------+---------+---------------+-----+-----+\n",
      "\n",
      "+---------+---------------+-----+\n",
      "|id1_x|id2_x|row_id|\n",
      "+---------+---------------+-----+\n",
      "|        _|              a|    1|\n",
      "|        _|              b|    2|\n",
      "|        _|              c|    3|\n",
      "+---------+---------------+-----+\n",
      "\n",
      "+---------+---------------+----+----+----+\n",
      "|id1_y|id2_y|   1|   2|   3|\n",
      "+---------+---------------+----+----+----+\n",
      "|        a|              _| 9.1| 7.1|null|\n",
      "|        b|              _| 5.1|null| 7.1|\n",
      "|        c|              _|null| 1.1|null|\n",
      "|        d|              _|null|null| 8.1|\n",
      "+---------+---------------+----+----+----+\n",
      "\n",
      "+---------+---------------+----+----+----+--------+\n",
      "|id1_y|id2_y|   1|   2|   3|assigned|\n",
      "+---------+---------------+----+----+----+--------+\n",
      "|        a|              _| 9.1| 7.1|null|       4|\n",
      "|        b|              _| 5.1|null| 7.1|       1|\n",
      "|        c|              _|null| 1.1|null|       2|\n",
      "|        d|              _|null|null| 8.1|       3|\n",
      "+---------+---------------+----+----+----+--------+\n",
      "\n",
      "['1', '2', '3']\n",
      "+---------+---------------+--------+-----+-----+--------+\n",
      "|id1_y|id2_y|assigned|row_id|distance|is_match|\n",
      "+---------+---------------+--------+-----+-----+--------+\n",
      "|        a|              _|       4|    1|  9.1|   false|\n",
      "|        a|              _|       4|    2|  7.1|   false|\n",
      "|        a|              _|       4|    3| null|   false|\n",
      "|        b|              _|       1|    1|  5.1|    true|\n",
      "|        b|              _|       1|    2| null|   false|\n",
      "|        b|              _|       1|    3|  7.1|   false|\n",
      "|        c|              _|       2|    1| null|   false|\n",
      "|        c|              _|       2|    2|  1.1|    true|\n",
      "|        c|              _|       2|    3| null|   false|\n",
      "|        d|              _|       3|    1| null|   false|\n",
      "|        d|              _|       3|    2| null|   false|\n",
      "|        d|              _|       3|    3|  8.1|    true|\n",
      "+---------+---------------+--------+-----+-----+--------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_1_assigned = emparejar(df_1)\n",
    "df_1_assigned.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n",
      "+---------+---------------+---------+---------------+-----+-----+\n",
      "|id1_x|id2_x|id1_y|id2_y|distance|row_id|\n",
      "+---------+---------------+---------+---------------+-----+-----+\n",
      "|        _|              a|        a|              _|  9.1|    1|\n",
      "|        _|              a|        b|              _|  5.1|    1|\n",
      "|        _|              a|        c|              _|  7.1|    1|\n",
      "|        _|              b|        a|              _|  1.1|    2|\n",
      "|        _|              b|        b|              _|  8.1|    2|\n",
      "|        _|              b|        c|              _|  7.1|    2|\n",
      "+---------+---------------+---------+---------------+-----+-----+\n",
      "\n",
      "+---------+---------------+-----+\n",
      "|id1_x|id2_x|row_id|\n",
      "+---------+---------------+-----+\n",
      "|        _|              a|    1|\n",
      "|        _|              b|    2|\n",
      "+---------+---------------+-----+\n",
      "\n",
      "+---------+---------------+---+---+\n",
      "|id1_y|id2_y|  1|  2|\n",
      "+---------+---------------+---+---+\n",
      "|        a|              _|9.1|1.1|\n",
      "|        b|              _|5.1|8.1|\n",
      "|        c|              _|7.1|7.1|\n",
      "+---------+---------------+---+---+\n",
      "\n",
      "+---------+---------------+---+---+--------+\n",
      "|id1_y|id2_y|  1|  2|assigned|\n",
      "+---------+---------------+---+---+--------+\n",
      "|        a|              _|9.1|1.1|       2|\n",
      "|        b|              _|5.1|8.1|       1|\n",
      "|        c|              _|7.1|7.1|       3|\n",
      "+---------+---------------+---+---+--------+\n",
      "\n",
      "['1', '2']\n",
      "+---------+---------------+--------+-----+-----+--------+\n",
      "|id1_y|id2_y|assigned|row_id|distance|is_match|\n",
      "+---------+---------------+--------+-----+-----+--------+\n",
      "|        a|              _|       2|    1|  9.1|   false|\n",
      "|        a|              _|       2|    2|  1.1|    true|\n",
      "|        b|              _|       1|    1|  5.1|    true|\n",
      "|        b|              _|       1|    2|  8.1|   false|\n",
      "|        c|              _|       3|    1|  7.1|   false|\n",
      "|        c|              _|       3|    2|  7.1|   false|\n",
      "+---------+---------------+--------+-----+-----+--------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_2_assigned = emparejar(df_2)\n",
    "df_2_assigned.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PySpark",
   "language": "python",
   "name": "spark-python"
  },
  "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.6.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}