What you will be able to do
- Add a new column or replace an existing one with withColumn, and know when to use a single select instead
- Rename columns with withColumnRenamed and recognise its silent no-op behaviour
- Drop columns by name or by Column expression, and predict the result after a join
- Split a string column into an array with split, choose a limit, and extract parts with getItem or split_part
Key concept
Column operations return a new DataFrame — withColumn, withColumnRenamed and drop never change the DataFrame you call them on. Each one returns a new DataFrame with a different schema, so you assign or chain the result, and every withColumn call adds another step to the query plan. The same holds when you shape columns inside a select: you name the result with alias (or AS in SQL), and split returns one array column that you then pick apart with getItem or split_part.
1.Adding or replacing a column with withColumn
Reshaping a table usually starts with its columns. In PySpark the main tool is withColumn(colName, col). It takes two arguments: the column name as a string, and a Column expression that computes the values. A new name adds a column. A name that already exists replaces that column, which is how you overwrite values (to clean or recompute them) without changing the list of columns. Either way you get back a new DataFrame. The original stays as it was.
df = spark.createDataFrame([(2, "Alice"), (5, "Bob")], schema=["age", "name"])
df.withColumn('age2', df.age + 2).show()
# +---+-----+----+
# |age| name|age2|
# +---+-----+----+
# | 2|Alice| 4|
# | 5| Bob| 7|
# +---+-----+----+The documentation advises against it. withColumn adds a projection to the plan internally, and calling it many times, for example in a loop, "can generate big plans which can cause performance issues and even StackOverflowException." The fix it gives is to build all the new columns in one select call instead of chaining many withColumn calls. Use withColumn for one or two columns. When you are adding many at once, use a single select.
Checkpoint 1 of 5· Check yourself
A job adds 50 columns by calling withColumn inside a for loop, and the driver fails with StackOverflowException. What fix does the withColumn documentation recommend?
Each withColumn call adds a projection, so a loop builds a very large plan. The documented fix is one select that creates all the columns at once.
“To avoid this, use select with multiple columns at once.”Source: docs.databricks.com
Sources1
2.Renaming columns with withColumnRenamed
Raw sources often use column names you would not keep: abbreviations, inconsistent casing, or names that clash after a join. withColumnRenamed(existing, new) takes the current name and the new name, both as strings, and returns a new DataFrame. There is one behaviour worth predicting before you read about it.
Renaming a column that doesn't exist is a no-op: the method returns the DataFrame unchanged and raises no error. So a typo in the existing argument won't fail. The old name simply stays. To rename several columns, chain the calls. Each call returns a DataFrame that the next call can rename further.
df.withColumnRenamed("non_existing", "new_name").show()
# +---+-----+
# |age| name|
# +---+-----+
# | 2|Alice|
# | 5| Bob|
# +---+-----+
df.withColumnRenamed("age", "age2").withColumnRenamed("name", "name2").show()
# +----+-----+
# |age2|name2|
# +----+-----+
# | 2|Alice|
# | 5| Bob|
# +----+-----+You can also name a column inside a select. In the DataFrame API that is .alias(...), and in SQL it is AS. The Databricks documentation uses exactly this to name the column that explode produces from an array.
df.select(explode("array_name").alias("column_name"))SELECT explode(array_name) AS column_name FROM table_nameCheckpoint 2 of 5· Check yourself
A job calls withColumnRenamed("custmer_id", "customer_id"), but the DataFrame's column is really named customer_id_raw, so "custmer_id" is not in the schema. What happens?
The documentation says the rename is a no-op when the schema doesn't contain the given name, so the typo passes silently and the old column name stays.
“This is a no-op if the schema doesn't contain the given column name.”Source: docs.databricks.com
3.Dropping columns by name or by Column
drop(*cols) returns a new DataFrame without the listed columns, and because it takes *cols you can list several columns in one call. Like renaming, it is a no-op when the schema doesn't contain a listed name. It accepts either a string name or a Column object, and the two are matched differently. A string is matched literally against column names. A Column is matched as an equivalent expression. The documentation states that drop(colName) and drop(col(colName)) have different semantics.
| Call | How the argument is matched |
|---|---|
| df.drop('age') | As a column name, treated literally without further interpretation |
| df.drop(df.age) | As a Column, matched to the equivalent expression |
| A name the schema doesn't contain | No-op: the DataFrame is returned unchanged |
The examples in this section use the DataFrame df below, with three rows of age and name. For a plain DataFrame with no join, dropping by string and dropping by Column give the same result.
df = spark.createDataFrame(
[(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"])
df.drop('age').show()
# +-----+
# | name|
# +-----+
# | Tom|
# |Alice|
# | Bob|
# +-----+
df.drop(df.age).show()
# +-----+
# | name|
# +-----+
# | Tom|
# |Alice|
# | Bob|
# +-----+The difference matters most after a join. Both inputs below have a name column, and both copies are still present after the join. Look at the output to see what dropping by string does to them. This df is the three-row DataFrame defined above (Tom, Alice, Bob), which is why the ages are 14 and 16. The documentation shows the join output only for the string form, so don't assume a result for drop(col(...)) after a join.
df2 = spark.createDataFrame([(80, "Tom"), (85, "Bob")], ["height", "name"])
df.join(df2, df.name == df2.name).drop('name').sort('age').show()
# +---+------+
# |age|height|
# +---+------+
# | 14| 80|
# | 16| 85|
# +---+------+Checkpoint 3 of 5· Check yourself
df(age, name) is joined to df2(height, name) on df.name == df2.name, then .drop('name') is called. According to the documented output, which columns remain?
A string argument is matched literally against column names, and the documented output keeps only age and height. Neither name column survives.
“When an input is a column name, it is treated literally without further interpretation.”Source: docs.databricks.com
Sources4
4.Splitting a string column with split
The last column operation divides one value into several. pyspark.sql.functions.split(str, pattern, limit) splits a string around matches of pattern and returns a column of type "array of separated strings". The pattern is not a plain delimiter: "The regex string should be a Java regular expression." So '[ABC]' is a character class that splits on A, B or C. In Databricks SQL, split(str, regex [, limit]) works the same way and returns ARRAY<STRING>. The optional limit controls how many times the pattern is applied.
| limit | Behaviour | Result |
|---|---|---|
| omitted (defaults to 0) | Regex applied as many times as possible | [one,two,three,] |
| 2 | Array length at most 2; the last entry holds the rest of the input | [one,twoBthreeC] |
| -1 | Regex applied as many times as possible | [one,two,three,] |
> SELECT split('oneAtwoBthreeC', '[ABC]', 2);
[one,twoBthreeC]
> SELECT split('oneAtwoBthreeC', '[ABC]', -1);
[one,two,three,]Checkpoint 4 of 5· Check yourself
What does SELECT split('oneAtwoBthreeC', '[ABC]', 2) return?
With limit 2 the array has at most two elements, and the last element holds everything after the first match.
“the resulting array's last entry will contain all input beyond the last matched regex.”Source: docs.databricks.com
split gives you one array column, not several new columns. To turn its pieces into separate columns, take individual elements. Column.getItem(key) returns an item from an array or map column (for an array, the key is an index, so df.l.getItem(0) is the first element). The documentation's example selects the first element of an array column and a value from a map column in one select.
df = spark.createDataFrame([([1, 2], {"key": "value"})], ["l", "d"])
df.select(df.l.getItem(0), df.d.getItem("key")).show()
# +----+------+
# |l[0]|d[key]|
# +----+------+
# | 1| value|
# +----+------+If you only need one part, split_part(src, delimiter, partNum) returns it directly. It "Splits str by delimiter and return requested part of the split (1-based)." Unlike getItem, it counts from 1. A negative partNum counts backwards from the end, an out-of-range partNum returns an empty string, and a partNum of 0 throws an error.
from pyspark.sql import functions as dbf
df = spark.createDataFrame([("11.12.13", ".", 3,)], ["a", "b", "c"])
df.select("*", dbf.split_part("a", "b", "c")).show()
df.select("*", dbf.split_part(df.a, df.b, dbf.lit(-2))).show()Checkpoint 5 of 5· Exam question
A data engineer is working with a Spark DataFrame `orders_df` that has columns `price` and `quantity`. They need to add a new column `total` equal to `price` multiplied by `quantity`, then reassign the result to `orders_df` for later steps in the same notebook. Which line of code correctly does this?
Correct answer: A — `orders_df = orders_df.withColumn("total", col("price") * col("quantity"))`
- A. `withColumn` returns a new DataFrame with the named column added when it does not already exist, or replaced when it does, so reassigning the result correctly produces the computed field. This is the standard way to derive a new column from an expression built out of existing columns.
- B. Spark DataFrames are immutable and do not support item assignment the way pandas DataFrames do, so this bracket-assignment syntax raises an error rather than mutating or creating a column. There is no `__setitem__` behavior defined on `DataFrame`.
- C. `withColumnRenamed` only accepts two string arguments, the existing column name and the new name, and does not accept a third expression argument, so this call fails with a signature error. It also renames a column rather than computing a new value from an expression.
- D. Passing the bare string `"price * quantity"` to `select` treats it as a plain string literal column name rather than a SQL expression, so this does not compute the product; the string would need to be wrapped in `expr(...)` to be parsed as an expression.
Exam traps
Each one states something that sounds right. Open it to see what is actually true.
1.Calling withColumn in a loop is the normal way to add many columns, and the only cost is a little extra code.Why is that wrong?
Each call adds a projection to the plan. Many calls can make the plan large enough to cause performance problems or even a StackOverflowException, so add many columns in one select.
Covered in Adding or replacing a column with withColumn
2.withColumnRenamed raises an error if the existing column name is misspelled.Why is that wrong?
Renaming a missing column is a no-op. The DataFrame comes back unchanged and no error is raised.
Covered in Renaming columns with withColumnRenamed
3.split's pattern is a literal delimiter string, so '[ABC]' splits on the text '[ABC]'.Why is that wrong?
The pattern is a Java regular expression. '[ABC]' is a character class that splits on A, B or C.
Covered in Splitting a string column with split
Sources
Every claim above is drawn from one of these pages, quoted as it was written on the date shown.
- 1.
“Returns a new DataFrame by adding a column or replacing the existing column that has the same name.”
↩︎ Adding or replacing a column with withColumn“can generate big plans which can cause performance issues and even StackOverflowException.”
↩︎ Adding or replacing a column with withColumn“Returns a new DataFrame by adding a column or replacing the existing column that has the same name.”
↩︎ Key concept“calling it multiple times, for instance, via loops in order to add multiple columns can generate big plans”
↩︎ Exam trap 1“To avoid this, use select with multiple columns at once.”
↩︎ Checkpoint - 2.https://docs.databricks.com/aws/en/pyspark/reference/classes/dataframe/withColumnRenamedOfficial docs
“Returns a new DataFrame by renaming an existing column.”
↩︎ Renaming columns with withColumnRenamed“This is a no-op if the schema doesn't contain the given column name.”
↩︎ Renaming columns with withColumnRenamed“This is a no-op if the schema doesn't contain the given column name.”
↩︎ Exam trap 2 - 3.
“SELECT explode(array_name) AS column_name FROM table_name”
↩︎ Renaming columns with withColumnRenamed - 4.
“So dropping a column by its name drop(colName) has a different semantic with directly dropping the column drop(col(colName)).”
↩︎ Dropping columns by name or by Column“This is a no-op if the schema doesn't contain the given column name(s).”
↩︎ Dropping columns by name or by Column“When an input is a column name, it is treated literally without further interpretation.”
↩︎ Checkpoint - 5.
“pyspark.sql.Column: array of separated strings.”
↩︎ Splitting a string column with split“The regex string should be a Java regular expression.”
↩︎ Exam trap 3 - 6.
“Get an item from an array or map column.”
↩︎ Splitting a string column with split - 7.
“Splits str by delimiter and return requested part of the split (1-based).”
↩︎ Splitting a string column with split“If partNum is 0, throws an error.”
↩︎ Splitting a string column with split - 8.
“limit: An optional INTEGER expression defaulting to 0 (no limit).”
↩︎ Splitting a string column with split“the resulting array's last entry will contain all input beyond the last matched regex.”
↩︎ Checkpoint