When working with large PySpark DataFrames, there are scenarios where you need to split the data into smaller, equally-sized chunks - for example, to process each chunk in parallel, batch data for API calls, write data in manageable partitions, or distribute workloads across workers. If the operation on each chunk is independent of other rows, splitting and processing in parallel can significantly improve efficiency.
In this guide, you will learn how to split a PySpark DataFrame into equal parts using limit() and subtract(), a more reliable approach using monotonically_increasing_id(), and how to process each split independently before combining results.
Setting Up the Example DataFrame
from pyspark.sql import SparkSession
spark = SparkSession.builder.appName("split_demo").getOrCreate()
columns = ["Brand", "Product"]
data = [
("HP", "Laptop"),
("Lenovo", "Mouse"),
("Dell", "Keyboard"),
("Samsung", "Monitor"),
("MSI", "Graphics Card"),
("Asus", "Motherboard"),
("Gigabyte", "Motherboard"),
("Zebronics", "Cabinet"),
("Adata", "RAM"),
("Transcend", "SSD"),
("Kingston", "HDD"),
("Toshiba", "DVD Writer"),
]
df = spark.createDataFrame(data=data, schema=columns)
df.show()Output:
+---------+-------------+
| Brand| Product|
+---------+-------------+
| HP| Laptop|
| Lenovo| Mouse|
| Dell| Keyboard|
| Samsung| Monitor|
| MSI|Graphics Card|
| Asus| Motherboard|
| Gigabyte| Motherboard|
|Zebronics| Cabinet|
| Adata| RAM|
|Transcend| SSD|
| Kingston| HDD|
| Toshiba| DVD Writer|
+---------+-------------+The DataFrame has 12 rows and 2 columns. We will split it into equal chunks.
Method 1: Using limit() and subtract()
The simplest approach uses limit() to take the first n rows and subtract() to remove those rows from the remaining data, repeating until all rows are distributed.
def split_dataframe_limit(df, n_splits):
"""Split a DataFrame into n equal parts using limit() and subtract()."""
each_len = df.count() // n_splits
remaining_df = df
splits = []
for i in range(n_splits):
if i == n_splits - 1:
splits.append(remaining_df)
else:
chunk = remaining_df.limit(each_len)
remaining_df = remaining_df.subtract(chunk)
splits.append(chunk)
return splits
chunks = split_dataframe_limit(df, n_splits=4)
for i, chunk in enumerate(chunks):
print(f"Chunk {i + 1} ({chunk.count()} rows):")
chunk.show(truncate=False)Output:
Chunk 1 (3 rows):
+---------+-------------+
|Brand |Product |
+---------+-------------+
|HP |Laptop |
|Dell |Keyboard |
|Lenovo |Mouse |
+---------+-------------+
Chunk 2 (3 rows):
+-------+-------------+
|Brand |Product |
+-------+-------------+
|Samsung|Monitor |
|Asus |Motherboard |
|MSI |Graphics Card|
+-------+-------------+
Chunk 3 (3 rows):
+---------+-----------+
|Brand |Product |
+---------+-----------+
|Gigabyte |Motherboard|
|Zebronics|Cabinet |
|Adata |RAM |
+---------+-----------+
Chunk 4 (3 rows):
+---------+----------+
|Brand |Product |
+---------+----------+
|Transcend|SSD |
|Kingston |HDD |
|Toshiba |DVD Writer|
+---------+----------+Method 2: Using monotonically_increasing_id() (Recommended)
A more reliable and efficient approach assigns a sequential row index to each row, then uses integer division to determine which split each row belongs to.
from pyspark.sql.functions import monotonically_increasing_id, spark_partition_id
from pyspark.sql.window import Window
from pyspark.sql.functions import row_number
def split_dataframe(df, n_splits):
"""Split a DataFrame into n roughly equal parts using row numbering."""
df_with_index = df.withColumn(
"row_idx",
row_number().over(Window.orderBy(monotonically_increasing_id())) - 1
)
total_rows = df_with_index.count()
each_len = total_rows // n_splits
splits = []
for i in range(n_splits):
start = i * each_len
if i == n_splits - 1:
chunk = df_with_index.filter(
df_with_index.row_idx >= start
)
else:
end = start + each_len
chunk = df_with_index.filter(
(df_with_index.row_idx >= start) &
(df_with_index.row_idx < end)
)
splits.append(chunk.drop("row_idx"))
return splits
chunks = split_dataframe(df, n_splits=4)
for i, chunk in enumerate(chunks):
print(f"Chunk {i + 1} ({chunk.count()} rows):")
chunk.show(truncate=False)Output:
Chunk 1 (3 rows):
+------+--------+
|Brand |Product |
+------+--------+
|HP |Laptop |
|Lenovo|Mouse |
|Dell |Keyboard|
+------+--------+
Chunk 2 (3 rows):
+-------+-------------+
|Brand |Product |
+-------+-------------+
|Samsung|Monitor |
|MSI |Graphics Card|
|Asus |Motherboard |
+-------+-------------+
Chunk 3 (3 rows):
+---------+-----------+
|Brand |Product |
+---------+-----------+
|Gigabyte |Motherboard|
|Zebronics|Cabinet |
|Adata |RAM |
+---------+-----------+
Chunk 4 (3 rows):
+---------+----------+
|Brand |Product |
+---------+----------+
|Transcend|SSD |
|Kingston |HDD |
|Toshiba |DVD Writer|
+---------+----------+Handling Uneven Splits
When the total row count is not evenly divisible by the number of splits, you need to decide what happens to the remainder rows. Both methods above assign remaining rows to the last chunk.
chunks = split_dataframe(df, n_splits=5)
for i, chunk in enumerate(chunks):
print(f"Chunk {i + 1}: {chunk.count()} rows")Output:
Chunk 1: 2 rows
Chunk 2: 2 rows
Chunk 3: 2 rows
Chunk 4: 2 rows
Chunk 5: 4 rowsFor a more balanced distribution of remainder rows, you can modify the function to spread extra rows across the first few chunks:
def split_dataframe_balanced(df, n_splits):
"""Split DataFrame into n parts with balanced distribution of remainder rows."""
df_with_index = df.withColumn(
"row_idx",
row_number().over(Window.orderBy(monotonically_increasing_id())) - 1
)
total_rows = df_with_index.count()
base_size = total_rows // n_splits
remainder = total_rows % n_splits
splits = []
start = 0
for i in range(n_splits):
chunk_size = base_size + (1 if i < remainder else 0)
end = start + chunk_size
chunk = df_with_index.filter(
(df_with_index.row_idx >= start) &
(df_with_index.row_idx < end)
).drop("row_idx")
splits.append(chunk)
start = end
return splits
chunks = split_dataframe_balanced(df, n_splits=5)
for i, chunk in enumerate(chunks):
print(f"Chunk {i + 1}: {chunk.count()} rows")Output:
Chunk 1: 3 rows
Chunk 2: 3 rows
Chunk 3: 2 rows
Chunk 4: 2 rows
Chunk 5: 2 rowsProcessing Each Split and Combining Results
A common pattern is to split a DataFrame, apply a transformation to each chunk independently, and then combine the results using union().
from pyspark.sql.functions import concat, col, lit
from pyspark.sql.types import StructType, StructField, StringType
chunks = split_dataframe(df, n_splits=4)
def transform_chunk(chunk_df):
"""Concatenate Brand and Product into a single column."""
return chunk_df.select(
concat(col("Brand"), lit(" - "), col("Product")).alias("Brand_Product")
)
result_df = None
for i, chunk in enumerate(chunks):
transformed = transform_chunk(chunk)
print(f"Processed Chunk {i + 1}:")
transformed.show(truncate=False)
if result_df is None:
result_df = transformed
else:
result_df = result_df.union(transformed)
print("Combined Result:")
result_df.show(truncate=False)Output:
Processed Chunk 1:
+-------------+
|Brand_Product|
+-------------+
|HP - Laptop |
|Lenovo - Mouse|
|Dell - Keyboard|
+-------------+
...
Combined Result:
+------------------------+
|Brand_Product |
+------------------------+
|HP - Laptop |
|Lenovo - Mouse |
|Dell - Keyboard |
|Samsung - Monitor |
|MSI - Graphics Card |
|Asus - Motherboard |
|Gigabyte - Motherboard |
|Zebronics - Cabinet |
|Adata - RAM |
|Transcend - SSD |
|Kingston - HDD |
|Toshiba - DVD Writer |
+------------------------+Alternative: Using randomSplit()
PySpark provides a built-in randomSplit() method that splits a DataFrame by approximate proportions. While not guaranteed to produce exact split sizes, it is simpler and avoids the overhead of row numbering.
chunks = df.randomSplit([0.25, 0.25, 0.25, 0.25], seed=42)
for i, chunk in enumerate(chunks):
print(f"Chunk {i + 1}: {chunk.count()} rows")Output (approximate):
Chunk 1: 2 rows
Chunk 2: 4 rows
Chunk 3: 3 rows
Chunk 4: 3 rowsComparison of Methods
| Method | Exact Splits | Deterministic | Performance | Handles Duplicates |
|---|---|---|---|---|
limit() + subtract() |
✅ Yes | ❌ No | ❌ Slow (shuffles) | ❌ Problematic |
row_number() + filter() |
✅ Yes | ✅ Yes | ✅ Good | ✅ Yes |
randomSplit() |
❌ Approximate | ✅ With seed | ✅ Good | ✅ Yes |
Summary
Splitting a PySpark DataFrame into equal parts is useful for batch processing, parallel execution, and managing large datasets. Key takeaways:
- Use
row_number()withmonotonically_increasing_id()for reliable, deterministic, and exact splits - this is the recommended approach. - Use
randomSplit()when approximate split sizes are acceptable and simplicity is preferred. - Avoid
limit()+subtract()in production due to performance overhead and non-deterministic behavior. - Handle remainder rows by either assigning them to the last chunk or distributing them across the first few chunks for balanced splits.
- Use
union()orfunctools.reduceto combine processed chunks back into a single DataFrame.