Notebooks/Spark-Example-10-Working-With-Dataframes-Array.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import findspark\n",
    "findspark.init()\n",
    "from pyspark import SparkContext\n",
    "from pyspark.sql import SparkSession\n",
    "sc = SparkContext(\"local\")\n",
    "spark = SparkSession.builder.getOrCreate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import *\n",
    "from pyspark.sql.types import *\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+---------------+\n",
      "| id|       features|\n",
      "+---+---------------+\n",
      "|  1|[2.0, 2.0, 3.0]|\n",
      "|  1|[2.0, 3.0, 3.0]|\n",
      "|  2|[3.0, 2.0, 3.0]|\n",
      "|  2|[3.0, 3.0, 3.0]|\n",
      "+---+---------------+\n",
      "\n",
      "root\n",
      " |-- id: long (nullable = true)\n",
      " |-- features: array (nullable = true)\n",
      " |    |-- element: double (containsNull = true)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Let us use python Array\n",
    "data2 = [(1, ([2.0, 2.0, 3.0]),),\n",
    "         (1, ([2.0, 3.0, 3.0]),),\n",
    "         (2, ([3.0, 2.0, 3.0]),),\n",
    "         (2, ([3.0, 3.0, 3.0]),)]\n",
    "\n",
    "df = spark.createDataFrame(data2, [\"id\", \"features\"])\n",
    "\n",
    "df.show()\n",
    "df.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+---------------+\n",
      "| id| features_array|\n",
      "+---+---------------+\n",
      "|  1|[2.0, 2.0, 3.0]|\n",
      "|  1|[2.0, 3.0, 3.0]|\n",
      "|  2|[3.0, 2.0, 3.0]|\n",
      "|  2|[3.0, 3.0, 3.0]|\n",
      "+---+---------------+\n",
      "\n",
      "root\n",
      " |-- id: long (nullable = true)\n",
      " |-- features_array: array (nullable = true)\n",
      " |    |-- element: float (containsNull = false)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Convert the python Array type to Dataframe array Type. \n",
    "\n",
    "convert_Array_udf = udf(lambda x: x, ArrayType(FloatType(), containsNull=False))\n",
    "\n",
    "\n",
    "df = df.withColumn('features_array', convert_Array_udf('features')).drop('features')\n",
    "\n",
    "df.show()\n",
    "df.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+---------------+--------------------+\n",
      "| id| features_array|            multiply|\n",
      "+---+---------------+--------------------+\n",
      "|  1|[2.0, 2.0, 3.0]|[10.0, 4.0, 4.0, ...|\n",
      "|  1|[2.0, 3.0, 3.0]|[10.0, 4.0, 9.0, ...|\n",
      "|  2|[3.0, 2.0, 3.0]|[10.0, 9.0, 4.0, ...|\n",
      "|  2|[3.0, 3.0, 3.0]|[10.0, 9.0, 9.0, ...|\n",
      "+---+---------------+--------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# get the final dataframe for gradient descent\n",
    "\n",
    "step = 10.0 \n",
    "multiplyPlusStep_udf = udf(lambda x, y: [step]+np.multiply(x, y).tolist(), ArrayType(FloatType(), containsNull=False))\n",
    "\n",
    "df2=df.withColumn('multiply', multiplyPlusStep_udf('features_array', 'features_array'))\n",
    "\n",
    "df2.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+---------------+---------+\n",
      "| id| features_array|      bin|\n",
      "+---+---------------+---------+\n",
      "|  1|[2.0, 2.0, 3.0]|[0, 0, 1]|\n",
      "|  1|[2.0, 3.0, 3.0]|[0, 1, 1]|\n",
      "|  2|[3.0, 2.0, 3.0]|[1, 0, 1]|\n",
      "|  2|[3.0, 3.0, 3.0]|[1, 1, 1]|\n",
      "+---+---------------+---------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# convert to 1 and 0\n",
    "\n",
    "binary_udf = udf(lambda x: np.where(np.array(x) ==3.0, 1, 0).tolist(), ArrayType(IntegerType(), containsNull=False) ) \n",
    "\n",
    "df.withColumn('bin', binary_udf('features_array')).show()  \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+---------------+--------+\n",
      "| id| features_array|multiply|\n",
      "+---+---------------+--------+\n",
      "|  1|[2.0, 2.0, 3.0]|    27.0|\n",
      "|  1|[2.0, 3.0, 3.0]|    32.0|\n",
      "|  2|[3.0, 2.0, 3.0]|    32.0|\n",
      "|  2|[3.0, 3.0, 3.0]|    37.0|\n",
      "+---+---------------+--------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Sum the rows\n",
    "# Note: You need to convert the results back to float\n",
    "\n",
    "# Define a UDF\n",
    "sumRows_udf = udf(lambda x: float(np.sum(x)), FloatType())\n",
    "\n",
    "# Run the UDF\n",
    "df3=df2.withColumn('multiply', sumRows_udf('multiply'))\n",
    "\n",
    "df3.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------------+\n",
      "|sum               |\n",
      "+------------------+\n",
      "|[10.0, 10.0, 12.0]|\n",
      "+------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Sum Column-Wise\n",
    "\n",
    "from pyspark.sql import functions as F\n",
    "\n",
    "n = len(df.select('features_array').first()[0])\n",
    "\n",
    "resultDF = df.agg(F.array(*[F.sum(F.col(\"features_array\")[i]) for i in range(n)]).alias(\"sum\"))\n",
    "\n",
    "resultDF.show(truncate=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------+-------------+\n",
      "|dummyKey|sum(multiply)|\n",
      "+--------+-------------+\n",
      "|     1.0|        128.0|\n",
      "+--------+-------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Now, we want to sum up the column 'multuply' and get the results. \n",
    "\n",
    "df4 = df3.withColumn('dummyKey', lit(1.0)).groupBy('dummyKey').agg({'multiply': 'sum'})\n",
    "\n",
    "df4.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "spark",
   "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.10.4 | packaged by conda-forge | (main, Mar 30 2022, 08:38:02) [MSC v.1916 64 bit (AMD64)]"
  },
  "vscode": {
   "interpreter": {
    "hash": "9db6cbf0fd79f8e79653fe7b0c50b956ca6e525ee712295da3c66f75e4fe96ce"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}