Adding and renaming columns in Apache Spark - Time & Space Complexity

Time Complexity: Adding and renaming columns

O(n)

Understanding Time Complexity

We want to understand how the time needed to add or rename columns in a Spark DataFrame changes as the data grows.

How does the work increase when the number of rows or columns gets bigger?

Scenario Under Consideration

Analyze the time complexity of the following code snippet.

from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()
df = spark.createDataFrame([(1, 2), (3, 4)], ['A', 'B'])

# Add a new column 'C' as sum of 'A' and 'B'
df = df.withColumn('C', df['A'] + df['B'])

# Rename column 'A' to 'X'
df = df.withColumnRenamed('A', 'X')

This code adds a new column by combining existing columns and then renames one column.

Identify Repeating Operations

Identify the loops, recursion, array traversals that repeat.

• Primary operation: Spark applies the column addition across all rows; renaming is a metadata change.
• How many times: Once per row for the new column calculation; renaming is a metadata change.

How Execution Grows With Input

Adding a column means Spark processes each row to compute the new value, so work grows with rows.

Input Size (n rows)	Approx. Operations
10	10 operations to add column
100	100 operations to add column
1000	1000 operations to add column

Renaming a column does not depend on row count; it stays constant.

Pattern observation: Adding columns scales linearly with rows; renaming is constant time.

Final Time Complexity

Time Complexity: O(n)

This means the time to add a column grows directly with the number of rows, while renaming is very fast and does not grow with data size.

Common Mistake

[X] Wrong: "Renaming a column takes as long as adding a column because both change the DataFrame."

[OK] Correct: Renaming only changes the column name in metadata, not the data itself, so it is much faster and does not depend on the number of rows.

Interview Connect

Understanding how operations scale with data size helps you write efficient Spark code and explain your reasoning clearly in conversations.

Self-Check

"What if we added multiple new columns at once? How would the time complexity change?"

Practice

1. What does the withColumn method do in Apache Spark DataFrame?

easy

A. Adds a new column or updates an existing column in the DataFrame

B. Deletes a column from the DataFrame

C. Renames a column in the DataFrame

D. Filters rows based on a condition

Solution

1. Step 1: Understand the purpose of withColumn

   The withColumn method is used to add a new column or replace an existing column in a Spark DataFrame.

2. Step 2: Compare with other methods

   Unlike withColumnRenamed which renames columns, withColumn changes or adds columns with new values or expressions.

3. Final Answer:

   Adds a new column or updates an existing column in the DataFrame -> Option A

4. Quick Check:

   withColumn = Add or update column [OK]

Hint: Use withColumn to add or change columns fast [OK]

Common Mistakes:

• Confusing withColumn with withColumnRenamed
• Thinking it deletes columns
• Assuming it filters rows

2. Which of the following is the correct syntax to rename a column named oldName to newName in a Spark DataFrame df?

easy

A. df.rename("oldName", "newName")

B. df.withColumn("oldName", "newName")

C. df.renameColumn("oldName", "newName")

D. df.withColumnRenamed("oldName", "newName")

Solution

1. Step 1: Identify the correct method for renaming

   The method to rename a column in Spark DataFrame is withColumnRenamed.

2. Step 2: Check syntax correctness

   The syntax is df.withColumnRenamed("oldName", "newName"). Other options are invalid methods or misuse of withColumn.

3. Final Answer:

   df.withColumnRenamed("oldName", "newName") -> Option D

4. Quick Check:

   Rename column = withColumnRenamed [OK]

Hint: Rename columns only with withColumnRenamed method [OK]

Common Mistakes:

• Using withColumn instead of withColumnRenamed
• Trying non-existent methods like renameColumn
• Passing arguments in wrong order

3. Given the following Spark code:

from pyspark.sql.functions import lit

df = spark.createDataFrame([(1, "Alice")], ["id", "name"])
df2 = df.withColumn("age", lit(30)).withColumnRenamed("name", "first_name")
df2.show()

What will be the output?

medium

A. +---+----------+---+
| id|first_name|age|
+---+----------+---+
| 1| Alice| 30|
+---+----------+---+

B. +---+-----+---+
| id| name|age|
+---+-----+---+
| 1|Alice| 30|
+---+-----+---+

C. +---+----------+
| id|first_name|
+---+----------+
| 1| Alice|
+---+----------+

D. Error: withColumnRenamed cannot be chained after withColumn

Solution

1. Step 1: Add new column 'age' with value 30

   The withColumn("age", lit(30)) adds a new column 'age' with constant value 30.

2. Step 2: Rename column 'name' to 'first_name'

   The withColumnRenamed("name", "first_name") renames the 'name' column to 'first_name'.

3. Step 3: Show the resulting DataFrame

   The DataFrame now has columns: 'id', 'first_name', and 'age' with one row (1, Alice, 30).

4. Final Answer:

   +---+----------+---+ | id|first_name|age| +---+----------+---+ | 1| Alice| 30| +---+----------+---+ -> Option A

5. Quick Check:

   withColumn + withColumnRenamed = added column + renamed column [OK]

Hint: Chain withColumn then withColumnRenamed to add and rename [OK]

Common Mistakes:

• Assuming withColumnRenamed cannot be chained
• Forgetting to import lit for constant values
• Expecting original column names unchanged

4. Consider this code snippet:

df = spark.createDataFrame([(1, "John")], ["id", "name"])
df2 = df.withColumnRenamed("Name", "first_name")
df2.show()

What is the issue and how to fix it?

medium

A. withColumnRenamed cannot rename columns; use withColumn instead

B. Column 'Name' does not exist; use correct case 'name' in withColumnRenamed

C. Syntax error: missing parentheses in withColumnRenamed

D. No issue; code runs correctly and renames column

Solution

1. Step 1: Check column names in DataFrame

   The DataFrame has columns 'id' and 'name' (lowercase). The code tries to rename 'Name' (uppercase N) which does not exist.

2. Step 2: Fix the column name case

   Change withColumnRenamed("Name", "first_name") to withColumnRenamed("name", "first_name") to match exact column name.

3. Final Answer:

   Column 'Name' does not exist; use correct case 'name' in withColumnRenamed -> Option B

4. Quick Check:

   Column names are case-sensitive in Spark [OK]

Hint: Match exact column name case when renaming [OK]

Common Mistakes:

• Ignoring case sensitivity of column names
• Using withColumn instead of withColumnRenamed to rename
• Assuming code runs without error despite wrong column name

5. You have a DataFrame df with columns ['id', 'score']. You want to add a new column grade based on score: if score ≥ 90 then 'A', else 'B'. Then rename score to exam_score. Which code correctly does this?

hard

A. df2 = df.withColumn('grade', when(df.score >= 90, 'A').otherwise('B')).renameColumn('score', 'exam_score')

B. df2 = df.withColumnRenamed('score', 'exam_score').withColumn('grade', when(df.score >= 90, 'A').otherwise('B'))

C. from pyspark.sql.functions import when

df2 = df.withColumn('grade', when(df.score >= 90, 'A').otherwise('B')).withColumnRenamed('score', 'exam_score')

D. df2 = df.withColumn('grade', when(df.score > 90, 'A').otherwise('B')).withColumnRenamed('score', 'exam_score')

Solution

1. Step 1: Add 'grade' column with condition on 'score'

   Use when to assign 'A' if score ≥ 90, else 'B'. This is done with withColumn('grade', when(df.score >= 90, 'A').otherwise('B')).

2. Step 2: Rename 'score' column to 'exam_score'

   Use withColumnRenamed('score', 'exam_score') to rename the column after adding 'grade'. The order matters to keep references correct.

3. Step 3: Check other options

   df2 = df.withColumnRenamed('score', 'exam_score').withColumn('grade', when(df.score >= 90, 'A').otherwise('B')) uses old column name after renaming, causing error. df2 = df.withColumn('grade', when(df.score >= 90, 'A').otherwise('B')).renameColumn('score', 'exam_score') uses non-existent renameColumn. df2 = df.withColumn('grade', when(df.score > 90, 'A').otherwise('B')).withColumnRenamed('score', 'exam_score') uses > 90 instead of ≥ 90, which changes logic.

4. Final Answer:

   from pyspark.sql.functions import when df2 = df.withColumn('grade', when(df.score >= 90, 'A').otherwise('B')).withColumnRenamed('score', 'exam_score') -> Option C

5. Quick Check:

   Add column with when, then rename column [OK]

Hint: Add conditional column first, then rename original column [OK]

Common Mistakes:

• Renaming column before adding new column using old name
• Using wrong comparison operator (> instead of >=)
• Using non-existent renameColumn method