add pyspark select vs withColumn

Author: khuyentran1401

Chapter5/spark.ipynb

@@ -116,6 +116,163 @@
     "In this example, PySpark performs a shuffle join behind the scenes to combine the two DataFrames. The process involves partitioning the data based on the join key (\"id\"), shuffling the partitions across the worker nodes, performing local joins on each worker node, and finally merging the results."
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "549e0e0a",
+   "metadata": {},
+   "source": [
+    "### PySpark DataFrame Transformations: select vs withColumn"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "322859d8",
+   "metadata": {
+    "tags": [
+     "hide-cell"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "!pip install 'pyspark[sql]'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "edbf287e",
+   "metadata": {
+    "tags": [
+     "hide-cell"
+    ]
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Setting default log level to \"WARN\".\n",
+      "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n",
+      "24/09/02 06:01:00 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n"
+     ]
+    }
+   ],
+   "source": [
+    "from pyspark.sql import SparkSession\n",
+    "\n",
+    "spark = SparkSession.builder.getOrCreate()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0b1dbabe",
+   "metadata": {},
+   "source": [
+    "PySpark's `select` and `withColumn` both can be used to add or modify existing columns. However, their behavior are different.\n",
+    "\n",
+    "To demonstrate this, let's start with creating a sample DataFrame:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "9561afc6",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "+-----+---+-------------+\n",
+      "| name|age|         city|\n",
+      "+-----+---+-------------+\n",
+      "|Alice| 28|     New York|\n",
+      "|  Bob| 35|San Francisco|\n",
+      "+-----+---+-------------+\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from pyspark.sql.functions import col, upper\n",
+    "\n",
+    "\n",
+    "data = [\n",
+    "    (\"Alice\", 28, \"New York\"),\n",
+    "    (\"Bob\", 35, \"San Francisco\"),\n",
+    "]\n",
+    "df = spark.createDataFrame(data, [\"name\", \"age\", \"city\"])\n",
+    "df.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "07eba50a",
+   "metadata": {},
+   "source": [
+    "\n",
+    "`select` only keeps specified columns."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "278dac0c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "+-------------+\n",
+      "|   upper_city|\n",
+      "+-------------+\n",
+      "|     NEW YORK|\n",
+      "|SAN FRANCISCO|\n",
+      "+-------------+\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "df_select = df.select(upper(col(\"city\")).alias(\"upper_city\"))\n",
+    "df_select.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3f1dec31",
+   "metadata": {},
+   "source": [
+    "`withColumn` retains all original columns plus the new/modified one."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "3b0a9103",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "+-----+---+-------------+-------------+\n",
+      "| name|age|         city|   upper_city|\n",
+      "+-----+---+-------------+-------------+\n",
+      "|Alice| 28|     New York|     NEW YORK|\n",
+      "|  Bob| 35|San Francisco|SAN FRANCISCO|\n",
+      "+-----+---+-------------+-------------+\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "df_withColumn = df.withColumn('upper_city', upper(col('city')))\n",
+    "df_withColumn.show()"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "d8dc8623",